SERVER-61046 TrialRunTracker support in HashAgg stage

11 files changed, 409 insertions, 11 deletions

diff --git a/src/mongo/db/exec/sbe/SConscript b/src/mongo/db/exec/sbe/SConscript
index 3a52248e9ac..e75bf78c705 100644
--- a/src/mongo/db/exec/sbe/SConscript
+++ b/src/mongo/db/exec/sbe/SConscript
@@ -164,6 +164,7 @@ env.CppUnitTest(
         'sbe_sorted_merge_test.cpp',
         'sbe_spool_test.cpp',
         'sbe_test.cpp',
+        'sbe_trial_run_tracker_test.cpp',
         'sbe_unique_test.cpp',
         'values/value_serialize_for_sorter_test.cpp',
         "values/value_test.cpp",
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
new file mode 100644
index 00000000000..6fb18efb84f
--- /dev/null
+++ b/src/mongo/db/exec/sbe/sbe_trial_run_tracker_test.cpp
@@ -0,0 +1,330 @@
+/**
+ *    Copyright (C) 2021-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+/**
+ * This file contains tests for sbe::HashAggStage.
+ */
+
+#include "mongo/platform/basic.h"
+
+
+#include "mongo/db/exec/sbe/sbe_plan_stage_test.h"
+#include "mongo/db/exec/sbe/stages/hash_agg.h"
+#include "mongo/db/exec/sbe/stages/scan.h"
+#include "mongo/db/exec/sbe/stages/sort.h"
+#include "mongo/db/exec/sbe/stages/union.h"
+#include "mongo/util/assert_util.h"
+
+namespace mongo::sbe {
+
+using TrialRunTrackerTest = PlanStageTestFixture;
+
+TEST_F(TrialRunTrackerTest, TrackerAttachesToStreamingStage) {
+    auto collUuid = CollectionUUID::parse("00000000-0000-0000-0000-000000000000").getValue();
+    auto scanStage = makeS<ScanStage>(collUuid,
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      boost::none,
+                                      std::vector<std::string>{"field"},
+                                      makeSV(generateSlotId()),
+                                      generateSlotId(),
+                                      true,
+                                      nullptr,
+                                      kEmptyPlanNodeId,
+                                      ScanCallbacks());
+
+    auto tracker = std::make_unique<TrialRunTracker>(size_t{0}, size_t{0});
+    ON_BLOCK_EXIT([&]() { scanStage->detachFromTrialRunTracker(); });
+
+    auto attachResult = scanStage->attachToTrialRunTracker(tracker.get());
+    ASSERT_EQ(attachResult, PlanStage::TrialRunTrackerAttachResultFlags::AttachedToStreamingStage);
+}
+
+TEST_F(TrialRunTrackerTest, TrackerAttachesToBlockingStage) {
+    auto sortStage = makeS<SortStage>(
+        makeS<LimitSkipStage>(
+            makeS<CoScanStage>(kEmptyPlanNodeId), 0, boost::none, kEmptyPlanNodeId),
+        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([&]() { sortStage->detachFromTrialRunTracker(); });
+
+    auto attachResult = sortStage->attachToTrialRunTracker(tracker.get());
+    ASSERT_EQ(attachResult, PlanStage::TrialRunTrackerAttachResultFlags::AttachedToBlockingStage);
+}
+
+TEST_F(TrialRunTrackerTest, TrackerAttachesToBothBlockingAndStreamingStages) {
+    auto collUuid = CollectionUUID::parse("00000000-0000-0000-0000-000000000000").getValue();
+    auto scanStage = makeS<ScanStage>(collUuid,
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      generateSlotId(),
+                                      boost::none,
+                                      std::vector<std::string>{"field"},
+                                      makeSV(generateSlotId()),
+                                      generateSlotId(),
+                                      true,
+                                      nullptr,
+                                      kEmptyPlanNodeId,
+                                      ScanCallbacks());
+
+    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::AttachedToStreamingStage |
+                  PlanStage::TrialRunTrackerAttachResultFlags::AttachedToBlockingStage);
+}
+
+TEST_F(TrialRunTrackerTest, TrialEndsDuringOpenPhaseOfBlockingStage) {
+    auto ctx = makeCompileCtx();
+
+    // Build a mock scan that will provide 9 values to its parent SortStage. The test
+    // TrialRunTracker will have a 'numResults' limit of 8, so it will reach its limit when the
+    // SortStage it attaches to reaches the last value.
+    const size_t numResultsLimit = 8;
+    auto [inputTag, inputVal] =
+        stage_builder::makeValue(BSON_ARRAY(1 << 2 << 3 << 4 << 5 << 6 << 7 << 8 << 9));
+    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,
+        kEmptyPlanNodeId);
+
+    auto tracker = std::make_unique<TrialRunTracker>(numResultsLimit, size_t{0});
+    auto attachResult = hashAggStage->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 HashAggStage, causing it to read documents from its
+    // child. Because the child provides more documents than the 'numResults' limit, we expect the
+    // open operation to be interrupted by a 'QueryTrialRunCompleted' exception.
+    ASSERT_THROWS_CODE(prepareTree(ctx.get(), hashAggStage.get(), countsSlot),
+                       DBException,
+                       ErrorCodes::QueryTrialRunCompleted);
+}
+
+TEST_F(TrialRunTrackerTest, OnlyDeepestNestedBlockingStageHasTrialRunTracker) {
+    auto ctx = makeCompileCtx();
+
+    // The contrived PlanStage tree constructed here allows us to observe what happens when a
+    // HashAgg stage has a SortStage in its subtree. A UnionStage injects extra documents into the
+    // HashAgg stage, allowing for a test scenario where the parent HashAgg stage sees many more
+    // documents than the SortStage.
+
+    // This "upperScan" gets unioned with the output of the SortStage and provides 10 values.
+    auto [upperScanSlot, upperScanStage] = [this]() {
+        auto [inputTag, inputVal] =
+            stage_builder::makeValue(BSON_ARRAY(1 << 2 << 3 << 4 << 5 << 6 << 7 << 8 << 9 << 10));
+        return generateVirtualScan(inputTag, inputVal);
+    }();
+
+    auto [sortSlot, sortStage] = [this]() {
+        auto [inputTag, inputVal] =
+            stage_builder::makeValue(BSON_ARRAY(1 << 2 << 3 << 4 << 5 << 6 << 7 << 8 << 9));
+        auto [scanSlot, scanStage] = generateVirtualScan(inputTag, inputVal);
+
+        auto sortStage =
+            makeS<SortStage>(std::move(scanStage),
+                             makeSV(scanSlot),
+                             std::vector<value::SortDirection>{value::SortDirection::Ascending},
+                             makeSV(),
+                             std::numeric_limits<std::size_t>::max(),
+                             204857600,
+                             false,
+                             kEmptyPlanNodeId);
+
+        return std::make_pair(scanSlot, std::move(sortStage));
+    }();
+
+    auto unionSlot = generateSlotId();
+    auto unionStage =
+        makeS<UnionStage>(makeSs(std::move(upperScanStage), std::move(sortStage)),
+                          std::vector<value::SlotVector>{makeSV(upperScanSlot), makeSV(sortSlot)},
+                          makeSV(unionSlot),
+                          kEmptyPlanNodeId);
+
+    auto countsSlot = generateSlotId();
+    auto hashAggStage = makeS<HashAggStage>(
+        std::move(unionStage),
+        makeSV(unionSlot),
+        makeEM(countsSlot,
+               stage_builder::makeFunction(
+                   "sum",
+                   makeE<EConstant>(value::TypeTags::NumberInt64, value::bitcastFrom<int64_t>(1)))),
+        makeSV(),  // Seek slot
+        true,
+        boost::none,
+        kEmptyPlanNodeId);
+
+    hashAggStage->prepare(*ctx);
+    hashAggStage->attachToOperationContext(opCtx());
+
+    {
+        // We expect the TrialRunTracker to attach to the SortStage but not the root HashAggStage.
+        auto tracker = std::make_unique<TrialRunTracker>(size_t{9}, size_t{0});
+        ON_BLOCK_EXIT([&]() { hashAggStage->detachFromTrialRunTracker(); });
+        auto attachResult = hashAggStage->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);
+
+        // In this scenario, the HashAggStage will see 10+ documents (the 10 documents from the
+        // "upperScan" plus the documents from the SortStage), which exceeds the 'numResults'
+        // requirement of the TrialRunTracker. The 'open()' call will _succeed_, however, because
+        // the TrialRunTracker is not attached to the HashAggStage.
+        hashAggStage->open(false);
+
+        hashAggStage->close();
+    }
+
+    {
+        // We expect the TrialRunTracker to attach to the SortStage but not the root HashAggStage.
+        auto tracker = std::make_unique<TrialRunTracker>(size_t{2}, size_t{0});
+        ON_BLOCK_EXIT([&]() { hashAggStage->detachFromTrialRunTracker(); });
+        auto attachResult = hashAggStage->attachToTrialRunTracker(tracker.get());
+
+        ASSERT_EQ(attachResult,
+                  PlanStage::TrialRunTrackerAttachResultFlags::AttachedToBlockingStage);
+
+        // In this scenario, the SortStage will see more documents than the 2 permitted by the
+        // 'numResults' requirement of the TrialRunTracker. The 'open()' call will _fail_, because
+        // the TrialRunTracker is attached to the SortStage.
+        ASSERT_THROWS_CODE(
+            hashAggStage->open(false), DBException, ErrorCodes::QueryTrialRunCompleted);
+
+        hashAggStage->close();
+    }
+}
+
+TEST_F(TrialRunTrackerTest, SiblingBlockingStagesBothGetTrialRunTracker) {
+    auto ctx = makeCompileCtx();
+
+    // This PlanStage tree allows us to observe what happens when we attach a TrialRunTracker to 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,
+            kEmptyPlanNodeId);
+
+        return std::make_pair(countsSlot, std::move(hashAggStage));
+    };
+
+    auto [leftCountsSlot, leftHashAggStage] = buildHashAgg();
+    auto [rightCountsSlot, rightHashAggStage] = buildHashAgg();
+
+    // The UnionStage allows us to establish the sibling relationship.
+    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 _both_ HashAgg stages but not to the SortStage.
+    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. If only one of the HashAgg stages were attached to the TrialRunTracker, it would not
+    // increment the 'numResults' metric enough to end the trial, but with both stages together,
+    // `numResults` gets incremented to 10 (with a limit of 9), resulting in a
+    // QueryTrialRunCompleted exception.
+    ASSERT_THROWS_CODE(prepareTree(ctx.get(), sortStage.get(), resultSlot),
+                       DBException,
+                       ErrorCodes::QueryTrialRunCompleted);
+}
+}  // namespace mongo::sbe
diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.cpp b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
index 2c5cd36f0e9..625d4c62585 100644
--- a/src/mongo/db/exec/sbe/stages/hash_agg.cpp
+++ b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
@@ -223,6 +223,16 @@ void HashAggStage::open(bool reOpen) {
                     makeTemporaryRecordStore();
                 }
             }
