SERVER-61281 Use memoization for Document::getApproximateSize

6 files changed, 32 insertions, 51 deletions

diff --git a/src/mongo/db/pipeline/accumulator_multi.cpp b/src/mongo/db/pipeline/accumulator_multi.cpp
index 4f63aaa4149..3dd7f02ce51 100644
--- a/src/mongo/db/pipeline/accumulator_multi.cpp
+++ b/src/mongo/db/pipeline/accumulator_multi.cpp
@@ -599,12 +599,6 @@ void AccumulatorTopBottomN<sense, single>::_processValue(const Value& val) {
         }
     }
 
-    // TODO SERVER-61281 consider removing this call to fillCache().
-    // Since Document caches fields the size of this cache and getApproximateSize() can vary
-    // depending on access. In order to avoid this and make sure we subtract the right amount if
-    // remove() ever gets called, we can fill the cache to get a consistent view of the size.
-    // Normally the outer window function code handles this, but _genKeyOutPair() makes a new
-    // document for sortKey, so its cache get reset.
     keyOutPair.first.fillCache();
     const auto memUsage = keyOutPair.first.getApproximateSize() +
         keyOutPair.second.getApproximateSize() + sizeof(KeyOutPair);
@@ -626,9 +620,6 @@ void AccumulatorTopBottomN<sense, single>::remove(const Value& val) {
     auto it = _map->lower_bound(keyOutPair.first);
     _map->erase(it);
 
-    // TODO SERVER-61281 consider removing this comment if its no longer relevant.
-    // After calling lower_bound() it uses SortKeyComparator and the sortKey's field cache should be
-    // fully populated so no need to call fillCache() again.
     _memUsageBytes -= keyOutPair.first.getApproximateSize() +
         keyOutPair.second.getApproximateSize() + sizeof(KeyOutPair);
 }
