SERVER-42979 Implement WorkingSet::extract() and WorkingSet::emplace().

Also replaces WorkingSetMember::isSuspicious with a scheme
that associates a snapshot id with every index key. This is
needed because extracted WorkingSetMembers are not
discoverable from the WorkingSet, and thus cannot be marked
as suspicious in preparation for yield.

11 files changed, 116 insertions, 102 deletions

diff --git a/src/mongo/db/exec/and_common.h b/src/mongo/db/exec/and_common.h
index ca063852aa6..93abdf84d96 100644
--- a/src/mongo/db/exec/and_common.h
+++ b/src/mongo/db/exec/and_common.h
@@ -93,9 +93,6 @@ public:
                 dest->keyData.push_back(src.keyData[i]);
             }
         }
-
-        if (src.isSuspicious)
-            dest->isSuspicious = true;
     }
 };
 
diff --git a/src/mongo/db/exec/delete.cpp b/src/mongo/db/exec/delete.cpp
index c31ce92b72d..702ecc68d5a 100644
--- a/src/mongo/db/exec/delete.cpp
+++ b/src/mongo/db/exec/delete.cpp
@@ -193,7 +193,6 @@ PlanStage::StageState DeleteStage::doWork(WorkingSetID* out) {
     // TODO: Do we want to buffer docs and delete them in a group rather than saving/restoring state
     // repeatedly?
 
-    WorkingSetCommon::prepareForSnapshotChange(_ws);
     try {
         child()->saveState();
     } catch (const WriteConflictException&) {
diff --git a/src/mongo/db/exec/distinct_scan.cpp b/src/mongo/db/exec/distinct_scan.cpp
index 24e45ee62d2..860e6323a40 100644
--- a/src/mongo/db/exec/distinct_scan.cpp
+++ b/src/mongo/db/exec/distinct_scan.cpp
@@ -122,7 +122,10 @@ PlanStage::StageState DistinctScan::doWork(WorkingSetID* out) {
             WorkingSetID id = _workingSet->allocate();
             WorkingSetMember* member = _workingSet->get(id);
             member->recordId = kv->loc;
-            member->keyData.push_back(IndexKeyDatum(_keyPattern, kv->key, workingSetIndexId()));
+            member->keyData.push_back(IndexKeyDatum(_keyPattern,
+                                                    kv->key,
+                                                    workingSetIndexId(),
+                                                    getOpCtx()->recoveryUnit()->getSnapshotId()));
             _workingSet->transitionToRecordIdAndIdx(id);
 
             *out = id;
diff --git a/src/mongo/db/exec/index_scan.cpp b/src/mongo/db/exec/index_scan.cpp
index 0e5d940e2f0..506ab6500e0 100644
--- a/src/mongo/db/exec/index_scan.cpp
+++ b/src/mongo/db/exec/index_scan.cpp
@@ -233,7 +233,8 @@ PlanStage::StageState IndexScan::doWork(WorkingSetID* out) {
     WorkingSetID id = _workingSet->allocate();
     WorkingSetMember* member = _workingSet->get(id);
     member->recordId = kv->loc;
-    member->keyData.push_back(IndexKeyDatum(_keyPattern, kv->key, workingSetIndexId()));
+    member->keyData.push_back(IndexKeyDatum(
+        _keyPattern, kv->key, workingSetIndexId(), getOpCtx()->recoveryUnit()->getSnapshotId()));
     _workingSet->transitionToRecordIdAndIdx(id);
 
     if (_addKeyMetadata) {
diff --git a/src/mongo/db/exec/projection_exec_test.cpp b/src/mongo/db/exec/projection_exec_test.cpp
index 684d3a07f35..f15094fd3aa 100644
--- a/src/mongo/db/exec/projection_exec_test.cpp
+++ b/src/mongo/db/exec/projection_exec_test.cpp
@@ -193,7 +193,7 @@ TEST(ProjectionExecTest, TransformCoveredDottedProjection) {
     ASSERT_EQ(boost::make_optional("{ b: { c: 2, d: 3, f: { g: 4, h: 5 } } }"s),
               project("{'b.c': 1, 'b.d': 1, 'b.f.g': 1, 'b.f.h': 1}",
                       "{}",
-                      IndexKeyDatum(keyPattern, keyData, 0)));
+                      IndexKeyDatum(keyPattern, keyData, 0, SnapshotId{})));
 }
 
 TEST(ProjectionExecTest, TransformNonCoveredDottedProjection) {
@@ -249,7 +249,7 @@ TEST(ProjectionExecTest, TransformMetaSortKeyCoveredNormal) {
     ASSERT_EQ(boost::make_optional("{ a: 5, b: { : 5 } }"s),
               project("{_id: 0, a: 1, b: {$meta: 'sortKey'}}",
                       "{}",
-                      IndexKeyDatum(BSON("a" << 1), BSON("" << 5), 0),
+                      IndexKeyDatum(BSON("a" << 1), BSON("" << 5), 0, SnapshotId{}),
                       boost::none,      // collator
                       BSON("" << 5)));  // sortKey
 }
@@ -258,7 +258,7 @@ TEST(ProjectionExecTest, TransformMetaSortKeyCoveredOverwrite) {
     ASSERT_EQ(boost::make_optional("{ a: { : 5 } }"s),
               project("{_id: 0, a: 1, a: {$meta: 'sortKey'}}",
                       "{}",
-                      IndexKeyDatum(BSON("a" << 1), BSON("" << 5), 0),
+                      IndexKeyDatum(BSON("a" << 1), BSON("" << 5), 0, SnapshotId{}),
                       boost::none,      // collator
                       BSON("" << 5)));  // sortKey
 }
@@ -267,7 +267,8 @@ TEST(ProjectionExecTest, TransformMetaSortKeyCoveredAdditionalData) {
     ASSERT_EQ(boost::make_optional("{ a: 5, c: 6, b: { : 5 } }"s),
               project("{_id: 0, a: 1, b: {$meta: 'sortKey'}, c: 1}",
                       "{}",
-                      IndexKeyDatum(BSON("a" << 1 << "c" << 1), BSON("" << 5 << "" << 6), 0),
+                      IndexKeyDatum(
+                          BSON("a" << 1 << "c" << 1), BSON("" << 5 << "" << 6), 0, SnapshotId{}),
                       boost::none,      // collator
                       BSON("" << 5)));  // sortKey
 }
@@ -276,7 +277,8 @@ TEST(ProjectionExecTest, TransformMetaSortKeyCoveredCompound) {
     ASSERT_EQ(boost::make_optional("{ a: 5, b: { : 5, : 6 } }"s),
               project("{_id: 0, a: 1, b: {$meta: 'sortKey'}}",
                       "{}",
-                      IndexKeyDatum(BSON("a" << 1 << "c" << 1), BSON("" << 5 << "" << 6), 0),
+                      IndexKeyDatum(
+                          BSON("a" << 1 << "c" << 1), BSON("" << 5 << "" << 6), 0, SnapshotId{}),
                       boost::none,                 // collator
                       BSON("" << 5 << "" << 6)));  // sortKey
 }
@@ -287,7 +289,8 @@ TEST(ProjectionExecTest, TransformMetaSortKeyCoveredCompound2) {
                       "{}",
                       IndexKeyDatum(BSON("a" << 1 << "b" << 1 << "c" << 1),
                                     BSON("" << 5 << "" << 6 << "" << 4),
-                                    0),
+                                    0,
+                                    SnapshotId{}),
                       boost::none,                 // collator
                       BSON("" << 5 << "" << 6)));  // sortKey
 }
@@ -298,7 +301,8 @@ TEST(ProjectionExecTest, TransformMetaSortKeyCoveredCompound3) {
                       "{}",
                       IndexKeyDatum(BSON("a" << 1 << "b" << 1 << "c" << 1 << "d" << 1),
                                     BSON("" << 5 << "" << 6 << "" << 4 << "" << 9000),
-                                    0),
+                                    0,
+                                    SnapshotId{}),
                       boost::none,                 // collator
                       BSON("" << 6 << "" << 4)));  // sortKey
 }
diff --git a/src/mongo/db/exec/update_stage.cpp b/src/mongo/db/exec/update_stage.cpp
index a9bddf56266..faef94ca7a2 100644
--- a/src/mongo/db/exec/update_stage.cpp
+++ b/src/mongo/db/exec/update_stage.cpp
@@ -749,8 +749,7 @@ PlanStage::StageState UpdateStage::doWork(WorkingSetID* out) {
         // is allowed to free the memory.
         member->makeObjOwnedIfNeeded();
 
-        // Save state before making changes
-        WorkingSetCommon::prepareForSnapshotChange(_ws);
+        // Save state before making changes.
         try {
             child()->saveState();
         } catch (const WriteConflictException&) {
diff --git a/src/mongo/db/exec/working_set.cpp b/src/mongo/db/exec/working_set.cpp
index 2aacba154ab..cfa6c22a748 100644
--- a/src/mongo/db/exec/working_set.cpp
+++ b/src/mongo/db/exec/working_set.cpp
@@ -76,14 +76,11 @@ void WorkingSet::clear() {
     // Since working set is now empty, the free list pointer should
     // point to nothing.
     _freeList = INVALID_ID;
-
-    _yieldSensitiveIds.clear();
 }
 
 void WorkingSet::transitionToRecordIdAndIdx(WorkingSetID id) {
     WorkingSetMember* member = get(id);
     member->_state = WorkingSetMember::RID_AND_IDX;
-    _yieldSensitiveIds.push_back(id);
 }
 
 void WorkingSet::transitionToRecordIdAndObj(WorkingSetID id) {
@@ -96,11 +93,17 @@ void WorkingSet::transitionToOwnedObj(WorkingSetID id) {
     member->transitionToOwnedObj();
 }
 
-std::vector<WorkingSetID> WorkingSet::getAndClearYieldSensitiveIds() {
-    std::vector<WorkingSetID> out;
-    // Clear '_yieldSensitiveIds' by swapping it into the set to be returned.
-    _yieldSensitiveIds.swap(out);
-    return out;
+WorkingSetMember WorkingSet::extract(WorkingSetID wsid) {
+    invariant(wsid < _data.size());
+    WorkingSetMember ret = std::move(_data[wsid].member);
+    free(wsid);
+    return ret;
+}
+
+WorkingSetID WorkingSet::emplace(WorkingSetMember&& wsm) {
+    auto wsid = allocate();
+    *get(wsid) = std::move(wsm);
+    return wsid;
 }
 
 //
@@ -210,6 +213,7 @@ void WorkingSetMember::serializeForSorter(BufBuilder& buf) const {
             indexKeyDatum.indexKeyPattern.serializeForSorter(buf);
             indexKeyDatum.keyData.serializeForSorter(buf);
             buf.appendNum(indexKeyDatum.indexId);
+            buf.appendNum(static_cast<unsigned long long>(indexKeyDatum.snapshotId.toNumber()));
         }
     }
 
@@ -244,14 +248,11 @@ WorkingSetMember WorkingSetMember::deserializeForSorter(BufReader& buf,
             auto indexKey =
                 BSONObj::deserializeForSorter(buf, BSONObj::SorterDeserializeSettings{}).getOwned();
             auto indexId = buf.read<LittleEndian<unsigned int>>();
-            wsm.keyData.push_back(
-                IndexKeyDatum{std::move(indexKeyPattern), std::move(indexKey), indexId});
+            auto snapshotIdRepr = buf.read<LittleEndian<uint64_t>>();
+            auto snapshotId = snapshotIdRepr ? SnapshotId{snapshotIdRepr} : SnapshotId{};
+            wsm.keyData.push_back(IndexKeyDatum{
+                std::move(indexKeyPattern), std::move(indexKey), indexId, snapshotId});
         }
-
-        // Mark any working set member representing an index key as suspicious on deserialization.
-        // This is needed because the member may have survived a yield while absent from the working
-        // set.
-        wsm.isSuspicious = true;
     }
 
     if (wsm.hasRecordId()) {
diff --git a/src/mongo/db/exec/working_set.h b/src/mongo/db/exec/working_set.h
index 7c5063f0785..0c60ff20de2 100644
--- a/src/mongo/db/exec/working_set.h
+++ b/src/mongo/db/exec/working_set.h
@@ -60,8 +60,9 @@ using WorkingSetRegisteredIndexId = unsigned int;
 struct IndexKeyDatum {
     IndexKeyDatum(const BSONObj& keyPattern,
                   const BSONObj& key,
-                  WorkingSetRegisteredIndexId indexId)
-        : indexKeyPattern(keyPattern), keyData(key), indexId(indexId) {}
+                  WorkingSetRegisteredIndexId indexId,
+                  SnapshotId snapshotId)
+        : indexKeyPattern(keyPattern), keyData(key), indexId(indexId), snapshotId(snapshotId) {}
 
     /**
      * getFieldDotted produces the field with the provided name based on index keyData. The return
@@ -95,6 +96,9 @@ struct IndexKeyDatum {
     // Associates this index key with an index that has been registered with the WorkingSet. Can be
     // used to recover pointers to catalog objects for this index from the WorkingSet.
     WorkingSetRegisteredIndexId indexId;
+
+    // Identifies the storage engine snapshot from which this index key was obtained.
+    SnapshotId snapshotId;
 };
 
 /**
@@ -152,10 +156,6 @@ public:
     Snapshotted<Document> doc;
     std::vector<IndexKeyDatum> keyData;
 
-    // True if this WSM has survived a yield in RID_AND_IDX state.
-    // TODO consider replacing by tracking SnapshotIds for IndexKeyDatums.
-    bool isSuspicious = false;
-
     bool hasRecordId() const;
     bool hasObj() const;
     bool hasOwnedObj() const;
@@ -302,20 +302,6 @@ public:
     void transitionToOwnedObj(WorkingSetID id);
 
     /**
-     * Returns the list of working set ids that have transitioned into the RID_AND_IDX state since
-     * the last yield. The members corresponding to these ids may have since transitioned to a
-     * different state or been freed, so these cases must be handled by the caller. The list may
-     * also contain duplicates.
-     *
-     * Execution stages are *not* responsible for managing this list, as working set ids are added
-     * to the set automatically by WorkingSet::transitionToRecordIdAndIdx().
-     *
-     * As a side effect, calling this method clears the list of flagged yield sensitive ids kept by
-     * the working set.
-     */
-    std::vector<WorkingSetID> getAndClearYieldSensitiveIds();
-
-    /**
      * Registers an IndexAccessMethod pointer with the WorkingSet, and returns a handle that can be
      * used to recover the IndexAccessMethod.
      */
@@ -329,6 +315,24 @@ public:
         return _registeredIndexes[indexId];
     }
 
+    /**
+     * Returns the WorkingSetMember with the given id after removing it from this WorkingSet. The
+     * WSM can be reinstated in the WorkingSet by calling 'emplace()'.
+     *
+     * WorkingSetMembers typically only temporarily live free of their WorkingSet, so calls to
+     * 'extract()' and 'emplace()' should come in pairs.
+     */
+    WorkingSetMember extract(WorkingSetID);
+
+    /**
+     * Puts the given WorkingSetMember into this WorkingSet. Assigns the WorkingSetMember an id and
+     * returns it. This id can be used later to obtain a pointer to the WSM using 'get()'.
+     *
+     * WorkingSetMembers typically only temporarily live free of their WorkingSet, so calls to
+     * 'extract()' and 'emplace()' should come in pairs.
+     */
+    WorkingSetID emplace(WorkingSetMember&&);
+
 private:
     struct MemberHolder {
         // Free list link if freed. Points to self if in use.
@@ -346,9 +350,6 @@ private:
     // If _freeList == INVALID_ID, the free list is empty and all elements in _data are in use.
     WorkingSetID _freeList;
 
-    // Contains ids of WSMs that may need to be adjusted when we next yield.
-    std::vector<WorkingSetID> _yieldSensitiveIds;
-
     // Holds IndexAccessMethods that have been registered with 'registerIndexAccessMethod()`. The
     // WorkingSetRegisteredIndexId is the offset into the vector.
     std::vector<const IndexAccessMethod*> _registeredIndexes;
diff --git a/src/mongo/db/exec/working_set_common.cpp b/src/mongo/db/exec/working_set_common.cpp
index e59d7c92d17..ef1fa322c42 100644
--- a/src/mongo/db/exec/working_set_common.cpp
+++ b/src/mongo/db/exec/working_set_common.cpp
@@ -40,20 +40,6 @@
 
 namespace mongo {
 
-void WorkingSetCommon::prepareForSnapshotChange(WorkingSet* workingSet) {
-    for (auto id : workingSet->getAndClearYieldSensitiveIds()) {
-        if (workingSet->isFree(id)) {
-            continue;
-        }
-
-        // We may see the same member twice, so anything we do here should be idempotent.
-        WorkingSetMember* member = workingSet->get(id);
-        if (member->getState() == WorkingSetMember::RID_AND_IDX) {
-            member->isSuspicious = true;
-        }
-    }
-}
-
 bool WorkingSetCommon::fetch(OperationContext* opCtx,
                              WorkingSet* workingSet,
                              WorkingSetID id,
@@ -69,23 +55,29 @@ bool WorkingSetCommon::fetch(OperationContext* opCtx,
         return false;
     }
 
-    member->resetDocument(opCtx->recoveryUnit()->getSnapshotId(), record->data.releaseToBson());
+    auto currentSnapshotId = opCtx->recoveryUnit()->getSnapshotId();
+    member->resetDocument(currentSnapshotId, record->data.releaseToBson());
 
-    if (member->isSuspicious) {
-        // Make sure that all of the keyData is still valid for this copy of the document.  This
-        // ensures both that index-provided filters and sort orders still hold.
-        //
-        // TODO provide a way for the query planner to opt out of this checking if it is unneeded
-        // due to the structure of the plan.
-        invariant(!member->keyData.empty());
+    // Make sure that all of the keyData is still valid for this copy of the document.  This ensures
+    // both that index-provided filters and sort orders still hold.
+    //
+    // TODO provide a way for the query planner to opt out of this checking if it is unneeded due to
+    // the structure of the plan.
+    if (member->getState() == WorkingSetMember::RID_AND_IDX) {
         for (size_t i = 0; i < member->keyData.size(); i++) {
+            auto&& memberKey = member->keyData[i];
+            // For storage engines that support document-level concurrency, if this key was obtained
+            // in the current snapshot, then move on to the next key.
+            if (supportsDocLocking() && memberKey.snapshotId == currentSnapshotId) {
+                continue;
+            }
+
             KeyStringSet keys;
             // There's no need to compute the prefixes of the indexed fields that cause the index to
             // be multikey when ensuring the keyData is still valid.
             KeyStringSet* multikeyMetadataKeys = nullptr;
             MultikeyPaths* multikeyPaths = nullptr;
-            auto indexId = member->keyData[i].indexId;
-            auto* iam = workingSet->retrieveIndexAccessMethod(indexId);
+            auto* iam = workingSet->retrieveIndexAccessMethod(memberKey.indexId);
             iam->getKeys(member->doc.value().toBson(),
                          IndexAccessMethod::GetKeysMode::kEnforceConstraints,
                          &keys,
@@ -93,7 +85,7 @@ bool WorkingSetCommon::fetch(OperationContext* opCtx,
                          multikeyPaths,
                          member->recordId);
             KeyString::HeapBuilder keyString(iam->getSortedDataInterface()->getKeyStringVersion(),
-                                             member->keyData[i].keyData,
+                                             memberKey.keyData,
                                              iam->getSortedDataInterface()->getOrdering(),
                                              member->recordId);
             if (!keys.count(keyString.release())) {
@@ -101,8 +93,6 @@ bool WorkingSetCommon::fetch(OperationContext* opCtx,
                 return false;
             }
         }
-
-        member->isSuspicious = false;
     }
 
     member->keyData.clear();
diff --git a/src/mongo/db/exec/working_set_common.h b/src/mongo/db/exec/working_set_common.h
index 034adb38fba..15657a87725 100644
--- a/src/mongo/db/exec/working_set_common.h
+++ b/src/mongo/db/exec/working_set_common.h
@@ -42,18 +42,6 @@ class SeekableRecordCursor;
 class WorkingSetCommon {
 public:
     /**
-     * This must be called as part of "saveState" operations after all nodes in the tree save their
-     * state.
-     *
-     * Iterates over WorkingSetIDs in 'workingSet' which are "sensitive to yield". These are ids
-     * that have transitioned into the RID_AND_IDX state since the previous yield.
-     *
-     * The RID_AND_IDX members are tagged as suspicious so that they can be handled properly in case
-     * the document keyed by the index key is deleted or updated during the yield.
-     */
-    static void prepareForSnapshotChange(WorkingSet* workingSet);
-
-    /**
      * Transitions the WorkingSetMember with WorkingSetID 'id' from the RID_AND_IDX state to the
      * RID_AND_OBJ state by fetching a document. Does the fetch using 'cursor'.
      *
diff --git a/src/mongo/db/exec/working_set_test.cpp b/src/mongo/db/exec/working_set_test.cpp
index 67513a75a27..7885e64a6d0 100644
--- a/src/mongo/db/exec/working_set_test.cpp
+++ b/src/mongo/db/exec/working_set_test.cpp
@@ -121,7 +121,8 @@ TEST_F(WorkingSetFixture, getFieldFromIndex) {
     string secondName = "y";
     int secondValue = 10;
 
-    member->keyData.push_back(IndexKeyDatum(BSON(firstName << 1), BSON("" << firstValue), 0));
+    member->keyData.push_back(
+        IndexKeyDatum(BSON(firstName << 1), BSON("" << firstValue), 0, SnapshotId{}));
     // Also a minor lie as RecordId is bogus.
     ws->transitionToRecordIdAndIdx(id);
     BSONElement elt;
@@ -131,7 +132,8 @@ TEST_F(WorkingSetFixture, getFieldFromIndex) {
     ASSERT_FALSE(member->getFieldDotted("foo", &elt));
 
     // Add another index datum.
-    member->keyData.push_back(IndexKeyDatum(BSON(secondName << 1), BSON("" << secondValue), 0));
+    member->keyData.push_back(
+        IndexKeyDatum(BSON(secondName << 1), BSON("" << secondValue), 0, SnapshotId{}));
     ASSERT_TRUE(member->getFieldDotted(secondName, &elt));
     ASSERT_EQUALS(elt.numberInt(), secondValue);
     ASSERT_TRUE(member->getFieldDotted(firstName, &elt));
@@ -144,7 +146,8 @@ TEST_F(WorkingSetFixture, getDottedFieldFromIndex) {
     string firstName = "x.y";
     int firstValue = 5;
 
-    member->keyData.push_back(IndexKeyDatum(BSON(firstName << 1), BSON("" << firstValue), 0));
+    member->keyData.push_back(
+        IndexKeyDatum(BSON(firstName << 1), BSON("" << firstValue), 0, SnapshotId{}));
     ws->transitionToRecordIdAndIdx(id);
     BSONElement elt;
     ASSERT_TRUE(member->getFieldDotted(firstName, &elt));
@@ -214,7 +217,6 @@ TEST_F(WorkingSetFixture, RecordIdAndObjStateCanRoundtripThroughSerialization) {
     ASSERT_DOCUMENT_EQ(roundtripped.doc.value(), doc);
     ASSERT_EQ(roundtripped.doc.snapshotId().toNumber(), 42u);
     ASSERT_EQ(roundtripped.recordId.repr(), 43);
-    ASSERT_FALSE(roundtripped.isSuspicious);
     ASSERT_FALSE(roundtripped.metadata());
 }
 
@@ -228,16 +230,15 @@ TEST_F(WorkingSetFixture, OwnedObjStateCanRoundtripThroughSerialization) {
     ASSERT_DOCUMENT_EQ(roundtripped.doc.value(), doc);
     ASSERT_EQ(roundtripped.doc.snapshotId().toNumber(), 42u);
     ASSERT(roundtripped.recordId.isNull());
-    ASSERT_FALSE(roundtripped.isSuspicious);
     ASSERT_FALSE(roundtripped.metadata());
 }
 
 TEST_F(WorkingSetFixture, RecordIdAndIdxStateCanRoundtripThroughSerialization) {
     member->recordId = RecordId{43};
-    member->keyData.emplace_back(BSON("a" << 1 << "b" << 1), BSON("" << 3 << "" << 4), 8u);
-    member->keyData.emplace_back(BSON("c" << -1), BSON("" << 5), 9u);
+    member->keyData.emplace_back(
+        BSON("a" << 1 << "b" << 1), BSON("" << 3 << "" << 4), 8u, SnapshotId{11u});
+    member->keyData.emplace_back(BSON("c" << -1), BSON("" << 5), 9u, SnapshotId{12u});
     ws->transitionToRecordIdAndIdx(id);
-    ASSERT_FALSE(member->isSuspicious);
 
     auto roundtripped = roundtripWsmThroughSerialization(*member);
     ASSERT_EQ(WorkingSetMember::RID_AND_IDX, roundtripped.getState());
@@ -247,12 +248,13 @@ TEST_F(WorkingSetFixture, RecordIdAndIdxStateCanRoundtripThroughSerialization) {
     ASSERT_BSONOBJ_EQ(roundtripped.keyData[0].indexKeyPattern, BSON("a" << 1 << "b" << 1));
     ASSERT_BSONOBJ_EQ(roundtripped.keyData[0].keyData, BSON("" << 3 << "" << 4));
     ASSERT_EQ(roundtripped.keyData[0].indexId, 8u);
+    ASSERT_EQ(roundtripped.keyData[0].snapshotId.toNumber(), 11u);
 
     ASSERT_BSONOBJ_EQ(roundtripped.keyData[1].indexKeyPattern, BSON("c" << -1));
     ASSERT_BSONOBJ_EQ(roundtripped.keyData[1].keyData, BSON("" << 5));
     ASSERT_EQ(roundtripped.keyData[1].indexId, 9u);
+    ASSERT_EQ(roundtripped.keyData[1].snapshotId.toNumber(), 12u);
 
-    ASSERT_TRUE(roundtripped.isSuspicious);
     ASSERT_FALSE(roundtripped.metadata());
 }
 
@@ -269,7 +271,6 @@ TEST_F(WorkingSetFixture, WsmWithMetadataCanRoundtripThroughSerialization) {
     ASSERT_FALSE(roundtripped.doc.value().metadata());
     ASSERT_TRUE(roundtripped.doc.snapshotId().isNull());
     ASSERT_TRUE(roundtripped.recordId.isNull());
-    ASSERT_FALSE(roundtripped.isSuspicious);
 
     ASSERT_TRUE(roundtripped.metadata());
     ASSERT_TRUE(roundtripped.metadata().hasTextScore());
@@ -280,4 +281,34 @@ TEST_F(WorkingSetFixture, WsmWithMetadataCanRoundtripThroughSerialization) {
     ASSERT_FALSE(roundtripped.metadata().hasGeoNearDistance());
 }
 
+TEST_F(WorkingSetFixture, WsmCanBeExtractedAndReinserted) {
+    Document doc{{"foo", Value{"bar"_sd}}};
+    member->doc.setValue(doc);
+    member->doc.setSnapshotId(SnapshotId{42u});
+    member->recordId = RecordId{43};
+    ws->transitionToRecordIdAndObj(id);
+    member = nullptr;
+    ASSERT_FALSE(ws->isFree(id));
+
+    auto extractedWsm = ws->extract(id);
+    ASSERT_TRUE(ws->isFree(id));
+    ASSERT_TRUE(extractedWsm.hasObj());
+    ASSERT_EQ(extractedWsm.getState(), WorkingSetMember::RID_AND_OBJ);
+    ASSERT_DOCUMENT_EQ(extractedWsm.doc.value(), doc);
+    ASSERT_EQ(extractedWsm.doc.snapshotId().toNumber(), 42u);
+    ASSERT_EQ(extractedWsm.recordId.repr(), 43);
+    ASSERT_FALSE(extractedWsm.metadata());
+
+    auto emplacedId = ws->emplace(std::move(extractedWsm));
+    ASSERT_FALSE(ws->isFree(emplacedId));
+
+    auto emplacedWsm = ws->get(emplacedId);
+    ASSERT_TRUE(emplacedWsm->hasObj());
+    ASSERT_EQ(emplacedWsm->getState(), WorkingSetMember::RID_AND_OBJ);
+    ASSERT_DOCUMENT_EQ(emplacedWsm->doc.value(), doc);
+    ASSERT_EQ(emplacedWsm->doc.snapshotId().toNumber(), 42u);
+    ASSERT_EQ(emplacedWsm->recordId.repr(), 43);
+    ASSERT_FALSE(emplacedWsm->metadata());
+}
+
 }  // namespace mongo