Make TcpTransport#openConnection fully async (#36095)

This is a follow-up to #35144. That commit made the underlying
connection opening process in TcpTransport asynchronous. However the
method still blocked on the process being complete before returning.
This commit moves the blocking to the ConnectionManager level. This is
another step towards the top-level TransportService api being async.

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

@@ -22,6 +22,7 @@
 import org.apache.logging.log4j.Logger;
 import org.apache.logging.log4j.message.ParameterizedMessage;
 import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.CheckedBiConsumer;
 import org.elasticsearch.common.lease.Releasable;
@@ -218,7 +219,18 @@ public void close() {
     }
 
     private Transport.Connection internalOpenConnection(DiscoveryNode node, ConnectionProfile connectionProfile) {
-        Transport.Connection connection = transport.openConnection(node, connectionProfile);
+        PlainActionFuture<Transport.Connection> future = PlainActionFuture.newFuture();
+        Releasable pendingConnection = transport.openConnection(node, connectionProfile, future);
+        Transport.Connection connection;
+        try {
+            connection = future.actionGet();
+        } catch (IllegalStateException e) {
+            // If the future was interrupted we must cancel the pending connection to avoid channels leaking
+            if (e.getCause() instanceof InterruptedException) {
+                pendingConnection.close();
+            }
+            throw e;
+        }
         try {
             connectionListener.onConnectionOpened(connection);
         } finally {

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

@@ -27,7 +27,6 @@
 import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.NotifyOnceListener;
-import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.Booleans;
 import org.elasticsearch.common.Strings;
@@ -46,6 +45,7 @@
 import org.elasticsearch.common.io.stream.NamedWriteableRegistry;
 import org.elasticsearch.common.io.stream.ReleasableBytesStreamOutput;
 import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.lease.Releasable;
 import org.elasticsearch.common.metrics.MeanMetric;
 import org.elasticsearch.common.network.CloseableChannel;
 import org.elasticsearch.common.network.NetworkAddress;
@@ -349,34 +349,24 @@ protected ConnectionProfile maybeOverrideConnectionProfile(ConnectionProfile con
     }
 
     @Override
-    public NodeChannels openConnection(DiscoveryNode node, ConnectionProfile connectionProfile) {
-        Objects.requireNonNull(connectionProfile, "connection profile cannot be null");
+    public Releasable openConnection(DiscoveryNode node, ConnectionProfile profile, ActionListener<Transport.Connection> listener) {
+        Objects.requireNonNull(profile, "connection profile cannot be null");
         if (node == null) {
             throw new ConnectTransportException(null, "can't open connection to a null node");
         }
-        connectionProfile = maybeOverrideConnectionProfile(connectionProfile);
+        ConnectionProfile finalProfile = maybeOverrideConnectionProfile(profile);
         closeLock.readLock().lock(); // ensure we don't open connections while we are closing
         try {
             ensureOpen();
-            PlainActionFuture<NodeChannels> connectionFuture = PlainActionFuture.newFuture();
-            List<TcpChannel> pendingChannels = initiateConnection(node, connectionProfile, connectionFuture);
-
-            try {
-                return connectionFuture.actionGet();
-            } catch (IllegalStateException e) {
-                // If the future was interrupted we can close the channels to improve the shutdown of the MockTcpTransport
-                if (e.getCause() instanceof InterruptedException) {
-                    CloseableChannel.closeChannels(pendingChannels, false);
-                }
-                throw e;
-            }
+            List<TcpChannel> pendingChannels = initiateConnection(node, finalProfile, listener);
+            return () -> CloseableChannel.closeChannels(pendingChannels, false);
         } finally {
             closeLock.readLock().unlock();
         }
     }
 
     private List<TcpChannel> initiateConnection(DiscoveryNode node, ConnectionProfile connectionProfile,
-                                                ActionListener<NodeChannels> listener) {
+                                                ActionListener<Transport.Connection> listener) {
         int numConnections = connectionProfile.getNumConnections();
         assert numConnections > 0 : "A connection profile must be configured with at least one connection";
 
@@ -432,7 +422,7 @@ public List<String> getLocalAddresses() {
 
     protected void bindServer(ProfileSettings profileSettings) {
         // Bind and start to accept incoming connections.
-        InetAddress hostAddresses[];
+        InetAddress[] hostAddresses;
         List<String> profileBindHosts = profileSettings.bindHosts;
         try {
             hostAddresses = networkService.resolveBindHostAddresses(profileBindHosts.toArray(Strings.EMPTY_ARRAY));
@@ -1581,11 +1571,11 @@ private final class ChannelsConnectedListener implements ActionListener<Void> {
         private final DiscoveryNode node;
         private final ConnectionProfile connectionProfile;
         private final List<TcpChannel> channels;
-        private final ActionListener<NodeChannels> listener;
+        private final ActionListener<Transport.Connection> listener;
         private final CountDown countDown;
 
         private ChannelsConnectedListener(DiscoveryNode node, ConnectionProfile connectionProfile, List<TcpChannel> channels,
-                                          ActionListener<NodeChannels> listener) {
+                                          ActionListener<Transport.Connection> listener) {
             this.node = node;
             this.connectionProfile = connectionProfile;
             this.channels = channels;

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

@@ -25,6 +25,7 @@
 import org.elasticsearch.common.breaker.CircuitBreaker;
 import org.elasticsearch.common.breaker.NoopCircuitBreaker;
 import org.elasticsearch.common.component.LifecycleComponent;
+import org.elasticsearch.common.lease.Releasable;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Setting.Property;
 import org.elasticsearch.common.transport.BoundTransportAddress;
@@ -86,10 +87,12 @@ default CircuitBreaker getInFlightRequestBreaker() {
     }
 
     /**
-     * Opens a new connection to the given node and returns it. The returned connection is not managed by
-     * the transport implementation. This connection must be closed once it's not needed anymore.
+     * Opens a new connection to the given node. When the connection is fully connected, the listener is
+     * called. A {@link Releasable} is returned representing the pending connection. If the caller of this
+     * method decides to move on before the listener is called with the completed connection, they should
+     * release the pending connection to prevent hanging connections.
      */
-    Connection openConnection(DiscoveryNode node, ConnectionProfile profile);
+    Releasable openConnection(DiscoveryNode node, ConnectionProfile profile, ActionListener<Transport.Connection> listener);
 
     TransportStats getStats();
 

server/src/test/java/org/elasticsearch/client/transport/FailAndRetryMockTransport.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.client.transport;
 
 import org.elasticsearch.Version;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.admin.cluster.node.liveness.LivenessResponse;
 import org.elasticsearch.action.admin.cluster.node.liveness.TransportLivenessAction;
 import org.elasticsearch.action.admin.cluster.state.ClusterStateAction;
@@ -30,6 +31,7 @@
 import org.elasticsearch.common.collect.MapBuilder;
 import org.elasticsearch.common.component.Lifecycle;
 import org.elasticsearch.common.component.LifecycleListener;
+import org.elasticsearch.common.lease.Releasable;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.transport.BoundTransportAddress;
 import org.elasticsearch.common.transport.TransportAddress;
@@ -79,8 +81,8 @@ abstract class FailAndRetryMockTransport<Response extends TransportResponse> imp
     protected abstract ClusterState getMockClusterState(DiscoveryNode node);
 
     @Override
-    public Connection openConnection(DiscoveryNode node, ConnectionProfile profile) {
-        return new CloseableConnection() {
+    public Releasable openConnection(DiscoveryNode node, ConnectionProfile profile, ActionListener<Connection> connectionListener) {
+        connectionListener.onResponse(new CloseableConnection() {
 
             @Override
             public DiscoveryNode getNode() {
@@ -134,7 +136,9 @@ public void sendRequest(long requestId, String action, TransportRequest request,
                     }
                 }
             }
-        };
+        });
+
+        return () -> {};
     }
 
     protected abstract Response newResponse();

server/src/test/java/org/elasticsearch/client/transport/TransportClientNodesServiceTests.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.action.admin.cluster.state.ClusterStateAction;
 import org.elasticsearch.action.admin.cluster.state.ClusterStateRequest;
 import org.elasticsearch.action.admin.cluster.state.ClusterStateResponse;
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.cluster.ClusterName;
 import org.elasticsearch.cluster.ClusterState;
 import org.elasticsearch.cluster.node.DiscoveryNode;
@@ -167,7 +168,9 @@ public <T extends TransportResponse> void sendRequest(Transport.Connection conne
             transportService.addNodeConnectedBehavior((connectionManager, discoveryNode) -> false);
             transportService.addGetConnectionBehavior((connectionManager, discoveryNode) -> {
                 // The FailAndRetryTransport does not use the connection profile
-                return transport.openConnection(discoveryNode, null);
+                PlainActionFuture<Transport.Connection> future = PlainActionFuture.newFuture();
+                transport.openConnection(discoveryNode, null, future);
+                return future.actionGet();
             });
             transportService.start();
             transportService.acceptIncomingRequests();
@@ -358,11 +361,19 @@ public void testSniffNodesSamplerClosesConnections() throws Exception {
             try (MockTransportService clientService = createNewService(clientSettings, Version.CURRENT, threadPool, null)) {
                 final List<Transport.Connection> establishedConnections = new CopyOnWriteArrayList<>();
 
-                clientService.addConnectBehavior(remoteService, (transport, discoveryNode, profile) -> {
-                    Transport.Connection connection = transport.openConnection(discoveryNode, profile);
-                    establishedConnections.add(connection);
-                    return connection;
-                });
+                clientService.addConnectBehavior(remoteService, (transport, discoveryNode, profile, listener) ->
+                    transport.openConnection(discoveryNode, profile, new ActionListener<Transport.Connection>() {
+                        @Override
+                        public void onResponse(Transport.Connection connection) {
+                            establishedConnections.add(connection);
+                            listener.onResponse(connection);
+                        }
+
+                        @Override
+                        public void onFailure(Exception e) {
+                            listener.onFailure(e);
+                        }
+                    }));
 
 
                 clientService.start();

server/src/test/java/org/elasticsearch/cluster/NodeConnectionsServiceTests.java

@@ -26,6 +26,7 @@
 import org.elasticsearch.common.UUIDs;
 import org.elasticsearch.common.component.Lifecycle;
 import org.elasticsearch.common.component.LifecycleListener;
+import org.elasticsearch.common.lease.Releasable;
 import org.elasticsearch.common.settings.ClusterSettings;
 import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.common.transport.BoundTransportAddress;
@@ -187,8 +188,6 @@ public HandshakeResponse handshake(Transport.Connection connection, long timeout
     private final class MockTransport implements Transport {
         private ResponseHandlers responseHandlers = new ResponseHandlers();
         private volatile boolean randomConnectionExceptions = false;
-        private TransportMessageListener listener = new TransportMessageListener() {
-        };
 
         @Override
         public <Request extends TransportRequest> void registerRequestHandler(RequestHandlerRegistry<Request> reg) {
@@ -201,7 +200,6 @@ public RequestHandlerRegistry getRequestHandler(String action) {
 
         @Override
         public void addMessageListener(TransportMessageListener listener) {
-            this.listener = listener;
         }
 
         @Override
@@ -225,13 +223,14 @@ public TransportAddress[] addressesFromString(String address, int perAddressLimi
         }
 
         @Override
-        public Connection openConnection(DiscoveryNode node, ConnectionProfile connectionProfile) {
-            if (connectionProfile == null) {
+        public Releasable openConnection(DiscoveryNode node, ConnectionProfile profile, ActionListener<Connection> listener) {
+            if (profile == null) {
                 if (randomConnectionExceptions && randomBoolean()) {
-                    throw new ConnectTransportException(node, "simulated");
+                    listener.onFailure(new ConnectTransportException(node, "simulated"));
+                    return () -> {};
                 }
             }
-            Connection connection = new Connection() {
+            listener.onResponse(new Connection() {
                 @Override
                 public DiscoveryNode getNode() {
                     return node;
@@ -257,8 +256,8 @@ public void close() {
                 public boolean isClosed() {
                     return false;
                 }
-            };
-            return connection;
+            });
+            return () -> {};
         }
 
         @Override

server/src/test/java/org/elasticsearch/transport/ConnectionManagerTests.java

@@ -20,6 +20,7 @@
 package org.elasticsearch.transport;
 
 import org.elasticsearch.Version;
+import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.cluster.node.DiscoveryNode;
 import org.elasticsearch.common.CheckedBiConsumer;
 import org.elasticsearch.common.settings.Settings;
@@ -35,8 +36,10 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
-import static org.mockito.Mockito.when;
 
 public class ConnectionManagerTests extends ESTestCase {
 
@@ -82,7 +85,11 @@ public void onNodeDisconnected(DiscoveryNode node) {
 
         DiscoveryNode node = new DiscoveryNode("", new TransportAddress(InetAddress.getLoopbackAddress(), 0), Version.CURRENT);
         Transport.Connection connection = new TestConnect(node);
-        when(transport.openConnection(node, connectionProfile)).thenReturn(connection);
+        doAnswer(invocationOnMock -> {
+            ActionListener<Transport.Connection> listener = (ActionListener<Transport.Connection>) invocationOnMock.getArguments()[2];
+            listener.onResponse(connection);
+            return null;
+        }).when(transport).openConnection(eq(node), eq(connectionProfile), any(ActionListener.class));
 
         assertFalse(connectionManager.nodeConnected(node));
 
@@ -126,7 +133,11 @@ public void onNodeDisconnected(DiscoveryNode node) {
 
         DiscoveryNode node = new DiscoveryNode("", new TransportAddress(InetAddress.getLoopbackAddress(), 0), Version.CURRENT);
         Transport.Connection connection = new TestConnect(node);
-        when(transport.openConnection(node, connectionProfile)).thenReturn(connection);
+        doAnswer(invocationOnMock -> {
+            ActionListener<Transport.Connection> listener = (ActionListener<Transport.Connection>) invocationOnMock.getArguments()[2];
+            listener.onResponse(connection);
+            return null;
+        }).when(transport).openConnection(eq(node), eq(connectionProfile), any(ActionListener.class));
 
         assertFalse(connectionManager.nodeConnected(node));
 

server/src/test/java/org/elasticsearch/transport/RemoteClusterConnectionTests.java

@@ -1441,7 +1441,7 @@ public static Transport getProxyTransport(ThreadPool threadPool, Map<String, Map
 
         StubbableTransport stubbableTransport = new StubbableTransport(MockTransportService.newMockTransport(Settings.EMPTY, Version
             .CURRENT, threadPool));
-        stubbableTransport.setDefaultConnectBehavior((t, node,  profile) -> {
+        stubbableTransport.setDefaultConnectBehavior((t, node,  profile, listener) -> {
                 Map<String, DiscoveryNode> proxyMapping = nodeMap.get(node.getAddress().toString());
                 if (proxyMapping == null) {
                     throw new IllegalStateException("no proxy mapping for node: " + node);
@@ -1455,34 +1455,44 @@ public static Transport getProxyTransport(ThreadPool threadPool, Map<String, Map
                     // route by seed hostname
                     proxyNode = proxyMapping.get(node.getHostName());
                 }
-                Transport.Connection connection = t.openConnection(proxyNode, profile);
-                return new Transport.Connection() {
+                return t.openConnection(proxyNode, profile, new ActionListener<Transport.Connection>() {
                     @Override
-                    public DiscoveryNode getNode() {
-                        return node;
-                    }
+                    public void onResponse(Transport.Connection connection) {
+                        Transport.Connection proxyConnection = new Transport.Connection() {
+                            @Override
+                            public DiscoveryNode getNode() {
+                                return node;
+                            }
 
-                    @Override
-                    public void sendRequest(long requestId, String action, TransportRequest request, TransportRequestOptions options)
-                        throws IOException, TransportException {
-                        connection.sendRequest(requestId, action, request, options);
-                    }
+                            @Override
+                            public void sendRequest(long requestId, String action, TransportRequest request,
+                                                    TransportRequestOptions options) throws IOException, TransportException {
+                                connection.sendRequest(requestId, action, request, options);
+                            }
 
-                    @Override
-                    public void addCloseListener(ActionListener<Void> listener) {
-                        connection.addCloseListener(listener);
-                    }
+                            @Override
+                            public void addCloseListener(ActionListener<Void> listener) {
+                                connection.addCloseListener(listener);
+                            }
 
-                    @Override
-                    public boolean isClosed() {
-                        return connection.isClosed();
+                            @Override
+                            public boolean isClosed() {
+                                return connection.isClosed();
+                            }
+
+                            @Override
+                            public void close() {
+                                connection.close();
+                            }
+                        };
+                        listener.onResponse(proxyConnection);
                     }
 
                     @Override
-                    public void close() {
-                        connection.close();
+                    public void onFailure(Exception e) {
+                        listener.onFailure(e);
                     }
-                };
+                });
             });
         return stubbableTransport;
     }