SERVER-61453 Fix crash in Simple8bBuilder when inserting large value after flush

2 files changed, 26 insertions, 2 deletions

diff --git a/src/mongo/bson/util/simple8b.cpp b/src/mongo/bson/util/simple8b.cpp
index db1c8be1fcd..3dc77036e34 100644
--- a/src/mongo/bson/util/simple8b.cpp
+++ b/src/mongo/bson/util/simple8b.cpp
@@ -549,7 +549,6 @@ bool Simple8bBuilder<T>::_appendValue(T value, bool tryRle) {
         PendingValue lastPendingValue = _pendingValues.back();
         do {
             uint64_t simple8bWord = _encodeLargestPossibleWord(_lastValidExtensionType);
-            isSelectorPossible = {true, true, true, true};
             _writeFn(simple8bWord);
         } while (!(_doesIntegerFitInCurrentWord(pendingValue)));
 
@@ -577,7 +576,6 @@ void Simple8bBuilder<T>::_appendSkip(bool tryRle) {
             // Form simple8b word if skip can not fit with last selector
             uint64_t simple8bWord = _encodeLargestPossibleWord(_lastValidExtensionType);
             _writeFn(simple8bWord);
-            isSelectorPossible = {true, true, true, true};
             _lastValidExtensionType = kBaseSelector;
         }
 
@@ -703,6 +701,8 @@ int64_t Simple8bBuilder<T>::_encodeLargestPossibleWord(uint8_t extensionType) {
     for (auto val : _pendingValues) {
         _updateSimple8bCurrentState(val);
     }
+    // Reset which selectors are possible to use for next word
+    isSelectorPossible.fill(true);
     return encodedWord;
 }
 
diff --git a/src/mongo/bson/util/simple8b_test.cpp b/src/mongo/bson/util/simple8b_test.cpp
index 385e0a7f6b8..361196a21de 100644
--- a/src/mongo/bson/util/simple8b_test.cpp
+++ b/src/mongo/bson/util/simple8b_test.cpp
@@ -1169,6 +1169,30 @@ TEST(Simple8b, RleFlushResetsRle) {
     assertValuesEqual(s8b, std::vector<boost::optional<uint64_t>>(121, 1));
 }
 
+TEST(Simple8b, RleFlushResetsPossibleSelectors) {
+    BufBuilder buffer;
+    Simple8bBuilder<uint64_t> builder([&buffer](uint64_t simple8bBlock) {
+        buffer.appendNum(simple8bBlock);
+        return true;
+    });
+
+    // Write a large value with many trailing zeros that does not fit in the base selector, we then
+    // flush and make sure that we can write a value that only fits in the base selector. We should
+    // have reset possible selectors as part of the flush.
+    std::vector<boost::optional<uint64_t>> expectedInts = {0x8000000000000000, 0x0FFFFFFFFFFFFFFE};
+
+    ASSERT_TRUE(builder.append(*expectedInts[0]));
+    builder.flush();
+    ASSERT_TRUE(builder.append(*expectedInts[1]));
+    builder.flush();
+
+    auto size = buffer.len();
+    auto sharedBuffer = buffer.release();
+
+    Simple8b<uint64_t> s8b(sharedBuffer.get(), size);
+    assertValuesEqual(s8b, expectedInts);
+}
+
 TEST(Simple8b, EightSelectorLargeMax) {
     // Selector 8 value
     // 1111 + 124 zeros