SERVER-61480 Improve efficiency of change stream pre-image loading

9 files changed, 60 insertions, 29 deletions

diff --git a/src/mongo/db/pipeline/document_source_change_stream_add_pre_image.cpp b/src/mongo/db/pipeline/document_source_change_stream_add_pre_image.cpp
index dad3a18b306..97b252c9829 100644
--- a/src/mongo/db/pipeline/document_source_change_stream_add_pre_image.cpp
+++ b/src/mongo/db/pipeline/document_source_change_stream_add_pre_image.cpp
@@ -134,34 +134,22 @@ boost::optional<Document> DocumentSourceChangeStreamAddPreImage::lookupPreImage(
         return change_stream_legacy::legacyLookupPreImage(pExpCtx, preImageId);
     }
 
-    // We need the pre-images UUID for lookup, so obtain the collection info via
-    // MongoProcessInterface.
-    auto preImagesCollectionInfo = pExpCtx->mongoProcessInterface->getCollectionOptions(
-        pExpCtx->opCtx, NamespaceString::kChangeStreamPreImagesNamespace);
-
-    // Return boost::none if pre-images collection doesn't exist.
-    if (preImagesCollectionInfo.isEmpty()) {
-        return boost::none;
-    }
-
-    // Extract the UUID from the collection information. We should always have a valid uuid here.
-    auto preImagesCollectionUUID = invariantStatusOK(UUID::parse(preImagesCollectionInfo["uuid"]));
-
-    // Look up the pre-image using the pre-image id as the query filter.
-    auto lookedUpDoc = pExpCtx->mongoProcessInterface->lookupSingleDocument(
+    // Look up the pre-image document on the local node by id.
+    auto lookedUpDoc = pExpCtx->mongoProcessInterface->lookupSingleDocumentLocally(
         pExpCtx,
         NamespaceString::kChangeStreamPreImagesNamespace,
-        preImagesCollectionUUID,
-        Document{{"_id", Value(preImageId)}},
-        boost::none);
+        Document{{ChangeStreamPreImage::kIdFieldName, preImageId}});
 
     // Return boost::none to signify that we failed to find the pre-image.
     if (!lookedUpDoc) {
         return boost::none;
     }
 
-    // Return preImage field from the document.
-    return lookedUpDoc->getField(ChangeStreamPreImage::kPreImageFieldName).getDocument().getOwned();
+    // Return "preImage" field value from the document.
+    auto preImageField = lookedUpDoc->getField(ChangeStreamPreImage::kPreImageFieldName);
+    tassert(
+        6148000, "Pre-image document must contain the 'preImage' field", !preImageField.nullish());
+    return preImageField.getDocument().getOwned();
 }
 
 Value DocumentSourceChangeStreamAddPreImage::serialize(
diff --git a/src/mongo/db/pipeline/process_interface/SConscript b/src/mongo/db/pipeline/process_interface/SConscript
index 87c6fe453dd..44b0afe1591 100644
--- a/src/mongo/db/pipeline/process_interface/SConscript
+++ b/src/mongo/db/pipeline/process_interface/SConscript
@@ -48,6 +48,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/catalog/database_holder',
         '$BUILD_DIR/mongo/db/collection_index_usage_tracker',
         '$BUILD_DIR/mongo/db/concurrency/flow_control_ticketholder',
+        '$BUILD_DIR/mongo/db/dbhelpers',
         '$BUILD_DIR/mongo/db/index_builds_coordinator_mongod',
         '$BUILD_DIR/mongo/db/multitenancy',
         '$BUILD_DIR/mongo/db/repl/primary_only_service',
diff --git a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp
index 2cf8fa19807..797f4ddd0d8 100644
--- a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp
+++ b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.cpp
@@ -48,6 +48,7 @@
 #include "mongo/db/curop.h"
 #include "mongo/db/cursor_manager.h"
 #include "mongo/db/db_raii.h"
+#include "mongo/db/dbhelpers.h"
 #include "mongo/db/index/index_descriptor.h"
 #include "mongo/db/pipeline/document_source_cursor.h"
 #include "mongo/db/pipeline/lite_parsed_pipeline.h"
@@ -734,4 +735,17 @@ void CommonMongodProcessInterface::truncateRecordStore(
     });
 }
 
