SERVER-15886 SubplanStage tries to plan each branch from the cache

Fixes a bug in which the SubplanStage could put a bad plan into the plan cache.

1 files changed, 59 insertions, 5 deletions

diff --git a/src/mongo/dbtests/query_stage_subplan.cpp b/src/mongo/dbtests/query_stage_subplan.cpp
index 18cc6e9b2f9..309e7516952 100644
--- a/src/mongo/dbtests/query_stage_subplan.cpp
+++ b/src/mongo/dbtests/query_stage_subplan.cpp
@@ -82,19 +82,19 @@ namespace QueryStageSubplan {
             BSONObj query = fromjson("{$or: [{a: {$geoWithin: {$centerSphere: [[0,0],10]}}},"
                                             "{a: {$geoWithin: {$centerSphere: [[1,1],10]}}}]}");
 
-            CanonicalQuery* cq;
-            ASSERT_OK(CanonicalQuery::canonicalize(ns(), query, &cq));
-            boost::scoped_ptr<CanonicalQuery> killCq(cq);
+            CanonicalQuery* rawCq;
+            ASSERT_OK(CanonicalQuery::canonicalize(ns(), query, &rawCq));
+            boost::scoped_ptr<CanonicalQuery> cq(rawCq);
 
             Collection* collection = ctx.getCollection();
 
             // Get planner params.
             QueryPlannerParams plannerParams;
-            fillOutPlannerParams(&_txn, collection, cq, &plannerParams);
+            fillOutPlannerParams(&_txn, collection, cq.get(), &plannerParams);
 
             WorkingSet ws;
             boost::scoped_ptr<SubplanStage> subplan(new SubplanStage(&_txn, collection, &ws,
-                                                                     plannerParams, cq));
+                                                                     plannerParams, cq.get()));
 
             // NULL means that 'subplan' will not yield during plan selection. Plan selection
             // should succeed due to falling back on regular planning.
@@ -102,12 +102,66 @@ namespace QueryStageSubplan {
         }
     };
 
+    /**
+     * Test the SubplanStage's ability to plan an individual branch using the plan cache.
+     */
+    class QueryStageSubplanPlanFromCache : public QueryStageSubplanBase {
+    public:
+        void run() {
+            Client::WriteContext ctx(&_txn, ns());
+
+            addIndex(BSON("a" << 1 << "b" << 1));
+            addIndex(BSON("a" << 1 << "c" << 1));
+
+            for (int i = 0; i < 10; i++) {
+                insert(BSON("a" << 1 << "b" << i << "c" << i));
+            }
+
+            // This query should result in a plan cache entry for the first branch. The second
+            // branch should tie, meaning that nothing is inserted into the plan cache.
+            BSONObj query = fromjson("{$or: [{a: 1, b: 3}, {a: 1}]}");
+
+            Collection* collection = ctx.getCollection();
+
+            CanonicalQuery* rawCq;
+            ASSERT_OK(CanonicalQuery::canonicalize(ns(), query, &rawCq));
+            boost::scoped_ptr<CanonicalQuery> cq(rawCq);
+
+            // Get planner params.
+            QueryPlannerParams plannerParams;
+            fillOutPlannerParams(&_txn, collection, cq.get(), &plannerParams);
+
+            WorkingSet ws;
+            boost::scoped_ptr<SubplanStage> subplan(new SubplanStage(&_txn, collection, &ws,
+                                                                     plannerParams, cq.get()));
+
+            // NULL means that 'subplan' should not yield during plan selection.
+            ASSERT_OK(subplan->pickBestPlan(NULL));
+
+            // Nothing is in the cache yet, so neither branch should have been planned from
+            // the plan cache.
+            ASSERT_FALSE(subplan->branchPlannedFromCache(0));
+            ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+
+            // If we repeat the same query, then the first branch should come from the cache,
+            // but the second is re-planned due to tying on the first run.
+            ws.clear();
+            subplan.reset(new SubplanStage(&_txn, collection, &ws, plannerParams, cq.get()));
+
+            ASSERT_OK(subplan->pickBestPlan(NULL));
+
+            ASSERT_TRUE(subplan->branchPlannedFromCache(0));
+            ASSERT_FALSE(subplan->branchPlannedFromCache(1));
+        }
+    };
+
     class All : public Suite {
     public:
         All() : Suite("query_stage_subplan") {}
 
         void setupTests() {
             add<QueryStageSubplanGeo2dOr>();
+            add<QueryStageSubplanPlanFromCache>();
         }
     };