SERVER-10026 add maxScan support for new query

12 files changed, 54 insertions, 13 deletions

diff --git a/jstests/maxscan.js b/jstests/maxscan.js
index c455efbf3ea..3d15b26f638 100644
--- a/jstests/maxscan.js
+++ b/jstests/maxscan.js
@@ -12,3 +12,7 @@ assert.eq( 50 , t.find()._addSpecial( "$maxScan" , 50 ).itcount() , "B" )
 
 assert.eq( 10 , t.find( { x : 2 } ).itcount() , "C" )
 assert.eq( 5 , t.find( { x : 2 } )._addSpecial( "$maxScan" , 50 ).itcount() , "D" )
+
+t.ensureIndex({x: 1});
+assert.eq( 10, t.find( { x : 2 } ).hint({x:1})._addSpecial( "$maxScan" , N ).itcount() , "E" )
+assert.eq( 0, t.find( { x : 2 } ).hint({x:1})._addSpecial( "$maxScan" , 1 ).itcount() , "E" )
diff --git a/src/mongo/db/exec/collection_scan.cpp b/src/mongo/db/exec/collection_scan.cpp
index 6bddeb67e02..f7785ef5423 100644
--- a/src/mongo/db/exec/collection_scan.cpp
+++ b/src/mongo/db/exec/collection_scan.cpp
@@ -43,7 +43,10 @@ namespace mongo {
     CollectionScan::CollectionScan(const CollectionScanParams& params,
                                    WorkingSet* workingSet,
                                    const MatchExpression* filter)
-        : _workingSet(workingSet), _filter(filter), _params(params), _nsDropped(false) { }
+        : _workingSet(workingSet),
+          _filter(filter),
+          _params(params),
+          _nsDropped(false) { }
 
     PlanStage::StageState CollectionScan::work(WorkingSetID* out) {
         ++_commonStats.works;
@@ -105,6 +108,9 @@ namespace mongo {
     }
 
     bool CollectionScan::isEOF() {
+        if ((0 != _params.maxScan) && (_specificStats.docsTested >= _params.maxScan)) {
+            return true;
+        }
         if (_nsDropped) { return true; }
         if (NULL == _iter) { return false; }
         return _iter->isEOF();
diff --git a/src/mongo/db/exec/collection_scan.h b/src/mongo/db/exec/collection_scan.h
index b9075a3ad97..44a0300c8c1 100644
--- a/src/mongo/db/exec/collection_scan.h
+++ b/src/mongo/db/exec/collection_scan.h
@@ -46,7 +46,8 @@ namespace mongo {
      */
     class CollectionScan : public PlanStage {
     public:
-        CollectionScan(const CollectionScanParams& params, WorkingSet* workingSet,
+        CollectionScan(const CollectionScanParams& params,
+                       WorkingSet* workingSet,
                        const MatchExpression* filter);
 
         virtual StageState work(WorkingSetID* out);
diff --git a/src/mongo/db/exec/collection_scan_common.h b/src/mongo/db/exec/collection_scan_common.h
index 880e7b56fc7..9c88e7653f6 100644
--- a/src/mongo/db/exec/collection_scan_common.h
+++ b/src/mongo/db/exec/collection_scan_common.h
@@ -40,7 +40,8 @@ namespace mongo {
 
         CollectionScanParams() : start(DiskLoc()),
                                  direction(FORWARD),
-                                 tailable(false) { }
+                                 tailable(false),
+                                 maxScan(0) { }
 
         // What collection?
         string ns;
@@ -53,6 +54,9 @@ namespace mongo {
 
         // Do we want the scan to be 'tailable'?  Only meaningful if the collection is capped.
         bool tailable;
+
+        // If non-zero, how many documents will we look at?
+        size_t maxScan;
     };
 
 }  // namespace mongo
diff --git a/src/mongo/db/exec/index_scan.cpp b/src/mongo/db/exec/index_scan.cpp
index d0d4d0a1244..333e1845503 100644
--- a/src/mongo/db/exec/index_scan.cpp
+++ b/src/mongo/db/exec/index_scan.cpp
@@ -204,6 +204,13 @@ namespace mongo {
             return false;
         }
 
+        // If there's a limit on how many keys we can scan, we may be EOF when we hit that.
+        if (0 != _params.maxScan) {
+            if (_specificStats.keysExamined >= _params.maxScan) {
+                return true;
+            }
+        }
+
         return _hitEnd || _indexCursor->isEOF();
     }
 
diff --git a/src/mongo/db/exec/index_scan.h b/src/mongo/db/exec/index_scan.h
index 891a97c6149..8bbe7111f38 100644
--- a/src/mongo/db/exec/index_scan.h
+++ b/src/mongo/db/exec/index_scan.h
@@ -45,8 +45,12 @@ namespace mongo {
     class WorkingSet;
 
     struct IndexScanParams {
-        IndexScanParams() : descriptor(NULL), direction(1), limit(0),
-                            forceBtreeAccessMethod(false), doNotDedup(false) { }
+        IndexScanParams() : descriptor(NULL),
+                            direction(1),
+                            limit(0),
+                            forceBtreeAccessMethod(false),
+                            doNotDedup(false),
+                            maxScan(0) { }
 
         IndexDescriptor* descriptor;
 
@@ -61,6 +65,9 @@ namespace mongo {
         bool forceBtreeAccessMethod;
 
         bool doNotDedup;
+
+        // How many keys will we look at?
+        size_t maxScan;
     };
 
     /**
diff --git a/src/mongo/db/query/new_find.cpp b/src/mongo/db/query/new_find.cpp
index 097e7bb59c6..0e79b3fc949 100644
--- a/src/mongo/db/query/new_find.cpp
+++ b/src/mongo/db/query/new_find.cpp
@@ -115,8 +115,7 @@ namespace mongo {
         const LiteParsedQuery& pq = cq->getParsed();
 
         // We fail to deal well with obscure arguments to .find().
-        if (pq.returnKey() || pq.showDiskLoc() || (0 != pq.getMaxScan()) || !pq.getMin().isEmpty()
-            || !pq.getMax().isEmpty()) {
+        if (pq.returnKey() || pq.showDiskLoc() || !pq.getMin().isEmpty() || !pq.getMax().isEmpty()) {
             QLOG() << "rejecting wacky query args query\n";
             return false;
         }
diff --git a/src/mongo/db/query/planner_access.cpp b/src/mongo/db/query/planner_access.cpp
index bb63b5106ef..132afb7de64 100644
--- a/src/mongo/db/query/planner_access.cpp
+++ b/src/mongo/db/query/planner_access.cpp
@@ -51,6 +51,7 @@ namespace mongo {
         csn->name = query.ns();
         csn->filter.reset(query.root()->shallowClone());
         csn->tailable = tailable;
+        csn->maxScan = query.getParsed().getMaxScan();
 
         // If the sort is {$natural: +-1} this changes the direction of the collection scan.
         const BSONObj& sortObj = query.getParsed().getSort();
@@ -74,7 +75,8 @@ namespace mongo {
     }
 
     // static
-    QuerySolutionNode* QueryPlannerAccess::makeLeafNode(const IndexEntry& index,
+    QuerySolutionNode* QueryPlannerAccess::makeLeafNode(const CanonicalQuery& query,
+                                                        const IndexEntry& index,
                                                         MatchExpression* expr,
                                                         IndexBoundsBuilder::BoundsTightness* tightnessOut) {
         // QLOG() << "making leaf node for " << expr->toString() << endl;
@@ -133,6 +135,7 @@ namespace mongo {
             isn->indexKeyPattern = index.keyPattern;
             isn->indexIsMultiKey = index.multikey;
             isn->bounds.fields.resize(index.keyPattern.nFields());
+            isn->maxScan = query.getParsed().getMaxScan();
 
             IndexBoundsBuilder::translate(expr, index.keyPattern.firstElement(),
                                           &isn->bounds.fields[0], tightnessOut);
@@ -474,7 +477,7 @@ namespace mongo {
                 currentIndexNumber = ixtag->index;
 
                 IndexBoundsBuilder::BoundsTightness tightness = IndexBoundsBuilder::INEXACT_FETCH;
-                currentScan.reset(makeLeafNode(indices[currentIndexNumber],
+                currentScan.reset(makeLeafNode(query, indices[currentIndexNumber],
                                                 child, &tightness));
 
                 if (tightness == IndexBoundsBuilder::EXACT && !inArrayOperator) {
@@ -727,7 +730,7 @@ namespace mongo {
                 IndexTag* tag = static_cast<IndexTag*>(root->getTag());
 
                 IndexBoundsBuilder::BoundsTightness tightness = IndexBoundsBuilder::EXACT;
-                QuerySolutionNode* soln = makeLeafNode(indices[tag->index], root,
+                QuerySolutionNode* soln = makeLeafNode(query, indices[tag->index], root,
                                                        &tightness);
                 verify(NULL != soln);
                 stringstream ss;
@@ -827,6 +830,7 @@ namespace mongo {
         isn->indexKeyPattern = index.keyPattern;
         isn->indexIsMultiKey = index.multikey;
         isn->bounds.fields.resize(index.keyPattern.nFields());
+        isn->maxScan = query.getParsed().getMaxScan();
 
         // TODO: can we use simple bounds with this compound idx?
         BSONObjIterator it(isn->indexKeyPattern);
diff --git a/src/mongo/db/query/planner_access.h b/src/mongo/db/query/planner_access.h
index d4d490f5154..57715fa2637 100644
--- a/src/mongo/db/query/planner_access.h
+++ b/src/mongo/db/query/planner_access.h
@@ -127,7 +127,8 @@ namespace mongo {
          * If the node is a geo node, grab the geo data from 'expr' and stuff it into the
          * geo solution node of the appropriate type.
          */
-        static QuerySolutionNode* makeLeafNode(const IndexEntry& index,
+        static QuerySolutionNode* makeLeafNode(const CanonicalQuery& query,
+                                               const IndexEntry& index,
                                                MatchExpression* expr,
                                                IndexBoundsBuilder::BoundsTightness* tightnessOut);
 
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index 9723654a650..af3f08ca4ef 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -77,7 +77,7 @@ namespace mongo {
     // CollectionScanNode
     //
 
-    CollectionScanNode::CollectionScanNode() : tailable(false), direction(1) { }
+    CollectionScanNode::CollectionScanNode() : tailable(false), direction(1), maxScan(0) { }
 
     void CollectionScanNode::appendToString(stringstream* ss, int indent) const {
         addIndent(ss, indent);
@@ -307,7 +307,7 @@ namespace mongo {
     //
 
     IndexScanNode::IndexScanNode()
-        : indexIsMultiKey(false), limit(0), direction(1) { }
+        : indexIsMultiKey(false), limit(0), direction(1), maxScan(0) { }
 
     void IndexScanNode::appendToString(stringstream* ss, int indent) const {
         addIndent(ss, indent);
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index bbc1da0d886..a3fdceabbd8 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -227,6 +227,9 @@ namespace mongo {
         bool tailable;
 
         int direction;
+
+        // maxScan option to .find() limits how many docs we look at.
+        int maxScan;
     };
 
     struct AndHashNode : public QuerySolutionNode {
@@ -352,6 +355,9 @@ namespace mongo {
 
         int direction;
 
+        // maxScan option to .find() limits how many docs we look at.
+        int maxScan;
+
         // BIG NOTE:
         // If you use simple bounds, we'll use whatever index access method the keypattern implies.
         // If you use the complex bounds, we force Btree access.
diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp
index 5983fb94052..2332ac7de04 100644
--- a/src/mongo/db/query/stage_builder.cpp
+++ b/src/mongo/db/query/stage_builder.cpp
@@ -58,6 +58,7 @@ namespace mongo {
             params.tailable = csn->tailable;
             params.direction = (csn->direction == 1) ? CollectionScanParams::FORWARD
                                                      : CollectionScanParams::BACKWARD;
+            params.maxScan = csn->maxScan;
             return new CollectionScan(params, ws, csn->filter.get());
         }
         else if (STAGE_IXSCAN == root->getType()) {
@@ -87,6 +88,7 @@ namespace mongo {
             params.bounds = ixn->bounds;
             params.direction = ixn->direction;
             params.limit = ixn->limit;
+            params.maxScan = ixn->maxScan;
             return new IndexScan(params, ws, ixn->filter.get());
         }
         else if (STAGE_FETCH == root->getType()) {