Add infrastructure to manage network connections outside of Transport/TransportService (#22194)

Some expert users like UnicastZenPing today establishes real connections to nodes during it's ping
phase that can be used by other parts of the system. Yet, this is potentially dangerous
and undesirable unless the nodes have been fully verified and should be connected to in the
case of a cluster state update or if we join a newly elected master. For use-cases like this, this change adds the infrastructure to manually handle connections that are not publicly available on the node ie. should not be managed by `Transport`/`TransportSerivce`

Author: s1monw

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

@@ -294,6 +294,9 @@ public List<String> getLocalAddresses() {
         return transport.getLocalAddresses();
     }
 
+    /**
+     * Returns <code>true</code> iff the given node is already connected.
+     */
     public boolean nodeConnected(DiscoveryNode node) {
         return node.equals(localNode) || transport.nodeConnected(node);
     }
@@ -315,6 +318,20 @@ public void connectToNode(final DiscoveryNode node, ConnectionProfile connection
         transport.connectToNode(node, connectionProfile);
     }
 
+    /**
+     * Establishes and returns a new connection to the given node. The connection is NOT maintained by this service, it's the callers
+     * responsibility to close the connection once it goes out of scope.
+     * @param node the node to connect to
+     * @param profile the connection profile to use
+     */
+    public Transport.Connection openConnection(final DiscoveryNode node, ConnectionProfile profile) throws IOException {
+        if (node.equals(localNode)) {
+            return localNodeConnection;
+        } else {
+            return transport.openConnection(node, profile);
+        }
+    }
+
     /**
      * Lightly connect to the specified node, returning updated node
      * information. The handshake will fail if the cluster name on the
@@ -341,7 +358,19 @@ public DiscoveryNode connectToNodeAndHandshake(
         return handshakeNode;
     }
 
-    private DiscoveryNode handshake(
+    /**
+     * Executes a high-level handshake using the given connection
+     * and returns the discovery node of the node the connection
+     * was established with. The handshake will fail if the cluster
+     * name on the target node mismatches the local cluster name.
+     *
+     * @param connection       the connection to a specific node
+     * @param handshakeTimeout handshake timeout
+     * @return the connected node
+     * @throws ConnectTransportException if the connection failed
+     * @throws IllegalStateException if the handshake failed
+     */
+    public DiscoveryNode handshake(
             final Transport.Connection connection,
             final long handshakeTimeout) throws ConnectTransportException {
         final HandshakeResponse response;
@@ -474,7 +503,7 @@ public final <T extends TransportResponse> void sendRequest(final DiscoveryNode
         }
     }
 
-    final <T extends TransportResponse> void sendRequest(final Transport.Connection connection, final String action,
+    public final <T extends TransportResponse> void sendRequest(final Transport.Connection connection, final String action,
                                                                 final TransportRequest request,
                                                                 final TransportRequestOptions options,
                                                                 TransportResponseHandler<T> handler) {
@@ -486,7 +515,7 @@ final <T extends TransportResponse> void sendRequest(final Transport.Connection
      * Returns either a real transport connection or a local node connection if we are using the local node optimization.
      * @throws NodeNotConnectedException if the given node is not connected
      */
-    private Transport.Connection getConnection(DiscoveryNode node) {
+    public Transport.Connection getConnection(DiscoveryNode node) {
         if (Objects.requireNonNull(node, "node must be non-null").equals(localNode)) {
             return localNodeConnection;
         } else {

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

@@ -41,6 +41,7 @@
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.test.VersionUtils;
 import org.elasticsearch.test.junit.annotations.TestLogging;
+import org.elasticsearch.test.transport.MockTransportService;
 import org.elasticsearch.threadpool.TestThreadPool;
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.transport.MockTcpTransport;
@@ -573,7 +574,7 @@ private NetworkHandle startServices(
         final BiFunction<Settings, Version, Transport> supplier) {
         final Transport transport = supplier.apply(settings, version);
         final TransportService transportService =
-            new TransportService(settings, transport, threadPool, TransportService.NOOP_TRANSPORT_INTERCEPTOR, null);
+            new MockTransportService(settings, transport, threadPool, TransportService.NOOP_TRANSPORT_INTERCEPTOR, null);
         transportService.start();
         transportService.acceptIncomingRequests();
         final ConcurrentMap<TransportAddress, AtomicInteger> counters = ConcurrentCollections.newConcurrentMap();

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

@@ -113,14 +113,24 @@ public void testConnectToNodeLight() throws IOException {
             emptyMap(),
             emptySet(),
             Version.CURRENT.minimumCompatibilityVersion());
+        try (Transport.Connection connection = handleA.transportService.openConnection(discoveryNode, ConnectionProfile.LIGHT_PROFILE)){
+            DiscoveryNode connectedNode = handleA.transportService.handshake(connection, timeout);
+            assertNotNull(connectedNode);
+            // the name and version should be updated
+            assertEquals(connectedNode.getName(), "TS_B");
+            assertEquals(connectedNode.getVersion(), handleB.discoveryNode.getVersion());
+            assertFalse(handleA.transportService.nodeConnected(discoveryNode));
+        }
+
         DiscoveryNode connectedNode =
-                handleA.transportService.connectToNodeAndHandshake(discoveryNode, timeout);
+            handleA.transportService.connectToNodeAndHandshake(discoveryNode, timeout);
         assertNotNull(connectedNode);
 
         // the name and version should be updated
         assertEquals(connectedNode.getName(), "TS_B");
         assertEquals(connectedNode.getVersion(), handleB.discoveryNode.getVersion());
         assertTrue(handleA.transportService.nodeConnected(discoveryNode));
+
     }
 
     public void testMismatchedClusterName() {
@@ -133,8 +143,12 @@ public void testMismatchedClusterName() {
             emptyMap(),
             emptySet(),
             Version.CURRENT.minimumCompatibilityVersion());
-        IllegalStateException ex = expectThrows(IllegalStateException.class, () -> handleA.transportService.connectToNodeAndHandshake(
-                discoveryNode, timeout));
+        IllegalStateException ex = expectThrows(IllegalStateException.class, () -> {
+            try (Transport.Connection connection = handleA.transportService.openConnection(discoveryNode,
+                ConnectionProfile.LIGHT_PROFILE)) {
+                handleA.transportService.handshake(connection, timeout);
+            }
+        });
         assertThat(ex.getMessage(), containsString("handshake failed, mismatched cluster name [Cluster [b]]"));
         assertFalse(handleA.transportService.nodeConnected(discoveryNode));
 }
@@ -150,8 +164,12 @@ public void testIncompatibleVersions() {
             emptyMap(),
             emptySet(),
             Version.CURRENT.minimumCompatibilityVersion());
-        IllegalStateException ex = expectThrows(IllegalStateException.class, () -> handleA.transportService.connectToNodeAndHandshake(
-            discoveryNode, timeout));
+        IllegalStateException ex = expectThrows(IllegalStateException.class, () -> {
+            try (Transport.Connection connection = handleA.transportService.openConnection(discoveryNode,
+                ConnectionProfile.LIGHT_PROFILE)) {
+                handleA.transportService.handshake(connection, timeout);
+            }
+        });
         assertThat(ex.getMessage(), containsString("handshake failed, incompatible version"));
         assertFalse(handleA.transportService.nodeConnected(discoveryNode));
     }

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

@@ -56,11 +56,13 @@
 import java.net.UnknownHostException;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Queue;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.LinkedBlockingDeque;
@@ -79,6 +81,8 @@
  */
 public final class MockTransportService extends TransportService {
 
+    private final Map<DiscoveryNode, List<Transport.Connection>> openConnections = new HashMap<>();
+
     public static class TestPlugin extends Plugin {
         @Override
         public List<Setting<?>> getSettings() {
@@ -560,9 +564,7 @@ public void stop() {
         }
 
         @Override
-        public void close() {
-            transport.close();
-        }
+        public void close() { transport.close(); }
 
         @Override
         public Map<String, BoundTransportAddress> profileBoundAddresses() {
@@ -708,4 +710,41 @@ public Transport getOriginalTransport() {
         }
         return transport;
     }
+
+    @Override
+    public Transport.Connection openConnection(DiscoveryNode node, ConnectionProfile profile) throws IOException {
+        FilteredConnection filteredConnection = new FilteredConnection(super.openConnection(node, profile)) {
+            final AtomicBoolean closed = new AtomicBoolean(false);
+            @Override
+            public void close() throws IOException {
+                try {
+                    super.close();
+                } finally {
+                    if (closed.compareAndSet(false, true)) {
+                        synchronized (openConnections) {
+                            List<Transport.Connection> connections = openConnections.get(node);
+                            boolean remove = connections.remove(this);
+                            assert remove;
+                            if (connections.isEmpty()) {
+                                openConnections.remove(node);
+                            }
+                        }
+                    }
+                }
+
+            }
+        };
+        synchronized (openConnections) {
+            List<Transport.Connection> connections = openConnections.computeIfAbsent(node,
+                (n) -> new CopyOnWriteArrayList<>());
+            connections.add(filteredConnection);
+        }
+        return filteredConnection;
+    }
+
+    @Override
+    protected void doClose() {
+        super.doClose();
+        assert openConnections.size() == 0 : "still open connections: " + openConnections;
+    }
 }

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

@@ -1347,8 +1347,8 @@ public void handleException(TransportException exp) {
             // all is well
         }
 
-        try {
-            serviceB.connectToNodeAndHandshake(nodeA, 100);
+        try (Transport.Connection connection = serviceB.openConnection(nodeA, ConnectionProfile.LIGHT_PROFILE)){
+            serviceB.handshake(connection, 100);
             fail("exception should be thrown");
         } catch (IllegalStateException e) {
             // all is well
@@ -1405,8 +1405,8 @@ public void handleException(TransportException exp) {
             // all is well
         }
 
-        try {
-            serviceB.connectToNodeAndHandshake(nodeA, 100);
+        try (Transport.Connection connection = serviceB.openConnection(nodeA, ConnectionProfile.LIGHT_PROFILE)){
+            serviceB.handshake(connection, 100);
             fail("exception should be thrown");
         } catch (IllegalStateException e) {
             // all is well