PYTHON-2677 Better wait queue timeout errors for load balanced clusters (#639)

Remove checkout argument in favor of SocketInfo.pin_txn/pin_cursor()

(cherry picked from commit 4c77d7c)

Author: ShaneHarvey

pymongo/change_stream.py

@@ -181,7 +181,7 @@ def _run_aggregation_cmd(self, session, explicit_session):
 
         return self._client._retryable_read(
             cmd.get_cursor, self._target._read_preference_for(session),
-            session, pin=self._client._should_pin_cursor(session))
+            session)
 
     def _create_cursor(self):
         with self._client._tmp_session(self._session, close=False) as s:

pymongo/client_session.py

@@ -313,7 +313,7 @@ def pin(self, server, sock_info):
         self.sharded = True
         self.pinned_address = server.description.address
         if server.description.server_type == SERVER_TYPE.LoadBalancer:
-            sock_info.pinned = True
+            sock_info.pin_txn()
             self.sock_mgr = _SocketManager(sock_info, False)
 
     def unpin(self):

pymongo/collection.py

@@ -2331,8 +2331,7 @@ def _cmd(session, server, sock_info, slave_ok):
             return cmd_cursor
 
         return self.__database.client._retryable_read(
-            _cmd, read_pref, session,
-            pin=self.__database.client._should_pin_cursor(session))
+            _cmd, read_pref, session)
 
     def index_information(self, session=None):
         """Get information on this collection's indexes.
@@ -2423,8 +2422,7 @@ def _aggregate(self, aggregation_command, pipeline, cursor_class, session,
             user_fields={'cursor': {'firstBatch': 1}}, use_cursor=use_cursor)
         return self.__database.client._retryable_read(
             cmd.get_cursor, cmd.get_read_preference(session), session,
-            retryable=not cmd._performs_write,
-            pin=self.database.client._should_pin_cursor(session))
+            retryable=not cmd._performs_write)
 
     def aggregate(self, pipeline, session=None, **kwargs):
         """Perform an aggregation using the aggregation framework on this

pymongo/command_cursor.py

@@ -139,6 +139,7 @@ def _maybe_pin_connection(self, sock_info):
         if not client._should_pin_cursor(self.__session):
             return
         if not self.__sock_mgr:
+            sock_info.pin_cursor()
             sock_mgr = _SocketManager(sock_info, False)
             # Ensure the connection gets returned when the entire result is
             # returned in the first batch.

pymongo/database.py

