SERVER-13715 plan subqueries on construction of a SubplanRunner

This enforces the error behavior expected by aggregation.
(cherry picked from commit 20263ce4e9c345a2d3748ef2071f6dc860a092f1)

4 files changed, 156 insertions, 26 deletions

diff --git a/jstests/aggregation/bugs/server13715.js b/jstests/aggregation/bugs/server13715.js
new file mode 100644
index 00000000000..608e94d37a4
--- /dev/null
+++ b/jstests/aggregation/bugs/server13715.js
@@ -0,0 +1,25 @@
+// Make sure that aggregation and query agree on how NO_BLOCKING_SORT should fail
+// if only blocking sort solutions are available.
+
+var t = db.jstests_server13715;
+t.drop();
+
+t.save({_id: 0, name: "red", value: 2});
+t.save({_id: 1, name: "blue", value: 1});
+
+var cursor = t.aggregate([
+    {$match: {$or: [{name: "red"}, {name: "blue"}]}},
+    {$sort: {value: 1}}
+]);
+assert.eq(1, cursor.next()["_id"]);
+assert.eq(0, cursor.next()["_id"]);
+
+// Repeat the test with an index.
+t.ensureIndex({name: 1});
+
+cursor = t.aggregate([
+    {$match: {$or: [{name: "red"}, {name: "blue"}]}},
+    {$sort: {value: 1}}
+]);
+assert.eq(1, cursor.next()["_id"]);
+assert.eq(0, cursor.next()["_id"]);
diff --git a/src/mongo/db/query/get_runner.cpp b/src/mongo/db/query/get_runner.cpp
index f062fbe578e..2ef4c4b5260 100644
--- a/src/mongo/db/query/get_runner.cpp
+++ b/src/mongo/db/query/get_runner.cpp
@@ -346,7 +346,15 @@ namespace mongo {
 
             QLOG() << "Running query as sub-queries: " << canonicalQuery->toStringShort();
             LOG(2) << "Running query as sub-queries: " << canonicalQuery->toStringShort();
-            *out = new SubplanRunner(collection, plannerParams, canonicalQuery.release());
+
+            SubplanRunner* runner;
+            Status runnerStatus = SubplanRunner::make(collection, plannerParams,
+                                                      canonicalQuery.release(), &runner);
+            if (!runnerStatus.isOK()) {
+                return runnerStatus;
+            }
+
+            *out = runner;
             return Status::OK();
         }
 
diff --git a/src/mongo/db/query/subplan_runner.cpp b/src/mongo/db/query/subplan_runner.cpp
index d5c41541c70..9a901314ef1 100644
--- a/src/mongo/db/query/subplan_runner.cpp
+++ b/src/mongo/db/query/subplan_runner.cpp
@@ -90,6 +90,21 @@ namespace mongo {
         return true;
     }
 
+    // static
+    Status SubplanRunner::make(Collection* collection,
+                               const QueryPlannerParams& params,
+                               CanonicalQuery* cq,
+                               SubplanRunner** out) {
+        auto_ptr<SubplanRunner> autoRunner(new SubplanRunner(collection, params, cq));
+        Status planningStatus = autoRunner->planSubqueries();
+        if (!planningStatus.isOK()) {
+            return planningStatus;
+        }
+
+        *out = autoRunner.release();
+        return Status::OK();
+    }
+
     SubplanRunner::SubplanRunner(Collection* collection,
                                  const QueryPlannerParams& params,
                                  CanonicalQuery* cq)
