Notify stubbable transport behaviors on clear (#77774)

If `MockTransportService#addUnresponsiveRule` adds a rule that drops requests to a node then we must notify the rule on removal so that the requests don't leak. Today we notify rules when removing them from a specific address but not if clearing all rules. This commit addresses that. Closes #77751
elastic · Oct 25, 2021 · ad975ca · ad975ca
1 parent a23f58f
commit ad975ca
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 4 deletions.
diff --git a/server/src/internalClusterTest/java/org/elasticsearch/discovery/DiscoveryDisruptionIT.java b/server/src/internalClusterTest/java/org/elasticsearch/discovery/DiscoveryDisruptionIT.java
@@ -154,7 +154,6 @@ public void testElectMasterWithLatestVersion() throws Exception {
      * sure that the node is removed form the cluster, that the node start pinging and that
      * the cluster reforms when healed.
      */
-    @AwaitsFix(bugUrl = "https://github.com/elastic/elasticsearch/issues/77751")
     public void testNodeNotReachableFromMaster() throws Exception {
         startCluster(3);
 

diff --git a/test/framework/src/main/java/org/elasticsearch/test/transport/MockTransportService.java b/test/framework/src/main/java/org/elasticsearch/test/transport/MockTransportService.java
@@ -281,7 +281,8 @@ public void addUnresponsiveRule(TransportAddress transportAddress) {
             listener.onFailure(new ConnectTransportException(discoveryNode, "UNRESPONSIVE: simulated")));
 
         transport().addSendBehavior(transportAddress, new StubbableTransport.SendRequestBehavior() {
-            private Set<Transport.Connection> toClose = ConcurrentHashMap.newKeySet();
+            private final Set<Transport.Connection> toClose = ConcurrentHashMap.newKeySet();
+
             @Override
             public void sendRequest(Transport.Connection connection, long requestId, String action,
                                     TransportRequest request, TransportRequestOptions options) {

diff --git a/test/framework/src/main/java/org/elasticsearch/test/transport/StubbableTransport.java b/test/framework/src/main/java/org/elasticsearch/test/transport/StubbableTransport.java
@@ -29,10 +29,13 @@
 
 import java.io.IOException;
 import java.net.UnknownHostException;
+import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
+import static junit.framework.TestCase.assertTrue;
+
 public class StubbableTransport implements Transport {
 
     private final ConcurrentHashMap<TransportAddress, SendRequestBehavior> sendBehaviors = new ConcurrentHashMap<>();
@@ -94,9 +97,22 @@ void clearInboundBehaviors() {
 
     void clearOutboundBehaviors() {
         this.defaultSendRequest = null;
-        sendBehaviors.clear();
+        final Iterator<SendRequestBehavior> sendBehaviorIterator = sendBehaviors.values().iterator();
+        while (sendBehaviorIterator.hasNext()) {
+            final SendRequestBehavior behavior = sendBehaviorIterator.next();
+            sendBehaviorIterator.remove();
+            behavior.clearCallback();
+        }
+        assertTrue(sendBehaviors.isEmpty());
+
         this.defaultConnectBehavior = null;
-        connectBehaviors.clear();
+        final Iterator<OpenConnectionBehavior> connectBehaviorIterator = connectBehaviors.values().iterator();
+        while (connectBehaviorIterator.hasNext()) {
+            final OpenConnectionBehavior behavior = connectBehaviorIterator.next();
+            connectBehaviorIterator.remove();
+            behavior.clearCallback();
+        }
+        assertTrue(connectBehaviors.isEmpty());
     }
 
     void clearOutboundBehaviors(TransportAddress transportAddress) {