use finer-grained locks

Author: blink1073

pymongo/asynchronous/pool.py

@@ -733,8 +733,11 @@ def __init__(
         # from the right side.
         self.conns: collections.deque[AsyncConnection] = collections.deque()
         self.active_contexts: set[_CancellationContext] = set()
+        # The main lock for the pool.  The lock should only be used to protect
+        # updating attributes.
+        # If possible, avoid any additional work while holding the lock.
+        # If looping over an attribute, copy the container and do not take the lock.
         self.lock = _async_create_lock()
-        self._max_connecting_cond = _async_create_condition(self.lock)
         self.active_sockets = 0
         # Monotonically increasing connection ID required for CMAP Events.
         self.next_connection_id = 1
@@ -760,15 +763,19 @@ def __init__(
         # The first portion of the wait queue.
         # Enforces: maxPoolSize
         # Also used for: clearing the wait queue
-        self.size_cond = _async_create_condition(self.lock)
+        # Use a different lock to prevent lock contention.  This lock protects
+        # "requests".
+        self.size_cond = _async_create_condition(_async_create_lock())
         self.requests = 0
         self.max_pool_size = self.opts.max_pool_size
         if not self.max_pool_size:
             self.max_pool_size = float("inf")
         # The second portion of the wait queue.
         # Enforces: maxConnecting
         # Also used for: clearing the wait queue
-        self._max_connecting_cond = _async_create_condition(self.lock)
+        # Use a different lock to prevent lock contention.  This lock protects
+        # "_pending".
+        self._max_connecting_cond = _async_create_condition(_async_create_lock())
         self._max_connecting = self.opts.max_connecting
         self._pending = 0
         self._client_id = client_id
@@ -797,20 +804,24 @@ def __init__(
 
     async def ready(self) -> None:
         # Take the lock to avoid the race condition described in PYTHON-2699.
-        async with self.lock:
-            if self.state != PoolState.READY:
+        state_changed = False
+        if self.state != PoolState.READY:
+            async with self.lock:
                 self.state = PoolState.READY
-                if self.enabled_for_cmap:
-                    assert self.opts._event_listeners is not None
-                    self.opts._event_listeners.publish_pool_ready(self.address)
-                if self.enabled_for_logging and _CONNECTION_LOGGER.isEnabledFor(logging.DEBUG):
-                    _debug_log(
-                        _CONNECTION_LOGGER,
-                        message=_ConnectionStatusMessage.POOL_READY,
-                        clientId=self._client_id,
-                        serverHost=self.address[0],
-                        serverPort=self.address[1],
-                    )
+                state_changed = True
+        if not state_changed:
+            return
+        if self.enabled_for_cmap:
+            assert self.opts._event_listeners is not None
+            self.opts._event_listeners.publish_pool_ready(self.address)
+        if self.enabled_for_logging and _CONNECTION_LOGGER.isEnabledFor(logging.DEBUG):
+            _debug_log(
+                _CONNECTION_LOGGER,
+                message=_ConnectionStatusMessage.POOL_READY,
+                clientId=self._client_id,
+                serverHost=self.address[0],
+                serverPort=self.address[1],
+            )
 
     @property
     def closed(self) -> bool:
@@ -824,38 +835,45 @@ async def _reset(
         interrupt_connections: bool = False,
     ) -> None:
         old_state = self.state
-        async with self.size_cond:
-            if self.closed:
-                return
-            if self.opts.pause_enabled and pause and not self.opts.load_balanced:
+        if self.closed:
+            return
+        if self.opts.pause_enabled and pause and not self.opts.load_balanced:
+            async with self.lock:
                 old_state, self.state = self.state, PoolState.PAUSED
+
+        with self.lock:
             self.gen.inc(service_id)
-            newpid = os.getpid()
-            if self.pid != newpid:
-                self.pid = newpid
-                self.active_sockets = 0
-                self.operation_count = 0
-            if service_id is None:
-                sockets, self.conns = self.conns, collections.deque()
-            else:
-                discard: collections.deque = collections.deque()  # type: ignore[type-arg]
-                keep: collections.deque = collections.deque()  # type: ignore[type-arg]
-                for conn in self.conns:
-                    if conn.service_id == service_id:
-                        discard.append(conn)
-                    else:
-                        keep.append(conn)
-                sockets = discard
+        newpid = os.getpid()
+        if self.pid != newpid:
+            self.pid = newpid
+        with self.lock:
+            self.active_sockets = 0
+            self.operation_count = 0
+        if service_id is None:
+            new_conns = collections.deque()
+            with self.lock:
+                sockets, self.conns = self.conns, new_conns
+        else:
+            discard: collections.deque = collections.deque()  # type: ignore[type-arg]
+            keep: collections.deque = collections.deque()  # type: ignore[type-arg]
+            for conn in self.conns.copy():
+                if conn.service_id == service_id:
+                    discard.append(conn)
+                else:
+                    keep.append(conn)
+            sockets = discard
+            with self.lock:
                 self.conns = keep
 
             if close:
-                self.state = PoolState.CLOSED
+                with self.lock:
+                    self.state = PoolState.CLOSED
             # Clear the wait queue
             self._max_connecting_cond.notify_all()
             self.size_cond.notify_all()
 
             if interrupt_connections:
-                for context in self.active_contexts:
+                for context in self.active_contexts.copy():
                     context.cancel()
 
         listeners = self.opts._event_listeners
@@ -914,9 +932,8 @@ async def update_is_writable(self, is_writable: Optional[bool]) -> None:
         Pool.
         """
         self.is_writable = is_writable
-        async with self.lock:
-            for _socket in self.conns:
-                _socket.update_is_writable(self.is_writable)  # type: ignore[arg-type]
+        for _socket in self.conns.copy():
+            _socket.update_is_writable(self.is_writable)  # type: ignore[arg-type]
 
     async def reset(
         self, service_id: Optional[ObjectId] = None, interrupt_connections: bool = False
@@ -947,12 +964,9 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
 
         if self.opts.max_idle_time_seconds is not None:
             close_conns = []
-            async with self.lock:
-                while (
-                    self.conns
-                    and self.conns[-1].idle_time_seconds() > self.opts.max_idle_time_seconds
-                ):
-                    close_conns.append(self.conns.pop())
+            conns = self.conns.copy()
+            while conns and conns[-1].idle_time_seconds() > self.opts.max_idle_time_seconds:
+                close_conns.append(self.conns.pop())
             if not _IS_SYNC:
                 await asyncio.gather(
                     *[conn.close_conn(ConnectionClosedReason.IDLE) for conn in close_conns],  # type: ignore[func-returns-value]
@@ -963,12 +977,12 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                     await conn.close_conn(ConnectionClosedReason.IDLE)
 
         while True:
+            # There are enough sockets in the pool.
+            if len(self.conns) + self.active_sockets >= self.opts.min_pool_size:
+                return
+            if self.requests >= self.opts.min_pool_size:
+                return
             async with self.size_cond:
-                # There are enough sockets in the pool.
-                if len(self.conns) + self.active_sockets >= self.opts.min_pool_size:
-                    return
-                if self.requests >= self.opts.min_pool_size:
-                    return
                 self.requests += 1
             incremented = False
             try:
@@ -978,15 +992,15 @@ async def remove_stale_sockets(self, reference_generation: int) -> None:
                     if self._pending >= self._max_connecting:
                         return
                     self._pending += 1
-                    incremented = True
+                incremented = True
                 conn = await self.connect()
                 close_conn = False
-                async with self.lock:
-                    # Close connection and return if the pool was reset during
-                    # socket creation or while acquiring the pool lock.
-                    if self.gen.get_overall() != reference_generation:
-                        close_conn = True
-                    if not close_conn:
+                # Close connection and return if the pool was reset during
+                # socket creation or while acquiring the pool lock.
+                if self.gen.get_overall() != reference_generation:
+                    close_conn = True
+                if not close_conn:
+                    async with self.lock:
                         self.conns.appendleft(conn)
                         self.active_contexts.discard(conn.cancel_context)
                 if close_conn:
@@ -1011,11 +1025,11 @@ async def connect(self, handler: Optional[_MongoClientErrorHandler] = None) -> A
         Note that the pool does not keep a reference to the socket -- you
         must call checkin() when you're done with it.
         """
+        # Use a temporary context so that interrupt_connections can cancel creating the socket.
+        tmp_context = _CancellationContext()
+        conn_id = self.next_connection_id
         async with self.lock:
-            conn_id = self.next_connection_id
             self.next_connection_id += 1
-            # Use a temporary context so that interrupt_connections can cancel creating the socket.
-            tmp_context = _CancellationContext()
             self.active_contexts.add(tmp_context)
 
         listeners = self.opts._event_listeners
@@ -1254,7 +1268,7 @@ async def _get_conn(
         try:
             async with self.lock:
                 self.active_sockets += 1
-                incremented = True
+            incremented = True
             while conn is None:
                 # CMAP: we MUST wait for either maxConnecting OR for a socket
                 # to be checked back into the pool.
@@ -1272,7 +1286,8 @@ async def _get_conn(
                         self._raise_if_not_ready(checkout_started_time, emit_event=False)
 
                     try:
-                        conn = self.conns.popleft()
+                        async with self.lock:
+                            conn = self.conns.popleft()
                     except IndexError:
                         self._pending += 1
                 if conn:  # We got a socket from the pool
@@ -1291,10 +1306,11 @@ async def _get_conn(
             if conn:
                 # We checked out a socket but authentication failed.
                 await conn.close_conn(ConnectionClosedReason.ERROR)
+            if incremented:
+                async with self.lock:
+                    self.active_sockets -= 1
             async with self.size_cond:
                 self.requests -= 1
-                if incremented:
-                    self.active_sockets -= 1
                 self.size_cond.notify()
 
             if not emitted_event:
@@ -1330,7 +1346,8 @@ async def checkin(self, conn: AsyncConnection) -> None:
         conn.active = False
         conn.pinned_txn = False
         conn.pinned_cursor = False
-        self.__pinned_sockets.discard(conn)
+        async with self.lock:
+            self.__pinned_sockets.discard(conn)
         listeners = self.opts._event_listeners
         async with self.lock:
             self.active_contexts.discard(conn.cancel_context)
@@ -1371,28 +1388,32 @@ async def checkin(self, conn: AsyncConnection) -> None:
                     )
             else:
                 close_conn = False
-                async with self.lock:
-                    # Hold the lock to ensure this section does not race with
-                    # Pool.reset().
-                    if self.stale_generation(conn.generation, conn.service_id):
-                        close_conn = True
-                    else:
-                        conn.update_last_checkin_time()
-                        conn.update_is_writable(bool(self.is_writable))
+                conn.update_last_checkin_time()
+                conn.update_is_writable(bool(self.is_writable))
+                if self.stale_generation(conn.generation, conn.service_id):
+                    close_conn = True
+                else:
+                    with self.lock:
                         self.conns.appendleft(conn)
+                    with self._max_connecting_cond:
                         # Notify any threads waiting to create a connection.
                         self._max_connecting_cond.notify()
                 if close_conn:
                     await conn.close_conn(ConnectionClosedReason.STALE)
 
-        async with self.size_cond:
-            if txn:
+        async with self.lock:
+            self.active_sockets -= 1
+            self.operation_count -= 1
+
+        if txn:
+            async with self.lock:
                 self.ntxns -= 1
-            elif cursor:
+        elif cursor:
+            async with self.lock:
                 self.ncursors -= 1
+
+        async with self.size_cond:
             self.requests -= 1
-            self.active_sockets -= 1
-            self.operation_count -= 1
             self.size_cond.notify()
 
     async def _perished(self, conn: AsyncConnection) -> bool: