SERVER-55290 Add memory tracking to sum and avg window functions

2 files changed, 10 insertions, 2 deletions

diff --git a/src/mongo/db/pipeline/window_function/window_function_avg.h b/src/mongo/db/pipeline/window_function/window_function_avg.h
index 549b2994efe..2dc0d09e4be 100644
--- a/src/mongo/db/pipeline/window_function/window_function_avg.h
+++ b/src/mongo/db/pipeline/window_function/window_function_avg.h
@@ -38,7 +38,11 @@ namespace mongo {
 
 class WindowFunctionAvg final : public RemovableSum {
 public:
-    explicit WindowFunctionAvg(ExpressionContext* const expCtx) : RemovableSum(expCtx), _count(0) {}
+    explicit WindowFunctionAvg(ExpressionContext* const expCtx) : RemovableSum(expCtx), _count(0) {
+        // Note that RemovableSum manages the memory usage tracker directly for calls to add/remove.
+        // Here we only add the members that this class holds.
+        _memUsageBytes += sizeof(long long);
+    }
 
     static std::unique_ptr<WindowFunctionState> create(ExpressionContext* const expCtx) {
         return std::make_unique<WindowFunctionAvg>(expCtx);
@@ -92,6 +96,7 @@ public:
     void reset() {
         RemovableSum::reset();
         _count = 0;
+        _memUsageBytes += sizeof(long long);
     }
 
 private:
diff --git a/src/mongo/db/pipeline/window_function/window_function_sum.h b/src/mongo/db/pipeline/window_function/window_function_sum.h
index f2bafc4a262..6f4b6bfaca9 100644
--- a/src/mongo/db/pipeline/window_function/window_function_sum.h
+++ b/src/mongo/db/pipeline/window_function/window_function_sum.h
@@ -45,7 +45,9 @@ protected:
           _negInfiniteValueCount(0),
           _nanCount(0),
           _doubleCount(0),
-          _decimalCount(0) {}
+          _decimalCount(0) {
+        _memUsageBytes = sizeof(*this) + _sumAcc->getMemUsage();
+    }
 
 public:
     static Value getDefault() {
@@ -69,6 +71,7 @@ public:
         _nanCount = 0;
         _doubleCount = 0;
         _decimalCount = 0;
+        _memUsageBytes = sizeof(*this) + _sumAcc->getMemUsage();
     }
 
 private: