SERVER-45854 Hold WorkingSetMembers at a distance for sorting.

This change improves the performance of blocking sort
operations executed in the PlanStage layer using the
SortStage.

8 files changed, 86 insertions, 31 deletions

diff --git a/src/mongo/db/exec/sort.cpp b/src/mongo/db/exec/sort.cpp
index fa6955bd79f..c3779314d84 100644
--- a/src/mongo/db/exec/sort.cpp
+++ b/src/mongo/db/exec/sort.cpp
@@ -65,11 +65,11 @@ PlanStage::StageState SortStage::doWork(WorkingSetID* out) {
             auto member = _ws->get(id);
             invariant(member->metadata().hasSortKey());
 
-            auto&& extractedMember = _ws->extract(id);
+            SortableWorkingSetMember extractedMember{_ws->extract(id)};
 
             try {
-                auto sortKey = extractedMember.metadata().getSortKey();
-                _sortExecutor.add(std::move(sortKey), std::move(extractedMember));
+                auto sortKey = extractedMember->metadata().getSortKey();
+                _sortExecutor.add(sortKey, extractedMember);
             } catch (const AssertionException&) {
                 // Propagate runtime errors using the FAILED status code.
                 *out = WorkingSetCommon::allocateStatusMember(_ws, exceptionToStatus());
@@ -103,7 +103,7 @@ PlanStage::StageState SortStage::doWork(WorkingSetID* out) {
         return PlanStage::IS_EOF;
     }
 
-    *out = _ws->emplace(std::move(*nextWsm));
+    *out = _ws->emplace(nextWsm->extract());
     return PlanStage::ADVANCED;
 }
 
diff --git a/src/mongo/db/exec/sort.h b/src/mongo/db/exec/sort.h
index 3f8f04a43b5..c0bdd6deaed 100644
--- a/src/mongo/db/exec/sort.h
+++ b/src/mongo/db/exec/sort.h
@@ -80,7 +80,7 @@ private:
     // Not owned by us.
     WorkingSet* _ws;
 
-    SortExecutor<WorkingSetMember> _sortExecutor;
+    SortExecutor<SortableWorkingSetMember> _sortExecutor;
 
     // Whether or not we have finished loading data into '_sortExecutor'.
     bool _populated = false;
diff --git a/src/mongo/db/exec/sort_executor.cpp b/src/mongo/db/exec/sort_executor.cpp
index 3cf2521c13d..10eb98db4ce 100644
--- a/src/mongo/db/exec/sort_executor.cpp
+++ b/src/mongo/db/exec/sort_executor.cpp
@@ -29,9 +29,11 @@
 
 #include "mongo/platform/basic.h"
 
-#include "mongo/db/exec/document_value/value_comparator.h"
 #include "mongo/db/exec/sort_executor.h"
 
+#include "mongo/db/exec/document_value/value_comparator.h"
+#include "mongo/db/exec/working_set.h"
+
 namespace mongo {
 namespace {
 /**
@@ -57,5 +59,5 @@ MONGO_CREATE_SORTER(mongo::Value,
                     mongo::Document,
                     mongo::SortExecutor<mongo::Document>::Comparator);
 MONGO_CREATE_SORTER(mongo::Value,
-                    mongo::WorkingSetMember,
-                    mongo::SortExecutor<mongo::WorkingSetMember>::Comparator);
+                    mongo::SortableWorkingSetMember,
+                    mongo::SortExecutor<mongo::SortableWorkingSetMember>::Comparator);
diff --git a/src/mongo/db/exec/sort_executor.h b/src/mongo/db/exec/sort_executor.h
index 6f0bac4b5c9..a9c155aadca 100644
--- a/src/mongo/db/exec/sort_executor.h
+++ b/src/mongo/db/exec/sort_executor.h
@@ -126,11 +126,11 @@ public:
      * Add data item to be sorted of type T with sort key specified by Value to the sort executor.
      * Should only be called before 'loadingDone()' is called.
      */
-    void add(Value sortKey, T data) {
+    void add(const Value& sortKey, const T& data) {
         if (!_sorter) {
             _sorter.reset(DocumentSorter::make(makeSortOptions(), Comparator(_sortPattern)));
         }
-        _sorter->add(std::move(sortKey), std::move(data));
+        _sorter->add(sortKey, data);
 
         _stats.totalDataSizeBytes += data.memUsageForSorter();
     }
diff --git a/src/mongo/db/exec/working_set.cpp b/src/mongo/db/exec/working_set.cpp
index 6529df45b07..c88775669df 100644
--- a/src/mongo/db/exec/working_set.cpp
+++ b/src/mongo/db/exec/working_set.cpp
@@ -202,7 +202,7 @@ void WorkingSetMember::resetDocument(SnapshotId snapshot, const BSONObj& obj) {
     doc.value() = md.freeze();
 }
 
-void WorkingSetMember::serializeForSorter(BufBuilder& buf) const {
+void WorkingSetMember::serialize(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());
@@ -232,8 +232,7 @@ void WorkingSetMember::serializeForSorter(BufBuilder& buf) const {
     _metadata.serializeForSorter(buf);
 }
 
-WorkingSetMember WorkingSetMember::deserializeForSorter(BufReader& buf,
-                                                        const SorterDeserializeSettings&) {
+WorkingSetMember WorkingSetMember::deserialize(BufReader& buf) {
     WorkingSetMember wsm;
 
     // First decode the state, which instructs us on how to interpret the rest of the buffer.
@@ -272,6 +271,12 @@ WorkingSetMember WorkingSetMember::deserializeForSorter(BufReader& buf,
     return wsm;
 }
 
+SortableWorkingSetMember SortableWorkingSetMember::getOwned() const {
+    auto ret = *this;
+    ret._holder->makeObjOwnedIfNeeded();
+    return ret;
+}
+
 WorkingSetRegisteredIndexId WorkingSet::registerIndexAccessMethod(
     const IndexAccessMethod* indexAccess) {
     for (WorkingSetRegisteredIndexId i = 0; i < _registeredIndexes.size(); ++i) {
diff --git a/src/mongo/db/exec/working_set.h b/src/mongo/db/exec/working_set.h
index 7021ed5eb62..20112ddbdcf 100644
--- a/src/mongo/db/exec/working_set.h
+++ b/src/mongo/db/exec/working_set.h
@@ -129,9 +129,7 @@ public:
         OWNED_OBJ,
     };
 
-    struct SorterDeserializeSettings {};
-
-    static WorkingSetMember deserializeForSorter(BufReader& buf, const SorterDeserializeSettings&);
+    static WorkingSetMember deserialize(BufReader& buf);
 
     /**
      * Reset to an "empty" state.
@@ -224,24 +222,72 @@ public:
      */
     void resetDocument(SnapshotId snapshot, const BSONObj& obj);
 
-    void serializeForSorter(BufBuilder& buf) const;
+    void serialize(BufBuilder& buf) const;
+
+private:
+    friend class WorkingSet;
+
+    MemberState _state = WorkingSetMember::INVALID;
+
+    DocumentMetadataFields _metadata;
+};
+
+/**
+ * A variant of WorkingSetMember that is designed to be compatible with the SortExecutor. Objects of
+ * this type are small, acting only as a handle to the underlying WorkingSetMember. This means that
+ * they can be efficiently copied or moved during the sorting process while the actual
+ * WorkingSetMember data remains in place.
+ *
+ * A SortableWorkingSetMember supports serialization and deserialization so that objects of this
+ * type can be flushed to disk and later recovered.
+ */
+class SortableWorkingSetMember {
+public:
+    struct SorterDeserializeSettings {};
+
+    static SortableWorkingSetMember deserializeForSorter(BufReader& buf,
+                                                         const SorterDeserializeSettings&) {
+        return WorkingSetMember::deserialize(buf);
+    }
+
+    /**
+     * Constructs an empty SortableWorkingSetMember.
+     */
+    SortableWorkingSetMember() = default;
+
+    /**
+     * Constructs a SortableWorkingSetMember from a regular WorkingSetMember. Supports implicit
+     * conversion from WorkingSetMember.
+     */
+    /* implicit */ SortableWorkingSetMember(WorkingSetMember&& wsm)
+        : _holder(std::make_shared<WorkingSetMember>(std::move(wsm))) {}
+
+    void serializeForSorter(BufBuilder& buf) const {
+        _holder->serialize(buf);
+    }
 
     int memUsageForSorter() const {
-        return getMemUsage();
+        return _holder->getMemUsage();
     }
 
-    WorkingSetMember getOwned() const {
-        auto ret = *this;
-        ret.makeObjOwnedIfNeeded();
-        return ret;
+    /**
+     * Extracts and returns the underlying WorkingSetMember held by this SortableWorkingSetMember.
+     */
+    WorkingSetMember extract() {
+        return std::move(*_holder);
     }
 
-private:
-    friend class WorkingSet;
+    /**
+     * Returns a reference to the underlying WorkingSetMember.
+     */
+    const WorkingSetMember* operator->() const {
+        return _holder.get();
+    }
 
-    MemberState _state = WorkingSetMember::INVALID;
+    SortableWorkingSetMember getOwned() const;
 
-    DocumentMetadataFields _metadata;
+private:
+    std::shared_ptr<WorkingSetMember> _holder;
 };
 
 /**
diff --git a/src/mongo/db/exec/working_set_test.cpp b/src/mongo/db/exec/working_set_test.cpp
index 6fbad0911f4..230c2d81a1a 100644
--- a/src/mongo/db/exec/working_set_test.cpp
+++ b/src/mongo/db/exec/working_set_test.cpp
@@ -197,12 +197,14 @@ TEST_F(WorkingSetFixture, MetadataCanBeCorrectlyTransferredBackAndForthFromDocum
 namespace {
 // Serializes the given working set member to a buffer, then returns a working set member resulting
 // from deserializing this buffer.
-WorkingSetMember roundtripWsmThroughSerialization(const WorkingSetMember& wsm) {
+WorkingSetMember roundtripWsmThroughSerialization(WorkingSetMember wsm) {
+    SortableWorkingSetMember sortableWsm{std::move(wsm)};
     BufBuilder builder{};
-    wsm.serializeForSorter(builder);
+    sortableWsm.serializeForSorter(builder);
     BufReader reader{builder.buf(), static_cast<unsigned>(builder.len())};
-    return WorkingSetMember::deserializeForSorter(reader,
-                                                  WorkingSetMember::SorterDeserializeSettings{});
+    auto deserializedSortableWsm = SortableWorkingSetMember::deserializeForSorter(
+        reader, SortableWorkingSetMember::SorterDeserializeSettings{});
+    return deserializedSortableWsm.extract();
 }
 }  // namespace
 
diff --git a/src/mongo/db/pipeline/document_source_sort.cpp b/src/mongo/db/pipeline/document_source_sort.cpp
index 605d5022c3d..cb8355f8bae 100644
--- a/src/mongo/db/pipeline/document_source_sort.cpp
+++ b/src/mongo/db/pipeline/document_source_sort.cpp
@@ -226,7 +226,7 @@ void DocumentSourceSort::loadDocument(Document&& doc) {
     // already computed the sort key we'd have split the pipeline there, would be merging presorted
     // documents, and wouldn't use this method.
     std::tie(sortKey, docForSorter) = extractSortKey(std::move(doc));
-    _sortExecutor->add(sortKey, std::move(docForSorter));
+    _sortExecutor->add(sortKey, docForSorter);
 }
 
 void DocumentSourceSort::loadingDone() {