+
+            if (_tracker && _tracker->trackProgress<TrialRunTracker::kNumResults>(1)) {
+                // During trial runs, we want to limit the amount of work done by opening a blocking
+                // stage, like this one. The blocking stage tracks the number of documents it has
+                // read from its child, and if the TrialRunTracker ends the trial, a special
+                // exception returns control back to the planner.
+                _tracker = nullptr;
+                _children[0]->close();
+                uasserted(ErrorCodes::QueryTrialRunCompleted, "Trial run early exit in group");
+            }
         }
 
         if (_optimizedClose) {
@@ -374,5 +384,22 @@ size_t HashAggStage::estimateCompileTimeSize() const {
     return size;
 }
 
+void HashAggStage::doDetachFromTrialRunTracker() {
+    _tracker = nullptr;
+}
+
+PlanStage::TrialRunTrackerAttachResultMask HashAggStage::doAttachToTrialRunTracker(
+    TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) {
+    // The HashAggStage only tracks the "numResults" metric when it is the most deeply nested
+    // blocking stage.
+    if (!(childrenAttachResult & TrialRunTrackerAttachResultFlags::AttachedToBlockingStage)) {
+        _tracker = tracker;
+    }
+
+    // Return true to indicate that the tracker is attached to a blocking stage: either this stage
+    // or one of its descendent stages.
+    return childrenAttachResult | TrialRunTrackerAttachResultFlags::AttachedToBlockingStage;
+}
+
 }  // namespace sbe
 }  // namespace mongo
diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.h b/src/mongo/db/exec/sbe/stages/hash_agg.h
index 1061eb00195..557f95c7f32 100644
--- a/src/mongo/db/exec/sbe/stages/hash_agg.h
+++ b/src/mongo/db/exec/sbe/stages/hash_agg.h
@@ -89,6 +89,11 @@ public:
     std::vector<DebugPrinter::Block> debugPrint() const final;
     size_t estimateCompileTimeSize() const final;
 
+protected:
+    void doDetachFromTrialRunTracker() override;
+    TrialRunTrackerAttachResultMask doAttachToTrialRunTracker(
+        TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) override;
+
 private:
     void makeTemporaryRecordStore();
 
@@ -138,6 +143,10 @@ private:
 
     // Used when spilling to disk.
     std::unique_ptr<TemporaryRecordStore> _recordStore;
+
+    // If provided, used during a trial run to accumulate certain execution stats. Once the trial
+    // run is complete, this pointer is reset to nullptr.
+    TrialRunTracker* _tracker{nullptr};
 };
 
 }  // namespace sbe
diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp
index b913c5575d9..3c9937fda79 100644
--- a/src/mongo/db/exec/sbe/stages/ix_scan.cpp
+++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp
@@ -237,8 +237,10 @@ void IndexScanStage::doDetachFromTrialRunTracker() {
     _tracker = nullptr;
 }
 
