SERVER-53771 Facet summary stats added

10 files changed, 54 insertions, 31 deletions

diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h
index fd662242b72..0c2aac4425a 100644
--- a/src/mongo/db/exec/plan_stats.h
+++ b/src/mongo/db/exec/plan_stats.h
@@ -864,6 +864,24 @@ struct UnionWithStats final : public SpecificStats {
     PlanSummaryStats planSummaryStats;
 };
 
+struct DocumentSourceFacetStats : public SpecificStats {
+    std::unique_ptr<SpecificStats> clone() const final {
+        return std::make_unique<DocumentSourceFacetStats>(*this);
+    }
+
+    uint64_t estimateObjectSizeInBytes() const {
+        return sizeof(*this) +
+            (planSummaryStats.estimateObjectSizeInBytes() - sizeof(planSummaryStats));
+    }
+
+    void accumulate(PlanSummaryStats& summary) const final {
+        summary.accumulate(planSummaryStats);
+    }
+
+    // Tracks the cumulative summary stats across all facets.
+    PlanSummaryStats planSummaryStats;
+};
+
 struct UnpackTimeseriesBucketStats final : public SpecificStats {
     std::unique_ptr<SpecificStats> clone() const final {
         return std::make_unique<UnpackTimeseriesBucketStats>(*this);
diff --git a/src/mongo/db/pipeline/document_source.cpp b/src/mongo/db/pipeline/document_source.cpp
index b4cf3a59362..f86e99022da 100644
--- a/src/mongo/db/pipeline/document_source.cpp
+++ b/src/mongo/db/pipeline/document_source.cpp
@@ -74,6 +74,15 @@ struct ParserRegistration {
 static StringMap<ParserRegistration> parserMap;
 }  // namespace
 
+void accumulatePipelinePlanSummaryStats(const Pipeline& pipeline,
+                                        PlanSummaryStats& planSummaryStats) {
+    for (auto&& source : pipeline.getSources()) {
+        if (auto specificStats = source->getSpecificStats()) {
+            specificStats->accumulate(planSummaryStats);
+        }
+    }
+}
+
 void DocumentSource::registerParser(
     string name,
     Parser parser,
diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index b75fc71ee15..cd86955a142 100644
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -709,4 +709,11 @@ private:
         boost::optional<ExplainOptions::Verbosity> explain = boost::none) const = 0;
 };
 
+/**
+ * Method to accumulate the plan summary stats from all stages of the pipeline into the given
+ * `planSummaryStats` object.
+ */
+void accumulatePipelinePlanSummaryStats(const Pipeline& pipeline,
+                                        PlanSummaryStats& planSummaryStats);
+
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_facet.cpp b/src/mongo/db/pipeline/document_source_facet.cpp
index 4d8da6de141..60f1394c726 100644
--- a/src/mongo/db/pipeline/document_source_facet.cpp
+++ b/src/mongo/db/pipeline/document_source_facet.cpp
@@ -172,6 +172,7 @@ DocumentSource::GetNextResult DocumentSourceFacet::doGetNext() {
                 results[facetId].emplace_back(next.releaseDocument());
             }
             allPipelinesEOF = allPipelinesEOF && next.isEOF();
+            accumulatePipelinePlanSummaryStats(*pipeline, _stats.planSummaryStats);
         }
     }
 
