SERVER-41720 Overload SortedDataInterface::Cursor::seek and seekExact to accept KeyStrings

10 files changed, 471 insertions, 71 deletions

diff --git a/src/mongo/db/storage/biggie/biggie_sorted_impl.cpp b/src/mongo/db/storage/biggie/biggie_sorted_impl.cpp
index f6a5003edde..96bf3988fe6 100644
--- a/src/mongo/db/storage/biggie/biggie_sorted_impl.cpp
+++ b/src/mongo/db/storage/biggie/biggie_sorted_impl.cpp
@@ -100,14 +100,12 @@ std::string createKeyString(const BSONObj& key,
 }
 
 std::string createKeyString(const KeyString::Value& keyString,
+                            const size_t size,
                             const RecordId& loc,
                             std::string prefixToUse,
                             bool isUnique) {
     KeyString::Builder ks(KeyString::Version::V1);
-    auto sizeWithoutRecordId =
-        KeyString::sizeWithoutRecordIdAtEnd(keyString.getBuffer(), keyString.getSize());
-    ks.resetFromBuffer(keyString.getBuffer(), sizeWithoutRecordId);
-
+    ks.resetFromBuffer(keyString.getBuffer(), size);
     prefixKeyString(&ks, loc, prefixToUse, isUnique);
     return std::string(ks.getBuffer(), ks.getSize());
 }
@@ -124,7 +122,7 @@ bool keysAreIdentical(std::string ks1, std::string ks2, bool isUnique) {
 }
 
 /**
- * This function converts a KeyString::Builder into an IndexKeyEntry. We don't need to store the
+ * This function converts a std::string KeyString into a Key. We don't need to store the
  * typebits for the outer key string (the one consisting of the prefix, the key, and the recordId)
  * since those will not be used. However, we do need to store the typebits for the internal
  * keystring (made from the key itself), as those typebits are potentially important.
@@ -132,23 +130,20 @@ bool keysAreIdentical(std::string ks1, std::string ks2, bool isUnique) {
  * The data which is serialized as a byte array, has the following structure:
  *     [RecordId][TypeBits of internal keystring]
  */
