PYTHON-4414 interruptInUseConnections should cancel pending connectio…

…ns too (#2010)
mongodb · Nov 19, 2024 · a7c1090 · a7c1090
1 parent a3bdc13
commit a7c1090
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py
@@ -1249,6 +1249,9 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
         async with self.lock:
             conn_id = self.next_connection_id
             self.next_connection_id += 1
+            # Use a temporary context so that interrupt_connections can cancel creating the socket.
+            tmp_context = _CancellationContext()
+            self.active_contexts.add(tmp_context)
 
         listeners = self.opts._event_listeners
         if self.enabled_for_cmap:
@@ -1267,6 +1270,8 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
         try:
             sock = await _configured_socket(self.address, self.opts)
         except BaseException as error:
+            async with self.lock:
+                self.active_contexts.discard(tmp_context)
             if self.enabled_for_cmap:
                 assert listeners is not None
                 listeners.publish_connection_closed(
@@ -1292,6 +1297,9 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
         conn = AsyncConnection(sock, self, self.address, conn_id)  # type: ignore[arg-type]
         async with self.lock:
             self.active_contexts.add(conn.cancel_context)
+            self.active_contexts.discard(tmp_context)
+        if tmp_context.cancelled:
+            conn.cancel_context.cancel()
         try:
             if self.handshake:
                 await conn.hello()
@@ -1301,6 +1309,8 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
 
             await conn.authenticate()
         except BaseException:
+            async with self.lock:
+                self.active_contexts.discard(conn.cancel_context)
             conn.close_conn(ConnectionClosedReason.ERROR)
             raise
 

diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py
@@ -1243,6 +1243,9 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
         with self.lock:
             conn_id = self.next_connection_id
             self.next_connection_id += 1
+            # Use a temporary context so that interrupt_connections can cancel creating the socket.
+            tmp_context = _CancellationContext()
+            self.active_contexts.add(tmp_context)
 
         listeners = self.opts._event_listeners
         if self.enabled_for_cmap:
@@ -1261,6 +1264,8 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
         try:
             sock = _configured_socket(self.address, self.opts)
         except BaseException as error:
+            with self.lock:
+                self.active_contexts.discard(tmp_context)
             if self.enabled_for_cmap:
                 assert listeners is not None
                 listeners.publish_connection_closed(
@@ -1286,6 +1291,9 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
         conn = Connection(sock, self, self.address, conn_id)  # type: ignore[arg-type]
         with self.lock:
             self.active_contexts.add(conn.cancel_context)
+            self.active_contexts.discard(tmp_context)
+        if tmp_context.cancelled:
+            conn.cancel_context.cancel()
         try:
             if self.handshake:
                 conn.hello()
@@ -1295,6 +1303,8 @@ def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> Connect
 
             conn.authenticate()
         except BaseException:
+            with self.lock:
+                self.active_contexts.discard(conn.cancel_context)
             conn.close_conn(ConnectionClosedReason.ERROR)
             raise
 

diff --git a/test/test_connection_monitoring.py b/test/test_connection_monitoring.py
@@ -216,11 +216,6 @@ def set_fail_point(self, command_args):
 
     def run_scenario(self, scenario_def, test):
         """Run a CMAP spec test."""
-        if (
-            scenario_def["description"]
-            == "clear with interruptInUseConnections = true closes pending connections"
-        ):
-            self.skipTest("Skip pending PYTHON-4414")
         self.logs: list = []
         self.assertEqual(scenario_def["version"], 1)
         self.assertIn(scenario_def["style"], ["unit", "integration"])