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

5 files changed, 87 insertions, 10 deletions

diff --git a/src/mongo/db/exec/geo_near.cpp b/src/mongo/db/exec/geo_near.cpp
index c5367706062..9026afc7239 100644
--- a/src/mongo/db/exec/geo_near.cpp
+++ b/src/mongo/db/exec/geo_near.cpp
@@ -303,6 +303,8 @@ namespace mongo {
                                    WorkingSetID* out,
                                    double* estimatedDistance);
 
+        void saveState();
+        void restoreState(OperationContext* txn);
         void invalidate(OperationContext* txn, const RecordId& dl, InvalidationType type);
 
     private:
@@ -405,6 +407,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) {
@@ -469,6 +483,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) {
@@ -995,6 +1021,8 @@ namespace mongo {
                                    WorkingSetID* out,
                                    double* estimatedDistance);
 
+        void saveState();
+        void restoreState(OperationContext* txn);
         void invalidate(OperationContext* txn, const RecordId& dl, InvalidationType type);
 
     private:
@@ -1096,6 +1124,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) {
@@ -1136,6 +1176,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 6bb449b10ad..af7352cbe47 100644
--- a/src/mongo/db/exec/geo_near.h
+++ b/src/mongo/db/exec/geo_near.h
@@ -89,10 +89,16 @@ namespace mongo {
                                                  Collection* collection,
                                                  WorkingSetID* out);
 
+    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
@@ -139,10 +145,16 @@ namespace mongo {
                                                  Collection* collection,
                                                  WorkingSetID* out);
 
+    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 46b345ad5ca..cb1abf498a5 100644
--- a/src/mongo/db/exec/near.cpp
+++ b/src/mongo/db/exec/near.cpp
@@ -322,6 +322,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) {
@@ -331,6 +334,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 60384679c00..d7d27a78331 100644
--- a/src/mongo/db/exec/near.h
+++ b/src/mongo/db/exec/near.h
@@ -147,15 +147,20 @@ namespace mongo {
                                       Collection* collection,
                                       WorkingSetID* out) = 0;
 
-        /**
-         * 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 4063c911f64..bb07fce894e 100644
--- a/src/mongo/dbtests/query_stage_near.cpp
+++ b/src/mongo/dbtests/query_stage_near.cpp
@@ -183,11 +183,13 @@ namespace {
             return IS_EOF;
         }
 
+        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: