SERVER-56630 Unittest OpObserverImpl and chunk migration for retryable findAndModify

(cherry picked from commit d78928f90492739ce5401d112b0fe89e278d3954)
(cherry picked from commit 68c6fbfe69283f70ff321cfaa7cb2a533747ac36)

1 files changed, 89 insertions, 15 deletions

diff --git a/src/mongo/db/s/session_catalog_migration_source_test.cpp b/src/mongo/db/s/session_catalog_migration_source_test.cpp
index 32e534ffeda..2de65446e15 100644
--- a/src/mongo/db/s/session_catalog_migration_source_test.cpp
+++ b/src/mongo/db/s/session_catalog_migration_source_test.cpp
@@ -31,6 +31,7 @@
 
 #include "mongo/db/dbdirectclient.h"
 #include "mongo/db/logical_session_id.h"
+#include "mongo/db/repl/image_collection_entry_gen.h"
 #include "mongo/db/repl/mock_repl_coord_server_fixture.h"
 #include "mongo/db/repl/replication_process.h"
 #include "mongo/db/s/session_catalog_migration_source.h"
@@ -63,8 +64,10 @@ repl::OplogEntry makeOplogEntry(repl::OpTime opTime,
                                 const std::vector<StmtId>& stmtIds,
                                 repl::OpTime prevWriteOpTimeInTransaction,
                                 boost::optional<repl::OpTime> preImageOpTime,
-                                boost::optional<repl::OpTime> postImageOpTime) {
-    return {repl::DurableOplogEntry(
+                                boost::optional<repl::OpTime> postImageOpTime,
+                                boost::optional<OperationSessionInfo> osi,
+                                boost::optional<repl::RetryImageEnum> needsRetryImage) {
+    return repl::DurableOplogEntry(
         opTime,                           // optime
         0,                                // hash
         opType,                           // opType
@@ -74,7 +77,7 @@ repl::OplogEntry makeOplogEntry(repl::OpTime opTime,
         repl::OplogEntry::kOplogVersion,  // version
         object,                           // o
         object2,                          // o2
-        {},                               // sessionInfo
+        osi.get_value_or({}),             // sessionInfo
         boost::none,                      // upsert
         wallClockTime,                    // wall clock time
         stmtIds,                          // statement ids
@@ -82,18 +85,22 @@ repl::OplogEntry makeOplogEntry(repl::OpTime opTime,
         preImageOpTime,                   // pre-image optime
         postImageOpTime,                  // post-image optime
         boost::none,                      // ShardId of resharding recipient
-        boost::none)};                    // _id
+        boost::none,                      // _id
+        needsRetryImage);
 }
 
-repl::OplogEntry makeOplogEntry(repl::OpTime opTime,
-                                repl::OpTypeEnum opType,
-                                BSONObj object,
-                                boost::optional<BSONObj> object2,
-                                Date_t wallClockTime,
-                                const std::vector<StmtId>& stmtIds,
-                                repl::OpTime prevWriteOpTimeInTransaction,
-                                boost::optional<repl::OpTime> preImageOpTime = boost::none,
-                                boost::optional<repl::OpTime> postImageOpTime = boost::none) {
+repl::OplogEntry makeOplogEntry(
+    repl::OpTime opTime,
+    repl::OpTypeEnum opType,
+    BSONObj object,
+    boost::optional<BSONObj> object2,
+    Date_t wallClockTime,
+    const std::vector<StmtId>& stmtIds,
+    repl::OpTime prevWriteOpTimeInTransaction,
+    boost::optional<repl::OpTime> preImageOpTime = boost::none,
+    boost::optional<repl::OpTime> postImageOpTime = boost::none,
+    boost::optional<OperationSessionInfo> osi = boost::none,
+    boost::optional<repl::RetryImageEnum> needsRetryImage = boost::none) {
     return makeOplogEntry(opTime,
                           opType,
                           kNs,
@@ -103,7 +110,9 @@ repl::OplogEntry makeOplogEntry(repl::OpTime opTime,
                           stmtIds,
                           prevWriteOpTimeInTransaction,
                           preImageOpTime,
-                          postImageOpTime);
+                          postImageOpTime,
+                          osi,
+                          needsRetryImage);
 }
 
 TEST_F(SessionCatalogMigrationSourceTest, NoSessionsToTransferShouldNotHaveOplog) {
@@ -467,6 +476,69 @@ TEST_F(SessionCatalogMigrationSourceTest,
     ASSERT_FALSE(migrationSource.hasMoreOplog());
 }
 
+TEST_F(SessionCatalogMigrationSourceTest, ForgeImageEntriesWhenFetchingEntriesWithNeedsRetryImage) {
+    repl::ImageEntry imageEntry;
+    const auto preImage = BSON("_id" << 1 << "x" << 50);
+    const auto lsid = makeLogicalSessionIdForTest();
+    const repl::OpTime imageEntryOpTime = repl::OpTime(Timestamp(52, 346), 2);
+    const auto txnNumber = 1LL;
+    imageEntry.set_id(lsid);
+    imageEntry.setTxnNumber(txnNumber);
+    imageEntry.setTs(imageEntryOpTime.getTimestamp());
+    imageEntry.setImageKind(repl::RetryImageEnum::kPreImage);
+    imageEntry.setImage(preImage);
+
+    OperationSessionInfo osi;
+    osi.setTxnNumber(txnNumber);
+    osi.setSessionId(lsid);
+
+    DBDirectClient client(opCtx());
+    client.insert(NamespaceString::kConfigImagesNamespace.ns(), imageEntry.toBSON());
+
+    // Insert an oplog entry with a non-null needsRetryImage field.
+    auto entry = makeOplogEntry(
+        repl::OpTime(Timestamp(52, 346), 2),  // optime
+        repl::OpTypeEnum::kDelete,            // op type
+        BSON("x" << 50),                      // o
+        boost::none,                          // o2
+        Date_t::now(),                        // wall clock time
+        {1},                                  // statement id
+        repl::OpTime(Timestamp(0, 0), 0),     // optime of previous write within same transaction
+        boost::none,                          // pre-image optime
+        boost::none,                          // post-image optime
+        osi,                                  // operation session info
+        repl::RetryImageEnum::kPreImage);     // needsRetryImage
+    insertOplogEntry(entry);
+
+    SessionTxnRecord sessionRecord;
+    sessionRecord.setSessionId(lsid);
+    sessionRecord.setTxnNum(1);
+    sessionRecord.setLastWriteOpTime(entry.getOpTime());
+    sessionRecord.setLastWriteDate(entry.getWallClockTime());
+
+    client.insert(NamespaceString::kSessionTransactionsTableNamespace.ns(), sessionRecord.toBSON());
+
+    SessionCatalogMigrationSource migrationSource(opCtx(), kNs, kChunkRange, kShardKey);
+    // The next oplog entry should be the forged preImage entry.
+    ASSERT_TRUE(migrationSource.fetchNextOplog(opCtx()));
+    ASSERT_TRUE(migrationSource.hasMoreOplog());
+    auto nextOplogResult = migrationSource.getLastFetchedOplog();
+    ASSERT_FALSE(nextOplogResult.shouldWaitForMajority);
+    // Check that the key fields are what we expect. The destination will overwrite any unneeded
+    // fields when it processes the incoming entries.
+    ASSERT_BSONOBJ_EQ(preImage, nextOplogResult.oplog->getObject());
+    ASSERT_EQUALS(txnNumber, nextOplogResult.oplog->getTxnNumber().get());
+    ASSERT_EQUALS(lsid, nextOplogResult.oplog->getSessionId().get());
+    ASSERT_EQUALS("n", repl::OpType_serializer(nextOplogResult.oplog->getOpType()));
+    ASSERT_EQUALS(0LL, nextOplogResult.oplog->getStatementIds().front());
+
+    // The next oplog entry should be the original entry that generated the image entry.
+    ASSERT_TRUE(migrationSource.hasMoreOplog());
+    ASSERT_TRUE(migrationSource.fetchNextOplog(opCtx()));
+    nextOplogResult = migrationSource.getLastFetchedOplog();
+    ASSERT_BSONOBJ_EQ(entry.getEntry().toBSON(), nextOplogResult.oplog->getEntry().toBSON());
+}
+
 TEST_F(SessionCatalogMigrationSourceTest, OplogWithOtherNsShouldBeIgnored) {
     auto entry1 = makeOplogEntry(
         repl::OpTime(Timestamp(52, 345), 2),  // optime
@@ -499,7 +571,9 @@ TEST_F(SessionCatalogMigrationSourceTest, OplogWithOtherNsShouldBeIgnored) {
         {1},                                 // statement ids
         repl::OpTime(Timestamp(0, 0), 0),    // optime of previous write within same transaction
         boost::none,                         // pre-image optime
-        boost::none);                        // post-image optime
+        boost::none,                         // post-image optime
+        boost::none,                         // operation session info
+        boost::none);                        // needsRetryImage
     insertOplogEntry(entry2);
 
     SessionTxnRecord sessionRecord2;