SERVER-14174 enforce ntoreturn during plan ranking

3 files changed, 59 insertions, 8 deletions

diff --git a/jstests/core/batch_size.js b/jstests/core/batch_size.js
index 6cbc45dc803..5af59ab5391 100644
--- a/jstests/core/batch_size.js
+++ b/jstests/core/batch_size.js
@@ -73,3 +73,33 @@ var explain = t.find({a: {$gte: 50}}).sort({b: 1}).hint({a: 1}).limit(6).explain
 assert.lte(explain.nscanned, 60, 'S');
 assert.lte(explain.nscannedObjects, 60, 'T');
 assert.eq(explain.n, 6, 'U');
+
+
+// -------
+
+
+// During plan ranking, we treat ntoreturn as a limit. This prevents us from buffering
+// too much data in a blocking sort stage during plan ranking.
+t.drop();
+
+// Generate big string to use in the object - 1MB+ String
+var bigStr = "ABCDEFGHIJKLMNBOPQRSTUVWXYZ012345687890";
+while (bigStr.length < 1000000) { bigStr = bigStr + "::" + bigStr; }
+
+// Insert enough documents to exceed the 32 MB in-memory sort limit.
+for (var i = 0; i < 40; i++) {
+    var doc = {x: 1, y: 1, z: i, big: bigStr};
+    t.insert(doc);
+}
+
+// Two indices needed in order to trigger plan ranking. Neither index provides
+// the sort order.
+t.ensureIndex({x: 1});
+t.ensureIndex({y: 1});
+
+// We should only buffer 3 docs in memory.
+var cursor = t.find({x: 1, y: 1}).sort({z: -1}).limit(3);
+assert.eq(39, cursor.next()["z"]);
+assert.eq(38, cursor.next()["z"]);
+assert.eq(37, cursor.next()["z"]);
+assert(!cursor.hasNext());
diff --git a/src/mongo/db/query/multi_plan_runner.cpp b/src/mongo/db/query/multi_plan_runner.cpp
index 709e08d8414..41a9332b4da 100644
--- a/src/mongo/db/query/multi_plan_runner.cpp
+++ b/src/mongo/db/query/multi_plan_runner.cpp
@@ -337,10 +337,23 @@ namespace mongo {
                                 size_t(fraction * _collection->numRecords()));
         }
 
+        // We treat ntoreturn as though it is a limit during plan ranking.
+        // This means that ranking might not be great for sort + batchSize.
+        // But it also means that we don't buffer too much data for sort + limit.
+        // See SERVER-14174 for details.
+        size_t numToReturn = _query->getParsed().getNumToReturn();
+
+        // Determine the number of results which we will produce during the plan
+        // ranking phase before stopping.
+        size_t numResults = (size_t)internalQueryPlanEvaluationMaxResults;
+        if (numToReturn > 0) {
+            numResults = std::min(numToReturn, numResults);
+        }
+
         // Work the plans, stopping when a plan hits EOF or returns some
         // fixed number of results.
         for (size_t i = 0; i < numWorks; ++i) {
-            bool moreToDo = workAllPlans(objOut);
+            bool moreToDo = workAllPlans(objOut, numResults);
             if (!moreToDo) { break; }
         }
 
@@ -461,7 +474,7 @@ namespace mongo {
         return NULL != _backupPlan;
     }
 
-    bool MultiPlanRunner::workAllPlans(BSONObj* objOut) {
+    bool MultiPlanRunner::workAllPlans(BSONObj* objOut, size_t numResults) {
         bool doneWorking = false;
 
         for (size_t i = 0; i < _candidates.size(); ++i) {
@@ -484,8 +497,7 @@ namespace mongo {
                 candidate.results.push_back(id);
 
                 // Once a plan returns enough results, stop working.
-                if (candidate.results.size()
-                    >= size_t(internalQueryPlanEvaluationMaxResults)) {
+                if (candidate.results.size() >= numResults) {
                     doneWorking = true;
                 }
             }
diff --git a/src/mongo/db/query/multi_plan_runner.h b/src/mongo/db/query/multi_plan_runner.h
index 01dbc3091c3..9f8eb22e477 100644
--- a/src/mongo/db/query/multi_plan_runner.h
+++ b/src/mongo/db/query/multi_plan_runner.h
@@ -120,13 +120,22 @@ namespace mongo {
                                PlanInfo** planInfo) const;
 
     private:
+        //
+        // Have all our candidate plans do something.
+        // If all our candidate plans fail, *objOut will contain
+        // information on the failure.
+        //
+
         /**
-         * Have all our candidate plans do something.
-         * If all our candidate plans fail, *objOut will contain
-         * information on the failure.
+         * Calls work on each child plan in a round-robin fashion. We stop when any plan
+         * hits EOF or returns 'numResults' results.
+         *
+         * Returns true if we need to keep working the plans and false otherwise.
          */
-        bool workAllPlans(BSONObj* objOut);
+        bool workAllPlans(BSONObj* objOut, size_t numResults);
+
         void allPlansSaveState();
+
         void allPlansRestoreState();
 
         const Collection* _collection;