PYTHON-1650 Always increment txnNumber before starting a retryable write

Author: ShaneHarvey

pymongo/bulk.py

@@ -160,6 +160,7 @@ def __init__(self, collection, ordered, bypass_document_validation):
  self.uses_array_filters = False
  self.is_retryable = True
  self.retrying = False
+ self.started_retryable_write = False
  # Extra state so that we know where to pick up on a retry attempt.
  self.current_run = None
 
@@ -275,6 +276,11 @@ def _execute_command(self, generator, write_concern, session,
 
  while run.idx_offset < len(run.ops):
  if session:
+ # Start a new retryable write unless one was already
+ # started for this command.
+ if retryable and not self.started_retryable_write:
+ session._start_retryable_write()
+ self.started_retryable_write = True
  session._apply_to(cmd, retryable, ReadPreference.PRIMARY)
  sock_info.send_cluster_time(cmd, session, client)
  check_keys = run.op_type == _INSERT
@@ -300,6 +306,8 @@ def _execute_command(self, generator, write_concern, session,
  _merge_command(run, full_result, run.idx_offset, result)
  # We're no longer in a retry once a command succeeds.
  self.retrying = False
+ self.started_retryable_write = False
+
  if self.ordered and "writeErrors" in result:
  break
  run.idx_offset += len(to_send)

pymongo/client_session.py

@@ -367,7 +367,7 @@ def start_transaction(self, read_concern=None, write_concern=None,
  self._transaction.opts = TransactionOptions(
  read_concern, write_concern, read_preference)
  self._transaction.state = _TxnState.STARTING
- self._server_session._transaction_id += 1
+ self._start_retryable_write()
  self._transaction.transaction_id = self._server_session.transaction_id
  return _TransactionContext(self)
 
@@ -544,7 +544,6 @@ def _apply_to(self, command, is_retryable, read_preference):
  self._transaction.state = _TxnState.NONE
 
  if is_retryable:
- self._server_session._transaction_id += 1
  command['txnNumber'] = self._server_session.transaction_id
  return
 
@@ -574,9 +573,9 @@ def _apply_to(self, command, is_retryable, read_preference):
  command['txnNumber'] = self._server_session.transaction_id
  command['autocommit'] = False
 
- def _retry_transaction_id(self):
+ def _start_retryable_write(self):
  self._check_ended()
- self._server_session.retry_transaction_id()
+ self._server_session.inc_transaction_id()
 
 
 class _ServerSession(object):
@@ -597,8 +596,8 @@ def transaction_id(self):
  """Positive 64-bit integer."""
  return Int64(self._transaction_id)
 
- def retry_transaction_id(self):
- self._transaction_id -= 1
+ def inc_transaction_id(self):
+ self._transaction_id += 1
 
 
 class _ServerSessionPool(collections.deque):

pymongo/mongo_client.py

@@ -1197,6 +1197,14 @@ def _retry_with_session(self, retryable, func, session, bulk):
 
  def is_retrying():
  return bulk.retrying if bulk else retrying
+ # Increment the transaction id up front to ensure any retry attempt
+ # will use the proper txnNumber, even if server or socket selection
+ # fails before the command can be sent.
+ if retryable:
+ session._start_retryable_write()
+ if bulk:
+ bulk.started_retryable_write = True
+
  while True:
  try:
  server = self._get_topology().select_server(
@@ -1211,9 +1219,6 @@ def is_retrying():
  # not support sessions raise the last error.
  raise last_error
  retryable = False
- if is_retrying():
- # Reset the transaction id and retry the operation.
- session._retry_transaction_id()
  return func(session, sock_info, retryable)
  except ServerSelectionTimeoutError:
  if is_retrying():

test/test_retryable_writes.py

@@ -14,6 +14,7 @@
 
 """Test retryable writes."""
 
+import copy
 import json
 import os
 import sys
@@ -169,7 +170,7 @@ def create_tests():
 create_tests()
 
 
-def retryable_single_statement_ops(coll):
+def _retryable_single_statement_ops(coll):
  return [
  (coll.bulk_write, [[InsertOne({}), InsertOne({})]], {}),
  (coll.bulk_write, [[InsertOne({}),
@@ -188,6 +189,11 @@ def retryable_single_statement_ops(coll):
  (coll.find_one_and_replace, [{}, {'a': 3}], {}),
  (coll.find_one_and_update, [{}, {'$set': {'a': 1}}], {}),
  (coll.find_one_and_delete, [{}, {}], {}),
+ ]
+
+
+def retryable_single_statement_ops(coll):
+ return _retryable_single_statement_ops(coll) + [
  # Deprecated methods.
  # Insert with single or multiple documents.
  (coll.insert, [{}], {}),
@@ -500,5 +506,46 @@ def test_batch_splitting_retry_fails(self):
  self.assertEqual(coll.find_one(projection={'_id': True}), {'_id': 1})
 
 
+# TODO: Make this a real integration test where we stepdown the primary.
+class TestRetryableWritesTxnNumber(IgnoreDeprecationsTest):
+ @client_context.require_version_min(3, 6)
+ @client_context.require_replica_set
+ def test_increment_transaction_id_without_sending_command(self):
+ """Test that the txnNumber field is properly incremented, even when
+ the first attempt fails before sending the command.
+ """
+ listener = OvertCommandListener()
+ client = rs_or_single_client(
+ retryWrites=True, event_listeners=[listener])
+ topology = client._topology
+ select_server = topology.select_server
+
+ def raise_connection_err_select_server(*args, **kwargs):
+ # Raise ConnectionFailure on the first attempt and perform
+ # normal selection on the retry attempt.
+ topology.select_server = select_server
+ raise ConnectionFailure('Connection refused')
+
+ for method, args, kwargs in _retryable_single_statement_ops(
+ client.db.retryable_write_test):
+ listener.results.clear()
+ topology.select_server = raise_connection_err_select_server
+ with client.start_session() as session:
+ kwargs = copy.deepcopy(kwargs)
+ kwargs['session'] = session
+ msg = '%s(*%r, **%r)' % (method.__name__, args, kwargs)
+ initial_txn_id = session._server_session.transaction_id
+
+ # Each operation should fail on the first attempt and succeed
+ # on the second.
+ method(*args, **kwargs)
+ self.assertEqual(len(listener.results['started']), 1, msg)
+ retry_cmd = listener.results['started'][0].command
+ sent_txn_id = retry_cmd['txnNumber']
+ final_txn_id = session._server_session.transaction_id
+ self.assertEqual(Int64(initial_txn_id + 1), sent_txn_id, msg)
+ self.assertEqual(sent_txn_id, final_txn_id, msg)
+
+
 if __name__ == '__main__':
  unittest.main()