SERVER-35973 Report MULTI_ITERATOR common exec stats correctly.

3 files changed, 62 insertions, 8 deletions

diff --git a/jstests/core/explain_sample.js b/jstests/core/explain_sample.js
new file mode 100644
index 00000000000..97d8aa98c7e
--- /dev/null
+++ b/jstests/core/explain_sample.js
@@ -0,0 +1,35 @@
+/**
+ * Tests for explaining an aggregation pipeline which uses the $sample stage. Only run on WT, since
+ * currently only the WT storage engine uses a storage-engine supported random cursor for $sample.
+ * @tags: [requires_wiredtiger]
+ */
+(function() {
+    "use strict";
+
+    load("jstests/libs/analyze_plan.js");
+
+    const coll = db.explain_sample;
+    coll.drop();
+
+    let docsToInsert = [];
+    for (let i = 0; i < 1000; ++i) {
+        docsToInsert.push({_id: i});
+    }
+    assert.commandWorked(coll.insert(docsToInsert));
+
+    // Verify that explain reports execution stats for the MULTI_ITERATOR stage. This is designed to
+    // reproduce SERVER-35973.
+    const explain =
+        assert.commandWorked(coll.explain("allPlansExecution").aggregate([{$sample: {size: 10}}]));
+    const multiIteratorStages = getAggPlanStages(explain, "MULTI_ITERATOR");
+    assert.gt(multiIteratorStages.length, 0, tojson(explain));
+    assert.gt(multiIteratorStages.reduce((acc, stage) => acc + stage.nReturned, 0),
+              0,
+              tojson(multiIteratorStages));
+    assert.gt(multiIteratorStages.reduce((acc, stage) => acc + stage.advanced, 0),
+              0,
+              tojson(multiIteratorStages));
+    assert.gt(multiIteratorStages.reduce((acc, stage) => acc + stage.works, 0),
+              0,
+              tojson(multiIteratorStages));
+}());
diff --git a/jstests/libs/analyze_plan.js b/jstests/libs/analyze_plan.js
index 6a719d611f6..1df538bd71d 100644
--- a/jstests/libs/analyze_plan.js
+++ b/jstests/libs/analyze_plan.js
@@ -130,8 +130,17 @@ function getAggPlanStages(root, stage) {
         assert(root.stages[0].hasOwnProperty("$cursor"));
         assert(root.stages[0].$cursor.hasOwnProperty("queryPlanner"));
         assert(root.stages[0].$cursor.queryPlanner.hasOwnProperty("winningPlan"));
-        results =
-            results.concat(getPlanStages(root.stages[0].$cursor.queryPlanner.winningPlan, stage));
+
+        // If execution stats are available, then use the execution stats tree. Otherwise use the
+        // plan info from the "queryPlanner" section.
+        if (root.stages[0].$cursor.hasOwnProperty("executionStats")) {
+            assert(root.stages[0].$cursor.executionStats.hasOwnProperty("executionStages"));
+            results = results.concat(
+                getPlanStages(root.stages[0].$cursor.executionStats.executionStages, stage));
+        } else {
+            results = results.concat(
+                getPlanStages(root.stages[0].$cursor.queryPlanner.winningPlan, stage));
+        }
     }
 
     if (root.hasOwnProperty("shards")) {
@@ -140,11 +149,21 @@ function getAggPlanStages(root, stage) {
 
             results = results.concat(getDocumentSources(root.shards[elem].stages));
 
-            assert(root.shards[elem].stages[0].hasOwnProperty("$cursor"));
-            assert(root.shards[elem].stages[0].$cursor.hasOwnProperty("queryPlanner"));
-            assert(root.shards[elem].stages[0].$cursor.queryPlanner.hasOwnProperty("winningPlan"));
-            results = results.concat(
-                getPlanStages(root.shards[elem].stages[0].$cursor.queryPlanner.winningPlan, stage));
+            const firstStage = root.shards[elem].stages[0];
+            assert(firstStage.hasOwnProperty("$cursor"));
+            assert(firstStage.$cursor.hasOwnProperty("queryPlanner"));
+            assert(firstStage.$cursor.queryPlanner.hasOwnProperty("winningPlan"));
+
+            // If execution stats are available, then use the execution stats tree. Otherwise use
+            // the plan info from the "queryPlanner" section.
+            if (firstStage.$cursor.hasOwnProperty("executionStats")) {
+                assert(firstStage.$cursor.executionStats.hasOwnProperty("executionStages"));
+                results = results.concat(
+                    getPlanStages(firstStage.$cursor.executionStats.executionStages, stage));
+            } else {
+                results = results.concat(
+                    getPlanStages(firstStage.$cursor.queryPlanner.winningPlan, stage));
+            }
         }
     }
 
diff --git a/src/mongo/db/exec/multi_iterator.cpp b/src/mongo/db/exec/multi_iterator.cpp
index 69d83775543..6652ab1c798 100644
--- a/src/mongo/db/exec/multi_iterator.cpp
+++ b/src/mongo/db/exec/multi_iterator.cpp
@@ -148,7 +148,7 @@ void MultiIteratorStage::doInvalidate(OperationContext* opCtx,
 
 unique_ptr<PlanStageStats> MultiIteratorStage::getStats() {
     unique_ptr<PlanStageStats> ret =
-        make_unique<PlanStageStats>(CommonStats(kStageType), STAGE_MULTI_ITERATOR);
+        make_unique<PlanStageStats>(_commonStats, STAGE_MULTI_ITERATOR);
     ret->specific = make_unique<CollectionScanStats>();
     return ret;
 }