SERVER-42858 Index cursor seek should not append RecordID to saved KeyString

9 files changed, 191 insertions, 2 deletions

diff --git a/src/mongo/db/storage/biggie/biggie_sorted_impl_test.cpp b/src/mongo/db/storage/biggie/biggie_sorted_impl_test.cpp
index 18e6da89d58..1a518225caa 100644
--- a/src/mongo/db/storage/biggie/biggie_sorted_impl_test.cpp
+++ b/src/mongo/db/storage/biggie/biggie_sorted_impl_test.cpp
@@ -52,6 +52,23 @@ private:
 
 public:
     SortedDataInterfaceTestHarnessHelper() : _order(Ordering::make(BSONObj())) {}
+
+    std::unique_ptr<mongo::SortedDataInterface> newIdIndexSortedDataInterface() final {
+        std::string ns = "test.biggie";
+        OperationContextNoop opCtx(newRecoveryUnit().release());
+
+        BSONObj spec = BSON("key" << BSON("_id" << 1) << "name"
+                                  << "_id_"
+                                  << "v" << static_cast<int>(IndexDescriptor::kLatestIndexVersion)
+                                  << "unique" << true);
+
+        auto collection = std::make_unique<CollectionMock>(NamespaceString(ns));
+        IndexDescriptor desc(collection.get(), "", spec);
+        invariant(desc.isIdIndex());
+
+        return std::make_unique<SortedDataInterface>(&opCtx, "ident"_sd, &desc);
+    }
+
     std::unique_ptr<mongo::SortedDataInterface> newSortedDataInterface(bool unique,
                                                                        bool partial) final {
         std::string ns = "test.biggie";
diff --git a/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_btree_impl_test.cpp b/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_btree_impl_test.cpp
index a2c4cfcd857..eb6c06367c9 100644
--- a/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_btree_impl_test.cpp
+++ b/src/mongo/db/storage/ephemeral_for_test/ephemeral_for_test_btree_impl_test.cpp
@@ -32,6 +32,7 @@
 #include <memory>
 
 #include "mongo/base/init.h"
+#include "mongo/db/index/index_descriptor.h"
 #include "mongo/db/storage/ephemeral_for_test/ephemeral_for_test_recovery_unit.h"
 #include "mongo/db/storage/sorted_data_interface_test_harness.h"
 #include "mongo/unittest/unittest.h"
@@ -44,6 +45,21 @@ class EphemeralForBtreeImplTestHarnessHelper final
 public:
     EphemeralForBtreeImplTestHarnessHelper() : _order(Ordering::make(BSONObj())) {}
 
+    std::unique_ptr<SortedDataInterface> newIdIndexSortedDataInterface() final {
+        BSONObj spec = BSON("key" << BSON("_id" << 1) << "name"
+                                  << "_id_"
+                                  << "v" << static_cast<int>(IndexDescriptor::kLatestIndexVersion)
+                                  << "unique" << true);
+
+        return std::unique_ptr<SortedDataInterface>(
+            getEphemeralForTestBtreeImpl(_order,
+                                         true /* unique */,
+                                         NamespaceString("test.EphemeralForTest"),
+                                         "indexName",
+                                         spec,
+                                         &_data));
+    }
+
     std::unique_ptr<SortedDataInterface> newSortedDataInterface(bool unique, bool partial) final {
         return std::unique_ptr<SortedDataInterface>(
             getEphemeralForTestBtreeImpl(_order,
diff --git a/src/mongo/db/storage/mobile/mobile_index.cpp b/src/mongo/db/storage/mobile/mobile_index.cpp
index 1f7172c5286..f66050f6b72 100644
--- a/src/mongo/db/storage/mobile/mobile_index.cpp
+++ b/src/mongo/db/storage/mobile/mobile_index.cpp
@@ -486,7 +486,11 @@ public:
                                                          _index.getOrdering(),
                                                          _savedKey.getTypeBits());
         if (_savedKey.getSize() == sizeWithoutRecordId) {
-            _savedKey.appendRecordId(_savedRecId);
+            // 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);
     }
diff --git a/src/mongo/db/storage/mobile/mobile_index_test.cpp b/src/mongo/db/storage/mobile/mobile_index_test.cpp
index f11052b3eb1..ca6811a657a 100644
--- a/src/mongo/db/storage/mobile/mobile_index_test.cpp
+++ b/src/mongo/db/storage/mobile/mobile_index_test.cpp
@@ -54,6 +54,10 @@ public:
         _sessionPool.reset(new MobileSessionPool(_fullPath));
     }
 
+    std::unique_ptr<SortedDataInterface> newIdIndexSortedDataInterface() final {
+        return newSortedDataInterface(true /* unique */, false /* partial */);
+    }
+
     std::unique_ptr<SortedDataInterface> newSortedDataInterface(bool isUnique, bool isPartial) {
         std::string ident("index_" + std::to_string(inc++));
         OperationContextNoop opCtx(newRecoveryUnit().release());
diff --git a/src/mongo/db/storage/sorted_data_interface_test_cursor_saverestore.cpp b/src/mongo/db/storage/sorted_data_interface_test_cursor_saverestore.cpp
index 30d207d5031..aa2eb8231ca 100644
--- a/src/mongo/db/storage/sorted_data_interface_test_cursor_saverestore.cpp
+++ b/src/mongo/db/storage/sorted_data_interface_test_cursor_saverestore.cpp
@@ -132,6 +132,98 @@ TEST(SortedDataInterface, SaveAndRestorePositionWhileIterateCursorReversed) {
     }
 }
 
+// Insert multiple keys on the _id index and try to iterate through all of them using a forward
+// cursor while calling save() and restore() in succession.
+TEST(SortedDataInterface, SaveAndRestorePositionWhileIterateCursorOnIdIndex) {
+    const auto harnessHelper(newSortedDataInterfaceHarnessHelper());
+    const std::unique_ptr<SortedDataInterface> sorted(
+        harnessHelper->newIdIndexSortedDataInterface());
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT(sorted->isEmpty(opCtx.get()));
+    }
+
+    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);
+            ASSERT_OK(sorted->insert(opCtx.get(), key, loc, 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));
+        int i = nToInsert - 1;
+        for (auto entry = cursor->seek(kMaxBSONKey, true); entry; i--, entry = cursor->next()) {
+            ASSERT_GTE(i, 0);
+            ASSERT_EQ(entry, IndexKeyEntry(BSON("" << i), RecordId(42, i * 2)));
+
+            cursor->save();
+            cursor->restore();
+        }
+        ASSERT(!cursor->next());
+        ASSERT_EQ(i, -1);
+    }
+}
+
+// Insert multiple keys on the _id index and try to iterate through all of them using a reverse
+// cursor while calling save() and restore() in succession.
+TEST(SortedDataInterface, SaveAndRestorePositionWhileIterateCursorReversedOnIdIndex) {
+    const auto harnessHelper(newSortedDataInterfaceHarnessHelper());
+    const std::unique_ptr<SortedDataInterface> sorted(
+        harnessHelper->newIdIndexSortedDataInterface());
+
+    {
+        const ServiceContext::UniqueOperationContext opCtx(harnessHelper->newOperationContext());
+        ASSERT(sorted->isEmpty(opCtx.get()));
+    }
+
+    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);
+            ASSERT_OK(sorted->insert(opCtx.get(), key, loc, 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));
+        int i = nToInsert - 1;
+        for (auto entry = cursor->seek(kMaxBSONKey, true); entry; i--, entry = cursor->next()) {
+            ASSERT_GTE(i, 0);
+            ASSERT_EQ(entry, IndexKeyEntry(BSON("" << i), RecordId(42, i * 2)));
+
+            cursor->save();
+            cursor->restore();
+        }
+        ASSERT(!cursor->next());
+        ASSERT_EQ(i, -1);
+    }
+}
+
 // Insert the same key multiple times and try to iterate through each
 // occurrence using a forward cursor while calling savePosition() and
 // restorePosition() in succession. Verify that the RecordId is saved
