SERVER-50744 Support generation of summaryStats in SBE

13 files changed, 196 insertions, 45 deletions

diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h
index a67b2b6f452..5baf6c46990 100644
--- a/src/mongo/db/exec/plan_stats.h
+++ b/src/mongo/db/exec/plan_stats.h
@@ -36,6 +36,7 @@
 
 #include "mongo/db/index/multikey_paths.h"
 #include "mongo/db/jsobj.h"
+#include "mongo/db/query/plan_summary_stats.h"
 #include "mongo/db/query/stage_types.h"
 #include "mongo/util/container_size_helper.h"
 #include "mongo/util/time_support.h"
@@ -54,6 +55,8 @@ struct SpecificStats {
     virtual SpecificStats* clone() const = 0;
 
     virtual uint64_t estimateObjectSizeInBytes() const = 0;
+
+    virtual void accumulate(PlanSummaryStats& summary) const {}
 };
 
 // Every stage has CommonStats.
@@ -619,6 +622,14 @@ struct SortStats : public SpecificStats {
         return sortPattern.objsize() + sizeof(*this);
     }
 
+    void accumulate(PlanSummaryStats& summary) const final {
+        summary.hasSortStage = true;
+
+        if (wasDiskUsed) {
+            summary.usedDisk = true;
+        }
+    }
+
     // The pattern according to which we are sorting.
     BSONObj sortPattern;
 