diff --git a/src/mongo/db/pipeline/lookup_set_cache_test.cpp b/src/mongo/db/pipeline/lookup_set_cache_test.cpp
index 6670b5b5971..ec78b0b0572 100644
--- a/src/mongo/db/pipeline/lookup_set_cache_test.cpp
+++ b/src/mongo/db/pipeline/lookup_set_cache_test.cpp
@@ -214,7 +214,7 @@ TEST(LookupSetCacheTest, DocumentWithStorageCachePopulated) {
     // initialization.
     BSONObj input = BSON("a" << 1);
     const auto doc1 = Document(input);
-    const auto sizeOfDoc1Before = doc1.getApproximateSize();
+    const auto sizeOfDoc1Before = doc1.getCurrentApproximateSize();
     auto key = Value("foo"_sd);
 
     // Insert a cache entry and verify that both the key and the document are accounted for in the
@@ -228,14 +228,14 @@ TEST(LookupSetCacheTest, DocumentWithStorageCachePopulated) {
     auto prevCacheSize = cache.getMemoryUsage();
     const auto doc2 = Document({{"a", 2}});
     cache.insert(key, doc2);
-    ASSERT_EQ(cache.getMemoryUsage(), prevCacheSize + doc2.getApproximateSize());
+    ASSERT_EQ(cache.getMemoryUsage(), prevCacheSize + doc2.getCurrentApproximateSize());
 
     // Calling serializeForSorter() should grow the overall document size. Verify that growing the
     // size of the 'Document' object does not have impact on the size stored in 'cache'.
     prevCacheSize = cache.getMemoryUsage();
     BufBuilder builder;
     doc1.serializeForSorter(builder);
-    ASSERT_LT(sizeOfDoc1Before, doc1.getApproximateSize());
+    ASSERT_LT(sizeOfDoc1Before, doc1.getCurrentApproximateSize());
     ASSERT_EQ(prevCacheSize, cache.getMemoryUsage());
 
     cache.evictOne();
diff --git a/src/mongo/db/pipeline/memory_usage_tracker.h b/src/mongo/db/pipeline/memory_usage_tracker.h
index 9847b41e3cb..6887e91289a 100644
--- a/src/mongo/db/pipeline/memory_usage_tracker.h
+++ b/src/mongo/db/pipeline/memory_usage_tracker.h
@@ -50,13 +50,11 @@ public:
         PerFunctionMemoryTracker() = delete;
 
         void update(long long diff) {
-
-            // TODO SERVER-61281: this is a temporary measure in tackling the problem in this
-            // ticket. It prevents the underflow from happening but doesn't address the cause
-            // which is inaccurate tracking. Once inaccurate tracking is resolved, the underflow
-            // assertion could be restored. See SERVER-62856 and SERVER-65473 for details.
-
-            set(std::max(_currentMemoryBytes + diff, 0LL));
+            tassert(5578603,
+                    str::stream() << "Underflow on memory tracking, attempting to add " << diff
+                                  << " but only " << _currentMemoryBytes << " available",
+                    diff >= 0 || _currentMemoryBytes >= std::abs(diff));
+            set(_currentMemoryBytes + diff);
         }
 
         void set(long long total) {
@@ -147,13 +145,11 @@ public:
      * Updates total memory usage.
      */
     void update(long long diff) {
-
-        // TODO SERVER-61281: this is a temporary measure in tackling the problem in this
-        // ticket. It prevents the underflow from happening but doesn't address the cause
-        // which is inaccurate tracking. Once inaccurate tracking is resolved, the underflow
-        // assertion could be restored. See SERVER-62856 and SERVER-65473 for details.
-
-        set(std::max(_memoryUsageBytes + diff, 0LL));
+        tassert(5578602,
+                str::stream() << "Underflow on memory tracking, attempting to add " << diff
+                              << " but only " << _memoryUsageBytes << " available",
+                diff >= 0 || (int)_memoryUsageBytes >= -1 * diff);
+        set(_memoryUsageBytes + diff);
     }
 
     auto currentMemoryBytes() const {
diff --git a/src/mongo/db/pipeline/memory_usage_tracker_test.cpp b/src/mongo/db/pipeline/memory_usage_tracker_test.cpp
index 81ca06a1381..9d383a8a468 100644
--- a/src/mongo/db/pipeline/memory_usage_tracker_test.cpp
+++ b/src/mongo/db/pipeline/memory_usage_tracker_test.cpp
@@ -99,36 +99,26 @@ TEST_F(MemoryUsageTrackerTest, UpdateUsageUpdatesGlobal) {
     ASSERT_EQ(_tracker.maxMemoryBytes(), 150LL);
 }
 
-// TODO SERVER-61281: Need to add back the DEATH_TEST_F removed in SERVER-65473, once
-// MemoryUsageTracker.update() no longer underflows.
-
-TEST_F(MemoryUsageTrackerTest, UpdateFunctionUsageToNegativeIsDisallowed) {
+DEATH_TEST_F(MemoryUsageTrackerTest,
+             UpdateFunctionUsageToNegativeIsDisallowed,
+             "Underflow on memory tracking") {
     _funcTracker.set(50LL);
     ASSERT_EQ(_funcTracker.currentMemoryBytes(), 50LL);
     ASSERT_EQ(_funcTracker.maxMemoryBytes(), 50LL);
     ASSERT_EQ(_tracker.currentMemoryBytes(), 50LL);
     ASSERT_EQ(_tracker.maxMemoryBytes(), 50LL);
 
-    // TODO SERVER-61281: This should throw an assertion once accurate tracking is implemented and
-    // no underflow should happen. That assertion would make sure "Underflow in memory tracking"
-    // is reported.
-
     _funcTracker.update(-100);
-    ASSERT_EQ(_funcTracker.currentMemoryBytes(), 0LL);
-    ASSERT_EQ(_tracker.currentMemoryBytes(), 0LL);
 }
 
-TEST_F(MemoryUsageTrackerTest, UpdateMemUsageToNegativeIsDisallowed) {
+DEATH_TEST_F(MemoryUsageTrackerTest,
+             UpdateMemUsageToNegativeIsDisallowed,
+             "Underflow on memory tracking") {
     _tracker.set(50LL);
     ASSERT_EQ(_tracker.currentMemoryBytes(), 50LL);
     ASSERT_EQ(_tracker.maxMemoryBytes(), 50LL);
 
-    // TODO SERVER-61281: This should throw an assertion once accurate tracking is implemented and
-    // no underflow should happen. That assertion would make sure "Underflow in memory tracking"
-    // is reported.
-
     _tracker.update(-100);
-    ASSERT_EQ(_tracker.currentMemoryBytes(), 0LL);
 }
 
 }  // namespace
diff --git a/src/mongo/db/pipeline/window_function/partition_iterator_test.cpp b/src/mongo/db/pipeline/window_function/partition_iterator_test.cpp
index cb6b5bcf1bf..4ba0a692e47 100644
--- a/src/mongo/db/pipeline/window_function/partition_iterator_test.cpp
+++ b/src/mongo/db/pipeline/window_function/partition_iterator_test.cpp
@@ -499,7 +499,7 @@ TEST_F(PartitionIteratorTest, MemoryUsageAccountsForDocumentIteratorCache) {
     const auto mock = DocumentSourceMock::createForTest(docs, getExpCtx());
 
     [[maybe_unused]] auto accessor = makeDefaultAccessor(mock, boost::none);
-    size_t initialDocSize = docs[0].getDocument().getApproximateSize();
+    size_t initialDocSize = docs[0].getDocument().getCurrentApproximateSize();
 
     // Pull in the first document, and verify the reported size of the iterator is roughly double
     // the size of the document. The size of the iterator is double the size of the document because
@@ -525,7 +525,7 @@ TEST_F(PartitionIteratorTest, MemoryUsageAccountsForArraysInDocumentIteratorCach
     const auto mock = DocumentSourceMock::createForTest(docs, getExpCtx());
 
     [[maybe_unused]] auto accessor = makeDefaultAccessor(mock, boost::none);
-    size_t initialDocSize = docs[0].getDocument().getApproximateSize();
+    size_t initialDocSize = docs[0].getDocument().getCurrentApproximateSize();
 
     // Pull in the first document, and verify the reported size of the iterator is roughly
     // triple the size of the document. The reason for this is that 'largeStr' is cached twice; once
@@ -550,7 +550,7 @@ TEST_F(PartitionIteratorTest, MemoryUsageAccountsForNestedArraysInDocumentIterat
     const auto mock = DocumentSourceMock::createForTest(docs, getExpCtx());
 
     [[maybe_unused]] auto accessor = makeDefaultAccessor(mock, boost::none);
-    size_t initialDocSize = docs[0].getDocument().getApproximateSize();
+    size_t initialDocSize = docs[0].getDocument().getCurrentApproximateSize();
 
     // Pull in the first document, and verify the reported size of the iterator is roughly
     // triple the size of the document. The reason for this is that 'largeStr' is cached twice; once
@@ -575,7 +575,7 @@ TEST_F(PartitionIteratorTest, MemoryUsageAccountsForNestedObjInDocumentIteratorC
     const auto mock = DocumentSourceMock::createForTest(docs, getExpCtx());
 
     [[maybe_unused]] auto accessor = makeDefaultAccessor(mock, boost::none);
-    size_t initialDocSize = docs[0].getDocument().getApproximateSize();
+    size_t initialDocSize = docs[0].getDocument().getCurrentApproximateSize();
 
     // Pull in the first document, and verify the reported size. TODO SERVER-57011: The approximate
     // size should not double count the nested strings.
@@ -592,7 +592,7 @@ TEST_F(PartitionIteratorTest, MemoryUsageAccountsForReleasedDocuments) {
     const auto mock = DocumentSourceMock::createForTest(docs, getExpCtx());
 
     auto accessor = makeDefaultAccessor(mock, boost::none);
-    size_t initialDocSize = docs[0].getDocument().getApproximateSize();
+    size_t initialDocSize = docs[0].getDocument().getCurrentApproximateSize();
 
     // Pull in the first document, and verify the reported size of the iterator is roughly double
     // the size of the document.
diff --git a/src/mongo/db/pipeline/window_function/window_function_top_bottom_n.h b/src/mongo/db/pipeline/window_function/window_function_top_bottom_n.h
index fbf77c59fa2..64d3319efb9 100644
--- a/src/mongo/db/pipeline/window_function/window_function_top_bottom_n.h
+++ b/src/mongo/db/pipeline/window_function/window_function_top_bottom_n.h
@@ -56,17 +56,17 @@ public:
     explicit WindowFunctionTopBottomN(ExpressionContext* const expCtx, SortPattern sp, long long n)
         : WindowFunctionState(expCtx), _acc(expCtx, std::move(sp), true) {
         _acc.startNewGroup(Value(n));
-        _memUsageBytes = sizeof(*this);
+        updateMemUsage();
     }
 
     void add(Value value) final {
         _acc.process(value, false);
-        _memUsageBytes = _acc.getMemUsage();
+        updateMemUsage();
     }
 
     void remove(Value value) final {
         _acc.remove(value);
-        _memUsageBytes = _acc.getMemUsage();
+        updateMemUsage();
     }
 
     Value getValue() const final {
@@ -75,10 +75,14 @@ public:
 
     void reset() final {
         _acc.reset();
-        _memUsageBytes = _acc.getMemUsage();
+        updateMemUsage();
     }
 
 private:
+    void updateMemUsage() {
+        _memUsageBytes = sizeof(*this) + _acc.getMemUsage();
+    }
+
     AccumulatorTopBottomN<sense, single> _acc;
 };