SERVER-53501 Fix incorrect reporting of scannedObjects in serverStatus() for $unionWith operation

5 files changed, 116 insertions, 10 deletions

diff --git a/jstests/aggregation/sources/unionWith/unionWith_query_stats.js b/jstests/aggregation/sources/unionWith/unionWith_query_stats.js
new file mode 100644
index 00000000000..113e5a99ac0
--- /dev/null
+++ b/jstests/aggregation/sources/unionWith/unionWith_query_stats.js
@@ -0,0 +1,56 @@
+/**
+ * Tests that the queryExecutor stats are correctly returned when $unionWith is performed on
+ * collections.
+ *
+ * @tags: [
+ *     assumes_unsharded_collection,
+ *     do_not_wrap_aggregations_in_facets,
+ *     assumes_read_preference_unchanged,
+ *     assumes_read_concern_unchanged,
+ *     assumes_against_mongod_not_mongos
+ * ]
+ */
+(function() {
+"use strict";
+
+const testDB = db.getSiblingDB("union_with_query_stats");
+testDB.dropDatabase();
+
+const collData = [
+    ["firstCol", Array.from({length: 10}, (_, i) => ({_id: i, "aField": i}))],
+    ["secondColl", Array.from({length: 20}, (_, i) => ({_id: i, "aField": i}))],
+    ["thirdColl", Array.from({length: 30}, (_, i) => ({_id: i, "aField": i}))],
+    ["forthColl", Array.from({length: 40}, (_, i) => ({_id: i, "aField": i}))]
+];
+
+const colls = Array.from(collData, elem => testDB.getCollection(elem[0]));
+
+for (let idx = 0; idx < collData.length; idx++) {
+    const coll = colls[idx];
+    const collDocs = collData[idx][1];
+    assert.commandWorked(coll.insert(collDocs));
+}
+
+(function testUnionWithQueryStats() {
+    // Collect the previous scannedObjects before running aggregation. This value will be
+    // subtracted from the current scannedObjects and will prevent test case from failing if it is
+    // run on existing mongoD instance.
+    const prevScannedObjects = testDB.serverStatus().metrics.queryExecutor.scannedObjects;
+
+    const pipeline = [
+        {$unionWith: {coll: collData[1][0]}},
+        {$unionWith: {coll: collData[2][0], pipeline: [{$unionWith: {coll: collData[3][0]}}]}},
+        {$sort: {_id: 1}}
+    ];
+
+    const output = colls[0].aggregate(pipeline).toArray();
+
+    // Concatenate and sort arrays by '_id'.
+    let expectedOutput = [].concat(collData[0][1], collData[1][1], collData[2][1], collData[3][1])
+                             .sort((elem1, elem2) => elem1._id - elem2._id);
+
+    assert.eq(output, expectedOutput);
+    assert.eq(expectedOutput.length,
+              testDB.serverStatus().metrics.queryExecutor.scannedObjects - prevScannedObjects);
+})();
+})();
diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h
index 19335e3d16a..2bb25ad0c6a 100644
--- a/src/mongo/db/exec/plan_stats.h
+++ b/src/mongo/db/exec/plan_stats.h
@@ -893,4 +893,22 @@ struct DocumentSourceLookupStats : public SpecificStats {
     PlanSummaryStats planSummaryStats;
 };
 
+struct UnionWithStats final : public SpecificStats {
+    SpecificStats* clone() const final {
+        return new UnionWithStats(*this);
+    }
+
+    uint64_t estimateObjectSizeInBytes() const {
+        return sizeof(*this) +
+            (planSummaryStats.estimateObjectSizeInBytes() - sizeof(planSummaryStats));
+    }
+
+    void accumulate(PlanSummaryStats& summary) const final {
+        summary.accumulate(planSummaryStats);
+    }
+
+    // Tracks the summary stats of the subpipeline.
+    PlanSummaryStats planSummaryStats;
+};
+
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/document_source_union_with.cpp b/src/mongo/db/pipeline/document_source_union_with.cpp
index 3191468d65b..bb89ec69a0d 100644
--- a/src/mongo/db/pipeline/document_source_union_with.cpp
+++ b/src/mongo/db/pipeline/document_source_union_with.cpp
@@ -210,6 +210,9 @@ DocumentSource::GetNextResult DocumentSourceUnionWith::doGetNext() {
     if (res)
         return std::move(*res);
 
+    // Record the plan summary stats after $unionWith operation is done.
+    recordPlanSummaryStats(*_pipeline);
+
     _executionState = ExecutionProgress::kFinished;
     return GetNextResult::makeEOF();
 }
