SERVER-71340 Change helper function in write_one_without_shard_key

5 files changed, 143 insertions, 60 deletions

diff --git a/src/mongo/s/write_ops/SConscript b/src/mongo/s/write_ops/SConscript
index 2a6c1d5ef29..4974289c61a 100644
--- a/src/mongo/s/write_ops/SConscript
+++ b/src/mongo/s/write_ops/SConscript
@@ -30,6 +30,7 @@ env.Library(
         'write_without_shard_key_util.cpp',
     ],
     LIBDEPS_PRIVATE=[
+        '$BUILD_DIR/mongo/db/internal_transactions_feature_flag',
         '$BUILD_DIR/mongo/s/grid',
     ],
 )
diff --git a/src/mongo/s/write_ops/batch_write_op.cpp b/src/mongo/s/write_ops/batch_write_op.cpp
index fbc1179ae03..f42b1d7f894 100644
--- a/src/mongo/s/write_ops/batch_write_op.cpp
+++ b/src/mongo/s/write_ops/batch_write_op.cpp
@@ -433,36 +433,33 @@ StatusWith<bool> BatchWriteOp::targetBatch(
 
         // Check if an updateOne or deleteOne necessitates using the two phase write in the case
         // where the query does not contain a shard key or _id to target by.
-        if (feature_flags::gFeatureFlagUpdateOneWithoutShardKey.isEnabled(
-                serverGlobalParams.featureCompatibility)) {
-            if (auto writeItem = writeOp.getWriteItem();
-                writeItem.getOpType() == BatchedCommandRequest::BatchType_Update ||
-                writeItem.getOpType() == BatchedCommandRequest::BatchType_Delete) {
-
-                bool isMultiWrite = false;
-                BSONObj query;
-
-                if (writeItem.getOpType() == BatchedCommandRequest::BatchType_Update) {
-                    isMultiWrite = writeItem.getUpdate().getMulti();
-                    query = writeItem.getUpdate().getQ();
-                } else {
-                    isMultiWrite = writeItem.getDelete().getMulti();
-                    query = writeItem.getDelete().getQ();
-                }
+        if (auto writeItem = writeOp.getWriteItem();
+            writeItem.getOpType() == BatchedCommandRequest::BatchType_Update ||
+            writeItem.getOpType() == BatchedCommandRequest::BatchType_Delete) {
 
-                if (!isMultiWrite &&
-                    !write_without_shard_key::canTargetQueryByShardKeyOrId(
-                        _opCtx, targeter.getNS(), query)) {
-
-                    // Writes without shard key should be in their own batch.
-                    if (!batchMap.empty()) {
-                        writeOp.cancelWrites(nullptr);
-                        break;
-                    } else {
-                        isWriteWithoutShardKeyOrId = true;
-                    }
-                };
+            bool isMultiWrite = false;
+            BSONObj query;
+
+            if (writeItem.getOpType() == BatchedCommandRequest::BatchType_Update) {
+                isMultiWrite = writeItem.getUpdate().getMulti();
+                query = writeItem.getUpdate().getQ();
+            } else {
+                isMultiWrite = writeItem.getDelete().getMulti();
+                query = writeItem.getDelete().getQ();
             }
+
+            if (!isMultiWrite &&
+                write_without_shard_key::useTwoPhaseProtocol(
+                    _opCtx, targeter.getNS(), true /* isUpdateOrDelete */, query)) {
+
+                // Writes without shard key should be in their own batch.
+                if (!batchMap.empty()) {
+                    writeOp.cancelWrites(nullptr);
+                    break;
+                } else {
+                    isWriteWithoutShardKeyOrId = true;
+                }
+            };
         }
 
         //
diff --git a/src/mongo/s/write_ops/write_without_shard_key_util.cpp b/src/mongo/s/write_ops/write_without_shard_key_util.cpp
index 0d944641c35..816159c58b5 100644
--- a/src/mongo/s/write_ops/write_without_shard_key_util.cpp
+++ b/src/mongo/s/write_ops/write_without_shard_key_util.cpp
@@ -28,6 +28,7 @@
  */
 
 #include "mongo/bson/bsonobj.h"
+#include "mongo/db/internal_transactions_feature_flag_gen.h"
 #include "mongo/logv2/log.h"
 #include "mongo/s/catalog_cache.h"
 #include "mongo/s/grid.h"
@@ -37,24 +38,31 @@
 namespace mongo {
 namespace write_without_shard_key {
 
-bool canTargetQueryByShardKeyOrId(OperationContext* opCtx,
-                                  NamespaceString ns,
-                                  const BSONObj& query) {
-    auto cm =
-        uassertStatusOK(Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, ns));
+bool useTwoPhaseProtocol(OperationContext* opCtx,
+                         NamespaceString nss,
+                         bool isUpdateOrDelete,
+                         const BSONObj& query) {
+    if (!feature_flags::gFeatureFlagUpdateOneWithoutShardKey.isEnabled(
+            serverGlobalParams.featureCompatibility)) {
+        return false;
+    }
 
-    // If a collection is unsharded then the query is always targetable to the primary shard.
-    if (!cm.isSharded()) {
-        return true;
+    // updateOne and deleteOne do not use the two phase protocol for single writes that specify _id
+    // in their queries.
+    if (isUpdateOrDelete && query.hasField("_id")) {
+        return false;
     }
-    auto shardKey =
-        uassertStatusOK(cm.getShardKeyPattern().extractShardKeyFromQuery(opCtx, ns, query));
 
-    if (shardKey.isEmpty() && !query.hasField("_id")) {
-        LOGV2(6970800, "A write without shard key or _id detected");
+    auto cm =
+        uassertStatusOK(Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss));
+
+    // Unsharded collections always target the primary shard.
+    if (!cm.isSharded()) {
         return false;
     }
-    return true;
+    auto shardKey =
+        uassertStatusOK(cm.getShardKeyPattern().extractShardKeyFromQuery(opCtx, nss, query));
+    return shardKey.isEmpty();
 }
 }  // namespace write_without_shard_key
 }  // namespace mongo
diff --git a/src/mongo/s/write_ops/write_without_shard_key_util.h b/src/mongo/s/write_ops/write_without_shard_key_util.h
index 9de749d3f5c..adcc604ca2b 100644
--- a/src/mongo/s/write_ops/write_without_shard_key_util.h
+++ b/src/mongo/s/write_ops/write_without_shard_key_util.h
@@ -37,12 +37,13 @@ namespace mongo {
 namespace write_without_shard_key {
 
 /**
- * Returns true if the write with the given query has enough information to target a single
- * document, meaning the query contains a shard key, an _id, or both.
- */
-bool canTargetQueryByShardKeyOrId(OperationContext* opCtx,
-                                  NamespaceString ns,
-                                  const BSONObj& query);
+ * Returns true if we can use the two phase protocol to complete a single write without shard
+ * key.
+ **/
+bool useTwoPhaseProtocol(OperationContext* opCtx,
+                         NamespaceString ns,
+                         bool isUpdateOrDelete,
+                         const BSONObj& query);
 
 }  // namespace write_without_shard_key
 }  // namespace mongo
diff --git a/src/mongo/s/write_ops/write_without_shard_key_util_test.cpp b/src/mongo/s/write_ops/write_without_shard_key_util_test.cpp
index 3c29c0d6d98..ce163fc1c9d 100644
--- a/src/mongo/s/write_ops/write_without_shard_key_util_test.cpp
+++ b/src/mongo/s/write_ops/write_without_shard_key_util_test.cpp
@@ -27,6 +27,7 @@
  *    it in the license file.
  */
 
+#include "mongo/idl/server_parameter_test_util.h"
 #include "mongo/logv2/log.h"
 #include "mongo/s/catalog_cache_test_fixture.h"
 #include "mongo/s/write_ops/write_without_shard_key_util.h"
@@ -64,30 +65,105 @@ private:
     boost::optional<ChunkManager> _cm;
 };
 
+class UnshardedCollectionTest : public CatalogCacheTestFixture {
+protected:
+    void setUp() override {
+        CatalogCacheTestFixture::setUp();
+        setupNShards(2);
+    }
+
+    OperationContext* getOpCtx() {
+        return operationContext();
+    }
+};
+
 TEST_F(WriteWithoutShardKeyUtilTest, WriteQueryContainingFullShardKeyCanTargetSingleDocument) {
-    auto hasTargetingInfo = write_without_shard_key::canTargetQueryByShardKeyOrId(
-        getOpCtx(), kNss, BSON("a" << 1 << "b" << 1));
-    ASSERT_EQ(hasTargetingInfo, true);
+    RAIIServerParameterControllerForTest featureFlagController(
+        "featureFlagUpdateOneWithoutShardKey", true);
+    auto useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, true /* isUpdateOrDelete */, BSON("a" << 1 << "b" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, false);
+
+    useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, false /* isUpdateOrDelete */, BSON("a" << 1 << "b" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, false);
 }
 
 TEST_F(WriteWithoutShardKeyUtilTest,
        WriteQueryContainingPartialShardKeyCannotTargetSingleDocument) {
-    auto hasTargetingInfo =
-        write_without_shard_key::canTargetQueryByShardKeyOrId(getOpCtx(), kNss, BSON("a" << 1));
-    ASSERT_EQ(hasTargetingInfo, false);
+    RAIIServerParameterControllerForTest featureFlagController(
+        "featureFlagUpdateOneWithoutShardKey", true);
+    auto useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, true /* isUpdateOrDelete */, BSON("a" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, true);
+
+    useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, false /* isUpdateOrDelete */, BSON("a" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, true);
 }
 
-TEST_F(WriteWithoutShardKeyUtilTest, WriteQueryContainingUnderscoreIdCanTargetSingleDocument) {
-    auto hasTargetingInfo =
-        write_without_shard_key::canTargetQueryByShardKeyOrId(getOpCtx(), kNss, BSON("_id" << 1));
-    ASSERT_EQ(hasTargetingInfo, true);
+TEST_F(WriteWithoutShardKeyUtilTest,
+       UpdateAndDeleteQueryContainingUnderscoreIdCanTargetSingleDocument) {
+    RAIIServerParameterControllerForTest featureFlagController(
+        "featureFlagUpdateOneWithoutShardKey", true);
+    auto useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, true /* isUpdateOrDelete */, BSON("_id" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, false);
 }
 
 TEST_F(WriteWithoutShardKeyUtilTest,
        WriteQueryWithoutShardKeyOrUnderscoreIdCannotTargetSingleDocument) {
-    auto hasTargetingInfo =
-        write_without_shard_key::canTargetQueryByShardKeyOrId(getOpCtx(), kNss, BSON("x" << 1));
-    ASSERT_EQ(hasTargetingInfo, false);
+    RAIIServerParameterControllerForTest featureFlagController(
+        "featureFlagUpdateOneWithoutShardKey", true);
+    auto useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, true /* isUpdateOrDelete */, BSON("x" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, true);
+
+    useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, false /* isUpdateOrDelete */, BSON("x" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, true);
+}
+
+TEST_F(WriteWithoutShardKeyUtilTest, FindAndModifyQueryWithOnlyIdMustUseTwoPhaseProtocol) {
+    RAIIServerParameterControllerForTest featureFlagController(
+        "featureFlagUpdateOneWithoutShardKey", true);
+    auto useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, false /* isUpdateOrDelete */, BSON("_id" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, true);
+}
+
+TEST_F(WriteWithoutShardKeyUtilTest, FindAndModifyQueryWithoutShardKeyMustUseTwoPhaseProtocol) {
+    RAIIServerParameterControllerForTest featureFlagController(
+        "featureFlagUpdateOneWithoutShardKey", true);
+    auto useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, false /* isUpdateOrDelete */, BSON("x" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, true);
+}
+
+TEST_F(WriteWithoutShardKeyUtilTest, QueryWithFeatureFlagDisabledDoesNotUseTwoPhaseProtocol) {
+    RAIIServerParameterControllerForTest featureFlagController(
+        "featureFlagUpdateOneWithoutShardKey", false);
+    auto useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, false /* isUpdateOrDelete */, BSON("x" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, false);
+}
+
+TEST_F(UnshardedCollectionTest, UnshardedCollectionDoesNotUseTwoPhaseProtocol) {
+    RAIIServerParameterControllerForTest featureFlagController(
+        "featureFlagUpdateOneWithoutShardKey", true);
+
+    auto future = scheduleRoutingInfoUnforcedRefresh(kNss);
+    expectGetDatabase(kNss);
+
+    // Return an empty collection
+    expectFindSendBSONObjVector(kConfigHostAndPort, {});
+
+    auto cm = *future.default_timed_get();
+    ASSERT(!cm.isSharded());
+
+    auto useTwoPhaseProtocol = write_without_shard_key::useTwoPhaseProtocol(
+        getOpCtx(), kNss, true /* isUpdateOrDelete */, BSON("x" << 1));
+    ASSERT_EQ(useTwoPhaseProtocol, false);
 }
 
 }  // namespace