diff --git a/src/mongo/db/exec/sbe/stages/plan_stats.h b/src/mongo/db/exec/sbe/stages/plan_stats.h
index b2332af62a5..26d6d8ec9d0 100644
--- a/src/mongo/db/exec/sbe/stages/plan_stats.h
+++ b/src/mongo/db/exec/sbe/stages/plan_stats.h
@@ -62,19 +62,23 @@ struct CommonStats {
 
 using PlanStageStats = BasePlanStageStats<CommonStats>;
 
-struct ScanStats : public SpecificStats {
+struct ScanStats final : public SpecificStats {
     SpecificStats* clone() const final {
         return new ScanStats(*this);
     }
 
-    uint64_t estimateObjectSizeInBytes() const {
+    uint64_t estimateObjectSizeInBytes() const final {
         return sizeof(*this);
     }
 
+    void accumulate(PlanSummaryStats& stats) const final {
+        stats.totalDocsExamined += numReads;
+    }
+
     size_t numReads{0};
 };
 
-struct IndexScanStats : public SpecificStats {
+struct IndexScanStats final : public SpecificStats {
     SpecificStats* clone() const final {
         return new IndexScanStats(*this);
     }
@@ -83,27 +87,31 @@ struct IndexScanStats : public SpecificStats {
         return sizeof(*this);
     }
 
+    void accumulate(PlanSummaryStats& stats) const final {
+        stats.totalKeysExamined += numReads;
+    }
+
     size_t numReads{0};
 };
 
-struct FilterStats : public SpecificStats {
+struct FilterStats final : public SpecificStats {
     SpecificStats* clone() const final {
         return new FilterStats(*this);
     }
 
-    uint64_t estimateObjectSizeInBytes() const {
+    uint64_t estimateObjectSizeInBytes() const final {
         return sizeof(*this);
     }
 
     size_t numTested{0};
 };
 
-struct LimitSkipStats : public SpecificStats {
+struct LimitSkipStats final : public SpecificStats {
     SpecificStats* clone() const final {
         return new LimitSkipStats(*this);
     }
 
-    uint64_t estimateObjectSizeInBytes() const {
+    uint64_t estimateObjectSizeInBytes() const final {
         return sizeof(*this);
     }
 
diff --git a/src/mongo/db/exec/sbe/stages/sort.cpp b/src/mongo/db/exec/sbe/stages/sort.cpp
index 81634b1d828..2e89692d031 100644
--- a/src/mongo/db/exec/sbe/stages/sort.cpp
+++ b/src/mongo/db/exec/sbe/stages/sort.cpp
@@ -58,14 +58,15 @@ SortStage::SortStage(std::unique_ptr<PlanStage> input,
       _obs(std::move(obs)),
       _dirs(std::move(dirs)),
       _vals(std::move(vals)),
-      _limit(limit),
-      _memoryLimit(memoryLimit),
       _allowDiskUse(allowDiskUse),
       _mergeData({0, 0}),
       _tracker(tracker) {
     _children.emplace_back(std::move(input));
 
     invariant(_obs.size() == _dirs.size());
+
+    _specificStats.limit = limit;
+    _specificStats.maxMemoryUsageBytes = memoryLimit;
 }
 
 SortStage ::~SortStage() {}
@@ -75,8 +76,8 @@ std::unique_ptr<PlanStage> SortStage::clone() const {
                                        _obs,
                                        _dirs,
                                        _vals,
-                                       _limit,
-                                       _memoryLimit,
+                                       _specificStats.limit,
+                                       _specificStats.maxMemoryUsageBytes,
                                        _allowDiskUse,
                                        _tracker,
                                        _commonStats.nodeId);
@@ -121,9 +122,10 @@ value::SlotAccessor* SortStage::getAccessor(CompileCtx& ctx, value::SlotId slot)
 void SortStage::makeSorter() {
     SortOptions opts;
     opts.tempDir = storageGlobalParams.dbpath + "/_tmp";
-    opts.maxMemoryUsageBytes = _memoryLimit;
+    opts.maxMemoryUsageBytes = _specificStats.maxMemoryUsageBytes;
     opts.extSortAllowed = _allowDiskUse;
-    opts.limit = _limit != std::numeric_limits<size_t>::max() ? _limit : 0;
+    opts.limit =
+        _specificStats.limit != std::numeric_limits<size_t>::max() ? _specificStats.limit : 0;
 
     auto comp = [&](const SorterData& lhs, const SorterData& rhs) {
         auto size = lhs.first.size();
@@ -169,6 +171,7 @@ void SortStage::open(bool reOpen) {
             vals.reset(idx++, true, tag, val);
         }
 
+        // TODO SERVER-51815: count total mem usage for specificStats.
         _sorter->emplace(std::move(keys), std::move(vals));
 
         if (_tracker && _tracker->trackProgress<TrialRunProgressTracker::kNumResults>(1)) {
@@ -187,6 +190,7 @@ void SortStage::open(bool reOpen) {
     }
 
     _mergeIt.reset(_sorter->done());
+    _specificStats.wasDiskUsed = _specificStats.wasDiskUsed || _sorter->usedDisk();
 
     _children[0]->close();
 }
@@ -210,12 +214,13 @@ void SortStage::close() {
 
 std::unique_ptr<PlanStageStats> SortStage::getStats() const {
     auto ret = std::make_unique<PlanStageStats>(_commonStats);
+    ret->specific = std::make_unique<SortStats>(_specificStats);
     ret->children.emplace_back(_children[0]->getStats());
     return ret;
 }
 
 const SpecificStats* SortStage::getSpecificStats() const {
-    return nullptr;
+    return &_specificStats;
 }
 
 std::vector<DebugPrinter::Block> SortStage::debugPrint() const {
@@ -241,8 +246,8 @@ std::vector<DebugPrinter::Block> SortStage::debugPrint() const {
     }
     ret.emplace_back("`]");
 
-    if (_limit != std::numeric_limits<size_t>::max()) {
-        ret.emplace_back(std::to_string(_limit));
+    if (_specificStats.limit != std::numeric_limits<size_t>::max()) {
+        ret.emplace_back(std::to_string(_specificStats.limit));
     }
 
     DebugPrinter::addNewLine(ret);
diff --git a/src/mongo/db/exec/sbe/stages/sort.h b/src/mongo/db/exec/sbe/stages/sort.h
index 2e487de7cdb..e61f2e6c9ff 100644
--- a/src/mongo/db/exec/sbe/stages/sort.h
+++ b/src/mongo/db/exec/sbe/stages/sort.h
@@ -75,9 +75,8 @@ private:
     const value::SlotVector _obs;
     const std::vector<value::SortDirection> _dirs;
     const value::SlotVector _vals;
-    const size_t _limit;
-    const size_t _memoryLimit;
     const bool _allowDiskUse;
+    SortStats _specificStats;
 
     std::vector<value::SlotAccessor*> _inKeyAccessors;
     std::vector<value::SlotAccessor*> _inValueAccessors;
diff --git a/src/mongo/db/exec/sbe/stages/stages.cpp b/src/mongo/db/exec/sbe/stages/stages.cpp
index 423142a2c98..fcb436ef7df 100644
--- a/src/mongo/db/exec/sbe/stages/stages.cpp
+++ b/src/mongo/db/exec/sbe/stages/stages.cpp
@@ -75,5 +75,16 @@ void CanChangeState::restoreState() {
 
     doRestoreState();
 }
+
+void CanTrackStats::accumulate(PlanNodeId nodeId, PlanSummaryStats& summary) const {
+    if (auto stats = getSpecificStats();
+        stats && (nodeId == kEmptyPlanNodeId || _commonStats.nodeId == nodeId)) {
+        stats->accumulate(summary);
+    }
+
+    for (auto&& child : _stage->_children) {
+        child->accumulate(nodeId, summary);
+    }
+}
 }  // namespace sbe
 }  // namespace mongo
diff --git a/src/mongo/db/exec/sbe/stages/stages.h b/src/mongo/db/exec/sbe/stages/stages.h
index c94875d3cf5..3e1712dd1e2 100644
--- a/src/mongo/db/exec/sbe/stages/stages.h
+++ b/src/mongo/db/exec/sbe/stages/stages.h
@@ -135,7 +135,8 @@ private:
  */
 class CanTrackStats {
 public:
-    CanTrackStats(StringData stageType, PlanNodeId nodeId) : _commonStats(stageType, nodeId) {}
+    CanTrackStats(PlanStage* stage, StringData stageType, PlanNodeId nodeId)
+        : _stage(stage), _commonStats(stageType, nodeId) {}
 
     /**
      * Returns a tree of stats. If the stage has any children it must propagate the request for
@@ -164,6 +165,14 @@ public:
         return &_commonStats;
     }
 
+    /**
+     * Populates plan 'summary' object by walking through the entire PlanStage tree and for each
+     * node whose plan node ID equals to the given 'nodeId', or if 'nodeId' is 'kEmptyPlanNodeId',
+     * invoking 'accumulate(summary)' on the SpecificStats instance obtained by calling
+     * 'getSpecificStats()'.
+     */
+    void accumulate(PlanNodeId nodeId, PlanSummaryStats& summary) const;
+
 protected:
     PlanState trackPlanState(PlanState state) {
         if (state == PlanState::IS_EOF) {
@@ -175,6 +184,7 @@ protected:
         return state;
     }
 
+    PlanStage* const _stage;
     CommonStats _commonStats;
 };
 
@@ -227,7 +237,7 @@ public:
     PlanStage(StringData stageType, PlanYieldPolicy* yieldPolicy, PlanNodeId nodeId)
         : CanSwitchOperationContext{this},
           CanChangeState{this},
-          CanTrackStats{stageType, nodeId},
+          CanTrackStats{this, stageType, nodeId},
           CanInterrupt{yieldPolicy} {}
 
     PlanStage(StringData stageType, PlanNodeId nodeId) : PlanStage(stageType, nullptr, nodeId) {}
@@ -282,6 +292,7 @@ public:
 
     friend class CanSwitchOperationContext;
     friend class CanChangeState;
+    friend class CanTrackStats;
 
 protected:
     std::vector<std::unique_ptr<PlanStage>> _children;
diff --git a/src/mongo/db/query/classic_stage_builder.cpp b/src/mongo/db/query/classic_stage_builder.cpp
index 6390a25dde6..a6ed5c573ed 100644
--- a/src/mongo/db/query/classic_stage_builder.cpp
+++ b/src/mongo/db/query/classic_stage_builder.cpp
@@ -45,6 +45,7 @@
 #include "mongo/db/exec/count_scan.h"
 #include "mongo/db/exec/distinct_scan.h"
 #include "mongo/db/exec/ensure_sorted.h"
+#include "mongo/db/exec/eof.h"
 #include "mongo/db/exec/fetch.h"
 #include "mongo/db/exec/geo_near.h"
 #include "mongo/db/exec/index_scan.h"
@@ -350,10 +351,12 @@ std::unique_ptr<PlanStage> ClassicStageBuilder::build(const QuerySolutionNode* r
             return std::make_unique<EnsureSortedStage>(
                 expCtx, esn->pattern, _ws, std::move(childStage));
         }
+        case STAGE_EOF: {
+            return std::make_unique<EOFStage>(expCtx);
+        }
         case STAGE_CACHED_PLAN:
         case STAGE_COUNT:
         case STAGE_DELETE:
-        case STAGE_EOF:
         case STAGE_IDHACK:
         case STAGE_MOCK:
         case STAGE_MULTI_ITERATOR:
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index 20afe2d5b03..b38296378fe 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -550,7 +550,15 @@ public:
                         "Collection does not exist. Using EOF plan",
                         "namespace"_attr = _cq->ns(),
                         "canonicalQuery"_attr = redact(_cq->toStringShort()));
-            return buildEofPlan();
+
+            auto solution = std::make_unique<QuerySolution>();
+            solution->setRoot(std::make_unique<EofNode>());
+
+            auto root = buildExecutableTree(*solution);
+
+            auto result = makeResult();
+            result->emplace(std::move(root), std::move(solution));
+            return std::move(result);
         }
 
         // Fill out the planning params.  We use these for both cached solutions and non-cached.
@@ -690,11 +698,6 @@ protected:
     virtual PlanStageType buildExecutableTree(const QuerySolution& solution) const = 0;
 
     /**
-     * Constructs a special PlanStage tree to return EOF immediately on the first call to getNext.
-     */
-    virtual std::unique_ptr<ResultType> buildEofPlan() = 0;
-
-    /**
      * If supported, constructs a special PlanStage tree for fast-path document retrievals via the
      * _id index. Otherwise, nullptr should be returned and  this helper will fall back to the
      * normal plan generation.
@@ -763,12 +766,6 @@ protected:
         return stage_builder::buildClassicExecutableTree(_opCtx, _collection, *_cq, solution, _ws);
     }
 
-    std::unique_ptr<ClassicPrepareExecutionResult> buildEofPlan() final {
-        auto result = makeResult();
-        result->emplace(std::make_unique<EOFStage>(_cq->getExpCtxRaw()), nullptr);
-        return result;
-    }
-
     std::unique_ptr<ClassicPrepareExecutionResult> buildIdHackPlan(
         const IndexDescriptor* descriptor, QueryPlannerParams* plannerParams) final {
         auto result = makeResult();
@@ -906,16 +903,6 @@ protected:
         return buildExecutableTree(solution, false);
     }
 
-    std::unique_ptr<SlotBasedPrepareExecutionResult> buildEofPlan() final {
-        auto result = makeResult();
-        result->emplace(
-            {sbe::makeS<sbe::LimitSkipStage>(
-                 sbe::makeS<sbe::CoScanStage>(kEmptyPlanNodeId), 0, boost::none, kEmptyPlanNodeId),
-             stage_builder::PlanStageData{std::make_unique<sbe::RuntimeEnvironment>()}},
-            nullptr);
-        return result;
-    }
-
     std::unique_ptr<SlotBasedPrepareExecutionResult> buildIdHackPlan(
         const IndexDescriptor* descriptor, QueryPlannerParams* plannerParams) final {
         uassert(4822862,
diff --git a/src/mongo/db/query/plan_explainer_sbe.cpp b/src/mongo/db/query/plan_explainer_sbe.cpp
index 1b930e0bb3c..c91addc4199 100644
--- a/src/mongo/db/query/plan_explainer_sbe.cpp
+++ b/src/mongo/db/query/plan_explainer_sbe.cpp
@@ -105,7 +105,78 @@ std::string PlanExplainerSBE::getPlanSummary() const {
 }
 
 void PlanExplainerSBE::getSummaryStats(PlanSummaryStats* statsOut) const {
-    // TODO: SERVER-50744
+    invariant(statsOut);
+
+    if (!_solution || !_root) {
+        return;
+    }
+
+    auto common = _root->getCommonStats();
+    statsOut->nReturned = common->advances;
+    statsOut->fromMultiPlanner = isMultiPlan();
+    statsOut->totalKeysExamined = 0;
+    statsOut->totalDocsExamined = 0;
+
+    // Collect cumulative execution stats for the plan.
+    _root->accumulate(kEmptyPlanNodeId, *statsOut);
+
+    std::queue<const QuerySolutionNode*> queue;
+    queue.push(_solution->root());
+
+    // Look through the QuerySolution to collect some static stat details.
+    //
+    // TODO SERVER-51138: handle replan reason for cached plan.
+    while (!queue.empty()) {
+        auto node = queue.front();
+        queue.pop();
+        invariant(node);
+
+        switch (node->getType()) {
+            case STAGE_COUNT_SCAN: {
+                auto csn = static_cast<const CountScanNode*>(node);
+                statsOut->indexesUsed.insert(csn->index.identifier.catalogName);
+                break;
+            }
+            case STAGE_DISTINCT_SCAN: {
+                auto dn = static_cast<const DistinctNode*>(node);
+                statsOut->indexesUsed.insert(dn->index.identifier.catalogName);
+                break;
+            }
+            case STAGE_GEO_NEAR_2D: {
+                auto geo2d = static_cast<const GeoNear2DNode*>(node);
+                statsOut->indexesUsed.insert(geo2d->index.identifier.catalogName);
+                break;
+            }
+            case STAGE_GEO_NEAR_2DSPHERE: {
+                auto geo2dsphere = static_cast<const GeoNear2DSphereNode*>(node);
+                statsOut->indexesUsed.insert(geo2dsphere->index.identifier.catalogName);
+                break;
+            }
+            case STAGE_IXSCAN: {
+                auto ixn = static_cast<const IndexScanNode*>(node);
+                statsOut->indexesUsed.insert(ixn->index.identifier.catalogName);
+                break;
+            }
+            case STAGE_TEXT: {
+                auto tn = static_cast<const TextNode*>(node);
+                statsOut->indexesUsed.insert(tn->index.identifier.catalogName);
+                break;
+            }
+            case STAGE_COLLSCAN: {
+                statsOut->collectionScans++;
+                auto csn = static_cast<const CollectionScanNode*>(node);
+                if (!csn->tailable) {
+                    statsOut->collectionScansNonTailable++;
+                }
+            }
+            default:
+                break;
+        }
+
+        for (auto&& child : node->children) {
+            queue.push(child);
+        }
+    }
 }
 
 PlanExplainer::PlanStatsDetails PlanExplainerSBE::getWinningPlanStats(
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index d45567d7550..f0146a19e3f 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -1322,4 +1322,19 @@ QuerySolutionNode* EnsureSortedNode::clone() const {
     return copy;
 }
 
+//
+// EofNode
+//
+
+void EofNode::appendToString(str::stream* ss, int indent) const {
+    addIndent(ss, indent);
+    *ss << "EOF\n";
+}
+
+QuerySolutionNode* EofNode::clone() const {
+    auto copy = new EofNode();
+    cloneBaseData(copy);
+    return copy;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index d9fb47db6e0..5e4280582f8 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -1166,4 +1166,27 @@ struct EnsureSortedNode : public QuerySolutionNode {
     BSONObj pattern;
 };
 
+struct EofNode : public QuerySolutionNodeWithSortSet {
+    EofNode() {}
+
+    virtual StageType getType() const {
+        return STAGE_EOF;
+    }
+
+    virtual void appendToString(str::stream* ss, int indent) const;
+
+    bool fetched() const {
+        return false;
+    }
+
+    FieldAvailability getFieldAvailability(const std::string& field) const {
+        return FieldAvailability::kNotProvided;
+    }
+
+    bool sortedByDiskLoc() const {
+        return false;
+    }
+
+    QuerySolutionNode* clone() const;
+};
 }  // namespace mongo
diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp
index 48a32f3546c..cdc908e14e2 100644
--- a/src/mongo/db/query/sbe_stage_builder.cpp
+++ b/src/mongo/db/query/sbe_stage_builder.cpp
@@ -681,6 +681,11 @@ std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::makeUnionForTailableCollS
     return unionStage;
 }
 
+std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::buildEof(const QuerySolutionNode* root) {
+    return sbe::makeS<sbe::LimitSkipStage>(
+        sbe::makeS<sbe::CoScanStage>(root->nodeId()), 0, boost::none, root->nodeId());
+}
+
 // Returns a non-null pointer to the root of a plan tree, or a non-OK status if the PlanStage tree
 // could not be constructed.
 std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::build(const QuerySolutionNode* root) {
@@ -701,7 +706,8 @@ std::unique_ptr<sbe::PlanStage> SlotBasedStageBuilder::build(const QuerySolution
             {STAGE_PROJECTION_COVERED, std::mem_fn(&SlotBasedStageBuilder::buildProjectionCovered)},
             {STAGE_OR, &SlotBasedStageBuilder::buildOr},
             {STAGE_TEXT, &SlotBasedStageBuilder::buildText},
-            {STAGE_RETURN_KEY, &SlotBasedStageBuilder::buildReturnKey}};
+            {STAGE_RETURN_KEY, &SlotBasedStageBuilder::buildReturnKey},
+            {STAGE_EOF, &SlotBasedStageBuilder::buildEof}};
 
     uassert(4822884,
             str::stream() << "Can't build exec tree for node: " << root->toString(),
diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h
index bb94951462d..2a66eb7a332 100644
--- a/src/mongo/db/query/sbe_stage_builder.h
+++ b/src/mongo/db/query/sbe_stage_builder.h
@@ -123,6 +123,7 @@ private:
     std::unique_ptr<sbe::PlanStage> buildOr(const QuerySolutionNode* root);
     std::unique_ptr<sbe::PlanStage> buildText(const QuerySolutionNode* root);
     std::unique_ptr<sbe::PlanStage> buildReturnKey(const QuerySolutionNode* root);
+    std::unique_ptr<sbe::PlanStage> buildEof(const QuerySolutionNode* root);
 
     std::unique_ptr<sbe::PlanStage> makeLoopJoinForFetch(
         std::unique_ptr<sbe::PlanStage> inputStage,