diff --git a/src/mongo/db/storage/sorted_data_interface_test_harness.h b/src/mongo/db/storage/sorted_data_interface_test_harness.h
index 0df65e0a220..b907e926101 100644
--- a/src/mongo/db/storage/sorted_data_interface_test_harness.h
+++ b/src/mongo/db/storage/sorted_data_interface_test_harness.h
@@ -91,6 +91,7 @@ public:
     virtual std::unique_ptr<SortedDataInterface> newSortedDataInterface(bool unique,
                                                                         bool partial) = 0;
 
+    virtual std::unique_ptr<SortedDataInterface> newIdIndexSortedDataInterface() = 0;
     /**
      * Creates a new SDI with some initial data.
      *
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp
index f55f34ff06e..af7509fe5c7 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_index.cpp
@@ -902,10 +902,17 @@ 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) {
-            _key.appendRecordId(_id);
+            // 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);
     }
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_prefixed_index_test.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_prefixed_index_test.cpp
index b86e2d60a85..dda3f3dd1ea 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_prefixed_index_test.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_prefixed_index_test.cpp
@@ -70,7 +70,31 @@ public:
         _conn->close(_conn, nullptr);
     }
 
+    std::unique_ptr<SortedDataInterface> newIdIndexSortedDataInterface() final {
+        std::string ns = "test.wt";
+        OperationContextNoop opCtx(newRecoveryUnit().release());
+
+        BSONObj spec = BSON("key" << BSON("_id" << 1) << "name"
+                                  << "_id_"
+                                  << "v" << static_cast<int>(IndexDescriptor::kLatestIndexVersion)
+                                  << "unique" << true);
+        auto collection = std::make_unique<CollectionMock>(NamespaceString(ns));
+        IndexDescriptor desc(collection.get(), "", spec);
+        invariant(desc.isIdIndex());
+
+        KVPrefix prefix = KVPrefix::generateNextPrefix();
+        StatusWith<std::string> result = WiredTigerIndex::generateCreateString(
+            kWiredTigerEngineName, "", "", desc, prefix.isPrefixed());
+        ASSERT_OK(result.getStatus());
+
+        string uri = "table:" + ns;
+        invariantWTOK(WiredTigerIndex::Create(&opCtx, uri, result.getValue()));
+
+        return std::make_unique<WiredTigerIndexUnique>(&opCtx, uri, &desc, prefix);
+    }
+
     std::unique_ptr<SortedDataInterface> newSortedDataInterface(bool unique, bool partial) final {
+
         std::string ns = "test.wt";
         OperationContextNoop opCtx(newRecoveryUnit().release());
 
diff --git a/src/mongo/db/storage/wiredtiger/wiredtiger_standard_index_test.cpp b/src/mongo/db/storage/wiredtiger/wiredtiger_standard_index_test.cpp
index 1f683424598..b7d1ae7e932 100644
--- a/src/mongo/db/storage/wiredtiger/wiredtiger_standard_index_test.cpp
+++ b/src/mongo/db/storage/wiredtiger/wiredtiger_standard_index_test.cpp
@@ -70,6 +70,30 @@ public:
         _conn->close(_conn, nullptr);
     }
 
+    std::unique_ptr<SortedDataInterface> newIdIndexSortedDataInterface() final {
+        std::string ns = "test.wt";
+        OperationContextNoop opCtx(newRecoveryUnit().release());
+
+        BSONObj spec = BSON("key" << BSON("_id" << 1) << "name"
+                                  << "_id_"
+                                  << "v" << static_cast<int>(IndexDescriptor::kLatestIndexVersion)
+                                  << "unique" << true);
+
+        auto collection = std::make_unique<CollectionMock>(NamespaceString(ns));
+        IndexDescriptor desc(collection.get(), "", spec);
+        invariant(desc.isIdIndex());
+
+        KVPrefix prefix = KVPrefix::kNotPrefixed;
+        StatusWith<std::string> result = WiredTigerIndex::generateCreateString(
+            kWiredTigerEngineName, "", "", desc, prefix.isPrefixed());
+        ASSERT_OK(result.getStatus());
+
+        string uri = "table:" + ns;
+        invariantWTOK(WiredTigerIndex::Create(&opCtx, uri, result.getValue()));
+
+        return std::make_unique<WiredTigerIndexUnique>(&opCtx, uri, &desc, prefix);
+    }
+
     std::unique_ptr<SortedDataInterface> newSortedDataInterface(bool unique, bool partial) final {
         std::string ns = "test.wt";
         OperationContextNoop opCtx(newRecoveryUnit().release());