PYTHON-5406 - AsyncPeriodicExecutor must reset CSOT contextvars before executing

Author: NoahStapp

pymongo/_asyncio_task.py

@@ -43,6 +43,7 @@ def cancelling(self) -> int:
         return self._cancel_requests
 
 
+# We can directly pass an empty Context() object to create_task in Python >= 3.11
 def create_task(coro: Coroutine[Any, Any, Any], *, name: Optional[str] = None) -> asyncio.Task:
     if sys.version_info >= (3, 11):
         return asyncio.create_task(coro, name=name)

pymongo/_csot.py

@@ -32,6 +32,12 @@
 DEADLINE: ContextVar[float] = ContextVar("DEADLINE", default=float("inf"))
 
 
+def reset_all() -> None:
+    TIMEOUT.set(None)
+    RTT.set(0.0)
+    DEADLINE.set(float("inf"))
+
+
 def get_timeout() -> Optional[float]:
     return TIMEOUT.get(None)
 

pymongo/periodic_executor.py

@@ -23,6 +23,7 @@
 import weakref
 from typing import Any, Optional
 
+from pymongo import _csot
 from pymongo._asyncio_task import create_task
 from pymongo.lock import _create_lock
 
@@ -93,6 +94,8 @@ def skip_sleep(self) -> None:
         self._skip_sleep = True
 
     async def _run(self) -> None:
+        # The CSOT contextvars must be cleared inside the executor task before execution begins
+        _csot.reset_all()
         while not self._stopped:
             if self._task and self._task.cancelling():  # type: ignore[unused-ignore, attr-defined]
                 raise asyncio.CancelledError