SERVER-55789 Add explain metric for peak memory usage of $setWindowFields stage

3 files changed, 96 insertions, 16 deletions

diff --git a/jstests/aggregation/sources/setWindowFields/explain.js b/jstests/aggregation/sources/setWindowFields/explain.js
index fbae93889d8..c820d267130 100644
--- a/jstests/aggregation/sources/setWindowFields/explain.js
+++ b/jstests/aggregation/sources/setWindowFields/explain.js
@@ -1,6 +1,8 @@
 /**
  * Tests that $setWindowFields stage reports memory footprint per function when explain is run
- * with verbosities "executionStats" and "allPlansExecution".
+ * with verbosities "executionStats" and "allPlansExecution". Also tests that the explain output
+ * includes a metric for peak memory usage across the entire stage, including each individual
+ * function as well as any other internal state.
  *
  * @tags: [assumes_against_mongod_not_mongos]
  */
@@ -35,9 +37,10 @@ assert.commandWorked(bulk.execute());
  * expected.
  * - 'stages' is an array of the explain output of $setWindowFields stages.
  * - 'expectedFunctionMemUsages' is used to check the memory footprint stats for each function.
+ * - 'expectedTotalMemUsage' is used to check the peak memory footprint for the entire stage.
  * - 'verbosity' indicates the explain verbosity used.
  */
-function checkExplainResult(pipeline, expectedFunctionMemUsages, verbosity) {
+function checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotalMemUsage, verbosity) {
     const stages =
         getAggPlanStages(coll.explain(verbosity).aggregate(pipeline), "$_internalSetWindowFields");
     for (let stage of stages) {
@@ -58,8 +61,15 @@ function checkExplainResult(pipeline, expectedFunctionMemUsages, verbosity) {
                               "mismatch for function '" + field + "': " + tojson(stage));
                 }
             }
+            assert.gt(stage["maxTotalMemoryUsageBytes"],
+                      expectedTotalMemUsage,
+                      "Incorrect total mem usage: " + tojson(stage));
+            assert.lt(stage["maxTotalMemoryUsageBytes"],
+                      2 * expectedTotalMemUsage,
+                      "Incorrect total mem usage: " + tojson(stage));
         } else {
             assert(!stage.hasOwnProperty("maxFunctionMemoryUsageBytes"), stage);
+            assert(!stage.hasOwnProperty("maxTotalMemoryUsageBytes"), stage);
         }
     }
 }
@@ -73,7 +83,7 @@ function checkExplainResult(pipeline, expectedFunctionMemUsages, verbosity) {
     ];
     const stages = getAggPlanStages(coll.explain("queryPlanner").aggregate(pipeline),
                                     "$_internalSetWindowFields");
-    checkExplainResult(stages, {}, "queryPlanner");
+    checkExplainResult(stages, {}, 0, "queryPlanner");
 })();
 
 (function testUnboundedMemUsage() {
@@ -93,8 +103,15 @@ function checkExplainResult(pipeline, expectedFunctionMemUsages, verbosity) {
         set: 1024,
     };
 
-    checkExplainResult(pipeline, expectedFunctionMemUsages, "executionStats");
-    checkExplainResult(pipeline, expectedFunctionMemUsages, "allPlansExecution");
+    // The total mem usage for unbounded windows is the total from each function as well as the size
+    // of all documents in the partition.
+    let expectedTotal = nDocs * docSize;
+    for (let func in expectedFunctionMemUsages) {
+        expectedTotal += expectedFunctionMemUsages[func];
+    }
+
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "executionStats");
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "allPlansExecution");
 
     // Test that the memory footprint is reduced with partitioning.
     pipeline = [
@@ -110,9 +127,13 @@ function checkExplainResult(pipeline, expectedFunctionMemUsages, verbosity) {
         push: (nDocs / nPartitions) * 1024,
         set: 1024,
     };
+    expectedTotal = (nDocs / nPartitions) * docSize;
+    for (let func in expectedFunctionMemUsages) {
+        expectedTotal += expectedFunctionMemUsages[func];
+    }
 
-    checkExplainResult(pipeline, expectedFunctionMemUsages, "executionStats");
-    checkExplainResult(pipeline, expectedFunctionMemUsages, "allPlansExecution");
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "executionStats");
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "allPlansExecution");
 })();
 
 (function testSlidingWindowMemUsage() {
@@ -130,8 +151,16 @@ function checkExplainResult(pipeline, expectedFunctionMemUsages, verbosity) {
             160,  // 10x64-bit integer values per window, and 160 for the $sum state.
     };
 
-    checkExplainResult(pipeline, expectedFunctionMemUsages, "executionStats");
-    checkExplainResult(pipeline, expectedFunctionMemUsages, "allPlansExecution");
+    // TODO SERVER-55786: Fix memory tracking in PartitionIterator when documents are
+    // released from the in-memory cache. This should be proportional to the size of the window not
+    // the number of documents in the partition.
+    let expectedTotal = nDocs * docSize;
+    for (let func in expectedFunctionMemUsages) {
+        expectedTotal += expectedFunctionMemUsages[func];
+    }
+
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "executionStats");
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "allPlansExecution");
 
     // Adding partitioning doesn't change the peak memory usage.
     pipeline = [
@@ -143,7 +172,41 @@ function checkExplainResult(pipeline, expectedFunctionMemUsages, verbosity) {
             }
         },
     ];
