SERVER-56630: Unittest OpObserverImpl::onUpdate for non-transaction calls.

4 files changed, 298 insertions, 7 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 799e6670b72..35d82f2606c 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -2429,6 +2429,7 @@ if wiredtiger:
             'range_arithmetic',
             'read_write_concern_defaults_mock',
             'record_id_helpers',
+            'repl/image_collection_entry',
             'repl/mock_repl_coord_server_fixture',
             'repl/oplog_interface_local',
             'repl/repl_coordinator_interface',
diff --git a/src/mongo/db/dbhelpers.cpp b/src/mongo/db/dbhelpers.cpp
index 0a1a5e44fd2..7090e277550 100644
--- a/src/mongo/db/dbhelpers.cpp
+++ b/src/mongo/db/dbhelpers.cpp
@@ -71,6 +71,21 @@ bool Helpers::findOne(OperationContext* opCtx,
     return true;
 }
 
+BSONObj Helpers::findOneForTesting(OperationContext* opCtx,
+                                   const CollectionPtr& collection,
+                                   const BSONObj& query,
+                                   const bool invariantOnError) {
+    BSONObj ret;
+    const bool requiresIndex = true;
+    bool found = findOne(opCtx, collection, query, ret, requiresIndex);
+    if (invariantOnError) {
+        invariant(found);
+    }
+
+    return ret.getOwned();
+}
+
+
 /* fetch a single object from collection ns that matches query
    set your db SavedContext first
 */
diff --git a/src/mongo/db/dbhelpers.h b/src/mongo/db/dbhelpers.h
index 0d624e5769a..f1caed50dfe 100644
--- a/src/mongo/db/dbhelpers.h
+++ b/src/mongo/db/dbhelpers.h
@@ -66,6 +66,15 @@ struct Helpers {
                         bool requireIndex = false);
 
     /**
+     * If `invariantOnError` is true, an error (e.g: no document found) will crash the
+     * process. Otherwise the empty BSONObj will be returned.
+     */
+    static BSONObj findOneForTesting(OperationContext* opCtx,
+                                     const CollectionPtr& collection,
+                                     const BSONObj& query,
+                                     const bool invariantOnError = true);
+
+    /**
      * Similar to the 'findOne()' overload above, except returns the RecordId of the first matching
      * document, or a null RecordId if no such document exists.
      */
diff --git a/src/mongo/db/op_observer_impl_test.cpp b/src/mongo/db/op_observer_impl_test.cpp
index a26bdebe285..3a2355251ba 100644
--- a/src/mongo/db/op_observer_impl_test.cpp
+++ b/src/mongo/db/op_observer_impl_test.cpp
@@ -27,23 +27,30 @@
  *    it in the license file.
  */
 
+#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kTest
+
 #include "mongo/platform/basic.h"
 
 #include "mongo/db/catalog/import_collection_oplog_entry_gen.h"
 #include "mongo/db/client.h"
 #include "mongo/db/concurrency/locker_noop.h"
+#include "mongo/db/concurrency/write_conflict_exception.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/db/dbdirectclient.h"
+#include "mongo/db/dbhelpers.h"
 #include "mongo/db/keys_collection_client_sharded.h"
 #include "mongo/db/keys_collection_manager.h"
 #include "mongo/db/logical_time_validator.h"
 #include "mongo/db/op_observer_impl.h"
+#include "mongo/db/op_observer_registry.h"
 #include "mongo/db/read_write_concern_defaults.h"
 #include "mongo/db/read_write_concern_defaults_cache_lookup_mock.h"
+#include "mongo/db/repl/image_collection_entry_gen.h"
 #include "mongo/db/repl/oplog.h"
 #include "mongo/db/repl/oplog_entry.h"
 #include "mongo/db/repl/oplog_interface_local.h"
 #include "mongo/db/repl/repl_client_info.h"
+#include "mongo/db/repl/repl_server_parameters_gen.h"
 #include "mongo/db/repl/replication_coordinator_mock.h"
 #include "mongo/db/repl/storage_interface_impl.h"
 #include "mongo/db/service_context_d_test_fixture.h"
@@ -51,6 +58,7 @@
 #include "mongo/db/transaction_participant.h"
 #include "mongo/db/transaction_participant_gen.h"
 #include "mongo/idl/server_parameter_test_util.h"
+#include "mongo/logv2/log.h"
 #include "mongo/rpc/get_status_from_command_result.h"
 #include "mongo/unittest/death_test.h"
 #include "mongo/util/clock_source_mock.h"
@@ -85,18 +93,74 @@ public:
         ReadWriteConcernDefaults::create(getServiceContext(), _lookupMock.getFetchDefaultsFn());
     }
 
+    void resetOplogAndTransactions(OperationContext* opCtx) const {
+        ::mongo::writeConflictRetry(
+            opCtx, "deleteAll", NamespaceString::kRsOplogNamespace.ns(), [&] {
+                opCtx->recoveryUnit()->setTimestampReadSource(
+                    RecoveryUnit::ReadSource::kNoTimestamp);
+                opCtx->recoveryUnit()->abandonSnapshot();
+
+                WriteUnitOfWork wunit(opCtx);
+                {
+                    AutoGetCollection collRaii(
+                        opCtx, NamespaceString::kRsOplogNamespace, LockMode::MODE_X);
+                    invariant(collRaii.getWritableCollection()->truncate(opCtx).isOK());
+                }
+                {
+                    AutoGetCollection collRaii(opCtx,
+                                               NamespaceString::kSessionTransactionsTableNamespace,
+                                               LockMode::MODE_X);
+                    if (collRaii) {
+                        invariant(collRaii.getWritableCollection()->truncate(opCtx).isOK());
+                    } else {
+                        auto db = collRaii.ensureDbExists();
+                        invariant(db->createCollection(
+                            opCtx, NamespaceString::kSessionTransactionsTableNamespace));
+                    }
+                }
+                {
+                    AutoGetCollection collRaii(
+                        opCtx, NamespaceString::kConfigImagesNamespace, LockMode::MODE_X);
+                    if (collRaii) {
+                        invariant(collRaii.getWritableCollection()->truncate(opCtx).isOK());
+                    } else {
+                        auto db = collRaii.ensureDbExists();
+                        invariant(
+                            db->createCollection(opCtx, NamespaceString::kConfigImagesNamespace));
+                    }
+                }
+                wunit.commit();
+            });
+    }
+
 protected:
     // Assert that the oplog has the expected number of entries, and return them
-    std::vector<BSONObj> getNOplogEntries(OperationContext* opCtx, int n) {
-        std::vector<BSONObj> result(n);
+    std::vector<BSONObj> getNOplogEntries(OperationContext* opCtx, int numExpected) {
+        std::vector<BSONObj> allOplogEntries;
         repl::OplogInterfaceLocal oplogInterface(opCtx);
         auto oplogIter = oplogInterface.makeIterator();
-        for (int i = n - 1; i >= 0; i--) {
+        while (true) {
+            StatusWith<std::pair<BSONObj, RecordId>> swEntry = oplogIter->next();
+            if (swEntry.getStatus() == ErrorCodes::CollectionIsEmpty) {
+                break;
+            }
+            allOplogEntries.push_back(swEntry.getValue().first);
+        }
+        if (allOplogEntries.size() != static_cast<std::size_t>(numExpected)) {
+            LOGV2(5739903,
+                  "Incorrect number of oplog entries made",
+                  "numExpected"_attr = numExpected,
+                  "numFound"_attr = allOplogEntries.size(),
+                  "entries"_attr = allOplogEntries);
+        }
+        ASSERT_EQUALS(allOplogEntries.size(), numExpected);
+
+        std::vector<BSONObj> ret(numExpected);
+        for (int idx = numExpected - 1; idx >= 0; idx--) {
             // The oplogIterator returns the entries in reverse order.
-            auto opEntry = unittest::assertGet(oplogIter->next());
-            result[i] = opEntry.first;
+            ret[idx] = allOplogEntries[numExpected - idx - 1];
         }
-        ASSERT_EQUALS(ErrorCodes::CollectionIsEmpty, oplogIter->next().getStatus());
+
         // Some unittests reuse the same OperationContext to read the oplog and end up acquiring the
         // RSTL lock after using the OplogInterfaceLocal. This is a hack to make sure we do not hold
         // RSTL lock for prepared transactions.
@@ -104,7 +168,7 @@ protected:
             TransactionParticipant::get(opCtx).transactionIsPrepared()) {
             opCtx->lockState()->unlockRSTLforPrepare();
         }
-        return result;
+        return ret;
     }
 
     // Assert that oplog only has a single entry and return that oplog entry.
@@ -1485,6 +1549,208 @@ TEST_F(OpObserverRetryableFindAndModifyTest, RetryableFindAndModifyDeleteHasNeed
                   "preImage"_sd);
 }
 