-IndexKeyEntry keyStringToIndexKeyEntry(const std::string keyString,
-                                       std::string data,
-                                       const Ordering order) {
-    int64_t ridRepr;
-    std::memcpy(&ridRepr, data.data(), sizeof(int64_t));
-    RecordId rid(ridRepr);
-
+BSONObj getKeyFromKeyString(const std::string& keyString,
+                            const std::string& data,
+                            const Ordering& order) {
     std::string typeBitsString(data.length() - sizeof(int64_t), '\0');
     std::memcpy(&typeBitsString[0], data.data() + sizeof(int64_t), data.length() - sizeof(int64_t));
 
+    BufReader brTbInternal(typeBitsString.c_str(), typeBitsString.length());
+
     KeyString::Version version = KeyString::Version::V1;
     KeyString::TypeBits tbInternal = KeyString::TypeBits(version);
-    KeyString::TypeBits tbOuter = KeyString::TypeBits(version);
 
-    BufReader brTbInternal(typeBitsString.c_str(), typeBitsString.length());
     tbInternal.resetFromBuffer(&brTbInternal);
 
+    KeyString::TypeBits tbOuter = KeyString::TypeBits(version);
     BSONObj bsonObj =
         KeyString::toBsonSafe(keyString.c_str(), keyString.length(), allAscending, tbOuter);
 
@@ -165,8 +160,31 @@ IndexKeyEntry keyStringToIndexKeyEntry(const std::string keyString,
 
     BSONObj key(ConstSharedBuffer{sb});
 
+    return key;
+}
+
+IndexKeyEntry keyStringToIndexKeyEntry(const std::string keyString,
+                                       std::string data,
+                                       const Ordering order) {
+
+    auto key = getKeyFromKeyString(keyString, data, order);
+    int64_t ridRepr;
+    std::memcpy(&ridRepr, data.data(), sizeof(int64_t));
+    RecordId rid(ridRepr);
     return IndexKeyEntry(key, rid);
 }
+
+boost::optional<KeyStringEntry> keyStringToKeyStringEntry(const std::string keyString,
+                                                          std::string data,
+                                                          const Ordering order) {
+    auto key = getKeyFromKeyString(keyString, data, order);
+    int64_t ridRepr;
+    std::memcpy(&ridRepr, data.data(), sizeof(int64_t));
+    RecordId rid(ridRepr);
+    KeyString::Builder ksFinal(KeyString::Version::V1, key, order);
+    ksFinal.appendRecordId(rid);
+    return KeyStringEntry(ksFinal.getValueCopy(), rid);
+}
 }  // namespace
 
 SortedDataBuilderInterface::SortedDataBuilderInterface(OperationContext* opCtx,
@@ -227,7 +245,7 @@ Status SortedDataBuilderInterface::addKey(const KeyString::Value& keyString, con
     }
 
     std::string workingCopyInsertKey =
-        createKeyString(keyString, loc, _prefix, /* isUnique */ _unique);
+        createKeyString(keyString, sizeWithoutRecordId, loc, _prefix, /* isUnique */ _unique);
 
     if (twoKeyCmp == 0 && twoRIDCmp != 0) {
         if (!_dupsAllowed) {
@@ -236,7 +254,8 @@ Status SortedDataBuilderInterface::addKey(const KeyString::Value& keyString, con
         }
         // Duplicate index entries are allowed on this unique index, so we put the RecordId in the
         // KeyString until the unique constraint is resolved.
-        workingCopyInsertKey = createKeyString(keyString, loc, _prefix, /* isUnique */ false);
+        workingCopyInsertKey =
+            createKeyString(keyString, sizeWithoutRecordId, loc, _prefix, /* isUnique */ false);
     }
 
     std::string internalTbString(keyString.getTypeBits().getBuffer(),
@@ -324,8 +343,10 @@ Status SortedDataInterface::insert(OperationContext* opCtx,
     dassert(loc == KeyString::decodeRecordIdAtEnd(keyString.getBuffer(), keyString.getSize()));
 
     StringStore* workingCopy(RecoveryUnit::get(opCtx)->getHead());
-    std::string insertKeyString = createKeyString(keyString, loc, _prefix, _isUnique);
-
+    auto sizeWithoutRecordId =
+        KeyString::sizeWithoutRecordIdAtEnd(keyString.getBuffer(), keyString.getSize());
+    std::string insertKeyString =
+        createKeyString(keyString, sizeWithoutRecordId, loc, _prefix, _isUnique);
     // For unique indexes, if duplicate keys are allowed then we do the following:
     //   - Create the KeyString without the RecordId in it and see if anything exists with that.
     //     - If the cursor didn't find anything, we index with this KeyString.
@@ -343,8 +364,11 @@ Status SortedDataInterface::insert(OperationContext* opCtx,
                 if (dupsAllowed) {
                     // Duplicate index entries are allowed on this unique index, so we put the
                     // RecordId in the KeyString until the unique constraint is resolved.
-                    insertKeyString =
-                        createKeyString(keyString, loc, _prefix, /* isUnique */ false);
+                    insertKeyString = createKeyString(keyString,
+                                                      sizeWithoutRecordId,
+                                                      loc,
+                                                      _prefix,
+                                                      /* isUnique */ false);
                 } else {
                     // There was an attempt to create an index entry with a different RecordId while
                     // dups were not allowed.
@@ -398,6 +422,8 @@ void SortedDataInterface::unindex(OperationContext* opCtx,
     std::string removeKeyString;
     bool erased;
 
+    auto sizeWithoutRecordId =
+        KeyString::sizeWithoutRecordIdAtEnd(keyString.getBuffer(), keyString.getSize());
     if (_isUnique) {
         // For unique indexes, to unindex them we do the following:
         //   - Create the KeyString with or without the RecordId in it depending on dupsAllowed
@@ -407,9 +433,11 @@ void SortedDataInterface::unindex(OperationContext* opCtx,
         //       RecordId in it.
         // This is required because of the way we insert on unique indexes when dups are allowed.
         if (dupsAllowed)
-            removeKeyString = createKeyString(keyString, loc, _prefix, /* isUnique */ false);
+            removeKeyString =
+                createKeyString(keyString, sizeWithoutRecordId, loc, _prefix, /* isUnique */ false);
         else
-            removeKeyString = createKeyString(keyString, loc, _prefix, /* isUnique */ true);
+            removeKeyString =
+                createKeyString(keyString, sizeWithoutRecordId, loc, _prefix, /* isUnique */ true);
 
         // Check that the record id matches when using partial indexes. We may be called to unindex
         // records that are not present in the index due to the partial filter expression.
@@ -422,16 +450,19 @@ void SortedDataInterface::unindex(OperationContext* opCtx,
             // the RecordId in it, and erase that. This could only happen on unique indexes where
             // duplicate index entries were/are allowed.
             if (dupsAllowed)
-                removeKeyString = createKeyString(keyString, loc, _prefix, /* isUnique */ true);
+                removeKeyString = createKeyString(
+                    keyString, sizeWithoutRecordId, loc, _prefix, /* isUnique */ true);
             else
-                removeKeyString = createKeyString(keyString, loc, _prefix, /* isUnique */ false);
+                removeKeyString = createKeyString(
+                    keyString, sizeWithoutRecordId, loc, _prefix, /* isUnique */ false);
 
             if (!ifPartialCheckRecordIdEquals(opCtx, removeKeyString, loc))
                 return;
             erased = workingCopy->erase(removeKeyString);
         }
     } else {
-        removeKeyString = createKeyString(keyString, loc, _prefix, /* isUnique */ false);
+        removeKeyString =
+            createKeyString(keyString, sizeWithoutRecordId, loc, _prefix, /* isUnique */ false);
         erased = workingCopy->erase(removeKeyString);
     }
 
@@ -724,15 +755,29 @@ boost::optional<IndexKeyEntry> SortedDataInterface::Cursor::seekAfterProcessing(
                                                                                 bool inclusive) {
     std::string workingCopyBound;
 
+    KeyString::Builder ks(KeyString::Version::V1, finalKey, _order);
+    auto ksEntry = seekAfterProcessing(ks.getValueCopy(), inclusive);
+
+    const BSONObj bson = KeyString::toBson(ksEntry->keyString.getBuffer(),
+                                           ksEntry->keyString.getSize(),
+                                           _order,
+                                           ksEntry->keyString.getTypeBits());
+    return IndexKeyEntry(bson, ksEntry->loc);
+}
+
+boost::optional<KeyStringEntry> SortedDataInterface::Cursor::seekAfterProcessing(
+    const KeyString::Value& keyStringVal, bool inclusive) {
+    std::string workingCopyBound;
     // Similar to above, if forward and inclusive or reverse and not inclusive, then use min() for
     // recordId. Else, we should use max().
     if (_forward == inclusive) {
-        workingCopyBound = createKeyString(finalKey, RecordId::min(), _prefix, _order, _isUnique);
+        workingCopyBound = createKeyString(
+            keyStringVal, keyStringVal.getSize(), RecordId::min(), _prefix, _isUnique);
     } else {
-        workingCopyBound = createKeyString(finalKey, RecordId::max(), _prefix, _order, _isUnique);
+        workingCopyBound = createKeyString(
+            keyStringVal, keyStringVal.getSize(), RecordId::max(), _prefix, _isUnique);
     }
-
-    if (finalKey.isEmpty()) {
+    if (keyStringVal.isEmpty()) {
         // If the key is empty and it's not inclusive, then no elements satisfy this seek.
         if (!inclusive) {
             _atEOF = true;
@@ -784,31 +829,83 @@ boost::optional<IndexKeyEntry> SortedDataInterface::Cursor::seekAfterProcessing(
     }
 
     // Everything checks out, so we have successfullly seeked and now return.
+    boost::optional<IndexKeyEntry> indexKeyEntry;
     if (_forward) {
-        return keyStringToIndexKeyEntry(_forwardIt->first, _forwardIt->second, _order);
+        return keyStringToKeyStringEntry(_forwardIt->first, _forwardIt->second, _order);
     }
-    return keyStringToIndexKeyEntry(_reverseIt->first, _reverseIt->second, _order);
+    return keyStringToKeyStringEntry(_reverseIt->first, _reverseIt->second, _order);
 }
 
 boost::optional<IndexKeyEntry> SortedDataInterface::Cursor::seek(const BSONObj& key,
                                                                  bool inclusive,
-                                                                 RequestedInfo parts) {
+                                                                 RequestedInfo) {
     BSONObj finalKey = BSONObj::stripFieldNames(key);
-    _lastMoveWasRestore = false;
-    _atEOF = false;
-
-    return seekAfterProcessing(finalKey, inclusive);
+    KeyString::Builder keyString(KeyString::Version::V1, finalKey, _order);
+    auto ksValue = seek(keyString.getValueCopy(), inclusive);
+    if (ksValue) {
+        BSONObj bson = KeyString::toBson(ksValue->keyString.getBuffer(),
+                                         ksValue->keyString.getSize(),
+                                         _order,
+                                         ksValue->keyString.getTypeBits());
+        return IndexKeyEntry(bson, ksValue->loc);
+    }
+    return boost::none;
 }
 
 boost::optional<IndexKeyEntry> SortedDataInterface::Cursor::seek(const IndexSeekPoint& seekPoint,
                                                                  RequestedInfo parts) {
     const BSONObj key = IndexEntryComparison::makeQueryObject(seekPoint, _forward);
-    _atEOF = false;
-    bool inclusive = true;
-    BSONObj finalKey = key;
+    KeyString::Builder keyString(KeyString::Version::V1, key, _order);
+    auto ksValue = seek(keyString.getValueCopy(), true);
+    if (ksValue) {
+        BSONObj bson = KeyString::toBson(ksValue->keyString.getBuffer(),
+                                         ksValue->keyString.getSize(),
+                                         _order,
+                                         ksValue->keyString.getTypeBits());
+        return IndexKeyEntry(bson, ksValue->loc);
+    }
+    return boost::none;
+}
+
+boost::optional<KeyStringEntry> SortedDataInterface::Cursor::seek(
+    const KeyString::Value& keyStringValue, bool inclusive) {
     _lastMoveWasRestore = false;
+    _atEOF = false;
+    return seekAfterProcessing(keyStringValue, inclusive);
+}
 
-    return seekAfterProcessing(finalKey, inclusive);
+boost::optional<IndexKeyEntry> SortedDataInterface::Cursor::seekExact(const BSONObj& key,
+                                                                      RequestedInfo) {
+    BSONObj finalKey = BSONObj::stripFieldNames(key);
+    KeyString::Builder keyString(KeyString::Version::V1, finalKey, _order);
+    auto ksEntry = seekExact(keyString.getValueCopy());
+    if (ksEntry) {
+        const BSONObj bson = KeyString::toBson(ksEntry->keyString.getBuffer(),
+                                               ksEntry->keyString.getSize(),
+                                               _order,
+                                               ksEntry->keyString.getTypeBits());
+        auto kv = seekAfterProcessing(bson, true);
+        if (kv) {
+            return kv;
+        }
+    }
+    return {};
+}
+
+boost::optional<KeyStringEntry> SortedDataInterface::Cursor::seekExact(
+    const KeyString::Value& keyStringValue) {
+    auto ksEntry = seek(keyStringValue, true);
+    if (!ksEntry) {
+        return {};
+    }
+    if (KeyString::compare(ksEntry->keyString.getBuffer(),
+                           keyStringValue.getBuffer(),
+                           KeyString::sizeWithoutRecordIdAtEnd(ksEntry->keyString.getBuffer(),
+                                                               ksEntry->keyString.getSize()),
+                           keyStringValue.getSize()) == 0) {
+        return KeyStringEntry(ksEntry->keyString, ksEntry->loc);
+    }
+    return {};
 }
 
 void SortedDataInterface::Cursor::save() {
diff --git a/src/mongo/db/storage/biggie/biggie_sorted_impl.h b/src/mongo/db/storage/biggie/biggie_sorted_impl.h
index d5183bf530c..466c2a1975b 100644
--- a/src/mongo/db/storage/biggie/biggie_sorted_impl.h
+++ b/src/mongo/db/storage/biggie/biggie_sorted_impl.h
@@ -134,6 +134,12 @@ public:
                                                     RequestedInfo parts = kKeyAndLoc) override;
         virtual boost::optional<IndexKeyEntry> seek(const IndexSeekPoint& seekPoint,
                                                     RequestedInfo parts = kKeyAndLoc) override;
+        virtual boost::optional<KeyStringEntry> seek(const KeyString::Value& keyStringValue,
+                                                     bool inclusive) override;
+        virtual boost::optional<IndexKeyEntry> seekExact(const BSONObj& key,
+                                                         RequestedInfo parts = kKeyAndLoc) override;
+        virtual boost::optional<KeyStringEntry> seekExact(
+            const KeyString::Value& keyStringValue) override;
         virtual void save() override;
         virtual void restore() override;
         virtual void detachFromOperationContext() override;
@@ -146,6 +152,8 @@ public:
         bool checkCursorValid();
         // This is a helper function for seek.
         boost::optional<IndexKeyEntry> seekAfterProcessing(BSONObj finalKey, bool inclusive);
+        boost::optional<KeyStringEntry> seekAfterProcessing(const KeyString::Value& keyString,
+                                                            bool inclusive);
         OperationContext* _opCtx;
         // This is the "working copy" of the master "branch" in the git analogy.
         StringStore* _workingCopy;
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 73de310adf0..7a0396998b8 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
@@ -253,11 +253,16 @@ public:
 
     class Cursor final : public SortedDataInterface::Cursor {
     public:
-        Cursor(OperationContext* opCtx, const IndexSet& data, bool isForward, bool isUnique)
+        Cursor(OperationContext* opCtx,
+               const IndexSet& data,
+               bool isForward,
+               bool isUnique,
+               const Ordering ordering)
             : _opCtx(opCtx),
               _data(data),
               _forward(isForward),
               _isUnique(isUnique),
+              _ordering(ordering),
               _it(data.end()) {}
 
         boost::optional<IndexKeyEntry> next(RequestedInfo parts) override {
@@ -291,7 +296,7 @@ public:
 
         boost::optional<IndexKeyEntry> seek(const BSONObj& key,
                                             bool inclusive,
-                                            RequestedInfo parts) override {
+                                            RequestedInfo) override {
             if (key.isEmpty()) {
                 _it = inclusive ? _data.begin() : _data.end();
                 _isEOF = (_it == _data.end());
@@ -312,7 +317,7 @@ public:
         }
 
         boost::optional<IndexKeyEntry> seek(const IndexSeekPoint& seekPoint,
-                                            RequestedInfo parts) override {
+                                            RequestedInfo) override {
             // Query encodes exclusive case so it can be treated as an inclusive query.
             const BSONObj query = IndexEntryComparison::makeQueryObject(seekPoint, _forward);
             locate(query, _forward ? RecordId::min() : RecordId::max());
@@ -323,6 +328,45 @@ public:
             return *_it;
         }
 
+        boost::optional<KeyStringEntry> seek(const KeyString::Value& keyStringValue,
+                                             bool inclusive) override {
+            const BSONObj query = KeyString::toBson(keyStringValue.getBuffer(),
+                                                    keyStringValue.getSize(),
+                                                    _ordering,
+                                                    keyStringValue.getTypeBits());
+            auto kv = seek(query, inclusive, kKeyAndLoc);
+            if (kv) {
+                KeyString::Builder ks(KeyString::Version::V1, kv->key, _ordering);
+                ks.appendRecordId(kv->loc);
+                return KeyStringEntry(ks.getValueCopy(), kv->loc);
+            } else {
+                return {};
+            }
+        }
+
+        boost::optional<IndexKeyEntry> seekExact(const BSONObj& key, RequestedInfo) {
+            auto kv = seek(key, true, kKeyAndLoc);
+            if (kv && kv->key.woCompare(key, BSONObj(), /*considerFieldNames*/ false) == 0)
+                return kv;
+            return {};
+        }
+
+        boost::optional<KeyStringEntry> seekExact(const KeyString::Value& keyStringValue) {
+            const BSONObj query = KeyString::toBson(keyStringValue.getBuffer(),
+                                                    keyStringValue.getSize(),
+                                                    _ordering,
+                                                    keyStringValue.getTypeBits());
+            auto kv = seekExact(query, kKeyAndLoc);
+            if (kv) {
+                // We have retrived a valid result from seekExact(). Convert to KeyString
+                // and return
+                KeyString::Builder ks(KeyString::Version::V1, kv->key, _ordering);
+                ks.appendRecordId(kv->loc);
+                return KeyStringEntry(ks.getValueCopy(), kv->loc);
+            }
+            return {};
+        }
+
         void save() override {
             // Keep original position if we haven't moved since the last restore.
             _opCtx = nullptr;
@@ -478,6 +522,7 @@ public:
         const IndexSet& _data;
         const bool _forward;
         const bool _isUnique;
+        const Ordering _ordering;
         bool _isEOF = true;
         IndexSet::const_iterator _it;
 
@@ -502,7 +547,7 @@ public:
 
     virtual std::unique_ptr<SortedDataInterface::Cursor> newCursor(OperationContext* opCtx,
                                                                    bool isForward) const {
-        return std::make_unique<Cursor>(opCtx, *_data, isForward, _isUnique);
+        return std::make_unique<Cursor>(opCtx, *_data, isForward, _isUnique, _ordering);
     }
 
     virtual Status initAsEmpty(OperationContext* opCtx) {
diff --git a/src/mongo/db/storage/index_entry_comparison.h b/src/mongo/db/storage/index_entry_comparison.h
index 41ff05f3dd3..f1c011746f9 100644
--- a/src/mongo/db/storage/index_entry_comparison.h
+++ b/src/mongo/db/storage/index_entry_comparison.h
@@ -37,6 +37,7 @@
 #include "mongo/db/jsobj.h"
 #include "mongo/db/namespace_string.h"
 #include "mongo/db/record_id.h"
+#include "mongo/db/storage/key_string.h"
 
 namespace mongo {
 
@@ -81,6 +82,18 @@ inline bool operator!=(const IndexKeyEntry& lhs, const IndexKeyEntry& rhs) {
 }
 
 /**
+ * Represents KeyString struct containing a KeyString::Value and its RecordId
+ */
+struct KeyStringEntry {
+    KeyStringEntry(KeyString::Value ks, RecordId loc) : keyString(ks), loc(loc) {
+        invariant(loc == KeyString::decodeRecordIdAtEnd(ks.getBuffer(), ks.getSize()));
+    }
+
+    const KeyString::Value keyString;
+    const RecordId loc;
+};
+
+/**
  * Describes a query that can be compared against an IndexKeyEntry in a way that allows
  * expressing exclusiveness on a prefix of the key. This is mostly used to express a location to
  * seek to in an index that may not be representable as a valid key.
diff --git a/src/mongo/db/storage/key_string.h b/src/mongo/db/storage/key_string.h
index 48529162659..4707c6671d1 100644
--- a/src/mongo/db/storage/key_string.h
+++ b/src/mongo/db/storage/key_string.h
@@ -437,8 +437,6 @@ public:
      */
     Value getValueCopy() {
         _doneAppending();
-        invariant(_state == BuildState::kEndAdded || _state == BuildState::kAppendedRecordID ||
-                  _state == BuildState::kAppendedTypeBits);
         BufBuilder newBuf(_buffer.len());
         newBuf.appendBuf(_buffer.buf(), _buffer.len());
         return {version, _typeBits, static_cast<size_t>(newBuf.len()), newBuf.release()};
@@ -682,6 +680,10 @@ inline typename std::enable_if<isKeyString<T>::value, std::ostream&>::type opera
     return stream << value.toString();
 }
 
+/**
+ * Given a KeyString which may or may not have a RecordId, returns the length of the section without
+ * the RecordId. More expensive than sizeWithoutRecordIdAtEnd
+ */
 size_t getKeySize(const char* buffer, size_t len, Ordering ord, const TypeBits& typeBits);
 
 /**
diff --git a/src/mongo/db/storage/mobile/mobile_index.cpp b/src/mongo/db/storage/mobile/mobile_index.cpp
index 96dd6b690e0..b98548c96e0 100644
--- a/src/mongo/db/storage/mobile/mobile_index.cpp
+++ b/src/mongo/db/storage/mobile/mobile_index.cpp
@@ -391,6 +391,7 @@ public:
           _opCtx(opCtx),
           _isForward(isForward),
           _savedKey(index.getKeyStringVersion()),
+          _savedKeyWithoutDiscriminator(index.getKeyStringVersion()),
           _savedRecId(0),
           _savedTypeBits(index.getKeyStringVersion()),
           _startPosition(index.getKeyStringVersion()) {
@@ -439,17 +440,15 @@ public:
     boost::optional<IndexKeyEntry> seek(const BSONObj& key,
                                         bool inclusive,
                                         RequestedInfo parts) override {
-        const BSONObj startKey = BSONObj::stripFieldNames(key);
-        // By using a discriminator other than kInclusive, there is no need to distinguish
-        // unique vs non-unique key formats since both start with the key.
         const auto discriminator = _isForward == inclusive
             ? KeyString::Discriminator::kExclusiveBefore
             : KeyString::Discriminator::kExclusiveAfter;
-        _startPosition.resetToKey(startKey, _index.getOrdering(), discriminator);
-
-        _doSeek();
-        _updatePosition();
-
+        _startPosition.resetToKey(
+            BSONObj::stripFieldNames(key), _index.getOrdering(), discriminator);
+        seek(_startPosition.getValueCopy(), inclusive /* unused by implementation */);
+        if (_isEOF) {
+            return {};
+        }
         return getCurrentEntry(parts);
     }
 
@@ -461,10 +460,83 @@ public:
                                               : KeyString::Discriminator::kExclusiveAfter;
         _startPosition.resetToKey(startKey, _index.getOrdering(), discriminator);
 
+        seek(_startPosition.getValueCopy(), true /* unused by implementation */);
+        if (_isEOF) {
+            return {};
+        }
+        return getCurrentEntry(parts);
+    }
+
+    boost::optional<KeyStringEntry> seek(const KeyString::Value& keyStringValue, bool) override {
+        _startPosition.resetFromBuffer(keyStringValue.getBuffer(), keyStringValue.getSize());
         _doSeek();
         _updatePosition();
+        if (_isEOF) {
+            return {};
+        }
+        auto sizeWithoutRecordId = KeyString::getKeySize(_savedKey.getBuffer(),
+                                                         _savedKey.getSize(),
+                                                         _index.getOrdering(),
+                                                         _savedKey.getTypeBits());
+        if (_savedKey.getSize() == sizeWithoutRecordId) {
+            _savedKey.appendRecordId(_savedRecId);
+        }
+        return KeyStringEntry(_savedKey.getValueCopy(), _savedRecId);
+    }
 
-        return getCurrentEntry(parts);
+    boost::optional<IndexKeyEntry> seekExact(const BSONObj& key, RequestedInfo) override {
+        // Create a separate KeyString that doesn't have a discriminator so that we can
+        // do a comparison with the retrieved KeyString.
+        if (!_isForward) {
+            _savedKeyWithoutDiscriminator.resetToKey(BSONObj::stripFieldNames(key),
+                                                     _index.getOrdering(),
+                                                     KeyString::Discriminator::kInclusive);
+        }
+        // If it's a reverse cursor, a kExclusiveAfter discriminator is included to ensure that
+        // the KeyString we construct will always be greater than the KeyString that we retrieve
+        // (as it might have a RecordId). So if our cursor lands on the exact match,
+        // it does not advance to the next key (in the reverse direction)
+        const auto discriminator = _isForward ? KeyString::Discriminator::kInclusive
+                                              : KeyString::Discriminator::kExclusiveAfter;
+        KeyString::Builder keyString(_index.getKeyStringVersion(),
+                                     BSONObj::stripFieldNames(key),
+                                     _index.getOrdering(),
+                                     discriminator);
+        auto ksEntry = seekExact(keyString.getValueCopy());
+        if (ksEntry) {
+            auto kv = getCurrentEntry(kKeyAndLoc);
+            invariant(kv);
+            return kv;
+        }
+        return {};
+    }
+
+    boost::optional<KeyStringEntry> seekExact(const KeyString::Value& keyStringValue) override {
+        auto ksEntry = seek(keyStringValue, true);
+        if (!ksEntry) {
+            return {};
+        }
+        // If it's a reverse cursor, we compare the KeyString we retrieved with the KeyString that
+        // doesn't have a discriminator.
+        if (!_isForward) {
+            if (KeyString::compare(
+                    ksEntry->keyString.getBuffer(),
+                    _savedKeyWithoutDiscriminator.getBuffer(),
+                    KeyString::sizeWithoutRecordIdAtEnd(ksEntry->keyString.getBuffer(),
+                                                        ksEntry->keyString.getSize()),
+                    _savedKeyWithoutDiscriminator.getSize()) == 0) {
+                return KeyStringEntry(ksEntry->keyString, ksEntry->loc);
+            }
+            return {};
+        }
+        if (KeyString::compare(ksEntry->keyString.getBuffer(),
+                               keyStringValue.getBuffer(),
+                               KeyString::sizeWithoutRecordIdAtEnd(ksEntry->keyString.getBuffer(),
+                                                                   ksEntry->keyString.getSize()),
+                               keyStringValue.getSize()) == 0) {
+            return KeyStringEntry(ksEntry->keyString, ksEntry->loc);
+        }
+        return {};
     }
 
     // All work is done in restore().
@@ -611,6 +683,7 @@ protected:
     bool _isEOF = true;
 
     KeyString::Builder _savedKey;
+    KeyString::Builder _savedKeyWithoutDiscriminator;
     RecordId _savedRecId;
     KeyString::TypeBits _savedTypeBits;
 
diff --git a/src/mongo/db/storage/mobile/mobile_index.h b/src/mongo/db/storage/mobile/mobile_index.h
index 8855f0c7e1f..6ee6e9816d8 100644
--- a/src/mongo/db/storage/mobile/mobile_index.h
+++ b/src/mongo/db/storage/mobile/mobile_index.h
@@ -108,7 +108,7 @@ public:
                     const ValueType& value,
                     bool isTransactional = true);
 
-    bool isUnique() {
+    bool isUnique() const {
         return _isUnique;
     }
 
diff --git a/src/mongo/db/storage/sorted_data_interface.h b/src/mongo/db/storage/sorted_data_interface.h
index 1f0408d688a..ceed0466b3c 100644
--- a/src/mongo/db/storage/sorted_data_interface.h
+++ b/src/mongo/db/storage/sorted_data_interface.h
@@ -320,18 +320,17 @@ public:
         virtual boost::optional<IndexKeyEntry> seek(const IndexSeekPoint& seekPoint,
                                                     RequestedInfo parts = kKeyAndLoc) = 0;
 
+        virtual boost::optional<KeyStringEntry> seek(const KeyString::Value& keyString,
+                                                     bool inclusive) = 0;
         /**
          * Seeks to a key with a hint to the implementation that you only want exact matches. If
          * an exact match can't be found, boost::none will be returned and the resulting
          * position of the cursor is unspecified.
          */
         virtual boost::optional<IndexKeyEntry> seekExact(const BSONObj& key,
-                                                         RequestedInfo parts = kKeyAndLoc) {
-            auto kv = seek(key, true, kKeyAndLoc);
-            if (kv && kv->key.woCompare(key, BSONObj(), /*considerFieldNames*/ false) == 0)
-                return kv;
-            return {};
-        }
+                                                         RequestedInfo parts = kKeyAndLoc) = 0;
+
+        virtual boost::optional<KeyStringEntry> seekExact(const KeyString::Value& keyString) = 0;
 
         //
         // Saving and restoring state
diff --git a/src/mongo/db/storage/sorted_data_interface_test_insert.cpp b/src/mongo/db/storage/sorted_data_interface_test_insert.cpp
index 5b5ac187b55..efd183aef93 100644
--- a/src/mongo/db/storage/sorted_data_interface_test_insert.cpp
+++ b/src/mongo/db/storage/sorted_data_interface_test_insert.cpp
@@ -489,6 +489,100 @@ TEST(SortedDataInterface, InsertMultipleKeyStrings) {
     }
 }
 
+/*
+ * Insert multiple KeyStrings and seek to the inserted KeyStrings
+ */
+TEST(SortedDataInterface, InsertAndSeekKeyString) {
+    const auto harnessHelper(newSortedDataInterfaceHarnessHelper());
+    const std::unique_ptr<SortedDataInterface> sorted(
+        harnessHelper->newSortedDataInterface(/*unique=*/true, /*partial=*/false));
+
+    KeyString::Builder keyString1(sorted->getKeyStringVersion(), key1, sorted->getOrdering(), loc1);
+    KeyString::Builder keyString2(sorted->getKeyStringVersion(), key2, sorted->getOrdering(), loc2);
+
+    KeyString::Builder keyString1WithoutRecordId(
+        sorted->getKeyStringVersion(), key1, sorted->getOrdering());
+    KeyString::Builder keyString2WithoutRecordId(
+        sorted->getKeyStringVersion(), key2, sorted->getOrdering());
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT(sorted->isEmpty(opCtx.get()));
+    }
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        {
+            WriteUnitOfWork uow(opCtx.get());
+            ASSERT_OK(sorted->insert(opCtx.get(), keyString1.getValueCopy(), loc1, false));
+            ASSERT_OK(sorted->insert(opCtx.get(), keyString2.getValueCopy(), loc2, false));
+            uow.commit();
+        }
+    }
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT_EQUALS(2, sorted->numEntries(opCtx.get()));
+
+        const std::unique_ptr<SortedDataInterface::Cursor> cursor(sorted->newCursor(opCtx.get()));
+
+        auto ksEntry1 = cursor->seek(keyString1WithoutRecordId.getValueCopy(), true);
+        ASSERT_EQUALS(ksEntry1->keyString.compare(keyString1.getValueCopy()), 0);
+        ASSERT_EQUALS(ksEntry1->keyString.compare(keyString2.getValueCopy()), -1);
+
+        auto ksEntry2 = cursor->seek(keyString2WithoutRecordId.getValueCopy(), true);
+        ASSERT_EQUALS(ksEntry2->keyString.compare(keyString2.getValueCopy()), 0);
+        ASSERT_EQUALS(ksEntry2->keyString.compare(keyString1.getValueCopy()), 1);
+    }
+}
+
+/*
+ * Insert multiple KeyStrings and seekExact to the inserted KeyStrings
+ */
+TEST(SortedDataInterface, InsertAndSeekExactKeyString) {
+    const auto harnessHelper(newSortedDataInterfaceHarnessHelper());
+    const std::unique_ptr<SortedDataInterface> sorted(
+        harnessHelper->newSortedDataInterface(/*unique=*/true, /*partial=*/false));
+
+    KeyString::Builder keyString1(sorted->getKeyStringVersion(), key1, sorted->getOrdering(), loc1);
+    KeyString::Builder keyString2(sorted->getKeyStringVersion(), key2, sorted->getOrdering(), loc2);
+
+    KeyString::Builder keyString1WithoutRecordId(
+        sorted->getKeyStringVersion(), key1, sorted->getOrdering());
+    KeyString::Builder keyString2WithoutRecordId(
+        sorted->getKeyStringVersion(), key2, sorted->getOrdering());
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT(sorted->isEmpty(opCtx.get()));
+    }
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        {
+            WriteUnitOfWork uow(opCtx.get());
+            ASSERT_OK(sorted->insert(opCtx.get(), keyString1.getValueCopy(), loc1, false));
+            ASSERT_OK(sorted->insert(opCtx.get(), keyString2.getValueCopy(), loc2, false));
+            uow.commit();
+        }
+    }
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT_EQUALS(2, sorted->numEntries(opCtx.get()));
+
+        const std::unique_ptr<SortedDataInterface::Cursor> cursor(sorted->newCursor(opCtx.get()));
+
+        auto ksEntry1 = cursor->seekExact(keyString1WithoutRecordId.getValueCopy());
+        ASSERT_EQUALS(ksEntry1->keyString.compare(keyString1.getValueCopy()), 0);
+        ASSERT_EQUALS(ksEntry1->keyString.compare(keyString2.getValueCopy()), -1);
+
+        auto ksEntry2 = cursor->seekExact(keyString2WithoutRecordId.getValueCopy());
+        ASSERT_EQUALS(ksEntry2->keyString.compare(keyString2.getValueCopy()), 0);
+        ASSERT_EQUALS(ksEntry2->keyString.compare(keyString1.getValueCopy()), 1);
+    }
+}
+
 // Insert multiple compound keys and verify that the number of entries
 // in the index equals the number that were inserted.
 TEST(SortedDataInterface, InsertMultipleCompoundKeys) {
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp
index 94cbf637f95..3a98224068d 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp
@@ -824,6 +824,7 @@ public:
           _key(idx.getKeyStringVersion()),
           _typeBits(idx.getKeyStringVersion()),
           _query(idx.getKeyStringVersion()),
+          _queryWithoutDiscriminator(idx.getKeyStringVersion()),
           _prefix(prefix) {
         _cursor.emplace(_idx.uri(), _idx.tableId(), false, _opCtx);
     }
@@ -864,12 +865,8 @@ public:
         const auto discriminator = _forward == inclusive
             ? KeyString::Discriminator::kExclusiveBefore
             : KeyString::Discriminator::kExclusiveAfter;
-
-        // By using a discriminator other than kInclusive, there is no need to distinguish
-        // unique vs non-unique key formats since both start with the key.
         _query.resetToKey(finalKey, _idx.getOrdering(), discriminator);
-        seekWTCursor(_query);
-        updatePosition();
+        seek(_query.getValueCopy(), inclusive /* unused by implementation */);
         return curr(parts);
     }
 
@@ -883,11 +880,82 @@ public:
         const auto discriminator = _forward ? KeyString::Discriminator::kExclusiveBefore
                                             : KeyString::Discriminator::kExclusiveAfter;
         _query.resetToKey(key, _idx.getOrdering(), discriminator);
-        seekWTCursor(_query);
-        updatePosition();
+        seek(_query.getValueCopy(), true /* unused by implementation */);
         return curr(parts);
     }
 
+    boost::optional<KeyStringEntry> seek(const KeyString::Value& keyStringValue, bool) override {
+        dassert(_opCtx->lockState()->isReadLocked());
+        seekWTCursor(keyStringValue);
+
+        updatePosition();
+        if (_eof)
+            return {};
+
+        dassert(!atOrPastEndPointAfterSeeking());
+        dassert(!_id.isNull());
+
+        auto sizeWithoutRecordId = KeyString::getKeySize(
+            _key.getBuffer(), _key.getSize(), _idx.getOrdering(), _key.getTypeBits());
+        if (_key.getSize() == sizeWithoutRecordId) {
+            _key.appendRecordId(_id);
+        }
+        return KeyStringEntry(_key.getValueCopy(), _id);
+    }
+
+    boost::optional<IndexKeyEntry> seekExact(const BSONObj& key, RequestedInfo) override {
+        // Create a separate KeyString that doesn't have a discriminator so that
+        // we can do a comparison with the retrieved KeyString.
+        if (!_forward) {
+            _queryWithoutDiscriminator.resetToKey(BSONObj::stripFieldNames(key),
+                                                  _idx.getOrdering(),
+                                                  KeyString::Discriminator::kInclusive);
+        }
+        // If it's a reverse cursor, a kExclusiveAfter discriminator is included to ensure that
+        // the KeyString we construct will always be greater than the KeyString that we retrieve
+        // (as it might have a RecordId). So if our cursor lands on the exact match,
+        // it does not advance to the next key (in the reverse direction)
+        const auto discriminator = _forward ? KeyString::Discriminator::kInclusive
+                                            : KeyString::Discriminator::kExclusiveAfter;
+        _query.resetToKey(BSONObj::stripFieldNames(key), _idx.getOrdering(), discriminator);
+        auto ksEntry = seekExact(_query.getValueCopy());
+        if (ksEntry) {
+            auto kv = curr(kKeyAndLoc);
+            invariant(kv);
+            return kv;
+        }
+        return {};
+    }
+
+    boost::optional<KeyStringEntry> seekExact(const KeyString::Value& keyStringValue) override {
+        auto ksEntry = seek(keyStringValue, true);
+        if (!ksEntry) {
+            return {};
+        }
+
+        // If it's a reverse cursor, we compare the KeyString we retrieved with the KeyString that
+        // doesn't have a discriminator.
+        if (!_forward) {
+            if (KeyString::compare(
+                    ksEntry->keyString.getBuffer(),
+                    _queryWithoutDiscriminator.getBuffer(),
+                    KeyString::sizeWithoutRecordIdAtEnd(ksEntry->keyString.getBuffer(),
+                                                        ksEntry->keyString.getSize()),
+                    _queryWithoutDiscriminator.getSize()) == 0) {
+                return KeyStringEntry(ksEntry->keyString, ksEntry->loc);
+            }
+            return {};
+        }
+        if (KeyString::compare(ksEntry->keyString.getBuffer(),
+                               keyStringValue.getBuffer(),
+                               KeyString::sizeWithoutRecordIdAtEnd(ksEntry->keyString.getBuffer(),
+                                                                   ksEntry->keyString.getSize()),
+                               keyStringValue.getSize()) == 0) {
+            return KeyStringEntry(ksEntry->keyString, ksEntry->loc);
+        }
+        return {};
+    }
+
     void save() override {
         try {
             if (_cursor)
@@ -921,7 +989,7 @@ public:
             //
             // Unique indexes can have both kinds of KeyStrings, ie with or without the record id.
             // Restore for unique indexes gets handled separately in it's own implementation.
-            _lastMoveSkippedKey = !seekWTCursor(_key);
+            _lastMoveSkippedKey = !seekWTCursor(_key.getValueCopy());
             TRACE_CURSOR << "restore _lastMoveSkippedKey:" << _lastMoveSkippedKey;
         }
     }
@@ -1041,7 +1109,7 @@ protected:
     }
 
     // Seeks to query. Returns true on exact match.
-    bool seekWTCursor(const KeyString::Builder& query) {
+    bool seekWTCursor(const KeyString::Value& query) {
         WT_CURSOR* c = _cursor->get();
 
         int cmp = -1;
@@ -1151,6 +1219,7 @@ protected:
     bool _lastMoveSkippedKey = false;
 
     KeyString::Builder _query;
+    KeyString::Builder _queryWithoutDiscriminator;
     KVPrefix _prefix;
 
     std::unique_ptr<KeyString::Builder> _endPosition;