PYTHON-1673 Mongos pinning for sharded transactions

In a sharded transaction, a session is pinned to the mongos server
selected for the initial command. All subsequent commands in the same
transaction are routed to the pinned mongos server.

Author: ShaneHarvey

pymongo/bulk.py

@@ -515,7 +515,7 @@ def execute(self, write_concern, session):
 
         client = self.collection.database.client
         if not write_concern.acknowledged:
-            with client._socket_for_writes() as sock_info:
+            with client._socket_for_writes(session) as sock_info:
                 self.execute_no_results(sock_info, generator)
         else:
             return self.execute_command(generator, write_concern, session)

pymongo/change_stream.py

@@ -107,7 +107,8 @@ def _run_aggregation_cmd(self, session, explicit_session):
         """
         read_preference = self._target._read_preference_for(session)
         client = self._database.client
-        with client._socket_for_reads(read_preference) as (sock_info, slave_ok):
+        with client._socket_for_reads(
+                read_preference, session) as (sock_info, slave_ok):
             pipeline = self._full_pipeline()
             cmd = SON([("aggregate", self._aggregation_target),
                        ("pipeline", pipeline),

pymongo/client_session.py

@@ -246,6 +246,7 @@ def __init__(self, opts):
         self.opts = opts
         self.state = _TxnState.NONE
         self.transaction_id = 0
+        self.pinned_address = None
 
     def active(self):
         return self.state in (_TxnState.STARTING, _TxnState.IN_PROGRESS)
@@ -369,6 +370,7 @@ def start_transaction(self, read_concern=None, write_concern=None,
         self._transaction.state = _TxnState.STARTING
         self._start_retryable_write()
         self._transaction.transaction_id = self._server_session.transaction_id
+        self._transaction.pinned_address = None
         return _TransactionContext(self)
 
     def commit_transaction(self):
@@ -388,6 +390,10 @@ def commit_transaction(self):
         elif state is _TxnState.ABORTED:
             raise InvalidOperation(
                 "Cannot call commitTransaction after calling abortTransaction")
+        elif state is _TxnState.COMMITTED:
+            # We're rerunning the commit, move the state back to "in progress"
+            # so that _in_transaction returns true.
+            self._transaction.state = _TxnState.IN_PROGRESS
 
         try:
             self._finish_transaction_with_retry("commitTransaction")
@@ -441,12 +447,10 @@ def abort_transaction(self):
             self._transaction.state = _TxnState.ABORTED
 
     def _finish_transaction(self, command_name):
-        with self._client._socket_for_writes() as sock_info:
+        with self._client._socket_for_writes(self) as sock_info:
             return self._client.admin._command(
                 sock_info,
                 command_name,
-                txnNumber=self._transaction.transaction_id,
-                autocommit=False,
                 session=self,
                 write_concern=self._transaction.opts.write_concern,
                 parse_write_concern_error=True)
@@ -528,6 +532,17 @@ def _in_transaction(self):
         """True if this session has an active multi-statement transaction."""
         return self._transaction.active()
 
+    @property
+    def _pinned_address(self):
+        """The mongos address this transaction was created on."""
+        if self._transaction.active():
+            return self._transaction.pinned_address
+        return None
+
+    def _pin_mongos(self, server):
+        """Pin this session to the given mongos Server."""
+        self._transaction.pinned_address = server.description.address
+
     def _txn_read_preference(self):
         """Return read preference of this transaction or None."""
         if self._in_transaction:
@@ -542,6 +557,7 @@ def _apply_to(self, command, is_retryable, read_preference):
 
         if not self._in_transaction:
             self._transaction.state = _TxnState.NONE
+            self._transaction.pinned_address = None
 
         if is_retryable:
             command['txnNumber'] = self._server_session.transaction_id

pymongo/collection.py

@@ -185,15 +185,16 @@ def __init__(self, database, name, create=False, codec_options=None,
 
     def _socket_for_reads(self, session):
         return self.__database.client._socket_for_reads(
-            self._read_preference_for(session))
+            self._read_preference_for(session), session)
 
     def _socket_for_primary_reads(self, session):
         read_pref = ((session and session._txn_read_preference())
                      or ReadPreference.PRIMARY)
-        return self.__database.client._socket_for_reads(read_pref), read_pref
+        return self.__database.client._socket_for_reads(
+            read_pref, session), read_pref
 
-    def _socket_for_writes(self):
-        return self.__database.client._socket_for_writes()
+    def _socket_for_writes(self, session):
+        return self.__database.client._socket_for_writes(session)
 
     def _command(self, sock_info, command, slave_ok=False,
                  read_preference=None,
@@ -251,7 +252,7 @@ def __create(self, options, collation, session):
             if "size" in options:
                 options["size"] = float(options["size"])
             cmd.update(options)
-        with self._socket_for_writes() as sock_info:
+        with self._socket_for_writes(session) as sock_info:
             self._command(
                 sock_info, cmd, read_preference=ReadPreference.PRIMARY,
                 write_concern=self._write_concern_for(session),
@@ -1804,7 +1805,7 @@ def create_indexes(self, indexes, session=None, **kwargs):
         """
         common.validate_list('indexes', indexes)
         names = []
