3 files changed, 278 insertions, 48 deletions

diff --git a/src/mongo/db/storage/key_string.cpp b/src/mongo/db/storage/key_string.cpp
index 9155149b189..40786dd18c7 100644
--- a/src/mongo/db/storage/key_string.cpp
+++ b/src/mongo/db/storage/key_string.cpp
@@ -326,45 +326,27 @@ string readInvertedCStringWithNuls(BufReader* reader) {
 }  // namespace
 
 void Builder::resetToKey(const BSONObj& obj, Ordering ord, RecordId recordId) {
-    resetToEmpty();
-    _appendAllElementsForIndexing(obj, ord, kInclusive);
+    resetToEmpty(ord);
+    _appendAllElementsForIndexing(obj, kInclusive);
     appendRecordId(recordId);
 }
 
 void Builder::resetToKey(const BSONObj& obj, Ordering ord, Discriminator discriminator) {
-    resetToEmpty();
-    _appendAllElementsForIndexing(obj, ord, discriminator);
+    resetToEmpty(ord, discriminator);
+    _appendAllElementsForIndexing(obj, discriminator);
 }
 
-// ----------------------------------------------------------------------
-// -----------   APPEND CODE  -------------------------------------------
-// ----------------------------------------------------------------------
-
-void Builder::_appendAllElementsForIndexing(const BSONObj& obj,
-                                            Ordering ord,
-                                            Discriminator discriminator) {
-    int elemCount = 0;
-    BSONObjIterator it(obj);
-    while (auto elem = it.next()) {
-        const int elemIdx = elemCount++;
-        const bool invert = (ord.get(elemIdx) == -1);
+void Builder::appendBSONElement(const BSONElement& elem) {
+    const int elemIdx = _elemCount++;
+    const bool invert = (_ordering.get(elemIdx) == -1);
 
-        _appendBsonValue(elem, invert, nullptr);
-
-        dassert(elem.fieldNameSize() < 3);  // fieldNameSize includes the NUL
-
-        // IndexEntryComparison::makeQueryObject() encodes a discriminator in the first byte of
-        // the field name. This discriminator overrides the passed in one. Normal elements only
-        // have the NUL byte terminator. Entries stored in an index are not allowed to have a
-        // discriminator.
-        if (char ch = *elem.fieldName()) {
-            // l for less / g for greater.
-            invariant(ch == 'l' || ch == 'g');
-            discriminator = ch == 'l' ? kExclusiveBefore : kExclusiveAfter;
-            invariant(!it.more());
-        }
+    if (_state == kEmpty) {
+        _transition(kAppendingBSONElements);
     }
+    _appendBsonValue(elem, invert, nullptr);
+}
 
+void Builder::_appendDiscriminator(const Discriminator discriminator) {
     // The discriminator forces this KeyString to compare Less/Greater than any KeyString with
     // the same prefix of keys. As an example, this can be used to land on the first key in the
     // index with the value "a" regardless of the RecordId. In compound indexes it can use a
@@ -382,10 +364,40 @@ void Builder::_appendAllElementsForIndexing(const BSONObj& obj,
 
     // TODO consider omitting kEnd when using a discriminator byte. It is not a storage format
     // change since keystrings with discriminators are not allowed to be stored.
+    _appendEnd();
+}
+// ----------------------------------------------------------------------
+// -----------   APPEND CODE  -------------------------------------------
+// ----------------------------------------------------------------------
+
+void Builder::_appendEnd() {
+    _transition(kEndAdded);
     _append(kEnd, false);
 }
 
+void Builder::_appendAllElementsForIndexing(const BSONObj& obj, Discriminator discriminator) {
+    _transition(kAppendingBSONElements);
+    BSONObjIterator it(obj);
+    while (auto elem = it.next()) {
+        appendBSONElement(elem);
+        dassert(elem.fieldNameSize() < 3);  // fieldNameSize includes the NUL
+
+        // IndexEntryComparison::makeQueryObject() encodes a discriminator in the first byte of
+        // the field name. This discriminator overrides the passed in one. Normal elements only
+        // have the NUL byte terminator. Entries stored in an index are not allowed to have a
+        // discriminator.
+        if (char ch = *elem.fieldName()) {
+            // l for less / g for greater.
+            invariant(ch == 'l' || ch == 'g');
+            discriminator = ch == 'l' ? kExclusiveBefore : kExclusiveAfter;
+            invariant(!it.more());
+        }
+    }
+    _appendDiscriminator(discriminator);
+}
+
 void Builder::appendRecordId(RecordId loc) {
+    _transition(kAppendedRecordID);
     // The RecordId encoding must be able to determine the full length starting from the last
     // byte, without knowing where the first byte is since it is stored at the end of a
     // KeyString, and we need to be able to read the RecordId without decoding the whole thing.
@@ -432,6 +444,7 @@ void Builder::appendRecordId(RecordId loc) {
 }
 
 void Builder::appendTypeBits(const TypeBits& typeBits) {
+    _transition(kAppendedTypeBits);
     // As an optimization, encode AllZeros as a single 0 byte.
     if (typeBits.isAllZeros()) {
         _append(uint8_t(0), false);
diff --git a/src/mongo/db/storage/key_string.h b/src/mongo/db/storage/key_string.h
index e2d62f2c094..c6baef2d894 100644
--- a/src/mongo/db/storage/key_string.h
+++ b/src/mongo/db/storage/key_string.h
@@ -51,6 +51,8 @@ static StringData keyStringVersionToString(Version version) {
     return version == Version::V0 ? "V0" : "V1";
 }
 
+static const Ordering ALL_ASCENDING = Ordering::make(BSONObj());
+
 /**
  * Encodes info needed to restore the original BSONTypes from a KeyString. They cannot be
  * stored in place since we don't want them to affect the ordering (1 and 1.0 compare as
@@ -73,9 +75,10 @@ public:
         _buf.reset();
         _buf.appendBuf(tb._buf.buf(), tb._buf.len());
     }
+
     TypeBits& operator=(const TypeBits& tb) = delete;
     TypeBits(TypeBits&&) = default;
-    TypeBits& operator=(TypeBits&&) = delete;
+    TypeBits& operator=(TypeBits&&) = default;
 
     /**
      * If there are no bytes remaining, assumes AllZeros. Otherwise, reads bytes out of the
@@ -248,7 +251,7 @@ public:
         const TypeBits& _typeBits;
     };
 
-    const Version version;
+    Version version;
 
 private:
     static uint32_t readSizeFromBuffer(BufReader* reader);
@@ -277,7 +280,7 @@ private:
      * the TypeBits size is in long encoding range(>127), all the bytes are used for the long
      * encoding format (first byte + 4 size bytes + data bytes).
      */
-    StackBufBuilder _buf;
+    BufBuilder _buf;
 };
 
 
@@ -343,6 +346,15 @@ public:
         kExclusiveAfter,
     };
 
+    enum BuildState {
+        kEmpty,                  // Buffer is empty.
+        kAppendingBSONElements,  // In the process of appending BSON Elements
+        kEndAdded,               // Finished appedning BSON Elements.
+        kAppendedRecordID,       // Finished appending a RecordID.
+        kAppendedTypeBits,       // Finished appending a TypeBits.
+        kReleased                // Released the buffer and so the buffer is no longer valid.
+    };
+
     /**
      * Encodes the kind of NumberDecimal that is stored.
      */
@@ -353,10 +365,18 @@ public:
         kDCMHasContinuationLargerThanDoubleRoundedUpTo15Digits = 0x3
     };
 
-    explicit Builder(Version version) : version(version), _typeBits(version) {}
+    Builder(Version version, Ordering ord, Discriminator discriminator)
+        : version(version),
+          _typeBits(version),
+          _state(kEmpty),
+          _elemCount(0),
+          _ordering(ord),
+          _discriminator(discriminator) {}
+    Builder(Version version, Ordering ord) : Builder(version, ord, kInclusive) {}
+    explicit Builder(Version version) : Builder(version, ALL_ASCENDING, kInclusive) {}
 
     Builder(Version version, const BSONObj& obj, Ordering ord, RecordId recordId)
-        : Builder(version) {
+        : Builder(version, ord) {
         resetToKey(obj, ord, recordId);
     }
 
@@ -364,19 +384,34 @@ public:
             const BSONObj& obj,
             Ordering ord,
             Discriminator discriminator = kInclusive)
-        : Builder(version) {
+        : Builder(version, ord) {
         resetToKey(obj, ord, discriminator);
     }
 
-    Builder(Version version, RecordId rid) : version(version), _typeBits(version) {
+    Builder(Version version, RecordId rid) : Builder(version) {
         appendRecordId(rid);
     }
 
     /**
+     * Releases the data held in this buffer into a Value type, releasing and transfering ownership
+     * of the buffer _buffer and TypeBits _typeBits to the returned Value object from the current
+     * Builder.
+     */
+    Value release() {
+        _prepareForRelease();
+        _transition(kReleased);
+        return {
+            version, std::move(_typeBits), static_cast<size_t>(_buffer.len()), _buffer.release()};
+    }
+
+    /**
      * Copies the data held in this buffer into a Value type that holds and owns a copy of the
      * buffer.
      */
-    Value getValue() {
+    Value getValueCopy() {
+        _prepareForRelease();
+        invariant(_state == kEndAdded || _state == kAppendedRecordID ||
+                  _state == kAppendedTypeBits);
         BufBuilder newBuf;
         newBuf.appendBuf(_buffer.buf(), _buffer.len());
         return {version, TypeBits(_typeBits), static_cast<size_t>(newBuf.len()), newBuf.release()};
@@ -384,14 +419,25 @@ public:
 
     void appendRecordId(RecordId loc);
     void appendTypeBits(const TypeBits& bits);
+    void appendBSONElement(const BSONElement& elem);
 
     /**
      * Resets to an empty state.
-     * Equivalent to but faster than *this = Builder()
+     * Equivalent to but faster than *this = Builder(ord, discriminator)
      */
-    void resetToEmpty() {
-        _buffer.reset();
-        _typeBits.reset();
+    void resetToEmpty(Ordering ord = ALL_ASCENDING,
+                      Discriminator discriminator = kInclusive) {
+        if (_state == kReleased) {
+            _buffer = BufBuilder();
+            _typeBits = TypeBits(version);
+        } else {
+            _buffer.reset();
+            _typeBits.reset();
+        }
+        _elemCount = 0;
+        _ordering = ord;
+        _discriminator = discriminator;
+        _transition(kEmpty);
     }
 
     void resetToKey(const BSONObj& obj, Ordering ord, RecordId recordId);
@@ -402,18 +448,22 @@ public:
     }
 
     const char* getBuffer() const {
+        invariant(_state != kReleased);
         return _buffer.buf();
     }
 
     size_t getSize() const {
+        invariant(_state != kReleased);
         return _buffer.len();
     }
 
     bool isEmpty() const {
+        invariant(_state != kReleased);
         return _buffer.len() == 0;
     }
 
     const TypeBits& getTypeBits() const {
+        invariant(_state != kReleased);
         return _typeBits;
     }
 
@@ -432,9 +482,7 @@ public:
     const Version version;
 
 private:
-    void _appendAllElementsForIndexing(const BSONObj& obj,
-                                       Ordering ord,
-                                       Discriminator discriminator);
+    void _appendAllElementsForIndexing(const BSONObj& obj, Discriminator discriminator);
 
     void _appendBool(bool val, bool invert);
     void _appendDate(Date_t val, bool invert);
@@ -471,13 +519,64 @@ private:
     void _appendDoubleWithoutTypeBits(const double num, DecimalContinuationMarker dcm, bool invert);
     void _appendHugeDecimalWithoutTypeBits(const Decimal128 dec, bool invert);
     void _appendTinyDecimalWithoutTypeBits(const Decimal128 dec, const double bin, bool invert);
+    void _appendDiscriminator(const Discriminator discriminator);
+    void _appendEnd();
 
     template <typename T>
     void _append(const T& thing, bool invert);
     void _appendBytes(const void* source, size_t bytes, bool invert);
 
+    void _prepareForRelease() {
+        if (_state == kAppendingBSONElements) {
+            _appendDiscriminator(_discriminator);
+        }
+    }
+
+    void _transition(BuildState to) {
+        // We can empty at any point since it just means that we are clearing the buffer.
+        if (to == kEmpty) {
+            _state = to;
+            return;
+        }
+
+        switch (_state) {
+            case kEmpty:
+                invariant(to == kAppendingBSONElements || to == kAppendedRecordID);
+                break;
+            case kAppendingBSONElements:
+                invariant(to == kEndAdded);
+                break;
+            case kEndAdded:
+                invariant(to == kAppendedRecordID || to == kReleased);
+                break;
+            case kAppendedRecordID:
+                // This first case is the special case in
+                // WiredTigerIndexUnique::_insertTimestampUnsafe.
+                // The third case is the case when we are appending a list of RecordIDs as in the
+                // KeyString test RecordIDs.
+                invariant(to == kAppendedTypeBits || to == kReleased || to == kAppendedRecordID);
+                break;
+            case kAppendedTypeBits:
+                // This first case is the special case in
+                // WiredTigerIndexUnique::_insertTimestampUnsafe.
+                invariant(to == kAppendedRecordID || to == kReleased);
+                break;
+            case kReleased:
+                invariant(to == kEmpty);
+                break;
+            default:
+                MONGO_UNREACHABLE;
+        }
+        _state = to;
+    }
+
+
     TypeBits _typeBits;
-    StackBufBuilder _buffer;
+    BufBuilder _buffer;
+    enum BuildState _state;
+    int _elemCount;
+    Ordering _ordering;
+    Discriminator _discriminator;
 };
 
 inline bool operator<(const Builder& lhs, const Builder& rhs) {
diff --git a/src/mongo/db/storage/key_string_test.cpp b/src/mongo/db/storage/key_string_test.cpp
index 75bab217857..46716b908c7 100644
--- a/src/mongo/db/storage/key_string_test.cpp
+++ b/src/mongo/db/storage/key_string_test.cpp
@@ -71,6 +71,15 @@ BSONObj toBsonAndCheckKeySize(const KeyString::Builder& ks, Ordering ord) {
     return KeyString::toBson(ks.getBuffer(), KeyStringBuilderSize, ord, ks.getTypeBits());
 }
 
+BSONObj toBsonAndCheckKeySize(const KeyString::Value& ks, Ordering ord) {
+    auto KeyStringSize = ks.getSize();
+
+    // Validate size of the key in KeyString::Builder.
+    ASSERT_EQUALS(KeyStringSize,
+                  KeyString::getKeySize(ks.getBuffer(), KeyStringSize, ord, ks.getTypeBits()));
+    return KeyString::toBson(ks.getBuffer(), KeyStringSize, ord, ks.getTypeBits());
+}
+
 Ordering ALL_ASCENDING = Ordering::make(BSONObj());
 Ordering ONE_ASCENDING = Ordering::make(BSON("a" << 1));
 Ordering ONE_DESCENDING = Ordering::make(BSON("a" << -1));
@@ -462,10 +471,10 @@ TEST_F(KeyStringBuilderTest, KeyStringValue) {
     // Test that KeyStringBuilder is releasable into a Value type that is comparable. Once
     // released, it is reusable once reset.
     KeyString::Builder ks1(KeyString::Version::V1, BSON("" << 1), ALL_ASCENDING);
-    KeyString::Value data1 = ks1.getValue();
+    KeyString::Value data1 = ks1.release();
 
     KeyString::Builder ks2(KeyString::Version::V1, BSON("" << 2), ALL_ASCENDING);
-    KeyString::Value data2 = ks2.getValue();
+    KeyString::Value data2 = ks2.release();
 
     ASSERT(data2.compare(data1) > 0);
     ASSERT(data1.compare(data2) < 0);
@@ -480,6 +489,115 @@ TEST_F(KeyStringBuilderTest, KeyStringValue) {
     ASSERT(data2.compare(dataCopy) == 0);
 }
 
+#define COMPARE_KS_BSON(ks, bson, order)                             \
+    do {                                                             \
+        const BSONObj _converted = toBsonAndCheckKeySize(ks, order); \
+        ASSERT_BSONOBJ_EQ(_converted, bson);                         \
+        ASSERT(_converted.binaryEqual(bson));                        \
+    } while (0)
+
+TEST_F(KeyStringBuilderTest, KeyStringValueReleaseReusableTest) {
+    // Test that KeyStringBuilder is reusable once reset.
+    BSONObj doc1 = BSON("fieldA" << 1 << "fieldB" << 2);
+    BSONObj doc2 = BSON("fieldA" << 2 << "fieldB" << 3);
+    BSONObj bson1 = BSON("" << 1 << "" << 2);
+    BSONObj bson2 = BSON("" << 2 << "" << 3);
+    KeyString::Builder ks1(KeyString::Version::V1);
+    ks1.appendBSONElement(doc1["fieldA"]);
+    ks1.appendBSONElement(doc1["fieldB"]);
+    KeyString::Value data1 = ks1.release();
+
+    ks1.resetToEmpty();
+    ks1.appendBSONElement(doc2["fieldA"]);
+    ks1.appendBSONElement(doc2["fieldB"]);
+    KeyString::Value data2 = ks1.release();
+    COMPARE_KS_BSON(data1, bson1, ALL_ASCENDING);
+    COMPARE_KS_BSON(data2, bson2, ALL_ASCENDING);
+}
+
+TEST_F(KeyStringBuilderTest, KeyStringGetValueCopyTest) {
+    // Test that KeyStringGetValueCopyTest creates a copy.
+    BSONObj doc = BSON("fieldA" << 1);
+    KeyString::Builder ks(KeyString::Version::V1, ALL_ASCENDING);
+    ks.appendBSONElement(doc["fieldA"]);
+    KeyString::Value data1 = ks.getValueCopy();
+    KeyString::Value data2 = ks.release();
+
+    // Assert that a copy was actually made and they don't share a buffer.
+    ASSERT_NOT_EQUALS(data1.getBuffer(), data2.getBuffer());
+
+    COMPARE_KS_BSON(data1, BSON("" << 1), ALL_ASCENDING);
+    COMPARE_KS_BSON(data2, BSON("" << 1), ALL_ASCENDING);
+}
+
+TEST_F(KeyStringBuilderTest, KeyStringBuilderAppendBsonElement) {
+    // Test that appendBsonElement works.
+    {
+        BSONObj doc = BSON("fieldA" << 1 << "fieldB" << 2);
+        KeyString::Builder ks(KeyString::Version::V1, ALL_ASCENDING);
+        ks.appendBSONElement(doc["fieldA"]);
+        ks.appendBSONElement(doc["fieldB"]);
+        KeyString::Value data = ks.release();
+        COMPARE_KS_BSON(data, BSON("" << 1 << "" << 2), ALL_ASCENDING);
+    }
+
+    {
+        BSONObj doc = BSON("fieldA" << 1 << "fieldB" << 2);
+        KeyString::Builder ks(KeyString::Version::V1, ONE_DESCENDING);
+        ks.appendBSONElement(doc["fieldA"]);
+        ks.appendBSONElement(doc["fieldB"]);
+        KeyString::Value data = ks.release();
+        COMPARE_KS_BSON(data, BSON("" << 1 << "" << 2), ONE_DESCENDING);
+    }
+
+    {
+        BSONObj doc = BSON("fieldA"
+                           << "value1"
+                           << "fieldB"
+                           << "value2");
+        KeyString::Builder ks(KeyString::Version::V1, ONE_DESCENDING);
+        ks.appendBSONElement(doc["fieldA"]);
+        ks.appendBSONElement(doc["fieldB"]);
+        KeyString::Value data = ks.release();
+        COMPARE_KS_BSON(data,
+                        BSON(""
+                             << "value1"
+                             << ""
+                             << "value2"),
+                        ONE_DESCENDING);
+    }
+}
+
+TEST_F(KeyStringBuilderTest, KeyStringBuilderOrdering) {
+    // Test that ordering works.
+    BSONObj doc = BSON("fieldA" << 1);
+    KeyString::Builder ks1(KeyString::Version::V1, ALL_ASCENDING);
+    ks1.appendBSONElement(doc["fieldA"]);
+    KeyString::Builder ks2(KeyString::Version::V1, ONE_DESCENDING);
+    ks2.appendBSONElement(doc["fieldA"]);
+    KeyString::Value data1 = ks1.release();
+    KeyString::Value data2 = ks2.release();
+
+    ASSERT_EQUALS(data1.getSize(), data2.getSize());
+    // Confirm that the buffers are different, indicating that the data is stored inverted in the
+    // second.
+    ASSERT_NE(0, memcmp(data1.getBuffer(), data2.getBuffer(), data1.getSize()));
+}
+
+TEST_F(KeyStringBuilderTest, KeyStringBuilderDiscriminator) {
+    // test that when passed in a Discriminator it gets added.
+    BSONObj doc = BSON("fieldA" << 1 << "fieldB" << 2);
+    KeyString::Builder ks(
+        KeyString::Version::V1, ALL_ASCENDING, KeyString::Builder::kExclusiveBefore);
+    ks.appendBSONElement(doc["fieldA"]);
+    ks.appendBSONElement(doc["fieldB"]);
+    KeyString::Value data = ks.release();
+    uint8_t appendedDescriminator = (uint8_t)(*(data.getBuffer() + (data.getSize() - 2)));
+    uint8_t end = (uint8_t)(*(data.getBuffer() + (data.getSize() - 1)));
+    ASSERT_EQ((uint8_t)'\001', appendedDescriminator);
+    ASSERT_EQ((uint8_t)'\004', end);
+}
+
 TEST_F(KeyStringBuilderTest, LotsOfNumbers1) {
     for (int i = 0; i < 64; i++) {
         int64_t x = 1LL << i;