From d46eeae93ccb814688f0058320f597e8f1592205 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Mon, 18 Nov 2024 16:08:51 -0800 Subject: [PATCH 1/2] PYTHON-4414 interruptInUseConnections should cancel pending connections too --- pymongo/asynchronous/pool.py | 10 ++++++++++ pymongo/synchronous/pool.py | 10 ++++++++++ test/test_connection_monitoring.py | 5 ----- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index a9f02d650a..1163018775 100644 --- 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( @@ -1291,7 +1296,10 @@ 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.discard(tmp_context) self.active_contexts.add(conn.cancel_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 index eb007a3471..ea52848e61 100644 --- 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( @@ -1285,7 +1290,10 @@ 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.discard(tmp_context) self.active_contexts.add(conn.cancel_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 index d576a1184a..05411d17ba 100644 --- 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"]) From 25ca20710a73961594f8a8763656a09e94596c57 Mon Sep 17 00:00:00 2001 From: Shane Harvey Date: Mon, 18 Nov 2024 16:22:38 -0800 Subject: [PATCH 2/2] PYTHON-4414 Fix order, add then discard is safer --- pymongo/asynchronous/pool.py | 2 +- pymongo/synchronous/pool.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/pymongo/asynchronous/pool.py b/pymongo/asynchronous/pool.py index 1163018775..ca0cebd417 100644 --- a/pymongo/asynchronous/pool.py +++ b/pymongo/asynchronous/pool.py @@ -1296,8 +1296,8 @@ 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.discard(tmp_context) self.active_contexts.add(conn.cancel_context) + self.active_contexts.discard(tmp_context) if tmp_context.cancelled: conn.cancel_context.cancel() try: diff --git a/pymongo/synchronous/pool.py b/pymongo/synchronous/pool.py index ea52848e61..86baf15b9a 100644 --- a/pymongo/synchronous/pool.py +++ b/pymongo/synchronous/pool.py @@ -1290,8 +1290,8 @@ 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.discard(tmp_context) self.active_contexts.add(conn.cancel_context) + self.active_contexts.discard(tmp_context) if tmp_context.cancelled: conn.cancel_context.cancel() try: