SERVER-56374 Add ability to write retryable findAndModify updates to `config.image_collection`

12 files changed, 176 insertions, 9 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 2a2ed84c589..adcc2cfae90 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -663,6 +663,7 @@ env.Library(
     ],
     LIBDEPS_PRIVATE=[
         'commands/server_status',
+        'op_observer',
         'update/update_driver',
     ]
 )
@@ -760,6 +761,7 @@ env.Library(
     ],
     LIBDEPS=[
         '$BUILD_DIR/mongo/base',
+        'server_parameters',
     ],
 )
 
@@ -781,6 +783,7 @@ env.Library(
     ],
     LIBDEPS_PRIVATE=[
         '$BUILD_DIR/mongo/db/commands/mongod_fcv',
+        '$BUILD_DIR/mongo/db/repl/image_collection_entry',
     ],
 )
 
diff --git a/src/mongo/db/namespace_string.cpp b/src/mongo/db/namespace_string.cpp
index 85b28afd6e8..06dc0ff5c5c 100644
--- a/src/mongo/db/namespace_string.cpp
+++ b/src/mongo/db/namespace_string.cpp
@@ -68,6 +68,9 @@ const NamespaceString NamespaceString::kSystemKeysNamespace(NamespaceString::kAd
                                                             "system.keys");
 const NamespaceString NamespaceString::kRsOplogNamespace(NamespaceString::kLocalDb, "oplog.rs");
 
+const NamespaceString NamespaceString::kConfigImagesNamespace(NamespaceString::kConfigDb,
+                                                              "image_collection");
+
 bool NamespaceString::isListCollectionsCursorNS() const {
     return coll() == listCollectionsCursorCol;
 }
diff --git a/src/mongo/db/namespace_string.h b/src/mongo/db/namespace_string.h
index 85462caf35a..d58e779af5a 100644
--- a/src/mongo/db/namespace_string.h
+++ b/src/mongo/db/namespace_string.h
@@ -93,6 +93,9 @@ public:
     // Namespace of the the oplog collection.
     static const NamespaceString kRsOplogNamespace;
 
+    // Namespace used for storing retryable findAndModify images.
+    static const NamespaceString kConfigImagesNamespace;
+
     /**
      * Constructs an empty NamespaceString.
      */
diff --git a/src/mongo/db/op_observer.cpp b/src/mongo/db/op_observer.cpp
index d390bc639f9..22cc46cab8c 100644
--- a/src/mongo/db/op_observer.cpp
+++ b/src/mongo/db/op_observer.cpp
@@ -33,8 +33,11 @@
 #include "mongo/db/op_observer.h"
 
 #include "mongo/db/operation_context.h"
