package com.alibaba.jstorm.yarn.thrift;
import java.io.IOException;
import java.net.InetAddress;
import java.net.Socket;
import java.net.UnknownHostException;
import java.util.Map;
import javax.security.auth.login.Configuration;
import org.apache.thrift7.TException;
import org.apache.thrift7.TProcessor;
import org.apache.thrift7.protocol.TBinaryProtocol;
import org.apache.thrift7.protocol.TProtocol;
import org.apache.thrift7.server.THsHaServer;
import org.apache.thrift7.server.TServer;
import org.apache.thrift7.transport.TFramedTransport;
import org.apache.thrift7.transport.TMemoryInputTransport;
import org.apache.thrift7.transport.TNonblockingServerSocket;
import org.apache.thrift7.transport.TSocket;
import org.apache.thrift7.transport.TTransport;
import org.apache.thrift7.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
/**
* Simple transport for Thrift plugin.
*
* This plugin is designed to be backward compatible with existing Storm code.
*/
public class SimpleTransportPlugin implements ITransportPlugin {
protected Configuration login_conf;
private static final Logger LOG = LoggerFactory.getLogger(SimpleTransportPlugin.class);
/**
* Invoked once immediately after construction
* @param conf Storm configuration
* @param login_conf login configuration
*/
public void prepare(Map storm_conf, Configuration login_conf) {
this.login_conf = login_conf;
}
/**
* We will let Thrift to apply default transport factory
*/
public TServer getServer(int port, TProcessor processor) throws IOException, TTransportException {
TNonblockingServerSocket serverTransport = new TNonblockingServerSocket(port);
THsHaServer.Args server_args = new THsHaServer.Args(serverTransport).
processor(new SimpleWrapProcessor(processor)).
workerThreads(64).
protocolFactory(new TBinaryProtocol.Factory());
//construct THsHaServer
return new THsHaServer(server_args);
}
/**
* Connect to the specified server via framed transport
* @param transport The underlying Thrift transport.
*/
public TTransport connect(TTransport transport, String serverHost) throws TTransportException {
//create a framed transport
TTransport conn = new TFramedTransport(transport);
//connect
conn.open();
LOG.debug("Simple client transport has been established");
return conn;
}
/**
* Processor that populate simple transport info into ReqContext, and then invoke a service handler
*/
private class SimpleWrapProcessor implements TProcessor {
final TProcessor wrapped;
SimpleWrapProcessor(TProcessor wrapped) {
this.wrapped = wrapped;
}
public boolean process(final TProtocol inProt, final TProtocol outProt) throws TException {
//populating request context
ReqContext req_context = ReqContext.context();
TTransport trans = inProt.getTransport();
if (trans instanceof TMemoryInputTransport) {
try {
req_context.setRemoteAddress(InetAddress.getLocalHost());
} catch (UnknownHostException e) {
throw new RuntimeException(e);
}
} else if (trans instanceof TSocket) {
TSocket tsocket = (TSocket)trans;
//remote address
Socket socket = tsocket.getSocket();
req_context.setRemoteAddress(socket.getInetAddress());
}
//anonymous user
req_context.setSubject(null);
//invoke service handler
return wrapped.process(inProt, outProt);
}
}
}