org.apache.thrift.TServiceClient Java Examples

/**
 * 创建原始连接的方法
 * 
 * @throws ThriftConnectionPoolException
 *             创建连接出现问题时抛出该异常
 */
@SuppressWarnings("unchecked")
private void createConnection() throws ThriftConnectionPoolException {
	try {
		transport = new TSocket(host, port, connectionTimeOut);
		transport.open();
		TProtocol protocol = createTProtocol(transport);
		// 反射实例化客户端对象
		Constructor<? extends TServiceClient> clientConstructor = clientClass.getConstructor(TProtocol.class);
		client = (T) clientConstructor.newInstance(protocol);
		if (logger.isDebugEnabled()) {
			logger.debug("创建新连接成功:" + host + " 端口：" + port);
		}
	} catch (Exception e) {
		throw new ThriftConnectionPoolException("无法连接服务器：" + host + " 端口：" + port);
	}
}
 

@SuppressWarnings("unchecked")
private void addPoolAdvice(ProxyFactory proxyFactory, ThriftClientKey key) {
    proxyFactory.addAdvice((MethodInterceptor) methodInvocation -> {
        Object[] args = methodInvocation.getArguments();

        TServiceClient thriftClient = null;

        try {
            thriftClient = thriftClientsPool.borrowObject(key);
            return ReflectionUtils.invokeMethod(methodInvocation.getMethod(), thriftClient, args);
        } catch (UndeclaredThrowableException e) {
            if (TException.class.isAssignableFrom(e.getUndeclaredThrowable().getClass()))
                throw (TException) e.getUndeclaredThrowable();
            throw e;
        } finally {
            if (null != thriftClient)
                thriftClientsPool.returnObject(key, thriftClient);
        }
    });
}
 

public <X extends TServiceClient> X iface(Class<X> ifaceClass) {
if (this.loadBalance == Constant.LoadBalance.HASH) {
	throw new ValidationException("Can not use HASH without a key.");
}

switch (this.loadBalance) {
case Constant.LoadBalance.RANDOM:
	return getRandomClient(ifaceClass);
case Constant.LoadBalance.ROUND_ROBIN:
	return getRRClient(ifaceClass);
case Constant.LoadBalance.WEIGHT:
	return getWeightClient(ifaceClass);
default:
	return getRandomClient(ifaceClass);
}
  }
 

@Override
public void passivateObject(ThriftClientKey key, PooledObject<TServiceClient> p) throws Exception {
    ThriftClientPooledObject<TServiceClient> pooledObject = (ThriftClientPooledObject<TServiceClient>) p;
    TTransport transport = pooledObject.getObject().getOutputProtocol().getTransport();

    if (transport instanceof THttpClient) {
        ((THttpClient) transport).setCustomHeaders(null);
    } else {
        ((TLoadBalancerClient) transport).setCustomHeaders(null);
    }

    resetAndClose(p);

    super.passivateObject(key, pooledObject);

    if (pooledObject.getSpan() != null) {
        Span span = pooledObject.getSpan();
        span.finish();
    }
}
 

@Override
public Object invoke(Object proxy, Method method, Object[] args) throws Throwable {

    int remainRetryTimes = clientConfig.getRetryTimes();
    String exceptionMsg = null;

    do {

        ObjectPool<TTransport> connPool = null;
        TTransport transport = null;
        InvokeConn conn = null;
        try {
            Invocation invocation = new Invocation(serviceInfo.getInterfaceClazz().getName(), method.getName());
            conn = clientConfig.getLoadBalanceStrategy().select(PROVIDER_CONN_LIST, invocation);
            connPool = conn.getConnPool();
            transport = connPool.borrowObject();
            TProtocol protocol = new TBinaryProtocol(transport);
            TServiceClient client = serviceClientConstructor.newInstance(protocol);

            return method.invoke(client, args);
        } catch (Exception e) {
            // 服务多次重试连接不上,则直接将该服务对应信息移除
            if (e instanceof OureaConnCreateException) {
                if (PROVIDER_CONN_LIST.remove(conn) && conn != null && conn.getConnPool() != null){
                    conn.getConnPool().close();
                    conn = null;//help GC
                }
            }
            LOGGER.warn("invoke thrift rpc provider fail.e:", e);
            exceptionMsg = e.getMessage();
        } finally {
            if (connPool != null && transport != null) {
                connPool.invalidateObject(transport);
            }
        }
    } while (remainRetryTimes-- > 0);

    throw new OureaException("invoke fail.msg:" + exceptionMsg);
}
 

