SERVER-65197 Disable trial run tracking for the right side of pushed down $lookup plans

3 files changed, 156 insertions, 2 deletions

diff --git a/src/mongo/db/exec/sbe/sbe_trial_run_tracker_test.cpp b/src/mongo/db/exec/sbe/sbe_trial_run_tracker_test.cpp
index b3aa236dd8d..54fadbe1f15 100644
--- a/src/mongo/db/exec/sbe/sbe_trial_run_tracker_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_trial_run_tracker_test.cpp
@@ -330,4 +330,132 @@ TEST_F(TrialRunTrackerTest, SiblingBlockingStagesBothGetTrialRunTracker) {
                        DBException,
                        ErrorCodes::QueryTrialRunCompleted);
 }
+
+TEST_F(TrialRunTrackerTest, TrialRunTrackingCanBeDisabled) {
+    auto scanStage =
+        makeS<ScanStage>(UUID::parse("00000000-0000-0000-0000-000000000000").getValue(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         boost::none,
+                         std::vector<std::string>{"field"},
+                         makeSV(generateSlotId()),
+                         generateSlotId(),
+                         true,
+                         nullptr,
+                         kEmptyPlanNodeId,
+                         ScanCallbacks());
+
+    scanStage->disableTrialRunTracking();
+    auto tracker = std::make_unique<TrialRunTracker>(size_t{0}, size_t{0});
+    auto attachResult = scanStage->attachToTrialRunTracker(tracker.get());
+    ASSERT_EQ(attachResult, PlanStage::TrialRunTrackerAttachResultFlags::NoAttachment);
+}
+
+TEST_F(TrialRunTrackerTest, DisablingTrackingForChildDoesNotInhibitTrackingForParent) {
+    auto scanStage =
+        makeS<ScanStage>(UUID::parse("00000000-0000-0000-0000-000000000000").getValue(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         generateSlotId(),
+                         boost::none,
+                         std::vector<std::string>{"field"},
+                         makeSV(generateSlotId()),
+                         generateSlotId(),
+                         true,
+                         nullptr,
+                         kEmptyPlanNodeId,
+                         ScanCallbacks());
+
+    // Disable tracking for 'scanStage'. We should still attach the tracker for 'rootSortStage'.
+    scanStage->disableTrialRunTracking();
+
+    auto rootSortStage = makeS<SortStage>(std::move(scanStage),
+                                          makeSV(),
+                                          std::vector<value::SortDirection>{},
+                                          makeSV(),
+                                          std::numeric_limits<std::size_t>::max(),
+                                          204857600,
+                                          false,
+                                          kEmptyPlanNodeId);
+
+
+    auto tracker = std::make_unique<TrialRunTracker>(size_t{0}, size_t{0});
+    ON_BLOCK_EXIT([&]() { rootSortStage->detachFromTrialRunTracker(); });
+
+    auto attachResult = rootSortStage->attachToTrialRunTracker(tracker.get());
+    ASSERT_EQ(attachResult, PlanStage::TrialRunTrackerAttachResultFlags::AttachedToBlockingStage);
+}
+
+TEST_F(TrialRunTrackerTest, DisablingTrackingForAChildStagePreventsEarlyExit) {
+    auto ctx = makeCompileCtx();
+
+    // This PlanStage tree allows us to observe what happens when we attach a TrialRunTracker to
+    // one of two sibling HashAgg stages.
+    auto buildHashAgg = [&]() {
+        auto [inputTag, inputVal] = stage_builder::makeValue(BSON_ARRAY(1 << 2 << 3 << 4 << 5));
+        auto [scanSlot, scanStage] = generateVirtualScan(inputTag, inputVal);
+
+        // Build a HashAggStage, group by the scanSlot and compute a simple count.
+        auto countsSlot = generateSlotId();
+        auto hashAggStage = makeS<HashAggStage>(
+            std::move(scanStage),
+            makeSV(scanSlot),
+            makeEM(countsSlot,
+                   stage_builder::makeFunction("sum",
+                                               makeE<EConstant>(value::TypeTags::NumberInt64,
+                                                                value::bitcastFrom<int64_t>(1)))),
+            makeSV(),  // Seek slot
+            true,
+            boost::none,
+            false /* allowDiskUse */,
+            kEmptyPlanNodeId);
+
+        return std::make_pair(countsSlot, std::move(hashAggStage));
+    };
+
+    auto [leftCountsSlot, leftHashAggStage] = buildHashAgg();
+    auto [rightCountsSlot, rightHashAggStage] = buildHashAgg();
+
+    // Disable trial run tracking for the right HashAgg stage.
+    rightHashAggStage->disableTrialRunTracking();
+    auto resultSlot = generateSlotId();
+    auto unionStage = makeS<UnionStage>(
+        makeSs(std::move(leftHashAggStage), std::move(rightHashAggStage)),
+        std::vector<value::SlotVector>{makeSV(leftCountsSlot), makeSV(rightCountsSlot)},
+        makeSV(resultSlot),
+        kEmptyPlanNodeId);
+
+    // The blocking SortStage at the root ensures that both of the child HashAgg stages will be
+    // opened during the open phase of the root stage.
+    auto sortStage =
+        makeS<SortStage>(std::move(unionStage),
+                         makeSV(resultSlot),
+                         std::vector<value::SortDirection>{value::SortDirection::Ascending},
+                         makeSV(),
+                         std::numeric_limits<std::size_t>::max(),
+                         204857600,
+                         false,
+                         kEmptyPlanNodeId);
+
+    // We expect the TrialRunTracker to attach to _only_ the left child.
+    auto tracker = std::make_unique<TrialRunTracker>(size_t{9}, size_t{0});
+    auto attachResult = sortStage->attachToTrialRunTracker(tracker.get());
+
+    // Note: A scan is a streaming stage, but the "virtual scan" used here does not attach to the
+    // tracker.
+    ASSERT_EQ(attachResult, PlanStage::TrialRunTrackerAttachResultFlags::AttachedToBlockingStage);
+
+    // The 'prepareTree()' function opens the SortStage, causing it to read documents from its
+    // child. Because only one of the HashAgg stages is attached to the TrialRunTracker, the
+    // 'numResults' metric will not be incremented enough to end the trial. As such, this call to
+    // 'prepareTree()' will not end the trial.
+    prepareTree(ctx.get(), sortStage.get(), resultSlot);
+}
 }  // namespace mongo::sbe
diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h
index c44af3a26eb..5fdfb5650e0 100644
--- a/src/mongo/db/exec/sbe/stages/stages.h
+++ b/src/mongo/db/exec/sbe/stages/stages.h
@@ -274,6 +274,9 @@ public:
 
     TrialRunTrackerAttachResultMask attachToTrialRunTracker(TrialRunTracker* tracker) {
         TrialRunTrackerAttachResultMask result = TrialRunTrackerAttachResultFlags::NoAttachment;
+        if (!_participateInTrialRunTracking) {
+            return result;
+        }
 
         auto stage = static_cast<T*>(this);
         for (auto&& child : stage->_children) {
@@ -307,6 +310,10 @@ public:
         }
     }
 
+    void disableTrialRunTracking() {
+        _participateInTrialRunTracking = false;
+    }
+
 protected:
     PlanState trackPlanState(PlanState state) {
         if (state == PlanState::IS_EOF) {
@@ -320,7 +327,6 @@ protected:
         return state;
     }
 
-
     void trackClose() {
         _commonStats.closes++;
         _slotsAccessible = false;
@@ -347,12 +353,20 @@ protected:
 
 private:
     /**
-     * In general, accessors can be accesed only after getNext returns a row. It is most definitely
+     * In general, accessors can be accessed only after getNext returns a row. It is most definitely
      * not OK to access accessors in ANY other state; e.g. closed, not yet opened, after EOF. We
      * need this tracker to support unfortunate consequences of the internal yielding feature. Once
      * that feature is retired we can then simply revisit all stages and simplify them.
      */
     bool _slotsAccessible{false};
+
+    /**
+     * Flag which determines whether this node and its children can participate in trial run
+     * tracking. A stage and its children are not eligible for trial run tracking when they are
+     * planned deterministically (that is, the amount of work they perform is independent of
+     * other parts of the tree which are multiplanned).
+     */
+    bool _participateInTrialRunTracking{true};
 };
 
 /**
diff --git a/src/mongo/db/query/sbe_stage_builder_lookup.cpp b/src/mongo/db/query/sbe_stage_builder_lookup.cpp
index d3cd8d96bd5..83e506cd9e9 100644
--- a/src/mongo/db/query/sbe_stage_builder_lookup.cpp
+++ b/src/mongo/db/query/sbe_stage_builder_lookup.cpp
@@ -479,6 +479,10 @@ std::pair<SlotId /* matched docs */, std::unique_ptr<sbe::PlanStage>> buildNljLo
                                                                     slotIdGenerator,
                                                                     allowDiskUse);
 
+    // 'innerRootStage' should not participate in trial run tracking as the number of reads that
+    // it performs should not influence planning decisions made for 'outerRootStage'.
+    innerRootStage->disableTrialRunTracking();
+
     // Connect the two branches with a nested loop join. For each outer record with a corresponding
     // value in the 'localKeySlot', the inner branch will be executed and will place the result into
     // 'matchedRecordsSlot'.
@@ -791,6 +795,10 @@ std::pair<SlotId, std::unique_ptr<sbe::PlanStage>> buildIndexJoinLookupStage(
                                                                      slotIdGenerator,
                                                                      state.allowDiskUse);
 
+    // 'foreignGroupStage' should not participate in trial run tracking as the number of reads
+    // that it performs should not influence planning decisions for 'localKeysSetStage'.
+    foreignGroupStage->disableTrialRunTracking();
+
     // The top level loop join stage that joins each local field with the matched foreign
     // documents.
     auto nljStage = makeS<LoopJoinStage>(std::move(localKeysSetStage),
@@ -832,6 +840,10 @@ std::pair<SlotId /*matched docs*/, std::unique_ptr<sbe::PlanStage>> buildHashJoi
                                                          slotIdGenerator,
                                                          allowDiskUse);
 
+    // 'foreignKeyStage' should not participate in trial run tracking as the number of
+    // reads that it performs should not influence planning decisions for 'outerRootStage'.
+    foreignKeyStage->disableTrialRunTracking();
+
     // Build lookup stage that matches the local and foreign rows and aggregates the
     // foreign values in an array.
     auto lookupAggSlot = slotIdGenerator.generate();