@@ -525,8 +525,7 @@ def aggregate(self, pipeline, session=None, **kwargs):
                 user_fields={'cursor': {'firstBatch': 1}})
             return self.client._retryable_read(
                 cmd.get_cursor, cmd.get_read_preference(s), s,
-                retryable=not cmd._performs_write,
-                pin=self.client._should_pin_cursor(s))
+                retryable=not cmd._performs_write)
 
     def watch(self, pipeline=None, full_document=None, resume_after=None,
               max_await_time_ms=None, batch_size=None, collation=None,
@@ -839,8 +838,7 @@ def _cmd(session, server, sock_info, slave_okay):
                 **kwargs)
 
         return self.__client._retryable_read(
-            _cmd, read_pref, session,
-            pin=self.client._should_pin_cursor(session))
+            _cmd, read_pref, session)
 
     def list_collection_names(self, session=None, filter=None, **kwargs):
         """Get a list of all the collection names in this database.

pymongo/mongo_client.py

@@ -1291,16 +1291,15 @@ def _get_topology(self):
         return self._topology
 
     @contextlib.contextmanager
-    def _get_socket(self, server, session, pin=False):
+    def _get_socket(self, server, session):
         in_txn = session and session.in_transaction
         with _MongoClientErrorHandler(self, server, session) as err_handler:
             # Reuse the pinned connection, if it exists.
             if in_txn and session._pinned_connection:
                 yield session._pinned_connection
                 return
             with server.get_socket(
-                    self.__all_credentials, checkout=pin,
-                    handler=err_handler) as sock_info:
+                    self.__all_credentials, handler=err_handler) as sock_info:
                 # Pin this session to the selected server or connection.
                 if (in_txn and server.description.server_type in (
                         SERVER_TYPE.Mongos, SERVER_TYPE.LoadBalancer)):
@@ -1351,7 +1350,7 @@ def _socket_for_writes(self, session):
         return self._get_socket(server, session)
 
     @contextlib.contextmanager
-    def _slaveok_for_server(self, read_preference, server, session, pin=False):
+    def _slaveok_for_server(self, read_preference, server, session):
         assert read_preference is not None, "read_preference must not be None"
         # Get a socket for a server matching the read preference, and yield
         # sock_info, slave_ok. Server Selection Spec: "slaveOK must be sent to
@@ -1362,7 +1361,7 @@ def _slaveok_for_server(self, read_preference, server, session, pin=False):
         topology = self._get_topology()
         single = topology.description.topology_type == TOPOLOGY_TYPE.Single
 
-        with self._get_socket(server, session, pin=pin) as sock_info:
+        with self._get_socket(server, session) as sock_info:
             slave_ok = (single and not sock_info.is_mongos) or (
                 read_preference != ReadPreference.PRIMARY)
             yield sock_info, slave_ok
@@ -1389,7 +1388,7 @@ def _should_pin_cursor(self, session):
         return (self.__options.load_balanced and
                 not (session and session.in_transaction))
 
-    def _run_operation(self, operation, unpack_res, pin=False, address=None):
+    def _run_operation(self, operation, unpack_res, address=None):
         """Run a _Query/_GetMore operation and return a Response.
 
         :Parameters:
@@ -1398,7 +1397,6 @@ def _run_operation(self, operation, unpack_res, pin=False, address=None):
           - `address` (optional): Optional address when sending a message
             to a specific server, used for getMore.
         """
-        pin = self._should_pin_cursor(operation.session) or operation.exhaust
         if operation.sock_mgr:
             server = self._select_server(
                 operation.read_preference, operation.session, address=address)
@@ -1409,17 +1407,16 @@ def _run_operation(self, operation, unpack_res, pin=False, address=None):
                     err_handler.contribute_socket(operation.sock_mgr.sock)
                     return server.run_operation(
                         operation.sock_mgr.sock, operation, True,
-                        self._event_listeners, pin, unpack_res)
+                        self._event_listeners, unpack_res)
 
         def _cmd(session, server, sock_info, slave_ok):
             return server.run_operation(
-                sock_info, operation, slave_ok, self._event_listeners, pin,
+                sock_info, operation, slave_ok, self._event_listeners,
                 unpack_res)
 
         return self._retryable_read(
             _cmd, operation.read_preference, operation.session,
-            address=address, retryable=isinstance(operation, message._Query),
-            pin=pin)
+            address=address, retryable=isinstance(operation, message._Query))
 
     def _retry_with_session(self, retryable, func, session, bulk):
         """Execute an operation with at most one consecutive retries
