SERVER-42583 Thread 'tempDir' through to SortStage.

This makes it possible to instantiate a SortExecutor inside
SortStage.

3 files changed, 38 insertions, 55 deletions

diff --git a/src/mongo/db/exec/sort.cpp b/src/mongo/db/exec/sort.cpp
index 1dfae92f7f6..b98870858bb 100644
--- a/src/mongo/db/exec/sort.cpp
+++ b/src/mongo/db/exec/sort.cpp
@@ -66,15 +66,21 @@ bool SortStage::WorkingSetComparator::operator()(const SortableDataItem& lhs,
     return lhs.recordId < rhs.recordId;
 }
 
-SortStage::SortStage(OperationContext* opCtx,
-                     const SortStageParams& params,
+SortStage::SortStage(boost::intrusive_ptr<ExpressionContext> expCtx,
                      WorkingSet* ws,
+                     BSONObj sortPattern,
+                     uint64_t limit,
+                     uint64_t maxMemoryUsageBytes,
                      std::unique_ptr<PlanStage> child)
-    : PlanStage(kStageType, opCtx),
+    : PlanStage(kStageType, expCtx->opCtx),
       _ws(ws),
-      _pattern(params.pattern),
-      _limit(params.limit),
-      _allowDiskUse(params.allowDiskUse),
+      _sortExecutor(SortPattern{sortPattern, expCtx},
+                    limit,
+                    maxMemoryUsageBytes,
+                    expCtx->tempDir,
+                    expCtx->allowDiskUse),
+      _pattern(std::move(sortPattern)),
+      _limit(limit),
       _sorted(false),
       _resultIterator(_data.end()),
       _memUsage(0) {
@@ -91,8 +97,6 @@ SortStage::SortStage(OperationContext* opCtx,
     }
 }
 
-SortStage::~SortStage() {}
-
 bool SortStage::isEOF() {
     // We're done when our child has no more results, we've sorted the child's results, and
     // we've returned all sorted results.
diff --git a/src/mongo/db/exec/sort.h b/src/mongo/db/exec/sort.h
index 8fc5827fe13..5a964e8bbc6 100644
--- a/src/mongo/db/exec/sort.h
+++ b/src/mongo/db/exec/sort.h
@@ -33,32 +33,13 @@
 #include <vector>
 
 #include "mongo/db/exec/plan_stage.h"
+#include "mongo/db/exec/sort_executor.h"
 #include "mongo/db/exec/sort_key_generator.h"
 #include "mongo/db/exec/working_set.h"
-#include "mongo/db/jsobj.h"
-#include "mongo/db/query/index_bounds.h"
 #include "mongo/db/record_id.h"
-#include "mongo/stdx/unordered_map.h"
 
 namespace mongo {
 
-class BtreeKeyGenerator;
-
-// Parameters that must be provided to a SortStage
-class SortStageParams {
-public:
-    SortStageParams() = default;
-
-    // How we're sorting.
-    BSONObj pattern;
-
-    // Equal to 0 for no limit.
-    size_t limit = 0;
-
-    // Whether we allow disk use, disabled by default.
-    bool allowDiskUse = false;
-};
-
 /**
  * Sorts the input received from the child according to the sort pattern provided.
  *
@@ -69,11 +50,12 @@ public:
  */
 class SortStage final : public PlanStage {
 public:
-    SortStage(OperationContext* opCtx,
-              const SortStageParams& params,
+    SortStage(boost::intrusive_ptr<ExpressionContext> expCtx,
               WorkingSet* ws,
+              BSONObj sortPattern,
+              uint64_t limit,
+              uint64_t maxMemoryUsageBytes,
               std::unique_ptr<PlanStage> child);
-    ~SortStage();
 
     bool isEOF() final;
     StageState doWork(WorkingSetID* out) final;
@@ -89,21 +71,18 @@ public:
     static const char* kStageType;
 
 private:
-    //
-    // Query Stage
-    //
-
     // Not owned by us.
     WorkingSet* _ws;
 
+    // TODO SERVER-42182: Use SortExecutor to implement 'doWork()'.
+    SortExecutor _sortExecutor;
+
     // The raw sort _pattern as expressed by the user
     BSONObj _pattern;
 
     // Equal to 0 for no limit.
     size_t _limit;
 
-    bool _allowDiskUse;
-
     //
     // Data storage
     //
diff --git a/src/mongo/db/exec/sort_test.cpp b/src/mongo/db/exec/sort_test.cpp
index 7ff0614d7eb..349d3a7a1af 100644
--- a/src/mongo/db/exec/sort_test.cpp
+++ b/src/mongo/db/exec/sort_test.cpp
@@ -51,6 +51,8 @@ namespace {
 
 class SortStageTest : public ServiceContextMongoDTest {
 public:
+    static constexpr uint64_t kMaxMemoryUsageBytes = 1024u * 1024u;
+
     SortStageTest() {
         getServiceContext()->setFastClockSource(std::make_unique<ClockSourceMock>());
         _opCtx = makeOperationContext();
@@ -62,11 +64,13 @@ public:
     }
 
     /**
-     * Test function to verify sort stage.
-     * SortStageParams will be initialized using patternStr, collator, and limit.
-     * inputStr represents the input data set in a BSONObj.
+     * Test function to verify sort stage. SortStage will be initialized using 'patternStr',
+     * 'collator', and 'limit;.
+     *
+     * 'inputStr' represents the input data set in a BSONObj.
      *     {input: [doc1, doc2, doc3, ...]}
-     * expectedStr represents the expected sorted data set.
+     *
+     * 'expectedStr; represents the expected sorted data set.
      *     {output: [docA, docB, docC, ...]}
      */
     void testWork(const char* patternStr,
@@ -97,20 +101,16 @@ public:
             queuedDataStage->pushBack(id);
         }
 
-        // Initialize SortStageParams
-        // Setting limit to 0 means no limit
-        SortStageParams params;
-        params.pattern = fromjson(patternStr);
-        params.limit = limit;
+        auto sortPattern = fromjson(patternStr);
 
         // Create an ExpressionContext for the SortKeyGeneratorStage.
-        boost::intrusive_ptr<ExpressionContext> pExpCtx(
-            new ExpressionContext(getOpCtx(), collator));
+        auto expCtx = make_intrusive<ExpressionContext>(getOpCtx(), collator);
 
         auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>(
-            pExpCtx, std::move(queuedDataStage), &ws, params.pattern);
+            expCtx, std::move(queuedDataStage), &ws, sortPattern);
 
-        SortStage sort(getOpCtx(), params, &ws, std::move(sortKeyGen));
+        SortStage sort(
+            expCtx, &ws, sortPattern, limit, kMaxMemoryUsageBytes, std::move(sortKeyGen));
 
         WorkingSetID id = WorkingSet::INVALID_ID;
         PlanStage::StageState state = PlanStage::NEED_TIME;
@@ -163,14 +163,14 @@ TEST_F(SortStageTest, SortEmptyWorkingSet) {
     WorkingSet ws;
 
     // Create an ExpressionContext for the SortKeyGeneratorStage.
-    boost::intrusive_ptr<ExpressionContext> pExpCtx(new ExpressionContext(getOpCtx(), nullptr));
+    auto expCtx = make_intrusive<ExpressionContext>(getOpCtx(), nullptr);
 
     // QueuedDataStage will be owned by SortStage.
     auto queuedDataStage = std::make_unique<QueuedDataStage>(getOpCtx(), &ws);
-    auto sortKeyGen = std::make_unique<SortKeyGeneratorStage>(
-        pExpCtx, std::move(queuedDataStage), &ws, BSONObj());
-    SortStageParams params;
-    SortStage sort(getOpCtx(), params, &ws, std::move(sortKeyGen));
+    auto sortKeyGen =
+        std::make_unique<SortKeyGeneratorStage>(expCtx, std::move(queuedDataStage), &ws, BSONObj());
+    auto sortPattern = BSON("a" << 1);
+    SortStage sort(expCtx, &ws, sortPattern, 0u, kMaxMemoryUsageBytes, std::move(sortKeyGen));
 
     // Check initial EOF state.
     ASSERT_FALSE(sort.isEOF());