SERVER-81812 BlockToRowStage will check for interrupt but won't yield

Crescentek · Oct 12, 2023 · 5bbb8a5 · 5bbb8a5
1 parent 9d8ecca
commit 5bbb8a5
Show file tree

Hide file tree

Showing 6 changed files with 20 additions and 7 deletions.
diff --git a/src/mongo/db/exec/sbe/stages/co_scan.cpp b/src/mongo/db/exec/sbe/stages/co_scan.cpp
@@ -55,7 +55,7 @@ void CoScanStage::open(bool reOpen) {
 PlanState CoScanStage::getNext() {
     auto optTimer(getOptTimer(_opCtx));
 
-    checkForInterrupt(_opCtx);
+    checkForInterruptAndYield(_opCtx);
 
     // Run forever.
     return trackPlanState(PlanState::ADVANCED);

diff --git a/src/mongo/db/exec/sbe/stages/column_scan.cpp b/src/mongo/db/exec/sbe/stages/column_scan.cpp
@@ -559,7 +559,7 @@ PlanState ColumnScanStage::getNext() {
     // call.
     disableSlotAccess();
 
-    checkForInterrupt(_opCtx);
+    checkForInterruptAndYield(_opCtx);
 
     if (_scanTracker.isScanningRowstore()) {
         _scanTracker.track();

diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp
@@ -298,7 +298,7 @@ PlanState IndexScanStageBase::getNext() {
     // state in case it yields as the state will be completely overwritten after the call.
     disableSlotAccess();
 
-    checkForInterrupt(_opCtx);
+    checkForInterruptAndYield(_opCtx);
 
     do {
         switch (_scanState) {

diff --git a/src/mongo/db/exec/sbe/stages/scan.cpp b/src/mongo/db/exec/sbe/stages/scan.cpp
@@ -457,7 +457,7 @@ PlanState ScanStage::getNext() {
     // PlanStage tree if a yield occurs. It's important that we call checkForInterrupt() before
     // checking '_needsToCheckCappedPositionLost' since a call to restoreState() may set
     // '_needsToCheckCappedPositionLost'.
-    checkForInterrupt(_opCtx);
+    checkForInterruptAndYield(_opCtx);
 
     if (_needsToCheckCappedPositionLost) {
         _cursor->save();
@@ -1097,7 +1097,7 @@ PlanState ParallelScanStage::getNext() {
         return trackPlanState(PlanState::IS_EOF);
     }
 
-    checkForInterrupt(_opCtx);
+    checkForInterruptAndYield(_opCtx);
 
     boost::optional<Record> nextRecord;
 

diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h
@@ -505,7 +505,7 @@ class CanInterrupt {
      * Checks for interrupt if necessary. If yielding has been enabled for this object, then also
      * performs a yield if necessary.
      */
-    void checkForInterrupt(OperationContext* opCtx) {
+    void checkForInterruptAndYield(OperationContext* opCtx) {
         invariant(opCtx);
 
         if (!_yieldPolicy) {
@@ -521,6 +521,19 @@ class CanInterrupt {
         }
     }
 
+    /**
+     * Checks for interrupt if necessary. Will never yield regardless of the yielding policy.
+     * Should only be used for ValueBlock stages.
+     */
+    void checkForInterrupt(OperationContext* opCtx) {
+        invariant(opCtx);
+
+        if (--_interruptCounter == 0) {
+            _interruptCounter = kInterruptCheckPeriod;
+            opCtx->checkForInterrupt();
+        }
+    }
+
     void attachNewYieldPolicy(PlanYieldPolicy* yieldPolicy) {
         auto stage = static_cast<T*>(this);
         for (auto&& child : stage->_children) {

diff --git a/src/mongo/db/exec/sbe/stages/window.cpp b/src/mongo/db/exec/sbe/stages/window.cpp
@@ -543,7 +543,7 @@ void WindowStage::open(bool reOpen) {
 
 PlanState WindowStage::getNext() {
     auto optTimer(getOptTimer(_opCtx));
-    checkForInterrupt(_opCtx);
+    checkForInterruptAndYield(_opCtx);
 
     // Fetch at least the current document into cache.
     _currId++;