SERVER-7694 Enable allowing disk use for sorts in the find command.

For find commands that request a sort, the query execution
engine can now spill data to disk if the memory requirements
would exceed 'internalQueryExecBlockingSortBytes' and the
'allowDiskUse' parameter is set to true. By default, the
memory threshold is currently 100MB. This allows
applications which need to sort server-side to no longer be
subject to an arbitrary data size threshold.

3 files changed, 100 insertions, 29 deletions

diff --git a/jstests/noPassthrough/external_sort_find.js b/jstests/noPassthrough/external_sort_find.js
new file mode 100644
index 00000000000..c4f3e09e4f9
--- /dev/null
+++ b/jstests/noPassthrough/external_sort_find.js
@@ -0,0 +1,100 @@
+/**
+ * Test that the find command can spill to disk while executing a blocking sort, if the client
+ * explicitly allows disk usage.
+ *
+ * Must be run with read commands enabled, since legacy OP_QUERY reads do not support the
+ * 'allowDiskUse' parameter.
+ * @tags: [requires_find_command]
+ */
+(function() {
+"use strict";
+
+load("jstests/libs/analyze_plan.js");
+
+// Only allow blocking sort execution to use 100 kB of memory.
+const kMaxMemoryUsageBytes = 100 * 1024;
+
+const kNumDocsWithinMemLimit = 70;
+const kNumDocsExceedingMemLimit = 100;
+
+const kMemoryLimitExceededErrCode = 16819;
+
+const options = {
+    setParameter: "internalQueryExecMaxBlockingSortBytes=" + kMaxMemoryUsageBytes
+};
+const conn = MongoRunner.runMongod(options);
+assert.neq(null, conn, "mongod was unable to start up with options: " + tojson(options));
+
+const testDb = conn.getDB("test");
+const collection = testDb.external_sort_find;
+
+// Construct a document that is just over 1 kB.
+const charToRepeat = "-";
+const templateDoc = {
+    padding: charToRepeat.repeat(1024)
+};
+
+// Insert data into the collection without exceeding the memory threshold.
+for (let i = 0; i < kNumDocsWithinMemLimit; ++i) {
+    templateDoc.sequenceNumber = i;
+    assert.commandWorked(collection.insert(templateDoc));
+}
+
+// We should be able to successfully sort the collection with or without disk use allowed.
+assert.eq(kNumDocsWithinMemLimit, collection.find().sort({sequenceNumber: -1}).itcount());
+assert.eq(kNumDocsWithinMemLimit,
+          collection.find().sort({sequenceNumber: -1}).allowDiskUse().itcount());
+
+function getSortStats(allowDiskUse) {
+    let cursor = collection.find().sort({sequenceNumber: -1});
+    if (allowDiskUse) {
+        cursor = cursor.allowDiskUse();
+    }
+    const explain = cursor.explain("executionStats");
+    return getPlanStage(explain.executionStats.executionStages, "SORT");
+}
+
+// Explain should report that less than 100 kB of memory was used, and we did not spill to disk.
+// Test that this result is the same whether or not 'allowDiskUse' is set.
+let sortStats = getSortStats(false);
+assert.eq(sortStats.memLimit, kMaxMemoryUsageBytes);
+assert.lt(sortStats.totalDataSizeSorted, kMaxMemoryUsageBytes);
+assert.eq(sortStats.usedDisk, false);
+sortStats = getSortStats(true);
+assert.eq(sortStats.memLimit, kMaxMemoryUsageBytes);
+assert.lt(sortStats.totalDataSizeSorted, kMaxMemoryUsageBytes);
+assert.eq(sortStats.usedDisk, false);
+
+// Add enough data to exceed the memory threshold.
+for (let i = kNumDocsWithinMemLimit; i < kNumDocsExceedingMemLimit; ++i) {
+    templateDoc.sequenceNumber = i;
+    assert.commandWorked(collection.insert(templateDoc));
+}
+
+// The sort should fail if disk use is not allowed, but succeed if disk use is allowed.
+assert.commandFailedWithCode(
+    testDb.runCommand({find: collection.getName(), sort: {sequenceNumber: -1}}),
+    kMemoryLimitExceededErrCode);
+assert.eq(kNumDocsExceedingMemLimit,
+          collection.find().sort({sequenceNumber: -1}).allowDiskUse().itcount());
+
+// Explain should report that the SORT stage failed if disk use is not allowed.
+sortStats = getSortStats(false);
+assert.eq(sortStats.failed, true);
+assert.eq(sortStats.usedDisk, false);
+assert.lt(sortStats.totalDataSizeSorted, kMaxMemoryUsageBytes);
+assert(!sortStats.inputStage.hasOwnProperty("failed"));
+
+// Explain should report that >=100 kB of memory was used, and that we spilled to disk.
+sortStats = getSortStats(true);
+assert.eq(sortStats.memLimit, kMaxMemoryUsageBytes);
+assert.gte(sortStats.totalDataSizeSorted, kMaxMemoryUsageBytes);
+assert.eq(sortStats.usedDisk, true);
+
+// If disk use is not allowed but there is a limit, we should be able to avoid exceeding the memory
+// limit.
+assert.eq(kNumDocsWithinMemLimit,
+          collection.find().sort({sequenceNumber: -1}).limit(kNumDocsWithinMemLimit).itcount());
+
+MongoRunner.stopMongod(conn);
+}());
diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp
index d8ab73b8116..d50d8a10f4d 100644
--- a/src/mongo/db/query/query_request.cpp
+++ b/src/mongo/db/query/query_request.cpp
@@ -263,10 +263,6 @@ 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;
diff --git a/src/mongo/db/query/query_request_test.cpp b/src/mongo/db/query/query_request_test.cpp
index c2624e65ceb..4c6a5cd2f41 100644
--- a/src/mongo/db/query/query_request_test.cpp
+++ b/src/mongo/db/query/query_request_test.cpp
@@ -35,7 +35,6 @@
 
 #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"
@@ -1082,10 +1081,6 @@ TEST(QueryRequestTest, DefaultQueryParametersCorrect) {
 }
 
 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;
@@ -1096,10 +1091,6 @@ TEST(QueryRequestTest, ParseCommandAllowDiskUseTrue) {
 }
 
 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;
@@ -1109,22 +1100,6 @@ TEST(QueryRequestTest, ParseCommandAllowDiskUseFalse) {
     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.");
-}
-
 //
 // Extra fields cause the parse to fail.
 //