Explicitly raise asyncio.CancelledError when tasks are being cancelled

Author: NoahStapp

pymongo/asynchronous/client_bulk.py

@@ -476,7 +476,6 @@ async def _process_results_cursor(
                         if op_type == "delete":
                             res = DeleteResult(doc, acknowledged=True)  # type: ignore[assignment]
                         full_result[f"{op_type}Results"][original_index] = res
-
             except Exception as exc:
                 # Attempt to close the cursor, then raise top-level error.
                 if cmd_cursor.alive:

pymongo/asynchronous/encryption.py

@@ -15,6 +15,7 @@
 """Support for explicit client-side field level encryption."""
 from __future__ import annotations
 
+import asyncio
 import contextlib
 import enum
 import socket
@@ -111,6 +112,8 @@ def _wrap_encryption_errors() -> Iterator[None]:
         # BSON encoding/decoding errors are unrelated to encryption so
         # we should propagate them unchanged.
         raise
+    except asyncio.CancelledError:
+        raise
     except Exception as exc:
         raise EncryptionError(exc) from exc
 
@@ -200,6 +203,8 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 conn.close()
         except (PyMongoError, MongoCryptError):
             raise  # Propagate pymongo errors directly.
+        except asyncio.CancelledError:
+            raise
         except Exception as error:
             # Wrap I/O errors in PyMongo exceptions.
             _raise_connection_failure((host, port), error)
@@ -722,6 +727,8 @@ async def create_encrypted_collection(
                 await database.create_collection(name=name, **kwargs),
                 encrypted_fields,
             )
+        except asyncio.CancelledError:
+            raise
         except Exception as exc:
             raise EncryptedCollectionError(exc, encrypted_fields) from exc
 

pymongo/asynchronous/monitor.py

@@ -238,6 +238,9 @@ async def _run(self) -> None:
         except ReferenceError:
             # Topology was garbage-collected.
             await self.close()
+        finally:
+            if self._executor._stopped:
+                await self._rtt_monitor.close()
 
     async def _check_server(self) -> ServerDescription:
         """Call hello or read the next streaming response.
@@ -254,6 +257,8 @@ async def _check_server(self) -> ServerDescription:
                 details = cast(Mapping[str, Any], exc.details)
                 await self._topology.receive_cluster_time(details.get("$clusterTime"))
                 raise
+        except asyncio.CancelledError:
+            raise
         except ReferenceError:
             raise
         except Exception as error:

pymongo/periodic_executor.py

@@ -61,9 +61,10 @@ def __repr__(self) -> str:
     def open(self) -> None:
         """Start. Multiple calls have no effect."""
         self._stopped = False
-        started = self._task and not self._task.done()
 
-        if not started:
+        if self._task is None or (
+            self._task.done() and not self._task.cancelled() and not self._task.cancelling()
+        ):
             self._task = asyncio.get_event_loop().create_task(self._run(), name=self._name)
 
     def close(self, dummy: Any = None) -> None:
@@ -83,7 +84,7 @@ async def join(self, timeout: Optional[int] = None) -> None:
                 pass
             except asyncio.exceptions.CancelledError:
                 # Task was already finished, or not yet started.
-                pass
+                raise
 
     def wake(self) -> None:
         """Execute the target function soon."""
@@ -97,6 +98,8 @@ def skip_sleep(self) -> None:
 
     async def _run(self) -> None:
         while not self._stopped:
+            if self._task and self._task.cancelling():
+                raise asyncio.CancelledError
             try:
                 if not await self._target():
                     self._stopped = True

pymongo/synchronous/client_bulk.py

@@ -474,7 +474,6 @@ def _process_results_cursor(
                         if op_type == "delete":
                             res = DeleteResult(doc, acknowledged=True)  # type: ignore[assignment]
                         full_result[f"{op_type}Results"][original_index] = res
-
             except Exception as exc:
                 # Attempt to close the cursor, then raise top-level error.
                 if cmd_cursor.alive:

pymongo/synchronous/encryption.py

@@ -15,6 +15,7 @@
 """Support for explicit client-side field level encryption."""
 from __future__ import annotations
 
+import asyncio
 import contextlib
 import enum
 import socket
@@ -111,6 +112,8 @@ def _wrap_encryption_errors() -> Iterator[None]:
         # BSON encoding/decoding errors are unrelated to encryption so
         # we should propagate them unchanged.
         raise
+    except asyncio.CancelledError:
+        raise
     except Exception as exc:
         raise EncryptionError(exc) from exc
 
@@ -200,6 +203,8 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 conn.close()
         except (PyMongoError, MongoCryptError):
             raise  # Propagate pymongo errors directly.
+        except asyncio.CancelledError:
+            raise
         except Exception as error:
             # Wrap I/O errors in PyMongo exceptions.
             _raise_connection_failure((host, port), error)
@@ -716,6 +721,8 @@ def create_encrypted_collection(
                 database.create_collection(name=name, **kwargs),
                 encrypted_fields,
             )
+        except asyncio.CancelledError:
+            raise
         except Exception as exc:
             raise EncryptedCollectionError(exc, encrypted_fields) from exc
 

pymongo/synchronous/monitor.py

@@ -238,6 +238,9 @@ def _run(self) -> None:
         except ReferenceError:
             # Topology was garbage-collected.
             self.close()
+        finally:
+            if self._executor._stopped:
+                self._rtt_monitor.close()
 
     def _check_server(self) -> ServerDescription:
         """Call hello or read the next streaming response.
@@ -254,6 +257,8 @@ def _check_server(self) -> ServerDescription:
                 details = cast(Mapping[str, Any], exc.details)
                 self._topology.receive_cluster_time(details.get("$clusterTime"))
                 raise
+        except asyncio.CancelledError:
+            raise
         except ReferenceError:
             raise
         except Exception as error:

test/__init__.py

@@ -868,8 +868,9 @@ def reset_client_context():
     if _IS_SYNC:
         # sync tests don't need to reset a client context
         return
-    client_context.client.close()
-    client_context.client = None
+    elif client_context.client is not None:
+        client_context.client.close()
+        client_context.client = None
     client_context._init_client()
 
 
@@ -1135,7 +1136,7 @@ class IntegrationTest(PyMongoTestCase):
 
     @client_context.require_connection
     def setUp(self) -> None:
-        if not _IS_SYNC and client_context.client is not None:
+        if not _IS_SYNC:
             reset_client_context()
         if client_context.load_balancer and not getattr(self, "RUN_ON_LOAD_BALANCER", False):
             raise SkipTest("this test does not support load balancers")
@@ -1210,7 +1211,6 @@ def teardown():
             c.drop_database("pymongo_test_mike")
             c.drop_database("pymongo_test_bernie")
         c.close()
-
     print_running_clients()
 
 

test/asynchronous/__init__.py

@@ -870,8 +870,9 @@ async def reset_client_context():
     if _IS_SYNC:
         # sync tests don't need to reset a client context
         return
-    await async_client_context.client.close()
-    async_client_context.client = None
+    elif async_client_context.client is not None:
+        await async_client_context.client.close()
+        async_client_context.client = None
     await async_client_context._init_client()
 
 
@@ -1153,7 +1154,7 @@ class AsyncIntegrationTest(AsyncPyMongoTestCase):
 
     @async_client_context.require_connection
     async def asyncSetUp(self) -> None:
-        if not _IS_SYNC and async_client_context.client is not None:
+        if not _IS_SYNC:
             await reset_client_context()
         if async_client_context.load_balancer and not getattr(self, "RUN_ON_LOAD_BALANCER", False):
             raise SkipTest("this test does not support load balancers")
@@ -1228,7 +1229,6 @@ async def async_teardown():
             await c.drop_database("pymongo_test_mike")
             await c.drop_database("pymongo_test_bernie")
         await c.close()
-
     print_running_clients()
 
 

test/asynchronous/test_client.py

@@ -1280,6 +1280,7 @@ async def get_x(db):
     async def test_server_selection_timeout(self):
         client = AsyncMongoClient(serverSelectionTimeoutMS=100, connect=False)
         self.assertAlmostEqual(0.1, client.options.server_selection_timeout)
+        await client.close()
 
         client = AsyncMongoClient(serverSelectionTimeoutMS=0, connect=False)
 
@@ -1292,18 +1293,22 @@ async def test_server_selection_timeout(self):
         self.assertRaises(
             ConfigurationError, AsyncMongoClient, serverSelectionTimeoutMS=None, connect=False
         )
+        await client.close()
 
         client = AsyncMongoClient(
             "mongodb://localhost/?serverSelectionTimeoutMS=100", connect=False
         )
         self.assertAlmostEqual(0.1, client.options.server_selection_timeout)
+        await client.close()
 
         client = AsyncMongoClient("mongodb://localhost/?serverSelectionTimeoutMS=0", connect=False)
         self.assertAlmostEqual(0, client.options.server_selection_timeout)
+        await client.close()
 
         # Test invalid timeout in URI ignored and set to default.
         client = AsyncMongoClient("mongodb://localhost/?serverSelectionTimeoutMS=-1", connect=False)
         self.assertAlmostEqual(30, client.options.server_selection_timeout)
+        await client.close()
 
         client = AsyncMongoClient("mongodb://localhost/?serverSelectionTimeoutMS=", connect=False)
         self.assertAlmostEqual(30, client.options.server_selection_timeout)