-void IndexScanStage::doAttachToTrialRunTracker(TrialRunTracker* tracker) {
+PlanStage::TrialRunTrackerAttachResultMask IndexScanStage::doAttachToTrialRunTracker(
+    TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) {
     _tracker = tracker;
+    return childrenAttachResult | TrialRunTrackerAttachResultFlags::AttachedToStreamingStage;
 }
 
 void IndexScanStage::open(bool reOpen) {
diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.h b/src/mongo/db/exec/sbe/stages/ix_scan.h
index 21c5afcc9c0..a215509bb2d 100644
--- a/src/mongo/db/exec/sbe/stages/ix_scan.h
+++ b/src/mongo/db/exec/sbe/stages/ix_scan.h
@@ -104,7 +104,8 @@ protected:
     void doDetachFromOperationContext() override;
     void doAttachToOperationContext(OperationContext* opCtx) override;
     void doDetachFromTrialRunTracker() override;
-    void doAttachToTrialRunTracker(TrialRunTracker* tracker) override;
+    TrialRunTrackerAttachResultMask doAttachToTrialRunTracker(
+        TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) override;
 
 private:
     /**
diff --git a/src/mongo/db/exec/sbe/stages/scan.cpp b/src/mongo/db/exec/sbe/stages/scan.cpp
index e39f3b5a65b..068d2690750 100644
--- a/src/mongo/db/exec/sbe/stages/scan.cpp
+++ b/src/mongo/db/exec/sbe/stages/scan.cpp
@@ -263,8 +263,10 @@ void ScanStage::doDetachFromTrialRunTracker() {
     _tracker = nullptr;
 }
 
-void ScanStage::doAttachToTrialRunTracker(TrialRunTracker* tracker) {
+PlanStage::TrialRunTrackerAttachResultMask ScanStage::doAttachToTrialRunTracker(
+    TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) {
     _tracker = tracker;
+    return childrenAttachResult | TrialRunTrackerAttachResultFlags::AttachedToStreamingStage;
 }
 
 void ScanStage::open(bool reOpen) {
diff --git a/src/mongo/db/exec/sbe/stages/scan.h b/src/mongo/db/exec/sbe/stages/scan.h
index 002bc658c8b..c94db992d30 100644
--- a/src/mongo/db/exec/sbe/stages/scan.h
+++ b/src/mongo/db/exec/sbe/stages/scan.h
@@ -128,7 +128,8 @@ protected:
     void doDetachFromOperationContext() override;
     void doAttachToOperationContext(OperationContext* opCtx) override;
     void doDetachFromTrialRunTracker() override;
-    void doAttachToTrialRunTracker(TrialRunTracker* tracker) override;
+    TrialRunTrackerAttachResultMask doAttachToTrialRunTracker(
+        TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) override;
 
 private:
     const CollectionUUID _collUuid;
diff --git a/src/mongo/db/exec/sbe/stages/sort.cpp b/src/mongo/db/exec/sbe/stages/sort.cpp
index a7476eafa0b..5acf73afe8d 100644
--- a/src/mongo/db/exec/sbe/stages/sort.cpp
+++ b/src/mongo/db/exec/sbe/stages/sort.cpp
@@ -154,8 +154,17 @@ void SortStage::doDetachFromTrialRunTracker() {
     _tracker = nullptr;
 }
 
-void SortStage::doAttachToTrialRunTracker(TrialRunTracker* tracker) {
-    _tracker = tracker;
+PlanStage::TrialRunTrackerAttachResultMask SortStage::doAttachToTrialRunTracker(
+    TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) {
+    // The SortStage only tracks the "numResults" metric when it is the most deeply nested blocking
+    // stage.
+    if (!(childrenAttachResult & TrialRunTrackerAttachResultFlags::AttachedToBlockingStage)) {
+        _tracker = tracker;
+    }
+
+    // Return true to indicate that the tracker is attached to a blocking stage: either this stage
+    // or one of its descendent stages.
+    return childrenAttachResult | TrialRunTrackerAttachResultFlags::AttachedToBlockingStage;
 }
 
 void SortStage::open(bool reOpen) {
diff --git a/src/mongo/db/exec/sbe/stages/sort.h b/src/mongo/db/exec/sbe/stages/sort.h
index 5599eedea64..2bfc9e1d9fb 100644
--- a/src/mongo/db/exec/sbe/stages/sort.h
+++ b/src/mongo/db/exec/sbe/stages/sort.h
@@ -89,7 +89,8 @@ public:
 
 protected:
     void doDetachFromTrialRunTracker() override;
-    void doAttachToTrialRunTracker(TrialRunTracker* tracker) override;
+    TrialRunTrackerAttachResultMask doAttachToTrialRunTracker(
+        TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) override;
 
 private:
     void makeSorter();
diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h
index 53e2f35a80b..555dc05aaa5 100644
--- a/src/mongo/db/exec/sbe/stages/stages.h
+++ b/src/mongo/db/exec/sbe/stages/stages.h
@@ -262,13 +262,25 @@ public:
         stage->doDetachFromTrialRunTracker();
     }
 
-    void attachToTrialRunTracker(TrialRunTracker* tracker) {
+    // Bit flags to indicate what kinds of stages a TrialRunTracker was attached to by a call to the
+    // 'attachToTrialRunTracker()' method.
+    enum TrialRunTrackerAttachResultFlags : uint32_t {
+        NoAttachment = 0x0,
+        AttachedToStreamingStage = 1 << 0,
+        AttachedToBlockingStage = 1 << 1,
+    };
+
+    using TrialRunTrackerAttachResultMask = uint32_t;
+
+    TrialRunTrackerAttachResultMask attachToTrialRunTracker(TrialRunTracker* tracker) {
+        TrialRunTrackerAttachResultMask result = TrialRunTrackerAttachResultFlags::NoAttachment;
+
         auto stage = static_cast<T*>(this);
         for (auto&& child : stage->_children) {
-            child->attachToTrialRunTracker(tracker);
+            result |= child->attachToTrialRunTracker(tracker);
         }
 
-        stage->doAttachToTrialRunTracker(tracker);
+        return result | stage->doAttachToTrialRunTracker(tracker, result);
     }
 
     /**
@@ -466,7 +478,10 @@ protected:
     virtual void doDetachFromOperationContext() {}
     virtual void doAttachToOperationContext(OperationContext* opCtx) {}
     virtual void doDetachFromTrialRunTracker() {}
-    virtual void doAttachToTrialRunTracker(TrialRunTracker* tracker) {}
+    virtual TrialRunTrackerAttachResultMask doAttachToTrialRunTracker(
+        TrialRunTracker* tracker, TrialRunTrackerAttachResultMask childrenAttachResult) {
+        return TrialRunTrackerAttachResultFlags::NoAttachment;
+    }
 
     Vector _children;
 };