PYTHON-3290 Support nested pymongo.timeout() calls (#962)

Author: ShaneHarvey

pymongo/__init__.py

@@ -127,6 +127,18 @@ def timeout(seconds: Optional[float]) -> ContextManager:
               NetworkTimeout) as exc:
           print(f"block timed out: {exc!r}")
 
+    When nesting :func:`~pymongo.timeout`, the nested block overrides the
+    timeout. When exiting the block, the previous deadline is restored::
+
+      with pymongo.timeout(5):
+          coll.find_one()  # Uses the 5 second deadline.
+          with pymongo.timeout(3):
+              coll.find_one() # Uses the 3 second deadline.
+          coll.find_one()  # Uses the original 5 second deadline.
+          with pymongo.timeout(10):
+              coll.find_one()  # Uses the 10 second deadline.
+          coll.find_one()  # Uses the original 5 second deadline.
+
     :Parameters:
       - `seconds`: A non-negative floating point number expressing seconds, or None.
 

pymongo/_csot.py

@@ -15,8 +15,8 @@
 """Internal helpers for CSOT."""
 
 import time
-from contextvars import ContextVar
-from typing import Optional
+from contextvars import ContextVar, Token
+from typing import Optional, Tuple
 
 TIMEOUT: ContextVar[Optional[float]] = ContextVar("TIMEOUT", default=None)
 RTT: ContextVar[float] = ContextVar("RTT", default=0.0)
@@ -39,11 +39,6 @@ def set_rtt(rtt: float) -> None:
     RTT.set(rtt)
 
 
-def set_timeout(timeout: Optional[float]) -> None:
-    TIMEOUT.set(timeout)
-    DEADLINE.set(time.monotonic() + timeout if timeout else float("inf"))
-
-
 def remaining() -> Optional[float]:
     if not get_timeout():
         return None
@@ -67,14 +62,24 @@ class _TimeoutContext(object):
           client.test.test.insert_one({})
     """
 
-    __slots__ = ("_timeout",)
+    __slots__ = ("_timeout", "_tokens")
 
     def __init__(self, timeout: Optional[float]):
         self._timeout = timeout
+        self._tokens: Optional[Tuple[Token, Token, Token]] = None
 
     def __enter__(self):
-        set_timeout(self._timeout)
+        timeout_token = TIMEOUT.set(self._timeout)
+        deadline_token = DEADLINE.set(
+            time.monotonic() + self._timeout if self._timeout else float("inf")
+        )
+        rtt_token = RTT.set(0.0)
+        self._tokens = (timeout_token, deadline_token, rtt_token)
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        set_timeout(None)
+        if self._tokens:
+            timeout_token, deadline_token, rtt_token = self._tokens
+            TIMEOUT.reset(timeout_token)
+            DEADLINE.reset(deadline_token)
+            RTT.reset(rtt_token)

pymongo/topology.py

@@ -270,7 +270,8 @@ def _select_server(self, selector, server_selection_timeout=None, address=None):
     def select_server(self, selector, server_selection_timeout=None, address=None):
         """Like select_servers, but choose a random server if several match."""
         server = self._select_server(selector, server_selection_timeout, address)
-        _csot.set_rtt(server.description.round_trip_time)
+        if _csot.get_timeout():
+            _csot.set_rtt(server.description.round_trip_time)
         return server
 
     def select_server_by_address(self, address, server_selection_timeout=None):

test/test_csot.py

@@ -19,14 +19,55 @@
 
 sys.path[0:0] = [""]
 
-from test import unittest
+from test import IntegrationTest, unittest
 from test.unified_format import generate_test_classes
 
+import pymongo
+from pymongo import _csot
+
 # Location of JSON test specifications.
 TEST_PATH = os.path.join(os.path.dirname(os.path.realpath(__file__)), "csot")
 
 # Generate unified tests.
 globals().update(generate_test_classes(TEST_PATH, module=__name__))
 
+
+class TestCSOT(IntegrationTest):
+    def test_timeout_nested(self):
+        coll = self.db.coll
+        self.assertEqual(_csot.get_timeout(), None)
+        self.assertEqual(_csot.get_deadline(), float("inf"))
+        self.assertEqual(_csot.get_rtt(), 0.0)
+        with pymongo.timeout(10):
+            coll.find_one()
+            self.assertEqual(_csot.get_timeout(), 10)
+            deadline_10 = _csot.get_deadline()
+
+            with pymongo.timeout(15):
+                coll.find_one()
+                self.assertEqual(_csot.get_timeout(), 15)
+                self.assertGreater(_csot.get_deadline(), deadline_10)
+
+            # Should be reset to previous values
+            self.assertEqual(_csot.get_timeout(), 10)
+            self.assertEqual(_csot.get_deadline(), deadline_10)
+            coll.find_one()
+
+            with pymongo.timeout(5):
+                coll.find_one()
+                self.assertEqual(_csot.get_timeout(), 5)
+                self.assertLess(_csot.get_deadline(), deadline_10)
+
+            # Should be reset to previous values
+            self.assertEqual(_csot.get_timeout(), 10)
+            self.assertEqual(_csot.get_deadline(), deadline_10)
+            coll.find_one()
+
+        # Should be reset to previous values
+        self.assertEqual(_csot.get_timeout(), None)
+        self.assertEqual(_csot.get_deadline(), float("inf"))
+        self.assertEqual(_csot.get_rtt(), 0.0)
+
+
 if __name__ == "__main__":
     unittest.main()