SERVER-98971 Decorate OperationContext with an OpMemoryUse Object (#30973)

GitOrigin-RevId: 86ebbc3

Author: erl-ang

src/mongo/BUILD.bazel

@@ -179,8 +179,6 @@ mongo_cc_library(
         "//src/mongo/util:assert_util.cpp",
         "//src/mongo/util:base64.cpp",
         "//src/mongo/util:boost_assert_impl.cpp",
-        "//src/mongo/util:chunked_memory_aggregator.cpp",
-        "//src/mongo/util:concurrent_memory_aggregator.cpp",
         "//src/mongo/util:duration.cpp",
         "//src/mongo/util:errno_util.cpp",
         "//src/mongo/util:exception_filter_win32.cpp",
@@ -337,9 +335,7 @@ mongo_cc_library(
         "//src/mongo/util:assert_util_core.h",
         "//src/mongo/util:base64.h",
         "//src/mongo/util:bufreader.h",
-        "//src/mongo/util:chunked_memory_aggregator.h",
         "//src/mongo/util:clock_source.h",
-        "//src/mongo/util:concurrent_memory_aggregator.h",
         "//src/mongo/util:ctype.h",
         "//src/mongo/util:debug_util.h",
         "//src/mongo/util:decimal_counter.h",
@@ -413,7 +409,6 @@ mongo_cc_library(
         "//conditions:default": [],
     }),
     deps = [
-        "//src/mongo/db/memory_tracking",
         "//src/mongo/stdx",
         "//src/mongo/util:boost_assert_shim",
         "//src/mongo/util:debugger",

src/mongo/db/commands/BUILD.bazel

@@ -644,6 +644,7 @@ mongo_cc_library(
         "//src/mongo/db/matcher/schema:encrypt_schema_types.h",
         "//src/mongo/db/matcher/schema:expression_internal_schema_allowed_properties.h",
         "//src/mongo/db/matcher/schema:json_pointer.h",
+        "//src/mongo/db/memory_tracking:memory_usage_tracker.h",
         "//src/mongo/db/pipeline:accumulation_statement.h",
         "//src/mongo/db/pipeline:accumulator.h",
         "//src/mongo/db/pipeline:accumulator_for_window_functions.h",

src/mongo/db/memory_tracking/BUILD.bazel

@@ -11,15 +11,28 @@ exports_files(
 
 mongo_cc_library(
     name = "memory_tracking",
+    srcs = [
+        "op_memory_use.cpp",
+    ],
     hdrs = [
         "memory_usage_tracker.h",
+        "op_memory_use.h",
+        "operation_memory_usage_tracker.h",
+    ],
+    deps = [
+        "//src/mongo/db:service_context",
     ],
 )
 
 mongo_cc_unit_test(
     name = "memory_tracking_test",
     srcs = [
         "memory_usage_tracker_test.cpp",
+        "op_memory_use_test.cpp",
     ],
     tags = ["mongo_unittest_first_group"],
+    deps = [
+        ":memory_tracking",
+        "//src/mongo/db:service_context_test_fixture",
+    ],
 )

src/mongo/db/memory_tracking/op_memory_use.cpp

@@ -0,0 +1,40 @@
+/**
+ *    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/op_memory_use.h"
+#include "mongo/base/initializer.h"
+#include "mongo/db/commands/server_status_metric.h"
+#include "mongo/db/memory_tracking/operation_memory_usage_tracker.h"
+#include "mongo/util/assert_util.h"
+
+namespace mongo {
+const OperationContext::Decoration<std::unique_ptr<OperationMemoryUsageTracker>>
+    OpMemoryUse::operationMemoryAggregator =
+        OperationContext::declareDecoration<std::unique_ptr<OperationMemoryUsageTracker>>();
+}  // namespace mongo

src/mongo/db/memory_tracking/op_memory_use.h

@@ -0,0 +1,49 @@
+/**
+ *    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.
+ */
+
+#pragma once
+
+#include "mongo/db/memory_tracking/operation_memory_usage_tracker.h"
+#include "mongo/db/operation_context.h"
+#include "mongo/db/service_context.h"
+
+namespace mongo {
+// OpMemoryUse will hold decorations which allow operations to roll up their memory use in the
+// operation context.
+class OpMemoryUse {
+public:
+    // This decoration will allow all code paths under one OperationContext to tally up their memory
+    // use to a single place. Because only the OperationContext should own this, we use a
+    // unique_ptr. This will also allow all the executors to have the same pointer to the
+    // operation memory tracker between the originating request and the call to getMore.
+    static const OperationContext::Decoration<std::unique_ptr<OperationMemoryUsageTracker>>
+        operationMemoryAggregator;
+};
+
+};  // namespace mongo

src/mongo/db/memory_tracking/op_memory_use_test.cpp

@@ -0,0 +1,62 @@
+/**
+ *    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/op_memory_use.h"
+#include "mongo/db/memory_tracking/operation_memory_usage_tracker.h"
+#include "mongo/db/service_context_test_fixture.h"
+
+namespace mongo {
+namespace {
+
+class MemoryTrackingTest : public ServiceContextTest {
+public:
+    using ServiceContextTest::ServiceContextTest;
+
+    auto makeOpCtx(std::string desc = "MemoryTrackingTest",
+                   std::shared_ptr<transport::Session> session = nullptr) {
+        auto client = getServiceContext()->getService()->makeClient(desc, session);
+        return client->makeOperationContext();
+    }
+};
+
+TEST_F(MemoryTrackingTest, OpMemoryUseCanAttachToOpCtx) {
+    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);
+    OperationMemoryUsageTracker* memoryTracker =
+        OpMemoryUse::operationMemoryAggregator(opCtx.get()).get();
+    ASSERT(memoryTracker);
+    ASSERT_EQ(memoryTracker->currentMemoryBytes(), 10);
+    ASSERT_EQ(memoryTracker->maxUsedMemoryBytes(), 15);
+}
+}  // namespace
+}  // namespace mongo

src/mongo/db/memory_tracking/operation_memory_usage_tracker.h

@@ -0,0 +1,54 @@
+/**
+ *    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.
+ */
+#pragma once
+#include <cstdint>
+
+namespace mongo {
+// TODO SERVER-98972 Replace this.
+class OperationMemoryUsageTracker {
+public:
+    OperationMemoryUsageTracker(int64_t currentMemoryBytes = 0, int64_t maxUsedMemoryBytes = 0)
+        : _currentMemoryBytes(currentMemoryBytes), _maxUsedMemoryBytes(maxUsedMemoryBytes) {}
+
+    int64_t currentMemoryBytes() const {
+        return _currentMemoryBytes;
+    }
+
+    int64_t maxUsedMemoryBytes() const {
+        return _maxUsedMemoryBytes;
+    }
+
+private:
+    // Current amount of memory in use by all blocking stages.
+    int64_t _currentMemoryBytes = 0;
+
+    // High-water mark: the highest amount of memory that has been allocated at one time so far.
+    int64_t _maxUsedMemoryBytes = 0;
+};
+}  // namespace mongo

src/mongo/util/BUILD.bazel

@@ -921,7 +921,6 @@ mongo_cc_unit_test(
         "base64_test.cpp",
         "cancellation_test.cpp",
         "clock_source_mock_test.cpp",
-        "concurrent_memory_aggregator_test.cpp",
         "container_size_helper_test.cpp",
         "ctype_test.cpp",
         "decimal_counter_test.cpp",