PYTHON-2484 Added lock sanitization for MongoClient and ObjectId (#985)

Author: bcwarner

bson/__init__.py

@@ -56,6 +56,7 @@
 
 import datetime
 import itertools
+import os
 import re
 import struct
 import sys
@@ -1336,3 +1337,16 @@ def decode(self, codec_options: "CodecOptions[_DocumentType]" = DEFAULT_CODEC_OP
 def has_c() -> bool:
     """Is the C extension installed?"""
     return _USE_C
+
+
+def _after_fork():
+    """Releases the ObjectID lock child."""
+    if ObjectId._inc_lock.locked():
+        ObjectId._inc_lock.release()
+
+
+if hasattr(os, "register_at_fork"):
+    # This will run in the same thread as the fork was called.
+    # If we fork in a critical region on the same thread, it should break.
+    # This is fine since we would never call fork directly from a critical region.
+    os.register_at_fork(after_in_child=_after_fork)

pymongo/cursor.py

@@ -14,7 +14,6 @@
 
 """Cursor class to iterate over Mongo query results."""
 import copy
-import threading
 import warnings
 from collections import deque
 from typing import (
@@ -45,6 +44,7 @@
     validate_is_mapping,
 )
 from pymongo.errors import ConnectionFailure, InvalidOperation, OperationFailure
+from pymongo.lock import _create_lock
 from pymongo.message import (
     _CursorAddress,
     _GetMore,
@@ -133,7 +133,7 @@ def __init__(self, sock, more_to_come):
         self.sock = sock
         self.more_to_come = more_to_come
         self.closed = False
-        self.lock = threading.Lock()
+        self.lock = _create_lock()
 
     def update_exhaust(self, more_to_come):
         self.more_to_come = more_to_come

pymongo/lock.py

@@ -0,0 +1,39 @@
+# 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.
+
+import os
+import threading
+import weakref
+
+_HAS_REGISTER_AT_FORK = hasattr(os, "register_at_fork")
+
+# References to instances of _create_lock
+_forkable_locks: weakref.WeakSet = weakref.WeakSet()
+
+
+def _create_lock():
+    """Represents a lock that is tracked upon instantiation using a WeakSet and
+    reset by pymongo upon forking.
+    """
+    lock = threading.Lock()
+    if _HAS_REGISTER_AT_FORK:
+        _forkable_locks.add(lock)
+    return lock
+
+
+def _release_locks() -> None:
+    # Completed the fork, reset all the locks in the child.
+    for lock in _forkable_locks:
+        if lock.locked():
+            lock.release()

pymongo/mongo_client.py

@@ -32,7 +32,7 @@
 """
 
 import contextlib
-import threading
+import os
 import weakref
 from collections import defaultdict
 from typing import (
@@ -82,6 +82,7 @@
     ServerSelectionTimeoutError,
     WaitQueueTimeoutError,
 )
+from pymongo.lock import _create_lock, _release_locks
 from pymongo.pool import ConnectionClosedReason
 from pymongo.read_preferences import ReadPreference, _ServerMode
 from pymongo.server_selectors import writable_server_selector
@@ -126,6 +127,7 @@ class MongoClient(common.BaseObject, Generic[_DocumentType]):
     # Define order to retrieve options from ClientOptions for __repr__.
     # No host/port; these are retrieved from TopologySettings.
     _constructor_args = ("document_class", "tz_aware", "connect")
+    _clients: weakref.WeakValueDictionary = weakref.WeakValueDictionary()
 
     def __init__(
         self,
@@ -788,7 +790,7 @@ def __init__(
         self.__options = options = ClientOptions(username, password, dbase, opts)
 
         self.__default_database_name = dbase
-        self.__lock = threading.Lock()
+        self.__lock = _create_lock()
         self.__kill_cursors_queue: List = []
 
         self._event_listeners = options.pool_options._event_listeners
@@ -817,6 +819,23 @@ def __init__(
             srv_max_hosts=srv_max_hosts,
         )
 
+        self._init_background()
+
+        if connect:
+            self._get_topology()
+
+        self._encrypter = None
+        if self.__options.auto_encryption_opts:
+            from pymongo.encryption import _Encrypter
+
+            self._encrypter = _Encrypter(self, self.__options.auto_encryption_opts)
+        self._timeout = self.__options.timeout
+
+        # Add this client to the list of weakly referenced items.
+        # This will be used later if we fork.
+        MongoClient._clients[self._topology._topology_id] = self
+
+    def _init_background(self):
         self._topology = Topology(self._topology_settings)
 
         def target():
@@ -838,15 +857,9 @@ def target():
         self_ref: Any = weakref.ref(self, executor.close)
         self._kill_cursors_executor = executor
 
-        if connect:
-            self._get_topology()
-
-        self._encrypter = None
-        if self.__options.auto_encryption_opts:
-            from pymongo.encryption import _Encrypter
-
-            self._encrypter = _Encrypter(self, self.__options.auto_encryption_opts)
-        self._timeout = options.timeout
+    def _after_fork(self):
+        """Resets topology in a child after successfully forking."""
+        self._init_background()
 
     def _duplicate(self, **kwargs):
         args = self.__init_kwargs.copy()
@@ -2150,3 +2163,22 @@ def __enter__(self):
 
     def __exit__(self, exc_type, exc_val, exc_tb):
         return self.handle(exc_type, exc_val)
+
+
+def _after_fork_child():
+    """Releases the locks in child process and resets the
+    topologies in all MongoClients.
+    """
+    # Reinitialize locks
+    _release_locks()
+
+    # Perform cleanup in clients (i.e. get rid of topology)
+    for _, client in MongoClient._clients.items():
+        client._after_fork()
+
+
+if hasattr(os, "register_at_fork"):
+    # This will run in the same thread as the fork was called.
+    # If we fork in a critical region on the same thread, it should break.
+    # This is fine since we would never call fork directly from a critical region.
+    os.register_at_fork(after_in_child=_after_fork_child)

pymongo/monitor.py

@@ -15,14 +15,14 @@
 """Class to monitor a MongoDB server on a background thread."""
 
 import atexit
-import threading
 import time
 import weakref
 from typing import Any, Mapping, cast
 
 from pymongo import common, periodic_executor
 from pymongo.errors import NotPrimaryError, OperationFailure, _OperationCancelled
 from pymongo.hello import Hello
+from pymongo.lock import _create_lock
 from pymongo.periodic_executor import _shutdown_executors
 from pymongo.read_preferences import MovingAverage
 from pymongo.server_description import ServerDescription
@@ -350,7 +350,7 @@ def __init__(self, topology, topology_settings, pool):
 
         self._pool = pool
         self._moving_average = MovingAverage()
-        self._lock = threading.Lock()
+        self._lock = _create_lock()
 
     def close(self):
         self.gc_safe_close()

pymongo/ocsp_cache.py

@@ -16,7 +16,8 @@
 
 from collections import namedtuple
 from datetime import datetime as _datetime
-from threading import Lock
+
+from pymongo.lock import _create_lock
 
 
 class _OCSPCache(object):
@@ -30,7 +31,7 @@ class _OCSPCache(object):
     def __init__(self):
         self._data = {}
         # Hold this lock when accessing _data.
-        self._lock = Lock()
+        self._lock = _create_lock()
 
     def _get_cache_key(self, ocsp_request):
         return self.CACHE_KEY_TYPE(

pymongo/periodic_executor.py

@@ -19,6 +19,8 @@
 import weakref
 from typing import Any, Optional
 
+from pymongo.lock import _create_lock
+
 
 class PeriodicExecutor(object):
     def __init__(self, interval, min_interval, target, name=None):
@@ -45,9 +47,8 @@ def __init__(self, interval, min_interval, target, name=None):
         self._thread: Optional[threading.Thread] = None
         self._name = name
         self._skip_sleep = False
-
         self._thread_will_exit = False
-        self._lock = threading.Lock()
+        self._lock = _create_lock()
 
     def __repr__(self):
         return "<%s(name=%s) object at 0x%x>" % (self.__class__.__name__, self._name, id(self))

pymongo/pool.py

@@ -56,6 +56,7 @@
     _CertificateError,
 )
 from pymongo.hello import Hello, HelloCompat
+from pymongo.lock import _create_lock
 from pymongo.monitoring import ConnectionCheckOutFailedReason, ConnectionClosedReason
 from pymongo.network import command, receive_message
 from pymongo.read_preferences import ReadPreference
@@ -1152,7 +1153,7 @@ def __init__(self, address, options, handshake=True):
         # and returned to pool from the left side. Stale sockets removed
         # from the right side.
         self.sockets: collections.deque = collections.deque()
-        self.lock = threading.Lock()
+        self.lock = _create_lock()
         self.active_sockets = 0
         # Monotonically increasing connection ID required for CMAP Events.
         self.next_connection_id = 1

pymongo/topology.py

@@ -17,7 +17,6 @@
 import os
 import queue
 import random
-import threading
 import time
 import warnings
 import weakref
@@ -37,6 +36,7 @@
     WriteError,
 )
 from pymongo.hello import Hello
+from pymongo.lock import _HAS_REGISTER_AT_FORK, _create_lock
 from pymongo.monitor import SrvMonitor
 from pymongo.pool import PoolOptions
 from pymongo.server import Server
@@ -127,7 +127,7 @@ def __init__(self, topology_settings):
         self._seed_addresses = list(topology_description.server_descriptions())
         self._opened = False
         self._closed = False
-        self._lock = threading.Lock()
+        self._lock = _create_lock()
         self._condition = self._settings.condition_class(self._lock)
         self._servers = {}
         self._pid = None
@@ -174,12 +174,13 @@ def open(self):
             self._pid = pid
         elif pid != self._pid:
             self._pid = pid
-            warnings.warn(
-                "MongoClient opened before fork. Create MongoClient only "
-                "after forking. See PyMongo's documentation for details: "
-                "https://pymongo.readthedocs.io/en/stable/faq.html#"
-                "is-pymongo-fork-safe"
-            )
+            if not _HAS_REGISTER_AT_FORK:
+                warnings.warn(
+                    "MongoClient opened before fork. May not be entirely fork-safe, "
+                    "proceed with caution. See PyMongo's documentation for details: "
+                    "https://pymongo.readthedocs.io/en/stable/faq.html#"
+                    "is-pymongo-fork-safe"
+                )
             with self._lock:
                 # Close servers and clear the pools.
                 for server in self._servers.values():