From d269dd586a99b2e6c05710543c66f958afa7ffe4 Mon Sep 17 00:00:00 2001
From: Dan Larkin-York <dan.larkin-york@mongodb.com>
Date: Wed, 22 Dec 2021 16:23:18 +0000
Subject: SERVER-61996 Reserve memory for MutableDocument in
 BucketUnpacker::getNext()

---
 src/mongo/db/exec/bucket_unpacker.cpp | 26 +++++++++++++++++++++++++-
 1 file changed, 25 insertions(+), 1 deletion(-)

diff --git a/src/mongo/db/exec/bucket_unpacker.cpp b/src/mongo/db/exec/bucket_unpacker.cpp
index 642116cd3e6..4c447bb6d9d 100644
--- a/src/mongo/db/exec/bucket_unpacker.cpp
+++ b/src/mongo/db/exec/bucket_unpacker.cpp
@@ -54,6 +54,13 @@ public:
                                           const BSONElement& metaValue,
                                           bool includeTimeField,
                                           bool includeMetaField) = 0;
+
+    // Provides an upper bound on the number of fields in each measurement.
+    virtual std::size_t numberOfFields() = 0;
+
+protected:
+    // Data field count is variable, but time and metadata are fixed.
+    constexpr static std::size_t kFixedFieldNumber = 2;
 };
 
 namespace {
@@ -98,6 +105,7 @@ public:
                                   const BSONElement& metaValue,
                                   bool includeTimeField,
                                   bool includeMetaField) override;
+    std::size_t numberOfFields() override;
 
 private:
     // Iterates the timestamp section of the bucket to drive the unpacking iteration.
@@ -202,6 +210,12 @@ void BucketUnpackerV1::extractSingleMeasurement(MutableDocument& measurement,
     }
 }
 
+std::size_t BucketUnpackerV1::numberOfFields() {
+    // The data fields are tracked by _fieldIters, but we need to account also for the time field
+    // and possibly the meta field.
+    return kFixedFieldNumber + _fieldIters.size();
+}
+
 // Unpacker for V2 compressed buckets
 class BucketUnpackerV2 : public BucketUnpacker::UnpackingImpl {
 public:
@@ -222,6 +236,7 @@ public:
                                   const BSONElement& metaValue,
                                   bool includeTimeField,
                                   bool includeMetaField) override;
+    std::size_t numberOfFields() override;
 
 private:
     struct ColumnStore {
@@ -320,6 +335,12 @@ void BucketUnpackerV2::extractSingleMeasurement(MutableDocument& measurement,
     }
 }
 
+std::size_t BucketUnpackerV2::numberOfFields() {
+    // The data fields are tracked by _fieldColumns, but we need to account also for the time field
+    // and possibly the meta field.
+    return kFixedFieldNumber + _fieldColumns.size();
+}
+
 /**
  * Erase computed meta projection fields if they are present in the exclusion field set.
  */
@@ -365,7 +386,10 @@ Document BucketUnpacker::getNext() {
     tassert(5521503, "'getNext()' requires the bucket to be owned", _bucket.isOwned());
     tassert(5422100, "'getNext()' was called after the bucket has been exhausted", hasNext());
 
-    auto measurement = MutableDocument{};
+    // MutableDocument reserves memory based on the number of fields, but uses a fixed size of 25
+    // bytes plus an allowance of 7 characters for the field name. Doubling the number of fields
+    // should give us enough overhead for longer field names without wasting too much memory.
+    auto measurement = MutableDocument{2 * _unpackingImpl->numberOfFields()};
     _hasNext = _unpackingImpl->getNext(
         measurement, _spec, _metaValue, _includeTimeField, _includeMetaField);
 
-- 
cgit v1.2.1