Use TransportChannel in TransportHandshaker (elastic#54684)

Currently the TransportHandshaker has a specialized codepath for sending
a response. In other work, we are going to start having handshakes
contribute to circuit breaking (while not being breakable). This commit
moves in that direction by allowing the handshaker to responding using a
standard TcpTransportChannel similar to other requests.

Author: Tim-Brooks

server/src/main/java/org/elasticsearch/transport/InboundHandler.java

@@ -157,7 +157,10 @@ private <T extends TransportRequest> void handleRequest(TcpChannel channel, Head
         try {
             messageListener.onRequestReceived(requestId, action);
             if (header.isHandshake()) {
-                handshaker.handleHandshake(version, features, channel, requestId, stream);
+                // Handshakes are not currently circuit broken
+                transportChannel = new TcpTransportChannel(outboundHandler, channel, action, requestId, version, features,
+                    circuitBreakerService, 0, header.isCompressed(), header.isHandshake());
+                handshaker.handleHandshake(transportChannel, requestId, stream);
             } else {
                 final RequestHandlerRegistry<T> reg = getRequestHandler(action);
                 if (reg == null) {
@@ -170,7 +173,7 @@ private <T extends TransportRequest> void handleRequest(TcpChannel channel, Head
                     breaker.addWithoutBreaking(messageLengthBytes);
                 }
                 transportChannel = new TcpTransportChannel(outboundHandler, channel, action, requestId, version, features,
-                    circuitBreakerService, messageLengthBytes, header.isCompressed());
+                    circuitBreakerService, messageLengthBytes, header.isCompressed(), header.isHandshake());
                 final T request = reg.newRequest(stream);
                 request.remoteAddress(new TransportAddress(channel.getRemoteAddress()));
                 // in case we throw an exception, i.e. when the limit is hit, we don't want to verify
@@ -186,7 +189,7 @@ private <T extends TransportRequest> void handleRequest(TcpChannel channel, Head
             // the circuit breaker tripped
             if (transportChannel == null) {
                 transportChannel = new TcpTransportChannel(outboundHandler, channel, action, requestId, version, features,
-                    circuitBreakerService, 0, header.isCompressed());
+                    circuitBreakerService, 0, header.isCompressed(), header.isHandshake());
             }
             try {
                 transportChannel.sendResponse(e);

server/src/main/java/org/elasticsearch/transport/TcpTransport.java

@@ -159,9 +159,7 @@ public TcpTransport(Settings settings, Version version, ThreadPool threadPool, P
         this.handshaker = new TransportHandshaker(version, threadPool,
             (node, channel, requestId, v) -> outboundHandler.sendRequest(node, channel, requestId,
                 TransportHandshaker.HANDSHAKE_ACTION_NAME, new TransportHandshaker.HandshakeRequest(version),
-                TransportRequestOptions.EMPTY, v, false, true),
-            (v, features1, channel, response, requestId) -> outboundHandler.sendResponse(v, features1, channel, requestId,
-                TransportHandshaker.HANDSHAKE_ACTION_NAME, response, false, true));
+                TransportRequestOptions.EMPTY, v, false, true));
         this.keepAlive = new TransportKeepAlive(threadPool, this.outboundHandler::sendBytes);
         this.inboundHandler = new InboundHandler(threadPool, outboundHandler, namedWriteableRegistry, circuitBreakerService, handshaker,
             keepAlive);

server/src/main/java/org/elasticsearch/transport/TcpTransportChannel.java

@@ -39,9 +39,11 @@ public final class TcpTransportChannel implements TransportChannel {
     private final CircuitBreakerService breakerService;
     private final long reservedBytes;
     private final boolean compressResponse;
+    private final boolean isHandshake;
 
     TcpTransportChannel(OutboundHandler outboundHandler, TcpChannel channel, String action, long requestId, Version version,
-                        Set<String> features, CircuitBreakerService breakerService, long reservedBytes, boolean compressResponse) {
+                        Set<String> features, CircuitBreakerService breakerService, long reservedBytes, boolean compressResponse,
+                        boolean isHandshake) {
         this.version = version;
         this.features = features;
         this.channel = channel;
@@ -51,6 +53,7 @@ public final class TcpTransportChannel implements TransportChannel {
         this.breakerService = breakerService;
         this.reservedBytes = reservedBytes;
         this.compressResponse = compressResponse;
+        this.isHandshake = isHandshake;
     }
 
     @Override
@@ -61,7 +64,7 @@ public String getProfileName() {
     @Override
     public void sendResponse(TransportResponse response) throws IOException {
         try {
-            outboundHandler.sendResponse(version, features, channel, requestId, action, response, compressResponse, false);
+            outboundHandler.sendResponse(version, features, channel, requestId, action, response, compressResponse, isHandshake);
         } finally {
             release(false);
         }
@@ -102,6 +105,5 @@ public Version getVersion() {
     public TcpChannel getChannel() {
         return channel;
     }
-
 }
 

server/src/main/java/org/elasticsearch/transport/TransportHandshaker.java

@@ -31,7 +31,6 @@
 
 import java.io.EOFException;
 import java.io.IOException;
-import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicBoolean;
@@ -49,14 +48,11 @@ final class TransportHandshaker {
     private final Version version;
     private final ThreadPool threadPool;
     private final HandshakeRequestSender handshakeRequestSender;
-    private final HandshakeResponseSender handshakeResponseSender;
 
-    TransportHandshaker(Version version, ThreadPool threadPool, HandshakeRequestSender handshakeRequestSender,
-                        HandshakeResponseSender handshakeResponseSender) {
+    TransportHandshaker(Version version, ThreadPool threadPool, HandshakeRequestSender handshakeRequestSender) {
         this.version = version;
         this.threadPool = threadPool;
         this.handshakeRequestSender = handshakeRequestSender;
-        this.handshakeResponseSender = handshakeResponseSender;
     }
 
     void sendHandshake(long requestId, DiscoveryNode node, TcpChannel channel, TimeValue timeout, ActionListener<Version> listener) {
@@ -88,16 +84,15 @@ void sendHandshake(long requestId, DiscoveryNode node, TcpChannel channel, TimeV
         }
     }
 
-    void handleHandshake(Version version, Set<String> features, TcpChannel channel, long requestId, StreamInput stream) throws IOException {
+    void handleHandshake(TransportChannel channel, long requestId, StreamInput stream) throws IOException {
         // Must read the handshake request to exhaust the stream
         HandshakeRequest handshakeRequest = new HandshakeRequest(stream);
         final int nextByte = stream.read();
         if (nextByte != -1) {
             throw new IllegalStateException("Handshake request not fully read for requestId [" + requestId + "], action ["
                 + TransportHandshaker.HANDSHAKE_ACTION_NAME + "], available [" + stream.available() + "]; resetting");
         }
-        HandshakeResponse response = new HandshakeResponse(this.version);
-        handshakeResponseSender.sendResponse(version, features, channel, response, requestId);
+        channel.sendResponse(new HandshakeResponse(this.version));
     }
 
     TransportResponseHandler<HandshakeResponse> removeHandlerForHandshake(long requestId) {
@@ -228,11 +223,4 @@ interface HandshakeRequestSender {
 
         void sendRequest(DiscoveryNode node, TcpChannel channel, long requestId, Version version) throws IOException;
     }
-
-    @FunctionalInterface
-    interface HandshakeResponseSender {
-
-        void sendResponse(Version version, Set<String> features, TcpChannel channel, TransportResponse response, long requestId)
-            throws IOException;
-    }
 }

server/src/test/java/org/elasticsearch/action/support/broadcast/node/TransportBroadcastByNodeActionTests.java

@@ -57,7 +57,7 @@
 import org.elasticsearch.test.transport.CapturingTransport;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
-import org.elasticsearch.transport.TransportChannel;
+import org.elasticsearch.transport.TestTransportChannel;
 import org.elasticsearch.transport.TransportResponse;
 import org.elasticsearch.transport.TransportService;
 import org.junit.After;
@@ -366,14 +366,15 @@ public void testOperationExecution() throws Exception {
         final TransportBroadcastByNodeAction.BroadcastByNodeTransportRequestHandler handler =
                 action.new BroadcastByNodeTransportRequestHandler();
 
-        TestTransportChannel channel = new TestTransportChannel();
+        final PlainActionFuture<TransportResponse> future = PlainActionFuture.newFuture();
+        TestTransportChannel channel = new TestTransportChannel(future);
 
         handler.messageReceived(action.new NodeRequest(nodeId, new Request(), new ArrayList<>(shards)), channel, null);
 
         // check the operation was executed only on the expected shards
         assertEquals(shards, action.getResults().keySet());
 
-        TransportResponse response = channel.getCapturedResponse();
+        TransportResponse response = future.actionGet();
         assertTrue(response instanceof TransportBroadcastByNodeAction.NodeResponse);
         TransportBroadcastByNodeAction.NodeResponse nodeResponse = (TransportBroadcastByNodeAction.NodeResponse) response;
 
@@ -469,32 +470,4 @@ public void testResultAggregation() throws ExecutionException, InterruptedExcept
         assertEquals("failed shards", totalFailedShards, response.getFailedShards());
         assertEquals("accumulated exceptions", totalFailedShards, response.getShardFailures().length);
     }
-
-    public class TestTransportChannel implements TransportChannel {
-        private TransportResponse capturedResponse;
-
-        public TransportResponse getCapturedResponse() {
-            return capturedResponse;
-        }
-
-        @Override
-        public String getProfileName() {
-            return "";
-        }
-
-        @Override
-        public void sendResponse(TransportResponse response) throws IOException {
-            capturedResponse = response;
-        }
-
-        @Override
-        public void sendResponse(Exception exception) throws IOException {
-        }
-
-        @Override
-        public String getChannelType() {
-            return "test";
-        }
-
-    }
 }

server/src/test/java/org/elasticsearch/action/support/replication/TransportReplicationActionTests.java

@@ -78,6 +78,7 @@
 import org.elasticsearch.test.transport.MockTransportService;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.transport.TestTransportChannel;
 import org.elasticsearch.transport.Transport;
 import org.elasticsearch.transport.TransportChannel;
 import org.elasticsearch.transport.TransportException;
@@ -817,7 +818,7 @@ public void testSeqNoIsSetOnPrimary() {
         Request request = new Request(shardId);
         TransportReplicationAction.ConcreteShardRequest<Request> concreteShardRequest =
             new TransportReplicationAction.ConcreteShardRequest<>(request, routingEntry.allocationId().getId(), primaryTerm);
-        PlainActionFuture<TestResponse> listener = new PlainActionFuture<>();
+        PlainActionFuture<TransportResponse> listener = new PlainActionFuture<>();
 
 
         final IndexShard shard = mockIndexShard(shardId, clusterService);
@@ -981,7 +982,7 @@ public void testPrimaryActionRejectsWrongAidOrWrongTerm() throws Exception {
         setState(clusterService, state(index, true, ShardRoutingState.STARTED));
         final ShardRouting primary = clusterService.state().routingTable().shardRoutingTable(shardId).primaryShard();
         final long primaryTerm = clusterService.state().metadata().index(shardId.getIndexName()).primaryTerm(shardId.id());
-        PlainActionFuture<TestResponse> listener = new PlainActionFuture<>();
+        PlainActionFuture<TransportResponse> listener = new PlainActionFuture<>();
         final boolean wrongAllocationId = randomBoolean();
         final long requestTerm = wrongAllocationId && randomBoolean() ? primaryTerm : primaryTerm + randomIntBetween(1, 10);
         Request request = new Request(shardId).timeout("1ms");
@@ -1018,7 +1019,7 @@ public void testReplicaActionRejectsWrongAid() throws Exception {
         state = ClusterState.builder(state).nodes(DiscoveryNodes.builder(state.nodes()).localNodeId(replica.currentNodeId())).build();
         setState(clusterService, state);
 
-        PlainActionFuture<TestResponse> listener = new PlainActionFuture<>();
+        PlainActionFuture<TransportResponse> listener = new PlainActionFuture<>();
         Request request = new Request(shardId).timeout("1ms");
         action.handleReplicaRequest(
             new TransportReplicationAction.ConcreteReplicaRequest<>(request, "_not_a_valid_aid_", randomNonNegativeLong(),
@@ -1062,7 +1063,7 @@ protected ReplicaResult shardOperationOnReplica(Request request, IndexShard repl
                 return new ReplicaResult();
             }
         };
-        final PlainActionFuture<TestResponse> listener = new PlainActionFuture<>();
+        final PlainActionFuture<TransportResponse> listener = new PlainActionFuture<>();
         final Request request = new Request(shardId);
         final long checkpoint = randomNonNegativeLong();
         final long maxSeqNoOfUpdatesOrDeletes = randomNonNegativeLong();
@@ -1130,7 +1131,7 @@ protected ReplicaResult shardOperationOnReplica(Request request, IndexShard repl
                 return new ReplicaResult();
             }
         };
-        final PlainActionFuture<TestResponse> listener = new PlainActionFuture<>();
+        final PlainActionFuture<TransportResponse> listener = new PlainActionFuture<>();
         final Request request = new Request(shardId);
         final long checkpoint = randomNonNegativeLong();
         final long maxSeqNoOfUpdates = randomNonNegativeLong();
@@ -1371,29 +1372,8 @@ private IndexShard mockIndexShard(ShardId shardId, ClusterService clusterService
     /**
      * Transport channel that is needed for replica operation testing.
      */
-    public TransportChannel createTransportChannel(final PlainActionFuture<TestResponse> listener) {
-        return new TransportChannel() {
-
-            @Override
-            public String getProfileName() {
-                return "";
-            }
-
-            @Override
-            public void sendResponse(TransportResponse response) {
-                listener.onResponse(((TestResponse) response));
-            }
-
-            @Override
-            public void sendResponse(Exception exception) {
-                listener.onFailure(exception);
-            }
-
-            @Override
-            public String getChannelType() {
-                return "replica_test";
-            }
-        };
+    public TransportChannel createTransportChannel(final PlainActionFuture<TransportResponse> listener) {
+        return new TestTransportChannel(listener);
     }
 
 }