SERVER-39842 Allow update to shard key value if retryable write

3 files changed, 257 insertions, 312 deletions

diff --git a/jstests/multiVersion/update_shard_key_disallowed_fcv40.js b/jstests/multiVersion/update_shard_key_disallowed_fcv40.js
index 65c23d0bf1a..1949aac101f 100644
--- a/jstests/multiVersion/update_shard_key_disallowed_fcv40.js
+++ b/jstests/multiVersion/update_shard_key_disallowed_fcv40.js
@@ -5,6 +5,7 @@
     "use strict";
 
     load("jstests/libs/feature_compatibility_version.js");
+    load("jstests/sharding/libs/update_shard_key_helpers.js");
 
     let st = new ShardingTest({
         shards: [{binVersion: "latest"}, {binVersion: "latest"}],
@@ -42,7 +43,9 @@
         let sessionDB = session.getDatabase(kDbName);
 
         // Updates to full shard key
-        shardCollectionAndMoveChunks([{x: 30}, {x: 50}, {x: 80}], {x: 1}, {x: 50}, {x: 80});
+        shardCollectionMoveChunks(
+            st, kDbName, ns, {x: 1}, [{x: 30}, {x: 50}, {x: 80}], {x: 50}, {x: 80});
+        cleanupOrphanedDocs(st, ns);
 
         // Assert that updating the shard key when the doc would remain on the same shard fails for
         // both modify and replacement updates
@@ -60,12 +63,12 @@
 
         // Assert that updating the shard key when the doc would move shards fails for both modify
         // and replacement updates
-        assert.writeError(sessionDB.foo.update({x: 30}, {$set: {x: 100}}));
+        assert.writeError(sessionDB.foo.update({x: 80}, {$set: {x: 3}}));
         // TODO: SERVER-39158. Currently, this update will not fail but will not update the doc.
         // After SERVER-39158 is finished, this should fail.
-        assert.writeOK(sessionDB.foo.update({x: 30}, {x: 100}));
-        assert.eq(1, mongos.getDB(kDbName).foo.find({x: 30}).toArray().length);
-        assert.eq(0, mongos.getDB(kDbName).foo.find({x: 100}).toArray().length);
+        assert.writeOK(sessionDB.foo.update({x: 80}, {x: 3}));
+        assert.eq(1, mongos.getDB(kDbName).foo.find({x: 80}).toArray().length);
+        assert.eq(0, mongos.getDB(kDbName).foo.find({x: 3}).toArray().length);
 
         assert.throws(function() {
             sessionDB.foo.findAndModify({query: {x: 80}, update: {$set: {x: 3}}});
@@ -77,10 +80,14 @@
         mongos.getDB(kDbName).foo.drop();
 
         // Updates to partial shard key
-        shardCollectionAndMoveChunks([{x: 30, y: 4}, {x: 50, y: 50}, {x: 80, y: 100}],
-                                     {x: 1, y: 1},
-                                     {x: 50, y: 50},
-                                     {x: 80, y: 100});
+        shardCollectionMoveChunks(st,
+                                  kDbName,
+                                  ns,
+                                  {x: 1, y: 1},
+                                  [{x: 30, y: 4}, {x: 50, y: 50}, {x: 80, y: 100}],
+                                  {x: 50, y: 50},
+                                  {x: 80, y: 100});
+        cleanupOrphanedDocs(st, ns);
 
         // Assert that updating the shard key when the doc would remain on the same shard fails for
         // both modify and replacement updates
@@ -98,8 +105,8 @@
 
         // Assert that updating the shard key when the doc would move shards fails for both modify
         // and replacement updates
-        assert.writeError(sessionDB.foo.update({x: 30}, {$set: {x: 100}}));
-        assert.writeError(sessionDB.foo.update({x: 30}, {x: 100}));
+        assert.writeError(sessionDB.foo.update({x: 80}, {$set: {x: 3}}));
+        assert.writeError(sessionDB.foo.update({x: 80}, {x: 3}));
         assert.throws(function() {
             sessionDB.foo.findAndModify({query: {x: 80}, update: {$set: {x: 3}}});
         });
@@ -114,7 +121,9 @@
         sessionDB = session.getDatabase(kDbName);
 
         // Updates to full shard key
-        shardCollectionAndMoveChunks([{x: 30}, {x: 50}, {x: 80}], {x: 1}, {x: 50}, {x: 80});
+        shardCollectionMoveChunks(
+            st, kDbName, ns, {x: 1}, [{x: 30}, {x: 50}, {x: 80}], {x: 50}, {x: 80});
+        cleanupOrphanedDocs(st, ns);
 
         // Assert that updating the shard key when the doc would remain on the same shard fails for
         // both modify and replacement updates
diff --git a/jstests/sharding/libs/update_shard_key_helpers.js b/jstests/sharding/libs/update_shard_key_helpers.js
index ce1e39cc7d1..afb79512e01 100644
--- a/jstests/sharding/libs/update_shard_key_helpers.js
+++ b/jstests/sharding/libs/update_shard_key_helpers.js
@@ -106,16 +106,23 @@ function runFindAndModifyCmdSuccess(
     }
 }
 
-function runUpdateCmdFail(st, kDbName, session, sessionDB, inTxn, query, update, multiParamSet) {
+function runUpdateCmdFail(
+    st, kDbName, session, sessionDB, inTxn, query, update, multiParamSet, errorCode) {
     let res;
     if (inTxn) {
         session.startTransaction();
         res = sessionDB.foo.update(query, update, {multi: multiParamSet});
         assert.writeError(res);
+        if (errorCode) {
+            assert.commandFailedWithCode(res, errorCode);
+        }
         session.abortTransaction();
     } else {
         res = sessionDB.foo.update(query, update, {multi: multiParamSet});
         assert.writeError(res);
+        if (errorCode) {
+            assert.commandFailedWithCode(res, errorCode);
+        }
     }
 
     let updatedVal = update["$set"] ? update["$set"] : update;
@@ -495,7 +502,7 @@ function assertCannotUpdateInBulkOpWhenDocsMoveShards(
     } else {
         bulkOp = sessionDB.foo.initializeUnorderedBulkOp();
     }
-    bulkOp.find({"x": 300}).updateOne({$set: {"x": 10}});
+    bulkOp.find({"x": 300}).updateOne({"$set": {"x": 10}});
     bulkOp.find({"x": 4}).updateOne({"$inc": {"a": 1}});
     bulkRes = assert.throws(function() {
         bulkOp.execute();
@@ -512,6 +519,7 @@ function assertCannotUpdateInBulkOpWhenDocsMoveShards(
         assert.eq(0, st.s.getDB(kDbName).foo.find({"x": 4, "a": 3}).itcount());
         assert.eq(1, st.s.getDB(kDbName).foo.find({"x": 4, "a": 4}).itcount());
 
+        assert.writeErrorWithCode(bulkRes, ErrorCodes.InvalidOptions);
         assert.eq(0, bulkRes.nUpserted);
         assert.eq(1, bulkRes.nMatched);
         assert.eq(1, bulkRes.nModified);
@@ -547,7 +555,7 @@ function assertCannotUpdateInBulkOpWhenDocsMoveShards(
         bulkOp = sessionDB.foo.initializeUnorderedBulkOp();
     }
     bulkOp.find({"x": 4}).updateOne({"$inc": {"a": 1}});
-    bulkOp.find({"x": 300}).updateOne({$set: {"x": 10}});
+    bulkOp.find({"x": 300}).updateOne({"$set": {"x": 10}});
     bulkRes = assert.throws(function() {
         bulkOp.execute();
     });
@@ -609,8 +617,8 @@ function assertCannotUpdateInBulkOpWhenDocsMoveShards(
     if (!inTxn) {
         assert.writeErrorWithCode(bulkRes, ErrorCodes.InvalidOptions);
         assert.eq(0, bulkRes.nUpserted);
-        assert.eq(1, bulkRes.nMatched);
-        assert.eq(1, bulkRes.nModified);
+        assert.eq(0, bulkRes.nMatched);
+        assert.eq(0, bulkRes.nModified);
     }
 
     sessionDB.foo.drop();
@@ -629,7 +637,7 @@ function assertCannotUpdateInBulkOpWhenDocsMoveShards(
         bulkOp = sessionDB.foo.initializeUnorderedBulkOp();
     }
     bulkOp.find({"x": 500}).updateOne({"$inc": {"a": 1}});
-    bulkOp.find({"x": 300}).updateOne({$set: {"x": 10}});
+    bulkOp.find({"x": 300}).updateOne({"$set": {"x": 10}});
     bulkRes = assert.throws(function() {
         bulkOp.execute();
     });
@@ -637,17 +645,23 @@ function assertCannotUpdateInBulkOpWhenDocsMoveShards(
         session.abortTransaction();
     }
 
-    // Both updates target the same shard, so neither write will be commited when the second errors.
     assert.eq(1, st.s.getDB(kDbName).foo.find({"x": 300}).itcount());
     assert.eq(0, st.s.getDB(kDbName).foo.find({"x": 10}).itcount());
-    assert.eq(1, st.s.getDB(kDbName).foo.find({"x": 500, "a": 6}).itcount());
-    assert.eq(0, st.s.getDB(kDbName).foo.find({"x": 500, "a": 7}).itcount());
 
     if (!inTxn) {
+        // The first write will succeed
+        assert.eq(0, st.s.getDB(kDbName).foo.find({"x": 500, "a": 6}).itcount());
+        assert.eq(1, st.s.getDB(kDbName).foo.find({"x": 500, "a": 7}).itcount());
+
         assert.writeErrorWithCode(bulkRes, ErrorCodes.InvalidOptions);
         assert.eq(0, bulkRes.nUpserted);
         assert.eq(0, bulkRes.nMatched);
         assert.eq(0, bulkRes.nModified);
+    } else {
+        // Both updates target the same shard, so neither write will be commited when the second
+        // errors.
+        assert.eq(1, st.s.getDB(kDbName).foo.find({"x": 500, "a": 6}).itcount());
+        assert.eq(0, st.s.getDB(kDbName).foo.find({"x": 500, "a": 7}).itcount());
     }
 
     sessionDB.foo.drop();
diff --git a/jstests/sharding/update_shard_key_doc_moves_shards.js b/jstests/sharding/update_shard_key_doc_moves_shards.js
index 37646d79ad9..61c8f739fbf 100644
--- a/jstests/sharding/update_shard_key_doc_moves_shards.js
+++ b/jstests/sharding/update_shard_key_doc_moves_shards.js
@@ -19,311 +19,231 @@
     assert.commandWorked(mongos.adminCommand({enableSharding: kDbName}));
     st.ensurePrimaryShard(kDbName, shard0);
 
-    // ---------------------------------------
-    // Update shard key in multi statement txn
-    // ---------------------------------------
-
-    let session = st.s.startSession();
-    let sessionDB = session.getDatabase(kDbName);
-
-    // ----Single writes in txn----
-
-    // Modify updates
-
-    // upsert : false
-    assertCanUpdatePrimitiveShardKey(st,
-                                     kDbName,
-                                     ns,
-                                     session,
-                                     sessionDB,
-                                     true,
-                                     false,
-                                     [{"x": 300}, {"x": 4}],
-                                     [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                     false);
-    assertCanUpdateDottedPath(st,
-                              kDbName,
-                              ns,
-                              session,
-                              sessionDB,
-                              true,
-                              false,
-                              [{"x.a": 300}, {"x.a": 4}],
-                              [{"$set": {"x": {"a": 30}}}, {"$set": {"x": {"a": 600}}}],
-                              false);
-    assertCanUpdatePartialShardKey(st,
-                                   kDbName,
-                                   ns,
-                                   session,
-                                   sessionDB,
-                                   true,
-                                   false,
-                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
-                                   [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                   false);
-
-    // upsert : true
-    assertCanUpdatePrimitiveShardKey(st,
-                                     kDbName,
-                                     ns,
-                                     session,
-                                     sessionDB,
-                                     true,
-                                     false,
-                                     [{"x": 300}, {"x": 4}],
-                                     [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                     true);
-    assertCanUpdateDottedPath(st,
-                              kDbName,
-                              ns,
-                              session,
-                              sessionDB,
-                              true,
-                              false,
-                              [{"x.a": 300}, {"x.a": 4}],
-                              [{"$set": {"x": {"a": 30}}}, {"$set": {"x": {"a": 600}}}],
-                              true);
-    assertCanUpdatePartialShardKey(st,
-                                   kDbName,
-                                   ns,
-                                   session,
-                                   sessionDB,
-                                   true,
-                                   false,
-                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
-                                   [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                   true);
-
-    // failing cases
-    assertCannotUpdate_id(
-        st, kDbName, ns, session, sessionDB, true, false, {"_id": 300}, {"$set": {"_id": 30}});
-    assertCannotUpdate_idDottedPath(
-        st, kDbName, ns, session, sessionDB, true, false, {"_id.a": 300}, {
-            "$set": {"_id": {"a": 30}}
-        });
-    assertCannotUpdateWithMultiTrue(
-        st, kDbName, ns, session, sessionDB, true, {"x": 300}, {"$set": {"x": 30}});
-    assertCannotUpdateSKToArray(
-        st, kDbName, ns, session, sessionDB, true, false, {"x": 300}, {"$set": {"x": [30]}});
-    assertCannotUnsetSKField(
-        st, kDbName, ns, session, sessionDB, true, false, {"x": 300}, {"$unset": {"x": 1}});
-
-    let docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
-    shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
-    cleanupOrphanedDocs(st, ns);
-    // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
-    mongos.getDB(kDbName).foo.insert({"x": 505});
+    // TODO SERVER-39158: Add tests that replaement style updates work as well.
 
-    // Assert that the document is not updated when the delete fails
-    assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
-        configureFailPoint: "failCommand",
-        mode: "alwaysOn",
-        data: {
-            errorCode: ErrorCodes.WriteConflict,
-            failCommands: ["delete"],
-            failInternalCommands: true
+    function changeShardKeyWhenFailpointsSet(session, sessionDB, runInTxn, isFindAndModify) {
+        let docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
+        shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
+        cleanupOrphanedDocs(st, ns);
+        // TODO: Remove once SERVER-37677 is done. Read so mongos doesn't get ssv causing shard to
+        // abort txn
+        mongos.getDB(kDbName).foo.insert({"x": 505});
+
+        // Assert that the document is not updated when the delete fails
+        assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
+            configureFailPoint: "failCommand",
+            mode: "alwaysOn",
+            data: {
+                errorCode: ErrorCodes.WriteConflict,
+                failCommands: ["delete"],
+                failInternalCommands: true
+            }
+        }));
+        if (isFindAndModify) {
+            runFindAndModifyCmdFail(
+                st, kDbName, session, sessionDB, runInTxn, {"x": 300}, {"$set": {"x": 30}}, false);
+        } else {
+            runUpdateCmdFail(st,
+                             kDbName,
+                             session,
+                             sessionDB,
+                             runInTxn,
+                             {"x": 300},
+                             {"$set": {"x": 30}},
+                             false,
+                             ErrorCodes.WriteConflict);
         }
-    }));
-    session.startTransaction();
-    assert.commandFailedWithCode(sessionDB.foo.update({"x": 300}, {"$set": {"x": 30}}),
-                                 ErrorCodes.WriteConflict);
-    session.abortTransaction();
-    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 300}).itcount());
-    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 30}).itcount());
-    assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
-        configureFailPoint: "failCommand",
-        mode: "off",
-    }));
-
-    // Assert that the document is not updated when the insert fails
-    assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
-        configureFailPoint: "failCommand",
-        mode: "alwaysOn",
-        data: {
-            errorCode: ErrorCodes.NamespaceNotFound,
-            failCommands: ["insert"],
-            failInternalCommands: true
+        assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
+            configureFailPoint: "failCommand",
+            mode: "off",
+        }));
+
+        // Assert that the document is not updated when the insert fails
+        assert.commandWorked(st.rs0.getPrimary().getDB(kDbName).adminCommand({
+            configureFailPoint: "failCommand",
+            mode: "alwaysOn",
+            data: {
+                errorCode: ErrorCodes.NamespaceNotFound,
+                failCommands: ["insert"],
+                failInternalCommands: true
+            }
+        }));
+        if (isFindAndModify) {
+            runFindAndModifyCmdFail(
+                st, kDbName, session, sessionDB, runInTxn, {"x": 300}, {"$set": {"x": 30}}, false);
+        } else {
+            runUpdateCmdFail(st,
+                             kDbName,
+                             session,
+                             sessionDB,
+                             runInTxn,
+                             {"x": 300},
+                             {"$set": {"x": 30}},
+                             false,
+                             ErrorCodes.NamespaceNotFound);
+        }
+        assert.commandWorked(st.rs0.getPrimary().getDB(kDbName).adminCommand({
+            configureFailPoint: "failCommand",
+            mode: "off",
+        }));
+
+        // Assert that the shard key update is not committed when there are no write errors and the
+        // transaction is explicity aborted.
+        if (runInTxn) {
+            session.startTransaction();
+            if (isFindAndModify) {
+                sessionDB.foo.findAndModify({query: {"x": 300}, update: {"$set": {"x": 30}}});
+            } else {
+                assert.commandWorked(sessionDB.foo.update({"x": 300}, {"$set": {"x": 30}}));
+            }
+            session.abortTransaction();
+            assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 300}).itcount());
+            assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 30}).itcount());
         }
