Fix race condition on concurrent modification of alive nodes in NodePool

elastic_transport/_node/_urllib3_chain_certs.py

@@ -15,13 +15,13 @@
 #  specific language governing permissions and limitations
 #  under the License.
 
+import _ssl  # type: ignore
 import hashlib
 import sys
 from binascii import hexlify, unhexlify
 from hmac import compare_digest
 from typing import Any, List, Optional
 
-import _ssl  # type: ignore
 import urllib3
 import urllib3.connection
 

elastic_transport/_node_pool.py

@@ -358,9 +358,10 @@ def get(self) -> BaseNode:
             return self._all_nodes[self._seed_nodes[0]]
 
         # Filter nodes in 'alive_nodes' to ones not marked as removed.
+        # Iterate over a copy as other threads may mark nodes dead or alive meanwhile.
         nodes = [
             node
-            for node_config, node in self._alive_nodes.items()
+            for node_config, node in self._alive_nodes.copy().items()
             if node_config not in self._removed_nodes
         ]
 

tests/test_node_pool.py

@@ -258,3 +258,20 @@ def run(self) -> None:
     [thread.join() for thread in threads]
 
     assert sum(thread.nodes_gotten for thread in threads) >= 10000
+
+
+def test_get_when_alive_nodes_change_during_iteration():
+    node_configs = [NodeConfig("http", "localhost", x) for x in range(3)]
+    pool = NodePool(node_configs, node_class=Urllib3HttpNode)
+    node = pool.all()[0]
+
+    class MarkDeadOnFirstLookup(set):
+        # Simulates another thread calling mark_dead() while get() iterates over the alive nodes.
+        def __contains__(self, item):
+            if node.config in pool._alive_nodes:
+                pool.mark_dead(node)
+            return super().__contains__(item)
+
+    pool._removed_nodes = MarkDeadOnFirstLookup()
+
+    assert pool.get().config in node_configs