SERVER-42077 Add 'allowDiskUse' option to find command

11 files changed, 152 insertions, 3 deletions

diff --git a/src/mongo/db/exec/sort.cpp b/src/mongo/db/exec/sort.cpp
index 89de29e69f9..9bf39c08773 100644
--- a/src/mongo/db/exec/sort.cpp
+++ b/src/mongo/db/exec/sort.cpp
@@ -75,6 +75,7 @@ SortStage::SortStage(OperationContext* opCtx,
       _ws(ws),
       _pattern(params.pattern),
       _limit(params.limit),
+      _allowDiskUse(params.allowDiskUse),
       _sorted(false),
       _resultIterator(_data.end()),
       _memUsage(0) {
diff --git a/src/mongo/db/exec/sort.h b/src/mongo/db/exec/sort.h
index 85c4aa444b9..adadc9ad121 100644
--- a/src/mongo/db/exec/sort.h
+++ b/src/mongo/db/exec/sort.h
@@ -54,6 +54,9 @@ public:
 
     // Equal to 0 for no limit.
     size_t limit = 0;
+
+    // Whether we allow disk use, disabled by default.
+    bool allowDiskUse = false;
 };
 
 /**
@@ -99,6 +102,8 @@ private:
     // Equal to 0 for no limit.
     size_t _limit;
 
+    bool _allowDiskUse;
+
     //
     // Data storage
     //
diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript
index dbba9fe6c18..dafb9ca9d4c 100644
--- a/src/mongo/db/query/SConscript
+++ b/src/mongo/db/query/SConscript
@@ -135,9 +135,12 @@ env.Library(
     ],
     LIBDEPS=[
         "$BUILD_DIR/mongo/base",
-        "$BUILD_DIR/mongo/db/repl/read_concern_args",
         "$BUILD_DIR/mongo/db/catalog/collection_catalog",
-        "$BUILD_DIR/mongo/db/pipeline/runtime_constants_idl"
+        # TODO: This dependency edge can be removed when the 'allowDiskUse' option no longer depends
+        # on enabling test commands.
+        "$BUILD_DIR/mongo/db/commands/test_commands_enabled",
+        "$BUILD_DIR/mongo/db/pipeline/runtime_constants_idl",
+        "$BUILD_DIR/mongo/db/repl/read_concern_args",
     ],
 )
 
diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 2be65fbf8d2..d64f43b4611 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -659,6 +659,7 @@ QuerySolutionNode* QueryPlannerAnalysis::analyzeSort(const CanonicalQuery& query
     sort->pattern = sortObj;
     sort->children.push_back(solnRoot);
     solnRoot = sort;
+    sort->allowDiskUse = qr.allowDiskUse();
     // When setting the limit on the sort, we need to consider both
     // the limit N and skip count M. The sort should return an ordered list
     // N + M items so that the skip stage can discard the first M results.
diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp
index 5072c883430..4fc31cd4965 100644
--- a/src/mongo/db/query/query_request.cpp
+++ b/src/mongo/db/query/query_request.cpp
@@ -63,6 +63,8 @@ const string QueryRequest::metaRecordId("recordId");
 const string QueryRequest::metaSortKey("sortKey");
 const string QueryRequest::metaTextScore("textScore");
 
+const string QueryRequest::kAllowDiskUseField("allowDiskUse");
+
 const long long QueryRequest::kDefaultBatchSize = 101;
 
 namespace {
@@ -262,6 +264,17 @@ StatusWith<unique_ptr<QueryRequest>> QueryRequest::parseFromFindCommand(unique_p
             }
 
             qr->_wantMore = !el.boolean();
+        } else if (fieldName == kAllowDiskUseField) {
+            if (!getTestCommandsEnabled()) {
+                return Status(ErrorCodes::FailedToParse,
+                              "allowDiskUse is not allowed unless test commands are enabled.");
+            }
+            Status status = checkFieldType(el, Bool);
+            if (!status.isOK()) {
+                return status;
+            }
+
+            qr->_allowDiskUse = el.boolean();
         } else if (fieldName == kCommentField) {
             Status status = checkFieldType(el, String);
             if (!status.isOK()) {
@@ -495,6 +508,10 @@ void QueryRequest::asFindCommandInternal(BSONObjBuilder* cmdBuilder) const {
         cmdBuilder->append(kLimitField, *_limit);
     }
 
+    if (_allowDiskUse) {
+        cmdBuilder->append(kAllowDiskUseField, true);
+    }
+
     if (_batchSize) {
         cmdBuilder->append(kBatchSizeField, *_batchSize);
     }
@@ -1132,6 +1149,9 @@ StatusWith<BSONObj> QueryRequest::asAggregationCommand() const {
     if (!_unwrappedReadPref.isEmpty()) {
         aggregationBuilder.append(QueryRequest::kUnwrappedReadPrefField, _unwrappedReadPref);
     }
+    if (_allowDiskUse) {
+        aggregationBuilder.append(QueryRequest::kAllowDiskUseField, _allowDiskUse);
+    }
     if (_runtimeConstants) {
         BSONObjBuilder rtcBuilder(aggregationBuilder.subobjStart(kRuntimeConstantsField));
         _runtimeConstants->serialize(&rtcBuilder);
diff --git a/src/mongo/db/query/query_request.h b/src/mongo/db/query/query_request.h
index ade62c52ddf..7b0cb74dd80 100644
--- a/src/mongo/db/query/query_request.h
+++ b/src/mongo/db/query/query_request.h
@@ -140,6 +140,9 @@ public:
     static const std::string metaSortKey;
     static const std::string metaTextScore;
 
+    // Allow using disk during the find command.
+    static const std::string kAllowDiskUseField;
+
     const NamespaceString& nss() const {
         invariant(!_nss.isEmpty());
         return _nss;
@@ -241,6 +244,14 @@ public:
         _wantMore = wantMore;
     }
 
+    bool allowDiskUse() const {
+        return _allowDiskUse;
+    }
+
+    void setAllowDiskUse(bool allowDiskUse) {
+        _allowDiskUse = allowDiskUse;
+    }
+
     bool isExplain() const {
         return _explain;
     }
@@ -504,6 +515,8 @@ private:
     // allowed.
     boost::optional<long long> _batchSize;
 
+    bool _allowDiskUse = false;
+
     // Set only when parsed from an OP_QUERY find message. The value is computed by driver or shell
     // and is set to be a min of batchSize and limit provided by user. QR can have set either
     // ntoreturn or batchSize / limit.
diff --git a/src/mongo/db/query/query_request_test.cpp b/src/mongo/db/query/query_request_test.cpp
index ed4d369602b..7ee502140f3 100644
--- a/src/mongo/db/query/query_request_test.cpp
+++ b/src/mongo/db/query/query_request_test.cpp
@@ -35,6 +35,7 @@
 
 #include "mongo/db/catalog/collection_catalog.h"
 #include "mongo/db/catalog/collection_mock.h"
+#include "mongo/db/commands/test_commands_enabled.h"
 #include "mongo/db/dbmessage.h"
 #include "mongo/db/json.h"
 #include "mongo/db/namespace_string.h"
@@ -1089,6 +1090,51 @@ TEST(QueryRequestTest, DefaultQueryParametersCorrect) {
     ASSERT_EQUALS(false, qr->isExhaust());
     ASSERT_EQUALS(false, qr->isAllowPartialResults());
     ASSERT_EQUALS(false, qr->getRuntimeConstants().has_value());
+    ASSERT_EQUALS(false, qr->allowDiskUse());
+}
+
+TEST(QueryRequestTest, ParseCommandAllowDiskUseTrue) {
+    const bool oldTestCommandsEnabledVal = getTestCommandsEnabled();
+    ON_BLOCK_EXIT([&] { setTestCommandsEnabled(oldTestCommandsEnabledVal); });
+    setTestCommandsEnabled(true);
+
+    BSONObj cmdObj = fromjson("{find: 'testns', allowDiskUse: true}");
+    const NamespaceString nss("test.testns");
+    const bool isExplain = false;
+    auto result = QueryRequest::makeFromFindCommand(nss, cmdObj, isExplain);
+
+    ASSERT_OK(result.getStatus());
+    ASSERT_EQ(true, result.getValue()->allowDiskUse());
+}
+
+TEST(QueryRequestTest, ParseCommandAllowDiskUseFalse) {
+    const bool oldTestCommandsEnabledVal = getTestCommandsEnabled();
+    ON_BLOCK_EXIT([&] { setTestCommandsEnabled(oldTestCommandsEnabledVal); });
+    setTestCommandsEnabled(true);
+
+    BSONObj cmdObj = fromjson("{find: 'testns', allowDiskUse: false}");
+    const NamespaceString nss("test.testns");
+    const bool isExplain = false;
+    auto result = QueryRequest::makeFromFindCommand(nss, cmdObj, isExplain);
+
+    ASSERT_OK(result.getStatus());
+    ASSERT_EQ(false, result.getValue()->allowDiskUse());
+}
+
+TEST(QueryRequestTest, ParseCommandAllowDiskUseTestCommandsDisabled) {
+    const bool oldTestCommandsEnabledVal = getTestCommandsEnabled();
+    ON_BLOCK_EXIT([&] { setTestCommandsEnabled(oldTestCommandsEnabledVal); });
+    setTestCommandsEnabled(false);
+
+    BSONObj cmdObj = fromjson("{find: 'testns', allowDiskUse: true}");
+    const NamespaceString nss("test.testns");
+    const bool isExplain = false;
+    auto result = QueryRequest::makeFromFindCommand(nss, cmdObj, isExplain);
+
+    ASSERT_NOT_OK(result.getStatus());
+    ASSERT_EQ(ErrorCodes::FailedToParse, result.getStatus().code());
+    ASSERT_STRING_CONTAINS(result.getStatus().toString(),
+                           "allowDiskUse is not allowed unless test commands are enabled.");
 }
 
 //
@@ -1384,6 +1430,46 @@ TEST(QueryRequestTest, ConvertToAggregationWithRuntimeConstantsSucceeds) {
     ASSERT_EQ(ar.getValue().getRuntimeConstants()->getClusterTime(), rtc.getClusterTime());
 }
 
+TEST(QueryRequestTest, ConvertToAggregationWithAllowDiskUseTrueSucceeds) {
+    QueryRequest qr(testns);
+    qr.setAllowDiskUse(true);
+    const auto aggCmd = qr.asAggregationCommand();
+    ASSERT_OK(aggCmd.getStatus());
+
+    auto ar = AggregationRequest::parseFromBSON(testns, aggCmd.getValue());
+    ASSERT_OK(ar.getStatus());
+    ASSERT_EQ(true, ar.getValue().shouldAllowDiskUse());
+}
+
+TEST(QueryRequestTest, ConvertToAggregationWithAllowDiskUseFalseSucceeds) {
+    QueryRequest qr(testns);
+    qr.setAllowDiskUse(false);
+    const auto aggCmd = qr.asAggregationCommand();
+    ASSERT_OK(aggCmd.getStatus());
+
+    auto ar = AggregationRequest::parseFromBSON(testns, aggCmd.getValue());
+    ASSERT_OK(ar.getStatus());
+    ASSERT_EQ(false, ar.getValue().shouldAllowDiskUse());
+}
+
+TEST(QueryRequestTest, ConvertToFindWithAllowDiskUseTrueSucceeds) {
+    QueryRequest qr(testns);
+    qr.setAllowDiskUse(true);
+    const auto findCmd = qr.asFindCommand();
+
+    BSONElement elem = findCmd[QueryRequest::kAllowDiskUseField];
+    ASSERT_EQ(true, elem.isBoolean());
+    ASSERT_EQ(true, elem.Bool());
+}
+
+TEST(QueryRequestTest, ConvertToFindWithAllowDiskUseFalseSucceeds) {
+    QueryRequest qr(testns);
+    qr.setAllowDiskUse(false);
+    const auto findCmd = qr.asFindCommand();
+
+    ASSERT_FALSE(findCmd[QueryRequest::kAllowDiskUseField]);
+}
+
 TEST(QueryRequestTest, ParseFromLegacyObjMetaOpComment) {
     BSONObj queryObj = fromjson(
         "{$query: {a: 1},"
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index 9f1951c14a9..331e94875dd 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -974,6 +974,7 @@ QuerySolutionNode* SortNode::clone() const {
     copy->_sorts = this->_sorts;
     copy->pattern = this->pattern;
     copy->limit = this->limit;
+    copy->allowDiskUse = this->allowDiskUse;
 
     return copy;
 }
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index a665b823529..16098b0146b 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -691,7 +691,10 @@ struct SortKeyGeneratorNode : public QuerySolutionNode {
 };
 
 struct SortNode : public QuerySolutionNode {
-    SortNode() : _sorts(SimpleBSONObjComparator::kInstance.makeBSONObjSet()), limit(0) {}
+    SortNode()
+        : _sorts(SimpleBSONObjComparator::kInstance.makeBSONObjSet()),
+          limit(0),
+          allowDiskUse(false) {}
 
     virtual ~SortNode() {}
 
@@ -731,6 +734,8 @@ struct SortNode : public QuerySolutionNode {
 
     // Sum of both limit and skip count in the parsed query.
     size_t limit;
+
+    bool allowDiskUse;
 };
 
 struct LimitNode : public QuerySolutionNode {
diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp
index 190ed408770..8ab1ca42cfb 100644
--- a/src/mongo/db/query/stage_builder.cpp
+++ b/src/mongo/db/query/stage_builder.cpp
@@ -129,6 +129,7 @@ PlanStage* buildStages(OperationContext* opCtx,
             SortStageParams params;
             params.pattern = sn->pattern;
             params.limit = sn->limit;
+            params.allowDiskUse = sn->allowDiskUse;
             return new SortStage(opCtx, params, ws, childStage);
         }
         case STAGE_SORT_KEY_GENERATOR: {
diff --git a/src/mongo/shell/query.js b/src/mongo/shell/query.js
index 4304903ca36..def25f8f65a 100644
--- a/src/mongo/shell/query.js
+++ b/src/mongo/shell/query.js
@@ -47,6 +47,7 @@ DBQuery.prototype.help = function() {
     print("\t.allowPartialResults()");
     print("\t.returnKey()");
     print("\t.showRecordId() - adds a $recordId field to each returned object");
+    print("\t.allowDiskUse() - allow using disk in completing the query");
 
     print("\nCursor methods");
     print("\t.toArray() - iterates through docs and returns an array of the results");
@@ -127,6 +128,10 @@ DBQuery.prototype._exec = function() {
                 throw new Error("collation requires use of read commands");
             }
 
+            if (this._special && this._query._allowDiskUse) {
+                throw new Error("allowDiskUse option requires use of read commands");
+            }
+
             this._cursor = this._mongo.find(this._ns,
                                             this._query,
                                             this._fields,
@@ -225,6 +230,10 @@ DBQuery.prototype._convertToCommand = function(canAttachReadPref) {
         cmd["collation"] = this._query.collation;
     }
 
+    if ("allowDiskUse" in this._query) {
+        cmd["allowDiskUse"] = this._query.allowDiskUse;
+    }
+
     if ((this._options & DBQuery.Option.tailable) != 0) {
         cmd["tailable"] = true;
     }
@@ -471,6 +480,10 @@ DBQuery.prototype.collation = function(collationSpec) {
     return this._addSpecial("collation", collationSpec);
 };
 
+DBQuery.prototype.allowDiskUse = function() {
+    return this._addSpecial("allowDiskUse", true);
+};
+
 /**
  * Sets the read preference for this cursor.
  *