-        with self._socket_for_writes() as sock_info:
+        with self._socket_for_writes(session) as sock_info:
             supports_collations = sock_info.max_wire_version >= 5
             def gen_indexes():
                 for index in indexes:
@@ -1844,7 +1845,7 @@ def __create_index(self, keys, index_options, session, **kwargs):
             index_options.pop('collation', None))
         index.update(index_options)
 
-        with self._socket_for_writes() as sock_info:
+        with self._socket_for_writes(session) as sock_info:
             if collation is not None:
                 if sock_info.max_wire_version < 5:
                     raise ConfigurationError(
@@ -2070,7 +2071,7 @@ def drop_index(self, index_or_name, session=None, **kwargs):
             self.__database.name, self.__name, name)
         cmd = SON([("dropIndexes", self.__name), ("index", name)])
         cmd.update(kwargs)
-        with self._socket_for_writes() as sock_info:
+        with self._socket_for_writes(session) as sock_info:
             self._command(sock_info,
                           cmd,
                           read_preference=ReadPreference.PRIMARY,
@@ -2106,7 +2107,7 @@ def reindex(self, session=None, **kwargs):
         """
         cmd = SON([("reIndex", self.__name)])
         cmd.update(kwargs)
-        with self._socket_for_writes() as sock_info:
+        with self._socket_for_writes(session) as sock_info:
             return self._command(
                 sock_info, cmd, read_preference=ReadPreference.PRIMARY,
                 session=session)
@@ -2606,7 +2607,7 @@ def rename(self, new_name, session=None, **kwargs):
         cmd.update(kwargs)
         write_concern = self._write_concern_for_cmd(cmd, session)
 
-        with self._socket_for_writes() as sock_info:
+        with self._socket_for_writes(session) as sock_info:
             with self.__database.client._tmp_session(session) as s:
                 return sock_info.command(
                     'admin', cmd,

pymongo/database.py

@@ -608,7 +608,7 @@ def command(self, command, value=1, check=True,
             read_preference = ((session and session._txn_read_preference())
                                or ReadPreference.PRIMARY)
         with self.__client._socket_for_reads(
-                read_preference) as (sock_info, slave_ok):
+                read_preference, session) as (sock_info, slave_ok):
             return self._command(sock_info, command, slave_ok, value,
                                  check, allowable_errors, read_preference,
                                  codec_options, session=session, **kwargs)
@@ -671,7 +671,7 @@ def list_collections(self, session=None, **kwargs):
         read_pref = ((session and session._txn_read_preference())
                      or ReadPreference.PRIMARY)
         with self.__client._socket_for_reads(
-                read_pref) as (sock_info, slave_okay):
+                read_pref, session) as (sock_info, slave_okay):
             return self._list_collections(
                 sock_info, slave_okay, session, read_preference=read_pref,
                 **kwargs)
@@ -745,7 +745,7 @@ def drop_collection(self, name_or_collection, session=None):
 
         self.__client._purge_index(self.__name, name)
 
-        with self.__client._socket_for_writes() as sock_info:
+        with self.__client._socket_for_writes(session) as sock_info:
             return self._command(
                 sock_info, 'drop', value=_unicode(name),
                 allowable_errors=['ns not found'],
@@ -826,7 +826,7 @@ def current_op(self, include_all=False, session=None):
            Added ``session`` parameter.
         """
         cmd = SON([("currentOp", 1), ("$all", include_all)])
-        with self.__client._socket_for_writes() as sock_info:
+        with self.__client._socket_for_writes(session) as sock_info:
             if sock_info.max_wire_version >= 4:
                 with self.__client._tmp_session(session) as s:
                     return sock_info.command("admin", cmd, session=s,

pymongo/mongo_client.py

@@ -1000,7 +1000,8 @@ def _end_sessions(self, session_ids):
             # Use SocketInfo.command directly to avoid implicitly creating
             # another session.
             with self._socket_for_reads(
-                    ReadPreference.PRIMARY_PREFERRED) as (sock_info, slave_ok):
+                    ReadPreference.PRIMARY_PREFERRED,
+                    None) as (sock_info, slave_ok):
                 if not sock_info.supports_sessions:
                     return
 
@@ -1102,12 +1103,40 @@ def _get_socket(self, server):
                 self.__reset_server(server.description.address)
             raise
 
-    def _socket_for_writes(self):
-        server = self._get_topology().select_server(writable_server_selector)
+    def _select_server(self, server_selector, session, address=None):
+        """Select a server to run an operation on this client.
+
+        :Parameters:
+          - `server_selector`: The server selector to use if the session is
+            not pinned and no address is given.
+          - `session`: The ClientSession for the next operation, or None. May
+            be pinned to a mongos server address.
+          - `address` (optional): Address when sending a message
+            to a specific server, used for getMore.
+        """
+        topology = self._get_topology()
+        address = address or (session and session._pinned_address)
+        if address:
+            # We're running a getMore or this session is pinned to a mongos.
+            server = topology.select_server_by_address(address)
+            if not server:
+                raise AutoReconnect('server %s:%d no longer available'
+                                    % address)
+        else:
+            server = topology.select_server(server_selector)
+            # Pin this session to the selected server if it's performing a
+            # sharded transaction.
+            if server.description.mongos and (session and
+                                              session._in_transaction):
+                session._pin_mongos(server)
+        return server
+
+    def _socket_for_writes(self, session):
+        server = self._select_server(writable_server_selector, session)
         return self._get_socket(server)
 
     @contextlib.contextmanager
-    def _socket_for_reads(self, read_preference):
+    def _socket_for_reads(self, read_preference, 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
@@ -1117,7 +1146,7 @@ def _socket_for_reads(self, read_preference):
         # Thread safe: if the type is single it cannot change.
         topology = self._get_topology()
         single = topology.description.topology_type == TOPOLOGY_TYPE.Single
-        server = topology.select_server(read_preference)
+        server = self._select_server(read_preference, session)
 
         with self._get_socket(server) as sock_info:
             slave_ok = (single and not sock_info.is_mongos) or (
@@ -1140,14 +1169,9 @@ def _send_message_with_response(self, operation, exhaust=False,
             # If needed, restart kill-cursors thread after a fork.
             self._kill_cursors_executor.open()
 
+        server = self._select_server(
+            operation.read_preference, operation.session, address=address)
         topology = self._get_topology()
-        if address:
-            server = topology.select_server_by_address(address)
-            if not server:
-                raise AutoReconnect('server %s:%d no longer available'
-                                    % address)
-        else:
-            server = topology.select_server(operation.read_preference)
 
         # If this is a direct connection to a mongod, *always* set the slaveOk
         # bit. See bullet point 2 in server-selection.rst#topology-type-single.
@@ -1207,8 +1231,7 @@ def is_retrying():
 
         while True:
             try:
-                server = self._get_topology().select_server(
-                    writable_server_selector)
+                server = self._select_server(writable_server_selector, session)
                 supports_session = (
                     session is not None and
                     server.description.retryable_writes_supported)
@@ -1737,7 +1760,7 @@ def drop_database(self, name_or_database, session=None):
                             "of %s or a Database" % (string_type.__name__,))
 
         self._purge_index(name)
-        with self._socket_for_writes() as sock_info:
+        with self._socket_for_writes(session) as sock_info:
             self[name]._command(
                 sock_info,
                 "dropDatabase",
@@ -1878,7 +1901,7 @@ def unlock(self, session=None):
            Added ``session`` parameter.
         """
         cmd = SON([("fsyncUnlock", 1)])
-        with self._socket_for_writes() as sock_info:
+        with self._socket_for_writes(session) as sock_info:
             if sock_info.max_wire_version >= 4:
                 try:
                     with self._tmp_session(session) as s:

pymongo/server_description.py

@@ -187,6 +187,10 @@ def is_writable(self):
     def is_readable(self):
         return self._is_readable
 
+    @property
+    def mongos(self):
+        return self._server_type == SERVER_TYPE.Mongos
+
     @property
     def is_server_type_known(self):
         return self.server_type != SERVER_TYPE.Unknown

test/__init__.py

@@ -162,6 +162,7 @@ def __init__(self):
         self.auth_enabled = False
         self.test_commands_enabled = False
         self.is_mongos = False
+        self.mongoses = []
         self.is_rs = False
         self.has_ipv6 = False
         self.ssl = False
@@ -295,6 +296,17 @@ def _init_client(self):
 
             self.is_mongos = (self.ismaster.get('msg') == 'isdbgrid')
             self.has_ipv6 = self._server_started_with_ipv6()
+            if self.is_mongos:
+                # Check for another mongos on the next port.
+                address = self.client.address
+                next_address = address[0], address[1] + 1
+                self.mongoses.append(address)
+                mongos_client = self._connect(*next_address,
+                                              **self.default_client_options)
+                if mongos_client:
+                    ismaster = mongos_client.admin.command('ismaster')
+                    if ismaster.get('msg') == 'isdbgrid':
+                        self.mongoses.append(next_address)
 
     def init(self):
         with self.conn_lock:
@@ -507,6 +519,13 @@ def require_mongos(self, func):
                              "Must be connected to a mongos",
                              func=func)
 
+    def require_multiple_mongoses(self, func):
+        """Run a test only if the client is connected to a sharded cluster
+        that has 2 mongos nodes."""
+        return self._require(lambda: len(self.mongoses) > 1,
+                             "Must have multiple mongoses available",
+                             func=func)
+
     def require_standalone(self, func):
         """Run a test only if the client is connected to a standalone."""
         return self._require(lambda: not (self.is_mongos or self.is_rs),
@@ -586,13 +605,25 @@ def require_sessions(self, func):
                              "Sessions not supported",
                              func=func)
 
+    def supports_transactions(self):
+        if self.version.at_least(4, 1, 6):
+            return self.is_mongos or self.is_rs
+
+        if self.version.at_least(4, 0):
+            return self.is_rs
+        return False
+
     def require_transactions(self, func):
         """Run a test only if the deployment might support transactions.
 
         *Might* because this does not test the storage engine or FCV.
         """
-        new_func = self.require_version_min(4, 0, 0, -1)(func)
-        return self.require_replica_set(new_func)
+        return self._require(self.supports_transactions,
+                             "Transactions are not supported",
+                             func=func)
+
+    def mongos_seeds(self):
+        return ','.join('%s:%s' % address for address in self.mongoses)
 
     @property
     def supports_reindex(self):

test/test_read_preferences.py

@@ -318,9 +318,9 @@ def __init__(self, *args, **kwargs):
         super(ReadPrefTester, self).__init__(*args, **client_options)
 
     @contextlib.contextmanager
-    def _socket_for_reads(self, read_preference):
+    def _socket_for_reads(self, read_preference, session):
         context = super(ReadPrefTester, self)._socket_for_reads(
-            read_preference)
+            read_preference, session)
         with context as (sock_info, slave_ok):
             self.record_a_read(sock_info.address)
             yield sock_info, slave_ok