SERVER-100671 Create an operation scoped memory tracker (#32606)

GitOrigin-RevId: df1bb59b70e8d922d40350c2808352f9281c5fe1

Author: wolffcm

.github/CODEOWNERS

@@ -1522,7 +1522,7 @@ WORKSPACE.bazel @10gen/devprod-build @svc-auto-approve-bot
 /src/mongo/db/memory_tracking/**/OWNERS.yml @10gen/query-execution-staff-leads @10gen/query-integration-staff-leads @svc-auto-approve-bot
 /src/mongo/db/memory_tracking/**/memory_usage_tracker* @10gen/query-execution-aggregation @svc-auto-approve-bot
 /src/mongo/db/memory_tracking/**/op_memory_use* @10gen/query-integration-observability @svc-auto-approve-bot
-/src/mongo/db/memory_tracking/**/operation_memory_usage_tracker.h @10gen/query-integration-observability @svc-auto-approve-bot
+/src/mongo/db/memory_tracking/**/operation_memory_usage_tracker* @10gen/query-integration-observability @svc-auto-approve-bot
 
 # The following patterns are parsed from ./src/mongo/db/modules/enterprise/OWNERS.yml
 /src/mongo/db/modules/enterprise/BUILD.bazel @10gen/devprod-build @svc-auto-approve-bot

jstests/aggregation/sources/setWindowFields/explain.js

@@ -90,7 +90,7 @@ function checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotalMe
     // each function. For the default [unbounded, unbounded] window type, each function uses memory
     // usage comparable to it's $group counterpart.
     let expectedFunctionMemUsages = {
-        count: 72,
+        count: 176,
         push: nDocs * 1024,
         set: 1024,
     };
@@ -115,7 +115,7 @@ function checkExplainResult(pipeline, expectedFunctionMemUsages, expectedTotalMe
         },
     ];
     expectedFunctionMemUsages = {
-        count: 72,
+        count: 176,
         push: (nDocs / nPartitions) * 1056 + 56,  // 56 constant state size. Uses 1056 per document.
         set: 1144,                                // 1024 for the string, rest is constant state.
     };

jstests/noPassthrough/query/spill_to_disk_server_status.js

@@ -36,7 +36,7 @@ assert.commandWorked(
     setParameter(db, "internalQuerySlotBasedExecutionHashAggApproxMemoryUseInBytesBeforeSpill", 1));
 // Spilling memory threshold for $setWindowFields
 assert.commandWorked(setParameter(
-    db, "internalDocumentSourceSetWindowFieldsMaxMemoryBytes", isSbeEnabled ? 129 : 264));
+    db, "internalDocumentSourceSetWindowFieldsMaxMemoryBytes", isSbeEnabled ? 129 : 392));
 // Spilling memory threshold for $lookup
 assert.commandWorked(setParameter(
     db, "internalQuerySlotBasedExecutionHashLookupApproxMemoryUseInBytesBeforeSpill", 1));

src/mongo/db/curop.cpp