@@ -1491,7 +1488,7 @@ def is_retrying():
                 last_error = exc
 
     def _retryable_read(self, func, read_pref, session, address=None,
-                        retryable=True, pin=False):
+                        retryable=True):
         """Execute an operation with at most one consecutive retries
 
         Returns func()'s return value on success. On error retries the same
@@ -1511,8 +1508,7 @@ def _retryable_read(self, func, read_pref, session, address=None,
                     read_pref, session, address=address)
                 if not server.description.retryable_reads_supported:
                     retryable = False
-                with self._slaveok_for_server(
-                        read_pref, server, session, pin=pin) as (
+                with self._slaveok_for_server(read_pref, server, session) as (
                             sock_info, slave_ok):
                     if retrying and not retryable:
                         # A retry is not possible because this server does

pymongo/pool.py

@@ -566,10 +566,19 @@ def __init__(self, sock, pool, address, id):
         self.service_id = None
         # When executing a transaction in load balancing mode, this flag is
         # set to true to indicate that the session now owns the connection.
-        self.pinned = False
+        self.pinned_txn = False
+        self.pinned_cursor = False
+        self.active = False
+
+    def pin_txn(self):
+        self.pinned_txn = True
+        assert not self.pinned_cursor
+
+    def pin_cursor(self):
+        self.pinned_cursor = True
+        assert not self.pinned_txn
 
     def unpin(self):
-        self.pinned = False
         pool = self.pool_ref()
         if pool:
             pool.return_socket(self)
@@ -1190,6 +1199,8 @@ def __init__(self, address, options, handshake=True):
         # from thinking that a cursor's pinned connection can be GC'd when the
         # cursor is GC'd (see PYTHON-2751).
         self.__pinned_sockets = set()
+        self.ncursors = 0
+        self.ntxns = 0
 
     def _reset(self, close, service_id=None):
         with self.lock:
@@ -1327,15 +1338,15 @@ def connect(self, all_credentials=None):
         return sock_info
 
     @contextlib.contextmanager
-    def get_socket(self, all_credentials, checkout=False, handler=None):
+    def get_socket(self, all_credentials, handler=None):
         """Get a socket from the pool. Use with a "with" statement.
 
         Returns a :class:`SocketInfo` object wrapping a connected
         :class:`socket.socket`.
 
         This method should always be used in a with-statement::
 
-            with pool.get_socket(credentials, checkout) as socket_info:
+            with pool.get_socket(credentials) as socket_info:
                 socket_info.send_message(msg)
                 data = socket_info.receive_message(op_code, request_id)
 
@@ -1347,17 +1358,13 @@ def get_socket(self, all_credentials, checkout=False, handler=None):
 
         :Parameters:
           - `all_credentials`: dict, maps auth source to MongoCredential.
-          - `checkout` (optional): keep socket checked out.
           - `handler` (optional): A _MongoClientErrorHandler.
         """
         listeners = self.opts.event_listeners
         if self.enabled_for_cmap:
             listeners.publish_connection_check_out_started(self.address)
 
         sock_info = self._get_socket(all_credentials)
-        if checkout:
-            self.__pinned_sockets.add(sock_info)
-
         if self.enabled_for_cmap:
             listeners.publish_connection_checked_out(
                 self.address, sock_info.id)
@@ -1368,20 +1375,25 @@ def get_socket(self, all_credentials, checkout=False, handler=None):
             # Note that when pinned is True, the session owns the
             # connection and it is responsible for checking the connection
             # back into the pool.
-            pinned = sock_info.pinned
+            pinned = sock_info.pinned_txn or sock_info.pinned_cursor
             if handler:
                 # Perform SDAM error handling rules while the connection is
                 # still checked out.
                 exc_type, exc_val, _ = sys.exc_info()
                 handler.handle(exc_type, exc_val)
-            if not pinned:
+            if not pinned and sock_info.active:
                 self.return_socket(sock_info)
             raise
-        else:
-            if sock_info.pinned:
+        if sock_info.pinned_txn:
+            with self.lock:
                 self.__pinned_sockets.add(sock_info)
-            elif not checkout:
-                self.return_socket(sock_info)
+                self.ntxns += 1
+        elif sock_info.pinned_cursor:
+            with self.lock:
+                self.__pinned_sockets.add(sock_info)
+                self.ncursors += 1
+        elif sock_info.active:
+            self.return_socket(sock_info)
 
     def _get_socket(self, all_credentials):
         """Get or create a SocketInfo. Can raise ConnectionFailure."""
@@ -1438,6 +1450,7 @@ def _get_socket(self, all_credentials):
                     self.address, ConnectionCheckOutFailedReason.CONN_ERROR)
             raise
 
+        sock_info.active = True
         return sock_info
 
     def return_socket(self, sock_info):
@@ -1446,8 +1459,12 @@ def return_socket(self, sock_info):
         :Parameters:
           - `sock_info`: The socket to check into the pool.
         """
