SERVER-44529 Query yield recovery after a stepdown should switch to read at the no-overlap time

After yielding and reacquiring locks in a query, 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.

2 files changed, 18 insertions, 0 deletions

diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp
index 290a7e924c8..f642f5be99b 100644
--- a/src/mongo/db/query/plan_executor_impl.cpp
+++ b/src/mongo/db/query/plan_executor_impl.cpp
@@ -80,6 +80,7 @@ namespace {
 
 MONGO_FAIL_POINT_DEFINE(planExecutorAlwaysFails);
 MONGO_FAIL_POINT_DEFINE(planExecutorHangBeforeShouldWaitForInserts);
+MONGO_FAIL_POINT_DEFINE(planExecutorHangWhileYieldedInWaitForInserts);
 
 /**
  * Constructs a PlanYieldPolicy based on 'policy'.
@@ -480,6 +481,12 @@ PlanExecutor::ExecState PlanExecutorImpl::_waitForInserts(CappedInsertNotifierDa
     auto yieldResult = _yieldPolicy->yieldOrInterrupt([opCtx, notifierData] {
         const auto deadline = awaitDataState(opCtx).waitForInsertsDeadline;
         notifierData->notifier->waitUntil(notifierData->lastEOFVersion, deadline);
+        if (MONGO_unlikely(planExecutorHangWhileYieldedInWaitForInserts.shouldFail())) {
+            LOGV2(4452903,
+                  "PlanExecutor - planExecutorHangWhileYieldedInWaitForInserts fail point "
+                  "enabled. Blocking until fail point is disabled");
+            planExecutorHangWhileYieldedInWaitForInserts.pauseWhileSet();
+        }
     });
     notifierData->lastEOFVersion = currentNotifierVersion;
 
diff --git a/src/mongo/db/query/plan_yield_policy.cpp b/src/mongo/db/query/plan_yield_policy.cpp
index 856d15c8b4c..448ac827797 100644
--- a/src/mongo/db/query/plan_yield_policy.cpp
+++ b/src/mongo/db/query/plan_yield_policy.cpp
@@ -37,6 +37,7 @@
 #include "mongo/db/operation_context.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"
 #include "mongo/util/scopeguard.h"
 #include "mongo/util/time_support.h"
@@ -208,6 +209,16 @@ void PlanYieldPolicy::_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);
+    }
 }
 
 }  // namespace mongo