SERVER-47692 SharedBufferFragmentAllocator should discard its buffer on destruction if the build was incomplete

3 files changed, 25 insertions, 0 deletions

diff --git a/src/mongo/bson/util/builder.h b/src/mongo/bson/util/builder.h
index 4bd4be51233..26d8a827aa9 100644
--- a/src/mongo/bson/util/builder.h
+++ b/src/mongo/bson/util/builder.h
@@ -130,6 +130,12 @@ class SharedBufferFragmentAllocator {
 public:
     SharedBufferFragmentAllocator(SharedBufferFragmentBuilder& fragmentBuilder)
         : _fragmentBuilder(fragmentBuilder) {}
+    ~SharedBufferFragmentAllocator() {
+        // Discard if the build was not finished at the time of destruction.
+        if (_fragmentBuilder.building()) {
+            free();
+        }
+    }
 
     // Allow moving but not copying. It would be an error for two SharedBufferFragmentAllocator to
     // use the same underlying builder at the same time.
diff --git a/src/mongo/db/storage/key_string_test.cpp b/src/mongo/db/storage/key_string_test.cpp
index dfba16a4d0f..7e0f802c164 100644
--- a/src/mongo/db/storage/key_string_test.cpp
+++ b/src/mongo/db/storage/key_string_test.cpp
@@ -768,6 +768,20 @@ TEST_F(KeyStringBuilderTest, ReasonableSize) {
     ASSERT_LTE(value5.memUsageForSorter(), 34);
 }
 
+TEST_F(KeyStringBuilderTest, DiscardIfNotReleased) {
+    SharedBufferFragmentBuilder fragmentBuilder(1024);
+    {
+        // Intentially not released, but the data should be discarded correctly.
+        KeyString::PooledBuilder pooledBuilder(
+            fragmentBuilder, KeyString::Version::kLatestVersion, BSONObj(), ALL_ASCENDING);
+    }
+    {
+        KeyString::PooledBuilder pooledBuilder(
+            fragmentBuilder, KeyString::Version::kLatestVersion, BSONObj(), ALL_ASCENDING);
+        pooledBuilder.release();
+    }
+}
+
 TEST_F(KeyStringBuilderTest, LotsOfNumbers1) {
     for (int i = 0; i < 64; i++) {
         int64_t x = 1LL << i;
diff --git a/src/mongo/util/shared_buffer_fragment.h b/src/mongo/util/shared_buffer_fragment.h
index 4824e38ea8b..b566d355287 100644
--- a/src/mongo/util/shared_buffer_fragment.h
+++ b/src/mongo/util/shared_buffer_fragment.h
@@ -185,6 +185,11 @@ public:
         return _buffer.get() + _offset;
     }
 
+    // Returns whether or not a memory fragment is currently being built.
+    bool building() const {
+        return _inUse;
+    }
+
 private:
     SharedBuffer _buffer;
     ptrdiff_t _offset;