+        txn = sock_info.pinned_txn
+        cursor = sock_info.pinned_cursor
+        sock_info.active = False
+        sock_info.pinned_txn = False
+        sock_info.pinned_cursor = False
         self.__pinned_sockets.discard(sock_info)
-        sock_info.pinned = False
         listeners = self.opts.event_listeners
         if self.enabled_for_cmap:
             listeners.publish_connection_checked_in(self.address, sock_info.id)
@@ -1464,7 +1481,6 @@ def return_socket(self, sock_info):
                         ConnectionClosedReason.ERROR)
             else:
                 with self.lock:
-                    assert sock_info not in self.sockets
                     # Hold the lock to ensure this section does not race with
                     # Pool.reset().
                     if self.stale_generation(sock_info.generation,
@@ -1477,6 +1493,10 @@ def return_socket(self, sock_info):
 
         self._socket_semaphore.release()
         with self.lock:
+            if txn:
+                self.ntxns -= 1
+            elif cursor:
+                self.ncursors -= 1
             self.active_sockets -= 1
 
     def _perished(self, sock_info):
@@ -1518,9 +1538,18 @@ def _raise_wait_queue_timeout(self):
         if self.enabled_for_cmap:
             listeners.publish_connection_check_out_failed(
                 self.address, ConnectionCheckOutFailedReason.TIMEOUT)
+        if self.opts.load_balanced:
+            other_ops = self.active_sockets - self.ncursors - self.ntxns
+            raise ConnectionFailure(
+                'Timeout waiting for connection from the connection pool. '
+                'maxPoolSize: %s, connections in use by cursors: %s, '
+                'connections in use by transactions: %s, connections in use '
+                'by other operations: %s, wait_queue_timeout: %s' % (
+                    self.opts.max_pool_size, self.ncursors, self.ntxns,
+                    other_ops, self.opts.wait_queue_timeout))
         raise ConnectionFailure(
-            'Timed out while checking out a connection from connection pool '
-            'with max_size %r and wait_queue_timeout %r' % (
+            'Timed out while checking out a connection from connection pool. '
+            'maxPoolSize: %s, wait_queue_timeout: %s' % (
                 self.opts.max_pool_size, self.opts.wait_queue_timeout))
 
     def __del__(self):

pymongo/server.py

@@ -69,7 +69,7 @@ def request_check(self):
         self._monitor.request_check()
 
     def run_operation(self, sock_info, operation, set_slave_okay, listeners,
-                      pin, unpack_res):
+                      unpack_res):
         """Run a _Query or _GetMore operation and return a Response object.
 
         This method is used only to run _Query/_GetMore operations from
@@ -81,7 +81,6 @@ def run_operation(self, sock_info, operation, set_slave_okay, listeners,
           - `set_slave_okay`: Pass to operation.get_message.
           - `all_credentials`: dict, maps auth source to MongoCredential.
           - `listeners`: Instance of _EventListeners or None.
-          - `pin`: If True, then this is a pinned cursor operation.
           - `unpack_res`: A callable that decodes the wire protocol response.
         """
         duration = None
@@ -170,7 +169,8 @@ def run_operation(self, sock_info, operation, set_slave_okay, listeners,
                 docs = _decode_all_selective(
                     decrypted, operation.codec_options, user_fields)
 
-        if pin:
+        if client._should_pin_cursor(operation.session) or operation.exhaust:
+            sock_info.pin_cursor()
             if isinstance(reply, _OpMsg):
                 # In OP_MSG, the server keeps sending only if the
                 # more_to_come flag is set.
@@ -200,8 +200,8 @@ def run_operation(self, sock_info, operation, set_slave_okay, listeners,
 
         return response
 
-    def get_socket(self, all_credentials, checkout=False, handler=None):
-        return self.pool.get_socket(all_credentials, checkout, handler)
+    def get_socket(self, all_credentials, handler=None):
+        return self.pool.get_socket(all_credentials, handler)
 
     @property
     def description(self):