1 files changed, 26 insertions, 4 deletions
diff --git a/src/mongo/db/storage/key_string_test.cpp b/src/mongo/db/storage/key_string_test.cpp
index ef9f7e744a5..e1a29dcd8a1 100644
--- a/src/mongo/db/storage/key_string_test.cpp
+++ b/src/mongo/db/storage/key_string_test.cpp
@@ -283,6 +283,18 @@ TEST_F(KeyStringBuilderTest, MaxElementsInCompoundKey) {
     KeyString::getKeySize(ks.getBuffer(), ks.getSize(), ALL_ASCENDING, ks.getTypeBits());
 }
 
+TEST_F(KeyStringBuilderTest, EmbeddedNullString) {
+    // Construct a KeyString where \x3c defines the type kStringLike then embedded with null
+    // characters and followed by \x00.
+    const char* data = "\x3c\x00\xff\x00";
+    const size_t size = 4;
+    KeyString::TypeBits typeBits(KeyString::Version::kLatestVersion);
+
+    // No exceptions should be thrown.
+    ASSERT_BSONOBJ_EQ(KeyString::toBson(data, size, ALL_ASCENDING, typeBits),
+                      BSON("" << StringData("\x00", 1)));
+};
+
 TEST_F(KeyStringBuilderTest, ExceededBSONDepth) {
     KeyString::Builder ks(KeyString::Version::V1);
 
@@ -746,7 +758,7 @@ TEST_F(KeyStringBuilderTest, InvalidInfinityDecimalV0) {
 
 TEST_F(KeyStringBuilderTest, ReasonableSize) {
     // Tests that KeyString::Builders do not use an excessive amount of memory for small key
-    // generation. These upper bounds were the calculate sizes of each type at the time this
+    // generation. These upper bounds were the calculated sizes of each type at the time this
     // test was written.
     KeyString::Builder stackBuilder(KeyString::Version::kLatestVersion, BSONObj(), ALL_ASCENDING);
     static_assert(sizeof(stackBuilder) <= 624);
@@ -755,8 +767,13 @@ TEST_F(KeyStringBuilderTest, ReasonableSize) {
         KeyString::Version::kLatestVersion, BSONObj(), ALL_ASCENDING);
     static_assert(sizeof(heapBuilder) <= 104);
 
-    // Use large 1KB blocks and verify that we use way less
-    SharedBufferFragmentBuilder fragmentBuilder(1024);
+    // Use a small block size to ensure we do not use more. Additionally, the minimum allocation
+    // size is 64.
+    const auto minSize = 64;
+    SharedBufferFragmentBuilder fragmentBuilder(
+        minSize,
+        SharedBufferFragmentBuilder::DoubleGrowStrategy(
+            SharedBufferFragmentBuilder::kDefaultMaxBlockSize));
     KeyString::PooledBuilder pooledBuilder(
         fragmentBuilder, KeyString::Version::kLatestVersion, BSONObj(), ALL_ASCENDING);
     static_assert(sizeof(pooledBuilder) <= 104);
@@ -776,11 +793,16 @@ TEST_F(KeyStringBuilderTest, ReasonableSize) {
 
     KeyString::Value value4 = pooledBuilder.getValueCopy();
     ASSERT_LTE(sizeof(value4), 32);
+    // This is safe because we are operating on a copy of the value and it is not shared elsewhere.
     ASSERT_LTE(value4.memUsageForSorter(), 34);
+    // We should still be using the initially-allocated size.
+    ASSERT_LTE(fragmentBuilder.memUsage(), 64);
 
+    // For values created with the pooledBuilder, it is invalid to call memUsageForSorter(). Instead
+    // we look at the mem usage of the builder itself.
     KeyString::Value value5 = pooledBuilder.release();
     ASSERT_LTE(sizeof(value5), 32);
-    ASSERT_LTE(value5.memUsageForSorter(), 34);
+    ASSERT_LTE(fragmentBuilder.memUsage(), 64);
 }
 
 TEST_F(KeyStringBuilderTest, DiscardIfNotReleased) {