From ee7a58e7af9fa14cf837a8f066eec784d1a9e2a6 Mon Sep 17 00:00:00 2001
From: jannaerin <golden.janna@gmail.com>
Date: Tue, 26 Jan 2021 18:36:40 +0000
Subject: SERVER-54040 Resharding may fail to throw WouldChangeOwningShard
 despite destined recipient having changed

---
 .../resharding_replicate_updates_as_insert_delete.js | 20 ++++++++++----------
 src/mongo/db/exec/update_stage.cpp                   |  9 +++++++--
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/jstests/sharding/resharding_replicate_updates_as_insert_delete.js b/jstests/sharding/resharding_replicate_updates_as_insert_delete.js
index 8ee26776719..ee8043d6e3d 100644
--- a/jstests/sharding/resharding_replicate_updates_as_insert_delete.js
+++ b/jstests/sharding/resharding_replicate_updates_as_insert_delete.js
@@ -19,14 +19,14 @@ reshardingTest.setup();
 const donorShardNames = reshardingTest.donorShardNames;
 const testColl = reshardingTest.createShardedCollection({
     ns: 'test.foo',
-    shardKeyPattern: {x: 1},
+    shardKeyPattern: {x: 1, s: 1},
     chunks: [
-        {min: {x: MinKey}, max: {x: 5}, shard: donorShardNames[0]},
-        {min: {x: 5}, max: {x: MaxKey}, shard: donorShardNames[1]},
+        {min: {x: MinKey, s: MinKey}, max: {x: 5, s: 5}, shard: donorShardNames[0]},
+        {min: {x: 5, s: 5}, max: {x: MaxKey, s: MaxKey}, shard: donorShardNames[1]},
     ],
 });
 
-const docToUpdate = ({_id: 0, x: 2, y: 2});
+const docToUpdate = ({_id: 0, x: 2, s: 2, y: 2});
 assert.commandWorked(testColl.insert(docToUpdate));
 
 let retryableWriteTs;
@@ -36,10 +36,10 @@ const mongos = testColl.getMongo();
 const recipientShardNames = reshardingTest.recipientShardNames;
 reshardingTest.withReshardingInBackground(  //
     {
-        newShardKeyPattern: {y: 1},
+        newShardKeyPattern: {y: 1, s: 1},
         newChunks: [
-            {min: {y: MinKey}, max: {y: 5}, shard: recipientShardNames[0]},
-            {min: {y: 5}, max: {y: MaxKey}, shard: recipientShardNames[1]},
+            {min: {y: MinKey, s: MinKey}, max: {y: 5, s: 5}, shard: recipientShardNames[0]},
+            {min: {y: 5, s: 5}, max: {y: MaxKey, s: MaxKey}, shard: recipientShardNames[1]},
         ],
     },
     (tempNs) => {
@@ -65,7 +65,7 @@ reshardingTest.withReshardingInBackground(  //
         // operationTime associated with the write and assert the generated oplog entry is a
         // delete+insert in an applyOps.
         assert.commandFailedWithCode(
-            testColl.update({_id: 0, x: 2}, {$set: {y: 10}}),
+            testColl.update({_id: 0, x: 2, s: 2}, {$set: {y: 10}}),
             ErrorCodes.IllegalOperation,
             'was able to update value under new shard key as ordinary write');
 
@@ -78,7 +78,7 @@ reshardingTest.withReshardingInBackground(  //
         let res;
         assert.soon(
             () => {
-                res = sessionColl.update({_id: 0, x: 2}, {$set: {y: 20}});
+                res = sessionColl.update({_id: 0, x: 2, s: 2}, {$set: {y: 20, s: 3}});
 
                 if (res.nModified === 1) {
                     assert.commandWorked(res);
@@ -95,7 +95,7 @@ reshardingTest.withReshardingInBackground(  //
 
         assert.soon(() => {
             session.startTransaction();
-            res = sessionColl.update({_id: 0, x: 2}, {$set: {y: -30}});
+            res = sessionColl.update({_id: 0, x: 2, s: 3}, {$set: {y: -30, s: 1}});
 
             if (res.nModified === 1) {
                 session.commitTransaction();
diff --git a/src/mongo/db/exec/update_stage.cpp b/src/mongo/db/exec/update_stage.cpp
index d3bb4392c2b..c95cbeda107 100644
--- a/src/mongo/db/exec/update_stage.cpp
+++ b/src/mongo/db/exec/update_stage.cpp
@@ -614,8 +614,13 @@ bool UpdateStage::checkUpdateChangesShardKeyFields(const boost::optional<BSONObj
     }
 
     const auto& newObj = newObjCopy ? *newObjCopy : _doc.getObject();
-    return wasExistingShardKeyUpdated(css, collDesc, newObj, oldObj) ||
-        wasReshardingKeyUpdated(collDesc, newObj, oldObj);
+
+    // It is possible that both the existing and new shard keys are being updated, so we do not want
+    // to short-circuit checking whether either is being modified.
+    const auto existingShardKeyUpdated = wasExistingShardKeyUpdated(css, collDesc, newObj, oldObj);
+    const auto reshardingKeyUpdated = wasReshardingKeyUpdated(collDesc, newObj, oldObj);
+
+    return existingShardKeyUpdated || reshardingKeyUpdated;
 }
 
 bool UpdateStage::wasExistingShardKeyUpdated(CollectionShardingState* css,
-- 
cgit v1.2.1