SERVER-50589 Change PersistentTaskStore::Update to assert that at least one document was affected

3 files changed, 60 insertions, 41 deletions

diff --git a/src/mongo/db/persistent_task_store.h b/src/mongo/db/persistent_task_store.h
index 3c256144f8e..74e806e0a88 100644
--- a/src/mongo/db/persistent_task_store.h
+++ b/src/mongo/db/persistent_task_store.h
@@ -29,6 +29,8 @@
 
 #pragma once
 
+#include <fmt/format.h>
+
 #include "mongo/db/dbdirectclient.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/operation_context.h"
@@ -39,6 +41,8 @@
 
 namespace mongo {
 
+using namespace fmt::literals;
+
 namespace WriteConcerns {
 
 const WriteConcernOptions kMajorityWriteConcern{WriteConcernOptions::kMajority,
@@ -81,27 +85,19 @@ public:
     void update(OperationContext* opCtx,
                 Query query,
                 const BSONObj& update,
-                const WriteConcernOptions& writeConcern = WriteConcerns::kMajorityWriteConcern,
-                bool upsert = false) {
-        DBDirectClient dbClient(opCtx);
-
-        auto commandResponse = dbClient.runCommand([&] {
-            write_ops::Update updateOp(_storageNss);
-            auto updateModification = write_ops::UpdateModification::parseFromClassicUpdate(update);
-            write_ops::UpdateOpEntry updateEntry(query.obj, updateModification);
-            updateEntry.setMulti(false);
-            updateEntry.setUpsert(upsert);
-            updateOp.setUpdates({updateEntry});
-
-            return updateOp.serialize({});
-        }());
-
-        const auto commandReply = commandResponse->getCommandReply();
-        uassertStatusOK(getStatusFromWriteCommandReply(commandReply));
+                const WriteConcernOptions& writeConcern = WriteConcerns::kMajorityWriteConcern) {
+        _update(opCtx, std::move(query), update, /* upsert */ false, writeConcern);
+    }
 
-        WriteConcernResult ignoreResult;
-        auto latestOpTime = repl::ReplClientInfo::forClient(opCtx->getClient()).getLastOp();
-        uassertStatusOK(waitForWriteConcern(opCtx, latestOpTime, writeConcern, &ignoreResult));
+    /**
+     * Upserts a document that matches the given query using the update modifier specified. Even if
+     * multiple documents match, at most one document will be updated.
+     */
+    void upsert(OperationContext* opCtx,
+                Query query,
+                const BSONObj& update,
+                const WriteConcernOptions& writeConcern = WriteConcerns::kMajorityWriteConcern) {
+        _update(opCtx, std::move(query), update, /* upsert */ true, writeConcern);
     }
 
     /**
@@ -168,6 +164,36 @@ public:
     }
 
 private:
+    void _update(OperationContext* opCtx,
+                 Query query,
+                 const BSONObj& update,
+                 bool upsert,
+                 const WriteConcernOptions& writeConcern = WriteConcerns::kMajorityWriteConcern) {
+        DBDirectClient dbClient(opCtx);
+
+        auto commandResponse = dbClient.runCommand([&] {
+            write_ops::Update updateOp(_storageNss);
+            auto updateModification = write_ops::UpdateModification::parseFromClassicUpdate(update);
+            write_ops::UpdateOpEntry updateEntry(query.obj, updateModification);
+            updateEntry.setMulti(false);
+            updateEntry.setUpsert(upsert);
+            updateOp.setUpdates({updateEntry});
+
+            return updateOp.serialize({});
+        }());
+
+        const auto commandReply = commandResponse->getCommandReply();
+        uassertStatusOK(getStatusFromWriteCommandReply(commandReply));
+
+        uassert(ErrorCodes::NoMatchingDocument,
+                "No matching document found for query {} on namespace {}"_format(
+                    query.toString(), _storageNss.toString()),
+                upsert || commandReply.getIntField("n") > 0);
+
+        WriteConcernResult ignoreResult;
+        auto latestOpTime = repl::ReplClientInfo::forClient(opCtx->getClient()).getLastOp();
+        uassertStatusOK(waitForWriteConcern(opCtx, latestOpTime, writeConcern, &ignoreResult));
+    }
     NamespaceString _storageNss;
 };
 
diff --git a/src/mongo/db/persistent_task_store_test.cpp b/src/mongo/db/persistent_task_store_test.cpp
index 7bf15f890f3..a8167ac9dbb 100644
--- a/src/mongo/db/persistent_task_store_test.cpp
+++ b/src/mongo/db/persistent_task_store_test.cpp
@@ -219,7 +219,7 @@ TEST_F(PersistentTaskStoreTest, TestUpdateOnlyUpdatesOneMatchingDocument) {
     ASSERT_EQ(store.count(opCtx, QUERY("key" << keyToMatch << "min" << expectedUpdatedMin)), 1);
 }
 
-TEST_F(PersistentTaskStoreTest, TestUpdateWithUpsert) {
+TEST_F(PersistentTaskStoreTest, TestUpsert) {
     auto opCtx = operationContext();
 
     PersistentTaskStore<TestTask> store(kNss);
@@ -230,18 +230,15 @@ TEST_F(PersistentTaskStoreTest, TestUpdateWithUpsert) {
     TestTask task(keyToMatch, 0, 0);
     BSONObj taskBson = task.toBSON();
 
-    // Test that the document is not created (upsert == false by default)
-    store.update(opCtx, query, taskBson);
-
     ASSERT_EQ(store.count(opCtx, query), 0);
 
-    // Test that the document is not created with upsert == false
-    store.update(opCtx, query, taskBson, WriteConcerns::kMajorityWriteConcern, false);
-
-    ASSERT_EQ(store.count(opCtx, query), 0);
+    // Test that an attempt to upsert from the update command throws an error.
+    ASSERT_THROWS_CODE(store.update(opCtx, query, taskBson, WriteConcerns::kMajorityWriteConcern),
+                       DBException,
+                       ErrorCodes::NoMatchingDocument);
 
-    // Test that the document is created with upsert == true
-    store.update(opCtx, query, taskBson, WriteConcerns::kMajorityWriteConcern, true);
+    // Test that the document is created when upserted.
+    store.upsert(opCtx, query, taskBson, WriteConcerns::kMajorityWriteConcern);
 
     ASSERT_EQ(store.count(opCtx, query), 1);
 
@@ -251,19 +248,16 @@ TEST_F(PersistentTaskStoreTest, TestUpdateWithUpsert) {
         return true;
     });
 
-    // Verify that updates happen as normal with upsert == true and upsert == false
-    store.update(
-        opCtx, query, BSON("$inc" << BSON("min" << 1)), WriteConcerns::kMajorityWriteConcern, true);
+    // Verify that updates happen as expected with upsert and update
+    store.upsert(
+        opCtx, query, BSON("$inc" << BSON("min" << 1)), WriteConcerns::kMajorityWriteConcern);
     store.forEach(opCtx, query, [&](const TestTask& t) {
         ASSERT_EQ(t.min, 1);
         return true;
     });
 
-    store.update(opCtx,
-                 query,
-                 BSON("$inc" << BSON("min" << 1)),
-                 WriteConcerns::kMajorityWriteConcern,
-                 false);
+    store.update(
+        opCtx, query, BSON("$inc" << BSON("min" << 1)), WriteConcerns::kMajorityWriteConcern);
     store.forEach(opCtx, query, [&](const TestTask& t) {
         ASSERT_EQ(t.min, 2);
         return true;
diff --git a/src/mongo/db/vector_clock_mongod.cpp b/src/mongo/db/vector_clock_mongod.cpp
index 363783685ad..b9d2d7b9ca2 100644
--- a/src/mongo/db/vector_clock_mongod.cpp
+++ b/src/mongo/db/vector_clock_mongod.cpp
@@ -351,11 +351,10 @@ Future<void> VectorClockMongoD::_doWhileQueueNotEmptyOrError(ServiceContext* ser
 
                     PersistentTaskStore<VectorClockDocument> store(
                         NamespaceString::kVectorClockNamespace);
-                    store.update(opCtx,
+                    store.upsert(opCtx,
                                  QUERY(VectorClockDocument::k_idFieldName << vcd.get_id()),
                                  vcd.toBSON(),
-                                 WriteConcerns::kMajorityWriteConcern,
-                                 true /* upsert */);
+                                 WriteConcerns::kMajorityWriteConcern);
                 } else {
                     // Persist as secondary, by asking the primary
                     auto const shardingState = ShardingState::get(opCtx);