SERVER-42981 Make WorkingSetMember compatible for use with the Sorter.

The bulk of this change is to implement serialization and
deserialization routines for WorkingSetMember, so that the
Sorter can spill WorkingSetMembers to disk. In addition,
this changes the SortExecutor to sort WorkingSetMembers
internally as opposed to sorting Documents.

1 files changed, 89 insertions, 1 deletions

diff --git a/src/mongo/db/exec/working_set.cpp b/src/mongo/db/exec/working_set.cpp
index 1c565cca8cd..2aacba154ab 100644
--- a/src/mongo/db/exec/working_set.cpp
+++ b/src/mongo/db/exec/working_set.cpp
@@ -88,7 +88,7 @@ void WorkingSet::transitionToRecordIdAndIdx(WorkingSetID id) {
 
 void WorkingSet::transitionToRecordIdAndObj(WorkingSetID id) {
     WorkingSetMember* member = get(id);
-    member->_state = WorkingSetMember::RID_AND_OBJ;
+    member->transitionToRecordIdAndObj();
 }
 
 void WorkingSet::transitionToOwnedObj(WorkingSetID id) {
@@ -123,6 +123,9 @@ void WorkingSetMember::transitionToOwnedObj() {
     _state = OWNED_OBJ;
 }
 
+void WorkingSetMember::transitionToRecordIdAndObj() {
+    _state = WorkingSetMember::RID_AND_OBJ;
+}
 
 bool WorkingSetMember::hasRecordId() const {
     return _state == RID_AND_IDX || _state == RID_AND_OBJ;
@@ -187,4 +190,89 @@ void WorkingSetMember::resetDocument(SnapshotId snapshot, const BSONObj& obj) {
     md.reset(obj, false);
     doc.value() = md.freeze();
 }
+
+void WorkingSetMember::serializeForSorter(BufBuilder& buf) const {
+    // It is not legal to serialize a Document which has metadata attached to it. Any metadata must
+    // reside directly in the WorkingSetMember.
+    invariant(!doc.value().metadata());
+
+    buf.appendChar(static_cast<char>(_state));
+
+    if (hasObj()) {
+        doc.value().serializeForSorter(buf);
+        buf.appendNum(static_cast<unsigned long long>(doc.snapshotId().toNumber()));
+    }
+
+    if (_state == RID_AND_IDX) {
+        // First append the number of index keys, and then encode them in series.
+        buf.appendNum(static_cast<char>(keyData.size()));
+        for (auto&& indexKeyDatum : keyData) {
+            indexKeyDatum.indexKeyPattern.serializeForSorter(buf);
+            indexKeyDatum.keyData.serializeForSorter(buf);
+            buf.appendNum(indexKeyDatum.indexId);
+        }
+    }
+
+    if (hasRecordId()) {
+        buf.appendNum(recordId.repr());
+    }
+
+    _metadata.serializeForSorter(buf);
+}
+
+WorkingSetMember WorkingSetMember::deserializeForSorter(BufReader& buf,
+                                                        const SorterDeserializeSettings&) {
+    WorkingSetMember wsm;
+
+    // First decode the state, which instructs us on how to interpret the rest of the buffer.
+    wsm._state = static_cast<MemberState>(buf.read<char>());
+
+    if (wsm.hasObj()) {
+        wsm.doc.setValue(
+            Document::deserializeForSorter(buf, Document::SorterDeserializeSettings{}));
+        auto snapshotIdRepr = buf.read<LittleEndian<uint64_t>>();
+        auto snapshotId = snapshotIdRepr ? SnapshotId{snapshotIdRepr} : SnapshotId{};
+        wsm.doc.setSnapshotId(snapshotId);
+    }
+
+    if (wsm.getState() == WorkingSetMember::RID_AND_IDX) {
+        auto numKeys = buf.read<char>();
+        wsm.keyData.reserve(numKeys);
+        for (auto i = 0; i < numKeys; ++i) {
+            auto indexKeyPattern =
+                BSONObj::deserializeForSorter(buf, BSONObj::SorterDeserializeSettings{}).getOwned();
+            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});
+        }
+
+        // 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()) {
+        wsm.recordId = RecordId{buf.read<LittleEndian<int64_t>>()};
+    }
+
+    DocumentMetadataFields::deserializeForSorter(buf, &wsm._metadata);
+
+    return wsm;
+}
+
+WorkingSetRegisteredIndexId WorkingSet::registerIndexAccessMethod(
+    const IndexAccessMethod* indexAccess) {
+    for (WorkingSetRegisteredIndexId i = 0; i < _registeredIndexes.size(); ++i) {
+        if (_registeredIndexes[i] == indexAccess) {
+            return i;
+        }
+    }
+
+    _registeredIndexes.push_back(indexAccess);
+    return _registeredIndexes.size() - 1;
+}
+
 }  // namespace mongo