private void setUserName(final TProtocol in) {
TTransport transport = in.getTransport();
if (transport instanceof TSaslServerTransport) {
String userName = ((TSaslServerTransport) transport).getSaslServer().getAuthorizationID();
SessionManager.setUserName(userName);
}
}
void send(Object[] args, TProtocol protocol, int seqid) throws TException {
protocol.writeMessageBegin(new TMessage(methodName, TMessageType.CALL, seqid));
protocol.writeStructBegin(argsStruct);
// method with no parameters
if (null != args) {
for (int i = 0, n = args.length; i < n; i++) {
FieldSpec fieldSpec = requestTypeList.get(i);
Object value = args[i];
if (value == null) {
if (fieldSpec.required) {
throw new TProtocolException(
"Required field '"
+ fieldSpec.name
+ "' was not present! Struct: "
+ argsStruct.name);
}
} else {
protocol.writeFieldBegin(fieldSpec.tField);
fieldSpec.typeAdapter.write(value, protocol);
protocol.writeFieldEnd();
}
}
}
protocol.writeFieldStop();
protocol.writeStructEnd();
protocol.writeMessageEnd();
protocol.getTransport().flush();
}
@Override
public boolean process(TProtocol in, TProtocol out) throws TException {
TMessage msg = in.readMessageBegin();
String name = msg.name;
int idx = name.indexOf('.');
if (idx != -1) {
name = name.substring(idx + 1);
}
TDynamicFunction f = descManager.getFunction(name);
if (f == null) throw new TException("unknow function '" + name + "'");
Reader reader = new Reader(f.getManager());
Xdom param = reader.readStruct(f.getParamStruct(), in);
in.readMessageEnd();
if (!f.isOneway()) {
XdomMap rep;
TDynamicField retf = f.getReturnStruct().getField("success");
if (retf.getTypeValue() != TType.VOID) {
Xdom ret = handle(name, param);
rep = new XdomMap();
rep.put("success", ret);
} else {
rep = new XdomMap();
}
TMessage repmsg = new TMessage(name, TMessageType.REPLY, msg.seqid);
out.writeMessageBegin(repmsg);
Writer writer = new Writer(f.getManager());
writer.writeStruct(f.getReturnStruct(), rep, out);
out.writeMessageEnd();
out.getTransport().flush();
}
return true;
}
protected void setIpAddress(final TProtocol in) {
TTransport transport = in.getTransport();
TSocket tSocket = getUnderlyingSocketFromTransport(transport);
if (tSocket != null) {
setIpAddress(tSocket.getSocket());
} else {
LOGGER.warn("Unknown Transport, cannot determine ipAddress");
}
}
@Test
public void testDestroy() {
TProtocol prot = Mockito.mock(TProtocol.class);
when(prot.getTransport()).thenReturn(Mockito.mock(TTransport.class));
when(client.getInputProtocol()).thenReturn(prot);
ThriftConnection<HelloClient> conn = new ThriftConnection<>(pool, host);
conn.destroy();
verify(pool).remove(conn);
}
/**
* Open the connection to the worker. And start the heartbeat thread.
*
* @return true if succeed, false otherwise
* @throws IOException
*/
private synchronized boolean connect() throws IOException {
if (!mConnected) {
NetAddress workerNetAddress = null;
try {
String localHostName = NetworkAddressUtils.getLocalHostName(mTachyonConf);
LOG.info("Trying to get local worker host : " + localHostName);
workerNetAddress = mMasterClient.user_getWorker(false, localHostName);
mIsLocal =
workerNetAddress
.getMHost()
.equals(InetAddress.getByName(localHostName).getHostAddress());
} catch (NoWorkerException e) {
LOG.info(e.getMessage());
workerNetAddress = null;
} catch (UnknownHostException e) {
LOG.info(e.getMessage());
workerNetAddress = null;
}
if (workerNetAddress == null) {
try {
workerNetAddress = mMasterClient.user_getWorker(true, "");
} catch (NoWorkerException e) {
LOG.info("No worker running in the system: " + e.getMessage());
mClient = null;
return false;
}
}
String host = NetworkAddressUtils.getFqdnHost(workerNetAddress);
int port = workerNetAddress.mPort;
mWorkerAddress = new InetSocketAddress(host, port);
mWorkerDataServerAddress = new InetSocketAddress(host, workerNetAddress.mSecondaryPort);
LOG.info("Connecting " + (mIsLocal ? "local" : "remote") + " worker @ " + mWorkerAddress);
mProtocol = new TBinaryProtocol(new TFramedTransport(new TSocket(host, port)));
mClient = new WorkerService.Client(mProtocol);
mHeartbeatExecutor = new WorkerClientHeartbeatExecutor(this, mMasterClient.getUserId());
String threadName = "worker-heartbeat-" + mWorkerAddress;
int interval = mTachyonConf.getInt(Constants.USER_HEARTBEAT_INTERVAL_MS, Constants.SECOND_MS);
mHeartbeat =
mExecutorService.submit(new HeartbeatThread(threadName, mHeartbeatExecutor, interval));
try {
mProtocol.getTransport().open();
} catch (TTransportException e) {
LOG.error(e.getMessage(), e);
return false;
}
mConnected = true;
}
return mConnected;
}
private void processTraceObject(
final Trace trace, Object target, Object[] args, Throwable throwable) {
// end spanEvent
try {
SpanEventRecorder recorder = trace.currentSpanEventRecorder();
// TODO Might need a way to collect and record method arguments
// trace.recordAttribute(...);
recorder.recordException(throwable);
recorder.recordApi(this.descriptor);
} catch (Throwable t) {
logger.warn("Error processing trace object. Cause:{}", t.getMessage(), t);
} finally {
trace.traceBlockEnd();
}
// end root span
SpanRecorder recorder = trace.getSpanRecorder();
String methodUri = getMethodUri(target);
recorder.recordRpcName(methodUri);
// retrieve connection information
String localIpPort = UNKNOWN_ADDRESS;
String remoteAddress = UNKNOWN_ADDRESS;
if (args.length == 2 && args[0] instanceof TProtocol) {
TProtocol inputProtocol = (TProtocol) args[0];
if (this.socketAccessor.isApplicable(inputProtocol.getTransport())) {
Socket socket = this.socketAccessor.get(inputProtocol.getTransport());
if (socket != null) {
localIpPort = ThriftUtils.getHostPort(socket.getLocalSocketAddress());
remoteAddress = ThriftUtils.getHost(socket.getRemoteSocketAddress());
}
}
}
if (localIpPort != UNKNOWN_ADDRESS) {
recorder.recordEndPoint(localIpPort);
}
if (remoteAddress != UNKNOWN_ADDRESS) {
recorder.recordRemoteAddress(remoteAddress);
}
}
public synchronized boolean open() {
if (!mIsConnected) {
try {
mProtocol.getTransport().open();
} catch (TTransportException e) {
LOG.error(e.getMessage(), e);
return false;
}
mIsConnected = true;
}
return mIsConnected;
}
/** Close the connection to worker. Shutdown the heartbeat thread. */
@Override
public synchronized void close() {
if (mConnected) {
try {
// Heartbeat to send the client metrics.
if (mHeartbeatExecutor != null) {
mHeartbeatExecutor.heartbeat();
}
mProtocol.getTransport().close();
} finally {
if (mHeartbeat != null) {
mHeartbeat.cancel(true);
}
}
mConnected = false;
}
}
public boolean process(TProtocol iprot, TProtocol oprot) throws TException {
String ip = getClientIp(iprot);
if (!whiteListHelper.check(ip)) {
System.out.println(
"------------ customized secure check, client ip (" + ip + ") not allowed");
TMessage msg = iprot.readMessageBegin();
TProtocolUtil.skip(iprot, TType.STRUCT);
iprot.readMessageEnd();
TApplicationException appException =
new TApplicationException(TApplicationException.UNKNOWN, "Invalid ip: " + ip);
oprot.writeMessageBegin(new TMessage(msg.name, TMessageType.EXCEPTION, msg.seqid));
appException.write(oprot);
oprot.writeMessageEnd();
oprot.getTransport().flush();
return false;
}
System.out.println("------------ customized secure check, client ip (" + ip + ") allowed");
boolean result = super.process(iprot, oprot);
return result;
}
public TTransport getTransport() {
return delegate.getTransport();
}
public InterningProtocol(TProtocol delegate) {
super(delegate.getTransport());
this.delegate = delegate;
}
public synchronized void close() {
mProtocol.getTransport().close();
mIsConnected = false;
}
private String getClientIp(TProtocol iprot) {
TSocket socket = (TSocket) iprot.getTransport();
InetAddress address = socket.getSocket().getInetAddress();
return address.getHostAddress();
}
@Override
public void before(Object target, Object[] args) {
if (isDebug) {
logger.beforeInterceptor(target, args);
}
if (target instanceof TServiceClient) {
TServiceClient client = (TServiceClient) target;
TProtocol oprot = client.getOutputProtocol();
TTransport transport = oprot.getTransport();
final Trace trace = traceContext.currentRawTraceObject();
if (trace == null) {
return;
}
ThriftRequestProperty parentTraceInfo = new ThriftRequestProperty();
final boolean shouldSample = trace.canSampled();
if (!shouldSample) {
if (isDebug) {
logger.debug("set Sampling flag=false");
}
parentTraceInfo.setShouldSample(shouldSample);
} else {
SpanEventRecorder recorder = trace.traceBlockBegin();
recorder.recordServiceType(ThriftConstants.THRIFT_CLIENT);
// retrieve connection information
String remoteAddress = ThriftConstants.UNKNOWN_ADDRESS;
if (transport instanceof SocketFieldAccessor) {
Socket socket = ((SocketFieldAccessor) transport)._$APM$_getSocket();
if (socket != null) {
remoteAddress = ThriftUtils.getHostPort(socket.getRemoteSocketAddress());
}
} else {
if (isDebug) {
logger.debug(
"Invalid target object. Need field accessor({}).",
SocketFieldAccessor.class.getName());
}
}
recorder.recordDestinationId(remoteAddress);
String methodName = ThriftConstants.UNKNOWN_METHOD_NAME;
if (args[0] instanceof String) {
methodName = (String) args[0];
}
String serviceName = ThriftUtils.getClientServiceName(client);
String thriftUrl = getServiceUrl(remoteAddress, serviceName, methodName);
recorder.recordAttribute(ThriftConstants.THRIFT_URL, thriftUrl);
TraceId nextId = trace.getTraceId().getNextTraceId();
recorder.recordNextSpanId(nextId.getSpanId());
parentTraceInfo.setTraceId(nextId.getTransactionId());
parentTraceInfo.setSpanId(nextId.getSpanId());
parentTraceInfo.setParentSpanId(nextId.getParentSpanId());
parentTraceInfo.setFlags(nextId.getFlags());
parentTraceInfo.setParentApplicationName(traceContext.getApplicationName());
parentTraceInfo.setParentApplicationType(traceContext.getServerTypeCode());
parentTraceInfo.setAcceptorHost(remoteAddress);
}
InterceptorGroupInvocation currentTransaction = this.group.getCurrentInvocation();
currentTransaction.setAttachment(parentTraceInfo);
}
}
protected void sendBase(String methodName, TBase args) throws TException {
oprot_.writeMessageBegin(new TMessage(methodName, TMessageType.CALL, ++seqid_));
args.write(oprot_);
oprot_.writeMessageEnd();
oprot_.getTransport().flush();
}
