SERVER-41002 Throw WouldChangeOwningShardError on upsert if shard does not own new doc

1 files changed, 17 insertions, 5 deletions

diff --git a/src/mongo/s/commands/document_shard_key_update_util.cpp b/src/mongo/s/commands/document_shard_key_update_util.cpp
index cf7692b5557..d10b9294110 100644
--- a/src/mongo/s/commands/document_shard_key_update_util.cpp
+++ b/src/mongo/s/commands/document_shard_key_update_util.cpp
@@ -48,12 +48,15 @@ MONGO_FAIL_POINT_DEFINE(hangBeforeInsertOnUpdateShardKey);
 
 /**
  * Calls into the command execution stack to run the given command. Will blindly uassert on any
- * error returned by a command.
+ * error returned by a command. If the original update was sent with {upsert: false}, returns
+ * whether or not we deleted the original doc and inserted the new one sucessfully. If the original
+ * update was sent with {upsert: true}, returns whether or not we inserted the new doc successfully.
  */
 bool executeOperationsAsPartOfShardKeyUpdate(OperationContext* opCtx,
                                              const BSONObj& deleteCmdObj,
                                              const BSONObj& insertCmdObj,
-                                             const StringData db) {
+                                             const StringData db,
+                                             const bool shouldUpsert) {
     auto deleteOpMsg = OpMsgRequest::fromDBAndBody(db, deleteCmdObj);
     auto deleteRequest = BatchedCommandRequest::parseDelete(deleteOpMsg);
 
@@ -62,9 +65,17 @@ bool executeOperationsAsPartOfShardKeyUpdate(OperationContext* opCtx,
 
     ClusterWriter::write(opCtx, deleteRequest, &deleteStats, &deleteResponse);
     uassertStatusOK(deleteResponse.toStatus());
-    // If we do not delete any document, this is essentially equivalent to not matching a doc.
-    if (deleteResponse.getN() != 1)
+    // If shouldUpsert is true, this means the original command specified {upsert: true} and did not
+    // match any docs, so we should not match any when doing this delete. If shouldUpsert is false
+    // and we do not delete any document, this is essentially equivalent to not matching a doc and
+    // we should not insert.
+    if (shouldUpsert) {
+        uassert(ErrorCodes::ConflictingOperationInProgress,
+                "Delete matched a document when it should not have.",
+                deleteResponse.getN() == 0);
+    } else if (deleteResponse.getN() != 1) {
         return false;
+    }
 
     if (MONGO_FAIL_POINT(hangBeforeInsertOnUpdateShardKey)) {
         log() << "Hit hangBeforeInsertOnUpdateShardKey failpoint";
@@ -127,7 +138,8 @@ bool updateShardKeyForDocument(OperationContext* opCtx,
     auto deleteCmdObj = constructShardKeyDeleteCmdObj(nss, updatePreImage, stmtId);
     auto insertCmdObj = constructShardKeyInsertCmdObj(nss, updatePostImage, stmtId);
 
-    return executeOperationsAsPartOfShardKeyUpdate(opCtx, deleteCmdObj, insertCmdObj, nss.db());
+    return executeOperationsAsPartOfShardKeyUpdate(
+        opCtx, deleteCmdObj, insertCmdObj, nss.db(), documentKeyChangeInfo.getShouldUpsert());
 }
 
 TransactionRouter* startTransactionForShardKeyUpdate(OperationContext* opCtx) {