Notify onConnectionClosed rather than onNodeDisconnect to prune trans…

…port handlers (#24639)

Today we prune transport handlers in TransportService when a node is disconnected.
This can cause connections to starve in the TransportService if the connection is
opened as a short living connection ie. without sharing the connection to a node
via registering in the transport itself. This change now moves to pruning based
on the connections cache key to ensure we notify handlers as soon as the connection
is closed for all connections not just for registered connections.

Relates to #24632
Relates to #24575
Relates to #24557

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

@@ -100,6 +100,7 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Consumer;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
@@ -368,8 +369,9 @@ public final class NodeChannels implements Connection {
         private final DiscoveryNode node;
         private final AtomicBoolean closed = new AtomicBoolean(false);
         private final Version version;
+        private final Consumer<Connection> onClose;
 
-        public NodeChannels(DiscoveryNode node, Channel[] channels, ConnectionProfile connectionProfile) {
+        public NodeChannels(DiscoveryNode node, Channel[] channels, ConnectionProfile connectionProfile, Consumer<Connection> onClose) {
             this.node = node;
             this.channels = channels;
             assert channels.length == connectionProfile.getNumConnections() : "expected channels size to be == "
@@ -380,6 +382,7 @@ public NodeChannels(DiscoveryNode node, Channel[] channels, ConnectionProfile co
                 typeMapping.put(type, handle);
             }
             version = node.getVersion();
+            this.onClose = onClose;
         }
 
         NodeChannels(NodeChannels channels, Version handshakeVersion) {
@@ -388,6 +391,7 @@ public NodeChannels(DiscoveryNode node, Channel[] channels, ConnectionProfile co
             this.typeMapping = channels.typeMapping;
             this.version = handshakeVersion;
             assert handshakeVersion != null : "handshakeVersion must not be null";
+            this.onClose = channels.onClose;
         }
 
         @Override
@@ -419,7 +423,11 @@ public Channel channel(TransportRequestOptions.Type type) {
         @Override
         public synchronized void close() throws IOException {
             if (closed.compareAndSet(false, true)) {
-                closeChannels(Arrays.stream(channels).filter(Objects::nonNull).collect(Collectors.toList()));
+                try {
+                    closeChannels(Arrays.stream(channels).filter(Objects::nonNull).collect(Collectors.toList()));
+                } finally {
+                    onClose.accept(this);
+                }
             }
         }
 
@@ -531,7 +539,7 @@ public final NodeChannels openConnection(DiscoveryNode node, ConnectionProfile c
                 final TimeValue handshakeTimeout = connectionProfile.getHandshakeTimeout() == null ?
                     connectTimeout : connectionProfile.getHandshakeTimeout();
                 final Version version = executeHandshake(node, channel, handshakeTimeout);
-                transportServiceAdapter.onConnectionOpened(node);
+                transportServiceAdapter.onConnectionOpened(nodeChannels);
                 if (version == null) {
                     // if we are talking to a pre 5.2 node we won't be able to retrieve the version since it doesn't implement the handshake
                     // we do since 5.2 - in this case we just go with the version provided by the node.
@@ -541,6 +549,7 @@ public final NodeChannels openConnection(DiscoveryNode node, ConnectionProfile c
                     nodeChannels = new NodeChannels(nodeChannels, version); // clone the channels - we now have the correct version
                     success = true;
                 }
+                success = true;
                 return nodeChannels;
             } catch (ConnectTransportException e) {
                 throw e;

core/src/main/java/org/elasticsearch/transport/Transport.java

@@ -140,5 +140,13 @@ void sendRequest(long requestId, String action, TransportRequest request, Transp
         default Version getVersion() {
             return getNode().getVersion();
         }
+
+        /**
+         * Returns a key that this connection can be cached on. Delegating subclasses must delegate method call to
+         * the original connection.
+         */
+        default Object getCacheKey() {
+            return this;
+        }
     }
 }

core/src/main/java/org/elasticsearch/transport/TransportConnectionListener.java

@@ -36,8 +36,14 @@ default void onNodeConnected(DiscoveryNode node) {}
      */
     default void onNodeDisconnected(DiscoveryNode node) {}
 
+    /**
+     * Called once a node connection is closed. The connection might not have been registered in the
+     * transport as a shared connection to a specific node
+     */
+    default void onConnectionClosed(Transport.Connection connection) {}
+
     /**
      * Called once a node connection is opened.
      */
-    default void onConnectionOpened(DiscoveryNode node) {}
+    default void onConnectionOpened(Transport.Connection connection) {}
 }

core/src/main/java/org/elasticsearch/transport/TransportService.java

@@ -556,7 +556,7 @@ private <T extends TransportResponse> void sendRequestInternal(final Transport.C
             }
             Supplier<ThreadContext.StoredContext> storedContextSupplier = threadPool.getThreadContext().newRestorableContext(true);
             TransportResponseHandler<T> responseHandler = new ContextRestoreResponseHandler<>(storedContextSupplier, handler);
-            clientHandlers.put(requestId, new RequestHolder<>(responseHandler, connection.getNode(), action, timeoutHandler));
+            clientHandlers.put(requestId, new RequestHolder<>(responseHandler, connection, action, timeoutHandler));
             if (lifecycle.stoppedOrClosed()) {
                 // if we are not started the exception handling will remove the RequestHolder again and calls the handler to notify
                 // the caller. It will only notify if the toStop code hasn't done the work yet.
@@ -797,7 +797,7 @@ public TransportResponseHandler onResponseReceived(final long requestId) {
             }
             holder.cancelTimeout();
             if (traceEnabled() && shouldTraceAction(holder.action())) {
-                traceReceivedResponse(requestId, holder.node(), holder.action());
+                traceReceivedResponse(requestId, holder.connection().getNode(), holder.action());
             }
             return holder.handler();
         }
@@ -842,12 +842,12 @@ public void onNodeConnected(final DiscoveryNode node) {
         }
 
         @Override
-        public void onConnectionOpened(DiscoveryNode node) {
+        public void onConnectionOpened(Transport.Connection connection) {
             // capture listeners before spawning the background callback so the following pattern won't trigger a call
             // connectToNode(); connection is completed successfully
             // addConnectionListener(); this listener shouldn't be called
             final Stream<TransportConnectionListener> listenersToNotify = TransportService.this.connectionListeners.stream();
-            threadPool.generic().execute(() -> listenersToNotify.forEach(listener -> listener.onConnectionOpened(node)));
+            threadPool.generic().execute(() -> listenersToNotify.forEach(listener -> listener.onConnectionOpened(connection)));
         }
 
         @Override
@@ -858,20 +858,28 @@ public void onNodeDisconnected(final DiscoveryNode node) {
                         connectionListener.onNodeDisconnected(node);
                     }
                 });
+            } catch (EsRejectedExecutionException ex) {
+                logger.debug("Rejected execution on NodeDisconnected", ex);
+            }
+        }
+
+        @Override
+        public void onConnectionClosed(Transport.Connection connection) {
+            try {
                 for (Map.Entry<Long, RequestHolder> entry : clientHandlers.entrySet()) {
                     RequestHolder holder = entry.getValue();
-                    if (holder.node().equals(node)) {
+                    if (holder.connection().getCacheKey().equals(connection.getCacheKey())) {
                         final RequestHolder holderToNotify = clientHandlers.remove(entry.getKey());
                         if (holderToNotify != null) {
                             // callback that an exception happened, but on a different thread since we don't
                             // want handlers to worry about stack overflows
-                            threadPool.generic().execute(() -> holderToNotify.handler().handleException(new NodeDisconnectedException(node,
-                                holderToNotify.action())));
+                            threadPool.generic().execute(() -> holderToNotify.handler().handleException(new NodeDisconnectedException(
+                                connection.getNode(), holderToNotify.action())));
                         }
                     }
                 }
             } catch (EsRejectedExecutionException ex) {
-                logger.debug("Rejected execution on NodeDisconnected", ex);
+                logger.debug("Rejected execution on onConnectionClosed", ex);
             }
         }
 
@@ -916,13 +924,14 @@ public void run() {
             if (holder != null) {
                 // add it to the timeout information holder, in case we are going to get a response later
                 long timeoutTime = System.currentTimeMillis();
-                timeoutInfoHandlers.put(requestId, new TimeoutInfoHolder(holder.node(), holder.action(), sentTime, timeoutTime));
+                timeoutInfoHandlers.put(requestId, new TimeoutInfoHolder(holder.connection().getNode(), holder.action(), sentTime,
+                    timeoutTime));
                 // now that we have the information visible via timeoutInfoHandlers, we try to remove the request id
                 final RequestHolder removedHolder = clientHandlers.remove(requestId);
                 if (removedHolder != null) {
                     assert removedHolder == holder : "two different holder instances for request [" + requestId + "]";
                     removedHolder.handler().handleException(
-                        new ReceiveTimeoutTransportException(holder.node(), holder.action(),
+                        new ReceiveTimeoutTransportException(holder.connection().getNode(), holder.action(),
                             "request_id [" + requestId + "] timed out after [" + (timeoutTime - sentTime) + "ms]"));
                 } else {
                     // response was processed, remove timeout info.
@@ -977,15 +986,15 @@ static class RequestHolder<T extends TransportResponse> {
 
         private final TransportResponseHandler<T> handler;
 
-        private final DiscoveryNode node;
+        private final Transport.Connection connection;
 
         private final String action;
 
         private final TimeoutHandler timeoutHandler;
 
-        RequestHolder(TransportResponseHandler<T> handler, DiscoveryNode node, String action, TimeoutHandler timeoutHandler) {
+        RequestHolder(TransportResponseHandler<T> handler, Transport.Connection connection, String action, TimeoutHandler timeoutHandler) {
             this.handler = handler;
-            this.node = node;
+            this.connection = connection;
             this.action = action;
             this.timeoutHandler = timeoutHandler;
         }
@@ -994,8 +1003,8 @@ public TransportResponseHandler<T> handler() {
             return handler;
         }
 
-        public DiscoveryNode node() {
-            return this.node;
+        public Transport.Connection connection() {
+            return this.connection;
         }
 
         public String action() {

core/src/main/java/org/elasticsearch/transport/local/LocalTransport.java

@@ -193,6 +193,8 @@ public void disconnectFromNode(DiscoveryNode node) {
             LocalTransport removed = connectedNodes.remove(node);
             if (removed != null) {
                 transportServiceAdapter.onNodeDisconnected(node);
+                // here we notify that the connection to the node is closed -- use the remove transport as cache key
+                transportServiceAdapter.onConnectionClosed(getConnectionForTransport(removed, node, true));
             }
         }
     }
@@ -208,7 +210,7 @@ public Connection getConnection(DiscoveryNode node) {
         if (targetTransport == null) {
             throw new NodeNotConnectedException(node, "Node not connected");
         }
-        return getConnectionForTransport(targetTransport, node);
+        return getConnectionForTransport(targetTransport, node, true);
     }
 
     @Override
@@ -217,10 +219,10 @@ public Connection openConnection(DiscoveryNode node, ConnectionProfile profile)
         if (targetTransport == null) {
             throw new ConnectTransportException(node, "Failed to connect");
         }
-        return getConnectionForTransport(targetTransport, node);
+        return getConnectionForTransport(targetTransport, node, false);
     }
 
-    private Connection getConnectionForTransport(LocalTransport targetTransport, DiscoveryNode node) {
+    private Connection getConnectionForTransport(LocalTransport targetTransport, DiscoveryNode node, boolean registered) {
         return new Connection() {
             @Override
             public DiscoveryNode getNode() {
@@ -242,7 +244,16 @@ public void sendRequest(long requestId, String action, TransportRequest request,
             }
 
             @Override
-            public void close() throws IOException {}
+            public Object getCacheKey() {
+                // the cache key depends if it's registered connection or a single connection (not shared)
+                // we use the actual LocalTransport if it's a regsitered connection otherwise we use the identity
+                return registered ? targetTransport : this;
+            }
+
+            @Override
+            public void close() throws IOException {
+                transportServiceAdapter.onConnectionClosed(this);
+            }
         };
     }
 

core/src/test/java/org/elasticsearch/discovery/zen/UnicastZenPingTests.java

@@ -652,8 +652,8 @@ public void testResolveReuseExistingNodeConnections() throws ExecutionException,
         // install a listener to check that no new connections are made
         handleA.transportService.addConnectionListener(new TransportConnectionListener() {
             @Override
-            public void onConnectionOpened(DiscoveryNode node) {
-                fail("should not open any connections. got [" + node + "]");
+            public void onConnectionOpened(Transport.Connection connection) {
+                fail("should not open any connections. got [" + connection.getNode() + "]");
             }
         });
 

core/src/test/java/org/elasticsearch/transport/TCPTransportTests.java

@@ -197,7 +197,7 @@ protected void sendMessage(Object o, BytesReference reference, Runnable sendList
 
                 @Override
                 protected NodeChannels connectToChannels(DiscoveryNode node, ConnectionProfile profile) throws IOException {
-                    return new NodeChannels(node, new Object[profile.getNumConnections()], profile);
+                    return new NodeChannels(node, new Object[profile.getNumConnections()], profile, c -> {});
                 }
 
                 @Override
@@ -213,7 +213,7 @@ public long serverOpen() {
                 @Override
                 public NodeChannels getConnection(DiscoveryNode node) {
                     return new NodeChannels(node, new Object[MockTcpTransport.LIGHT_PROFILE.getNumConnections()],
-                        MockTcpTransport.LIGHT_PROFILE);
+                        MockTcpTransport.LIGHT_PROFILE, c -> {});
                 }
             };
             DiscoveryNode node = new DiscoveryNode("foo", new LocalTransportAddress("1"), Version.CURRENT);

modules/transport-netty3/src/main/java/org/elasticsearch/transport/netty3/Netty3Transport.java

@@ -347,7 +347,7 @@ public long serverOpen() {
     @Override
     protected NodeChannels connectToChannels(DiscoveryNode node, ConnectionProfile profile) {
         final Channel[] channels = new Channel[profile.getNumConnections()];
-        final NodeChannels nodeChannels = new NodeChannels(node, channels, profile);
+        final NodeChannels nodeChannels = new NodeChannels(node, channels, profile, transportServiceAdapter::onConnectionClosed);
         boolean success = false;
         try {
             final TimeValue connectTimeout;

modules/transport-netty4/src/main/java/org/elasticsearch/transport/netty4/Netty4Transport.java

@@ -331,7 +331,7 @@ public long serverOpen() {
     @Override
     protected NodeChannels connectToChannels(DiscoveryNode node, ConnectionProfile profile) {
         final Channel[] channels = new Channel[profile.getNumConnections()];
-        final NodeChannels nodeChannels = new NodeChannels(node, channels, profile);
+        final NodeChannels nodeChannels = new NodeChannels(node, channels, profile, transportServiceAdapter::onConnectionClosed);
         boolean success = false;
         try {
             final TimeValue connectTimeout;

test/framework/src/main/java/org/elasticsearch/test/transport/MockTransportService.java

@@ -820,6 +820,11 @@ public void sendRequest(long requestId, String action, TransportRequest request,
         public void close() throws IOException {
             connection.close();
         }
+
+        @Override
+        public Object getCacheKey() {
+            return connection.getCacheKey();
+        }
     }
 
     public Transport getOriginalTransport() {

test/framework/src/main/java/org/elasticsearch/transport/AbstractSimpleTransportTestCase.java

@@ -2106,9 +2106,6 @@ public void handleException(TransportException exp) {
 
             @Override
             public String executor() {
-                if (1 == 1)
-                    return "same";
-
                 return randomFrom(executors);
             }
         };
@@ -2118,4 +2115,59 @@ public String executor() {
         latch.await();
     }
 
+    public void testHandlerIsInvokedOnConnectionClose() throws IOException, InterruptedException {
+        List<String> executors = new ArrayList<>(ThreadPool.THREAD_POOL_TYPES.keySet());
+        CollectionUtil.timSort(executors); // makes sure it's reproducible
+        TransportService serviceC = build(Settings.builder().put("name", "TS_TEST").build(), version0, null, true);
+        serviceC.registerRequestHandler("action", TestRequest::new, ThreadPool.Names.SAME,
+            (request, channel) -> {
+                // do nothing
+            });
+        serviceC.start();
+        serviceC.acceptIncomingRequests();
+        CountDownLatch latch = new CountDownLatch(1);
+        TransportResponseHandler<TransportResponse> transportResponseHandler = new TransportResponseHandler<TransportResponse>() {
+            @Override
+            public TransportResponse newInstance() {
+                return TransportResponse.Empty.INSTANCE;
+            }
+
+            @Override
+            public void handleResponse(TransportResponse response) {
+                try {
+                    fail("no response expected");
+                } finally {
+                    latch.countDown();
+                }
+            }
+
+            @Override
+            public void handleException(TransportException exp) {
+                try {
+                    assertTrue(exp.getClass().toString(), exp instanceof NodeDisconnectedException);
+                } finally {
+                    latch.countDown();
+                }
+            }
+
+            @Override
+            public String executor() {
+                return randomFrom(executors);
+            }
+        };
+        ConnectionProfile.Builder builder = new ConnectionProfile.Builder();
+        builder.addConnections(1,
+            TransportRequestOptions.Type.BULK,
+            TransportRequestOptions.Type.PING,
+            TransportRequestOptions.Type.RECOVERY,
+            TransportRequestOptions.Type.REG,
+            TransportRequestOptions.Type.STATE);
+        Transport.Connection connection = serviceB.openConnection(serviceC.getLocalNode(), builder.build());
+        serviceB.sendRequest(connection, "action",  new TestRequest(randomFrom("fail", "pass")), TransportRequestOptions.EMPTY,
+            transportResponseHandler);
+        connection.close();
+        latch.await();
+        serviceC.close();
+    }
+
 }

test/framework/src/main/java/org/elasticsearch/transport/MockTcpTransport.java

@@ -182,7 +182,8 @@ private void readMessage(MockChannel mockChannel, StreamInput input) throws IOEx
     @Override
     protected NodeChannels connectToChannels(DiscoveryNode node, ConnectionProfile profile) throws IOException {
         final MockChannel[] mockChannels = new MockChannel[1];
-        final NodeChannels nodeChannels = new NodeChannels(node, mockChannels, LIGHT_PROFILE); // we always use light here
+        final NodeChannels nodeChannels = new NodeChannels(node, mockChannels, LIGHT_PROFILE,
+            transportServiceAdapter::onConnectionClosed); // we always use light here
         boolean success = false;
         final Socket socket = new Socket();
         try {