SERVER-73311 Add group spill stats to serverStatus for SBE

11 files changed, 82 insertions, 42 deletions

diff --git a/jstests/aggregation/spill_to_disk.js b/jstests/aggregation/spill_to_disk.js
index 568d973e988..8d37777c75b 100644
--- a/jstests/aggregation/spill_to_disk.js
+++ b/jstests/aggregation/spill_to_disk.js
@@ -45,7 +45,7 @@ function assertSpillingOccurredInSbeExplain(groupStats) {
     assert(groupStats);
     assert(groupStats.hasOwnProperty("usedDisk"), groupStats);
     assert(groupStats.usedDisk, groupStats);
-    assert.gt(groupStats.numSpills, 0, groupStats);
+    assert.gt(groupStats.spills, 0, groupStats);
     assert.gt(groupStats.spilledRecords, 0, groupStats);
     assert.gt(groupStats.spilledDataStorageSize, 0, groupStats);
 }
diff --git a/jstests/noPassthrough/explain_group_stage_exec_stats.js b/jstests/noPassthrough/explain_group_stage_exec_stats.js
index f2ce62ba9a5..8a8774f694a 100644
--- a/jstests/noPassthrough/explain_group_stage_exec_stats.js
+++ b/jstests/noPassthrough/explain_group_stage_exec_stats.js
@@ -67,8 +67,9 @@ function checkGroupStages(stage, expectedAccumMemUsages, isExecExplain, expected
 
     if (isExecExplain) {
         assert(stage.hasOwnProperty("maxAccumulatorMemoryUsageBytes"), stage);
-        assert(stage.hasOwnProperty("spillFileSizeBytes"), stage);
+        assert(stage.hasOwnProperty("spilledDataStorageSize"), stage);
         assert(stage.hasOwnProperty("numBytesSpilledEstimate"), stage);
+        assert(stage.hasOwnProperty("spilledRecords"), stage);
 
         const maxAccmMemUsages = stage["maxAccumulatorMemoryUsageBytes"];
         for (const field of Object.keys(maxAccmMemUsages)) {
@@ -83,19 +84,22 @@ function checkGroupStages(stage, expectedAccumMemUsages, isExecExplain, expected
             }
         }
 
-        const spillFileSizeBytes = stage["spillFileSizeBytes"];
+        const spilledDataStorageSize = stage["spilledDataStorageSize"];
         const numBytesSpilledEstimate = stage["numBytesSpilledEstimate"];
+        const spilledRecords = stage["spilledRecords"];
         if (stage.usedDisk) {
             // We cannot compute the size of the spill file, so assert that it is non-zero if we
             // have spilled.
-            assert.gt(spillFileSizeBytes, 0, stage);
+            assert.gt(spilledDataStorageSize, 0, stage);
 
             // The number of bytes spilled, on the other hand, is at least as much as the
             // accumulator memory usage.
             assert.gt(numBytesSpilledEstimate, totalAccumMemoryUsageBytes);
+            assert.gt(spilledRecords, 0, stage);
         } else {
-            assert.eq(spillFileSizeBytes, 0, stage);
+            assert.eq(spilledDataStorageSize, 0, stage);
             assert.eq(numBytesSpilledEstimate, 0, stage);
+            assert.eq(spilledRecords, 0, stage);
         }
 
         // Don't verify spill count for debug builds, since for debug builds a spill occurs on every
@@ -109,8 +113,9 @@ function checkGroupStages(stage, expectedAccumMemUsages, isExecExplain, expected
         assert(!stage.hasOwnProperty("usedDisk"), stage);
         assert(!stage.hasOwnProperty("spills"), stage);
         assert(!stage.hasOwnProperty("maxAccumulatorMemoryUsageBytes"), stage);
-        assert(!stage.hasOwnProperty("spillFileSizeBytes"), stage);
+        assert(!stage.hasOwnProperty("spilledDataStorageSize"), stage);
         assert(!stage.hasOwnProperty("numBytesSpilledEstimate"), stage);
+        assert(!stage.hasOwnProperty("spilledRecords"), stage);
     }
 
     // Add some wiggle room to the total memory used compared to the limit parameter since the check
diff --git a/jstests/noPassthrough/group_spill_metrics.js b/jstests/noPassthrough/group_spill_metrics.js
index 6a946ee25fc..89c6d6072d9 100644
--- a/jstests/noPassthrough/group_spill_metrics.js
+++ b/jstests/noPassthrough/group_spill_metrics.js
@@ -1,10 +1,17 @@
 /**
  * Tests that $group stage reports spill stats when serverStatus is run.
+ *
+ * @tags: [
+ *   # TODO SERVER-73757: Allow this test to run against the inMemory storage engine once ephemeral
+ *   # temporary record stores used for spilling report the correct storage size.
+ *   requires_persistence,
+ * ]
  */
 (function() {
 "use strict";
 
 load("jstests/libs/analyze_plan.js");  // For getAggPlanStage().
+load("jstests/libs/sbe_util.js");      // For checkSBEEnabled.
 
 const conn = MongoRunner.runMongod();
 const db = conn.getDB('test');
@@ -15,6 +22,7 @@ const bigStr = Array(1025).toString();  // 1KB of ','
 const maxMemoryLimitForGroupStage = 1024 * 300;
 const nDocs = 1000;
 const nGroups = 50;
+const isSbeEnabled = checkSBEEnabled(db);
 
 const bulk = coll.initializeUnorderedBulkOp();
 for (let i = 1; i <= nDocs; i++) {
@@ -23,7 +31,6 @@ for (let i = 1; i <= nDocs; i++) {
 assert.commandWorked(bulk.execute());
 
 const pipeline = [
-    {$_internalInhibitOptimization: {}},
     {$match: {a: {$gt: 0}}},
     {$sort: {b: 1}},
     {$group: {_id: "$b", count: {$sum: 1}, push: {$push: "$bigStr"}, set: {$addToSet: "$bigStr"}}},
@@ -34,27 +41,37 @@ const metricsBefore = db.serverStatus().metrics.query.group;
 // Set MaxMemory low to force spill to disk.
 assert.commandWorked(db.adminCommand(
     {setParameter: 1, internalDocumentSourceGroupMaxMemoryBytes: maxMemoryLimitForGroupStage}));
+assert.commandWorked(db.adminCommand({
+    setParameter: 1,
+    internalQuerySlotBasedExecutionHashAggApproxMemoryUseInBytesBeforeSpill:
+        maxMemoryLimitForGroupStage
+}));
 
-const result = getAggPlanStage(coll.explain("executionStats").aggregate(pipeline), "$group");
+const result = coll.explain("executionStats").aggregate(pipeline);
+
+const groupStage =
+    isSbeEnabled ? getAggPlanStage(result, "group") : getAggPlanStage(result, "$group");
 
 const metricsAfter = db.serverStatus().metrics.query.group;
 
-const expectedSpills = result.spills + metricsBefore.spills;
-const expectedSpillFileSizeBytes = result.spillFileSizeBytes + metricsBefore.spillFileSizeBytes;
-const expectedNumBytesSpilledEstimate =
-    result.numBytesSpilledEstimate + metricsBefore.numBytesSpilledEstimate;
+const expectedSpills = groupStage.spills;
+const expectedSpilledDataStorageSize = groupStage.spilledDataStorageSize;
+const expectedSpilledRecords = groupStage.spilledRecords;
 
 assert.gt(metricsAfter.spills, metricsBefore.spills, pipeline);
 
-assert.eq(metricsAfter.spills, expectedSpills, pipeline);
+assert.eq(metricsAfter.spills, expectedSpills + metricsBefore.spills, pipeline);
 
-assert.gt(metricsAfter.spillFileSizeBytes, metricsBefore.spillFileSizeBytes, pipeline);
+assert.gt(metricsAfter.spilledDataStorageSize, metricsBefore.spilledDataStorageSize, pipeline);
 
-assert.eq(metricsAfter.spillFileSizeBytes, expectedSpillFileSizeBytes, pipeline);
+assert.eq(metricsAfter.spilledDataStorageSize,
+          expectedSpilledDataStorageSize + metricsBefore.spilledDataStorageSize,
+          pipeline);
 
-assert.gt(metricsAfter.numBytesSpilledEstimate, metricsBefore.numBytesSpilledEstimate, pipeline);
+assert.gt(metricsAfter.spilledRecords, metricsBefore.spilledRecords, pipeline);
 
-assert.eq(metricsAfter.numBytesSpilledEstimate, expectedNumBytesSpilledEstimate, pipeline);
+assert.eq(
+    metricsAfter.spilledRecords, expectedSpilledRecords + metricsBefore.spilledRecords, pipeline);
 
 MongoRunner.stopMongod(conn);
 }());
diff --git a/jstests/noPassthrough/spill_to_disk_secondary_read.js b/jstests/noPassthrough/spill_to_disk_secondary_read.js
index 8ab9178b4c5..f143aa88524 100644
--- a/jstests/noPassthrough/spill_to_disk_secondary_read.js
+++ b/jstests/noPassthrough/spill_to_disk_secondary_read.js
@@ -93,7 +93,7 @@ const readColl = secondary.getDB("test").foo;
         assert.eq(hashAggGroup.spilledRecords, expectedSpilledRecords, hashAggGroup);
         // We expect each record to be individually spilled, so the number of spill events and the
         // number of spilled records should be equal.
-        assert.eq(hashAggGroup.numSpills, hashAggGroup.spilledRecords, hashAggGroup);
+        assert.eq(hashAggGroup.spills, hashAggGroup.spilledRecords, hashAggGroup);
         assert.gt(hashAggGroup.spilledDataStorageSize, expectedSpilledBytesAtLeast, hashAggGroup);
     } finally {
         assert.commandWorked(secondary.adminCommand({
diff --git a/src/mongo/db/exec/plan_stats.h b/src/mongo/db/exec/plan_stats.h
index 36919dc6be6..4801769c441 100644
--- a/src/mongo/db/exec/plan_stats.h
+++ b/src/mongo/db/exec/plan_stats.h
@@ -1043,13 +1043,15 @@ struct GroupStats : public SpecificStats {
     // The size of the file spilled to disk. Note that this is not the same as the number of bytes
     // spilled to disk, as any data spilled to disk will be compressed before being written to a
     // file.
-    uint64_t spillFileSizeBytes = 0u;
+    uint64_t spilledDataStorageSize = 0u;
 
     // The number of bytes evicted from memory and spilled to disk.
     uint64_t numBytesSpilledEstimate = 0u;
 
     // The number of times that we spilled data to disk while grouping the data.
     uint64_t spills = 0u;
+
+    uint64_t spilledRecords = 0u;
 };
 
 struct DocumentSourceCursorStats : public SpecificStats {
diff --git a/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp b/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp
index b48ed496798..8888929969f 100644
--- a/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp
+++ b/src/mongo/db/exec/sbe/sbe_hash_agg_test.cpp
@@ -447,7 +447,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpill) {
     // Check that the spilling behavior matches the expected.
     auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats());
     ASSERT_FALSE(stats->usedDisk);
-    ASSERT_EQ(0, stats->numSpills);
+    ASSERT_EQ(0, stats->spills);
     ASSERT_EQ(0, stats->spilledRecords);
 
     stage->close();
@@ -514,7 +514,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountSpill) {
     // spilling after estimating that the memory budget is exceeded. These two factors result in
     // fewer expected spills than there are input records, even though only one record fits in
     // memory at a time.
-    ASSERT_EQ(stats->numSpills, 3);
+    ASSERT_EQ(stats->spills, 3);
     // The input has one run of two consecutive values, so we expect to spill as many records as
     // there are input values minus one.
     ASSERT_EQ(stats->spilledRecords, 8);
@@ -588,7 +588,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpillIfNoMemCheck) {
     // Check that it did not spill.
     auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats());
     ASSERT_FALSE(stats->usedDisk);
-    ASSERT_EQ(0, stats->numSpills);
+    ASSERT_EQ(0, stats->spills);
     ASSERT_EQ(0, stats->spilledRecords);
 
     stage->close();
@@ -655,7 +655,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountSpillDouble) {
     // spilling after estimating that the memory budget is exceeded. These two factors result in
     // fewer expected spills than there are input records, even though only one record fits in
     // memory at a time.
-    ASSERT_EQ(stats->numSpills, 3);
+    ASSERT_EQ(stats->spills, 3);
     // The input has one run of two consecutive values, so we expect to spill as many records as
     // there are input values minus one.
     ASSERT_EQ(stats->spilledRecords, 8);
@@ -715,7 +715,7 @@ TEST_F(HashAggStageTest, HashAggBasicCountNoSpillWithNoGroupByDouble) {
     // Check that the spilling behavior matches the expected.
     auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats());
     ASSERT_FALSE(stats->usedDisk);
-    ASSERT_EQ(0, stats->numSpills);
+    ASSERT_EQ(0, stats->spills);
     ASSERT_EQ(0, stats->spilledRecords);
 
     stage->close();
@@ -792,7 +792,7 @@ TEST_F(HashAggStageTest, HashAggMultipleAccSpill) {
     // Check that the spilling behavior matches the expected.
     auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats());
     ASSERT_TRUE(stats->usedDisk);
-    ASSERT_EQ(stats->numSpills, 3);
+    ASSERT_EQ(stats->spills, 3);
     // The input has one run of two consecutive values, so we expect to spill as many records as
     // there are input values minus one.
     ASSERT_EQ(stats->spilledRecords, 8);
@@ -871,7 +871,7 @@ TEST_F(HashAggStageTest, HashAggMultipleAccSpillAllToDisk) {
     auto stats = static_cast<const HashAggStats*>(stage->getSpecificStats());
     ASSERT_TRUE(stats->usedDisk);
     // We expect each incoming value to result in a spill of a single record.
-    ASSERT_EQ(stats->numSpills, 9);
+    ASSERT_EQ(stats->spills, 9);
     ASSERT_EQ(stats->spilledRecords, 9);
 
     stage->close();
diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.cpp b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
index 6d33345ef7c..7d04af90d72 100644
--- a/src/mongo/db/exec/sbe/stages/hash_agg.cpp
+++ b/src/mongo/db/exec/sbe/stages/hash_agg.cpp
@@ -32,6 +32,7 @@
 #include "mongo/db/exec/sbe/expressions/compile_ctx.h"
 #include "mongo/db/exec/sbe/size_estimator.h"
 #include "mongo/db/exec/sbe/util/spilling.h"
+#include "mongo/db/stats/counters.h"
 #include "mongo/db/stats/resource_consumption_metrics.h"
 #include "mongo/db/storage/kv/kv_engine.h"
 #include "mongo/db/storage/storage_engine.h"
@@ -102,6 +103,12 @@ std::unique_ptr<PlanStage> HashAggStage::clone() const {
                                           _forceIncreasedSpilling);
 }
 
+HashAggStage::~HashAggStage() {
+    groupCounters.incrementGroupCounters(_specificStats.spills,
+                                         _specificStats.spilledDataStorageSize,
+                                         _specificStats.spilledRecords);
+}
+
 void HashAggStage::doSaveState(bool relinquishCursor) {
     if (relinquishCursor) {
         if (_rsCursor) {
@@ -313,7 +320,7 @@ void HashAggStage::spill(MemoryCheckData& mcd) {
 
     _ht->clear();
 
-    ++_specificStats.numSpills;
+    ++_specificStats.spills;
 }
 
 // Checks memory usage. Ideally, we'd want to know the exact size of already accumulated data, but
@@ -613,7 +620,7 @@ std::unique_ptr<PlanStageStats> HashAggStage::getStats(bool includeDebugInfo) co
 
         // Spilling stats.
         bob.appendBool("usedDisk", _specificStats.usedDisk);
-        bob.appendNumber("numSpills", _specificStats.numSpills);
+        bob.appendNumber("spills", _specificStats.spills);
         bob.appendNumber("spilledRecords", _specificStats.spilledRecords);
         bob.appendNumber("spilledDataStorageSize", _specificStats.spilledDataStorageSize);
 
diff --git a/src/mongo/db/exec/sbe/stages/hash_agg.h b/src/mongo/db/exec/sbe/stages/hash_agg.h
index 86f7337b7f7..db991952f49 100644
--- a/src/mongo/db/exec/sbe/stages/hash_agg.h
+++ b/src/mongo/db/exec/sbe/stages/hash_agg.h
@@ -93,6 +93,8 @@ public:
                  bool participateInTrialRunTracking = true,
                  bool forceIncreasedSpilling = false);
 
+    virtual ~HashAggStage();
+
     std::unique_ptr<PlanStage> clone() const final;
 
     void prepare(CompileCtx& ctx) final;
diff --git a/src/mongo/db/exec/sbe/stages/plan_stats.h b/src/mongo/db/exec/sbe/stages/plan_stats.h
index 2801c8abb07..cd093715921 100644
--- a/src/mongo/db/exec/sbe/stages/plan_stats.h
+++ b/src/mongo/db/exec/sbe/stages/plan_stats.h
@@ -320,7 +320,7 @@ struct HashAggStats : public SpecificStats {
 
     bool usedDisk{false};
     // The number of times that the entire hash table was spilled.
-    long long numSpills{0};
+    long long spills{0};
     // The number of individual records spilled to disk.
     long long spilledRecords{0};
     // An estimate, in bytes, of the size of the final spill table after all spill events have taken
diff --git a/src/mongo/db/pipeline/document_source_group_base.cpp b/src/mongo/db/pipeline/document_source_group_base.cpp
index 8532b4d6d19..791b0ce7151 100644
--- a/src/mongo/db/pipeline/document_source_group_base.cpp
+++ b/src/mongo/db/pipeline/document_source_group_base.cpp
@@ -83,7 +83,8 @@ using std::shared_ptr;
 using std::vector;
 
 DocumentSourceGroupBase::~DocumentSourceGroupBase() {
-    groupCounters.incrementGroupCounters(_stats);
+    groupCounters.incrementGroupCounters(
+        _stats.spills, _stats.spilledDataStorageSize, _stats.spilledRecords);
 }
 
 Value DocumentSourceGroupBase::serialize(boost::optional<ExplainOptions::Verbosity> explain) const {
@@ -134,9 +135,11 @@ Value DocumentSourceGroupBase::serialize(boost::optional<ExplainOptions::Verbosi
             Value(static_cast<long long>(_stats.totalOutputDataSizeBytes));
         out["usedDisk"] = Value(_stats.spills > 0);
         out["spills"] = Value(static_cast<long long>(_stats.spills));
-        out["spillFileSizeBytes"] = Value(static_cast<long long>(_stats.spillFileSizeBytes));
+        out["spilledDataStorageSize"] =
+            Value(static_cast<long long>(_stats.spilledDataStorageSize));
         out["numBytesSpilledEstimate"] =
             Value(static_cast<long long>(_stats.numBytesSpilledEstimate));
+        out["spilledRecords"] = Value(static_cast<long long>(_stats.spilledRecords));
     }
 
     return out.freezeToValue();
@@ -592,6 +595,7 @@ void DocumentSourceGroupBase::resetReadyGroups() {
 void DocumentSourceGroupBase::spill() {
     _stats.spills++;
     _stats.numBytesSpilledEstimate += _memoryTracker.currentMemoryBytes();
+    _stats.spilledRecords += _groups->size();
 
     vector<const GroupsMap::value_type*> ptrs;  // using pointers to speed sorting
     ptrs.reserve(_groups->size());
@@ -644,7 +648,7 @@ void DocumentSourceGroupBase::spill() {
 
     _sortedFiles.emplace_back(writer.done());
     if (_spillStats) {
-        _stats.spillFileSizeBytes = _spillStats->bytesSpilled();
+        _stats.spilledDataStorageSize = _spillStats->bytesSpilled();
     }
 }
 
diff --git a/src/mongo/db/stats/counters.h b/src/mongo/db/stats/counters.h
index 23452dbb8a8..ccefe81e6d0 100644
--- a/src/mongo/db/stats/counters.h
+++ b/src/mongo/db/stats/counters.h
@@ -433,19 +433,22 @@ class GroupCounters {
 public:
     GroupCounters() = default;
 
-    void incrementGroupCounters(const GroupStats& stats) {
-        spills.increment(stats.spills);
-        spillFileSizeBytes.increment(stats.spillFileSizeBytes);
-        numBytesSpilledEstimate.increment(stats.numBytesSpilledEstimate);
+    void incrementGroupCounters(uint64_t spills,
+                                uint64_t spilledDataStorageSize,
+                                uint64_t spilledRecords) {
+        groupSpills.increment(spills);
+        groupSpilledDataStorageSize.increment(spilledDataStorageSize);
+        groupSpilledRecords.increment(spilledRecords);
     }
 
-    CounterMetric spills{
+    // Counters tracking group stats across all execution engines.
+    CounterMetric groupSpills{
         "query.group.spills"};  // The total number of spills to disk from group stages.
-    CounterMetric spillFileSizeBytes{
-        "query.group.spillFileSizeBytes"};  // The size of the file spilled to disk.
-    CounterMetric numBytesSpilledEstimate{
-        "query.group.numBytesSpilledEstimate"};  // The number of bytes evicted from memory and
-                                                 // spilled to disk.
+    CounterMetric groupSpilledDataStorageSize{
+        "query.group.spilledDataStorageSize"};  // The size of the file or RecordStore spilled to
+                                                // disk.
+    CounterMetric groupSpilledRecords{
+        "query.group.spilledRecords"};  // The number of records spilled to disk.
 };
 extern GroupCounters groupCounters;