SERVER-15450 Avoid getStats() and keep interval stages as children in NearStage

1 files changed, 20 insertions, 18 deletions

diff --git a/src/mongo/db/exec/near.cpp b/src/mongo/db/exec/near.cpp
index 53304381eb1..c568b21414f 100644
--- a/src/mongo/db/exec/near.cpp
+++ b/src/mongo/db/exec/near.cpp
@@ -43,7 +43,8 @@ namespace mongo {
           _workingSet(workingSet),
           _collection(collection),
           _searchState(SearchState_Buffering),
-          _stats(stats) {
+          _stats(stats),
+          _nextInterval(NULL) {
 
         // Ensure we have specific distance search stats unless a child class specified their
         // own distance stats subclass
@@ -151,7 +152,9 @@ namespace mongo {
                 return PlanStage::IS_EOF;
             }
 
-            _nextInterval.reset(intervalStatus.getValue());
+            // CoveredInterval and its child stage are owned by _childrenIntervals
+            _childrenIntervals.push_back(intervalStatus.getValue());
+            _nextInterval = _childrenIntervals.back();
             _nextIntervalStats.reset(new IntervalStats());
         }
 
@@ -159,9 +162,7 @@ namespace mongo {
         PlanStage::StageState intervalState = _nextInterval->covering->work(&nextMemberID);
 
         if (PlanStage::IS_EOF == intervalState) {
-
-            _stats->children.push_back(_nextInterval->covering->getStats());
-            _nextInterval.reset();
+            _nextInterval = NULL;
             _nextIntervalSeen.clear();
 
             _searchState = SearchState_Advancing;
@@ -267,23 +268,23 @@ namespace mongo {
 
     void NearStage::saveState() {
         ++_stats->common.yields;
-        if (_nextInterval) {
-            _nextInterval->covering->saveState();
+        for (size_t i = 0; i < _childrenIntervals.size(); i++) {
+            _childrenIntervals[i]->covering->saveState();
         }
     }
 
     void NearStage::restoreState(OperationContext* opCtx) {
         _txn = opCtx;
         ++_stats->common.unyields;
-        if (_nextInterval) {
-            _nextInterval->covering->restoreState(opCtx);
+        for (size_t i = 0; i < _childrenIntervals.size(); i++) {
+            _childrenIntervals[i]->covering->restoreState(opCtx);
         }
     }
 
     void NearStage::invalidate(const DiskLoc& dl, InvalidationType type) {
         ++_stats->common.invalidates;
-        if (_nextInterval) {
-            _nextInterval->covering->invalidate(dl, type);
+        for (size_t i = 0; i < _childrenIntervals.size(); i++) {
+            _childrenIntervals[i]->covering->invalidate(dl, type);
         }
 
         // If a result is in _resultBuffer and has a DiskLoc it will be in _nextIntervalSeen as
@@ -304,17 +305,18 @@ namespace mongo {
     }
 
     vector<PlanStage*> NearStage::getChildren() const {
-        // TODO: Keep around all our interval stages and report as children
-        return vector<PlanStage*>();
+        vector<PlanStage*> children;
+        for (size_t i = 0; i < _childrenIntervals.size(); i++) {
+            children.push_back(_childrenIntervals[i]->covering.get());
+        }
+        return children;
     }
 
     PlanStageStats* NearStage::getStats() {
         PlanStageStats* statsClone = _stats->clone();
-
-        // If we have an active interval, add those child stats as well
-        if (_nextInterval)
-            statsClone->children.push_back(_nextInterval->covering->getStats());
-
+        for (size_t i = 0; i < _childrenIntervals.size(); ++i) {
+            statsClone->children.push_back(_childrenIntervals[i]->covering->getStats());
+        }
         return statsClone;
     }