@Override
public <K extends TServiceClient> K getClient(String serviceName, Class<K> clazz) {
	if (thriftConnection != null) {
		return thriftConnection.getClient(serviceName, clazz);
	}
	throw new IllegalStateException("连接代理类没有绑定的原始连接信息");
}
 

@Override
public void verifyTraces(PluginTestVerifier verifier, String expectedMessage) throws Exception {
    // refer to TServiceClientSendBaseInterceptor.getRemoteAddress(...)
    final InetSocketAddress socketAddress = this.environment.getServerAddress();
    final String hostName = SocketAddressUtils.getHostNameFirst(socketAddress);
    final String remoteAddress = HostAndPort.toHostAndPortString(hostName, socketAddress.getPort());

    // SpanEvent - TServiceClient.sendBase
    Method sendBase = TServiceClient.class.getDeclaredMethod("sendBase", String.class, TBase.class);

    ExpectedAnnotation thriftUrl = Expectations.annotation("thrift.url",
            remoteAddress + "/com/navercorp/pinpoint/plugin/thrift/dto/EchoService/echo");
    ExpectedAnnotation thriftArgs = Expectations.annotation("thrift.args",
            "echo_args(message:" + expectedMessage + ")");

    // SpanEvent - TServiceClient.receiveBase
    Method receiveBase = TServiceClient.class.getDeclaredMethod("receiveBase", TBase.class, String.class);
    ExpectedAnnotation thriftResult = Expectations.annotation("thrift.result", "echo_result(success:"
            + expectedMessage + ")");

    verifier.verifyDiscreteTrace(event("THRIFT_CLIENT", // ServiceType
            sendBase, // Method
            null, // rpc
            null, // endPoint
            remoteAddress, // destinationId
            thriftUrl, // Annotation("thrift.url")
            thriftArgs), // Annotation("thrift.args")
            event("THRIFT_CLIENT_INTERNAL", // ServiceType
                    receiveBase, // Method
                    thriftResult // Annotation("thrift.result")
            ));
}
 

@Override
public void verifyTraces(PluginTestVerifier verifier, String expectedMessage) throws Exception {
    final InetSocketAddress socketAddress = this.environment.getServerAddress();
    final String hostName = SocketAddressUtils.getHostNameFirst(socketAddress);
    // refer to com.navercorp.pinpoint.plugin.thrift.ThriftUtils#getHostPort
    final String remoteAddress = HostAndPort.toHostAndPortString(hostName, socketAddress.getPort());
    // ********** Asynchronous Traces
    // SpanEvent - Asynchronous Invocation
    ExpectedTrace asyncInvocationTrace = event("ASYNC", "Asynchronous Invocation");

    // SpanEvent - TAsyncMethodCall.cleanUpAndFireCallback
    Method cleanUpAndFireCallback = TAsyncMethodCall.class.getDeclaredMethod("cleanUpAndFireCallback",
            SelectionKey.class);
    ExpectedTrace cleanUpAndFireCallbackTrace = event("THRIFT_CLIENT_INTERNAL", cleanUpAndFireCallback);

    // SpanEvent - TServiceClient.receiveBase
    Method receiveBase = TServiceClient.class.getDeclaredMethod("receiveBase", TBase.class, String.class);
    ExpectedAnnotation thriftResult = Expectations.annotation("thrift.result", "echo_result(success:"
            + expectedMessage + ")");
    ExpectedTrace receiveBaseTrace = event("THRIFT_CLIENT_INTERNAL", // ServiceType
            receiveBase, // Method
            thriftResult // Annotation("thrift.result")
    );

    // ********** Root trace for Asynchronous traces
    // SpanEvent - TAsyncClientManager.call
    Method call = TAsyncClientManager.class.getDeclaredMethod("call", TAsyncMethodCall.class);
    ExpectedAnnotation thriftUrl = Expectations.annotation("thrift.url",
            remoteAddress + "/com/navercorp/pinpoint/plugin/thrift/dto/EchoService/echo");
    ExpectedTrace callTrace = event("THRIFT_CLIENT", // ServiceType
            call, // Method
            null, // rpc
            null, // endPoint
            remoteAddress, // destinationId
            thriftUrl // Annotation("thrift.url")
    );
    verifier.verifyTrace(async(callTrace, asyncInvocationTrace, cleanUpAndFireCallbackTrace, receiveBaseTrace));
}
 