@@ -264,10 +265,12 @@ StageConstraints DocumentSourceFacet::constraints(Pipeline::SplitState) const {
 
 bool DocumentSourceFacet::usedDisk() {
     for (auto&& facet : _facets) {
-        if (facet.pipeline->usedDisk())
-            return true;
+        if (facet.pipeline->usedDisk()) {
+            _stats.planSummaryStats.usedDisk = true;
+            break;
+        }
     }
-    return false;
+    return _stats.planSummaryStats.usedDisk;
 }
 
 DepsTracker::State DocumentSourceFacet::getDependencies(DepsTracker* deps) const {
diff --git a/src/mongo/db/pipeline/document_source_facet.h b/src/mongo/db/pipeline/document_source_facet.h
index f5c18e4352d..3cc29294b5a 100644
--- a/src/mongo/db/pipeline/document_source_facet.h
+++ b/src/mongo/db/pipeline/document_source_facet.h
@@ -141,6 +141,9 @@ public:
     void reattachToOperationContext(OperationContext* opCtx) final;
     StageConstraints constraints(Pipeline::SplitState pipeState) const final;
     bool usedDisk() final;
+    const SpecificStats* getSpecificStats() const final {
+        return &_stats;
+    }
 
 protected:
     /**
@@ -163,5 +166,7 @@ private:
     const size_t _maxOutputDocSizeBytes;
 
     bool _done = false;
+
+    DocumentSourceFacetStats _stats;
 };
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_lookup.cpp b/src/mongo/db/pipeline/document_source_lookup.cpp
index eefa1e13626..4edce8a721b 100644
--- a/src/mongo/db/pipeline/document_source_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup.cpp
@@ -480,7 +480,7 @@ DocumentSource::GetNextResult DocumentSourceLookUp::doGetNext() {
         results.emplace_back(std::move(*result));
     }
 
-    recordPlanSummaryStats(*pipeline);
+    accumulatePipelinePlanSummaryStats(*pipeline, _stats.planSummaryStats);
     MutableDocument output(std::move(inputDoc));
     output.setNestedField(_as, Value(std::move(results)));
     return output.freeze();
@@ -806,7 +806,7 @@ bool DocumentSourceLookUp::usedDisk() {
 
 void DocumentSourceLookUp::doDispose() {
     if (_pipeline) {
-        recordPlanSummaryStats(*_pipeline);
+        accumulatePipelinePlanSummaryStats(*_pipeline, _stats.planSummaryStats);
         _pipeline->dispose(pExpCtx->opCtx);
         _pipeline.reset();
     }
@@ -922,7 +922,7 @@ DocumentSource::GetNextResult DocumentSourceLookUp::unwindResult() {
         }
 
         if (_pipeline) {
-            recordPlanSummaryStats(*_pipeline);
+            accumulatePipelinePlanSummaryStats(*_pipeline, _stats.planSummaryStats);
             _pipeline->dispose(pExpCtx->opCtx);
         }
 
@@ -983,14 +983,6 @@ void DocumentSourceLookUp::initializeResolvedIntrospectionPipeline() {
     _fromExpCtx->stopExpressionCounters();
 }
 
-void DocumentSourceLookUp::recordPlanSummaryStats(const Pipeline& pipeline) {
-    for (auto&& source : pipeline.getSources()) {
-        if (auto specificStats = source->getSpecificStats()) {
-            specificStats->accumulate(_stats.planSummaryStats);
-        }
-    }
-}
-
 void DocumentSourceLookUp::appendSpecificExecStats(MutableDocument& doc) const {
     const PlanSummaryStats& stats = _stats.planSummaryStats;
     doc["totalDocsExamined"] = Value(static_cast<long long>(stats.totalDocsExamined));
diff --git a/src/mongo/db/pipeline/document_source_lookup.h b/src/mongo/db/pipeline/document_source_lookup.h
index cf22eca3073..74b71070fbe 100644
--- a/src/mongo/db/pipeline/document_source_lookup.h
+++ b/src/mongo/db/pipeline/document_source_lookup.h
@@ -308,11 +308,6 @@ private:
     }
 
     /**
-     * Method to accumulate the plan summary stats from all stages of the pipeline.
-     */
-    void recordPlanSummaryStats(const Pipeline& pipeline);
-
-    /**
      * Method to add a DocumentSourceSequentialDocumentCache stage and optimize the pipeline to
      * move the cache to its final position.
      */
diff --git a/src/mongo/db/pipeline/document_source_union_with.cpp b/src/mongo/db/pipeline/document_source_union_with.cpp
index 7d1e4ea5568..2ddb448b641 100644
--- a/src/mongo/db/pipeline/document_source_union_with.cpp
+++ b/src/mongo/db/pipeline/document_source_union_with.cpp
@@ -252,7 +252,7 @@ DocumentSource::GetNextResult DocumentSourceUnionWith::doGetNext() {
         return std::move(*res);
 
     // Record the plan summary stats after $unionWith operation is done.
-    recordPlanSummaryStats(*_pipeline);
+    accumulatePipelinePlanSummaryStats(*_pipeline, _stats.planSummaryStats);
 
     _executionState = ExecutionProgress::kFinished;
     return GetNextResult::makeEOF();
@@ -318,7 +318,7 @@ void DocumentSourceUnionWith::doDispose() {
         _pipeline.get_deleter().dismissDisposal();
         _stats.planSummaryStats.usedDisk =
             _stats.planSummaryStats.usedDisk || _pipeline->usedDisk();
-        recordPlanSummaryStats(*_pipeline);
+        accumulatePipelinePlanSummaryStats(*_pipeline, _stats.planSummaryStats);
 
         if (!_pipeline->getContext()->explain) {
             _pipeline->dispose(pExpCtx->opCtx);
@@ -412,12 +412,4 @@ void DocumentSourceUnionWith::addInvolvedCollections(
     collectionNames->merge(_pipeline->getInvolvedCollections());
 }
 
-void DocumentSourceUnionWith::recordPlanSummaryStats(const Pipeline& pipeline) {
-    for (auto&& source : pipeline.getSources()) {
-        if (auto specificStats = source->getSpecificStats()) {
-            specificStats->accumulate(_stats.planSummaryStats);
-        }
-    }
-}
-
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_union_with.h b/src/mongo/db/pipeline/document_source_union_with.h
index 9219e08c48c..1d9068a9f4c 100644
--- a/src/mongo/db/pipeline/document_source_union_with.h
+++ b/src/mongo/db/pipeline/document_source_union_with.h
@@ -170,8 +170,6 @@ private:
 
     void addViewDefinition(NamespaceString nss, std::vector<BSONObj> viewPipeline);
 
-    void recordPlanSummaryStats(const Pipeline& pipeline);
-
     void logStartingSubPipeline(const std::vector<BSONObj>& serializedPipeline);
     void logShardedViewFound(
         const ExceptionFor<ErrorCodes::CommandOnShardedViewNotSupportedOnMongod>& e);
diff --git a/src/mongo/db/pipeline/plan_explainer_pipeline.cpp b/src/mongo/db/pipeline/plan_explainer_pipeline.cpp
index 072643fd7e6..cd138db0466 100644
--- a/src/mongo/db/pipeline/plan_explainer_pipeline.cpp
+++ b/src/mongo/db/pipeline/plan_explainer_pipeline.cpp
@@ -32,6 +32,7 @@
 #include "mongo/db/pipeline/plan_explainer_pipeline.h"
 
 #include "mongo/db/pipeline/document_source_cursor.h"
+#include "mongo/db/pipeline/document_source_facet.h"
 #include "mongo/db/pipeline/document_source_lookup.h"
 #include "mongo/db/pipeline/document_source_sort.h"
 #include "mongo/db/pipeline/document_source_union_with.h"
@@ -88,6 +89,9 @@ void PlanExplainerPipeline::getSummaryStats(PlanSummaryStats* statsOut) const {
         } else if (auto docSourceUnionWith = dynamic_cast<DocumentSourceUnionWith*>(source.get())) {
             collectPlanSummaryStats<DocumentSourceUnionWith, UnionWithStats>(*docSourceUnionWith,
                                                                              statsOut);
+        } else if (auto docSourceFacet = dynamic_cast<DocumentSourceFacet*>(source.get())) {
+            collectPlanSummaryStats<DocumentSourceFacet, DocumentSourceFacetStats>(*docSourceFacet,
+                                                                                   statsOut);
         }
     }