SERVER-17280 geo near must propagate saveState()/restoreState() to the DensityEstimator

(cherry picked from commit 77dae5d4457c85e64fa92ebd19daa1ee8027dad6)

Conflicts:
	src/mongo/db/exec/geo_near.h
	src/mongo/dbtests/query_stage_near.cpp

5 files changed, 85 insertions, 10 deletions

diff --git a/src/mongo/db/exec/geo_near.cpp b/src/mongo/db/exec/geo_near.cpp
index 34596bc57f4..a472a3bedaf 100644
--- a/src/mongo/db/exec/geo_near.cpp
+++ b/src/mongo/db/exec/geo_near.cpp
@@ -302,6 +302,8 @@ namespace mongo {
                                    Collection* collection,
                                    double* estimatedDistance);
 
+        void saveState();
+        void restoreState(OperationContext* txn);
         void invalidate(OperationContext* txn, const RecordId& dl, InvalidationType type);
 
     private:
@@ -403,6 +405,18 @@ namespace mongo {
         return state;
     }
 
+    void GeoNear2DStage::DensityEstimator::saveState() {
+        if (_indexScan) {
+            _indexScan->saveState();
+        }
+    }
+
+    void GeoNear2DStage::DensityEstimator::restoreState(OperationContext* txn) {
+        if (_indexScan) {
+            _indexScan->restoreState(txn);
+        }
+    }
+
     void GeoNear2DStage::DensityEstimator::invalidate(OperationContext* txn,
                                                       const RecordId& dl,
                                                       InvalidationType type) {
@@ -465,6 +479,18 @@ namespace mongo {
     GeoNear2DStage::~GeoNear2DStage() {
     }
 
+    void GeoNear2DStage::finishSaveState() {
+        if (_densityEstimator) {
+            _densityEstimator->saveState();
+        }
+    }
+
+    void GeoNear2DStage::finishRestoreState(OperationContext* txn) {
+        if (_densityEstimator) {
+            _densityEstimator->restoreState(txn);
+        }
+    }
+
     void GeoNear2DStage::finishInvalidate(OperationContext* txn,
                                           const RecordId& dl,
                                           InvalidationType type) {
@@ -990,6 +1016,8 @@ namespace mongo {
                                    Collection* collection,
                                    double* estimatedDistance);
 
+        void saveState();
+        void restoreState(OperationContext* txn);
         void invalidate(OperationContext* txn, const RecordId& dl, InvalidationType type);
 
     private:
@@ -1090,6 +1118,18 @@ namespace mongo {
         return state;
     }
 
+    void GeoNear2DSphereStage::DensityEstimator::saveState() {
+        if (_indexScan) {
+            _indexScan->saveState();
+        }
+    }
+
+    void GeoNear2DSphereStage::DensityEstimator::restoreState(OperationContext* txn) {
+        if (_indexScan) {
+            _indexScan->restoreState(txn);
+        }
+    }
+
     void GeoNear2DSphereStage::DensityEstimator::invalidate(OperationContext* txn,
                                                             const RecordId& dl,
                                                             InvalidationType type) {
@@ -1128,6 +1168,18 @@ namespace mongo {
         return state;
     }
 
+    void GeoNear2DSphereStage::finishSaveState() {
+        if (_densityEstimator) {
+            _densityEstimator->saveState();
+        }
+    }
+
+    void GeoNear2DSphereStage::finishRestoreState(OperationContext* txn) {
+        if (_densityEstimator) {
+            _densityEstimator->restoreState(txn);
+        }
+    }
+
     void GeoNear2DSphereStage::finishInvalidate(OperationContext* txn,
                                                 const RecordId& dl,
                                                 InvalidationType type) {
diff --git a/src/mongo/db/exec/geo_near.h b/src/mongo/db/exec/geo_near.h
index 04c2b2d2ede..e832063a612 100644
--- a/src/mongo/db/exec/geo_near.h
+++ b/src/mongo/db/exec/geo_near.h
@@ -88,11 +88,16 @@ namespace mongo {
                                                  WorkingSet* workingSet,
                                                  Collection* collection);
 
+    private:
+
+        virtual void finishSaveState();
+
+        virtual void finishRestoreState(OperationContext* txn);
+
         virtual void finishInvalidate(OperationContext* txn,
                                       const RecordId& dl,
                                       InvalidationType type);
 
-    private:
         const GeoNearParams _nearParams;
 
         // The 2D index we're searching over
@@ -138,11 +143,16 @@ namespace mongo {
                                                  WorkingSet* workingSet,
                                                  Collection* collection);
 
+    private:
+
+        virtual void finishSaveState();
+
+        virtual void finishRestoreState(OperationContext* txn);
+
         virtual void finishInvalidate(OperationContext* txn,
                                       const RecordId& dl,
                                       InvalidationType type);
 
-    private:
         const GeoNearParams _nearParams;
 
         // The 2D index we're searching over
diff --git a/src/mongo/db/exec/near.cpp b/src/mongo/db/exec/near.cpp
index 22967ffe793..507f3849fbd 100644
--- a/src/mongo/db/exec/near.cpp
+++ b/src/mongo/db/exec/near.cpp
@@ -329,6 +329,9 @@ namespace mongo {
         for (size_t i = 0; i < _childrenIntervals.size(); i++) {
             _childrenIntervals[i]->covering->saveState();
         }
+
+        // Subclass specific saving, e.g. saving the 2d or 2dsphere density estimator.
+        finishSaveState();
     }
 
     void NearStage::restoreState(OperationContext* opCtx) {
@@ -338,6 +341,9 @@ namespace mongo {
         for (size_t i = 0; i < _childrenIntervals.size(); i++) {
             _childrenIntervals[i]->covering->restoreState(opCtx);
         }
+
+        // Subclass specific restoring, e.g. restoring the 2d or 2dsphere density estimator.
+        finishRestoreState(opCtx);
     }
 
     void NearStage::invalidate(OperationContext* txn, const RecordId& dl, InvalidationType type) {
diff --git a/src/mongo/db/exec/near.h b/src/mongo/db/exec/near.h
index 35acb01ac06..cc6d170f380 100644
--- a/src/mongo/db/exec/near.h
+++ b/src/mongo/db/exec/near.h
@@ -146,15 +146,20 @@ namespace mongo {
                                       WorkingSet* workingSet,
                                       Collection* collection);
 
-        /**
-         * Does any invalidation work specific to the search type.
-         */
+    private:
+
+        //
+        // Save/restore/invalidate work specific to the search type.
+        //
+
+        virtual void finishSaveState() = 0;
+
+        virtual void finishRestoreState(OperationContext* txn) = 0;
+
         virtual void finishInvalidate(OperationContext* txn,
                                       const RecordId& dl,
                                       InvalidationType type) = 0;
 
-    private:
-
         //
         // Generic methods for progressive search functionality
         //
diff --git a/src/mongo/dbtests/query_stage_near.cpp b/src/mongo/dbtests/query_stage_near.cpp
index 56f5e4003d4..f5b9bd9f434 100644
--- a/src/mongo/dbtests/query_stage_near.cpp
+++ b/src/mongo/dbtests/query_stage_near.cpp
@@ -176,11 +176,13 @@ namespace {
             return StatusWith<double>(member->obj.value()["distance"].numberDouble());
         }
 
+        virtual void finishSaveState() { }
+
+        virtual void finishRestoreState(OperationContext* txn) { }
+
         virtual void finishInvalidate(OperationContext* txn,
                                       const RecordId& dl,
-                                      InvalidationType type) {
-            invariant(!"MockNearStage should not receive invalidations");
-        }
+                                      InvalidationType type) { }
 
     private: