SERVER-50870 Establish consistent storage snapshot with the refreshed Collection after yielding

12 files changed, 336 insertions, 120 deletions

diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index 3ae923af8f6..a12cb0f9b76 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -1257,7 +1257,6 @@ env.Library(
         'commands/server_status_core',
         'kill_sessions',
         'stats/resource_consumption_metrics',
-        'storage/snapshot_helper',
     ],
 )
 
diff --git a/src/mongo/db/catalog/SConscript b/src/mongo/db/catalog/SConscript
index 84c782e70a8..2e7152c8a7d 100644
--- a/src/mongo/db/catalog/SConscript
+++ b/src/mongo/db/catalog/SConscript
@@ -272,6 +272,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/namespace_string',
         '$BUILD_DIR/mongo/db/profile_filter',
         '$BUILD_DIR/mongo/db/service_context',
+        '$BUILD_DIR/mongo/db/storage/snapshot_helper',
         'collection',
     ],
     LIBDEPS_PRIVATE=[
diff --git a/src/mongo/db/catalog/collection_catalog.cpp b/src/mongo/db/catalog/collection_catalog.cpp
index e6e28b0d920..671417978c8 100644
--- a/src/mongo/db/catalog/collection_catalog.cpp
+++ b/src/mongo/db/catalog/collection_catalog.cpp
@@ -38,6 +38,7 @@
 #include "mongo/db/concurrency/write_conflict_exception.h"
 #include "mongo/db/server_options.h"
 #include "mongo/db/storage/recovery_unit.h"
+#include "mongo/db/storage/snapshot_helper.h"
 #include "mongo/logv2/log.h"
 #include "mongo/util/assert_util.h"
 #include "mongo/util/uuid.h"
@@ -1037,7 +1038,21 @@ void CollectionCatalogStasher::reset() {
 
 const Collection* LookupCollectionForYieldRestore::operator()(OperationContext* opCtx,
                                                               CollectionUUID uuid) const {
-    return CollectionCatalog::get(opCtx)->lookupCollectionByUUID(opCtx, uuid).get();
+    auto collection = CollectionCatalog::get(opCtx)->lookupCollectionByUUID(opCtx, uuid).get();
+    if (!collection)
+        return nullptr;
+
+    // After yielding and reacquiring locks, the preconditions that were used to select our
+    // ReadSource initially need to be checked again. We select a ReadSource based on replication
+    // state. After a query yields its locks, the replication state may have changed, invalidating
+    // our current choice of ReadSource. Using the same preconditions, change our ReadSource if
+    // necessary.
+    auto newReadSource = SnapshotHelper::getNewReadSource(opCtx, collection->ns());
+    if (newReadSource) {
+        opCtx->recoveryUnit()->setTimestampReadSource(*newReadSource);
+    }
+
+    return collection;
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/catalog_raii.cpp b/src/mongo/db/catalog_raii.cpp
index ba7dbe54cc2..92e0b9c4f96 100644
--- a/src/mongo/db/catalog_raii.cpp
+++ b/src/mongo/db/catalog_raii.cpp
@@ -207,6 +207,7 @@ Collection* AutoGetCollection::getWritableCollection(CollectionCatalog::Lifetime
 
 AutoGetCollectionLockFree::AutoGetCollectionLockFree(OperationContext* opCtx,
                                                      const NamespaceStringOrUUID& nsOrUUID,
+                                                     RestoreFromYieldFn restoreFromYield,
                                                      AutoGetCollectionViewMode viewMode,
                                                      Date_t deadline)
     : _lockFreeReadsBlock(opCtx),
@@ -223,11 +224,13 @@ AutoGetCollectionLockFree::AutoGetCollectionLockFree(OperationContext* opCtx,
 
     // When we restore from yield on this CollectionPtr we will update _collection above and use its
     // new pointer in the CollectionPtr
-    _collectionPtr = CollectionPtr(
-        opCtx, _collection.get(), [this](OperationContext* opCtx, CollectionUUID uuid) {
-            _collection = CollectionCatalog::get(opCtx)->lookupCollectionByUUIDForRead(opCtx, uuid);
-            return _collection.get();
-        });
+    _collectionPtr = CollectionPtr(opCtx,
+                                   _collection.get(),
+                                   [this, restoreFromYield = std::move(restoreFromYield)](
+                                       OperationContext* opCtx, CollectionUUID uuid) {
+                                       restoreFromYield(_collection, opCtx, uuid);
+                                       return _collection.get();
+                                   });
 
     // TODO (SERVER-51319): add DatabaseShardingState::checkDbVersion somewhere.
 
diff --git a/src/mongo/db/catalog_raii.h b/src/mongo/db/catalog_raii.h
index f5e7a8e5f26..6ddde54c1c3 100644
--- a/src/mongo/db/catalog_raii.h
+++ b/src/mongo/db/catalog_raii.h
@@ -206,30 +206,30 @@ protected:
  * The collection references returned by this class will no longer be safe to retain after this
  * object goes out of scope. This object ensures the continued existence of a Collection reference,
  * if the collection exists when this object is instantiated.
+ *
+ * This class is only used by AutoGetCollectionForReadLockFree.
  */
 class AutoGetCollectionLockFree {
     AutoGetCollectionLockFree(const AutoGetCollectionLockFree&) = delete;
     AutoGetCollectionLockFree& operator=(const AutoGetCollectionLockFree&) = delete;
 
 public:
-    AutoGetCollectionLockFree(
-        OperationContext* opCtx,
-        const NamespaceStringOrUUID& nsOrUUID,
-        AutoGetCollectionViewMode viewMode = AutoGetCollectionViewMode::kViewsForbidden,
-        Date_t deadline = Date_t::max());
+    /**
+     * Function used to customize restore after yield behavior
+     */
+    using RestoreFromYieldFn =
+        std::function<void(std::shared_ptr<const Collection>&, OperationContext*, CollectionUUID)>;
 
     /**
-     * Same constructor as above except it accepts a fifth unused LockMode parameter in order to
-     * parallel AutoGetCollection and meet the templated AutoGetCollectionForReadBase class'
-     * type structure expectations.
+     * Used by AutoGetCollectionForReadLockFree where it provides implementation for restore after
+     * yield.
      */
     AutoGetCollectionLockFree(
         OperationContext* opCtx,
         const NamespaceStringOrUUID& nsOrUUID,
-        LockMode unused,  // unused
+        RestoreFromYieldFn restoreFromYield,
         AutoGetCollectionViewMode viewMode = AutoGetCollectionViewMode::kViewsForbidden,
-        Date_t deadline = Date_t::max())
-        : AutoGetCollectionLockFree(opCtx, nsOrUUID, viewMode, deadline) {}
+        Date_t deadline = Date_t::max());
 
     explicit operator bool() const {
         // Use the CollectionPtr because it is updated if it yields whereas _collection is not until
diff --git a/src/mongo/db/concurrency/SConscript b/src/mongo/db/concurrency/SConscript
index 892e7364cc8..37b0095d141 100644
--- a/src/mongo/db/concurrency/SConscript
+++ b/src/mongo/db/concurrency/SConscript
@@ -37,6 +37,16 @@ env.Library(
 )
 
 env.Library(
+    target='lock_manager_defs',
+    source=[
+        'lock_manager_defs.cpp',
+    ],
+    LIBDEPS=[
+        '$BUILD_DIR/mongo/base',
+    ],
+)
+
+env.Library(
     target='lock_manager',
     source=[
         'd_concurrency.cpp',
@@ -52,6 +62,7 @@ env.Library(
         '$BUILD_DIR/mongo/util/concurrency/spin_lock',
         '$BUILD_DIR/mongo/util/concurrency/ticketholder',
         '$BUILD_DIR/third_party/shim_boost',
+        'lock_manager_defs',
     ],
     LIBDEPS_PRIVATE=[
         '$BUILD_DIR/mongo/db/catalog/collection_catalog',
diff --git a/src/mongo/db/concurrency/lock_manager_defs.cpp b/src/mongo/db/concurrency/lock_manager_defs.cpp
new file mode 100644
index 00000000000..01ae4586cad
--- /dev/null
+++ b/src/mongo/db/concurrency/lock_manager_defs.cpp
@@ -0,0 +1,42 @@
+/**
+ *    Copyright (C) 2020-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.
+ */
+
+#include "lock_manager_defs.h"
+
+namespace mongo {
+
+// Hardcoded resource IDs.
+const ResourceId resourceIdLocalDB = ResourceId(RESOURCE_DATABASE, StringData("local"));
+const ResourceId resourceIdOplog = ResourceId(RESOURCE_COLLECTION, StringData("local.oplog.rs"));
+const ResourceId resourceIdAdminDB = ResourceId(RESOURCE_DATABASE, StringData("admin"));
+const ResourceId resourceIdGlobal = ResourceId(RESOURCE_GLOBAL, 1ULL);
+const ResourceId resourceIdParallelBatchWriterMode = ResourceId(RESOURCE_PBWM, 1ULL);
+const ResourceId resourceIdReplicationStateTransitionLock = ResourceId(RESOURCE_RSTL, 1ULL);
+
+}  // namespace mongo
diff --git a/src/mongo/db/concurrency/lock_state.cpp b/src/mongo/db/concurrency/lock_state.cpp
index 416c8c5e7e2..3be2cebb251 100644
--- a/src/mongo/db/concurrency/lock_state.cpp
+++ b/src/mongo/db/concurrency/lock_state.cpp
@@ -1109,12 +1109,4 @@ void resetGlobalLockStats() {
     globalStats.reset();
 }
 
-// Hardcoded resource IDs.
-const ResourceId resourceIdLocalDB = ResourceId(RESOURCE_DATABASE, StringData("local"));
-const ResourceId resourceIdOplog = ResourceId(RESOURCE_COLLECTION, StringData("local.oplog.rs"));
-const ResourceId resourceIdAdminDB = ResourceId(RESOURCE_DATABASE, StringData("admin"));
-const ResourceId resourceIdGlobal = ResourceId(RESOURCE_GLOBAL, 1ULL);
-const ResourceId resourceIdParallelBatchWriterMode = ResourceId(RESOURCE_PBWM, 1ULL);
-const ResourceId resourceIdReplicationStateTransitionLock = ResourceId(RESOURCE_RSTL, 1ULL);
-
 }  // namespace mongo
diff --git a/src/mongo/db/db_raii.cpp b/src/mongo/db/db_raii.cpp
index cb1d875cb1f..2b146c8ce26 100644
--- a/src/mongo/db/db_raii.cpp
+++ b/src/mongo/db/db_raii.cpp
@@ -62,6 +62,90 @@ bool supportsLockFreeRead(OperationContext* opCtx) {
         !opCtx->inMultiDocumentTransaction();
 }
 
+/**
+ * Helper function to acquire a collection and consistent snapshot without holding the RSTL or
+ * collection locks.
+ *
+ * GetCollectionAndEstablishReadSourceFunc is called before we open a snapshot, it needs to fetch
+ * the Collection from the catalog and select the read source.
+ *
+ * GetCollectionAfterSnapshotFunc is called after the snapshot is opened, it needs to fetch the
+ * Collection from the catalog that is used to compare consistency with the Collection returned by
+ * GetCollectionAndEstablishReadSourceFunc.
+ *
+ * ResetFunc is called when we failed to achieve consistency and need to retry.
+ */
+template <typename GetCollectionAndEstablishReadSourceFunc,
+          typename GetCollectionAfterSnapshotFunc,
+          typename ResetFunc>
+auto aquireCollectionAndConsistentSnapshot(
+    OperationContext* opCtx,
+    CollectionCatalogStasher& catalogStasher,
+    GetCollectionAndEstablishReadSourceFunc getCollectionAndEstablishReadSource,
+    GetCollectionAfterSnapshotFunc getCollectionAfterSnapshot,
+    ResetFunc reset) {
+    // Figure out what type of Collection GetCollectionAndEstablishReadSourceFunc returns. It needs
+    // to behave like a pointer.
+    using CollectionPtrT = decltype(std::declval<GetCollectionAndEstablishReadSourceFunc>()(
+        std::declval<OperationContext*>(), std::declval<const CollectionCatalog&>()));
+
+    CollectionPtrT collection;
+    catalogStasher.reset();
+    while (true) {
+        // AutoGetCollectionForReadBase can choose a read source based on the current replication
+        // state. Therefore we must fetch the repl state beforehand, to compare with afterwards.
+        long long replTerm = repl::ReplicationCoordinator::get(opCtx)->getTerm();
+
+        auto catalog = CollectionCatalog::get(opCtx);
+        collection = getCollectionAndEstablishReadSource(opCtx, *catalog);
+
+        // A lock request does not always find a collection to lock.
+        if (!collection)
+            break;
+
+        // We must open a storage snapshot consistent with the fetched in-memory Collection instance
+        // and chosen read source. The Collection instance and replication state after opening a
+        // snapshot will be compared with the previously acquired state. If either does not match,
+        // then this loop will retry lock acquisition and read source selection until there is a
+        // match.
+        //
+        // Note: getCollectionAndEstablishReadSource() may open a snapshot for PIT reads, so
+        // preallocateSnapshot() may be a no-op, but that is OK because the snapshot is established
+        // by getCollectionAndEstablishReadSource() after it fetches a Collection instance.
+        if (collection->ns().isOplog()) {
+            // Signal to the RecoveryUnit that the snapshot will be used for reading the oplog.
+            // Normally the snapshot is opened from a cursor that can take special action when
+            // reading from the oplog.
+            opCtx->recoveryUnit()->preallocateSnapshotForOplogRead();
+        } else {
+            opCtx->recoveryUnit()->preallocateSnapshot();
+        }
+
+        // The collection may have been dropped since the previous lookup, run the loop one more
+        // time to cleanup if newCollection is nullptr
+        auto newCatalog = CollectionCatalog::get(opCtx);
+        if (catalog == newCatalog) {
+            auto newCollection = getCollectionAfterSnapshot(opCtx, *catalog);
+            if (newCollection && catalog == newCatalog &&
+                collection->getMinimumVisibleSnapshot() ==
+                    newCollection->getMinimumVisibleSnapshot() &&
+                replTerm == repl::ReplicationCoordinator::get(opCtx)->getTerm()) {
+                catalogStasher.stash(std::move(catalog));
+                break;
+            }
+        }
+
+        LOGV2_DEBUG(5067701,
+                    3,
+                    "Retrying acquiring state for lock-free read because collection or replication "
+                    "state changed.");
+        reset();
+        opCtx->recoveryUnit()->abandonSnapshot();
+    }
+
+    return collection;
+}
+
 }  // namespace
 
 AutoStatsTracker::AutoStatsTracker(OperationContext* opCtx,
@@ -95,12 +179,10 @@ AutoStatsTracker::~AutoStatsTracker() {
                 curOp->getReadWriteType());
 }
 
-template <typename AutoGetCollectionType>
-AutoGetCollectionForReadBase<AutoGetCollectionType>::AutoGetCollectionForReadBase(
-    OperationContext* opCtx,
-    const NamespaceStringOrUUID& nsOrUUID,
-    AutoGetCollectionViewMode viewMode,
-    Date_t deadline) {
+template <typename AutoGetCollectionType, typename EmplaceAutoCollFunc>
+AutoGetCollectionForReadBase<AutoGetCollectionType, EmplaceAutoCollFunc>::
+    AutoGetCollectionForReadBase(OperationContext* opCtx,
+                                 const EmplaceAutoCollFunc& emplaceAutoColl) {
     // The caller was expecting to conflict with batch application before entering this function.
     // i.e. the caller does not currently have a ShouldNotConflict... block in scope.
     bool callerWasConflicting = opCtx->lockState()->shouldConflictWithSecondaryBatchApplication();
@@ -110,10 +192,7 @@ AutoGetCollectionForReadBase<AutoGetCollectionType>::AutoGetCollectionForReadBas
         _shouldNotConflictWithSecondaryBatchApplicationBlock.emplace(opCtx->lockState());
     }
 
-    // Multi-document transactions need MODE_IX locks, otherwise MODE_IS.
-    const auto collectionLockMode = getLockModeForQuery(opCtx, nsOrUUID.nss());
-
-    _autoColl.emplace(opCtx, nsOrUUID, collectionLockMode, viewMode, deadline);
+    emplaceAutoColl.emplace(_autoColl);
 
     repl::ReplicationCoordinator* const replCoord = repl::ReplicationCoordinator::get(opCtx);
     const auto readConcernLevel = repl::ReadConcernArgs::get(opCtx).getLevel();
@@ -253,10 +332,84 @@ AutoGetCollectionForReadBase<AutoGetCollectionType>::AutoGetCollectionForReadBas
             CurOp::get(opCtx)->yielded();
         }
 
-        _autoColl.emplace(opCtx, nsOrUUID, collectionLockMode, viewMode, deadline);
+        emplaceAutoColl.emplace(_autoColl);
     }
 }
 
+EmplaceAutoGetCollectionForRead::EmplaceAutoGetCollectionForRead(
+    OperationContext* opCtx,
+    const NamespaceStringOrUUID& nsOrUUID,
+    AutoGetCollectionViewMode viewMode,
+    Date_t deadline)
+    : _opCtx(opCtx), _nsOrUUID(nsOrUUID), _viewMode(viewMode), _deadline(deadline) {
+    // Multi-document transactions need MODE_IX locks, otherwise MODE_IS.
+    _collectionLockMode = getLockModeForQuery(opCtx, nsOrUUID.nss());
+}
+
+void EmplaceAutoGetCollectionForRead::emplace(boost::optional<AutoGetCollection>& autoColl) const {
+    autoColl.emplace(_opCtx, _nsOrUUID, _collectionLockMode, _viewMode, _deadline);
+}
+
+AutoGetCollectionForRead::AutoGetCollectionForRead(OperationContext* opCtx,
+                                                   const NamespaceStringOrUUID& nsOrUUID,
+                                                   AutoGetCollectionViewMode viewMode,
+                                                   Date_t deadline)
+    : AutoGetCollectionForReadBase(
+          opCtx, EmplaceAutoGetCollectionForRead(opCtx, nsOrUUID, viewMode, deadline)) {}
+
+AutoGetCollectionForReadLockFree::EmplaceHelper::EmplaceHelper(
+    OperationContext* opCtx,
+    CollectionCatalogStasher& catalogStasher,
+    const NamespaceStringOrUUID& nsOrUUID,
+    AutoGetCollectionViewMode viewMode,
+    Date_t deadline)
+    : _opCtx(opCtx),
+      _catalogStasher(catalogStasher),
+      _nsOrUUID(nsOrUUID),
+      _viewMode(viewMode),
+      _deadline(deadline) {}
+
+void AutoGetCollectionForReadLockFree::EmplaceHelper::emplace(
+    boost::optional<AutoGetCollectionLockFree>& autoColl) const {
+    autoColl.emplace(
+        _opCtx,
+        _nsOrUUID,
+        /* restoreFromYield */
+        [& catalogStasher = _catalogStasher](std::shared_ptr<const Collection>& collection,
+                                             OperationContext* opCtx,
+                                             CollectionUUID uuid) {
+            collection = aquireCollectionAndConsistentSnapshot(
+                opCtx,
+                catalogStasher,
+                /* GetCollectionAndEstablishReadSourceFunc */
+                [uuid](OperationContext* opCtx, const CollectionCatalog& catalog) {
+                    auto coll = catalog.lookupCollectionByUUIDForRead(opCtx, uuid);
+
+                    // After yielding and reacquiring locks, the preconditions that were used to
+                    // select our ReadSource initially need to be checked again. We select a
+                    // ReadSource based on replication state. After a query yields its locks, the
+                    // replication state may have changed, invalidating our current choice of
+                    // ReadSource. Using the same preconditions, change our ReadSource if necessary.
+                    if (coll) {
+                        auto newReadSource = SnapshotHelper::getNewReadSource(opCtx, coll->ns());
+                        if (newReadSource) {
+                            opCtx->recoveryUnit()->setTimestampReadSource(*newReadSource);
+                        }
+                    }
+
+                    return coll;
+                },
+                /* GetCollectionAfterSnapshotFunc */
+                [uuid](OperationContext* opCtx, const CollectionCatalog& catalog) {
+                    return catalog.lookupCollectionByUUIDForRead(opCtx, uuid);
+                },
+                /* ResetFunc */
+                []() {});
+        },
+        _viewMode,
+        _deadline);
+}
+
 AutoGetCollectionForReadLockFree::AutoGetCollectionForReadLockFree(
     OperationContext* opCtx,
     const NamespaceStringOrUUID& nsOrUUID,
@@ -267,59 +420,22 @@ AutoGetCollectionForReadLockFree::AutoGetCollectionForReadLockFree(
     // this helper. The storage snapshot and in-memory state fetched here must be consistent.
     invariant(supportsLockFreeRead(opCtx) && !opCtx->recoveryUnit()->isActive());
 
-    while (true) {
-        // AutoGetCollectionForReadBase can choose a read source based on the current replication
-        // state. Therefore we must fetch the repl state beforehand, to compare with afterwards.
-        long long replTerm = repl::ReplicationCoordinator::get(opCtx)->getTerm();
-
-        auto catalog = CollectionCatalog::get(opCtx);
-        _autoGetCollectionForReadBase.emplace(opCtx, nsOrUUID, viewMode, deadline);
-
-        // A lock request does not always find a collection to lock.
-        if (!_autoGetCollectionForReadBase.get()) {
-            break;
-        }
-
-        // We must open a storage snapshot consistent with the fetched in-memory Collection instance
-        // and chosen read source. The Collection instance and replication state after opening a
-        // snapshot will be compared with the previously acquired state. If either does not match,
-        // then this loop will retry lock acquisition and read source selection until there is a
-        // match.
-        //
-        // Note: AutoGetCollectionForReadBase may open a snapshot for PIT reads, so
-        // preallocateSnapshot() may be a no-op, but that is OK because the snapshot is established
-        // by _autoGetCollectionForReadBase after it fetches a Collection instance.
-
-        if (_autoGetCollectionForReadBase->getNss().isOplog()) {
-            // Signal to the RecoveryUnit that the snapshot will be used for reading the oplog.
-            // Normally the snapshot is opened from a cursor that can take special action when
-            // reading from the oplog.
-            opCtx->recoveryUnit()->preallocateSnapshotForOplogRead();
-        } else {
-            opCtx->recoveryUnit()->preallocateSnapshot();
-        }
-
-        auto newCatalog = CollectionCatalog::get(opCtx);
-        auto newCollection = newCatalog->lookupCollectionByUUIDForRead(
-            opCtx, _autoGetCollectionForReadBase.get()->uuid());
-
-        // The collection may have been dropped since the previous lookup, run the loop one more
-        // time to cleanup if newCollection is nullptr
-        if (newCollection && catalog == newCatalog &&
-            _autoGetCollectionForReadBase.get()->getMinimumVisibleSnapshot() ==
-                newCollection->getMinimumVisibleSnapshot() &&
-            replTerm == repl::ReplicationCoordinator::get(opCtx)->getTerm()) {
-            _catalogStash.stash(std::move(catalog));
-            break;
-        }
-
-        LOGV2_DEBUG(5067701,
-                    3,
-                    "Retrying acquiring state for lock-free read because collection or replication "
-                    "state changed.");
-        _autoGetCollectionForReadBase.reset();
-        opCtx->recoveryUnit()->abandonSnapshot();
-    }
+    EmplaceHelper emplaceFunc(opCtx, _catalogStash, nsOrUUID, viewMode, deadline);
+    aquireCollectionAndConsistentSnapshot(
+        opCtx,
+        _catalogStash,
+        /* GetCollectionAndEstablishReadSourceFunc */
+        [this, &emplaceFunc](OperationContext* opCtx, const CollectionCatalog&) {
+            _autoGetCollectionForReadBase.emplace(opCtx, emplaceFunc);
+            return _autoGetCollectionForReadBase->getCollection().get();
+        },
+        /* GetCollectionAfterSnapshotFunc */
+        [this](OperationContext* opCtx, const CollectionCatalog& catalog) {
+            return catalog.lookupCollectionByUUIDForRead(
+                opCtx, _autoGetCollectionForReadBase.get()->uuid());
+        },
+        /* ResetFunc */
+        [this]() { _autoGetCollectionForReadBase.reset(); });
 }
 
 AutoGetCollectionForReadMaybeLockFree::AutoGetCollectionForReadMaybeLockFree(
@@ -492,9 +608,10 @@ BlockSecondaryReadsDuringBatchApplication_DONT_USE::
     allowSecondaryReads->swap(_originalSettings);
 }
 
-template class AutoGetCollectionForReadBase<AutoGetCollection>;
+template class AutoGetCollectionForReadBase<AutoGetCollection, EmplaceAutoGetCollectionForRead>;
 template class AutoGetCollectionForReadCommandBase<AutoGetCollectionForRead>;
-template class AutoGetCollectionForReadBase<AutoGetCollectionLockFree>;
+template class AutoGetCollectionForReadBase<AutoGetCollectionLockFree,
+                                            AutoGetCollectionForReadLockFree::EmplaceHelper>;
 template class AutoGetCollectionForReadCommandBase<AutoGetCollectionForReadLockFree>;
 
 }  // namespace mongo
diff --git a/src/mongo/db/db_raii.h b/src/mongo/db/db_raii.h
index d8d7a0a3390..57c8d8a1753 100644
--- a/src/mongo/db/db_raii.h
+++ b/src/mongo/db/db_raii.h
@@ -84,17 +84,18 @@ private:
     const LogMode _logMode;
 };
 
-template <typename AutoGetCollectionType>
+/**
+ * Shared base class for AutoGetCollectionForRead and AutoGetCollectionForReadLockFree.
+ * Do not use directly.
+ */
+template <typename AutoGetCollectionType, typename EmplaceAutoGetCollectionFunc>
 class AutoGetCollectionForReadBase {
     AutoGetCollectionForReadBase(const AutoGetCollectionForReadBase&) = delete;
     AutoGetCollectionForReadBase& operator=(const AutoGetCollectionForReadBase&) = delete;
 
 public:
-    AutoGetCollectionForReadBase(
-        OperationContext* opCtx,
-        const NamespaceStringOrUUID& nsOrUUID,
-        AutoGetCollectionViewMode viewMode = AutoGetCollectionViewMode::kViewsForbidden,
-        Date_t deadline = Date_t::max());
+    AutoGetCollectionForReadBase(OperationContext* opCtx,
+                                 const EmplaceAutoGetCollectionFunc& emplaceAutoColl);
 
     explicit operator bool() const {
         return static_cast<bool>(getCollection());
@@ -133,6 +134,27 @@ protected:
 };
 
 /**
+ * Helper for AutoGetCollectionForRead below. Contains implementation on how contained
+ * AutoGetCollection is instantiated by AutoGetCollectionForReadBase.
+ */
+class EmplaceAutoGetCollectionForRead {
+public:
+    EmplaceAutoGetCollectionForRead(OperationContext* opCtx,
+                                    const NamespaceStringOrUUID& nsOrUUID,
+                                    AutoGetCollectionViewMode viewMode,
+                                    Date_t deadline);
+
+    void emplace(boost::optional<AutoGetCollection>& autoColl) const;
+
+private:
+    OperationContext* _opCtx;
+    const NamespaceStringOrUUID& _nsOrUUID;
+    AutoGetCollectionViewMode _viewMode;
+    Date_t _deadline;
+    LockMode _collectionLockMode;
+};
+
+/**
  * Same as calling AutoGetCollection with MODE_IS, but in addition ensures that the read will be
  * performed against an appropriately committed snapshot if the operation is using a readConcern of
  * 'majority'.
@@ -144,14 +166,14 @@ protected:
  * NOTE: Must not be used with any locks held, because it needs to block waiting on the committed
  * snapshot to become available, and can potentially release and reacquire locks.
  */
-class AutoGetCollectionForRead : public AutoGetCollectionForReadBase<AutoGetCollection> {
+class AutoGetCollectionForRead
+    : public AutoGetCollectionForReadBase<AutoGetCollection, EmplaceAutoGetCollectionForRead> {
 public:
     AutoGetCollectionForRead(
         OperationContext* opCtx,
         const NamespaceStringOrUUID& nsOrUUID,
         AutoGetCollectionViewMode viewMode = AutoGetCollectionViewMode::kViewsForbidden,
-        Date_t deadline = Date_t::max())
-        : AutoGetCollectionForReadBase(opCtx, nsOrUUID, viewMode, deadline) {}
+        Date_t deadline = Date_t::max());
 
     Database* getDb() const {
         return _autoColl->getDb();
@@ -196,7 +218,28 @@ public:
     }
 
 private:
-    boost::optional<AutoGetCollectionForReadBase<AutoGetCollectionLockFree>>
+    /**
+     * Helper for how AutoGetCollectionForReadBase instantiates its owned AutoGetCollectionLockFree.
+     */
+    class EmplaceHelper {
+    public:
+        EmplaceHelper(OperationContext* opCtx,
+                      CollectionCatalogStasher& catalogStasher,
+                      const NamespaceStringOrUUID& nsOrUUID,
+                      AutoGetCollectionViewMode viewMode,
+                      Date_t deadline);
+
+        void emplace(boost::optional<AutoGetCollectionLockFree>& autoColl) const;
+
+    private:
+        OperationContext* _opCtx;
+        CollectionCatalogStasher& _catalogStasher;
+        const NamespaceStringOrUUID& _nsOrUUID;
+        AutoGetCollectionViewMode _viewMode;
+        Date_t _deadline;
+    };
+
+    boost::optional<AutoGetCollectionForReadBase<AutoGetCollectionLockFree, EmplaceHelper>>
         _autoGetCollectionForReadBase;
     CollectionCatalogStasher _catalogStash;
 };
diff --git a/src/mongo/db/query/plan_yield_policy_impl.cpp b/src/mongo/db/query/plan_yield_policy_impl.cpp
index e5492c2939d..c1fdd445472 100644
--- a/src/mongo/db/query/plan_yield_policy_impl.cpp
+++ b/src/mongo/db/query/plan_yield_policy_impl.cpp
@@ -36,7 +36,6 @@
 #include "mongo/db/curop_failpoint_helpers.h"
 #include "mongo/db/query/query_knobs_gen.h"
 #include "mongo/db/service_context.h"
-#include "mongo/db/storage/snapshot_helper.h"
 #include "mongo/util/fail_point.h"
 
 namespace mongo {
@@ -119,13 +118,12 @@ void PlanYieldPolicyImpl::_yieldAllLocks(OperationContext* opCtx,
         opCtx->checkForInterrupt();  // throws
     }
 
-    ON_BLOCK_EXIT([yieldable]() {
-        if (yieldable)
-            yieldable->restore();
-    });
-
     if (!unlocked) {
-        // Nothing was unlocked, just return, yielding is pointless.
+        // Nothing was unlocked, just return, yielding is pointless. Restore the yieldable before
+        // returning.
+        if (yieldable) {
+            yieldable->restore();
+        }
         return;
     }
 
@@ -144,15 +142,11 @@ void PlanYieldPolicyImpl::_yieldAllLocks(OperationContext* opCtx,
 
     locker->restoreLockState(opCtx, snapshot);
 
-    // After yielding and reacquiring locks, the preconditions that were used to select our
-    // ReadSource initially need to be checked again. Queries hold an AutoGetCollectionForRead RAII
-    // lock for their lifetime, which may select a ReadSource based on state (e.g. replication
-    // state). After a query yields its locks, this state may have changed, invalidating our current
-    // choice of ReadSource. Using the same preconditions, change our ReadSource if necessary.
-    auto newReadSource = SnapshotHelper::getNewReadSource(opCtx, planExecNS);
-    if (newReadSource) {
-        opCtx->recoveryUnit()->setTimestampReadSource(*newReadSource);
-    }
+    // If we get this far we should have a yieldable instance.
+    invariant(yieldable);
+
+    // Yieldable restore may set a new read source if necessary
+    yieldable->restore();
 }
 
 void PlanYieldPolicyImpl::preCheckInterruptOnly(OperationContext* opCtx) {
diff --git a/src/mongo/db/storage/SConscript b/src/mongo/db/storage/SConscript
index 2b3a2703cb4..e138a247664 100644
--- a/src/mongo/db/storage/SConscript
+++ b/src/mongo/db/storage/SConscript
@@ -38,8 +38,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/namespace_string',
         ],
     LIBDEPS_PRIVATE=[
-        '$BUILD_DIR/mongo/db/catalog/collection_catalog',
-        '$BUILD_DIR/mongo/db/concurrency/lock_manager',
+        '$BUILD_DIR/mongo/db/concurrency/lock_manager_defs',
         '$BUILD_DIR/mongo/db/repl/read_concern_args',
         '$BUILD_DIR/mongo/db/repl/repl_coordinator_interface',
         'recovery_unit_base',