@Override
public void verifyTraces(PluginTestVerifier verifier, String expectedMessage) throws Exception {
    // ignore jdk http url connector traces
    verifier.ignoreServiceType("JDK_HTTPURLCONNECTOR");

    // refer to TServiceClientSendBaseInterceptor.getRemoteAddressForTHttpClient(...)
    URL url = new URL(environment.getHttpUrl());
    String hostAndPort = HostAndPort.toHostAndPortString(url.getHost(), url.getPort());

    // SpanEvent - TServiceClient.sendBase
    Method sendBaseMethod = TServiceClient.class.getDeclaredMethod("sendBase", String.class, TBase.class);
    ExpectedAnnotation thriftUrl = Expectations.annotation("thrift.url",
            hostAndPort + "/com/navercorp/pinpoint/plugin/thrift/dto/EchoService/echo");
    ExpectedAnnotation thriftArgs = Expectations.annotation("thrift.args",
            "echo_args(message:" + expectedMessage + ")");
    ExpectedTrace tServiceClientSendBaseTrace = event(
            "THRIFT_CLIENT_INTERNAL",
            sendBaseMethod,
            thriftUrl, thriftArgs);

    // SpanEvent - HttpURLConnection.connect (ignore)

    // SpanEvent - TServiceClient.receiveBase
    Method receiveBaseMethod = TServiceClient.class.getDeclaredMethod("receiveBase", TBase.class, String.class);
    ExpectedAnnotation thriftResult = Expectations.annotation(
            "thrift.result", "echo_result(success:" + expectedMessage + ")");
    ExpectedTrace tServiceClientReceiveBaseTrace = event(
            "THRIFT_CLIENT_INTERNAL",
            receiveBaseMethod,
            thriftResult);

    verifier.verifyDiscreteTrace(
            tServiceClientSendBaseTrace,
            tServiceClientReceiveBaseTrace);
}
 

protected <X extends TServiceClient> X iface(Class<X> ifaceClass, int index) {
    List<ThriftServer> serverList = getAvaliableServers();
    if (serverList == null || serverList.isEmpty()) {
        throw new NoServerAvailableException("No server available.");
    }
    index = Math.abs(index);
    final ThriftServer selected = serverList.get(index % serverList.size());
    return iface(ifaceClass, selected);
}
 

@SuppressWarnings("unchecked")
@Override
public <K extends TServiceClient> K getClient(String serviceName, Class<K> clazz) {
	T serviceClient = this.clients.get(serviceName);
	if (serviceClient == null) {
		throw new IllegalArgumentException("未知的服务名称：" + serviceName);
	}

	return (K) serviceClient;
}
 

/**
 * 创建原始连接的方法
 * 
 * @throws ThriftConnectionPoolException
 *             创建连接出现问题时抛出该异常
 */
@SuppressWarnings("unchecked")
private void createConnection() throws ThriftConnectionPoolException {
	try {
		transport = new TSocket(host, port, connectionTimeOut);
		transport.open();
		TProtocol protocol = createTProtocol(transport);

		Iterator<Entry<String, Class<? extends TServiceClient>>> iterator = thriftClientClasses.entrySet()
				.iterator();
		while (iterator.hasNext()) {
			Entry<String, Class<? extends TServiceClient>> entry = iterator.next();
			String serviceName = entry.getKey();
			Class<? extends TServiceClient> clientClass = entry.getValue();
			TMultiplexedProtocol multiProtocol = new TMultiplexedProtocol(protocol, serviceName);
			// 反射实例化客户端对象
			Constructor<? extends TServiceClient> clientConstructor = clientClass.getConstructor(TProtocol.class);
			T client = (T) clientConstructor.newInstance(multiProtocol);
			clients.put(serviceName, client);
			if (logger.isDebugEnabled()) {
				logger.debug("创建新连接成功:" + host + " 端口：" + port);
			}
		}

	} catch (Exception e) {
		e.printStackTrace();
		throw new ThriftConnectionPoolException("无法连接服务器：" + host + " 端口：" + port, e);
	}
}
 

