SERVER-16437 IndexScan optimize end checker for single interval scans

3 files changed, 128 insertions, 25 deletions

diff --git a/src/mongo/db/exec/index_scan.cpp b/src/mongo/db/exec/index_scan.cpp
index 2a8ab64305b..66a75d668b9 100644
--- a/src/mongo/db/exec/index_scan.cpp
+++ b/src/mongo/db/exec/index_scan.cpp
@@ -36,6 +36,7 @@
 #include "mongo/db/index/index_access_method.h"
 #include "mongo/db/index/index_cursor.h"
 #include "mongo/db/index/index_descriptor.h"
+#include "mongo/db/query/index_bounds_builder.h"
 #include "mongo/util/log.h"
 
 namespace {
@@ -64,8 +65,11 @@ namespace mongo {
           _filter(filter),
           _shouldDedup(true),
           _params(params),
+          _commonStats(kStageType),
           _btreeCursor(NULL),
-          _commonStats(kStageType) {
+          _keyEltsToUse(0),
+          _movePastKeyElts(false),
+          _endKeyInclusive(false) {
         _iam = _params.descriptor->getIndexCatalog()->getIndex(_params.descriptor);
         _keyPattern = _params.descriptor->keyPattern().getOwned();
 
@@ -116,24 +120,51 @@ namespace mongo {
             }
         }
         else {
-            // "Fast" Btree-specific navigation.
             _btreeCursor = static_cast<BtreeIndexCursor*>(_indexCursor.get());
-            _checker.reset(new IndexBoundsChecker(&_params.bounds,
-                                                  _keyPattern,
-                                                  _params.direction));
-
-            int nFields = _keyPattern.nFields();
-            vector<const BSONElement*> key;
-            vector<bool> inc;
-            key.resize(nFields);
-            inc.resize(nFields);
-            if (_checker->getStartKey(&key, &inc)) {
-                _btreeCursor->seek(key, inc);
-                _keyElts.resize(nFields);
-                _keyEltsInc.resize(nFields);
+
+            // For single intervals, we can use an optimized scan which checks against the position
+            // of an end cursor.  For all other index scans, we fall back on using
+            // IndexBoundsChecker to determine when we've finished the scan.
+            BSONObj startKey;
+            bool startKeyInclusive;
+            if (IndexBoundsBuilder::isSingleInterval(_params.bounds,
+                                                     &startKey,
+                                                     &startKeyInclusive,
+                                                     &_endKey,
+                                                     &_endKeyInclusive)) {
+                // We want to point at the start key if it's inclusive, and we want to point past
+                // the start key if it's exclusive.
+                _btreeCursor->seek(startKey, !startKeyInclusive);
+
+                IndexCursor* endCursor;
+                invariant(_iam->newCursor(_txn, cursorOptions, &endCursor).isOK());
+                invariant(endCursor);
+
+                // TODO: Is it a valid assumption that we can always make this cast safely?
+                // See SERVER-12397.
+                _endCursor.reset(static_cast<BtreeIndexCursor*>(endCursor));
+
+                // If the end key is inclusive, we want to point *past* it since that's the end.
+                _endCursor->seek(_endKey, _endKeyInclusive);
             }
             else {
-                _scanState = HIT_END;
+                _checker.reset(new IndexBoundsChecker(&_params.bounds,
+                                                      _keyPattern,
+                                                      _params.direction));
+
+                int nFields = _keyPattern.nFields();
+                vector<const BSONElement*> key;
+                vector<bool> inc;
+                key.resize(nFields);
+                inc.resize(nFields);
+                if (_checker->getStartKey(&key, &inc)) {
+                    _btreeCursor->seek(key, inc);
+                    _keyElts.resize(nFields);
+                    _keyEltsInc.resize(nFields);
+                }
+                else {
+                    _scanState = HIT_END;
+                }
             }
         }
 
@@ -251,6 +282,10 @@ namespace mongo {
             _savedLoc = _indexCursor->getValue();
         }
         _indexCursor->savePosition();
+
+        if (_endCursor) {
+            _endCursor->savePosition();
+        }
     }
 
     void IndexScan::restoreState(OperationContext* opCtx) {
@@ -267,6 +302,24 @@ namespace mongo {
             return;
         }
 
+        if (_endCursor) {
+            // Single interval case.
+            if (!_endCursor->restorePosition(opCtx).isOK()) {
+                _scanState = HIT_END;
+                return;
+            }
+
+            // If we were EOF when we yielded, we don't always want to have '_btreeCursor' run until
+            // EOF. New documents may have been inserted after our end key, and our end marker may
+            // be before them.
+            //
+            // As an example, say we're counting from 5 to 10 and the index only has keys for 6, 7,
+            // 8, and 9. '_btreeCursor' will point at key 6 at the start and '_endCursor' will be
+            // EOF. If we insert documents with keys 11 during a yield, we need to relocate
+            // '_endCursor' to point at them as the end key of our scan.
+            _endCursor->seek(_endKey, _endKeyInclusive);
+        }
+
         if (!_savedKey.binaryEqual(_indexCursor->getKey())
             || _savedLoc != _indexCursor->getValue()) {
             // Our restored position isn't the same as the saved position.  When we call work()
@@ -322,6 +375,20 @@ namespace mongo {
                 ++_specificStats.keysExamined;
             }
         }
+        else if (_endCursor) {
+            // We're in the single interval case, and we have a cursor pointing to the end position.
+            // We can check whether the scan is over by seeing if our cursor points at the same
+            // thing as the end cursor.
+            _scanState = GETTING_NEXT;
+            invariant(!_checker);
+
+            if (_endCursor->pointsAt(*_btreeCursor)) {
+                _scanState = HIT_END;
+            }
+            else {
+                ++_specificStats.keysExamined;
+            }
+        }
         else {
             verify(NULL != _btreeCursor);
             verify(NULL != _checker.get());
diff --git a/src/mongo/db/exec/index_scan.h b/src/mongo/db/exec/index_scan.h
index 54481c77ff1..33a6d8e7d21 100644
--- a/src/mongo/db/exec/index_scan.h
+++ b/src/mongo/db/exec/index_scan.h
@@ -162,17 +162,55 @@ namespace mongo {
 
         IndexScanParams _params;
 
-        // For our "fast" Btree-only navigation AKA the index bounds optimization.
-        scoped_ptr<IndexBoundsChecker> _checker;
+        // Stats
+        CommonStats _commonStats;
+        IndexScanStats _specificStats;
+
+        //
+        // Btree-specific navigation state.
+        //
+
+        // When _params.bounds.isSimpleRange is false, we assume that we are traversing a btree
+        // index and make use of the BtreeIndexCursor methods for navigation.  We do this because
+        // using the btree-specific navigation methods will yield better index scan performance.
+        //
+        // This assumption will need to change when we introduce new index types not based on the
+        // btree index.  This is being tracked in SERVER-12397.
         BtreeIndexCursor* _btreeCursor;
+
+        //
+        // If we have decided to use the BtreeIndexCursor methods for navigation, we make a decision
+        // to employ one of two different algorithms for determining when the index scan has reached
+        // the end:
+        //
+
+        //
+        // 1) If the index scan is not a single interval, then we use an IndexBoundsChecker to
+        //    determine when the index scan has reached the end.  In this case, _checker will be
+        //    non-NULL (and _endCursor will be NULL).
+        //
+
+        scoped_ptr<IndexBoundsChecker> _checker;
         int _keyEltsToUse;
         bool _movePastKeyElts;
         std::vector<const BSONElement*> _keyElts;
         std::vector<bool> _keyEltsInc;
 
-        // Stats
-        CommonStats _commonStats;
-        IndexScanStats _specificStats;
+        //
+        // 2) If the index scan is a single interval, then the scan can execute faster by
+        //    checking for the end via comparison against an end cursor, rather than repeatedly
+        //    doing BSON compares against scanned keys.  In this case, _endCursor will be non-NULL
+        //    (and _checker will be NULL).
+        //
+
+        // The end cursor.
+        boost::scoped_ptr<BtreeIndexCursor> _endCursor;
+
+        // The key that the end cursor should point to.
+        BSONObj _endKey;
+
+        // Is the end key included in the range?
+        bool _endKeyInclusive;
     };
 
 }  // namespace mongo
diff --git a/src/mongo/db/index/btree_index_cursor.cpp b/src/mongo/db/index/btree_index_cursor.cpp
index ea305a6e417..177e3d91253 100644
--- a/src/mongo/db/index/btree_index_cursor.cpp
+++ b/src/mongo/db/index/btree_index_cursor.cpp
@@ -74,10 +74,8 @@ namespace mongo {
     }
 
     void BtreeIndexCursor::seek(const BSONObj& position, bool afterKey) {
-        // XXX This used a hard-coded direction of 1 and is only correct in the forward direction.
-        invariant(_cursor->getDirection() == 1);
-        _cursor->locate(position, 
-                        afterKey ? RecordId::max() : RecordId::min());
+        const bool forward = (1 == _cursor->getDirection());
+        _cursor->locate(position, (afterKey == forward) ? RecordId::max() : RecordId::min());
     }
 
     bool BtreeIndexCursor::pointsAt(const BtreeIndexCursor& other) {