PYTHON-4482 Improve performance by making _ServerSessionPool lock-free (#1660) (#1671)

Author: ShaneHarvey

pymongo/client_session.py

@@ -515,9 +515,6 @@ def end_session(self) -> None:
 
         It is an error to use the session after the session has ended.
         """
-        self._end_session(lock=True)
-
-    def _end_session(self, lock: bool) -> None:
         if self._server_session is not None:
             try:
                 if self.in_transaction:
@@ -526,7 +523,7 @@ def _end_session(self, lock: bool) -> None:
                 # is in the committed state when the session is discarded.
                 self._unpin()
             finally:
-                self._client._return_server_session(self._server_session, lock)
+                self._client._return_server_session(self._server_session)
                 self._server_session = None
 
     def _check_ended(self) -> None:
@@ -537,7 +534,7 @@ def __enter__(self) -> ClientSession:
         return self
 
     def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        self._end_session(lock=True)
+        self.end_session()
 
     @property
     def client(self) -> MongoClient:
@@ -1097,7 +1094,7 @@ def inc_transaction_id(self) -> None:
 class _ServerSessionPool(collections.deque):
     """Pool of _ServerSession objects.
 
-    This class is not thread-safe, access it while holding the Topology lock.
+    This class is thread-safe.
     """
 
     def __init__(self, *args: Any, **kwargs: Any):
@@ -1110,8 +1107,11 @@ def reset(self) -> None:
 
     def pop_all(self) -> list[_ServerSession]:
         ids = []
-        while self:
-            ids.append(self.pop().session_id)
+        while True:
+            try:
+                ids.append(self.pop().session_id)
+            except IndexError:
+                break
         return ids
 
     def get_server_session(self, session_timeout_minutes: Optional[int]) -> _ServerSession:
@@ -1123,33 +1123,30 @@ def get_server_session(self, session_timeout_minutes: Optional[int]) -> _ServerS
         self._clear_stale(session_timeout_minutes)
 
         # The most recently used sessions are on the left.
-        while self:
-            s = self.popleft()
+        while True:
+            try:
+                s = self.popleft()
+            except IndexError:
+                break
             if not s.timed_out(session_timeout_minutes):
                 return s
 
         return _ServerSession(self.generation)
 
-    def return_server_session(
-        self, server_session: _ServerSession, session_timeout_minutes: Optional[int]
-    ) -> None:
-        if session_timeout_minutes is not None:
-            self._clear_stale(session_timeout_minutes)
-            if server_session.timed_out(session_timeout_minutes):
-                return
-        self.return_server_session_no_lock(server_session)
-
-    def return_server_session_no_lock(self, server_session: _ServerSession) -> None:
+    def return_server_session(self, server_session: _ServerSession) -> None:
         # Discard sessions from an old pool to avoid duplicate sessions in the
         # child process after a fork.
         if server_session.generation == self.generation and not server_session.dirty:
             self.appendleft(server_session)
 
     def _clear_stale(self, session_timeout_minutes: Optional[int]) -> None:
         # Clear stale sessions. The least recently used are on the right.
-        while self:
-            if self[-1].timed_out(session_timeout_minutes):
-                self.pop()
-            else:
+        while True:
+            try:
+                s = self.pop()
+            except IndexError:
+                break
+            if not s.timed_out(session_timeout_minutes):
+                self.append(s)
                 # The remaining sessions also haven't timed out.
                 break

pymongo/command_cursor.py

@@ -73,7 +73,7 @@ def __init__(
         self.__killed = self.__id == 0
         self.__comment = comment
         if self.__killed:
-            self.__end_session(True)
+            self.__end_session()
 
         if "ns" in cursor_info:  # noqa: SIM401
             self.__ns = cursor_info["ns"]
@@ -112,9 +112,9 @@ def __die(self, synchronous: bool = False) -> None:
             self.__session = None
         self.__sock_mgr = None
 
-    def __end_session(self, synchronous: bool) -> None:
+    def __end_session(self) -> None:
         if self.__session and not self.__explicit_session:
-            self.__session._end_session(lock=synchronous)
+            self.__session.end_session()
             self.__session = None
 
     def close(self) -> None:

pymongo/mongo_client.py

@@ -907,6 +907,8 @@ def target() -> bool:
     def _after_fork(self) -> None:
         """Resets topology in a child after successfully forking."""
         self._init_background(self._topology._pid)
+        # Reset the session pool to avoid duplicate sessions in the child process.
+        self._topology._session_pool.reset()
 
     def _duplicate(self, **kwargs: Any) -> MongoClient:
         args = self.__init_kwargs.copy()
@@ -1679,7 +1681,7 @@ def _cleanup_cursor(
             if cursor_id or conn_mgr:
                 self._close_cursor_soon(cursor_id, address, conn_mgr)
         if session and not explicit_session:
-            session._end_session(lock=locks_allowed)
+            session.end_session()
 
     def _close_cursor_soon(
         self,
@@ -1838,12 +1840,12 @@ def start_session(
         )
 
     def _return_server_session(
-        self, server_session: Union[_ServerSession, _EmptyServerSession], lock: bool
+        self, server_session: Union[_ServerSession, _EmptyServerSession]
     ) -> None:
         """Internal: return a _ServerSession to the pool."""
         if isinstance(server_session, _EmptyServerSession):
             return None
-        return self._topology.return_server_session(server_session, lock)
+        return self._topology.return_server_session(server_session)
 
     def _ensure_session(self, session: Optional[ClientSession] = None) -> Optional[ClientSession]:
         """If provided session is None, lend a temporary session."""

pymongo/topology.py

@@ -669,23 +669,14 @@ def description(self) -> TopologyDescription:
 
     def pop_all_sessions(self) -> list[_ServerSession]:
         """Pop all session ids from the pool."""
-        with self._lock:
-            return self._session_pool.pop_all()
+        return self._session_pool.pop_all()
 
     def get_server_session(self, session_timeout_minutes: Optional[int]) -> _ServerSession:
         """Start or resume a server session, or raise ConfigurationError."""
-        with self._lock:
-            return self._session_pool.get_server_session(session_timeout_minutes)
+        return self._session_pool.get_server_session(session_timeout_minutes)
 
-    def return_server_session(self, server_session: _ServerSession, lock: bool) -> None:
-        if lock:
-            with self._lock:
-                self._session_pool.return_server_session(
-                    server_session, self._description.logical_session_timeout_minutes
-                )
-        else:
-            # Called from a __del__ method, can't use a lock.
-            self._session_pool.return_server_session_no_lock(server_session)
+    def return_server_session(self, server_session: _ServerSession) -> None:
+        self._session_pool.return_server_session(server_session)
 
     def _new_selection(self) -> Selection:
         """A Selection object, initially including all known servers.