SERVER-60023 Make SBE cursor repositioning and copying optional

21 files changed, 122 insertions, 80 deletions

diff --git a/src/mongo/db/exec/sbe/stages/check_bounds.cpp b/src/mongo/db/exec/sbe/stages/check_bounds.cpp
index e5129b63fe1..d8d7f6612a9 100644
--- a/src/mongo/db/exec/sbe/stages/check_bounds.cpp
+++ b/src/mongo/db/exec/sbe/stages/check_bounds.cpp
@@ -188,8 +188,8 @@ size_t CheckBoundsStage::estimateCompileTimeSize() const {
     return size;
 }
 
-void CheckBoundsStage::doSaveState() {
-    if (!slotsAccessible()) {
+void CheckBoundsStage::doSaveState(bool relinquishCursor) {
+    if (!slotsAccessible() || !relinquishCursor) {
         return;
     }
 
diff --git a/src/mongo/db/exec/sbe/stages/check_bounds.h b/src/mongo/db/exec/sbe/stages/check_bounds.h
index 5201a41e2bb..c0399439ea7 100644
--- a/src/mongo/db/exec/sbe/stages/check_bounds.h
+++ b/src/mongo/db/exec/sbe/stages/check_bounds.h
@@ -86,7 +86,7 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() final;
+    void doSaveState(bool relinquishCursor) final;
 
 private:
     const CheckBoundsParams _params;
diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp
index cbe7e5867dd..10060aa033f 100644
--- a/src/mongo/db/exec/sbe/stages/ix_scan.cpp
+++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp
@@ -155,32 +155,34 @@ value::SlotAccessor* IndexScanStage::getAccessor(CompileCtx& ctx, value::SlotId
     return ctx.getAccessor(slot);
 }
 
-void IndexScanStage::doSaveState() {
-    if (slotsAccessible()) {
-        if (_recordAccessor) {
-            _recordAccessor->makeOwned();
-        }
-        if (_recordIdAccessor) {
-            _recordIdAccessor->makeOwned();
-        }
-        for (auto& accessor : _accessors) {
-            accessor.makeOwned();
+void IndexScanStage::doSaveState(bool relinquishCursor) {
+    if (relinquishCursor) {
+        if (slotsAccessible()) {
+            if (_recordAccessor) {
+                _recordAccessor->makeOwned();
+            }
+            if (_recordIdAccessor) {
+                _recordIdAccessor->makeOwned();
+            }
+            for (auto& accessor : _accessors) {
+                accessor.makeOwned();
+            }
         }
-    }
 
-    // Seek points are external to the index scan and must be accessible no matter what as long as
-    // the index scan is opened.
-    if (_open) {
-        if (_seekKeyLowHolder) {
-            _seekKeyLowHolder->makeOwned();
-        }
-        if (_seekKeyHighHolder) {
-            _seekKeyHighHolder->makeOwned();
+        // Seek points are external to the index scan and must be accessible no matter what as long
+        // as the index scan is opened.
+        if (_open) {
+            if (_seekKeyLowHolder) {
+                _seekKeyLowHolder->makeOwned();
+            }
+            if (_seekKeyHighHolder) {
+                _seekKeyHighHolder->makeOwned();
+            }
         }
-    }
 
-    if (_cursor) {
-        _cursor->save();
+        if (_cursor) {
+            _cursor->save();
+        }
     }
 
     _coll.reset();
@@ -196,7 +198,7 @@ void IndexScanStage::restoreCollectionAndIndex() {
             indexCatalogEntry && !indexCatalogEntry->isDropped());
 }
 
-void IndexScanStage::doRestoreState() {
+void IndexScanStage::doRestoreState(bool relinquishCursor) {
     invariant(_opCtx);
     invariant(!_coll);
 
@@ -206,7 +208,7 @@ void IndexScanStage::doRestoreState() {
     }
     restoreCollectionAndIndex();
 
-    if (_cursor) {
+    if (_cursor && relinquishCursor) {
         _cursor->restore();
     }
 
diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.h b/src/mongo/db/exec/sbe/stages/ix_scan.h
index 72c81bde98c..21c5afcc9c0 100644
--- a/src/mongo/db/exec/sbe/stages/ix_scan.h
+++ b/src/mongo/db/exec/sbe/stages/ix_scan.h
@@ -99,8 +99,8 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() override;
-    void doRestoreState() override;
+    void doSaveState(bool relinquishCursor) override;
+    void doRestoreState(bool relinquishCursor) override;
     void doDetachFromOperationContext() override;
     void doAttachToOperationContext(OperationContext* opCtx) override;
     void doDetachFromTrialRunTracker() override;
diff --git a/src/mongo/db/exec/sbe/stages/loop_join.cpp b/src/mongo/db/exec/sbe/stages/loop_join.cpp
index 54db378bdde..9f26e0b84d0 100644
--- a/src/mongo/db/exec/sbe/stages/loop_join.cpp
+++ b/src/mongo/db/exec/sbe/stages/loop_join.cpp
@@ -162,7 +162,7 @@ void LoopJoinStage::close() {
     _children[0]->close();
 }
 
-void LoopJoinStage::doSaveState() {
+void LoopJoinStage::doSaveState(bool relinquishCursor) {
     if (_isReadingLeftSide || _outerGetNext) {
         // If we yield while reading the left side, there is no need to makeOwned() data held in
         // the right side, since we will have to re-open it anyway.
diff --git a/src/mongo/db/exec/sbe/stages/loop_join.h b/src/mongo/db/exec/sbe/stages/loop_join.h
index b335f880b24..22e3a93ffa2 100644
--- a/src/mongo/db/exec/sbe/stages/loop_join.h
+++ b/src/mongo/db/exec/sbe/stages/loop_join.h
@@ -69,7 +69,7 @@ public:
     void open(bool reOpen) final;
     PlanState getNext() final;
     void close() final;
-    void doSaveState() final;
+    void doSaveState(bool relinquishCursor) final;
 
     std::unique_ptr<PlanStageStats> getStats(bool includeDebugInfo) const final;
     const SpecificStats* getSpecificStats() const final;
diff --git a/src/mongo/db/exec/sbe/stages/makeobj.cpp b/src/mongo/db/exec/sbe/stages/makeobj.cpp
index 36c022544b1..54e1ce05e89 100644
--- a/src/mongo/db/exec/sbe/stages/makeobj.cpp
+++ b/src/mongo/db/exec/sbe/stages/makeobj.cpp
@@ -420,8 +420,8 @@ size_t MakeObjStageBase<O>::estimateCompileTimeSize() const {
 }
 
 template <MakeObjOutputType O>
-void MakeObjStageBase<O>::doSaveState() {
-    if (!slotsAccessible()) {
+void MakeObjStageBase<O>::doSaveState(bool relinquishCursor) {
+    if (!slotsAccessible() || !relinquishCursor) {
         return;
     }
 
diff --git a/src/mongo/db/exec/sbe/stages/makeobj.h b/src/mongo/db/exec/sbe/stages/makeobj.h
index 8fd81185e08..1cf0755f1c5 100644
--- a/src/mongo/db/exec/sbe/stages/makeobj.h
+++ b/src/mongo/db/exec/sbe/stages/makeobj.h
@@ -103,7 +103,7 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() final;
+    void doSaveState(bool relinquishCursor) final;
 
 private:
     void projectField(value::Object* obj, size_t idx);
diff --git a/src/mongo/db/exec/sbe/stages/merge_join.cpp b/src/mongo/db/exec/sbe/stages/merge_join.cpp
index 37991cded4e..dec5080fadf 100644
--- a/src/mongo/db/exec/sbe/stages/merge_join.cpp
+++ b/src/mongo/db/exec/sbe/stages/merge_join.cpp
@@ -322,8 +322,8 @@ void MergeJoinStage::close() {
     _outerProjectsBuffer.clear();
 }
 
-void MergeJoinStage::doSaveState() {
-    if (!slotsAccessible()) {
+void MergeJoinStage::doSaveState(bool relinquishCursor) {
+    if (!slotsAccessible() || !relinquishCursor) {
         return;
     }
 
diff --git a/src/mongo/db/exec/sbe/stages/merge_join.h b/src/mongo/db/exec/sbe/stages/merge_join.h
index 42425394bc0..b0f61cd677c 100644
--- a/src/mongo/db/exec/sbe/stages/merge_join.h
+++ b/src/mongo/db/exec/sbe/stages/merge_join.h
@@ -78,7 +78,7 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() final;
+    void doSaveState(bool relinquishCursor) final;
 
 private:
     using MergeJoinBuffer = std::vector<value::MaterializedRow>;
diff --git a/src/mongo/db/exec/sbe/stages/project.cpp b/src/mongo/db/exec/sbe/stages/project.cpp
index ef6c6739ade..bcc51d5b7b7 100644
--- a/src/mongo/db/exec/sbe/stages/project.cpp
+++ b/src/mongo/db/exec/sbe/stages/project.cpp
@@ -154,8 +154,8 @@ size_t ProjectStage::estimateCompileTimeSize() const {
     return size;
 }
 
-void ProjectStage::doSaveState() {
-    if (!slotsAccessible()) {
+void ProjectStage::doSaveState(bool relinquishCursor) {
+    if (!slotsAccessible() || !relinquishCursor) {
         return;
     }
 
diff --git a/src/mongo/db/exec/sbe/stages/project.h b/src/mongo/db/exec/sbe/stages/project.h
index 24033b55aca..1754dd7d2a9 100644
--- a/src/mongo/db/exec/sbe/stages/project.h
+++ b/src/mongo/db/exec/sbe/stages/project.h
@@ -63,7 +63,7 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() final;
+    void doSaveState(bool relinquishCursor) final;
 
 private:
     const value::SlotMap<std::unique_ptr<EExpression>> _projects;
diff --git a/src/mongo/db/exec/sbe/stages/scan.cpp b/src/mongo/db/exec/sbe/stages/scan.cpp
index 1e1cc4b3744..d7e0d714323 100644
--- a/src/mongo/db/exec/sbe/stages/scan.cpp
+++ b/src/mongo/db/exec/sbe/stages/scan.cpp
@@ -161,27 +161,29 @@ value::SlotAccessor* ScanStage::getAccessor(CompileCtx& ctx, value::SlotId slot)
     return ctx.getAccessor(slot);
 }
 
-void ScanStage::doSaveState() {
+void ScanStage::doSaveState(bool fullSave) {
     if (slotsAccessible()) {
-        if (_recordAccessor) {
-            _recordAccessor->makeOwned();
-        }
-        if (_recordIdAccessor) {
-            _recordIdAccessor->makeOwned();
-        }
-        for (auto& [fieldName, accessor] : _fieldAccessors) {
-            accessor->makeOwned();
+        if (fullSave) {
+            if (_recordAccessor) {
+                _recordAccessor->makeOwned();
+            }
+            if (_recordIdAccessor) {
+                _recordIdAccessor->makeOwned();
+            }
+            for (auto& [fieldName, accessor] : _fieldAccessors) {
+                accessor->makeOwned();
+            }
         }
     }
 
-    if (_cursor) {
+    if (_cursor && fullSave) {
         _cursor->save();
     }
 
     _coll.reset();
 }
 
-void ScanStage::doRestoreState() {
+void ScanStage::doRestoreState(bool fullSave) {
     invariant(_opCtx);
     invariant(!_coll);
 
@@ -193,7 +195,7 @@ void ScanStage::doRestoreState() {
     tassert(5777408, "Catalog epoch should be initialized", _catalogEpoch);
     _coll = restoreCollection(_opCtx, *_collName, _collUuid, *_catalogEpoch);
 
-    if (_cursor) {
+    if (_cursor && fullSave) {
         const bool couldRestore = _cursor->restore();
         uassert(ErrorCodes::CappedPositionLost,
                 str::stream()
@@ -617,7 +619,7 @@ value::SlotAccessor* ParallelScanStage::getAccessor(CompileCtx& ctx, value::Slot
     return ctx.getAccessor(slot);
 }
 
-void ParallelScanStage::doSaveState() {
+void ParallelScanStage::doSaveState(bool fullSave) {
     if (slotsAccessible()) {
         if (_recordAccessor) {
             _recordAccessor->makeOwned();
@@ -637,7 +639,7 @@ void ParallelScanStage::doSaveState() {
     _coll.reset();
 }
 
-void ParallelScanStage::doRestoreState() {
+void ParallelScanStage::doRestoreState(bool fullSave) {
     invariant(_opCtx);
     invariant(!_coll);
 
@@ -649,7 +651,7 @@ void ParallelScanStage::doRestoreState() {
     tassert(5777409, "Catalog epoch should be initialized", _catalogEpoch);
     _coll = restoreCollection(_opCtx, *_collName, _collUuid, *_catalogEpoch);
 
-    if (_cursor) {
+    if (_cursor && fullSave) {
         const bool couldRestore = _cursor->restore();
         uassert(ErrorCodes::CappedPositionLost,
                 str::stream()
diff --git a/src/mongo/db/exec/sbe/stages/scan.h b/src/mongo/db/exec/sbe/stages/scan.h
index 23d8370ce47..8787afb58b6 100644
--- a/src/mongo/db/exec/sbe/stages/scan.h
+++ b/src/mongo/db/exec/sbe/stages/scan.h
@@ -122,8 +122,8 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() override;
-    void doRestoreState() override;
+    void doSaveState(bool fullSave) override;
+    void doRestoreState(bool fullSave) override;
     void doDetachFromOperationContext() override;
     void doAttachToOperationContext(OperationContext* opCtx) override;
     void doDetachFromTrialRunTracker() override;
@@ -232,8 +232,8 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() final;
-    void doRestoreState() final;
+    void doSaveState(bool fullSave) final;
+    void doRestoreState(bool fullSave) final;
     void doDetachFromOperationContext() final;
     void doAttachToOperationContext(OperationContext* opCtx) final;
 
diff --git a/src/mongo/db/exec/sbe/stages/spool.cpp b/src/mongo/db/exec/sbe/stages/spool.cpp
index a7841ca6714..bf8a94e7c58 100644
--- a/src/mongo/db/exec/sbe/stages/spool.cpp
+++ b/src/mongo/db/exec/sbe/stages/spool.cpp
@@ -273,8 +273,8 @@ PlanState SpoolLazyProducerStage::getNext() {
     return trackPlanState(state);
 }
 
-void SpoolLazyProducerStage::doSaveState() {
-    if (!slotsAccessible()) {
+void SpoolLazyProducerStage::doSaveState(bool relinquishCursor) {
+    if (!slotsAccessible() || !relinquishCursor) {
         return;
     }
 
diff --git a/src/mongo/db/exec/sbe/stages/spool.h b/src/mongo/db/exec/sbe/stages/spool.h
index 197ad48d3ba..a2dd6f81657 100644
--- a/src/mongo/db/exec/sbe/stages/spool.h
+++ b/src/mongo/db/exec/sbe/stages/spool.h
@@ -125,7 +125,7 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() final;
+    void doSaveState(bool relinquishCursor) final;
 
 private:
     std::shared_ptr<SpoolBuffer> _buffer{nullptr};
diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h
index 5054618ca1f..aa97caf938e 100644
--- a/src/mongo/db/exec/sbe/stages/stages.h
+++ b/src/mongo/db/exec/sbe/stages/stages.h
@@ -29,6 +29,7 @@
 
 #pragma once
 
+#include "mongo/config.h"
 #include "mongo/db/exec/sbe/stages/plan_stats.h"
 #include "mongo/db/exec/sbe/util/debug_print.h"
 #include "mongo/db/exec/sbe/values/slot.h"
@@ -124,18 +125,30 @@ public:
      * before the first call to open(), before execution of the plan has begun.
      *
      * Propagates to all children, then calls doSaveState().
+     *
+     * The 'relinquishCursor' parameter indicates whether cursors should be reset and all data
+     * should be copied.
+     *
+     * TODO SERVER-59620: Remove the 'relinquishCursor' parameter once all callers pass 'false'.
      */
-    void saveState() {
+    void saveState(bool relinquishCursor) {
         auto stage = static_cast<T*>(this);
         stage->_commonStats.yields++;
+#if defined(MONGO_CONFIG_DEBUG_BUILD)
+        invariant(_saveState == SaveState::kNotSaved);
+#endif
 
-        stage->doSaveState();
+        stage->doSaveState(relinquishCursor);
         // Save the children in a right to left order so dependent stages (i.e. one using correlated
         // slots) are saved first.
         auto& children = stage->_children;
         for (auto it = children.rbegin(); it != children.rend(); it++) {
-            (*it)->saveState();
+            (*it)->saveState(relinquishCursor);
         }
+
+#if defined(MONGO_CONFIG_DEBUG_BUILD)
+        _saveState = relinquishCursor ? SaveState::kSavedFull : SaveState::kSavedNotFull;
+#endif
     }
 
     /**
@@ -148,16 +161,41 @@ public:
      * Throws a UserException on failure to restore due to a conflicting event such as a
      * collection drop. May throw a WriteConflictException, in which case the caller may choose to
      * retry.
+     *
+     * The 'relinquishCursor' parameter indicates whether the stages are recovering from a "full
+     * save" or not, as discussed in saveState(). It is the caller's responsibility to pass the same
+     * value for 'relinquishCursor' as was passed in the previous call to saveState().
      */
-    void restoreState() {
+    void restoreState(bool relinquishCursor) {
         auto stage = static_cast<T*>(this);
         stage->_commonStats.unyields++;
+#if defined(MONGO_CONFIG_DEBUG_BUILD)
+        if (relinquishCursor) {
+            invariant(_saveState == SaveState::kSavedFull);
+        } else {
+            invariant(_saveState == SaveState::kSavedNotFull);
+        }
+#endif
+
         for (auto&& child : stage->_children) {
-            child->restoreState();
+            child->restoreState(relinquishCursor);
         }
 
-        stage->doRestoreState();
+        stage->doRestoreState(relinquishCursor);
+#if defined(MONGO_CONFIG_DEBUG_BUILD)
+        stage->_saveState = SaveState::kNotSaved;
+#endif
     }
+
+protected:
+    // We do not want to incur the overhead of tracking information about saved-ness
+    // per stage. This information is only used for sanity checking, so we only run these
+    // checks in debug builds.
+#if defined(MONGO_CONFIG_DEBUG_BUILD)
+    // TODO SERVER-59620: Remove this.
+    enum class SaveState { kNotSaved, kSavedFull, kSavedNotFull };
+    SaveState _saveState{SaveState::kNotSaved};
+#endif
 };
 
 /**
@@ -426,8 +464,8 @@ public:
 
 protected:
     // Derived classes can optionally override these methods.
-    virtual void doSaveState() {}
-    virtual void doRestoreState() {}
+    virtual void doSaveState(bool relinquishCursor) {}
+    virtual void doRestoreState(bool relinquishCursor) {}
     virtual void doDetachFromOperationContext() {}
     virtual void doAttachToOperationContext(OperationContext* opCtx) {}
     virtual void doDetachFromTrialRunTracker() {}
diff --git a/src/mongo/db/exec/sbe/stages/traverse.cpp b/src/mongo/db/exec/sbe/stages/traverse.cpp
index ec92c0525b8..7c117a0f7f9 100644
--- a/src/mongo/db/exec/sbe/stages/traverse.cpp
+++ b/src/mongo/db/exec/sbe/stages/traverse.cpp
@@ -276,7 +276,7 @@ void TraverseStage::close() {
     _children[0]->close();
 }
 
-void TraverseStage::doSaveState() {
+void TraverseStage::doSaveState(bool relinquishCursor) {
     if (_isReadingLeftSide) {
         // If we yield while reading the left side, there is no need to makeOwned() data held in
         // the right side, since we will have to re-open it anyway.
@@ -288,14 +288,14 @@ void TraverseStage::doSaveState() {
         _outFieldOutputAccessor.reset();
     }
 
-    if (!slotsAccessible()) {
+    if (!slotsAccessible() || !relinquishCursor) {
         return;
     }
 
     _outFieldOutputAccessor.makeOwned();
 }
 
-void TraverseStage::doRestoreState() {
+void TraverseStage::doRestoreState(bool relinquishCursor) {
     if (!slotsAccessible()) {
         return;
     }
diff --git a/src/mongo/db/exec/sbe/stages/traverse.h b/src/mongo/db/exec/sbe/stages/traverse.h
index f59dc1a6763..2b3fee33a47 100644
--- a/src/mongo/db/exec/sbe/stages/traverse.h
+++ b/src/mongo/db/exec/sbe/stages/traverse.h
@@ -90,8 +90,8 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() final;
-    void doRestoreState() final;
+    void doSaveState(bool relinquishCursor) final;
+    void doRestoreState(bool relinquishCursor) final;
 
 private:
     void openInner(value::TypeTags tag, value::Value val);
diff --git a/src/mongo/db/exec/sbe/stages/unwind.cpp b/src/mongo/db/exec/sbe/stages/unwind.cpp
index 8afc781044f..41fe4be59c2 100644
--- a/src/mongo/db/exec/sbe/stages/unwind.cpp
+++ b/src/mongo/db/exec/sbe/stages/unwind.cpp
@@ -201,8 +201,8 @@ std::vector<DebugPrinter::Block> UnwindStage::debugPrint() const {
     return ret;
 }
 
-void UnwindStage::doSaveState() {
-    if (!slotsAccessible()) {
+void UnwindStage::doSaveState(bool fullSave) {
+    if (!slotsAccessible() || !fullSave) {
         return;
     }
 
@@ -214,7 +214,7 @@ void UnwindStage::doSaveState() {
     }
 }
 
-void UnwindStage::doRestoreState() {
+void UnwindStage::doRestoreState(bool fullSave) {
     if (!slotsAccessible()) {
         return;
     }
diff --git a/src/mongo/db/exec/sbe/stages/unwind.h b/src/mongo/db/exec/sbe/stages/unwind.h
index 149130cf415..43a7bc88554 100644
--- a/src/mongo/db/exec/sbe/stages/unwind.h
+++ b/src/mongo/db/exec/sbe/stages/unwind.h
@@ -68,8 +68,8 @@ public:
     size_t estimateCompileTimeSize() const final;
 
 protected:
-    void doSaveState() final;
-    void doRestoreState() final;
+    void doSaveState(bool relinquishCursor) final;
+    void doRestoreState(bool relinquishCursor) final;
 
 private:
     const value::SlotId _inField;