diff --git a/server/base/src/main/java/org/apache/accumulo/server/rpc/ClientInfoProcessorFactory.java b/server/base/src/main/java/org/apache/accumulo/server/rpc/ClientInfoProcessorFactory.java deleted file mode 100644 index 94d0bd71414..00000000000 --- a/server/base/src/main/java/org/apache/accumulo/server/rpc/ClientInfoProcessorFactory.java +++ /dev/null @@ -1,51 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -package org.apache.accumulo.server.rpc; - -import org.apache.thrift.TProcessor; -import org.apache.thrift.TProcessorFactory; -import org.apache.thrift.transport.TSocket; -import org.apache.thrift.transport.TTransport; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * Sets the address of a client in a ThreadLocal to allow for more informative log messages. - */ -public class ClientInfoProcessorFactory extends TProcessorFactory { - private static final Logger log = LoggerFactory.getLogger(ClientInfoProcessorFactory.class); - - private final ThreadLocal clientAddress; - - public ClientInfoProcessorFactory(ThreadLocal clientAddress, TProcessor processor) { - super(processor); - this.clientAddress = clientAddress; - } - - @Override - public TProcessor getProcessor(TTransport trans) { - if (trans instanceof TSocket tsock) { - clientAddress.set( - tsock.getSocket().getInetAddress().getHostAddress() + ":" + tsock.getSocket().getPort()); - } else { - log.warn("Unable to extract clientAddress from transport of type {}", trans.getClass()); - } - return super.getProcessor(trans); - } -} diff --git a/server/base/src/main/java/org/apache/accumulo/server/rpc/CustomNonBlockingServer.java b/server/base/src/main/java/org/apache/accumulo/server/rpc/CustomNonBlockingServer.java deleted file mode 100644 index 6cbd5c6ca52..00000000000 --- a/server/base/src/main/java/org/apache/accumulo/server/rpc/CustomNonBlockingServer.java +++ /dev/null @@ -1,179 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -package org.apache.accumulo.server.rpc; - -import java.io.IOException; -import java.lang.reflect.Field; -import java.net.Socket; -import java.nio.channels.SelectionKey; - -import org.apache.thrift.server.THsHaServer; -import org.apache.thrift.server.TNonblockingServer; -import org.apache.thrift.transport.TNonblockingServerTransport; -import org.apache.thrift.transport.TNonblockingSocket; -import org.apache.thrift.transport.TNonblockingTransport; -import org.apache.thrift.transport.TTransportException; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -/** - * This class implements a custom non-blocking thrift server that stores the client address in - * thread-local storage for the invocation. - */ -public class CustomNonBlockingServer extends THsHaServer { - - private static final Logger log = LoggerFactory.getLogger(CustomNonBlockingServer.class); - private final Field selectAcceptThreadField; - - public CustomNonBlockingServer(Args args) { - super(args); - - try { - selectAcceptThreadField = TNonblockingServer.class.getDeclaredField("selectAcceptThread_"); - selectAcceptThreadField.setAccessible(true); - } catch (ReflectiveOperationException e) { - throw new IllegalStateException("Failed to access required field in Thrift code.", e); - } - } - - @Override - public void stop() { - super.stop(); - try { - getInvoker().shutdownNow(); - } catch (Exception e) { - log.error("Unable to call shutdownNow", e); - } - } - - @Override - protected boolean startThreads() { - // Yet another dirty/gross hack to get access to the client's address. - - // start the selector - try { - // Hack in our SelectAcceptThread impl - SelectAcceptThread selectAcceptThread_ = - new CustomSelectAcceptThread((TNonblockingServerTransport) serverTransport_); - // Set the private field before continuing. - selectAcceptThreadField.set(this, selectAcceptThread_); - - selectAcceptThread_.start(); - return true; - } catch (IOException e) { - LOGGER.error("Failed to start selector thread!", e); - return false; - } catch (IllegalAccessException | IllegalArgumentException e) { - throw new IllegalStateException("Exception setting customer select thread in Thrift"); - } - } - - /** - * Custom wrapper around {@link org.apache.thrift.server.TNonblockingServer.SelectAcceptThread} to - * create our {@link CustomFrameBuffer}. - */ - private class CustomSelectAcceptThread extends SelectAcceptThread { - - public CustomSelectAcceptThread(TNonblockingServerTransport serverTransport) - throws IOException { - super(serverTransport); - } - - @Override - protected FrameBuffer createFrameBuffer(final TNonblockingTransport trans, - final SelectionKey selectionKey, final AbstractSelectThread selectThread) - throws TTransportException { - if (processorFactory_.isAsyncProcessor()) { - throw new IllegalStateException("This implementation does not support AsyncProcessors"); - } - - return new CustomFrameBuffer(trans, selectionKey, selectThread); - } - } - - /** - * Custom wrapper around {@link org.apache.thrift.server.AbstractNonblockingServer.FrameBuffer} to - * extract the client's network location before accepting the request. - */ - private class CustomFrameBuffer extends FrameBuffer { - private final String clientAddress; - - public CustomFrameBuffer(TNonblockingTransport trans, SelectionKey selectionKey, - AbstractSelectThread selectThread) throws TTransportException { - super(trans, selectionKey, selectThread); - // Store the clientAddress in the buffer so it can be referenced for logging during read/write - this.clientAddress = getClientAddress(); - } - - @Override - public void invoke() { - // On invoke() set the clientAddress on the ThreadLocal so that it can be accessed elsewhere - // in the same thread that called invoke() on the buffer - TServerUtils.clientAddress.set(clientAddress); - super.invoke(); - } - - @Override - public boolean read() { - boolean result = super.read(); - if (!result) { - log.trace("CustomFrameBuffer.read returned false when reading data from client: {}", - clientAddress); - } - return result; - } - - @Override - public boolean write() { - boolean result = super.write(); - if (!result) { - log.trace("CustomFrameBuffer.write returned false when writing data to client: {}", - clientAddress); - } - return result; - } - - /* - * Helper method used to capture the client address inside the CustomFrameBuffer constructor so - * that it can be referenced inside the read/write methods for logging purposes. It previously - * was only set on the ThreadLocal in the invoke() method but that does not work because A) the - * method isn't called until after reading is finished so the value will be null inside of - * read() and B) The other problem is that invoke() is called on a different thread than - * read()/write() so even if the order was correct it would not be available. - * - * Since a new FrameBuffer is created for each request we can use it to capture the client - * address earlier in the constructor and not wait for invoke(). A FrameBuffer is used to read - * data and write a response back to the client and as part of creation of the buffer the - * TNonblockingSocket is stored as a final variable and won't change so we can safely capture - * the clientAddress in the constructor and use it for logging during read/write and then use - * the value inside of invoke() to set the ThreadLocal so the client address will still be - * available on the thread that called invoke(). - */ - private String getClientAddress() { - String clientAddress = null; - if (trans_ instanceof TNonblockingSocket tsock) { - Socket sock = tsock.getSocketChannel().socket(); - clientAddress = sock.getInetAddress().getHostAddress() + ":" + sock.getPort(); - log.trace("CustomFrameBuffer captured client address: {}", clientAddress); - } - return clientAddress; - } - } - -} diff --git a/server/base/src/main/java/org/apache/accumulo/server/rpc/CustomThreadedSelectorServer.java b/server/base/src/main/java/org/apache/accumulo/server/rpc/CustomThreadedSelectorServer.java deleted file mode 100644 index 455640ab543..00000000000 --- a/server/base/src/main/java/org/apache/accumulo/server/rpc/CustomThreadedSelectorServer.java +++ /dev/null @@ -1,74 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -package org.apache.accumulo.server.rpc; - -import java.lang.reflect.Field; -import java.net.Socket; - -import org.apache.thrift.server.TThreadedSelectorServer; -import org.apache.thrift.transport.TNonblockingSocket; -import org.apache.thrift.transport.TNonblockingTransport; -import org.slf4j.LoggerFactory; - -public class CustomThreadedSelectorServer extends TThreadedSelectorServer { - - private final Field fbTansportField; - - public CustomThreadedSelectorServer(Args args) { - super(args); - - try { - fbTansportField = FrameBuffer.class.getDeclaredField("trans_"); - fbTansportField.setAccessible(true); - } catch (SecurityException | NoSuchFieldException e) { - throw new IllegalStateException("Failed to access required field in Thrift code.", e); - } - } - - private TNonblockingTransport getTransport(FrameBuffer frameBuffer) { - try { - return (TNonblockingTransport) fbTansportField.get(frameBuffer); - } catch (IllegalAccessException e) { - throw new IllegalStateException(e); - } - } - - @Override - protected Runnable getRunnable(FrameBuffer frameBuffer) { - return () -> { - - try { - TNonblockingTransport transport = getTransport(frameBuffer); - - if (transport instanceof TNonblockingSocket tsock) { - // This block of code makes the client address available to the server side code that - // executes a RPC. It is made available for informational purposes. - Socket sock = tsock.getSocketChannel().socket(); - TServerUtils.clientAddress - .set(sock.getInetAddress().getHostAddress() + ":" + sock.getPort()); - } - } catch (Exception e) { - LoggerFactory.getLogger(CustomThreadedSelectorServer.class) - .warn("Failed to get client address ", e); - } - frameBuffer.invoke(); - }; - } - -} diff --git a/server/base/src/main/java/org/apache/accumulo/server/rpc/TServerUtils.java b/server/base/src/main/java/org/apache/accumulo/server/rpc/TServerUtils.java index a601eb8a66c..67503e57a1a 100644 --- a/server/base/src/main/java/org/apache/accumulo/server/rpc/TServerUtils.java +++ b/server/base/src/main/java/org/apache/accumulo/server/rpc/TServerUtils.java @@ -51,6 +51,7 @@ import org.apache.thrift.TProcessor; import org.apache.thrift.TProcessorFactory; import org.apache.thrift.protocol.TProtocolFactory; +import org.apache.thrift.server.THsHaServer; import org.apache.thrift.server.TThreadPoolServer; import org.apache.thrift.server.TThreadedSelectorServer; import org.apache.thrift.transport.TNonblockingServerSocket; @@ -75,8 +76,7 @@ public class TServerUtils { private static final Logger log = LoggerFactory.getLogger(TServerUtils.class); /** - * Static instance, passed to {@link ClientInfoProcessorFactory}, which will contain the client - * address of any incoming RPC. + * Static instance, which will contain the client address of any incoming RPC. */ public static final ThreadLocal clientAddress = new ThreadLocal<>(); @@ -190,7 +190,10 @@ private static ServerAddress createThreadedSelectorServer(HostAndPort address, address = HostAndPort.fromParts(address.getHost(), transport.getPort()); } - return new ServerAddress(new CustomThreadedSelectorServer(options), address); + final TThreadedSelectorServer server = new TThreadedSelectorServer(options); + server.setServerEventHandler(new ThriftServerEventHandler()); + + return new ServerAddress(server, address); } /** @@ -207,7 +210,7 @@ private static ServerAddress createNonBlockingServer(HostAndPort address, TProce .clientTimeout(0).maxFrameSize(Ints.saturatedCast(maxMessageSize)); final TNonblockingServerSocket transport = new TNonblockingServerSocket(args); - final CustomNonBlockingServer.Args options = new CustomNonBlockingServer.Args(transport); + THsHaServer.Args options = new THsHaServer.Args(transport); options.protocolFactory(protocolFactory); options.transportFactory(ThriftUtil.transportFactory(maxMessageSize)); @@ -225,7 +228,10 @@ private static ServerAddress createNonBlockingServer(HostAndPort address, TProce address = HostAndPort.fromParts(address.getHost(), transport.getPort()); } - return new ServerAddress(new CustomNonBlockingServer(options), address); + final THsHaServer server = new THsHaServer(options); + server.setServerEventHandler(new ThriftServerEventHandler()); + + return new ServerAddress(server, address); } /** @@ -295,6 +301,8 @@ private static ServerAddress createBlockingServer(HostAndPort address, TProcesso log.info("Blocking Server bound on {}", address); } + server.setServerEventHandler(new ThriftServerEventHandler()); + return new ServerAddress(server, address); } @@ -313,11 +321,14 @@ private static TThreadPoolServer createTThreadPoolServer(TServerTransport transp TThreadPoolServer.Args options = new TThreadPoolServer.Args(transport); options.protocolFactory(protocolFactory); options.transportFactory(transportFactory); - options.processorFactory(new ClientInfoProcessorFactory(clientAddress, processor)); if (service != null) { options.executorService(service); } - return new TThreadPoolServer(options); + + final TThreadPoolServer server = new TThreadPoolServer(options); + server.setServerEventHandler(new ThriftServerEventHandler()); + + return server; } /** @@ -395,9 +406,11 @@ private static ServerAddress createSslThreadPoolServer(HostAndPort address, TPro ThreadPoolExecutor pool = createSelfResizingThreadPool(numThreads, threadTimeOut, conf, timeBetweenThreadChecks); + TThreadPoolServer server = createTThreadPoolServer(transport, processor, + ThriftUtil.transportFactory(), protocolFactory, pool); + server.setServerEventHandler(new ThriftServerEventHandler()); - return new ServerAddress(createTThreadPoolServer(transport, processor, - ThriftUtil.transportFactory(), protocolFactory, pool), address); + return new ServerAddress(server, address); } private static ServerAddress createSaslThreadPoolServer(HostAndPort address, TProcessor processor, @@ -495,6 +508,7 @@ private static ServerAddress createSaslThreadPoolServer(HostAndPort address, TPr final TThreadPoolServer server = createTThreadPoolServer(transport, processor, ugiTransportFactory, protocolFactory, pool); + server.setServerEventHandler(new ThriftServerEventHandler()); return new ServerAddress(server, address); } diff --git a/server/base/src/main/java/org/apache/accumulo/server/rpc/ThriftServerEventHandler.java b/server/base/src/main/java/org/apache/accumulo/server/rpc/ThriftServerEventHandler.java new file mode 100644 index 00000000000..a4564653c4c --- /dev/null +++ b/server/base/src/main/java/org/apache/accumulo/server/rpc/ThriftServerEventHandler.java @@ -0,0 +1,71 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * https://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.accumulo.server.rpc; + +import java.net.SocketAddress; + +import org.apache.thrift.protocol.TProtocol; +import org.apache.thrift.server.ServerContext; +import org.apache.thrift.server.TServerEventHandler; +import org.apache.thrift.transport.TTransport; + +public class ThriftServerEventHandler implements TServerEventHandler { + + private static final ThriftServerContext context = new ThriftServerContext(); + + public static class ThriftServerContext implements ServerContext { + + @Override + public T unwrap(Class iface) { + throw new UnsupportedOperationException("This method has not been implemented"); + } + + @Override + public boolean isWrapperFor(Class iface) { + return false; + } + + @Override + public void setRemoteAddress(SocketAddress remoteAddress) { + TServerUtils.clientAddress.set(remoteAddress.toString()); + } + + public void clear() { + TServerUtils.clientAddress.set(""); + } + } + + @Override + public void preServe() {} + + @Override + public ServerContext createContext(TProtocol input, TProtocol output) { + return context; + } + + @Override + public void deleteContext(ServerContext serverContext, TProtocol input, TProtocol output) { + context.clear(); + } + + @Override + public void processContext(ServerContext serverContext, TTransport inputTransport, + TTransport outputTransport) {} + +}