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

2 files changed, 22 insertions, 5 deletions

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;