+#include "mongo/db/server_parameters.h"
 
 namespace mongo {
+
+MONGO_EXPORT_SERVER_PARAMETER(storeFindAndModifyImagesInSideCollection, bool, false);
 namespace {
 const auto getOpObserverTimes = OperationContext::declareDecoration<OpObserver::Times>();
 }  // namespace
diff --git a/src/mongo/db/op_observer.h b/src/mongo/db/op_observer.h
index e2eaefd4d79..3e7094ab97a 100644
--- a/src/mongo/db/op_observer.h
+++ b/src/mongo/db/op_observer.h
@@ -42,6 +42,8 @@
 
 namespace mongo {
 
+extern AtomicBool storeFindAndModifyImagesInSideCollection;
+
 struct InsertStatement;
 class OperationContext;
 struct OplogSlot;
diff --git a/src/mongo/db/op_observer_impl.cpp b/src/mongo/db/op_observer_impl.cpp
index 171427154d1..dfee2e4f4e5 100644
--- a/src/mongo/db/op_observer_impl.cpp
+++ b/src/mongo/db/op_observer_impl.cpp
@@ -42,15 +42,17 @@
 #include "mongo/db/commands/feature_compatibility_version.h"
 #include "mongo/db/commands/feature_compatibility_version_parser.h"
 #include "mongo/db/concurrency/d_concurrency.h"
+#include "mongo/db/db_raii.h"
+#include "mongo/db/dbhelpers.h"
 #include "mongo/db/index/index_descriptor.h"
 #include "mongo/db/logical_time_validator.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/operation_context.h"
+#include "mongo/db/repl/image_collection_entry_gen.h"
 #include "mongo/db/repl/oplog.h"
 #include "mongo/db/repl/replication_coordinator.h"
 #include "mongo/db/s/collection_sharding_state.h"
 #include "mongo/db/server_options.h"
-#include "mongo/db/server_parameters.h"
 #include "mongo/db/session_catalog.h"
 #include "mongo/db/views/durable_view_catalog.h"
 #include "mongo/s/client/shard_registry.h"
@@ -64,7 +66,6 @@ using repl::OplogEntry;
 namespace {
 
 MONGO_FAIL_POINT_DEFINE(failCollectionUpdates);
-MONGO_EXPORT_SERVER_PARAMETER(storeFindAndModifyImagesInSideCollection, bool, false);
 
 const auto documentKeyDecoration = OperationContext::declareDecoration<BSONObj>();
 
@@ -166,17 +167,31 @@ struct OpTimeBundle {
     Date_t wallClockTime;
 };
 
+void writeToImagesCollection(OperationContext* opCtx,
+                             BSONObj image,
+                             repl::RetryImageEnum imageKind,
+                             Timestamp ts) {
+    repl::ImageEntry imageEntry;
+    invariant(opCtx->getLogicalSessionId());
+    imageEntry.set_id(*opCtx->getLogicalSessionId());
+    imageEntry.setTs(ts);
+    imageEntry.setImage(std::move(image));
+    imageEntry.setImageKind(imageKind);
+    repl::UnreplicatedWritesBlock unreplicated(opCtx);
+    AutoGetCollection imageCollectionRaii(
+        opCtx, NamespaceString::kConfigImagesNamespace, LockMode::MODE_IX);
+    Helpers::upsert(opCtx, NamespaceString::kConfigImagesNamespace.toString(), imageEntry.toBSON());
+}
+
 /**
  * Write oplog entry(ies) for the update operation.
  */
 OpTimeBundle replLogUpdate(OperationContext* opCtx,
                            Session* session,
-                           const OplogUpdateEntryArgs& args) {
+                           const OplogUpdateEntryArgs& args,
+                           const bool storeImagesInSideCollection) {
     BSONObj storeObj;
     boost::optional<repl::RetryImageEnum> needsRetryImage;
-    const auto storeImagesInSideCollection = storeFindAndModifyImagesInSideCollection.load() &&
-        serverGlobalParams.featureCompatibility.getVersion() >=
-            ServerGlobalParams::FeatureCompatibility::Version::kUpgradingTo40;
     if (args.storeDocOption == OplogUpdateEntryArgs::StoreDocOption::PreImage) {
         invariant(args.preImageDoc);
         storeObj = *args.preImageDoc;
@@ -493,7 +508,26 @@ void OpObserverImpl::onUpdate(OperationContext* opCtx, const OplogUpdateEntryArg
             OplogEntry::makeUpdateOperation(args.nss, args.uuid, args.update, args.criteria);
         session->addTransactionOperation(opCtx, operation);
     } else {
-        opTime = replLogUpdate(opCtx, session, args);
+        const auto storeImagesInSideCollection = storeFindAndModifyImagesInSideCollection.load() &&
+            serverGlobalParams.featureCompatibility.isVersionInitialized() &&
+            serverGlobalParams.featureCompatibility.getVersion() >=
+                ServerGlobalParams::FeatureCompatibility::Version::kUpgradingTo40;
+        opTime = replLogUpdate(opCtx, session, args, storeImagesInSideCollection);
+        if (storeImagesInSideCollection && opCtx->getTxnNumber() &&
+            args.storeDocOption != OplogUpdateEntryArgs::StoreDocOption::None) {
+            BSONObj imageDoc;
+            repl::RetryImageEnum imageKind;
+            if (args.storeDocOption == OplogUpdateEntryArgs::StoreDocOption::PreImage) {
+                invariant(args.preImageDoc);
+                imageDoc = *args.preImageDoc;
+                imageKind = repl::RetryImageEnum::kPreImage;
+            } else {
+                invariant(args.storeDocOption == OplogUpdateEntryArgs::StoreDocOption::PostImage);
+                imageDoc = args.updatedDoc;
+                imageKind = repl::RetryImageEnum::kPostImage;
+            }
+            writeToImagesCollection(opCtx, imageDoc, imageKind, opTime.writeOpTime.getTimestamp());
+        }
         onWriteOpCompleted(opCtx,
                            args.nss,
                            session,
diff --git a/src/mongo/db/ops/update.cpp b/src/mongo/db/ops/update.cpp
index 349762b9382..53653c8f939 100644
--- a/src/mongo/db/ops/update.cpp
+++ b/src/mongo/db/ops/update.cpp
@@ -103,11 +103,22 @@ UpdateResult update(OperationContext* opCtx, Database* db, const UpdateRequest&
     OpDebug* const nullOpDebug = nullptr;
     auto exec = uassertStatusOK(getExecutorUpdate(opCtx, nullOpDebug, collection, &parsedUpdate));
 
+    BSONObj docImage;
+    PlanExecutor::ExecState state = PlanExecutor::ADVANCED;
+    if (request.shouldReturnAnyDocs()) {
+        state = exec->getNext(&docImage, nullptr);
+    }
+
+    while (state == PlanExecutor::ADVANCED) {
+        state = exec->getNext(nullptr, nullptr);
+    }
     uassertStatusOK(exec->executePlan());
 
     const UpdateStats* updateStats = UpdateStage::getUpdateStats(exec.get());
 
-    return UpdateStage::makeUpdateResult(updateStats);
+    UpdateResult result = UpdateStage::makeUpdateResult(updateStats);
+    result.requestedDocImage = docImage.getOwned();
+    return result;
 }
 
 BSONObj applyUpdateOperators(OperationContext* opCtx,
diff --git a/src/mongo/db/ops/update_result.h b/src/mongo/db/ops/update_result.h
index f91fe54e843..c52888ffa6e 100644
--- a/src/mongo/db/ops/update_result.h
+++ b/src/mongo/db/ops/update_result.h
@@ -57,6 +57,8 @@ struct UpdateResult {
 
     // if something was upserted, the new _id of the object
     BSONObj upserted;
+
+    BSONObj requestedDocImage;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/repl/SConscript b/src/mongo/db/repl/SConscript
index fdcd7645f47..2ecf6829075 100644
--- a/src/mongo/db/repl/SConscript
+++ b/src/mongo/db/repl/SConscript
@@ -14,6 +14,7 @@ env.Library(
     ],
     LIBDEPS_PRIVATE=[
         'dbcheck',
+        'image_collection_entry',
         'repl_coordinator_interface',
         'repl_settings',
         '$BUILD_DIR/mongo/base',
@@ -1713,3 +1714,16 @@ env.Library(
         'election_reason_counter',
     ],
 )
+
+env.Library(
+    target='image_collection_entry',
+    source=[
+        env.Idlc('image_collection_entry.idl')[0],
+    ],
+    LIBDEPS=[
+        '$BUILD_DIR/mongo/base',
+        '$BUILD_DIR/mongo/db/logical_session_id',
+        '$BUILD_DIR/mongo/idl/idl_parser',
+        'oplog_entry',
+    ],
+)
diff --git a/src/mongo/db/repl/image_collection_entry.idl b/src/mongo/db/repl/image_collection_entry.idl
new file mode 100644
index 00000000000..2cafe05ae8d
--- /dev/null
+++ b/src/mongo/db/repl/image_collection_entry.idl
@@ -0,0 +1,28 @@
+ # Oplog Entry IDL File
+
+global:
+    cpp_namespace: "mongo::repl"
+
+imports:
+    - "mongo/idl/basic_types.idl"
+    - "mongo/db/logical_session_id.idl"
+    - "mongo/db/repl/oplog_entry.idl"
+
+structs:
+    ImageEntry:
+        description: "Represents either a pre-image or post-image necessary for satisfying a
+                      retryable findAndModify."
+        strict: false
+        fields:
+            _id:
+                cpp_name: _id
+                type: LogicalSessionId
+            ts:
+                cpp_name: ts
+                type: timestamp
+            imageKind:
+                cpp_name: imageKind
+                type: RetryImage
+            image:
+                cpp_name: image
+                type: object
\ No newline at end of file
diff --git a/src/mongo/db/repl/oplog.cpp b/src/mongo/db/repl/oplog.cpp
index ce2cf2ed50b..0857070ccaf 100644
--- a/src/mongo/db/repl/oplog.cpp
+++ b/src/mongo/db/repl/oplog.cpp
@@ -77,6 +77,7 @@
 #include "mongo/db/repl/apply_ops.h"
 #include "mongo/db/repl/bgsync.h"
 #include "mongo/db/repl/dbcheck.h"
+#include "mongo/db/repl/image_collection_entry_gen.h"
 #include "mongo/db/repl/oplogreader.h"
 #include "mongo/db/repl/optime.h"
 #include "mongo/db/repl/repl_client_info.h"
@@ -312,6 +313,36 @@ void createIndexForApplyOps(OperationContext* opCtx,
     }
 }
 
+void writeToImageCollection(OperationContext* opCtx,
+                            const BSONObj& op,
+                            const BSONObj& image,
+                            repl::RetryImageEnum imageKind,
+                            bool* upsertConfigImage) {
+    AutoGetCollection autoColl(opCtx, NamespaceString::kConfigImagesNamespace, LockMode::MODE_IX);
+    repl::ImageEntry imageEntry;
+    LogicalSessionId sessionId =
+        LogicalSessionId::parse(IDLParserErrorContext("ParseSessionIdWhenWritingToImageCollection"),
+                                op.getField(OplogEntryBase::kSessionIdFieldName).Obj());
+    imageEntry.set_id(sessionId);
+    imageEntry.setTs(op["ts"].timestamp());
+    imageEntry.setImageKind(imageKind);
+    imageEntry.setImage(image);
+
+    UpdateRequest request(NamespaceString::kConfigImagesNamespace);
+    request.setQuery(
+        BSON("_id" << imageEntry.get_id().toBSON() << "ts" << BSON("$lt" << imageEntry.getTs())));
+    request.setUpsert(*upsertConfigImage);
+    request.setUpdates(imageEntry.toBSON());
+    request.setFromOplogApplication(true);
+    try {
+        ::mongo::update(opCtx, autoColl.getDb(), request);
+    } catch (const ExceptionFor<ErrorCodes::DuplicateKey>&) {
+        // We can get a duplicate key when two upserts race on inserting a document.
+        *upsertConfigImage = false;
+        throw WriteConflictException();
+    }
+}
+
 namespace {
 
 /**
@@ -1446,6 +1477,17 @@ Status applyOperation_inlock(OperationContext* opCtx,
         request.setUpdates(o);
         request.setUpsert(upsert);
         request.setFromOplogApplication(true);
+        boost::optional<repl::RetryImageEnum> imageKind;
+        if (op.hasField(OplogEntryBase::kNeedsRetryImageFieldName)) {
+            imageKind = repl::RetryImage_parse(
+                IDLParserErrorContext("applyUpdate"),
+                op.getField(OplogEntryBase::kNeedsRetryImageFieldName).String());
+            if (imageKind == repl::RetryImageEnum::kPreImage) {
+                request.setReturnDocs(UpdateRequest::ReturnDocOption::RETURN_OLD);
+            } else if (imageKind == repl::RetryImageEnum::kPostImage) {
+                request.setReturnDocs(UpdateRequest::ReturnDocOption::RETURN_NEW);
+            }
+        }
 
         UpdateLifecycleImpl updateLifecycle(requestNss);
         request.setLifecycle(&updateLifecycle);
@@ -1456,6 +1498,7 @@ Status applyOperation_inlock(OperationContext* opCtx,
         }
 
         const StringData ns = fieldNs.valuestrsafe();
+        bool upsertConfigImage = true;
         auto status = writeConflictRetry(opCtx, "applyOps_update", ns, [&] {
             WriteUnitOfWork wuow(opCtx);
             if (timestamp != Timestamp::min()) {
@@ -1501,7 +1544,11 @@ Status applyOperation_inlock(OperationContext* opCtx,
                     }
                 }
             }
-
+            if (op.hasField(OplogEntryBase::kNeedsRetryImageFieldName)) {
+                invariant(imageKind);
+                writeToImageCollection(
+                    opCtx, op, ur.requestedDocImage, *imageKind, &upsertConfigImage);
+            }
             wuow.commit();
             return Status::OK();
         });
diff --git a/src/mongo/db/session_catalog.cpp b/src/mongo/db/session_catalog.cpp
index 4b2355ce6d2..c9a6e615131 100644
--- a/src/mongo/db/session_catalog.cpp
+++ b/src/mongo/db/session_catalog.cpp
@@ -41,6 +41,7 @@
 #include "mongo/db/dbdirectclient.h"
 #include "mongo/db/kill_sessions_common.h"
 #include "mongo/db/namespace_string.h"
+#include "mongo/db/op_observer.h"
 #include "mongo/db/operation_context.h"
 #include "mongo/db/ops/write_ops.h"
 #include "mongo/db/service_context.h"
@@ -207,6 +208,22 @@ void SessionCatalog::onStepUp(OperationContext* opCtx) {
 
     BSONObj result;
 
+    if (storeFindAndModifyImagesInSideCollection.load() &&
+        serverGlobalParams.featureCompatibility.getVersion() >=
+            ServerGlobalParams::FeatureCompatibility::Version::kUpgradingTo40) {
+        BSONObj imageResult;
+        client.createCollection(NamespaceString::kConfigImagesNamespace.ns(),
+                                initialExtentSize,
+                                capped,
+                                maxSize,
+                                &imageResult);
+        const auto status = getStatusFromCommandResult(imageResult);
+        uassertStatusOKWithContext(status,
+                                   str::stream() << "Failed to create the "
+                                                 << NamespaceString::kConfigImagesNamespace.ns()
+                                                 << " collection");
+    }
+
     if (client.createCollection(NamespaceString::kSessionTransactionsTableNamespace.ns(),
                                 initialExtentSize,
                                 capped,