private Constructor<TServiceClient> getClientConstructorClazz() {

        String parentClazzName = StringUtils.substringBeforeLast(serviceInfo.getInterfaceClazz().getCanonicalName(),
                ".Iface");
        String clientClazzName = parentClazzName + "$Client";

        try {
            return ((Class<TServiceClient>) Class.forName(clientClazzName)).getConstructor(TProtocol.class);
        } catch (Exception e) {
            //
            LOGGER.error("get thrift client class constructor fail.e:", e);
            throw new IllegalArgumentException("invalid iface implement");
        }

    }
 

public DefaultThriftConnection(String host, int port, int connectionTimeOut, TProtocolType tProtocolType,
		Class<? extends TServiceClient> clientClass) throws ThriftConnectionPoolException {
	this.host = host;
	this.port = port;
	this.connectionTimeOut = connectionTimeOut;
	this.tProtocolType = tProtocolType;
	this.clientClass = clientClass;

	// 创建连接
	createConnection();
}
 

public MulitServiceThriftConnecion(String host, int port, int connectionTimeOut, TProtocolType tProtocolType,
		Map<String, Class<? extends TServiceClient>> thriftClientClasses) throws ThriftConnectionPoolException {
	this.host = host;
	this.port = port;
	this.connectionTimeOut = connectionTimeOut;
	this.tProtocolType = tProtocolType;
	this.thriftClientClasses = thriftClientClasses;

	// 创建连接
	createConnection();
}
 

/**
 * Reflectively wraps an already existing Thrift client.
 *
 * @param baseClient the client to wrap
 * @param options    options that control behavior of the reconnecting client
 * @param <T>        a class that extends TServiceClient
 * @param <C>        a client interface
 * @return a wrapped client interface
 */
public static <T extends TServiceClient, C> C wrap(T baseClient, Options options) {
    Class<?>[] interfaces = baseClient.getClass().getInterfaces();

    for (Class<?> iface : interfaces) {
        if (iface.getSimpleName().equals("Iface")
                && iface.getEnclosingClass().equals(baseClient.getClass().getEnclosingClass())) {
            return (C) wrap(baseClient, iface, options);
        }
    }

    throw new IllegalStateException(
            "Class needs to implement Iface directly. Use wrap(TServiceClient, Class) instead.");
}
 

public <X extends TServiceClient> X iface(Class<X> ifaceClass, String key) {
if (this.loadBalance != Constant.LoadBalance.HASH) {
	throw new ValidationException("Must use other load balance strategy.");
}

return getHashIface(ifaceClass, key);
  }
 

public void testMultiplexedService() throws Exception {
	ThriftConnectionPoolConfig config = new ThriftConnectionPoolConfig(ThriftServiceType.MULTIPLEXED_INTERFACE);
	config.setConnectTimeout(3000);
	config.setThriftProtocol(TProtocolType.BINARY);
	// 该端口不存在
	for (ThriftServerInfo thriftServerInfo : servers) {
		config.addThriftServer(thriftServerInfo.getHost(), thriftServerInfo.getPort());
	}
	config.addThriftClientClass("example", Example.Client.class);
	config.addThriftClientClass("other", Other.Client.class);

	config.setMaxConnectionPerServer(2);
	config.setMinConnectionPerServer(1);
	config.setIdleMaxAge(2, TimeUnit.SECONDS);
	config.setMaxConnectionAge(2);
	config.setLazyInit(false);
	config.setAcquireIncrement(2);
	config.setAcquireRetryDelay(2000);

	config.setAcquireRetryAttempts(1);
	config.setMaxConnectionCreateFailedCount(1);
	config.setConnectionTimeoutInMs(5000);

	config.check();

	ThriftConnectionPool<TServiceClient> pool = new ThriftConnectionPool<TServiceClient>(config);
	ThriftConnection<TServiceClient> connection = pool.getConnection();
	// example service
	com.wmz7year.thrift.pool.example.Example.Client exampleServiceClient = connection.getClient("example",
			Example.Client.class);
	exampleServiceClient.ping();

	// other service
	com.wmz7year.thrift.pool.example.Other.Client otherServiceClient = connection.getClient("other",
			Other.Client.class);
	otherServiceClient.ping();
	pool.close();
}
 

