SERVER-46283 Implement command mirroring for update

3 files changed, 213 insertions, 0 deletions

diff --git a/src/mongo/db/commands/SConscript b/src/mongo/db/commands/SConscript
index d2aff58f8aa..cd31c5daa87 100644
--- a/src/mongo/db/commands/SConscript
+++ b/src/mongo/db/commands/SConscript
@@ -550,6 +550,21 @@ env.CppUnitTest(
 )
 
 env.CppUnitTest(
+    target="command_mirroring_test",
+    source=[
+        'command_mirroring_test.cpp',
+    ],
+    LIBDEPS=[
+        '$BUILD_DIR/mongo/base',
+        '$BUILD_DIR/mongo/db/auth/authorization_manager_global',
+        '$BUILD_DIR/mongo/db/auth/authmocks',
+        '$BUILD_DIR/mongo/db/commands/standalone',
+        '$BUILD_DIR/mongo/db/service_context',
+        '$BUILD_DIR/mongo/unittest/unittest',
+    ],
+)
+
+env.CppUnitTest(
     target="db_commands_test",
     source=[
         "index_filter_commands_test.cpp",
diff --git a/src/mongo/db/commands/command_mirroring_test.cpp b/src/mongo/db/commands/command_mirroring_test.cpp
new file mode 100644
index 00000000000..e9677a77b9d
--- /dev/null
+++ b/src/mongo/db/commands/command_mirroring_test.cpp
@@ -0,0 +1,166 @@
+/**
+ *    Copyright (C) 2020-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include <memory>
+
+#include "mongo/bson/bsonobjbuilder.h"
+#include "mongo/db/client.h"
+#include "mongo/db/commands.h"
+#include "mongo/db/logical_session_id.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/service_context.h"
+#include "mongo/unittest/unittest.h"
+
+namespace mongo {
+namespace {
+
+class CommandMirroringTest : public unittest::Test {
+public:
+    CommandMirroringTest() : _lsid(makeLogicalSessionIdForTest()) {}
+
+    void setUp() override {
+        setGlobalServiceContext(ServiceContext::make());
+        Client::initThread("CommandMirroringTest"_sd);
+    }
+
+    void tearDown() override {
+        auto client = Client::releaseCurrent();
+        client.reset(nullptr);
+    }
+
+    virtual BSONObj makeCommand(std::string, std::vector<BSONObj>) = 0;
+
+    const LogicalSessionId& getLogicalSessionId() const {
+        return _lsid;
+    }
+
+    BSONObj getMirroredCommand(BSONObj& bson) {
+        auto request = OpMsgRequest::fromDBAndBody(kDB, bson);
+        auto cmd = globalCommandRegistry()->findCommand(request.getCommandName());
+        ASSERT(cmd);
+
+        auto opCtx = cc().makeOperationContext();
+        opCtx->setLogicalSessionId(_lsid);
+
+        auto invocation = cmd->parse(opCtx.get(), request);
+        ASSERT(invocation->supportsReadMirroring());
+
+        BSONObjBuilder bob;
+        invocation->appendMirrorableRequest(&bob);
+        return bob.obj();
+    }
+
+    static constexpr auto kDB = "test"_sd;
+
+private:
+    const LogicalSessionId _lsid;
+};
+
+class UpdateCommandTest : public CommandMirroringTest {
+public:
+    BSONObj makeCommand(std::string coll, std::vector<BSONObj> updates) override {
+        BSONObjBuilder bob;
+
+        bob << "update" << coll;
+        BSONArrayBuilder bab;
+        for (auto update : updates) {
+            bab << update;
+        }
+        bob << "updates" << bab.arr();
+        bob << "lsid" << getLogicalSessionId().toBSON();
+
+        return bob.obj();
+    }
+};
+
+TEST_F(UpdateCommandTest, NoQuery) {
+    auto update = BSON("q" << BSONObj() << "u" << BSON("$set" << BSON("_id" << 1)));
+    auto cmd = makeCommand("my_collection", {update});
+
+    auto mirroredObj = getMirroredCommand(cmd);
+
+    ASSERT_EQ(mirroredObj["find"].String(), "my_collection");
+    ASSERT_EQ(mirroredObj["filter"].Obj().toString(), "{}");
+    ASSERT(!mirroredObj.hasField("hint"));
+    ASSERT(mirroredObj["singleBatch"].Bool());
+    ASSERT_EQ(mirroredObj["batchSize"].Int(), 1);
+}
+
+TEST_F(UpdateCommandTest, SingleQuery) {
+    auto update =
+        BSON("q" << BSON("qty" << BSON("$lt" << 50.0)) << "u" << BSON("$inc" << BSON("qty" << 1)));
+    auto cmd = makeCommand("products", {update});
+
+    auto mirroredObj = getMirroredCommand(cmd);
+
+    ASSERT_EQ(mirroredObj["find"].String(), "products");
+    ASSERT_EQ(mirroredObj["filter"].Obj().toString(), "{ qty: { $lt: 50.0 } }");
+    ASSERT(!mirroredObj.hasField("hint"));
+    ASSERT(mirroredObj["singleBatch"].Bool());
+    ASSERT_EQ(mirroredObj["batchSize"].Int(), 1);
+}
+
+TEST_F(UpdateCommandTest, SingleQueryWithHintAndCollation) {
+    auto update = BSON("q" << BSON("price" << BSON("$gt" << 100)) << "hint" << BSON("price" << 1)
+                           << "collation"
+                           << BSON("locale"
+                                   << "fr")
+                           << "u" << BSON("$inc" << BSON("price" << 10)));
+    auto cmd = makeCommand("products", {update});
+
+    auto mirroredObj = getMirroredCommand(cmd);
+
+    ASSERT_EQ(mirroredObj["find"].String(), "products");
+    ASSERT_EQ(mirroredObj["filter"].Obj().toString(), "{ price: { $gt: 100 } }");
+    ASSERT_EQ(mirroredObj["hint"].Obj().toString(), "{ price: 1 }");
+    ASSERT_EQ(mirroredObj["collation"].Obj().toString(), "{ locale: \"fr\" }");
+    ASSERT(mirroredObj["singleBatch"].Bool());
+    ASSERT_EQ(mirroredObj["batchSize"].Int(), 1);
+}
+
+TEST_F(UpdateCommandTest, MultipleQueries) {
+    constexpr int kUpdatesQ = 10;
+    std::vector<BSONObj> updates;
+    for (auto i = 0; i < kUpdatesQ; i++) {
+        updates.emplace_back(BSON("q" << BSON("_id" << BSON("$eq" << i)) << "u"
+                                      << BSON("$inc" << BSON("qty" << 1))));
+    }
+    auto cmd = makeCommand("products", updates);
+
+    auto mirroredObj = getMirroredCommand(cmd);
+
+    ASSERT_EQ(mirroredObj["find"].String(), "products");
+    ASSERT_EQ(mirroredObj["filter"].Obj().toString(), "{ _id: { $eq: 0 } }");
+    ASSERT(!mirroredObj.hasField("hint"));
+    ASSERT(mirroredObj["singleBatch"].Bool());
+    ASSERT_EQ(mirroredObj["batchSize"].Int(), 1);
+}
+
+}  // namespace
+}  // namespace mongo
diff --git a/src/mongo/db/commands/write_commands/write_commands.cpp b/src/mongo/db/commands/write_commands/write_commands.cpp
index 0b1184238ae..172cf11d2f7 100644
--- a/src/mongo/db/commands/write_commands/write_commands.cpp
+++ b/src/mongo/db/commands/write_commands/write_commands.cpp
@@ -28,6 +28,7 @@
  */
 
 #include "mongo/base/init.h"
+#include "mongo/bson/bsonobjbuilder.h"
 #include "mongo/bson/mutable/document.h"
 #include "mongo/bson/mutable/element.h"
 #include "mongo/db/catalog/database_holder.h"
@@ -40,6 +41,7 @@
 #include "mongo/db/json.h"
 #include "mongo/db/lasterror.h"
 #include "mongo/db/matcher/extensions_callback_real.h"
+#include "mongo/db/namespace_string.h"
 #include "mongo/db/ops/delete_request.h"
 #include "mongo/db/ops/parsed_delete.h"
 #include "mongo/db/ops/parsed_update.h"
@@ -328,6 +330,36 @@ private:
         Invocation(const WriteCommand* cmd, const OpMsgRequest& request)
             : InvocationBase(cmd, request), _batch(UpdateOp::parse(request)) {}
 
+        bool supportsReadMirroring() const override {
+            // This would translate to a find command with no filter!
+            // Based on the documentation for `update`, this vector should never be empty.
+            return !_batch.getUpdates().empty();
+        }
+
+        void appendMirrorableRequest(BSONObjBuilder* bob) const override {
+            auto extractQueryDetails = [](const write_ops::Update& query,
+                                          BSONObjBuilder* bob) -> void {
+                auto updates = query.getUpdates();
+                // `supportsReadMirroring()` is responsible for this validation.
+                invariant(!updates.empty());
+
+                // Current design ignores contents of `updates` array except for the first entry.
+                // Assuming identical collation for all elements in `updates`, future design could
+                // use the disjunction primitive (i.e, `$or`) to compile all queries into a single
+                // filter. Such a design also requires a sound way of combining hints.
+                bob->append("filter", updates.front().getQ());
+                if (!updates.front().getHint().isEmpty())
+                    bob->append("hint", updates.front().getHint());
+                if (updates.front().getCollation())
+                    bob->append("collation", *updates.front().getCollation());
+            };
+
+            bob->append("find", _batch.getNamespace().coll());
+            extractQueryDetails(_batch, bob);
+            bob->append("batchSize", 1);
+            bob->append("singleBatch", true);
+        }
+
     private:
         NamespaceString ns() const override {
             return _batch.getNamespace();