@@ -234,14 +237,18 @@ Pipeline::SourceContainer::iterator DocumentSourceUnionWith::doOptimizeAt(
 
 bool DocumentSourceUnionWith::usedDisk() {
     if (_pipeline) {
-        _usedDisk = _usedDisk || _pipeline->usedDisk();
+        _stats.planSummaryStats.usedDisk =
+            _stats.planSummaryStats.usedDisk || _pipeline->usedDisk();
     }
-    return _usedDisk;
+    return _stats.planSummaryStats.usedDisk;
 }
 
 void DocumentSourceUnionWith::doDispose() {
     if (_pipeline) {
-        _usedDisk = _usedDisk || _pipeline->usedDisk();
+        _stats.planSummaryStats.usedDisk =
+            _stats.planSummaryStats.usedDisk || _pipeline->usedDisk();
+        recordPlanSummaryStats(*_pipeline);
+
         if (!_pipeline->getContext()->explain) {
             _pipeline->dispose(pExpCtx->opCtx);
             _pipeline.reset();
@@ -328,4 +335,12 @@ 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 38f5eeeb804..02dda56665b 100644
--- a/src/mongo/db/pipeline/document_source_union_with.h
+++ b/src/mongo/db/pipeline/document_source_union_with.h
@@ -126,6 +126,10 @@ public:
 
     bool usedDisk() final;
 
+    const SpecificStats* getSpecificStats() const final {
+        return &_stats;
+    }
+
 protected:
     GetNextResult doGetNext() final;
 
@@ -160,10 +164,12 @@ private:
 
     void addViewDefinition(NamespaceString nss, std::vector<BSONObj> viewPipeline);
 
+    void recordPlanSummaryStats(const Pipeline& pipeline);
+
     std::unique_ptr<Pipeline, PipelineDeleter> _pipeline;
     Pipeline::SourceContainer _cachedPipeline;
-    bool _usedDisk = false;
     ExecutionProgress _executionState = ExecutionProgress::kIteratingSource;
+    UnionWithStats _stats;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/plan_explainer_pipeline.cpp b/src/mongo/db/pipeline/plan_explainer_pipeline.cpp
index 64759c4ca3c..49e05ead917 100644
--- a/src/mongo/db/pipeline/plan_explainer_pipeline.cpp
+++ b/src/mongo/db/pipeline/plan_explainer_pipeline.cpp
@@ -34,10 +34,22 @@
 #include "mongo/db/pipeline/document_source_cursor.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"
 #include "mongo/db/pipeline/plan_executor_pipeline.h"
 #include "mongo/db/query/explain.h"
 
 namespace mongo {
+/**
+ * Templatized method to get plan summary stats from document source and aggregate it to 'statsOut'.
+ */
+template <typename DocSourceType, typename DocSourceStatType>
+void collectPlanSummaryStats(const DocSourceType& source, PlanSummaryStats* statsOut) {
+    auto specificStats = source.getSpecificStats();
+    invariant(specificStats);
+    auto& docSpecificStats = static_cast<const DocSourceStatType&>(*specificStats);
+    statsOut->accumulate(docSpecificStats.planSummaryStats);
+}
+
 std::string PlanExplainerPipeline::getPlanSummary() const {
     if (auto docSourceCursor =
             dynamic_cast<DocumentSourceCursor*>(_pipeline->getSources().front().get())) {
@@ -61,12 +73,11 @@ void PlanExplainerPipeline::getSummaryStats(PlanSummaryStats* statsOut) const {
         if (dynamic_cast<DocumentSourceSort*>(source.get())) {
             statsOut->hasSortStage = true;
         } else if (auto docSourceLookUp = dynamic_cast<DocumentSourceLookUp*>(source.get())) {
-            auto specificStats = docSourceLookUp->getSpecificStats();
-            invariant(specificStats);
-            auto lookupSpecificStats =
-                dynamic_cast<const DocumentSourceLookupStats*>(specificStats);
-            invariant(lookupSpecificStats);
-            statsOut->accumulate(lookupSpecificStats->planSummaryStats);
+            collectPlanSummaryStats<DocumentSourceLookUp, DocumentSourceLookupStats>(
+                *docSourceLookUp, statsOut);
+        } else if (auto docSourceUnionWith = dynamic_cast<DocumentSourceUnionWith*>(source.get())) {
+            collectPlanSummaryStats<DocumentSourceUnionWith, UnionWithStats>(*docSourceUnionWith,
+                                                                             statsOut);
         }
     }