public void testThriftConnectionPoolRemoveServerStopThread() throws Exception {

		ThriftConnectionPoolConfig config = new ThriftConnectionPoolConfig(ThriftServiceType.MULTIPLEXED_INTERFACE);
		config.setConnectTimeout(3000);
		config.setThriftProtocol(TProtocolType.BINARY);
		// 该端口不存在
		ThriftServerInfo thriftServerInfo = servers.get(0);
		config.addThriftServer(thriftServerInfo);
		config.addThriftClientClass("example", Example.Client.class);
		config.addThriftClientClass("other", Other.Client.class);

		config.setMaxConnectionPerServer(2);
		config.setMinConnectionPerServer(1);
		config.setIdleMaxAge(2, TimeUnit.SECONDS);
		config.setMaxConnectionAge(2);
		config.setLazyInit(false);
		config.setAcquireIncrement(2);
		config.setAcquireRetryDelay(2000);

		config.setAcquireRetryAttempts(1);
		config.setMaxConnectionCreateFailedCount(1);
		config.setConnectionTimeoutInMs(5000);

		ThriftConnectionPool<TServiceClient> pool = new ThriftConnectionPool<TServiceClient>(config);
		pool.getConnection().close();

		// 移除并且停止服务
		pool.removeThriftServer(thriftServerInfo);
		stopAllServers();
		pool.close();

	}
 

@Test
public void poolFactoryFullObjectLifecycleTest() throws Exception {
    PooledObject<TServiceClient> pooledObject = factory.makeObject(thriftClientKey);
    factory.activateObject(thriftClientKey, pooledObject);
    factory.validateObject(thriftClientKey, pooledObject);
    factory.passivateObject(thriftClientKey, pooledObject);
}
 

@Test
public void poolFactoryNonstandardObjectLifecycleTest() throws Exception {
    // Pooled object may be created and passivated right away
    // E.g. GenericKeyedObjectPool: returnObject() -> reuseCapacity() -> create() and addIdleObject(), which invokes passivate()
    PooledObject<TServiceClient> pooledObject = factory.makeObject(thriftClientKey);
    factory.passivateObject(thriftClientKey, pooledObject);
}
 

@Test
public void poolFactoryOverlappingObjectLifecyclesTest() throws Exception {
    PooledObject<TServiceClient> pooledObject1 = factory.makeObject(thriftClientKey);

    // activateObject leads to tracer.isEnabled() --> true
    factory.activateObject(thriftClientKey, pooledObject1);
    factory.validateObject(thriftClientKey, pooledObject1);

    // Create another pooled object and passivate it right away.
    // See poolFactoryNonstandardObjectLifecycleTest() for details
    PooledObject<TServiceClient> pooledObject2 = factory.makeObject(thriftClientKey);
    factory.passivateObject(thriftClientKey, pooledObject2);

}
 

@Override
public TServiceClient create(ThriftClientKey key) throws Exception {
    String serviceName = key.getServiceName();

    String endpoint = propertyResolver.getProperty(serviceName + ".endpoint");

    int connectTimeout = propertyResolver.getProperty(serviceName + ".connectTimeout", Integer.class, DEFAULT_CONNECTION_TIMEOUT);
    int readTimeout = propertyResolver.getProperty(serviceName + ".readTimeout", Integer.class, DEFAULT_READ_TIMEOUT);
    int maxRetries = propertyResolver.getProperty(serviceName + ".maxRetries", Integer.class, DEFAULT_MAX_RETRIES);

    TProtocol protocol;

    if (StringUtils.isEmpty(endpoint)) {
        final TLoadBalancerClient loadBalancerClient = new TLoadBalancerClient(
                this.loadBalancerClient,
                serviceName,
                propertyResolver.getProperty(serviceName + ".path", "") + key.getPath()
        );
        loadBalancerClient.setConnectTimeout(connectTimeout);
        loadBalancerClient.setReadTimeout(readTimeout);
        loadBalancerClient.setMaxRetries(maxRetries);

        protocol = protocolFactory.getProtocol(loadBalancerClient);
    } else {
        final THttpClient httpClient = new THttpClient(endpoint);
        httpClient.setConnectTimeout(connectTimeout);
        httpClient.setReadTimeout(readTimeout);

        protocol = protocolFactory.getProtocol(httpClient);
    }

    return BeanUtils.instantiateClass(
            key.getClazz().getConstructor(TProtocol.class),
            (TProtocol) protocol
    );
}
 

