SERVER-61127 Retry multi-writes that hit StaleConfig due to critical section on the shardr6.0.8-rc0r6.0.8

(cherry picked from commit 824b9b7e608687ba0db7af2d5ccc5b6811a46720)

6 files changed, 133 insertions, 36 deletions

diff --git a/src/mongo/db/exec/batched_delete_stage.cpp b/src/mongo/db/exec/batched_delete_stage.cpp
index 78bfd05e352..fb8be63cd8a 100644
--- a/src/mongo/db/exec/batched_delete_stage.cpp
+++ b/src/mongo/db/exec/batched_delete_stage.cpp
@@ -242,6 +242,16 @@ PlanStage::StageState BatchedDeleteStage::_deleteBatch(WorkingSetID* out) {
         wuow.commit();
     } catch (const WriteConflictException&) {
         return _prepareToRetryDrainAfterWCE(out, recordsThatNoLongerMatch);
+    } catch (const ExceptionFor<ErrorCodes::StaleConfig>& ex) {
+        if (ex->getVersionReceived() == ChunkVersion::IGNORED() && ex->getCriticalSectionSignal()) {
+            // If ChunkVersion is IGNORED and we encountered a critical section, then yield, wait
+            // for critical section to finish and then we'll resume the write from the point we had
+            // left. We do this to prevent large multi-writes from repeatedly failing due to
+            // StaleConfig and exhausting the mongos retry attempts.
+            planExecutorShardingCriticalSectionFuture(opCtx()) = ex->getCriticalSectionSignal();
+            return _prepareToRetryDrainAfterWCE(out, recordsThatNoLongerMatch);
+        }
+        throw;
     }
 
     incrementSSSMetricNoOverflow(batchedDeletesSSS.docs, docsDeleted);
