NETWORKING: MockTransportService Wait for Close (#35038)

* NETWORKING: MockTransportService Wait for Close

* Make `MockTransportService` wait `30s` for close listeners to run before failing the assertion
* Closes #34990

Author: original-brownbear

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

@@ -202,7 +202,7 @@ public void close() {
             threadPool.generic().execute(() -> {
                 closeLock.writeLock().lock();
                 try {
-                    // we are holding a write lock so nobody modifies the connectedNodes / openConnections map - it's safe to first close
+                    // we are holding a write lock so nobody adds to the connectedNodes / openConnections map - it's safe to first close
                     // all instances and then clear them maps
                     Iterator<Map.Entry<DiscoveryNode, Transport.Connection>> iterator = connectedNodes.entrySet().iterator();
                     while (iterator.hasNext()) {

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

@@ -20,6 +20,7 @@
 package org.elasticsearch.test.transport;
 
 import com.carrotsearch.randomizedtesting.SysGlobals;
+import java.util.concurrent.TimeUnit;
 import org.elasticsearch.Version;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.cluster.ClusterModule;
@@ -599,30 +600,37 @@ public Transport.Connection openConnection(DiscoveryNode node, ConnectionProfile
         Transport.Connection connection = super.openConnection(node, profile);
 
         synchronized (openConnections) {
-            List<Transport.Connection> connections = openConnections.computeIfAbsent(node,
-                (n) -> new CopyOnWriteArrayList<>());
-            connections.add(connection);
-        }
-
-        connection.addCloseListener(ActionListener.wrap(() -> {
-            synchronized (openConnections) {
-                List<Transport.Connection> connections = openConnections.get(node);
-                boolean remove = connections.remove(connection);
-                assert remove : "Should have removed connection";
-                if (connections.isEmpty()) {
-                    openConnections.remove(node);
+            openConnections.computeIfAbsent(node, n -> new CopyOnWriteArrayList<>()).add(connection);
+            connection.addCloseListener(ActionListener.wrap(() -> {
+                synchronized (openConnections) {
+                    List<Transport.Connection> connections = openConnections.get(node);
+                    boolean remove = connections.remove(connection);
+                    assert remove : "Should have removed connection";
+                    if (connections.isEmpty()) {
+                        openConnections.remove(node);
+                    }
+                    if (openConnections.isEmpty()) {
+                        openConnections.notifyAll();
+                    }
                 }
-            }
-        }));
+            }));
+        }
 
         return connection;
     }
 
     @Override
     protected void doClose() throws IOException {
         super.doClose();
-        synchronized (openConnections) {
-            assert openConnections.size() == 0 : "still open connections: " + openConnections;
+        try {
+            synchronized (openConnections) {
+                if (openConnections.isEmpty() == false) {
+                    openConnections.wait(TimeUnit.SECONDS.toMillis(30L));
+                }
+                assert openConnections.size() == 0 : "still open connections: " + openConnections;
+            }
+        } catch (InterruptedException e) {
+            throw new IllegalStateException(e);
         }
     }