SERVER-58060 Add new internal FindAndModifyImageLookup aggregation stage

1 files changed, 301 insertions, 0 deletions

diff --git a/src/mongo/db/pipeline/document_source_find_and_modify_image_lookup_test.cpp b/src/mongo/db/pipeline/document_source_find_and_modify_image_lookup_test.cpp
new file mode 100644
index 00000000000..5ec22e51a4c
--- /dev/null
+++ b/src/mongo/db/pipeline/document_source_find_and_modify_image_lookup_test.cpp
@@ -0,0 +1,301 @@
+/**
+ *    Copyright (C) 2021-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.
+ */
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kTest
+
+#include "mongo/platform/basic.h"
+
+#include <vector>
+
+#include "mongo/db/exec/document_value/document.h"
+#include "mongo/db/exec/document_value/document_value_test_util.h"
+#include "mongo/db/logical_session_id_helpers.h"
+#include "mongo/db/matcher/matcher.h"
+#include "mongo/db/namespace_string.h"
+#include "mongo/db/pipeline/aggregation_context_fixture.h"
+#include "mongo/db/pipeline/document_source.h"
+#include "mongo/db/pipeline/document_source_find_and_modify_image_lookup.h"
+#include "mongo/db/pipeline/document_source_mock.h"
+#include "mongo/db/pipeline/process_interface/stub_lookup_single_document_process_interface.h"
+#include "mongo/db/repl/image_collection_entry_gen.h"
+#include "mongo/logv2/log.h"
+#include "mongo/unittest/unittest.h"
+#include "mongo/util/assert_util.h"
+
+namespace mongo {
+namespace {
+
+/**
+ * Creates OplogEntry with given field values.
+ */
+repl::OplogEntry makeOplogEntry(
+    repl::OpTime opTime,
+    repl::OpTypeEnum opType,
+    NamespaceString nss,
+    BSONObj oField,
+    OperationSessionInfo sessionInfo,
+    boost::optional<repl::OpTime> preImageOpTime = boost::none,
+    boost::optional<repl::OpTime> postImageOpTime = boost::none,
+    boost::optional<repl::RetryImageEnum> needsRetryImage = boost::none) {
+    static const UUID* uuid = new UUID(UUID::gen());
+    return {
+        repl::DurableOplogEntry(opTime,                           // optime
+                                boost::none,                      // hash
+                                opType,                           // opType
+                                nss,                              // namespace
+                                *uuid,                            // uuid
+                                boost::none,                      // fromMigrate
+                                repl::OplogEntry::kOplogVersion,  // version
+                                oField,                           // o
+                                boost::none,                      // o2
+                                sessionInfo,                      // sessionInfo
+                                boost::none,                      // upsert
+                                Date_t(),                         // wall clock time
+                                {1},                              // statement ids
+                                boost::none,     // optime of previous write within same transaction
+                                preImageOpTime,  // pre-image optime
+                                postImageOpTime,    // post-image optime
+                                boost::none,        // ShardId of resharding recipient
+                                boost::none,        // _id
+                                needsRetryImage)};  // needsRetryImage
+}
+
+struct MockMongoInterface final : public StubMongoProcessInterface {
+    MockMongoInterface(std::vector<Document> documentsForLookup = {})
+        : _documentsForLookup{std::move(documentsForLookup)} {}
+
+    BSONObj getCollectionOptions(OperationContext* opCtx, const NamespaceString& nss) {
+        static const UUID* oplog_uuid = new UUID(UUID::gen());
+        return BSON("uuid" << *oplog_uuid);
+    }
+
+    boost::optional<Document> lookupSingleDocument(
+        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+        const NamespaceString& nss,
+        UUID collectionUUID,
+        const Document& documentKey,
+        boost::optional<BSONObj> readConcern) final {
+        Matcher matcher(documentKey.toBson(), expCtx);
+        auto it = std::find_if(_documentsForLookup.begin(),
+                               _documentsForLookup.end(),
+                               [&](const Document& lookedUpDoc) {
+                                   return matcher.matches(lookedUpDoc.toBson(), nullptr);
+                               });
+        return (it != _documentsForLookup.end() ? *it : boost::optional<Document>{});
+    }
+
+    // These documents are used to feed the 'lookupSingleDocument' method.
+    std::vector<Document> _documentsForLookup;
+};
+
+// This provides access to getExpCtx(), but we'll use a different name for this test suite.
+using FindAndModifyImageLookupTest = AggregationContextFixture;
+
+TEST_F(FindAndModifyImageLookupTest, NoopWhenEntryDoesNotHaveNeedsRetryImageField) {
+    auto imageLookup = DocumentSourceFindAndModifyImageLookup::create(getExpCtx());
+    const auto sessionId = makeLogicalSessionIdForTest();
+    OperationSessionInfo sessionInfo;
+    sessionInfo.setSessionId(sessionId);
+    sessionInfo.setTxnNumber(1);
+    const auto opTime = repl::OpTime(Timestamp(2, 1), 1);
+    const auto preImageOpTime = repl::OpTime(Timestamp(1, 1), 1);
+    const auto oplogEntryBson = makeOplogEntry(opTime,
+                                               repl::OpTypeEnum::kNoop,
+                                               NamespaceString("test.foo"),
+                                               BSON("a" << 1),
+                                               sessionInfo,
+                                               preImageOpTime)
+                                    .getEntry()
+                                    .toBSON();
+    auto mock = DocumentSourceMock::createForTest(Document(oplogEntryBson), getExpCtx());
+    imageLookup->setSource(mock.get());
+    // Mock out the foreign collection.
+    getExpCtx()->mongoProcessInterface =
+        std::make_unique<MockMongoInterface>(std::vector<Document>{});
+
+    auto next = imageLookup->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    Document expected = Document(oplogEntryBson);
+    ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected);
+
+    ASSERT_TRUE(imageLookup->getNext().isEOF());
+    ASSERT_TRUE(imageLookup->getNext().isEOF());
+    ASSERT_TRUE(imageLookup->getNext().isEOF());
+}
+
+TEST_F(FindAndModifyImageLookupTest, ShouldNotForgeImageEntryWhenImageDocMissing) {
+    auto imageLookup = DocumentSourceFindAndModifyImageLookup::create(getExpCtx());
+    const auto sessionId = makeLogicalSessionIdForTest();
+    OperationSessionInfo sessionInfo;
+    sessionInfo.setSessionId(sessionId);
+    sessionInfo.setTxnNumber(1);
+    const auto opTime = repl::OpTime(Timestamp(2, 1), 1);
+    const auto oplogEntryBson = makeOplogEntry(opTime,
+                                               repl::OpTypeEnum::kNoop,
+                                               NamespaceString("test.foo"),
+                                               BSON("a" << 1),
+                                               sessionInfo,
+                                               boost::none /* preImageOpTime */,
+                                               boost::none /* postImageOpTime */,
+                                               repl::RetryImageEnum::kPreImage)
+                                    .getEntry()
+                                    .toBSON();
+    auto mock = DocumentSourceMock::createForTest(Document(oplogEntryBson), getExpCtx());
+    imageLookup->setSource(mock.get());
+
+    // Mock out the foreign collection.
+    getExpCtx()->mongoProcessInterface =
+        std::make_unique<MockMongoInterface>(std::vector<Document>{});
+
+    auto next = imageLookup->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    Document expected = Document(oplogEntryBson);
+    ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected);
+
+    ASSERT_TRUE(imageLookup->getNext().isEOF());
+    ASSERT_TRUE(imageLookup->getNext().isEOF());
+    ASSERT_TRUE(imageLookup->getNext().isEOF());
+}
+
+TEST_F(FindAndModifyImageLookupTest, ShouldNotForgeImageEntryWhenImageDocHasDifferentTxnNumber) {
+    auto imageLookup = DocumentSourceFindAndModifyImageLookup::create(getExpCtx());
+    const auto sessionId = makeLogicalSessionIdForTest();
+    OperationSessionInfo sessionInfo;
+    sessionInfo.setSessionId(sessionId);
+    sessionInfo.setTxnNumber(1);
+    const auto ts = Timestamp(2, 1);
+    const auto opTime = repl::OpTime(ts, 1);
+    const auto oplogEntryBson = makeOplogEntry(opTime,
+                                               repl::OpTypeEnum::kNoop,
+                                               NamespaceString("test.foo"),
+                                               BSON("a" << 1),
+                                               sessionInfo,
+                                               boost::none /* preImageOpTime */,
+                                               boost::none /* postImageOpTime */,
+                                               repl::RetryImageEnum::kPreImage)
+                                    .getEntry()
+                                    .toBSON();
+    auto mock = DocumentSourceMock::createForTest(Document(oplogEntryBson), getExpCtx());
+    imageLookup->setSource(mock.get());
+
+    // Create an 'ImageEntry' with a higher 'txnNumber'.
+    const auto preImage = BSON("a" << 2);
+    repl::ImageEntry imageEntry;
+    imageEntry.set_id(sessionId);
+    imageEntry.setTxnNumber(2);
+    imageEntry.setTs(ts);
+    imageEntry.setImageKind(repl::RetryImageEnum::kPreImage);
+    imageEntry.setImage(preImage);
+    // Mock out the foreign collection.
+    getExpCtx()->mongoProcessInterface =
+        std::make_unique<MockMongoInterface>(std::vector<Document>{Document{imageEntry.toBSON()}});
+
+    auto next = imageLookup->getNext();
+    ASSERT_TRUE(next.isAdvanced());
+    Document expected = Document(oplogEntryBson);
+    ASSERT_DOCUMENT_EQ(next.releaseDocument(), expected);
+
+    ASSERT_TRUE(imageLookup->getNext().isEOF());
+    ASSERT_TRUE(imageLookup->getNext().isEOF());
+    ASSERT_TRUE(imageLookup->getNext().isEOF());
+}
+
+
+TEST_F(FindAndModifyImageLookupTest, ShouldForgeImageEntryWhenMatchingImageDocIsFound) {
+    std::vector<repl::RetryImageEnum> cases{repl::RetryImageEnum::kPreImage,
+                                            repl::RetryImageEnum::kPostImage};
+    for (auto imageType : cases) {
+        LOGV2(5806002,
+              "ForgeImageEntryTestCase",
+              "imageType"_attr = repl::RetryImage_serializer(imageType));
+        auto imageLookup = DocumentSourceFindAndModifyImageLookup::create(getExpCtx());
+        const auto sessionId = makeLogicalSessionIdForTest();
+        const auto txnNum = 1LL;
+        OperationSessionInfo sessionInfo;
+        sessionInfo.setSessionId(sessionId);
+        sessionInfo.setTxnNumber(txnNum);
+        const auto ts = Timestamp(2, 1);
+        const auto opTime = repl::OpTime(ts, 1);
+        const auto oplogEntryBson = makeOplogEntry(opTime,
+                                                   repl::OpTypeEnum::kUpdate,
+                                                   NamespaceString("test.foo"),
+                                                   BSON("a" << 1),
+                                                   sessionInfo,
+                                                   boost::none /* preImageOpTime */,
+                                                   boost::none /* postImageOpTime */,
+                                                   imageType)
+                                        .getEntry()
+                                        .toBSON();
+        auto mock = DocumentSourceMock::createForTest(Document(oplogEntryBson), getExpCtx());
+        imageLookup->setSource(mock.get());
+
+        const auto prePostImage = BSON("a" << 2);
+        repl::ImageEntry imageEntry;
+        imageEntry.set_id(sessionId);
+        imageEntry.setTxnNumber(txnNum);
+        imageEntry.setTs(ts);
+        imageEntry.setImageKind(imageType);
+        imageEntry.setImage(prePostImage);
+        // Mock out the foreign collection.
+        getExpCtx()->mongoProcessInterface = std::make_unique<MockMongoInterface>(
+            std::vector<Document>{Document{imageEntry.toBSON()}});
+
+        // The forged image oplog entry should be returned before the findAndModify oplog entry.
+        auto next = imageLookup->getNext();
+        ASSERT_TRUE(next.isAdvanced());
+        const auto forgedImageEntry =
+            repl::OplogEntry::parse(next.releaseDocument().toBson()).getValue();
+        ASSERT_BSONOBJ_EQ(prePostImage, forgedImageEntry.getObject());
+        ASSERT_EQUALS(txnNum, forgedImageEntry.getTxnNumber().get());
+        ASSERT_EQUALS(sessionId, forgedImageEntry.getSessionId().get());
+        ASSERT_EQUALS("n", repl::OpType_serializer(forgedImageEntry.getOpType()));
+        ASSERT_EQUALS(0LL, forgedImageEntry.getStatementIds().front());
+        ASSERT_EQUALS(ts, forgedImageEntry.getTimestamp());
+        ASSERT_EQUALS(1, forgedImageEntry.getTerm().get());
+
+        // The next doc should be the original findAndModify oplog entry with the 'needsRetryImage'
+        // field removed and 'preImageOpTime'/'postImageOpTime' field appended.
+        next = imageLookup->getNext();
+        MutableDocument expectedDownConvertedDoc{Document{oplogEntryBson}};
+        expectedDownConvertedDoc.remove(repl::OplogEntryBase::kNeedsRetryImageFieldName);
+        const auto expectedImageOpTimeFieldName = imageType == repl::RetryImageEnum::kPreImage
+            ? repl::OplogEntry::kPreImageOpTimeFieldName
+            : repl::OplogEntry::kPostImageOpTimeFieldName;
+        expectedDownConvertedDoc.setField(
+            expectedImageOpTimeFieldName,
+            Value{Document{{repl::OpTime::kTimestampFieldName.toString(), opTime.getTimestamp()},
+                           {repl::OpTime::kTermFieldName.toString(), opTime.getTerm()}}});
+        ASSERT_DOCUMENT_EQ(next.releaseDocument(), expectedDownConvertedDoc.freeze());
+
+        ASSERT_TRUE(imageLookup->getNext().isEOF());
+        ASSERT_TRUE(imageLookup->getNext().isEOF());
+        ASSERT_TRUE(imageLookup->getNext().isEOF());
+    }
+}
+}  // namespace
+}  // namespace mongo