com.github.ddth.thriftpool.ThriftClientPool Maven / Gradle / Ivy
package com.github.ddth.thriftpool;
import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.lang.reflect.Proxy;
import java.util.Random;
import java.util.Set;
import java.util.UUID;
import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.reflect.ConstructorUtils;
import org.apache.commons.pool2.BasePooledObjectFactory;
import org.apache.commons.pool2.ObjectPool;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;
import org.apache.commons.pool2.impl.GenericObjectPool;
import org.apache.thrift.TServiceClient;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.google.common.collect.Sets;
/**
* Pool of Thrift clients.
*
* @author Thanh Nguyen
*
* @param
* Thrift client class
* @param
* Thrift client interface
* @since 0.1.0
*/
public class ThriftClientPool {
private final Logger LOGGER = LoggerFactory.getLogger(ThriftClientPool.class);
private Class clientClass;
private Class clientInterface;
private PoolConfig poolConfig;
private RetryPolicy retryPolicy;
private ObjectPool thriftClientPool;
private ITProtocolFactory tprotocolFactory;
public ThriftClientPool() {
// EMPTY
}
public ThriftClientPool(Class clientClass, Class clientInterface,
ITProtocolFactory tprotocolFactory) {
this.clientClass = clientClass;
this.clientInterface = clientInterface;
this.tprotocolFactory = tprotocolFactory;
}
public ThriftClientPool(Class clientClass, Class clientInterface,
ITProtocolFactory tprotocolFactory, PoolConfig poolConfig) {
this.clientClass = clientClass;
this.clientInterface = clientInterface;
this.tprotocolFactory = tprotocolFactory;
this.poolConfig = poolConfig;
}
public ThriftClientPool(Class clientClass, Class clientInterface,
ITProtocolFactory tprotocolFactory, RetryPolicy retryPolicy) {
this.clientClass = clientClass;
this.clientInterface = clientInterface;
this.tprotocolFactory = tprotocolFactory;
this.retryPolicy = retryPolicy;
}
public ThriftClientPool(Class clientClass, Class clientInterface,
ITProtocolFactory tprotocolFactory, PoolConfig poolConfig, RetryPolicy retryPolicy) {
this.clientClass = clientClass;
this.clientInterface = clientInterface;
this.poolConfig = poolConfig;
this.tprotocolFactory = tprotocolFactory;
this.retryPolicy = retryPolicy;
}
/*----------------------------------------------------------------------*/
public Class getClientClass() {
return clientClass;
}
public ThriftClientPool setClientClass(Class clientClass) {
this.clientClass = clientClass;
return this;
}
public Class getClientInterface() {
return clientInterface;
}
public ThriftClientPool setClientInterface(Class clientInterface) {
this.clientInterface = clientInterface;
return this;
}
public PoolConfig getPoolConfig() {
return poolConfig;
}
public ThriftClientPool setPoolConfig(PoolConfig poolConfig) {
this.poolConfig = poolConfig;
return this;
}
public RetryPolicy getRetryPolicy() {
return retryPolicy;
}
public ThriftClientPool setRetryPolicy(RetryPolicy retryPolicy) {
this.retryPolicy = retryPolicy;
return this;
}
public ITProtocolFactory getTProtocolFactory() {
return tprotocolFactory;
}
public ThriftClientPool setTProtocolFactory(ITProtocolFactory tprotocolFactory) {
this.tprotocolFactory = tprotocolFactory;
return this;
}
synchronized public ThriftClientPool init() {
if (thriftClientPool == null) {
if (tprotocolFactory == null) {
throw new IllegalStateException("No ITProtocolFactory instance found!");
}
if (retryPolicy == null) {
retryPolicy = RetryPolicy.DEFAULT;
}
ThriftClientFactory factory = new ThriftClientFactory();
GenericObjectPool pool = new GenericObjectPool(factory);
pool.setBlockWhenExhausted(true);
pool.setTestOnReturn(false);
int maxActive = poolConfig != null ? poolConfig.getMaxActive()
: PoolConfig.DEFAULT_MAX_ACTIVE;
long maxWaitTime = poolConfig != null ? poolConfig.getMaxWaitTime()
: PoolConfig.DEFAULT_MAX_WAIT_TIME;
int maxIdle = poolConfig != null ? poolConfig.getMaxIdle()
: PoolConfig.DEFAULT_MAX_IDLE;
int minIdle = poolConfig != null ? poolConfig.getMinIdle()
: PoolConfig.DEFAULT_MIN_IDLE;
pool.setMaxTotal(maxActive);
pool.setMaxIdle(maxIdle);
pool.setMinIdle(minIdle);
pool.setMaxWaitMillis(maxWaitTime);
pool.setTestOnBorrow(poolConfig != null ? poolConfig.isTestOnBorrow() : false);
pool.setTestOnCreate(poolConfig != null ? poolConfig.isTestOnCreate() : false);
pool.setTestWhileIdle(poolConfig != null ? poolConfig.isTestWhileIdle() : false);
pool.setTimeBetweenEvictionRunsMillis(10000);
this.thriftClientPool = pool;
}
return this;
}
synchronized public void destroy() {
if (thriftClientPool != null) {
try {
thriftClientPool.close();
} finally {
thriftClientPool = null;
}
}
}
/*----------------------------------------------------------------------*/
/**
* Obtains a Thrift client object from pool.
*
* @return
* @throws Exception
*/
public I borrowObject() throws Exception {
return thriftClientPool.borrowObject();
}
/**
* Returns a borrowed Thrift client object back to pool.
*
* @param borrowedClient
* @throws Exception
*/
public void returnObject(I borrowedClient) throws Exception {
thriftClientPool.returnObject(borrowedClient);
}
/*----------------------------------------------------------------------*/
private final class ThriftClientFactory extends BasePooledObjectFactory {
@SuppressWarnings("unchecked")
@Override
public I create() throws Exception {
Object proxyObj = Proxy.newProxyInstance(clientInterface.getClassLoader(),
new Class[] { clientInterface },
new ReconnectingClientProxy(retryPolicy.clone()));
return (I) proxyObj;
}
@Override
public PooledObject wrap(I obj) {
return new DefaultPooledObject(obj);
}
/**
* {@inheritDoc}
*/
@SuppressWarnings("unchecked")
@Override
public void destroyObject(PooledObject pooledObj) throws Exception {
I obj = pooledObj.getObject();
if (Proxy.isProxyClass(obj.getClass())) {
InvocationHandler iv = Proxy.getInvocationHandler(obj);
if (iv instanceof ThriftClientPool.ReconnectingClientProxy) {
((ReconnectingClientProxy) iv).destroy();
}
}
}
}
/*----------------------------------------------------------------------*/
private static final Set RESTARTABLE_CAUSES = Sets.newHashSet(
TTransportException.NOT_OPEN, TTransportException.END_OF_FILE,
TTransportException.TIMED_OUT, TTransportException.UNKNOWN);
private Random random = new Random(System.currentTimeMillis());
/**
* Helper proxy class. Attempts to call method on proxy object wrapped in
* try/catch. If it fails, it attempts a reconnect and tries the method
* again.
*
*
* Credit: http://blog.liveramp.com/2014/04/10/reconnecting-thrift-client/
*
*
* @param
*/
private final class ReconnectingClientProxy implements InvocationHandler {
private RetryPolicy retryPolicy;
private UUID id = UUID.randomUUID();
private T clientObj;
public ReconnectingClientProxy(RetryPolicy retryPolicy) {
this.retryPolicy = retryPolicy;
}
/**
* {@inheritDoc}
*/
@SuppressWarnings("unchecked")
@Override
public boolean equals(Object obj) {
ReconnectingClientProxy other = (ReconnectingClientProxy) obj;
return id.equals(other.id);
}
/**
* {@inheritDoc}
*/
@Override
public int hashCode() {
return id.hashCode();
}
public void destroy() {
if (clientObj != null) {
try {
clientObj.getInputProtocol().getTransport().close();
} catch (Exception e) {
}
}
clientObj = null;
}
/**
* Creates a new thrift client object.
*
* @param serverIndexHash
* @return
* @throws Exception
*/
private T newClientObj(int serverIndexHash) throws Exception {
TProtocol protocol = tprotocolFactory.create(serverIndexHash);
T clientObj = ConstructorUtils.invokeConstructor(clientClass, protocol);
return clientObj;
}
private T getClientId(boolean renew, int serverIndexHash) throws Exception {
if (clientObj == null || renew) {
if (clientObj != null) {
try {
clientObj.getInputProtocol().getTransport().close();
} catch (Exception e) {
}
}
clientObj = newClientObj(serverIndexHash);
}
return clientObj;
}
/**
* {@inheritDoc}
*/
@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {
if (StringUtils.equals("hashCode", method.getName())) {
return hashCode();
}
if (StringUtils.equals("equals", method.getName())) {
return equals(args[0]);
}
if (StringUtils.equals("toString", method.getName())) {
return toString();
}
retryPolicy.reset();
return invokeWithRetries(proxy, method, args);
}
private Object invokeWithRetries(Object proxy, Method method, Object[] args)
throws Throwable {
boolean hasError = false;
while (!retryPolicy.exceedsMaxRetries()) {
int serverIndexHash = 0;
int numServer = tprotocolFactory.getNumServers();
switch (retryPolicy.getRetryType()) {
case FAILOVER:
serverIndexHash = retryPolicy.getCounter();
break;
case ROUND_ROBIN:
if (retryPolicy.getCounter() == 0) {
serverIndexHash = random.nextInt(Short.MAX_VALUE);
if (numServer > 1) {
serverIndexHash = serverIndexHash % numServer;
}
} else {
serverIndexHash = retryPolicy.getLastServerIndexHash() + 1;
}
retryPolicy.setLastServerIndexHash(serverIndexHash);
break;
case RANDOM_FAILOVER:
if (retryPolicy.getCounter() == 0 || numServer < 2) {
serverIndexHash = 0;
} else {
serverIndexHash = 1 + (random.nextInt(Short.MAX_VALUE) % (numServer - 1));
}
break;
case RANDOM:
default:
serverIndexHash = random.nextInt(Short.MAX_VALUE);
break;
}
try {
T clientObj = getClientId(hasError, serverIndexHash);
return method.invoke(clientObj, args);
} catch (InvocationTargetException e) {
hasError = true;
Throwable target = e.getTargetException();
if (target instanceof TTransportException) {
TTransportException cause = (TTransportException) target;
if (RESTARTABLE_CAUSES.contains(cause.getType())) {
LOGGER.info("Attempting to retry [" + (retryPolicy.getCounter() + 1)
+ "/" + retryPolicy.getNumRetries() + "]...");
retryPolicy.sleep();
if (retryPolicy.exceedsMaxRetries()) {
throw target;
} else {
continue;
}
}
}
throw e;
}
}
return null;
}
}
}
© 2015 - 2025 Weber Informatics LLC | Privacy Policy