-    }));
-    session.startTransaction();
-    assert.commandFailedWithCode(sessionDB.foo.update({"x": 4}, {"$set": {"x": 600}}),
-                                 ErrorCodes.NamespaceNotFound);
-    session.abortTransaction();
-    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 4}).itcount());
-    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 600}).itcount());
-    assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
-        configureFailPoint: "failCommand",
-        mode: "off",
-    }));
-
-    // Assert that the shard key update is not committed when there are no write errors and the
-    // transaction is explicity aborted.
-    session.startTransaction();
-    assert.writeOK(sessionDB.foo.update({"x": 300}, {"$set": {"x": 30}}));
-    session.abortTransaction();
-    assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 300}).itcount());
-    assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 30}).itcount());
 
-    mongos.getDB(kDbName).foo.drop();
+        mongos.getDB(kDbName).foo.drop();
+    }
+
+    // Test that changing the shard key works correctly when either the update or findAndModify
+    // command is used and when the command is run either as a retryable write or in a transaction.
+    // Pairs represent [shouldRunCommandInTxn, runUpdateAsFindAndModifyCmd]
+    //
+    // TODO: SERVER-39843 add [false, true] to run retryable write findAndModify commands that
+    // update the shard key
+    let changeShardKeyOptions = [[false, false], [true, false], [true, true]];
+    changeShardKeyOptions.forEach(function(updatePair) {
+        let runInTxn = updatePair[0];
+        let isFindAndModify = updatePair[1];
+
+        jsTestLog("Testing changing the shard key using " +
+                  (isFindAndModify ? "findAndModify command " : "update command ") +
+                  (runInTxn ? "in transaction " : "as retryable write"));
+
+        let session = st.s.startSession({retryWrites: runInTxn ? false : true});
+        let sessionDB = session.getDatabase(kDbName);
+
+        // Modify updates
+
+        // upsert : false
+        assertCanUpdatePrimitiveShardKey(st,
+                                         kDbName,
+                                         ns,
+                                         session,
+                                         sessionDB,
+                                         runInTxn,
+                                         isFindAndModify,
+                                         [{"x": 300}, {"x": 4}],
+                                         [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
+                                         false);
+        assertCanUpdateDottedPath(st,
+                                  kDbName,
+                                  ns,
+                                  session,
+                                  sessionDB,
+                                  runInTxn,
+                                  isFindAndModify,
+                                  [{"x.a": 300}, {"x.a": 4}],
+                                  [{"$set": {"x": {"a": 30}}}, {"$set": {"x": {"a": 600}}}],
+                                  false);
+        assertCanUpdatePartialShardKey(st,
+                                       kDbName,
+                                       ns,
+                                       session,
+                                       sessionDB,
+                                       runInTxn,
+                                       isFindAndModify,
+                                       [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                       [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
+                                       false);
+
+        // upsert : true
+        assertCanUpdatePrimitiveShardKey(st,
+                                         kDbName,
+                                         ns,
+                                         session,
+                                         sessionDB,
+                                         runInTxn,
+                                         isFindAndModify,
+                                         [{"x": 300}, {"x": 4}],
+                                         [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
+                                         true);
+        assertCanUpdateDottedPath(st,
+                                  kDbName,
+                                  ns,
+                                  session,
+                                  sessionDB,
+                                  runInTxn,
+                                  isFindAndModify,
+                                  [{"x.a": 300}, {"x.a": 4}],
+                                  [{"$set": {"x": {"a": 30}}}, {"$set": {"x": {"a": 600}}}],
+                                  true);
+        assertCanUpdatePartialShardKey(st,
+                                       kDbName,
+                                       ns,
+                                       session,
+                                       sessionDB,
+                                       runInTxn,
+                                       isFindAndModify,
+                                       [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
+                                       [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
+                                       true);
+
+        // failing cases
+        assertCannotUpdate_id(
+            st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"_id": 300}, {
+                "$set": {"_id": 30}
+            });
+        assertCannotUpdate_idDottedPath(
+            st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"_id.a": 300}, {
+                "$set": {"_id": {"a": 30}}
+            });
+        assertCannotUpdateSKToArray(
+            st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"x": 300}, {
+                "$set": {"x": [30]}
+            });
+        assertCannotUnsetSKField(
+            st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"x": 300}, {
+                "$unset": {"x": 1}
+            });
+
+        if (!isFindAndModify) {
+            assertCannotUpdateWithMultiTrue(
+                st, kDbName, ns, session, sessionDB, runInTxn, {"x": 300}, {"$set": {"x": 30}});
+        }
 
-    // TODO SERVER-39158: Add tests that replaement style updates work as well.
+        changeShardKeyWhenFailpointsSet(session, sessionDB, runInTxn, isFindAndModify);
+    });
 
-    // Modify style findAndModify
-
-    // upsert : false
-    assertCanUpdatePrimitiveShardKey(st,
-                                     kDbName,
-                                     ns,
-                                     session,
-                                     sessionDB,
-                                     true,
-                                     true,
-                                     [{"x": 300}, {"x": 4}],
-                                     [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                     false);
-    assertCanUpdateDottedPath(st,
-                              kDbName,
-                              ns,
-                              session,
-                              sessionDB,
-                              true,
-                              true,
-                              [{"x.a": 300}, {"x.a": 4}],
-                              [{"$set": {"x": {"a": 30}}}, {"$set": {"x": {"a": 600}}}],
-                              false);
-    assertCanUpdatePartialShardKey(st,
-                                   kDbName,
-                                   ns,
-                                   session,
-                                   sessionDB,
-                                   true,
-                                   true,
-                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
-                                   [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                   false);
-
-    // upsert : true
-    assertCanUpdatePrimitiveShardKey(st,
-                                     kDbName,
-                                     ns,
-                                     session,
-                                     sessionDB,
-                                     true,
-                                     true,
-                                     [{"x": 300}, {"x": 4}],
-                                     [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                     true);
-    assertCanUpdateDottedPath(st,
-                              kDbName,
-                              ns,
-                              session,
-                              sessionDB,
-                              true,
-                              true,
-                              [{"x.a": 300}, {"x.a": 4}],
-                              [{"$set": {"x": {"a": 30}}}, {"$set": {"x": {"a": 600}}}],
-                              true);
-    assertCanUpdatePartialShardKey(st,
-                                   kDbName,
-                                   ns,
-                                   session,
-                                   sessionDB,
-                                   true,
-                                   true,
-                                   [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
-                                   [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                   true);
-
-    // failing cases
-    assertCannotUpdate_id(
-        st, kDbName, ns, session, sessionDB, true, true, {"_id": 300}, {"$set": {"_id": 30}});
-    assertCannotUpdate_idDottedPath(
-        st, kDbName, ns, session, sessionDB, true, true, {"_id.a": 300}, {
-            "$set": {"_id": {"a": 30}}
-        });
-    assertCannotUpdateSKToArray(
-        st, kDbName, ns, session, sessionDB, true, true, {"x": 300}, {"$set": {"x": [30]}});
-    assertCannotUnsetSKField(
-        st, kDbName, ns, session, sessionDB, true, true, {"x": 300}, {"$unset": {"x": 1}});
-
-    // The update should fail when either of the delete or insert commands fail.
-    docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
-    shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
-    cleanupOrphanedDocs(st, ns);
-    // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
-    mongos.getDB(kDbName).foo.insert({"x": 505});
+    // ----Assert that updating the shard key in a batch with size > 1 fails----
 
-    // Assert that the document is not updated when the delete fails
-    assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
-        configureFailPoint: "failCommand",
-        mode: "alwaysOn",
-        data: {
-            errorCode: ErrorCodes.WriteConflict,
-            failCommands: ["delete"],
-            failInternalCommands: true
-        }
-    }));
-    session.startTransaction();
-    assert.throws(function() {
-        sessionDB.foo.findAndModify({query: {"x": 300}, update: {"$set": {"x": 30}}});
-    });
-    session.abortTransaction();
-    assert.eq(1, st.s.getDB(kDbName).foo.find({"x": 300}).itcount());
-    assert.eq(0, st.s.getDB(kDbName).foo.find({"x": 30}).itcount());
-    assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
-        configureFailPoint: "failCommand",
-        mode: "off",
-    }));
-
-    // Assert that the document is not updated when the insert fails
-    assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
-        configureFailPoint: "failCommand",
-        mode: "alwaysOn",
-        data: {
-            errorCode: ErrorCodes.NamespaceNotFound,
-            failCommands: ["insert"],
-            failInternalCommands: true
-        }
-    }));
-    session.startTransaction();
-    assert.throws(function() {
-        sessionDB.foo.findAndModify({query: {"x": 4}, update: {"$set": {"x": 600}}});
-    });
-    session.abortTransaction();
-    assert.eq(1, st.s.getDB(kDbName).foo.find({"x": 4}).itcount());
-    assert.eq(0, st.s.getDB(kDbName).foo.find({"x": 600}).itcount());
-    assert.commandWorked(st.rs1.getPrimary().getDB(kDbName).adminCommand({
-        configureFailPoint: "failCommand",
-        mode: "off",
-    }));
-
-    // Assert that the shard key update is not committed when there are no write errors and the
-    // transaction is explicity aborted.
-    session.startTransaction();
-    sessionDB.foo.findAndModify({query: {"x": 300}, update: {"$set": {"x": 30}}});
-    session.abortTransaction();
-    assert.eq(1, st.s.getDB(kDbName).foo.find({"x": 300}).itcount());
-    assert.eq(0, st.s.getDB(kDbName).foo.find({"x": 30}).itcount());
+    let session = st.s.startSession({retryWrites: true});
+    let sessionDB = session.getDatabase(kDbName);
 