diff --git a/src/mongo/db/exec/delete_stage.cpp b/src/mongo/db/exec/delete_stage.cpp
index 83941b91c4e..581ef2ef294 100644
--- a/src/mongo/db/exec/delete_stage.cpp
+++ b/src/mongo/db/exec/delete_stage.cpp
@@ -178,23 +178,38 @@ PlanStage::StageState DeleteStage::doWork(WorkingSetID* out) {
 
     bool writeToOrphan = false;
     if (!_params->isExplain && !_params->fromMigrate) {
-        const auto action = _preWriteFilter.computeAction(member->doc.value());
-        if (action == write_stage_common::PreWriteFilter::Action::kSkip) {
-            LOGV2_DEBUG(5983201,
-                        3,
-                        "Skipping delete operation to orphan document to prevent a wrong change "
-                        "stream event",
-                        "namespace"_attr = collection()->ns(),
-                        "record"_attr = member->doc.value());
-            return PlanStage::NEED_TIME;
-        } else if (action == write_stage_common::PreWriteFilter::Action::kWriteAsFromMigrate) {
-            LOGV2_DEBUG(6184700,
-                        3,
-                        "Marking delete operation to orphan document with the fromMigrate flag "
-                        "to prevent a wrong change stream event",
-                        "namespace"_attr = collection()->ns(),
-                        "record"_attr = member->doc.value());
-            writeToOrphan = true;
+        try {
+            const auto action = _preWriteFilter.computeAction(member->doc.value());
+            if (action == write_stage_common::PreWriteFilter::Action::kSkip) {
+                LOGV2_DEBUG(
+                    5983201,
+                    3,
+                    "Skipping delete operation to orphan document to prevent a wrong change "
+                    "stream event",
+                    "namespace"_attr = collection()->ns(),
+                    "record"_attr = member->doc.value());
+                return PlanStage::NEED_TIME;
+            } else if (action == write_stage_common::PreWriteFilter::Action::kWriteAsFromMigrate) {
+                LOGV2_DEBUG(6184700,
+                            3,
+                            "Marking delete operation to orphan document with the fromMigrate flag "
+                            "to prevent a wrong change stream event",
+                            "namespace"_attr = collection()->ns(),
+                            "record"_attr = member->doc.value());
+                writeToOrphan = true;
+            }
+        } catch (const ExceptionFor<ErrorCodes::StaleConfig>& ex) {
+            if (ex->getVersionReceived() == ChunkVersion::IGNORED() &&
+                ex->getCriticalSectionSignal()) {
+                // If ChunkVersion is IGNORED and we encountered a critical section, then yield,
+                // wait for the critical section to finish and then we'll resume the write from the
+                // point we had left. We do this to prevent large multi-writes from repeatedly
+                // failing due to StaleConfig and exhausting the mongos retry attempts.
+                planExecutorShardingCriticalSectionFuture(opCtx()) = ex->getCriticalSectionSignal();
+                memberFreer.dismiss();  // Keep this member around so we can retry deleting it.
+                return prepareToRetryWSM(id, out);
+            }
+            throw;
         }
     }
 
@@ -235,6 +250,18 @@ PlanStage::StageState DeleteStage::doWork(WorkingSetID* out) {
         } catch (const WriteConflictException&) {
             memberFreer.dismiss();  // Keep this member around so we can retry deleting it.
             return prepareToRetryWSM(id, out);
+        } catch (const ExceptionFor<ErrorCodes::StaleConfig>& ex) {
+            if (ex->getVersionReceived() == ChunkVersion::IGNORED() &&
+                ex->getCriticalSectionSignal()) {
+                // If ChunkVersion is IGNORED and we encountered a critical section, then yield,
+                // wait for the critical section to finish and then we'll resume the write from the
+                // point we had left. We do this to prevent large multi-writes from repeatedly
+                // failing due to StaleConfig and exhausting the mongos retry attempts.
+                planExecutorShardingCriticalSectionFuture(opCtx()) = ex->getCriticalSectionSignal();
+                memberFreer.dismiss();  // Keep this member around so we can retry deleting it.
+                return prepareToRetryWSM(id, out);
+            }
+            throw;
         }
     }
     _specificStats.docsDeleted += _params->numStatsForDoc ? _params->numStatsForDoc(bsonObjDoc) : 1;
diff --git a/src/mongo/db/exec/update_stage.cpp b/src/mongo/db/exec/update_stage.cpp
index 5a273dc2d89..0ed308dc79c 100644
--- a/src/mongo/db/exec/update_stage.cpp
+++ b/src/mongo/db/exec/update_stage.cpp
@@ -459,24 +459,41 @@ PlanStage::StageState UpdateStage::doWork(WorkingSetID* out) {
 
         bool writeToOrphan = false;
         if (!_params.request->explain() && _isUserInitiatedWrite) {
-            const auto action = _preWriteFilter.computeAction(member->doc.value());
-            if (action == write_stage_common::PreWriteFilter::Action::kSkip) {
-                LOGV2_DEBUG(
-                    5983200,
-                    3,
-                    "Skipping update operation to orphan document to prevent a wrong change "
-                    "stream event",
-                    "namespace"_attr = collection()->ns(),
-                    "record"_attr = member->doc.value());
-                return PlanStage::NEED_TIME;
-            } else if (action == write_stage_common::PreWriteFilter::Action::kWriteAsFromMigrate) {
-                LOGV2_DEBUG(6184701,
-                            3,
-                            "Marking update operation to orphan document with the fromMigrate flag "
-                            "to prevent a wrong change stream event",
-                            "namespace"_attr = collection()->ns(),
-                            "record"_attr = member->doc.value());
-                writeToOrphan = true;
+            try {
+                const auto action = _preWriteFilter.computeAction(member->doc.value());
+                if (action == write_stage_common::PreWriteFilter::Action::kSkip) {
+                    LOGV2_DEBUG(
+                        5983200,
+                        3,
+                        "Skipping update operation to orphan document to prevent a wrong change "
+                        "stream event",
+                        "namespace"_attr = collection()->ns(),
+                        "record"_attr = member->doc.value());
+                    return PlanStage::NEED_TIME;
+                } else if (action ==
+                           write_stage_common::PreWriteFilter::Action::kWriteAsFromMigrate) {
+                    LOGV2_DEBUG(
+                        6184701,
+                        3,
+                        "Marking update operation to orphan document with the fromMigrate flag "
+                        "to prevent a wrong change stream event",
+                        "namespace"_attr = collection()->ns(),
+                        "record"_attr = member->doc.value());
+                    writeToOrphan = true;
+                }
+            } catch (const ExceptionFor<ErrorCodes::StaleConfig>& ex) {
+                if (ex->getVersionReceived() == ChunkVersion::IGNORED() &&
+                    ex->getCriticalSectionSignal()) {
+                    // If ChunkVersion is IGNORED and we encountered a critical section, then yield,
+                    // wait for critical section to finish and then we'll resume the write from the
+                    // point we had left. We do this to prevent large multi-writes from repeatedly
+                    // failing due to StaleConfig and exhausting the mongos retry attempts.
+                    planExecutorShardingCriticalSectionFuture(opCtx()) =
+                        ex->getCriticalSectionSignal();
+                    memberFreer.dismiss();  // Keep this member around so we can retry deleting it.
+                    return prepareToRetryWSM(id, out);
+                }
+                throw;
             }
         }
 
@@ -506,6 +523,18 @@ PlanStage::StageState UpdateStage::doWork(WorkingSetID* out) {
         } catch (const WriteConflictException&) {
             memberFreer.dismiss();  // Keep this member around so we can retry updating it.
             return prepareToRetryWSM(id, out);
+        } catch (const ExceptionFor<ErrorCodes::StaleConfig>& ex) {
+            if (ex->getVersionReceived() == ChunkVersion::IGNORED() &&
+                ex->getCriticalSectionSignal()) {
+                // If ChunkVersion is IGNORED and we encountered a critical section, then yield,
+                // wait for critical section to finish and then we'll resume the write from the
+                // point we had left. We do this to prevent large multi-writes from repeatedly
+                // failing due to StaleConfig and exhausting the mongos retry attempts.
+                planExecutorShardingCriticalSectionFuture(opCtx()) = ex->getCriticalSectionSignal();
+                memberFreer.dismiss();  // Keep this member around so we can retry updating it.
+                return prepareToRetryWSM(id, out);
+            }
+            throw;
         }
 
         // Set member's obj to be the doc we want to return.
diff --git a/src/mongo/db/query/plan_executor.cpp b/src/mongo/db/query/plan_executor.cpp
index ee41d15d84c..99b2fd8fefa 100644
--- a/src/mongo/db/query/plan_executor.cpp
+++ b/src/mongo/db/query/plan_executor.cpp
@@ -38,6 +38,10 @@ namespace {
 MONGO_FAIL_POINT_DEFINE(planExecutorAlwaysFails);
 }  // namespace
 
+const OperationContext::Decoration<boost::optional<SharedSemiFuture<void>>>
+    planExecutorShardingCriticalSectionFuture =
+        OperationContext::declareDecoration<boost::optional<SharedSemiFuture<void>>>();
+
 std::string PlanExecutor::stateToStr(ExecState execState) {
     switch (execState) {
         case PlanExecutor::ADVANCED:
diff --git a/src/mongo/db/query/plan_executor.h b/src/mongo/db/query/plan_executor.h
index 33fbd075b93..069fb4b4608 100644
--- a/src/mongo/db/query/plan_executor.h
+++ b/src/mongo/db/query/plan_executor.h
@@ -56,6 +56,15 @@ class RecordId;
 extern const OperationContext::Decoration<repl::OpTime> clientsLastKnownCommittedOpTime;
 
 /**
+ * If a plan yielded because it encountered a sharding critical section,
+ * 'planExecutorShardingCriticalSectionFuture' will be set to a future that becomes ready when the
+ * critical section ends. This future can be waited on to hold off resuming the plan execution while
+ * the critical section is still active.
+ */
+extern const OperationContext::Decoration<boost::optional<SharedSemiFuture<void>>>
+    planExecutorShardingCriticalSectionFuture;
+
+/**
  * A PlanExecutor is the abstraction that knows how to crank a tree of stages into execution.
  * The executor is usually part of a larger abstraction that is interacting with the cache
  * and/or the query optimizer.
diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp
index c3ae2946a38..55253f6e02d 100644
--- a/src/mongo/db/query/plan_executor_impl.cpp
+++ b/src/mongo/db/query/plan_executor_impl.cpp
@@ -62,6 +62,7 @@
 #include "mongo/db/query/plan_yield_policy_impl.h"
 #include "mongo/db/query/yield_policy_callbacks_impl.h"
 #include "mongo/db/repl/replication_coordinator.h"
+#include "mongo/db/s/operation_sharding_state.h"
 #include "mongo/db/service_context.h"
 #include "mongo/logv2/log.h"
 #include "mongo/util/fail_point.h"
@@ -360,8 +361,25 @@ PlanExecutor::ExecState PlanExecutorImpl::_getNextImpl(Snapshotted<Document>* ob
         //   2) some stage requested a yield, or
         //   3) we need to yield and retry due to a WriteConflictException.
         // In all cases, the actual yielding happens here.
+
+        const auto whileYieldingFn = [&]() {
+            // If we yielded because we encountered a sharding critical section, wait for the
+            // critical section to end before continuing. By waiting for the critical section to be
+            // exited we avoid busy spinning immediately and encountering the same critical section
+            // again. It is important that this wait happens after having released the lock
+            // hierarchy -- otherwise deadlocks could happen, or the very least, locks would be
+            // unnecessarily held while waiting.
+            const auto& shardingCriticalSection = planExecutorShardingCriticalSectionFuture(_opCtx);
+            if (shardingCriticalSection) {
+                OperationShardingState::waitForCriticalSectionToComplete(_opCtx,
+                                                                         *shardingCriticalSection)
+                    .ignore();
+                planExecutorShardingCriticalSectionFuture(_opCtx).reset();
+            }
+        };
+
         if (_yieldPolicy->shouldYieldOrInterrupt(_opCtx)) {
-            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(_opCtx));
+            uassertStatusOK(_yieldPolicy->yieldOrInterrupt(_opCtx, whileYieldingFn));
         }
 
         WorkingSetID id = WorkingSet::INVALID_ID;