fix targetsbynode logic

Signed-off-by: Rahul Karajgikar <karajgik@amazon.com>

Author: Rahul Karajgikar

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())
             );
         }

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);

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);

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);

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);
     }

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);