@@ -482,13 +482,16 @@ void CurOp::setMemoryTrackingStats(const int64_t inUseMemoryBytes,
     tassert(9897000,
             "featureFlagQueryMemoryTracking must be turned on before writing memory stats to CurOp",
             feature_flags::gFeatureFlagQueryMemoryTracking.isEnabled());
-    _inUseMemoryBytes.fetchAndAdd(inUseMemoryBytes);
-
     // We recompute the max here (the memory tracker that calls this method will already computing
     // the max) in order to avoid writing to the atomic if the max does not change.
+    //
+    // Set the max first, so that we can maintain the invariant that the max is always equal to or
+    // greater than the current in-use tally.
     if (maxUsedMemoryBytes > _maxUsedMemoryBytes.load()) {
-        _maxUsedMemoryBytes.swap(maxUsedMemoryBytes);
+        _maxUsedMemoryBytes.store(maxUsedMemoryBytes);
     }
+
+    _inUseMemoryBytes.store(inUseMemoryBytes);
 }
 
 void CurOp::setNS(WithLock, NamespaceString nss) {

src/mongo/db/curop_test.cpp

@@ -412,12 +412,12 @@ TEST(CurOpTest, ShouldUpdateMemoryStats) {
     ASSERT_EQ(15, curop->getMaxUsedMemoryBytes());
 
     // The max memory usage is updated if the new max is greater than the current max.
-    curop->setMemoryTrackingStats(11 /*currentMemoryBytes*/, 20 /*maxUsedMemoryBytes*/);
+    curop->setMemoryTrackingStats(21 /*currentMemoryBytes*/, 20 /*maxUsedMemoryBytes*/);
     ASSERT_EQ(21, curop->getInUseMemoryBytes());
     ASSERT_EQ(20, curop->getMaxUsedMemoryBytes());
 
     // The max memory usage is not updated if the new max is not greater than the current max.
-    curop->setMemoryTrackingStats(10 /*currentMemoryBytes*/, 15 /*maxUsedMemoryBytes*/);
+    curop->setMemoryTrackingStats(31 /*currentMemoryBytes*/, 15 /*maxUsedMemoryBytes*/);
     ASSERT_EQ(31, curop->getInUseMemoryBytes());
     ASSERT_EQ(20, curop->getMaxUsedMemoryBytes());
 }

src/mongo/db/memory_tracking/BUILD.bazel

@@ -13,13 +13,16 @@ mongo_cc_library(
     name = "memory_tracking",
     srcs = [
         "op_memory_use.cpp",
+        "operation_memory_usage_tracker.cpp",
     ],
     hdrs = [
         "memory_usage_tracker.h",
         "op_memory_use.h",
         "operation_memory_usage_tracker.h",
     ],
     deps = [
+        "//src/mongo/db:commands",
+        "//src/mongo/db:server_feature_flags",
         "//src/mongo/db:service_context",
     ],
 )
@@ -29,10 +32,14 @@ mongo_cc_unit_test(
     srcs = [
         "memory_usage_tracker_test.cpp",
         "op_memory_use_test.cpp",
+        "operation_memory_usage_tracker_test.cpp",
     ],
     tags = ["mongo_unittest_second_group"],
     deps = [
         ":memory_tracking",
         "//src/mongo/db:service_context_test_fixture",
+        "//src/mongo/db/pipeline:aggregation_context_fixture",
+        "//src/mongo/db/pipeline:document_source_mock",
+        "//src/mongo/idl:server_parameter_test_util",
     ],
 )

src/mongo/db/memory_tracking/OWNERS.yml

@@ -14,6 +14,6 @@ filters:
   - "op_memory_use*":
     approvers:
       - 10gen/query-integration-observability
-  - "operation_memory_usage_tracker.h":
+  - "operation_memory_usage_tracker*":
     approvers:
       - 10gen/query-integration-observability

src/mongo/db/memory_tracking/memory_usage_tracker.h

@@ -53,9 +53,11 @@ class SimpleMemoryUsageTracker {
     SimpleMemoryUsageTracker(SimpleMemoryUsageTracker* base, int64_t maxAllowedMemoryUsageBytes)
         : _base(base), _maxAllowedMemoryUsageBytes(maxAllowedMemoryUsageBytes) {}
 
-    SimpleMemoryUsageTracker(int64_t maxAllowedMemoryUsageBytes)
+    explicit SimpleMemoryUsageTracker(int64_t maxAllowedMemoryUsageBytes)
         : SimpleMemoryUsageTracker(nullptr, maxAllowedMemoryUsageBytes) {}
 
+    SimpleMemoryUsageTracker() : SimpleMemoryUsageTracker(std::numeric_limits<int64_t>::max()) {}
+
     void add(int64_t diff) {
         _currentMemoryBytes += diff;
         tassert(6128100,
@@ -68,6 +70,10 @@ class SimpleMemoryUsageTracker {
         if (_base) {
             _base->add(diff);
         }
+
+        if (_doExtraBookkeeping) {
+            _doExtraBookkeeping(_currentMemoryBytes, _maxMemoryBytes);
+        }
     }
 
     void set(int64_t total) {
@@ -90,6 +96,15 @@ class SimpleMemoryUsageTracker {
         return _maxAllowedMemoryUsageBytes;
     }
 
+protected:
+    /**
+     * Provide an extra function that is called whenever add() is invoked. Let it be set via this
+     * method instead in the constructor to allow subclasses to capture "this."
+     */
+    void setDoExtraBookkeeping(std::function<void(int64_t, int64_t)> doExtraBookkeeping) {
+        _doExtraBookkeeping = std::move(doExtraBookkeeping);
+    }
+
 private:
     SimpleMemoryUsageTracker* _base = nullptr;
 
@@ -99,6 +114,11 @@ class SimpleMemoryUsageTracker {
     int64_t _currentMemoryBytes = 0;
 
     int64_t _maxAllowedMemoryUsageBytes;
+
+    // Allow for some extra bookkeeping to be done when add() is called. If set, this function will
+    // be invoked with _currentMemoryBytes and _maxMemoryBytes. This mechanism exists to avoid
+    // making add() virtual, since it has been shown to have an effect on performance in some cases.
+    std::function<void(int64_t, int64_t)> _doExtraBookkeeping;
 };
 
 /**
@@ -120,8 +140,13 @@ class SimpleMemoryUsageTracker {
  */
 class MemoryUsageTracker {
 public:
+    MemoryUsageTracker(SimpleMemoryUsageTracker* baseParent,
+                       bool allowDiskUse = false,
+                       int64_t maxMemoryUsageBytes = 0)
+        : _allowDiskUse(allowDiskUse), _baseTracker(baseParent, maxMemoryUsageBytes) {}
+
     MemoryUsageTracker(bool allowDiskUse = false, int64_t maxMemoryUsageBytes = 0)
-        : _allowDiskUse(allowDiskUse), _baseTracker(nullptr, maxMemoryUsageBytes) {}
+        : MemoryUsageTracker(nullptr, allowDiskUse, maxMemoryUsageBytes) {}
 
     /**
      * Sets the new total for 'name', and updates the current total memory usage.
@@ -198,10 +223,6 @@ class MemoryUsageTracker {
     }
 
 private:
-    SimpleMemoryUsageTracker& _tracker() {
-        return _baseTracker;
-    }
-
     static absl::string_view _key(StringData s) {
         return {s.rawData(), s.size()};
     }

src/mongo/db/memory_tracking/op_memory_use_test.cpp

@@ -30,6 +30,7 @@
 #include "mongo/db/memory_tracking/op_memory_use.h"
 #include "mongo/db/memory_tracking/operation_memory_usage_tracker.h"
 #include "mongo/db/service_context_test_fixture.h"
+#include "mongo/idl/server_parameter_test_util.h"
 
 namespace mongo {
 namespace {
@@ -46,17 +47,22 @@ class MemoryTrackingTest : public ServiceContextTest {
 };
 
 TEST_F(MemoryTrackingTest, OpMemoryUseCanAttachToOpCtx) {
+    RAIIServerParameterControllerForTest featureFlagController("featureFlagQueryMemoryTracking",
+                                                               true);
+
     auto opCtx = makeOpCtx();
 
-    // Attach mock memory usage stats to the opCtx and validate that the decoration was set.
-    auto mockMemoryStats = std::make_unique<OperationMemoryUsageTracker>(10 /*currentMemoryBytes*/,
-                                                                         15 /*maxUsedMemoryBytes*/);
-    OpMemoryUse::operationMemoryAggregator(opCtx.get()) = std::move(mockMemoryStats);
+    auto tracker = std::make_unique<OperationMemoryUsageTracker>(opCtx.get());
+    tracker->add(15);
+    tracker->add(-5);
+
+    OpMemoryUse::operationMemoryAggregator(opCtx.get()) = std::move(tracker);
     OperationMemoryUsageTracker* memoryTracker =
         OpMemoryUse::operationMemoryAggregator(opCtx.get()).get();
     ASSERT(memoryTracker);
     ASSERT_EQ(memoryTracker->currentMemoryBytes(), 10);
-    ASSERT_EQ(memoryTracker->maxUsedMemoryBytes(), 15);
+    ASSERT_EQ(memoryTracker->maxMemoryBytes(), 15);
 }
+
 }  // namespace
 }  // namespace mongo

src/mongo/db/memory_tracking/operation_memory_usage_tracker.cpp

@@ -0,0 +1,78 @@
+/**
+ *    Copyright (C) 2025-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/db/memory_tracking/operation_memory_usage_tracker.h"
+#include "mongo/db/memory_tracking/op_memory_use.h"
+
+namespace mongo {
+
+/**
+ * Return the OperationMemoryUsageTracker for this operation. If we haven't yet created one, do it
+ * now.
+ */
+OperationMemoryUsageTracker* OperationMemoryUsageTracker::getOperationMemoryUsageTracker(
+    OperationContext* opCtx) {
+    OperationMemoryUsageTracker* opTracker = OpMemoryUse::operationMemoryAggregator(opCtx).get();
+    if (!opTracker) {
+        auto uniqueTracker = std::make_unique<OperationMemoryUsageTracker>(opCtx);
+        opTracker = uniqueTracker.get();
+        opTracker->setDoExtraBookkeeping(
+            [opTracker](int64_t currentMemoryBytes, int64_t maxUsedMemoryBytes) {
+                CurOp::get(opTracker->_opCtx)
+                    ->setMemoryTrackingStats(currentMemoryBytes, maxUsedMemoryBytes);
+            });
+        OpMemoryUse::operationMemoryAggregator(opCtx) = std::move(uniqueTracker);
+    }
+
+    return opTracker;
+}
+
+SimpleMemoryUsageTracker OperationMemoryUsageTracker::createSimpleMemoryUsageTrackerForStage(
+    const ExpressionContext& expCtx, int64_t maxMemoryUsageBytes) {
+    if (!feature_flags::gFeatureFlagQueryMemoryTracking.isEnabled()) {
+        return SimpleMemoryUsageTracker{maxMemoryUsageBytes};
+    }
+
+    OperationContext* opCtx = expCtx.getOperationContext();
+    OperationMemoryUsageTracker* opTracker = getOperationMemoryUsageTracker(opCtx);
+    return SimpleMemoryUsageTracker{opTracker, maxMemoryUsageBytes};
+}
+
+MemoryUsageTracker OperationMemoryUsageTracker::createMemoryUsageTrackerForStage(
+    const ExpressionContext& expCtx, bool allowDiskUse, int64_t maxMemoryUsageBytes) {
+    if (!feature_flags::gFeatureFlagQueryMemoryTracking.isEnabled()) {
+        return MemoryUsageTracker{allowDiskUse, maxMemoryUsageBytes};
+    }
+
+    OperationContext* opCtx = expCtx.getOperationContext();
+    OperationMemoryUsageTracker* opTracker = getOperationMemoryUsageTracker(opCtx);
+    return MemoryUsageTracker{opTracker, allowDiskUse, maxMemoryUsageBytes};
+}
+
+}  // namespace mongo

src/mongo/db/memory_tracking/operation_memory_usage_tracker.h

@@ -29,26 +29,42 @@
 #pragma once
 #include <cstdint>
 
+#include "mongo/db/curop.h"
+#include "mongo/db/memory_tracking/memory_usage_tracker.h"
+
 namespace mongo {
-// TODO SERVER-100671 Replace this.
-class OperationMemoryUsageTracker {
+
+/**
+ * A memory usage tracker class that aggregates memory statistics for the entire operation. Stages
+ * that track memory will report to an instance of this class, which in turn will update statistics
+ * in CurOp.
+ */
+class OperationMemoryUsageTracker : public SimpleMemoryUsageTracker {
+    OperationMemoryUsageTracker() = delete;
+
 public:
-    OperationMemoryUsageTracker(int64_t currentMemoryBytes = 0, int64_t maxUsedMemoryBytes = 0)
-        : _currentMemoryBytes(currentMemoryBytes), _maxUsedMemoryBytes(maxUsedMemoryBytes) {}
+    /**
+     * When constructing a stage containing a SimpleMemoryUsageTracker, use this method to ensure
+     * that we aggregate operation-wide memory stats.
+     */
+    static SimpleMemoryUsageTracker createSimpleMemoryUsageTrackerForStage(
+        const ExpressionContext& expCtx,
+        int64_t maxMemoryUsageBytes = std::numeric_limits<int64_t>::max());
 
-    int64_t currentMemoryBytes() const {
-        return _currentMemoryBytes;
-    }
+    /**
+     * When constructing a stage containing a MemoryUsageTracker, use this method to ensure that we
+     * aggregate operation-wide memory stats.
+     */
+    static MemoryUsageTracker createMemoryUsageTrackerForStage(
+        const ExpressionContext& expCtx,
+        bool allowDiskUse = false,
+        int64_t maxMemoryUsageBytes = std::numeric_limits<int64_t>::max());
 
-    int64_t maxUsedMemoryBytes() const {
-        return _maxUsedMemoryBytes;
-    }
+    explicit OperationMemoryUsageTracker(OperationContext* opCtx) : _opCtx(opCtx) {}
 
 private:
-    // Current amount of memory in use by all blocking stages.
-    int64_t _currentMemoryBytes = 0;
+    static OperationMemoryUsageTracker* getOperationMemoryUsageTracker(OperationContext* opCtx);
 
-    // High-water mark: the highest amount of memory that has been allocated at one time so far.
-    int64_t _maxUsedMemoryBytes = 0;
+    OperationContext* _opCtx;
 };
 }  // namespace mongo