+
+enum class RetryableOptions { NotRetryable, WithOplog, WithSideCollection };
+using StoreDocOption = CollectionUpdateArgs::StoreDocOption;
+struct TestCase {
+    StoreDocOption imageType;
+    bool alwaysRecordPreImages;
+    RetryableOptions retryableOptions;
+
+    int numOutputOplogs;
+
+    std::string getImageTypeStr() const {
+        switch (imageType) {
+            case StoreDocOption::None:
+                return "None";
+            case StoreDocOption::PreImage:
+                return "PreImage";
+            case StoreDocOption::PostImage:
+                return "PostImage";
+        }
+        MONGO_UNREACHABLE;
+    }
+
+    std::string getRetryableOptionsStr() const {
+        switch (retryableOptions) {
+            case RetryableOptions::NotRetryable:
+                return "Not retryable";
+            case RetryableOptions::WithOplog:
+                return "Images in oplog";
+            case RetryableOptions::WithSideCollection:
+                return "Images in side collection";
+        }
+        MONGO_UNREACHABLE;
+    }
+};
+
+OplogEntry findByTimestamp(const std::vector<BSONObj>& oplogs, Timestamp ts) {
+    for (auto& oplog : oplogs) {
+        const auto& entry = assertGet(OplogEntry::parse(oplog));
+        if (entry.getTimestamp() == ts) {
+            return entry;
+        }
+    }
+
+    FAIL("Not found.");
+    // C++/clang isn't smart enough to know FAIL is guaranteed to throw.
+    MONGO_UNREACHABLE
+}
+
+TEST_F(OpObserverTest, TestFundamentalOnUpdateOutputs) {
+    // Create a registry that only registers the Impl. It can be challenging to call methods on the
+    // Impl directly. It falls into cases where `ReservedTimes` is expected to be instantiated. Due
+    // to strong encapsulation, we use the registry that managers the `ReservedTimes` on our behalf.
+    OpObserverRegistry opObserver;
+    opObserver.addObserver(std::make_unique<OpObserverImpl>());
+
+    auto opCtxRaii = cc().makeOperationContext();
+    OperationContext* opCtx = opCtxRaii.get();
+    NamespaceString nss("test", "coll");
+    CollectionUUID uuid = CollectionUUID::gen();
+
+    RAIIServerParameterControllerForTest ffRaii("featureFlagRetryableFindAndModify", true);
+
+    const bool kRecordPreImages = true;
+    const bool kDoNotRecordPreImages = false;
+
+    std::vector<TestCase> cases{
+        // Regular updates.
+        {StoreDocOption::None, kDoNotRecordPreImages, RetryableOptions::NotRetryable, 1},
+        {StoreDocOption::None, kRecordPreImages, RetryableOptions::NotRetryable, 2},
+        {StoreDocOption::None, kRecordPreImages, RetryableOptions::WithOplog, 2},
+        {StoreDocOption::None, kRecordPreImages, RetryableOptions::WithSideCollection, 2},
+        // FindAndModify asking for a preImage.
+        {StoreDocOption::PreImage, kDoNotRecordPreImages, RetryableOptions::NotRetryable, 1},
+        {StoreDocOption::PreImage, kDoNotRecordPreImages, RetryableOptions::WithOplog, 2},
+        {StoreDocOption::PreImage, kDoNotRecordPreImages, RetryableOptions::WithSideCollection, 1},
+        {StoreDocOption::PreImage, kRecordPreImages, RetryableOptions::NotRetryable, 2},
+        {StoreDocOption::PreImage, kRecordPreImages, RetryableOptions::WithOplog, 2},
+        {StoreDocOption::PreImage, kRecordPreImages, RetryableOptions::WithSideCollection, 2},
+        // FindAndModify asking for a postImage.
+        {StoreDocOption::PostImage, kDoNotRecordPreImages, RetryableOptions::NotRetryable, 1},
+        {StoreDocOption::PostImage, kDoNotRecordPreImages, RetryableOptions::WithOplog, 2},
+        {StoreDocOption::PostImage, kDoNotRecordPreImages, RetryableOptions::WithSideCollection, 1},
+        {StoreDocOption::PostImage, kRecordPreImages, RetryableOptions::NotRetryable, 2},
+        {StoreDocOption::PostImage, kRecordPreImages, RetryableOptions::WithOplog, 3},
+        {StoreDocOption::PostImage, kRecordPreImages, RetryableOptions::WithSideCollection, 2}};
+
+    for (std::size_t idx = 0; idx < cases.size(); ++idx) {
+        const auto& testCase = cases[idx];
+        LOGV2(5739902,
+              "TestCase",
+              "ImageType"_attr = testCase.getImageTypeStr(),
+              "AlwaysRecordPreImages"_attr = testCase.alwaysRecordPreImages,
+              "RetryableOptions"_attr = testCase.getRetryableOptionsStr(),
+              "ExpectedOplogEntries"_attr = testCase.numOutputOplogs);
+
+        // Phase 1: Clearing any state and setting up fixtures/the update call.
+        resetOplogAndTransactions(opCtx);
+
+        boost::optional<MongoDOperationContextSession> contextSession;
+        boost::optional<TransactionParticipant::Participant> txnParticipant;
+        CollectionUpdateArgs updateArgs;
+        switch (testCase.retryableOptions) {
+            case RetryableOptions::NotRetryable:
+                updateArgs.stmtIds = {kUninitializedStmtId};
+                break;
+            case RetryableOptions::WithOplog:
+                repl::gStoreFindAndModifyImagesInSideCollection.store(false);
+                updateArgs.stmtIds = {1};
+                break;
+            case RetryableOptions::WithSideCollection:
+                repl::gStoreFindAndModifyImagesInSideCollection.store(true);
+                updateArgs.stmtIds = {1};
+                break;
+        }
+        if (testCase.retryableOptions != RetryableOptions::NotRetryable) {
+            opCtx->setLogicalSessionId(makeLogicalSessionIdForTest());
+            opCtx->setTxnNumber(TxnNumber(idx));
+            contextSession.emplace(opCtx);
+            txnParticipant.emplace(TransactionParticipant::get(opCtx));
+            txnParticipant->beginOrContinue(opCtx, TxnNumber(idx), boost::none, boost::none);
+        }
+
+        if (testCase.imageType == StoreDocOption::None && !testCase.alwaysRecordPreImages) {
+            updateArgs.preImageDoc = boost::none;
+        } else {
+            updateArgs.preImageDoc = BSON("_id" << 0 << "preImage" << true);
+        }
+
+        updateArgs.updatedDoc = BSON("_id" << 0 << "postImage" << true);
+        updateArgs.update =
+            BSON("$set" << BSON("postImage" << true) << "$unset" << BSON("preImage" << 1));
+        updateArgs.criteria = BSON("_id" << 0);
+        updateArgs.storeDocOption = testCase.imageType;
+        updateArgs.preImageRecordingEnabledForCollection = testCase.alwaysRecordPreImages;
+        OplogUpdateEntryArgs update(std::move(updateArgs), nss, uuid);
+
+        // Phase 2: Call the code we're testing.
+        WriteUnitOfWork wuow(opCtx);
+        AutoGetCollection locks(opCtx, nss, LockMode::MODE_IX);
+        opObserver.onUpdate(opCtx, update);
+        wuow.commit();
+
+        // Phase 3: Analyze the results:
+
+        // This `getNOplogEntries` also asserts that all oplogs are retrieved.
+        std::vector<BSONObj> oplogs = getNOplogEntries(opCtx, testCase.numOutputOplogs);
+        // Entries are returned in ascending timestamp order.
+        const OplogEntry& actualOp = assertGet(OplogEntry::parse(oplogs.back()));
+
+        const bool checkPreImageInOplog = testCase.alwaysRecordPreImages ||
+            (testCase.imageType == StoreDocOption::PreImage &&
+             testCase.retryableOptions == RetryableOptions::WithOplog);
+        if (checkPreImageInOplog) {
+            ASSERT(actualOp.getPreImageOpTime());
+            const Timestamp preImageOpTime = actualOp.getPreImageOpTime()->getTimestamp();
+            ASSERT_FALSE(preImageOpTime.isNull());
+            OplogEntry preImage = findByTimestamp(oplogs, preImageOpTime);
+            ASSERT_EQUALS(0, update.updateArgs.preImageDoc->woCompare(preImage.getObject()));
+        }
+
+        const bool checkPostImageInOplog = testCase.imageType == StoreDocOption::PostImage &&
+            testCase.retryableOptions == RetryableOptions::WithOplog;
+        if (checkPostImageInOplog) {
+            ASSERT(actualOp.getPostImageOpTime());
+            const Timestamp postImageOpTime = actualOp.getPostImageOpTime()->getTimestamp();
+            ASSERT_FALSE(postImageOpTime.isNull());
+            OplogEntry postImage = findByTimestamp(oplogs, postImageOpTime);
+            ASSERT_EQUALS(0, update.updateArgs.updatedDoc.woCompare(postImage.getObject()));
+        }
+
+        bool checkSideCollection = testCase.imageType != StoreDocOption::None &&
+            testCase.retryableOptions == RetryableOptions::WithSideCollection;
+        if (checkSideCollection && testCase.alwaysRecordPreImages &&
+            testCase.imageType == StoreDocOption::PreImage) {
+            // When `alwaysRecordPreImages` is enabled for a collection, we always store an image in
+            // the oplog. To avoid unnecessary writes, we won't also store an image in the side
+            // collection.
+            checkSideCollection = false;
+        }
+
+        if (checkSideCollection) {
+            AutoGetCollection sideCollection(
+                opCtx, NamespaceString::kConfigImagesNamespace, LockMode::MODE_IS);
+            auto imageEntry = repl::ImageEntry::parse(
+                IDLParserErrorContext("image entry"),
+                Helpers::findOneForTesting(opCtx,
+                                           sideCollection.getCollection(),
+                                           BSON("_id" << actualOp.getSessionId()->toBSON())));
+            const BSONObj& expectedImage = testCase.imageType == StoreDocOption::PreImage
+                ? update.updateArgs.preImageDoc.get()
+                : update.updateArgs.updatedDoc;
+            ASSERT_EQUALS(0, expectedImage.woCompare(imageEntry.getImage()));
+            if (testCase.imageType == StoreDocOption::PreImage) {
+                ASSERT(imageEntry.getImageKind() == repl::RetryImageEnum::kPreImage);
+            } else {
+                ASSERT(imageEntry.getImageKind() == repl::RetryImageEnum::kPostImage);
+            }
+        }
+    }
+}
+
+
 TEST_F(OpObserverMultiEntryTransactionTest, TransactionalInsertTest) {
     const NamespaceString nss1("testDB", "testColl");
     const NamespaceString nss2("testDB2", "testColl2");