PYTHON-3288 Implement client side operation timeout (#954)

Add timeoutMS URI option and MongoClient keyword argument.
Add provisional/beta pymongo.timeout() api to set a deadline for a block of operations.

Author: ShaneHarvey

.evergreen/resync-specs.sh

@@ -105,6 +105,9 @@ do
     crud|CRUD)
       cpjson crud/tests/ crud
       ;;
+    csot|CSOT|client-side-operations-timeout)
+      cpjson client-side-operations-timeout/tests csot
+      ;;
     load-balancers|load_balancer)
       cpjson load-balancers/tests load_balancer
       ;;
@@ -150,6 +153,7 @@ do
       ;;
     uri|uri-options|uri_options)
       cpjson uri-options/tests uri_options
+      cp "$SPECS"/source/uri-options/tests/*.pem  $PYMONGO/test/uri_options
       ;;
     stable-api|versioned-api)
       cpjson versioned-api/tests versioned-api

doc/api/pymongo/index.rst

@@ -22,6 +22,8 @@
 
       The maximum wire protocol version PyMongo supports.
 
+   .. autofunction:: timeout
+
 Sub-modules:
 
 .. toctree::

doc/changelog.rst

@@ -6,6 +6,13 @@ Changes in Version 4.2
 
 .. warning:: PyMongo 4.2 drops support for Python 3.6: Python 3.7+ is now required.
 
+PyMongo 4.2 brings a number of improvements including:
+
+- Support for MongoDB 6.0.
+- Provisional (beta) support for :func:`pymongo.timeout` to apply a single timeout
+  to an entire block of pymongo operations.
+- Beta support for Queryable Encryption with MongoDB 6.0.
+
 Bug fixes
 .........
 

pymongo/__init__.py

@@ -14,7 +14,7 @@
 
 """Python driver for MongoDB."""
 
-from typing import Tuple, Union
+from typing import ContextManager, Optional, Tuple, Union
 
 ASCENDING = 1
 """Ascending sort order."""
@@ -69,6 +69,7 @@ def get_version_string() -> str:
 
 """Current version of PyMongo."""
 
+from pymongo import _csot
 from pymongo.collection import ReturnDocument  # noqa: F401
 from pymongo.common import (  # noqa: F401
     MAX_SUPPORTED_WIRE_VERSION,
@@ -97,3 +98,47 @@ def has_c() -> bool:
         return True
     except ImportError:
         return False
+
+
+def timeout(seconds: Optional[float]) -> ContextManager:
+    """**(Provisional)** Apply the given timeout for a block of operations.
+
+    .. note:: :func:`~pymongo.timeout` is currently provisional. Backwards
+       incompatible changes may occur before becoming officially supported.
+
+    Use :func:`~pymongo.timeout` in a with-statement::
+
+      with pymongo.timeout(5):
+          client.db.coll.insert_one({})
+          client.db.coll2.insert_one({})
+
+    When the with-statement is entered, a deadline is set for the entire
+    block. When that deadline is exceeded, any blocking pymongo operation
+    will raise a timeout exception. For example::
+
+      try:
+          with pymongo.timeout(5):
+              client.db.coll.insert_one({})
+              time.sleep(5)
+              # The deadline has now expired, the next operation will raise
+              # a timeout exception.
+              client.db.coll2.insert_one({})
+      except (ServerSelectionTimeoutError, ExecutionTimeout, WTimeoutError,
+              NetworkTimeout) as exc:
+          print(f"block timed out: {exc!r}")
+
+    :Parameters:
+      - `seconds`: A non-negative floating point number expressing seconds, or None.
+
+    :Raises:
+      - :py:class:`ValueError`: When `seconds` is negative.
+
+    .. versionadded:: 4.2
+    """
+    if not isinstance(seconds, (int, float, type(None))):
+        raise TypeError("timeout must be None, an int, or a float")
+    if seconds and seconds < 0:
+        raise ValueError("timeout cannot be negative")
+    if seconds is not None:
+        seconds = float(seconds)
+    return _csot._TimeoutContext(seconds)

pymongo/_csot.py

@@ -0,0 +1,80 @@
+# Copyright 2022-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.
+
+"""Internal helpers for CSOT."""
+
+import time
+from contextvars import ContextVar
+from typing import Optional
+
+TIMEOUT: ContextVar[Optional[float]] = ContextVar("TIMEOUT", default=None)
+RTT: ContextVar[float] = ContextVar("RTT", default=0.0)
+DEADLINE: ContextVar[float] = ContextVar("DEADLINE", default=float("inf"))
+
+
+def get_timeout() -> Optional[float]:
+    return TIMEOUT.get(None)
+
+
+def get_rtt() -> float:
+    return RTT.get()
+
+
+def get_deadline() -> float:
+    return DEADLINE.get()
+
+
+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
+    return DEADLINE.get() - time.monotonic()
+
+
+def clamp_remaining(max_timeout: float) -> float:
+    """Return the remaining timeout clamped to a max value."""
+    timeout = remaining()
+    if timeout is None:
+        return max_timeout
+    return min(timeout, max_timeout)
+
+
+class _TimeoutContext(object):
+    """Internal timeout context manager.
+
+    Use :func:`pymongo.timeout` instead::
+
+      with client.timeout(0.5):
+          client.test.test.insert_one({})
+    """
+
+    __slots__ = ("_timeout",)
+
+    def __init__(self, timeout: Optional[float]):
+        self._timeout = timeout
+
+    def __enter__(self):
+        set_timeout(self._timeout)
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        set_timeout(None)

pymongo/bulk.py

@@ -330,6 +330,8 @@ def _execute_command(
                     session._apply_to(cmd, retryable, ReadPreference.PRIMARY, sock_info)
                 sock_info.send_cluster_time(cmd, session, client)
                 sock_info.add_server_api(cmd)
+                # CSOT: apply timeout before encoding the command.
+                sock_info.apply_timeout(client, cmd)
                 ops = islice(run.ops, run.idx_offset, None)
 
                 # Run as many ops as possible in one command.

pymongo/client_options.py

@@ -14,6 +14,8 @@
 
 """Tools to parse mongo client options."""
 
+from typing import Optional
+
 from bson.codec_options import _parse_codec_options
 from pymongo import common
 from pymongo.auth import _build_credentials_tuple
@@ -195,6 +197,7 @@ def __init__(self, username, password, database, options):
         self.__server_selector = options.get("server_selector", any_server_selector)
         self.__auto_encryption_opts = options.get("auto_encryption_opts")
         self.__load_balanced = options.get("loadbalanced")
+        self.__timeout = options.get("timeoutms")
 
     @property
     def _options(self):
@@ -260,6 +263,14 @@ def read_concern(self):
         """A :class:`~pymongo.read_concern.ReadConcern` instance."""
         return self.__read_concern
 
+    @property
+    def timeout(self) -> Optional[float]:
+        """The timeout.
+
+        ..versionadded: 4.2
+        """
+        return self.__timeout
+
     @property
     def retry_writes(self):
         """If this instance should retry supported write operations."""

pymongo/client_session.py

@@ -150,6 +150,7 @@
 from bson.int64 import Int64
 from bson.son import SON
 from bson.timestamp import Timestamp
+from pymongo import _csot
 from pymongo.cursor import _SocketManager
 from pymongo.errors import (
     ConfigurationError,
@@ -826,7 +827,7 @@ def _finish_transaction(self, sock_info, command_name):
         wc = opts.write_concern
         cmd = SON([(command_name, 1)])
         if command_name == "commitTransaction":
-            if opts.max_commit_time_ms:
+            if opts.max_commit_time_ms and _csot.get_timeout() is None:
                 cmd["maxTimeMS"] = opts.max_commit_time_ms
 
             # Transaction spec says that after the initial commit attempt,

pymongo/collection.py

@@ -116,6 +116,7 @@ def __init__(
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional["ReadConcern"] = None,
         session: Optional["ClientSession"] = None,
+        timeout: Optional[float] = None,
         encrypted_fields: Optional[Mapping[str, Any]] = None,
         **kwargs: Any,
     ) -> None:
@@ -198,6 +199,7 @@ def __init__(
             read_preference or database.read_preference,
             write_concern or database.write_concern,
             read_concern or database.read_concern,
+            timeout if timeout is not None else database.timeout,
         )
         if not isinstance(name, str):
             raise TypeError("name must be an instance of str")
@@ -390,6 +392,7 @@ def with_options(
         read_preference: Optional[_ServerMode] = None,
         write_concern: Optional[WriteConcern] = None,
         read_concern: Optional["ReadConcern"] = None,
+        timeout: Optional[float] = None,
     ) -> "Collection[_DocumentType]":
         """Get a clone of this collection changing the specified settings.
 
@@ -428,6 +431,7 @@ def with_options(
             read_preference or self.read_preference,
             write_concern or self.write_concern,
             read_concern or self.read_concern,
+            timeout=timeout if timeout is not None else self.timeout,
         )
 
     def bulk_write(

pymongo/common.py

@@ -339,6 +339,15 @@ def validate_timeout_or_none_or_zero(option: Any, value: Any) -> Optional[float]
     return validate_positive_float(option, value) / 1000.0
 
 
+def validate_timeoutms(option: Any, value: Any) -> Optional[float]:
+    """Validates a timeout specified in milliseconds returning
+    a value in floating point seconds.
+    """
+    if value is None:
+        return None
+    return validate_positive_float_or_zero(option, value) / 1000.0
+
+
 def validate_max_staleness(option: str, value: Any) -> int:
     """Validates maxStalenessSeconds according to the Max Staleness Spec."""
     if value == -1 or value == "-1":
@@ -658,6 +667,7 @@ def validate_auto_encryption_opts_or_none(option: Any, value: Any) -> Optional[A
     "zlibcompressionlevel": validate_zlib_compression_level,
     "srvservicename": validate_string,
     "srvmaxhosts": validate_non_negative_integer,
+    "timeoutms": validate_timeoutms,
 }
 
 # Dictionary where keys are the names of URI options specific to pymongo,
@@ -821,8 +831,8 @@ def __init__(
         read_preference: _ServerMode,
         write_concern: WriteConcern,
         read_concern: ReadConcern,
+        timeout: Optional[float],
     ) -> None:
-
         if not isinstance(codec_options, CodecOptions):
             raise TypeError("codec_options must be an instance of bson.codec_options.CodecOptions")
         self.__codec_options = codec_options
@@ -845,6 +855,12 @@ def __init__(
             raise TypeError("read_concern must be an instance of pymongo.read_concern.ReadConcern")
         self.__read_concern = read_concern
 
+        if not isinstance(timeout, (int, float, type(None))):
+            raise TypeError("timeout must be None, an int, or a float")
+        if timeout and timeout < 0:
+            raise TypeError("timeout cannot be negative")
+        self.__timeout = float(timeout) if timeout else None
+
     @property
     def codec_options(self) -> CodecOptions:
         """Read only access to the :class:`~bson.codec_options.CodecOptions`
@@ -894,6 +910,14 @@ def read_concern(self) -> ReadConcern:
         """
         return self.__read_concern
 
+    @property
+    def timeout(self) -> Optional[float]:
+        """Read only access to the timeout of this instance.
+
+        .. versionadded:: 4.2
+        """
+        return self.__timeout
+
 
 class _CaseInsensitiveDictionary(abc.MutableMapping):
     def __init__(self, *args, **kwargs):