SERVER-54296 Fix incorrect cache size estimate logic in lookup_set_cache.h

(cherry picked from commit abf0352882f93e760537d612401eb9fb6e5a030e)

2 files changed, 62 insertions, 25 deletions

diff --git a/src/mongo/db/pipeline/lookup_set_cache.h b/src/mongo/db/pipeline/lookup_set_cache.h
index cd033caa968..8f9eb821dda 100644
--- a/src/mongo/db/pipeline/lookup_set_cache.h
+++ b/src/mongo/db/pipeline/lookup_set_cache.h
@@ -60,16 +60,21 @@ using boost::multi_index::sequenced;
  */
 class LookupSetCache {
 public:
-    using Cached = std::pair<Value, std::vector<Document>>;
+    struct Cached {
+        Value key;
+        std::vector<Document> docs;
+
+        // This includes the size of both key and the documents.
+        size_t approxCacheEntrySize = 0;
+    };
 
     // boost::multi_index_container provides a system for implementing a cache. Here, we create
-    // a container of std::pair<Value, std::vector<Document>>, that is both sequenced, and has a
-    // unique index on the Value. From this, we are able to evict the least-recently-used member,
-    // and maintain key uniqueness.
+    // a container of 'Cached', that is both sequenced, and has a unique index on the 'Cached::key'.
+    // From this, we are able to evict the least-recently-used member, and maintain key uniqueness.
     using IndexedContainer =
         multi_index_container<Cached,
                               indexed_by<sequenced<>,
-                                         hashed_unique<member<Cached, Value, &Cached::first>,
+                                         hashed_unique<member<Cached, Value, &Cached::key>,
                                                        ValueComparator::Hasher,
                                                        ValueComparator::EqualTo>>>;
 
@@ -81,7 +86,7 @@ public:
     explicit LookupSetCache(const ValueComparator& comparator)
         : _container(boost::make_tuple(IndexedContainer::nth_index<0>::type::ctor_args(),
                                        boost::make_tuple(0,
-                                                         member<Cached, Value, &Cached::first>(),
+                                                         member<Cached, Value, &Cached::key>(),
                                                          comparator.getHasher(),
                                                          comparator.getEqualTo()))) {}
 
@@ -102,12 +107,12 @@ public:
         auto it = _container.begin();
         std::advance(it, middle);
         const auto keySize = key.getApproximateSize();
-        const auto docSize = doc.getApproximateSize();
+        auto cacheEntrySizeIncreaseBy = doc.getApproximateSize();
 
         // Find the cache entry, or create one if it doesn't exist yet.
-        auto insertionResult = _container.insert(it, {std::move(key), {}});
+        auto insertionResult = _container.insert(it, {std::move(key), {}, 0});
         if (insertionResult.second) {
-            _memoryUsage += keySize;
+            cacheEntrySizeIncreaseBy += keySize;
         } else {
             // We did not insert due to a duplicate key. Update the cached doc, moving it to the
             // middle of the cache.
@@ -115,11 +120,11 @@ public:
         }
 
         // Add the doc to the cache entry.
-        _container.modify(insertionResult.first,
-                          [&doc](std::pair<Value, std::vector<Document>>& entry) {
-                              entry.second.push_back(std::move(doc));
-                          });
-        _memoryUsage += docSize;
+        _container.modify(insertionResult.first, [&doc, cacheEntrySizeIncreaseBy](Cached& entry) {
+            entry.docs.push_back(std::move(doc));
+            entry.approxCacheEntrySize += cacheEntrySizeIncreaseBy;
+        });
+        _memoryUsage += cacheEntrySizeIncreaseBy;
     }
 
     /**
@@ -130,17 +135,10 @@ public:
             return;
         }
 
-        const Cached& pair = _container.back();
-
-        size_t keySize = pair.first.getApproximateSize();
-        invariant(keySize <= _memoryUsage);
-        _memoryUsage -= keySize;
+        const auto cacheEntrySize = _container.back().approxCacheEntrySize;
+        invariant(cacheEntrySize <= _memoryUsage);
+        _memoryUsage -= cacheEntrySize;
 
-        for (auto&& elem : pair.second) {
-            size_t valueSize = static_cast<size_t>(elem.getApproximateSize());
-            invariant(valueSize <= _memoryUsage);
-            _memoryUsage -= valueSize;
-        }
         _container.erase(std::prev(_container.end()));
     }
 
@@ -186,11 +184,15 @@ public:
             boost::multi_index::get<0>(_container)
                 .relocate(boost::multi_index::get<0>(_container).begin(),
                           boost::multi_index::project<0>(_container, it));
-            return &it->second;
+            return &it->docs;
         }
         return nullptr;
     }
 
+    auto getMemoryUsage() const {
+        return _memoryUsage;
+    }
+
 private:
     IndexedContainer _container;
 
diff --git a/src/mongo/db/pipeline/lookup_set_cache_test.cpp b/src/mongo/db/pipeline/lookup_set_cache_test.cpp
index c0a2479e01a..6670b5b5971 100644
--- a/src/mongo/db/pipeline/lookup_set_cache_test.cpp
+++ b/src/mongo/db/pipeline/lookup_set_cache_test.cpp
@@ -207,4 +207,39 @@ TEST(LookupSetCacheTest, CachedValuesDontRespectCollation) {
     ASSERT_EQ(2U, fooResult->size());
 }
 
+TEST(LookupSetCacheTest, DocumentWithStorageCachePopulated) {
+    LookupSetCache cache(defaultComparator);
+
+    // Use a BSONObj backed Document so that the Document's hot storage isn't populated at
+    // initialization.
+    BSONObj input = BSON("a" << 1);
+    const auto doc1 = Document(input);
+    const auto sizeOfDoc1Before = doc1.getApproximateSize();
+    auto key = Value("foo"_sd);
+
+    // Insert a cache entry and verify that both the key and the document are accounted for in the
+    // cache size.
+    cache.insert(key, doc1);
+    ASSERT_EQ(cache.getMemoryUsage(), sizeOfDoc1Before + key.getApproximateSize());
+
+
+    // A second document with the same key should require additional memory usage to store the
+    // document, but not the key.
+    auto prevCacheSize = cache.getMemoryUsage();
+    const auto doc2 = Document({{"a", 2}});
+    cache.insert(key, doc2);
+    ASSERT_EQ(cache.getMemoryUsage(), prevCacheSize + doc2.getApproximateSize());
+
+    // 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_EQ(prevCacheSize, cache.getMemoryUsage());
+
+    cache.evictOne();
+    ASSERT_EQ(cache.getMemoryUsage(), 0);
+}
+
 }  // namespace mongo