@SuppressWarnings("unchecked")
private void addPoolAdvice(ProxyFactory proxyFactory, ThriftClient annotataion) {
    proxyFactory.addAdvice((MethodInterceptor) methodInvocation -> getObject(
            methodInvocation,
            getThriftClientKey(
                    (Class<? extends TServiceClient>) methodInvocation.getMethod().getDeclaringClass(),
                    annotataion
            )
    ));
}
 

@Override
public void activateObject(ThriftClientKey key, PooledObject<TServiceClient> p) throws Exception {
    super.activateObject(key, p);
    ThriftClientPooledObject<TServiceClient> pooledObject = (ThriftClientPooledObject<TServiceClient>) p;

    Span span = this.tracer
            .nextSpan()
            .name(key.getServiceName())
            .kind(Span.Kind.CLIENT)
            .start();

    pooledObject.setSpan(span);
    TTransport transport = pooledObject.getObject().getOutputProtocol().getTransport();
    injectTraceHeaders(span, transport);
}
 

private Object getObject(MethodInvocation methodInvocation, ThriftClientKey key) throws Exception {
    TServiceClient thriftClient = null;
    try {
        thriftClient = thriftClientsPool.borrowObject(key);
        return ReflectionUtils.invokeMethod(methodInvocation.getMethod(), thriftClient, methodInvocation.getArguments());
    } catch (UndeclaredThrowableException e) {
        if (TException.class.isAssignableFrom(e.getUndeclaredThrowable().getClass()))
            throw (TException) e.getUndeclaredThrowable();
        throw e;
    } finally {
        if (null != thriftClient) {
            thriftClientsPool.returnObject(key, thriftClient);
        }
    }
}
 

private static boolean isThriftClient(Class<?> param) {
    Class<?> superclass = param.getSuperclass();
    boolean result = false;
    while (superclass != null) {
        if (superclass == TServiceClient.class) {
            result = true;
            break;
        } else {
            superclass = superclass.getSuperclass();
        }
    }
    return result;
}
 

@Bean
@ConditionalOnMissingBean(name = "thriftClientsPool")
public KeyedObjectPool<ThriftClientKey, TServiceClient> thriftClientsPool() {
    GenericKeyedObjectPoolConfig poolConfig = new GenericKeyedObjectPoolConfig();
    poolConfig.setMaxTotal(maxTotalThreads);
    poolConfig.setMaxIdlePerKey(maxIdleThreads);
    poolConfig.setMaxTotalPerKey(maxThreads);
    poolConfig.setJmxEnabled(false); //cause spring will autodetect itself
    return new ThriftClientPool(thriftClientPoolFactory(), poolConfig);
}
 

private KeyedPooledObjectFactory<ThriftClientKey, TServiceClient> thriftClientPoolFactory() {
    return ThriftClientPooledObjectFactory
            .builder()
            .protocolFactory(protocolFactory)
            .propertyResolver(propertyResolver)
            .loadBalancerClient(loadBalancerClient)
            .tracing(tracing)
            .tracer(tracer)
            .build();
}
 

@SuppressWarnings("unchecked")
private void addPoolAdvice(ProxyFactory proxyFactory) {
    proxyFactory.addAdvice((MethodInterceptor) methodInvocation -> getObject(
            methodInvocation,
            getThriftClientKey(
                    (Class<? extends TServiceClient>) methodInvocation.getMethod().getDeclaringClass()
            )
    ));
}