-    checkExplainResult(pipeline, expectedFunctionMemUsages, "executionStats");
-    checkExplainResult(pipeline, expectedFunctionMemUsages, "allPlansExecution");
+
+    // TODO SERVER-55786: This should not be needed once the memory tracking is fixed in the
+    // iterator. For now, the iterator handles the tracking correctly across partition boundaries so
+    // we need to adjust the expected total.
+    expectedTotal = (nDocs / nPartitions) * docSize;
+    for (let func in expectedFunctionMemUsages) {
+        expectedTotal += expectedFunctionMemUsages[func];
+    }
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "executionStats");
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "allPlansExecution");
+})();
+
+(function testRangeBasedWindowMemUsage() {
+    const maxDocsInWindow = 20;
+    let pipeline = [
+        {
+            $setWindowFields: {
+                sortBy: {_id: 1},
+                output: {pushArray: {$push: "$bigStr", window: {range: [-10, 9]}}}
+            }
+        },
+    ];
+    // The memory usage is doubled since both the executor and the function state have copies of the
+    // large string.
+    const expectedFunctionMemUsages = {pushArray: 1024 * maxDocsInWindow * 2};
+
+    // TODO SERVER-55786: Fix memory tracking in PartitionIterator when documents are
+    // released from the in-memory cache. This should be proportional to the size of the window not
+    // the number of documents in the partition.
+    let expectedTotal = nDocs * docSize;
+    for (let func in expectedFunctionMemUsages) {
+        expectedTotal += expectedFunctionMemUsages[func];
+    }
+
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "executionStats");
+    checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotal, "allPlansExecution");
 })();
 }());
diff --git a/src/mongo/db/pipeline/document_source_set_window_fields.cpp b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
index 9a3e3d3498f..fae2ae4f4a3 100644
--- a/src/mongo/db/pipeline/document_source_set_window_fields.cpp
+++ b/src/mongo/db/pipeline/document_source_set_window_fields.cpp
@@ -262,6 +262,8 @@ Value DocumentSourceInternalSetWindowFields::serialize(
         }
 
         out["maxFunctionMemoryUsageBytes"] = Value(md.freezeToValue());
+        out["maxTotalMemoryUsageBytes"] =
+            Value(static_cast<long long>(_memoryTracker.maxMemoryBytes()));
     }
 
     return Value(out.freezeToValue());
@@ -351,7 +353,7 @@ DocumentSource::GetNextResult DocumentSourceInternalSetWindowFields::doGetNext()
         case PartitionIterator::AdvanceResult::kNewPartition:
             // We've advanced to a new partition, reset the state of every function as well as the
             // memory tracker.
-            _memoryTracker.reset();
+            _memoryTracker.resetCurrent();
             for (auto&& [fieldName, function] : _executableOutputs) {
                 function->reset();
             }
diff --git a/src/mongo/db/pipeline/memory_usage_tracker.h b/src/mongo/db/pipeline/memory_usage_tracker.h
index 76de09383ac..91770cdd2ad 100644
--- a/src/mongo/db/pipeline/memory_usage_tracker.h
+++ b/src/mongo/db/pipeline/memory_usage_tracker.h
@@ -82,7 +82,7 @@ public:
     void set(StringData functionName, uint64_t total) {
         auto oldFuncUsage = _functionMemoryTracker[functionName].currentMemoryBytes();
         _functionMemoryTracker[functionName].set(total);
-        _memoryUsageBytes += total - oldFuncUsage;
+        update(total - oldFuncUsage);
     }
 
     /**
@@ -90,12 +90,16 @@ public:
      */
     void set(uint64_t total) {
         _memoryUsageBytes = total;
+        if (_memoryUsageBytes > _maxMemoryUsageBytes) {
+            _maxMemoryUsageBytes = _memoryUsageBytes;
+        }
     }
 
     /**
-     * Resets both the total memory usage as well as the per-function memory usage.
+     * Resets both the total memory usage as well as the per-function memory usage, but retains the
+     * current value for maximum total memory usage.
      */
-    void reset() {
+    void resetCurrent() {
         _memoryUsageBytes = 0;
         for (auto& [_, funcTracker] : _functionMemoryTracker) {
             funcTracker.set(0);
@@ -119,12 +123,22 @@ public:
      */
     void update(StringData name, int diff) {
         _functionMemoryTracker[name].update(diff);
-        _memoryUsageBytes += diff;
+        update(diff);
+    }
+
+    /**
+     * Updates total memory usage.
+     */
+    void update(int diff) {
+        set(_memoryUsageBytes + diff);
     }
 
     auto currentMemoryBytes() const {
         return _memoryUsageBytes;
     }
+    auto maxMemoryBytes() const {
+        return _maxMemoryUsageBytes;
+    }
 
     const bool _allowDiskUse;
     const size_t _maxAllowedMemoryUsageBytes;
@@ -132,6 +146,7 @@ public:
 private:
     // Tracks current memory used.
     size_t _memoryUsageBytes = 0;
+    size_t _maxMemoryUsageBytes = 0;
 
     // Tracks memory consumption per function using the output field name as a key.
     StringMap<PerFunctionMemoryTracker> _functionMemoryTracker;