+boost::optional<Document> CommonMongodProcessInterface::lookupSingleDocumentLocally(
+    const boost::intrusive_ptr<ExpressionContext>& expCtx,
+    const NamespaceString& nss,
+    const Document& documentKey) {
+    AutoGetCollectionForRead autoColl(expCtx->opCtx, nss);
+    BSONObj document;
+    if (!Helpers::findById(
+            expCtx->opCtx, autoColl.getDb(), nss.ns(), documentKey.toBson(), document)) {
+        return boost::none;
+    }
+    return Document(document).getOwned();
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.h b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.h
index 62954640d4b..0890f0e1ba1 100644
--- a/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/common_mongod_process_interface.h
@@ -120,6 +120,11 @@ public:
     void truncateRecordStore(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                              RecordStore* rs) const final;
 
+    boost::optional<Document> lookupSingleDocumentLocally(
+        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+        const NamespaceString& nss,
+        const Document& documentKey) final;
+
 protected:
     BSONObj getCollectionOptionsLocally(OperationContext* opCtx, const NamespaceString& nss);
 
diff --git a/src/mongo/db/pipeline/process_interface/mongo_process_interface.h b/src/mongo/db/pipeline/process_interface/mongo_process_interface.h
index fba8fcfea3d..ae7a625675b 100644
--- a/src/mongo/db/pipeline/process_interface/mongo_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/mongo_process_interface.h
@@ -357,6 +357,17 @@ public:
         boost::optional<BSONObj> readConcern) = 0;
 
     /**
+     * Returns zero or one document with the document _id being equal to 'documentKey'. The document
+     * is looked up only on the current node. Returns boost::none if no matching documents were
+     * found, including cases where the given namespace does not exist. It is illegal to call this
+     * method on nodes other than mongod.
+     */
+    virtual boost::optional<Document> lookupSingleDocumentLocally(
+        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+        const NamespaceString& nss,
+        const Document& documentKey) = 0;
+
+    /**
      * Returns a vector of all idle (non-pinned) local cursors.
      */
     virtual std::vector<GenericCursor> getIdleCursors(
diff --git a/src/mongo/db/pipeline/process_interface/mongos_process_interface.cpp b/src/mongo/db/pipeline/process_interface/mongos_process_interface.cpp
index 35c6febec29..6f17c7a0121 100644
--- a/src/mongo/db/pipeline/process_interface/mongos_process_interface.cpp
+++ b/src/mongo/db/pipeline/process_interface/mongos_process_interface.cpp
@@ -218,6 +218,13 @@ boost::optional<Document> MongosProcessInterface::lookupSingleDocument(
     }
 }
 
+boost::optional<Document> MongosProcessInterface::lookupSingleDocumentLocally(
+    const boost::intrusive_ptr<ExpressionContext>& expCtx,
+    const NamespaceString& nss,
+    const Document& documentKey) {
+    MONGO_UNREACHABLE_TASSERT(6148001);
+}
+
 BSONObj MongosProcessInterface::_reportCurrentOpForClient(
     OperationContext* opCtx,
     Client* client,
diff --git a/src/mongo/db/pipeline/process_interface/mongos_process_interface.h b/src/mongo/db/pipeline/process_interface/mongos_process_interface.h
index e08a8169ed7..923bf0917a8 100644
--- a/src/mongo/db/pipeline/process_interface/mongos_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/mongos_process_interface.h
@@ -52,6 +52,11 @@ public:
         const Document& documentKey,
         boost::optional<BSONObj> readConcern) final;
 
+    boost::optional<Document> lookupSingleDocumentLocally(
+        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+        const NamespaceString& nss,
+        const Document& documentKey) final;
+
     std::vector<GenericCursor> getIdleCursors(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                                               CurrentOpUserMode userMode) const final;
 
diff --git a/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h b/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h
index 5fe7d78a26b..314c32c160f 100644
--- a/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/stub_mongo_process_interface.h
@@ -198,6 +198,13 @@ public:
         MONGO_UNREACHABLE;
     }
 
+    boost::optional<Document> lookupSingleDocumentLocally(
+        const boost::intrusive_ptr<ExpressionContext>& expCtx,
+        const NamespaceString& nss,
+        const Document& documentKey) {
+        MONGO_UNREACHABLE_TASSERT(6148002);
+    }
+
     std::vector<GenericCursor> getIdleCursors(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                                               CurrentOpUserMode userMode) const {
         MONGO_UNREACHABLE;
diff --git a/src/mongo/db/repl/oplog.cpp b/src/mongo/db/repl/oplog.cpp
index 232bf268931..0030db48c13 100644
--- a/src/mongo/db/repl/oplog.cpp
+++ b/src/mongo/db/repl/oplog.cpp
@@ -1553,15 +1553,8 @@ Status applyOperation_inlock(OperationContext* opCtx,
                     // document.
                     invariant(op.getObject2());
                     auto&& documentId = *op.getObject2();
-
-                    // Assume that either an index on _id exists or the collection is clustered by
-                    // _id.
-                    const bool requireIndex{false};
-
-                    // TODO: SERVER-61480 use a more efficient implementation - bypass the find
-                    // component.
-                    auto documentFound = Helpers::findOne(
-                        opCtx, collection, documentId, changeStreamPreImage, requireIndex);
+                    auto documentFound = Helpers::findById(
+                        opCtx, db, collection->ns().ns(), documentId, changeStreamPreImage);
                     invariant(documentFound);
                 }