SERVER-53299 Fix so AutoGetCollectionForReadLockFree only read replication state once when setting up read source.

4 files changed, 32 insertions, 21 deletions

diff --git a/src/mongo/db/catalog/collection_catalog.cpp b/src/mongo/db/catalog/collection_catalog.cpp
index 7ca30a9d778..3f9d434a3d7 100644
--- a/src/mongo/db/catalog/collection_catalog.cpp
+++ b/src/mongo/db/catalog/collection_catalog.cpp
@@ -1030,7 +1030,7 @@ const Collection* LookupCollectionForYieldRestore::operator()(OperationContext*
     // 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());
+    auto [newReadSource, _] = SnapshotHelper::shouldChangeReadSource(opCtx, collection->ns());
     if (newReadSource) {
         opCtx->recoveryUnit()->setTimestampReadSource(*newReadSource);
     }
diff --git a/src/mongo/db/db_raii.cpp b/src/mongo/db/db_raii.cpp
index 307e97994b9..1da574a91c1 100644
--- a/src/mongo/db/db_raii.cpp
+++ b/src/mongo/db/db_raii.cpp
@@ -219,7 +219,9 @@ AutoGetCollectionForReadBase<AutoGetCollectionType, EmplaceAutoCollFunc>::
 
         // Once we have our locks, check whether or not we should override the ReadSource that was
         // set before acquiring locks.
-        if (auto newReadSource = SnapshotHelper::getNewReadSource(opCtx, nss)) {
+        auto [newReadSource, shouldReadAtLastApplied] =
+            SnapshotHelper::shouldChangeReadSource(opCtx, nss);
+        if (newReadSource) {
             opCtx->recoveryUnit()->setTimestampReadSource(*newReadSource);
             readSource = *newReadSource;
         }
@@ -257,8 +259,7 @@ AutoGetCollectionForReadBase<AutoGetCollectionType, EmplaceAutoCollFunc>::
         // serially in oplog application, and therefore can be safely read without taking the PBWM
         // lock or reading at a timestamp.
         if (readSource == RecoveryUnit::ReadSource::kNoTimestamp && callerWasConflicting &&
-            !nss.mustBeAppliedInOwnOplogBatch() &&
-            SnapshotHelper::shouldReadAtLastApplied(opCtx, nss)) {
+            !nss.mustBeAppliedInOwnOplogBatch() && shouldReadAtLastApplied) {
             LOGV2_FATAL(4728700,
                         "Reading from replicated collection on a secondary without read timestamp "
                         "or PBWM lock",
@@ -391,7 +392,8 @@ void AutoGetCollectionForReadLockFree::EmplaceHelper::emplace(
                     // 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());
+                        auto [newReadSource, _] =
+                            SnapshotHelper::shouldChangeReadSource(opCtx, coll->ns());
                         if (newReadSource) {
                             opCtx->recoveryUnit()->setTimestampReadSource(*newReadSource);
                         }
diff --git a/src/mongo/db/storage/snapshot_helper.cpp b/src/mongo/db/storage/snapshot_helper.cpp
index 84af208d391..03f752b6e2d 100644
--- a/src/mongo/db/storage/snapshot_helper.cpp
+++ b/src/mongo/db/storage/snapshot_helper.cpp
@@ -56,9 +56,7 @@ bool canReadAtLastApplied(OperationContext* opCtx) {
     }
     return false;
 }
-}  // namespace
 
-namespace SnapshotHelper {
 bool shouldReadAtLastApplied(OperationContext* opCtx,
                              const NamespaceString& nss,
                              std::string* reason) {
@@ -120,15 +118,19 @@ bool shouldReadAtLastApplied(OperationContext* opCtx,
 
     return true;
 }
-boost::optional<RecoveryUnit::ReadSource> getNewReadSource(OperationContext* opCtx,
-                                                           const NamespaceString& nss) {
+}  // namespace
+
+namespace SnapshotHelper {
+
+ReadSourceChange shouldChangeReadSource(OperationContext* opCtx, const NamespaceString& nss) {
+    std::string reason;
+    const bool readAtLastApplied = shouldReadAtLastApplied(opCtx, nss, &reason);
+
     if (!canReadAtLastApplied(opCtx)) {
-        return boost::none;
+        return {boost::none, readAtLastApplied};
     }
 
     const auto existing = opCtx->recoveryUnit()->getTimestampReadSource();
-    std::string reason;
-    const bool readAtLastApplied = shouldReadAtLastApplied(opCtx, nss, &reason);
     if (existing == RecoveryUnit::ReadSource::kNoTimestamp) {
         // Shifting from reading without a timestamp to reading with a timestamp can be dangerous
         // because writes will appear to vanish. This case is intended for new reads on secondaries
@@ -138,7 +140,7 @@ boost::optional<RecoveryUnit::ReadSource> getNewReadSource(OperationContext* opC
         // at the lastApplied point because reading without a timestamp is not safe.
         if (readAtLastApplied) {
             LOGV2_DEBUG(4452901, 2, "Changing ReadSource to kLastApplied", logAttrs(nss));
-            return RecoveryUnit::ReadSource::kLastApplied;
+            return {RecoveryUnit::ReadSource::kLastApplied, readAtLastApplied};
         }
     } else if (existing == RecoveryUnit::ReadSource::kLastApplied) {
         // For some reason, we can no longer read at lastApplied.
@@ -156,10 +158,10 @@ boost::optional<RecoveryUnit::ReadSource> getNewReadSource(OperationContext* opC
             // manipulate their ReadSources after performing reads at an un-timetamped snapshot. The
             // only exception is callers of this function that may need to change from kNoTimestamp
             // to kLastApplied in the event of a catalog conflict or query yield.
-            return RecoveryUnit::ReadSource::kNoTimestamp;
+            return {RecoveryUnit::ReadSource::kNoTimestamp, readAtLastApplied};
         }
     }
-    return boost::none;
+    return {boost::none, readAtLastApplied};
 }
 
 bool collectionChangesConflictWithRead(boost::optional<Timestamp> collectionMin,
diff --git a/src/mongo/db/storage/snapshot_helper.h b/src/mongo/db/storage/snapshot_helper.h
index c24dfd16d8c..cc62cbec889 100644
--- a/src/mongo/db/storage/snapshot_helper.h
+++ b/src/mongo/db/storage/snapshot_helper.h
@@ -33,13 +33,20 @@
 
 namespace mongo {
 namespace SnapshotHelper {
-// Returns a ReadSource if we should change our current ReadSource. Returns boost::none otherwise.
-boost::optional<RecoveryUnit::ReadSource> getNewReadSource(OperationContext* opCtx,
-                                                           const NamespaceString& nss);
+struct ReadSourceChange {
+    boost::optional<RecoveryUnit::ReadSource> newReadSource;
+    bool shouldReadAtLastApplied;
+};
 
-bool shouldReadAtLastApplied(OperationContext* opCtx,
-                             const NamespaceString& nss,
-                             std::string* reason = nullptr);
+/**
+ * Returns a ReadSourceChange containing data necessary to decide if we need to change read source.
+ *
+ * For Lock-Free Reads, the decisions made within this function based on replication state may
+ * become invalid after it returns and multiple calls may yield different answers. Higher level code
+ * must validate the relevance of the outcome based on replication state before and after calling
+ * this function.
+ */
+ReadSourceChange shouldChangeReadSource(OperationContext* opCtx, const NamespaceString& nss);
 
 bool collectionChangesConflictWithRead(boost::optional<Timestamp> collectionMin,
                                        boost::optional<Timestamp> readTimestamp);