SERVER-59682 Recover SBE plans from the new plan cache

5 files changed, 32 insertions, 9 deletions

diff --git a/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp b/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp
index 7037bd1e86d..659c249b18d 100644
--- a/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp
@@ -64,7 +64,7 @@ void HashAggStageTest::performHashAggWithSpillChecking(
     value::ValueGuard expectedGuard{expectedTag, expectedVal};
 
     auto collatorSlot = generateSlotId();
-    auto shouldUseCollator = optionalCollator.get() != nullptr;
+    auto shouldUseCollator = optionalCollator != nullptr;
 
     auto makeStageFn = [this, collatorSlot, shouldUseCollator, shouldSpill](
                            value::SlotId scanSlot, std::unique_ptr<PlanStage> scanStage) {
@@ -92,7 +92,7 @@ void HashAggStageTest::performHashAggWithSpillChecking(
     if (shouldUseCollator) {
         ctx->pushCorrelated(collatorSlot, &collatorAccessor);
         collatorAccessor.reset(value::TypeTags::collator,
-                               value::bitcastFrom<CollatorInterface*>(optionalCollator.get()));
+                               value::bitcastFrom<CollatorInterface*>(optionalCollator.release()));
     }
 
     // Generate a mock scan from 'input' with a single output slot.
@@ -155,8 +155,8 @@ TEST_F(HashAggStageTest, HashAggMinMaxTest) {
 
     auto makeStageFn = [this, &collator](value::SlotId scanSlot,
                                          std::unique_ptr<PlanStage> scanStage) {
-        auto collExpr = makeE<EConstant>(value::TypeTags::collator,
-                                         value::bitcastFrom<CollatorInterface*>(collator.get()));
+        auto collExpr = makeE<EConstant>(
+            value::TypeTags::collator, value::bitcastFrom<CollatorInterface*>(collator.release()));
 
         // Build a HashAggStage that exercises the collMin() and collMax() aggregate functions.
         auto minSlot = generateSlotId();
@@ -222,8 +222,8 @@ TEST_F(HashAggStageTest, HashAggAddToSetTest) {
 
     auto makeStageFn = [this, &collator](value::SlotId scanSlot,
                                          std::unique_ptr<PlanStage> scanStage) {
-        auto collExpr = makeE<EConstant>(value::TypeTags::collator,
-                                         value::bitcastFrom<CollatorInterface*>(collator.get()));
+        auto collExpr = makeE<EConstant>(
+            value::TypeTags::collator, value::bitcastFrom<CollatorInterface*>(collator.release()));
 
         // Build a HashAggStage that exercises the collAddToSet() aggregate function.
         auto hashAggSlot = generateSlotId();
diff --git a/src/mongo/db/exec/sbe/sbe_hash_join_test.cpp b/src/mongo/db/exec/sbe/sbe_hash_join_test.cpp
index 3eb25c4ebf4..4bdcedb31a0 100644
--- a/src/mongo/db/exec/sbe/sbe_hash_join_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_hash_join_test.cpp
@@ -92,7 +92,7 @@ TEST_F(HashJoinStageTest, HashJoinCollationTest) {
         value::OwnedValueAccessor collatorAccessor;
         ctx->pushCorrelated(collatorSlot, &collatorAccessor);
         collatorAccessor.reset(value::TypeTags::collator,
-                               value::bitcastFrom<CollatorInterface*>(collator.get()));
+                               value::bitcastFrom<CollatorInterface*>(collator.release()));
 
         // Two separate virtual scans are needed since HashJoinStage needs two child stages.
         outerGuard.reset();
diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h
index 555dc05aaa5..9afd8f7de64 100644
--- a/src/mongo/db/exec/sbe/stages/stages.h
+++ b/src/mongo/db/exec/sbe/stages/stages.h
@@ -360,6 +360,7 @@ private:
  * Maintains an internal state to maintain the interrupt check period. Also responsible for
  * triggering yields if this object has been configured with a yield policy.
  */
+template <typename T>
 class CanInterrupt {
 public:
     /**
@@ -389,8 +390,17 @@ public:
         }
     }
 
+    void attachNewYieldPolicy(PlanYieldPolicy* yieldPolicy) {
+        auto stage = static_cast<T*>(this);
+        for (auto&& child : stage->_children) {
+            child->attachNewYieldPolicy(yieldPolicy);
+        }
+
+        _yieldPolicy = yieldPolicy;
+    }
+
 protected:
-    PlanYieldPolicy* const _yieldPolicy{nullptr};
+    PlanYieldPolicy* _yieldPolicy{nullptr};
 
 private:
     static const int kInterruptCheckPeriod = 128;
@@ -403,7 +413,7 @@ private:
 class PlanStage : public CanSwitchOperationContext<PlanStage>,
                   public CanChangeState<PlanStage>,
                   public CanTrackStats<PlanStage>,
-                  public CanInterrupt {
+                  public CanInterrupt<PlanStage> {
 public:
     using Vector = absl::InlinedVector<std::unique_ptr<PlanStage>, 2>;
 
@@ -470,6 +480,7 @@ public:
     friend class CanSwitchOperationContext<PlanStage>;
     friend class CanChangeState<PlanStage>;
     friend class CanTrackStats<PlanStage>;
+    friend class CanInterrupt<PlanStage>;
 
 protected:
     // Derived classes can optionally override these methods.
diff --git a/src/mongo/db/exec/sbe/values/value.cpp b/src/mongo/db/exec/sbe/values/value.cpp
index 0e2e19d08a5..8c7de0f1b8e 100644
--- a/src/mongo/db/exec/sbe/values/value.cpp
+++ b/src/mongo/db/exec/sbe/values/value.cpp
@@ -243,6 +243,11 @@ std::pair<TypeTags, Value> makeCopySortSpec(const SortSpec& ss) {
     return {TypeTags::sortSpec, ssCopy};
 }
 
+std::pair<TypeTags, Value> makeCopyCollator(const CollatorInterface& collator) {
+    auto collatorCopy = bitcastFrom<CollatorInterface*>(collator.clone().release());
+    return {TypeTags::collator, collatorCopy};
+}
+
 std::pair<TypeTags, Value> makeNewRecordId(int64_t rid) {
     auto val = bitcastFrom<RecordId*>(new RecordId(rid));
     return {TypeTags::RecordId, val};
@@ -311,6 +316,9 @@ void releaseValue(TypeTags tag, Value val) noexcept {
         case TypeTags::sortSpec:
             delete getSortSpecView(val);
             break;
+        case TypeTags::collator:
+            delete getCollatorView(val);
+            break;
         default:
             break;
     }
diff --git a/src/mongo/db/exec/sbe/values/value.h b/src/mongo/db/exec/sbe/values/value.h
index b5947391347..6a7cfe810e8 100644
--- a/src/mongo/db/exec/sbe/values/value.h
+++ b/src/mongo/db/exec/sbe/values/value.h
@@ -1359,6 +1359,8 @@ std::pair<TypeTags, Value> makeCopyFtsMatcher(const fts::FTSMatcher&);
 
 std::pair<TypeTags, Value> makeCopySortSpec(const SortSpec&);
 
+std::pair<TypeTags, Value> makeCopyCollator(const CollatorInterface& collator);
+
 /**
  * Releases memory allocated for the value. If the value does not have any memory allocated for it,
  * does nothing.
@@ -1434,6 +1436,8 @@ inline std::pair<TypeTags, Value> copyValue(TypeTags tag, Value val) {
             return makeCopyFtsMatcher(*getFtsMatcherView(val));
         case TypeTags::sortSpec:
             return makeCopySortSpec(*getSortSpecView(val));
+        case TypeTags::collator:
+            return makeCopyCollator(*getCollatorView(val));
         default:
             break;
     }