fix targetsbynode logic

Signed-off-by: Rahul Karajgikar <[email protected]>
opensearch-project · Sep 6, 2024 · 9e322c4 · 9e322c4
1 parent 1b9ae7c
commit 9e322c4
Show file tree

Hide file tree

Showing 6 changed files with 23 additions and 2 deletions.
diff --git a/server/src/main/java/org/opensearch/cluster/NodeConnectionsService.java b/server/src/main/java/org/opensearch/cluster/NodeConnectionsService.java
@@ -182,9 +182,9 @@ public void disconnectFromNodesExcept(DiscoveryNodes discoveryNodes) {
             // to avoid permanently blocking node joins
             // This situation should ideally not happen, this is just for extra safety
             transportService.removePendingDisconnections(
-                targetsByNode.keySet()
+                transportService.getPendingDisconnections()
                     .stream()
-                    .filter(discoveryNode -> !discoveryNodes.nodeExists(discoveryNode))
+                    .filter(discoveryNode -> !discoveryNodes.nodeExists(discoveryNode) && !targetsByNode.containsKey(discoveryNode))
                     .collect(Collectors.toSet())
             );
         }

diff --git a/server/src/main/java/org/opensearch/transport/ClusterConnectionManager.java b/server/src/main/java/org/opensearch/transport/ClusterConnectionManager.java
@@ -247,6 +247,11 @@ public void disconnectFromNode(DiscoveryNode node) {
         logger.debug("Removed node [{}] from pending disconnections list", node);
     }
 
+    @Override
+    public Set<DiscoveryNode> getPendingDisconnections() {
+        return pendingDisconnections;
+    }
+
     @Override
     public void setPendingDisconnection(DiscoveryNode node) {
         logger.debug("marking disconnection as pending for node: [{}]", node);

diff --git a/server/src/main/java/org/opensearch/transport/ConnectionManager.java b/server/src/main/java/org/opensearch/transport/ConnectionManager.java
@@ -65,6 +65,8 @@ void connectToNode(
 
     void disconnectFromNode(DiscoveryNode node);
 
+    Set<DiscoveryNode> getPendingDisconnections();
+
     void setPendingDisconnection(DiscoveryNode node);
 
     void removePendingDisconnection(DiscoveryNode node);

diff --git a/server/src/main/java/org/opensearch/transport/RemoteConnectionManager.java b/server/src/main/java/org/opensearch/transport/RemoteConnectionManager.java
@@ -114,6 +114,11 @@ public void disconnectFromNode(DiscoveryNode node) {
         delegate.disconnectFromNode(node);
     }
 
+    @Override
+    public Set<DiscoveryNode> getPendingDisconnections() {
+        return delegate.getPendingDisconnections();
+    }
+
     @Override
     public void setPendingDisconnection(DiscoveryNode node) {
         delegate.setPendingDisconnection(node);

diff --git a/server/src/main/java/org/opensearch/transport/TransportService.java b/server/src/main/java/org/opensearch/transport/TransportService.java
@@ -773,6 +773,10 @@ public void disconnectFromNode(DiscoveryNode node) {
         connectionManager.disconnectFromNode(node);
     }
 
+    public Set<DiscoveryNode> getPendingDisconnections() {
+        return connectionManager.getPendingDisconnections();
+    }
+
     public void setPendingDisconnections(Set<DiscoveryNode> nodes) {
         nodes.forEach(connectionManager::setPendingDisconnection);
     }

diff --git a/test/framework/src/main/java/org/opensearch/test/transport/StubbableConnectionManager.java b/test/framework/src/main/java/org/opensearch/test/transport/StubbableConnectionManager.java
@@ -123,6 +123,11 @@ public void disconnectFromNode(DiscoveryNode node) {
         delegate.disconnectFromNode(node);
     }
 
+    @Override
+    public Set<DiscoveryNode> getPendingDisconnections() {
+        return delegate.getPendingDisconnections();
+    }
+
     @Override
     public void setPendingDisconnection(DiscoveryNode node) {
         delegate.setPendingDisconnection(node);