-    mongos.getDB(kDbName).foo.drop();
+    let docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
 
-    // ----Multiple writes in txn-----
+    assertCannotUpdateInBulkOpWhenDocsMoveShards(st, kDbName, ns, session, sessionDB, false, true);
+    assertCannotUpdateInBulkOpWhenDocsMoveShards(st, kDbName, ns, session, sessionDB, false, false);
 
+    session = st.s.startSession({retryWrites: false});
+    sessionDB = session.getDatabase(kDbName);
     assertCannotUpdateInBulkOpWhenDocsMoveShards(st, kDbName, ns, session, sessionDB, true, true);
     assertCannotUpdateInBulkOpWhenDocsMoveShards(st, kDbName, ns, session, sessionDB, true, false);
 
+    // ----Multiple writes in txn-----
+
     // Update two docs, updating one twice
     shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
     cleanupOrphanedDocs(st, ns);
 
-    // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+    // TODO: Remove once SERVER-37677 is done. Read so mongos doesn't get ssv causing shard to abort
+    // txn
     mongos.getDB(kDbName).foo.insert({"x": 505});
 
     session.startTransaction();
     let id = mongos.getDB(kDbName).foo.find({"x": 500}).toArray()[0]._id;
-    assert.writeOK(sessionDB.foo.update({"x": 500}, {"$set": {"x": 30}}));
-    assert.writeOK(sessionDB.foo.update({"x": 30}, {"$set": {"x": 600}}));
-    assert.writeOK(sessionDB.foo.update({"x": 4}, {"$set": {"x": 50}}));
+    assert.commandWorked(sessionDB.foo.update({"x": 500}, {"$set": {"x": 30}}));
+    assert.commandWorked(sessionDB.foo.update({"x": 30}, {"$set": {"x": 600}}));
+    assert.commandWorked(sessionDB.foo.update({"x": 4}, {"$set": {"x": 50}}));
     session.commitTransaction();
 
     assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 500}).itcount());
