SERVER-43619 Overload SortedDataInterface::Cursor to return KeyString

9 files changed, 252 insertions, 67 deletions

diff --git a/src/mongo/db/storage/biggie/biggie_sorted_impl.cpp b/src/mongo/db/storage/biggie/biggie_sorted_impl.cpp
index 05166bc7307..d55ed2d1bf6 100644
--- a/src/mongo/db/storage/biggie/biggie_sorted_impl.cpp
+++ b/src/mongo/db/storage/biggie/biggie_sorted_impl.cpp
@@ -597,6 +597,35 @@ SortedDataInterface::Cursor::Cursor(OperationContext* opCtx,
       _KSForIdentStart(_KSForIdentStart),
       _KSForIdentEnd(identEndBSON) {}
 
+bool SortedDataInterface::Cursor::advanceNext() {
+    if (!_atEOF) {
+        // If the last move was restore, then we don't need to advance the cursor, since the user
+        // never got the value the cursor was pointing to in the first place. However,
+        // _lastMoveWasRestore will go through extra logic on a unique index, since unique indexes
+        // are not allowed to return the same key twice.
+        if (_lastMoveWasRestore) {
+            _lastMoveWasRestore = false;
+        } else {
+            // We basically just check to make sure the cursor is in the ident.
+            if (_forward && checkCursorValid()) {
+                ++_forwardIt;
+            } else if (!_forward && checkCursorValid()) {
+                ++_reverseIt;
+            }
+            // We check here to make sure that we are on the correct side of the end position, and
+            // that the cursor is still in the ident after advancing.
+            if (!checkCursorValid()) {
+                _atEOF = true;
+                return false;
+            }
+        }
+    } else {
+        _lastMoveWasRestore = false;
+        return false;
+    }
+    return true;
+}
+
 // This function checks whether or not the cursor end position was set by the user or not.
 bool SortedDataInterface::Cursor::endPosSet() {
     return (_forward && _endPos != boost::none) || (!_forward && _endPosReverse != boost::none);
@@ -691,34 +720,8 @@ void SortedDataInterface::Cursor::setEndPosition(const BSONObj& key, bool inclus
 }
 
 boost::optional<IndexKeyEntry> SortedDataInterface::Cursor::next(RequestedInfo parts) {
-    if (!_atEOF) {
-        // If the last move was restore, then we don't need to advance the cursor, since the user
-        // never got the value the cursor was pointing to in the first place. However,
-        // _lastMoveWasRestore will go through extra logic on a unique index, since unique indexes
-        // are not allowed to return the same key twice.
-        if (_lastMoveWasRestore) {
-            _lastMoveWasRestore = false;
-        } else {
-            // We basically just check to make sure the cursor is in the ident.
-            if (_forward) {
-                if (checkCursorValid()) {
-                    ++_forwardIt;
-                }
-            } else {
-                if (checkCursorValid()) {
-                    ++_reverseIt;
-                }
-            }
-            // We check here to make sure that we are on the correct side of the end position, and
-            // that the cursor is still in the ident after advancing.
-            if (!checkCursorValid()) {
-                _atEOF = true;
-                return boost::none;
-            }
-        }
-    } else {
-        _lastMoveWasRestore = false;
-        return boost::none;
+    if (!advanceNext()) {
+        return {};
     }
 
     if (_forward) {
@@ -727,6 +730,17 @@ boost::optional<IndexKeyEntry> SortedDataInterface::Cursor::next(RequestedInfo p
     return keyStringToIndexKeyEntry(_reverseIt->first, _reverseIt->second, _order);
 }
 
+boost::optional<KeyStringEntry> SortedDataInterface::Cursor::nextKeyString() {
+    if (!advanceNext()) {
+        return {};
+    }
+
+    if (_forward) {
+        return keyStringToKeyStringEntry(_forwardIt->first, _forwardIt->second, _order);
+    }
+    return keyStringToKeyStringEntry(_reverseIt->first, _reverseIt->second, _order);
+}
+
 boost::optional<IndexKeyEntry> SortedDataInterface::Cursor::seekAfterProcessing(BSONObj finalKey) {
     std::string workingCopyBound;
 
diff --git a/src/mongo/db/storage/biggie/biggie_sorted_impl.h b/src/mongo/db/storage/biggie/biggie_sorted_impl.h
index c2c6f638fd2..acf8414b9eb 100644
--- a/src/mongo/db/storage/biggie/biggie_sorted_impl.h
+++ b/src/mongo/db/storage/biggie/biggie_sorted_impl.h
@@ -118,6 +118,7 @@ public:
                std::string KSForIdentEnd);
         virtual void setEndPosition(const BSONObj& key, bool inclusive) override;
         virtual boost::optional<IndexKeyEntry> next(RequestedInfo parts = kKeyAndLoc) override;
+        virtual boost::optional<KeyStringEntry> nextKeyString() override;
         virtual boost::optional<IndexKeyEntry> seek(const KeyString::Value& keyString,
                                                     RequestedInfo parts = kKeyAndLoc) override;
         virtual boost::optional<KeyStringEntry> seekForKeyString(
@@ -132,6 +133,7 @@ public:
         virtual void reattachToOperationContext(OperationContext* opCtx) override;
 
     private:
+        bool advanceNext();
         // This is a helper function to check if the cursor was explicitly set by the user or not.
         bool endPosSet();
         // This is a helper function to check if the cursor is valid or not.
diff --git a/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_btree_impl.cpp b/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_btree_impl.cpp
index 4568596e939..8eb549629e2 100644
--- a/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_btree_impl.cpp
+++ b/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_btree_impl.cpp
@@ -283,6 +283,18 @@ public:
             return *_it;
         }
 
+        boost::optional<KeyStringEntry> nextKeyString() override {
+            boost::optional<IndexKeyEntry> indexKeyEntry = next(RequestedInfo::kKeyAndLoc);
+            if (!indexKeyEntry) {
+                return {};
+            }
+
+            KeyString::Builder builder(
+                KeyString::Version::kLatestVersion, indexKeyEntry->key, _ordering);
+            builder.appendRecordId(indexKeyEntry->loc);
+            return KeyStringEntry(builder.getValueCopy(), indexKeyEntry->loc);
+        }
+
         void setEndPosition(const BSONObj& key, bool inclusive) override {
             if (key.isEmpty()) {
                 // This means scan to end of index.
diff --git a/src/mongo/db/storage/key_string.h b/src/mongo/db/storage/key_string.h
index dc739804204..74fb80406e3 100644
--- a/src/mongo/db/storage/key_string.h
+++ b/src/mongo/db/storage/key_string.h
@@ -609,6 +609,11 @@ public:
         return _buffer.len() == 0;
     }
 
+    void setTypeBits(const TypeBits& typeBits) {
+        invariant(_state != BuildState::kReleased);
+        _typeBits = typeBits;
+    }
+
     const TypeBits& getTypeBits() const {
         invariant(_state != BuildState::kReleased);
         return _typeBits;
diff --git a/src/mongo/db/storage/mobile/mobile_index.cpp b/src/mongo/db/storage/mobile/mobile_index.cpp
index 7851a290bab..0c0808b5dc3 100644
--- a/src/mongo/db/storage/mobile/mobile_index.cpp
+++ b/src/mongo/db/storage/mobile/mobile_index.cpp
@@ -373,16 +373,22 @@ public:
     virtual ~CursorBase() {}
 
     boost::optional<IndexKeyEntry> next(RequestedInfo parts) override {
-        if (_isEOF) {
+        if (!_advanceNext()) {
             return {};
         }
-
-        _advance();
-        _updatePosition();
-
         return getCurrentEntry(parts);
     }
 
+    boost::optional<KeyStringEntry> nextKeyString() override {
+        if (!_advanceNext()) {
+            return {};
+        }
+        if (_isEOF) {
+            return {};
+        }
+        return _getKeyStringEntry();
+    }
+
     void setEndPosition(const BSONObj& key, bool inclusive) override {
         // Scan to end of index.
         if (key.isEmpty()) {
@@ -417,18 +423,7 @@ public:
         if (_isEOF) {
             return {};
         }
-        auto sizeWithoutRecordId = KeyString::getKeySize(_savedKey.getBuffer(),
-                                                         _savedKey.getSize(),
-                                                         _index.getOrdering(),
-                                                         _savedKey.getTypeBits());
-        if (_savedKey.getSize() == sizeWithoutRecordId) {
-            // Create a copy of _key with a RecordId. Because _key is used during cursor restore(),
-            // appending the RecordId would cause the cursor to be repositioned incorrectly.
-            KeyString::Builder keyWithRecordId(_savedKey);
-            keyWithRecordId.appendRecordId(_savedRecId);
-            return KeyStringEntry(keyWithRecordId.getValueCopy(), _savedRecId);
-        }
-        return KeyStringEntry(_savedKey.getValueCopy(), _savedRecId);
+        return _getKeyStringEntry();
     }
 
     boost::optional<IndexKeyEntry> seekExact(const KeyString::Value& keyStringValue,
@@ -516,6 +511,34 @@ public:
     }
 
 protected:
+    bool _advanceNext() {
+        if (_isEOF) {
+            return false;
+        }
+
+        _advance();
+        _updatePosition();
+        return true;
+    }
+
+    KeyStringEntry _getKeyStringEntry() {
+        auto sizeWithoutRecordId = KeyString::getKeySize(_savedKey.getBuffer(),
+                                                         _savedKey.getSize(),
+                                                         _index.getOrdering(),
+                                                         _savedKey.getTypeBits());
+        if (_savedKey.getSize() == sizeWithoutRecordId) {
+            // Create a copy of _key with a RecordId. Because _key is used during cursor restore(),
+            // appending the RecordId would cause the cursor to be repositioned incorrectly.
+            KeyString::Builder keyWithRecordId(_savedKey);
+            keyWithRecordId.appendRecordId(_savedRecId);
+            keyWithRecordId.setTypeBits(_savedTypeBits);
+            return KeyStringEntry(keyWithRecordId.getValueCopy(), _savedRecId);
+        }
+
+        _savedKey.setTypeBits(_savedTypeBits);
+        return KeyStringEntry(_savedKey.getValueCopy(), _savedRecId);
+    }
+
     /**
      * Advances the cursor and determines if end reached.
      */
diff --git a/src/mongo/db/storage/mobile/mobile_index.h b/src/mongo/db/storage/mobile/mobile_index.h
index cbea32d200d..a660ea232f5 100644
--- a/src/mongo/db/storage/mobile/mobile_index.h
+++ b/src/mongo/db/storage/mobile/mobile_index.h
@@ -105,6 +105,10 @@ public:
     }
 
 protected:
+    bool _advanceNext();
+
+    KeyStringEntry _getKeyStringEntry();
+
     bool _isDup(OperationContext* opCtx, const KeyString::Value& key);
 
     /**
diff --git a/src/mongo/db/storage/sorted_data_interface.h b/src/mongo/db/storage/sorted_data_interface.h
index d546b3b14c6..448d79bf52b 100644
--- a/src/mongo/db/storage/sorted_data_interface.h
+++ b/src/mongo/db/storage/sorted_data_interface.h
@@ -252,6 +252,7 @@ public:
          * If not positioned, returns boost::none.
          */
         virtual boost::optional<IndexKeyEntry> next(RequestedInfo parts = kKeyAndLoc) = 0;
+        virtual boost::optional<KeyStringEntry> nextKeyString() = 0;
 
         //
         // Seeking
diff --git a/src/mongo/db/storage/sorted_data_interface_test_cursor.cpp b/src/mongo/db/storage/sorted_data_interface_test_cursor.cpp
index 5bfc4467950..61a8cffd93a 100644
--- a/src/mongo/db/storage/sorted_data_interface_test_cursor.cpp
+++ b/src/mongo/db/storage/sorted_data_interface_test_cursor.cpp
@@ -129,6 +129,53 @@ TEST(SortedDataInterface, ExhaustCursor) {
     }
 }
 
+TEST(SortedDataInterface, ExhaustKeyStringCursor) {
+    const auto harnessHelper(newSortedDataInterfaceHarnessHelper());
+    const std::unique_ptr<SortedDataInterface> sorted(
+        harnessHelper->newSortedDataInterface(/*unique=*/false, /*partial=*/false));
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT(sorted->isEmpty(opCtx.get()));
+    }
+
+    std::vector<KeyString::Value> keyStrings;
+    int nToInsert = 10;
+    for (int i = 0; i < nToInsert; i++) {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        {
+            WriteUnitOfWork uow(opCtx.get());
+            BSONObj key = BSON("" << i);
+            RecordId loc(42, i * 2);
+            KeyString::Value ks = makeKeyString(sorted.get(), key, loc);
+            keyStrings.push_back(ks);
+            ASSERT_OK(sorted->insert(opCtx.get(), ks, true));
+            uow.commit();
+        }
+    }
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT_EQUALS(nToInsert, sorted->numEntries(opCtx.get()));
+    }
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        const std::unique_ptr<SortedDataInterface::Cursor> cursor(sorted->newCursor(opCtx.get()));
+        for (int i = 0; i < nToInsert; i++) {
+            auto entry = i == 0 ? cursor->seekForKeyString(
+                                      makeKeyStringForSeek(sorted.get(), BSONObj(), true, true))
+                                : cursor->nextKeyString();
+            ASSERT(entry);
+            ASSERT_EQ(entry->keyString, keyStrings.at(i));
+        }
+        ASSERT(!cursor->nextKeyString());
+
+        // Cursor at EOF should remain at EOF when advanced
+        ASSERT(!cursor->nextKeyString());
+    }
+}
+
 // Call advance() on a reverse cursor until it is exhausted.
 // When a cursor positioned at EOF is advanced, it stays at EOF.
 TEST(SortedDataInterface, ExhaustCursorReversed) {
@@ -175,6 +222,54 @@ TEST(SortedDataInterface, ExhaustCursorReversed) {
     }
 }
 
+TEST(SortedDataInterface, ExhaustKeyStringCursorReversed) {
+    const auto harnessHelper(newSortedDataInterfaceHarnessHelper());
+    const std::unique_ptr<SortedDataInterface> sorted(
+        harnessHelper->newSortedDataInterface(/*unique=*/false, /*partial=*/false));
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT(sorted->isEmpty(opCtx.get()));
+    }
+
+    std::vector<KeyString::Value> keyStrings;
+    int nToInsert = 10;
+    for (int i = 0; i < nToInsert; i++) {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        {
+            WriteUnitOfWork uow(opCtx.get());
+            BSONObj key = BSON("" << i);
+            RecordId loc(42, i * 2);
+            KeyString::Value ks = makeKeyString(sorted.get(), key, loc);
+            keyStrings.push_back(ks);
+            ASSERT_OK(sorted->insert(opCtx.get(), ks, true));
+            uow.commit();
+        }
+    }
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT_EQUALS(nToInsert, sorted->numEntries(opCtx.get()));
+    }
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        const std::unique_ptr<SortedDataInterface::Cursor> cursor(
+            sorted->newCursor(opCtx.get(), false));
+        for (int i = nToInsert - 1; i >= 0; i--) {
+            auto entry = (i == nToInsert - 1) ? cursor->seekForKeyString(makeKeyStringForSeek(
+                                                    sorted.get(), kMaxBSONKey, false, true))
+                                              : cursor->nextKeyString();
+            ASSERT(entry);
+            ASSERT_EQ(entry->keyString, keyStrings.at(i));
+        }
+        ASSERT(!cursor->nextKeyString());
+
+        // Cursor at EOF should remain at EOF when advanced
+        ASSERT(!cursor->nextKeyString());
+    }
+}
+
 void testBoundaries(bool unique, bool forward, bool inclusive) {
     const auto harnessHelper(newSortedDataInterfaceHarnessHelper());
     const std::unique_ptr<SortedDataInterface> sorted(
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp
index 50ed41b44bb..e8b74cb5193 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp
@@ -797,16 +797,23 @@ public:
     }
 
     boost::optional<IndexKeyEntry> next(RequestedInfo parts) override {
-        // Advance on a cursor at the end is a no-op
-        if (_eof)
+        if (!advanceNext()) {
             return {};
-
-        if (!_lastMoveSkippedKey)
-            advanceWTCursor();
-        updatePosition(true);
+        }
         return curr(parts);
     }
 
+    boost::optional<KeyStringEntry> nextKeyString() override {
+        if (!advanceNext()) {
+            return {};
+        }
+        if (_eof) {
+            return {};
+        }
+
+        return getKeyStringEntry();
+    }
+
     void setEndPosition(const BSONObj& key, bool inclusive) override {
         TRACE_CURSOR << "setEndPosition inclusive: " << inclusive << ' ' << key;
         if (key.isEmpty()) {
@@ -842,19 +849,7 @@ public:
         dassert(!atOrPastEndPointAfterSeeking());
         dassert(!_id.isNull());
 
-        // Most keys will have a RecordId appeneded to the end, with the exception of the _id index
-        // and timestamp unsafe unique indexes. The contract of this function is to always return a
-        // KeyString with a RecordId, so append one if it does not exists already.
-        auto sizeWithoutRecordId = KeyString::getKeySize(
-            _key.getBuffer(), _key.getSize(), _idx.getOrdering(), _key.getTypeBits());
-        if (_key.getSize() == sizeWithoutRecordId) {
-            // Create a copy of _key with a RecordId. Because _key is used during cursor restore(),
-            // appending the RecordId would cause the cursor to be repositioned incorrectly.
-            KeyString::Builder keyWithRecordId(_key);
-            keyWithRecordId.appendRecordId(_id);
-            return KeyStringEntry(keyWithRecordId.getValueCopy(), _id);
-        }
-        return KeyStringEntry(_key.getValueCopy(), _id);
+        return getKeyStringEntry();
     }
 
     boost::optional<KeyStringEntry> seekExactForKeyString(const KeyString::Value& key) override {
@@ -1147,6 +1142,40 @@ protected:
         updateIdAndTypeBits();
     }
 
+    bool advanceNext() {
+        // Advance on a cursor at the end is a no-op.
+        if (_eof) {
+            return false;
+        }
+        if (!_lastMoveSkippedKey) {
+            advanceWTCursor();
+        }
+        updatePosition(true);
+        return true;
+    }
+
+    KeyStringEntry getKeyStringEntry() {
+        // Most keys will have a RecordId appended to the end, with the exception of the _id index
+        // and timestamp unsafe unique indexes. The contract of this function is to always return a
+        // KeyString with a RecordId, so append one if it does not exists already.
+        auto sizeWithoutRecordId =
+            KeyString::getKeySize(_key.getBuffer(), _key.getSize(), _idx.getOrdering(), _typeBits);
+        if (_key.getSize() == sizeWithoutRecordId) {
+            // Create a copy of _key with a RecordId. Because _key is used during cursor restore(),
+            // appending the RecordId would cause the cursor to be repositioned incorrectly.
+            KeyString::Builder keyWithRecordId(_key);
+            keyWithRecordId.appendRecordId(_id);
+            keyWithRecordId.setTypeBits(_typeBits);
+
+            TRACE_CURSOR << " returning " << keyWithRecordId << ' ' << _id;
+            return KeyStringEntry(keyWithRecordId.getValueCopy(), _id);
+        }
+
+        _key.setTypeBits(_typeBits);
+        TRACE_CURSOR << " returning " << _key << ' ' << _id;
+        return KeyStringEntry(_key.getValueCopy(), _id);
+    }
+
     OperationContext* _opCtx;
     boost::optional<WiredTigerCursor> _cursor;
     const WiredTigerIndex& _idx;  // not owned
@@ -1200,8 +1229,8 @@ public:
             return;
         }
 
-        auto keySize = KeyString::getKeySize(
-            _key.getBuffer(), _key.getSize(), _idx.getOrdering(), _key.getTypeBits());
+        auto keySize =
+            KeyString::getKeySize(_key.getBuffer(), _key.getSize(), _idx.getOrdering(), _typeBits);
 
         if (_key.getSize() == keySize) {
             _updateIdAndTypeBitsFromValue();
@@ -1237,7 +1266,7 @@ public:
 
             // Get the size of the prefix key
             auto keySize = KeyString::getKeySize(
-                _key.getBuffer(), _key.getSize(), _idx.getOrdering(), _key.getTypeBits());
+                _key.getBuffer(), _key.getSize(), _idx.getOrdering(), _typeBits);
 
             // This check is only to avoid returning the same key again after a restore. Keys
             // shorter than _key cannot have "prefix key" same as _key. Therefore we care only about