PYTHON-2673 Connection pinning behavior for load balanced clusters (#630)

Tweak spec test because pymongo unpins cursors eagerly after errors.
Tweak spec test for PoolClearedEvent ordering when MongoDB handshake fails (see DRIVERS-1785).
Only skip killCursors for some error codes.
Rely on SDAM error handling to close the connection after a state change error.
Add service_id to various events.
Retain reference to pinned sockets to prevent premptive closure by CPython's cyclic GC.

(cherry picked from commit c8f32a7)

Author: ShaneHarvey

pymongo/aggregation.py

@@ -161,11 +161,13 @@ def get_cursor(self, session, server, sock_info, slave_ok):
             }
 
         # Create and return cursor instance.
-        return self._cursor_class(
+        cmd_cursor = self._cursor_class(
             self._cursor_collection(cursor), cursor, sock_info.address,
             batch_size=self._batch_size or 0,
             max_await_time_ms=self._max_await_time_ms,
             session=session, explicit_session=self._explicit_session)
+        cmd_cursor._maybe_pin_connection(sock_info)
+        return cmd_cursor
 
 
 class _CollectionAggregationCommand(_AggregationCommand):

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)
+            session, pin=self._client._should_pin_cursor(session))
 
     def _create_cursor(self):
         with self._client._tmp_session(self._session, close=False) as s:

pymongo/client_session.py

@@ -107,6 +107,7 @@
 from bson.timestamp import Timestamp
 
 from pymongo import monotonic
+from pymongo.cursor import _SocketManager
 from pymongo.errors import (ConfigurationError,
                             ConnectionFailure,
                             InvalidOperation,
@@ -116,6 +117,7 @@
 from pymongo.helpers import _RETRYABLE_ERROR_CODES
 from pymongo.read_concern import ReadConcern
 from pymongo.read_preferences import ReadPreference, _ServerMode
+from pymongo.server_type import SERVER_TYPE
 from pymongo.write_concern import WriteConcern
 
 
@@ -291,6 +293,7 @@ def __init__(self, opts):
         self.state = _TxnState.NONE
         self.sharded = False
         self.pinned_address = None
+        self.sock_mgr = None
         self.recovery_token = None
         self.attempt = 0
 
@@ -300,10 +303,29 @@ def active(self):
     def starting(self):
         return self.state == _TxnState.STARTING
 
+    @property
+    def pinned_conn(self):
+        if self.active() and self.sock_mgr:
+            return self.sock_mgr.sock
+        return None
+
+    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
+            self.sock_mgr = _SocketManager(sock_info, False)
+
+    def unpin(self):
+        self.pinned_address = None
+        if self.sock_mgr:
+            self.sock_mgr.close()
+        self.sock_mgr = None
+
     def reset(self):
+        self.unpin()
         self.state = _TxnState.NONE
         self.sharded = False
-        self.pinned_address = None
         self.recovery_token = None
         self.attempt = 0
 
@@ -374,6 +396,9 @@ def _end_session(self, lock):
             try:
                 if self.in_transaction:
                     self.abort_transaction()
+                # It's possible we're still pinned here when the transaction
+                # is in the committed state when the session is discarded.
+                self._unpin()
             finally:
                 self._client._return_server_session(self._server_session, lock)
                 self._server_session = None
@@ -779,14 +804,18 @@ def _pinned_address(self):
             return self._transaction.pinned_address
         return None
 
-    def _pin(self, server):
-        """Pin this session to the given Server."""
-        self._transaction.sharded = True
-        self._transaction.pinned_address = server.description.address
+    @property
+    def _pinned_connection(self):
+        """The connection this transaction was started on."""
+        return self._transaction.pinned_conn
+
+    def _pin(self, server, sock_info):
+        """Pin this session to the given Server or to the given connection."""
+        self._transaction.pin(server, sock_info)
 
     def _unpin(self):
         """Unpin this session from any pinned Server."""
-        self._transaction.pinned_address = None
+        self._transaction.unpin()
 
     def _txn_read_preference(self):
         """Return read preference of this transaction or None."""
@@ -800,9 +829,6 @@ def _apply_to(self, command, is_retryable, read_preference):
         self._server_session.last_use = monotonic.time()
         command['lsid'] = self._server_session.session_id
 
-        if not self.in_transaction:
-            self._transaction.reset()
-
         if is_retryable:
             command['txnNumber'] = self._server_session.transaction_id
             return

pymongo/collection.py

@@ -527,7 +527,8 @@ def _legacy_write(self, sock_info, name, cmd, op_id,
         if publish:
             duration = datetime.datetime.now() - start
             listeners.publish_command_start(
-                cmd, self.__database.name, rqst_id, sock_info.address, op_id)
+                cmd, self.__database.name, rqst_id, sock_info.address, op_id,
+                sock_info.service_id)
             start = datetime.datetime.now()
         try:
             result = sock_info.legacy_write(rqst_id, msg, max_size, False)
@@ -541,12 +542,14 @@ def _legacy_write(self, sock_info, name, cmd, op_id,
                         reply = message._convert_write_result(
                             name, cmd, details)
                         listeners.publish_command_success(
-                            dur, reply, name, rqst_id, sock_info.address, op_id)
+                            dur, reply, name, rqst_id, sock_info.address,
+                            op_id, sock_info.service_id)
                         raise
                 else:
                     details = message._convert_exception(exc)
                 listeners.publish_command_failure(
-                    dur, details, name, rqst_id, sock_info.address, op_id)
+                    dur, details, name, rqst_id, sock_info.address, op_id,
+                    sock_info.service_id)
             raise
         if publish:
             if result is not None:
@@ -556,7 +559,8 @@ def _legacy_write(self, sock_info, name, cmd, op_id,
                 reply = {'ok': 1}
             duration = (datetime.datetime.now() - start) + duration
             listeners.publish_command_success(
-                duration, reply, name, rqst_id, sock_info.address, op_id)
+                duration, reply, name, rqst_id, sock_info.address, op_id,
+                sock_info.service_id)
         return result
 
     def _insert_one(
@@ -2310,9 +2314,9 @@ def _cmd(session, server, sock_info, slave_ok):
                         if exc.code != 26:
                             raise
                         cursor = {'id': 0, 'firstBatch': []}
-                return CommandCursor(coll, cursor, sock_info.address,
-                                     session=s,
-                                     explicit_session=session is not None)
+                cmd_cursor = CommandCursor(
+                    coll, cursor, sock_info.address, session=s,
+                    explicit_session=session is not None)
             else:
                 res = message._first_batch(
                     sock_info, self.__database.name, "system.indexes",
@@ -2322,10 +2326,13 @@ def _cmd(session, server, sock_info, slave_ok):
                 cursor = res["cursor"]
                 # Note that a collection can only have 64 indexes, so there
                 # will never be a getMore call.
-                return CommandCursor(coll, cursor, sock_info.address)
+                cmd_cursor = CommandCursor(coll, cursor, sock_info.address)
+            cmd_cursor._maybe_pin_connection(sock_info)
+            return cmd_cursor
 
         return self.__database.client._retryable_read(
-            _cmd, read_pref, session)
+            _cmd, read_pref, session,
+            pin=self.__database.client._should_pin_cursor(session))
 
     def index_information(self, session=None):
         """Get information on this collection's indexes.
@@ -2416,7 +2423,8 @@ 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)
+            retryable=not cmd._performs_write,
+            pin=self.database.client._should_pin_cursor(session))
 
     def aggregate(self, pipeline, session=None, **kwargs):
         """Perform an aggregation using the aggregation framework on this

pymongo/command_cursor.py

@@ -18,13 +18,14 @@
 
 from bson import _convert_raw_document_lists_to_streams
 from bson.py3compat import integer_types
+from pymongo.cursor import _SocketManager, _CURSOR_CLOSED_ERRORS
 from pymongo.errors import (ConnectionFailure,
                             InvalidOperation,
-                            NotMasterError,
                             OperationFailure)
 from pymongo.message import (_CursorAddress,
                              _GetMore,
                              _RawBatchGetMore)
+from pymongo.response import PinnedResponse
 
 
 class CommandCursor(object):
@@ -38,6 +39,7 @@ def __init__(self, collection, cursor_info, address, retrieved=0,
 
         The parameter 'retrieved' is unused.
         """
+        self.__sock_mgr = None
         self.__collection = collection
         self.__id = cursor_info['id']
         self.__data = deque(cursor_info['firstBatch'])
@@ -76,11 +78,15 @@ def __die(self, synchronous=False):
                 self.__address, self.__collection.full_name)
             if synchronous:
                 self.__collection.database.client._close_cursor_now(
-                    self.__id, address, session=self.__session)
+                    self.__id, address, session=self.__session,
+                    sock_mgr=self.__sock_mgr)
             else:
                 # The cursor will be closed later in a different session.
                 self.__collection.database.client._close_cursor(
                     self.__id, address)
+        if self.__sock_mgr:
+            self.__sock_mgr.close()
+            self.__sock_mgr = None
         self.__end_session(synchronous)
 
     def __end_session(self, synchronous):
@@ -128,52 +134,58 @@ def _post_batch_resume_token(self):
         changeStream aggregate or getMore."""
         return self.__postbatchresumetoken
 
+    def _maybe_pin_connection(self, sock_info):
+        client = self.__collection.database.client
+        if not client._should_pin_cursor(self.__session):
+            return
+        if not self.__sock_mgr:
+            sock_mgr = _SocketManager(sock_info, False)
+            # Ensure the connection gets returned when the entire result is
+            # returned in the first batch.
+            if self.__id == 0:
+                sock_mgr.close()
+            else:
+                self.__sock_mgr = sock_mgr
+
     def __send_message(self, operation):
         """Send a getmore message and handle the response.
         """
-        def kill():
-            self.__killed = True
-            self.__end_session(True)
-
         client = self.__collection.database.client
         try:
-            response = client._run_operation_with_response(
+            response = client._run_operation(
                 operation, self._unpack_response, address=self.__address)
-        except OperationFailure:
-            kill()
-            raise
-        except NotMasterError:
-            # Don't send kill cursors to another server after a "not master"
-            # error. It's completely pointless.
-            kill()
+        except OperationFailure as exc:
+            if exc.code in _CURSOR_CLOSED_ERRORS:
+                # Don't send killCursors because the cursor is already closed.
+                self.__killed = True
+            # Return the session and pinned connection, if necessary.
+            self.close()
             raise
         except ConnectionFailure:
-            # Don't try to send kill cursors on another socket
-            # or to another server. It can cause a _pinValue
-            # assertion on some server releases if we get here
-            # due to a socket timeout.
-            kill()
+            # Don't send killCursors because the cursor is already closed.
+            self.__killed = True
+            # Return the session and pinned connection, if necessary.
+            self.close()
             raise
         except Exception:
-            # Close the cursor
-            self.__die()
+            self.close()
             raise
 
-        from_command = response.from_command
-        reply = response.data
-        docs = response.docs
-
-        if from_command:
-            cursor = docs[0]['cursor']
+        if isinstance(response, PinnedResponse):
+            if not self.__sock_mgr:
+                self.__sock_mgr = _SocketManager(response.socket_info,
+                                                 response.more_to_come)
+        if response.from_command:
+            cursor = response.docs[0]['cursor']
             documents = cursor['nextBatch']
             self.__postbatchresumetoken = cursor.get('postBatchResumeToken')
             self.__id = cursor['id']
         else:
-            documents = docs
-            self.__id = reply.cursor_id
+            documents = response.docs
+            self.__id = response.data.cursor_id
 
         if self.__id == 0:
-            kill()
+            self.__die(True)
         self.__data = deque(documents)
 
     def _unpack_response(self, response, cursor_id, codec_options,
@@ -204,10 +216,9 @@ def _refresh(self):
                                     self.__session,
                                     self.__collection.database.client,
                                     self.__max_await_time_ms,
-                                    False))
+                                    self.__sock_mgr, False))
         else:  # Cursor id is zero nothing else to return
-            self.__killed = True
-            self.__end_session(True)
+            self.__die(True)
 
         return len(self.__data)