@@ -340,13 +260,14 @@
     shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
     cleanupOrphanedDocs(st, ns);
 
-    // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+    // TODO: Remove once SERVER-37677 is done. Read so mongos doesn't get ssv causing shard to abort
+    // txn
     mongos.getDB(kDbName).foo.insert({"x": 505});
 
     session.startTransaction();
-    assert.writeOK(sessionDB.foo.update({"x": 500}, {"$inc": {"a": 1}}));
-    assert.writeOK(sessionDB.foo.update({"x": 500}, {"$set": {"x": 30}}));
-    assert.writeOK(sessionDB.foo.update({"x": 500}, {"$inc": {"a": 1}}));
+    assert.commandWorked(sessionDB.foo.update({"x": 500}, {"$inc": {"a": 1}}));
+    assert.commandWorked(sessionDB.foo.update({"x": 500}, {"$set": {"x": 30}}));
+    assert.commandWorked(sessionDB.foo.update({"x": 500}, {"$inc": {"a": 1}}));
     session.commitTransaction();
 
     assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 500}).itcount());
@@ -359,14 +280,15 @@
     shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
     cleanupOrphanedDocs(st, ns);
 
-    // TODO: Remove once SERVER-37677 is done. Read so don't get ssv causing shard to abort txn
+    // TODO: Remove once SERVER-37677 is done. Read so mongos doesn't get ssv causing shard to abort
+    // txn
     mongos.getDB(kDbName).foo.insert({"x": 505});
 
     // Insert and $inc before moving doc
     session.startTransaction();
     id = mongos.getDB(kDbName).foo.find({"x": 500}).toArray()[0]._id;
-    assert.writeOK(sessionDB.foo.insert({"x": 1, "a": 1}));
-    assert.writeOK(sessionDB.foo.update({"x": 500}, {"$inc": {"a": 1}}));
+    assert.commandWorked(sessionDB.foo.insert({"x": 1, "a": 1}));
+    assert.commandWorked(sessionDB.foo.update({"x": 500}, {"$inc": {"a": 1}}));
     sessionDB.foo.findAndModify({query: {"x": 500}, update: {$set: {"x": 20}}});
     session.commitTransaction();