@@ -101,7 +116,20 @@ namespace mongo {
           _policy(Runner::YIELD_MANUAL),
           _ns(cq->getParsed().ns()) { }
 
-    SubplanRunner::~SubplanRunner() { }
+    SubplanRunner::~SubplanRunner() {
+        while (!_solutions.empty()) {
+            vector<QuerySolution*> solns = _solutions.front();
+            for (size_t i = 0; i < solns.size(); i++) {
+                delete solns[i];
+            }
+            _solutions.pop();
+        }
+
+        while (!_cqs.empty()) {
+            delete _cqs.front();
+            _cqs.pop();
+        }
+    }
 
     Runner::RunnerState SubplanRunner::getNext(BSONObj* objOut, DiskLoc* dlOut) {
         if (_killed) {
@@ -162,19 +190,12 @@ namespace mongo {
         return _underlyingRunner->getNext(objOut, dlOut);
     }
 
-    bool SubplanRunner::runSubplans() {
-        // This is what we annotate with the index selections and then turn into a solution.
-        auto_ptr<OrMatchExpression> theOr(
-            static_cast<OrMatchExpression*>(_query->root()->shallowClone()));
+    Status SubplanRunner::planSubqueries() {
+        MatchExpression* theOr = _query->root();
 
-        // This is the skeleton of index selections that is inserted into the cache.
-        auto_ptr<PlanCacheIndexTree> cacheData(new PlanCacheIndexTree());
-
-        // We need this to extract cache-friendly index data from the index assignments.
-        map<BSONObj, size_t> indexMap;
         for (size_t i = 0; i < _plannerParams.indices.size(); ++i) {
             const IndexEntry& ie = _plannerParams.indices[i];
-            indexMap[ie.keyPattern] = i;
+            _indexMap[ie.keyPattern] = i;
             QLOG() << "Subplanner: index " << i << " is " << ie.toString() << endl;
         }
 
@@ -186,8 +207,10 @@ namespace mongo {
                                                                 orChild,
                                                                 &orChildCQ);
             if (!childCQStatus.isOK()) {
-                QLOG() << "Subplanner: Can't canonicalize subchild " << orChild->toString();
-                return false;
+                mongoutils::str::stream ss;
+                ss << "Subplanner: Can't canonicalize subchild " << orChild->toString()
+                   << " " << childCQStatus.reason();
+                return Status(ErrorCodes::BadValue, ss);
             }
 
             // Make sure it gets cleaned up.
@@ -202,17 +225,53 @@ namespace mongo {
             Status status = QueryPlanner::plan(*safeOrChildCQ, _plannerParams, &solutions);
 
             if (!status.isOK()) {
-                QLOG() << "Subplanner: Can't plan for subchild " << orChildCQ->toString();
-                return false;
+                mongoutils::str::stream ss;
+                ss << "Subplanner: Can't plan for subchild " << orChildCQ->toString()
+                   << " " << status.reason();
+                return Status(ErrorCodes::BadValue, ss);
             }
             QLOG() << "Subplanner: got " << solutions.size() << " solutions";
 
             if (0 == solutions.size()) {
                 // If one child doesn't have an indexed solution, bail out.
-                QLOG() << "Subplanner: No solutions for subchild " << orChildCQ->toString();
-                return false;
+                mongoutils::str::stream ss;
+                ss << "Subplanner: No solutions for subchild " << orChildCQ->toString();
+                return Status(ErrorCodes::BadValue, ss);
             }
-            else if (1 == solutions.size()) {
+
+            // Hang onto the canonicalized subqueries and the corresponding query solutions
+            // so that they can be used in subplan running later on.
+            _cqs.push(safeOrChildCQ.release());
+            _solutions.push(solutions);
+        }
+
+        return Status::OK();
+    }
+
+    bool SubplanRunner::runSubplans() {
+        // This is what we annotate with the index selections and then turn into a solution.
+        auto_ptr<OrMatchExpression> theOr(
+            static_cast<OrMatchExpression*>(_query->root()->shallowClone()));
+
+        // This is the skeleton of index selections that is inserted into the cache.
+        auto_ptr<PlanCacheIndexTree> cacheData(new PlanCacheIndexTree());
+
+        for (size_t i = 0; i < theOr->numChildren(); ++i) {
+            MatchExpression* orChild = theOr->getChild(i);
+
+            auto_ptr<CanonicalQuery> orChildCQ(_cqs.front());
+            _cqs.pop();
+
+            // 'solutions' is owned by the SubplanRunner instance until
+            // it is popped from the queue.
+            vector<QuerySolution*> solutions = _solutions.front();
+
+            // We already checked for zero solutions in planSubqueries(...).
+            invariant(!solutions.empty());
+
+            if (1 == solutions.size()) {
+                // There is only one solution. Transfer ownership to an auto_ptr.
+                _solutions.pop();
                 auto_ptr<QuerySolution> autoSoln(solutions[0]);
 
                 // We want a well-formed *indexed* solution.
@@ -230,7 +289,7 @@ namespace mongo {
 
                 // Add the index assignments to our original query.
                 Status tagStatus = QueryPlanner::tagAccordingToCache(
-                    orChild, autoSoln->cacheData->tree.get(), indexMap);
+                    orChild, autoSoln->cacheData->tree.get(), _indexMap);
 
                 if (!tagStatus.isOK()) {
                     QLOG() << "Subplanner: Failed to extract indices from subchild"
@@ -243,9 +302,11 @@ namespace mongo {
             }
             else {
                 // N solutions, rank them.  Takes ownership of safeOrChildCQ.
-                MultiPlanRunner* mpr = new MultiPlanRunner(_collection, safeOrChildCQ.release());
+                MultiPlanRunner* mpr = new MultiPlanRunner(_collection, orChildCQ.release());
 
-                // Dump all the solutions into the MPR.
+                // Dump all the solutions into the MPR. The MPR takes ownership of
+                // each solution.
+                _solutions.pop();
                 for (size_t i = 0; i < solutions.size(); ++i) {
                     WorkingSet* ws;
                     PlanStage* root;
@@ -287,7 +348,7 @@ namespace mongo {
 
                 // Add the index assignments to our original query.
                 Status tagStatus = QueryPlanner::tagAccordingToCache(
-                    orChild, bestSoln->cacheData->tree.get(), indexMap);
+                    orChild, bestSoln->cacheData->tree.get(), _indexMap);
 
                 if (!tagStatus.isOK()) {
                     QLOG() << "Subplanner: Failed to extract indices from subchild"
diff --git a/src/mongo/db/query/subplan_runner.h b/src/mongo/db/query/subplan_runner.h
index 6b6fe0d8401..7415b50c934 100644
--- a/src/mongo/db/query/subplan_runner.h
+++ b/src/mongo/db/query/subplan_runner.h
@@ -30,10 +30,12 @@
 
 #include <boost/scoped_ptr.hpp>
 #include <string>
+#include <queue>
 
 #include "mongo/base/status.h"
 #include "mongo/db/query/runner.h"
 #include "mongo/db/query/query_planner_params.h"
+#include "mongo/db/query/query_solution.h"
 
 namespace mongo {
 
@@ -45,9 +47,16 @@ namespace mongo {
 
     class SubplanRunner : public Runner {
     public:
-        SubplanRunner(Collection* collection,
-                      const QueryPlannerParams& params,
-                      CanonicalQuery* cq);
+        /**
+         * Used to create SubplanRunner instances. The caller owns the instance
+         * returned through 'out'.
+         *
+         * 'out' is valid only if an OK status is returned.
+         */
+        static Status make(Collection* collection,
+                           const QueryPlannerParams& params,
+                           CanonicalQuery* cq,
+                           SubplanRunner** out);
 
         static bool canUseSubplanRunner(const CanonicalQuery& query);
 
@@ -76,7 +85,21 @@ namespace mongo {
         virtual Status getInfo(TypeExplain** explain,
                                PlanInfo** planInfo) const;
 
+        /**
+         * Plan each branch of the $or independently, and store the resulting
+         * lists of query solutions in '_solutions'.
+         *
+         * Called from SubplanRunner::make so that getRunner can fail if
+         * subquery planning fails, rather than returning a runner and failing
+         * through getNext(...).
+         */
+        Status planSubqueries();
+
     private:
+        SubplanRunner(Collection* collection,
+                      const QueryPlannerParams& params,
+                      CanonicalQuery* cq);
+
         bool runSubplans();
 
         enum SubplanRunnerState {
@@ -99,6 +122,19 @@ namespace mongo {
         boost::scoped_ptr<Runner> _underlyingRunner;
 
         std::string _ns;
+
+        // We do the subquery planning up front, and keep the resulting
+        // query solutions here. Lists of query solutions are dequeued
+        // and ownership is transferred to the underlying runners one
+        // at a time.
+        std::queue< std::vector<QuerySolution*> > _solutions;
+
+        // Holds the canonicalized subqueries. Ownership is transferred
+        // to the underlying runners one at a time.
+        std::queue<CanonicalQuery*> _cqs;
+
+        // We need this to extract cache-friendly index data from the index assignments.
+        map<BSONObj, size_t> _indexMap;
     };
 
 }  // namespace mongo