Subclass asyncio.Task

Author: NoahStapp

pymongo/_asyncio_task.py

@@ -0,0 +1,40 @@
+# Copyright 2024-present MongoDB, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""A custom asyncio.Task that allows checking if a task has been sent a cancellation request.
+Can be removed once we drop Python 3.10 support in favor of asyncio.Task.cancelling."""
+
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Coroutine, Optional
+
+
+class _Task(asyncio.Task):
+    def __init__(self, coro: Coroutine[Any, Any, Any], *, name: Optional[str] = None) -> None:
+        super().__init__(coro, name=name)
+        self._cancelled: bool = False
+        asyncio._register_task(self)
+
+    def cancel(self, msg: Optional[str] = None) -> bool:
+        self._cancelled = True
+        return super().cancel(msg=msg)
+
+    def is_cancelled(self) -> bool:
+        return self._cancelled
+
+
+def create_task(coro: Coroutine[Any, Any, Any], *, name: Optional[str] = None) -> _Task:
+    return _Task(coro, name=name)

pymongo/network_layer.py

@@ -29,6 +29,7 @@
 )
 
 from pymongo import _csot, ssl_support
+from pymongo._asyncio_task import create_task
 from pymongo.errors import _OperationCancelled
 from pymongo.socket_checker import _errno_from_exception
 
@@ -259,12 +260,12 @@ async def async_receive_data(
 
     sock.settimeout(0.0)
     loop = asyncio.get_event_loop()
-    cancellation_task = asyncio.create_task(_poll_cancellation(conn))
+    cancellation_task = create_task(_poll_cancellation(conn))
     try:
         if _HAVE_SSL and isinstance(sock, (SSLSocket, _sslConn)):
-            read_task = asyncio.create_task(_async_receive_ssl(sock, length, loop))  # type: ignore[arg-type]
+            read_task = create_task(_async_receive_ssl(sock, length, loop))  # type: ignore[arg-type]
         else:
-            read_task = asyncio.create_task(_async_receive(sock, length, loop))  # type: ignore[arg-type]
+            read_task = create_task(_async_receive(sock, length, loop))  # type: ignore[arg-type]
         tasks = [read_task, cancellation_task]
         done, pending = await asyncio.wait(
             tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED

pymongo/periodic_executor.py

@@ -23,6 +23,7 @@
 import weakref
 from typing import Any, Optional
 
+from pymongo._asyncio_task import _Task, create_task
 from pymongo.lock import _create_lock
 
 _IS_SYNC = False
@@ -51,7 +52,7 @@ def __init__(
         self._min_interval = min_interval
         self._target = target
         self._stopped = False
-        self._task: Optional[asyncio.Task] = None
+        self._task: Optional[_Task] = None
         self._name = name
         self._skip_sleep = False
 
@@ -63,9 +64,9 @@ def open(self) -> None:
         self._stopped = False
 
         if self._task is None or (
-            self._task.done() and not self._task.cancelled() and not self._task.cancelling()
+            self._task.done() and not self._task.cancelled() and not self._task.is_cancelled()
         ):
-            self._task = asyncio.get_event_loop().create_task(self._run(), name=self._name)
+            self._task = create_task(self._run(), name=self._name)
 
     def close(self, dummy: Any = None) -> None:
         """Stop. To restart, call open().
@@ -98,7 +99,7 @@ def skip_sleep(self) -> None:
 
     async def _run(self) -> None:
         while not self._stopped:
-            if self._task and self._task.cancelling():
+            if self._task and self._task.is_cancelled():
                 raise asyncio.CancelledError
             try:
                 if not await self._target():