SERVER-15541 preserve chunk ranges during long-running yielding count/geoNear

3 files changed, 24 insertions, 7 deletions

diff --git a/src/mongo/db/commands/count.cpp b/src/mongo/db/commands/count.cpp
index 8687a56c184..38df9b45286 100644
--- a/src/mongo/db/commands/count.cpp
+++ b/src/mongo/db/commands/count.cpp
@@ -38,6 +38,7 @@
 #include "mongo/db/exec/count.h"
 #include "mongo/db/query/get_executor.h"
 #include "mongo/db/query/explain.h"
+#include "mongo/db/range_preserver.h"
 #include "mongo/db/repl/repl_coordinator_global.h"
 #include "mongo/util/log.h"
 
@@ -82,6 +83,10 @@ namespace mongo {
             AutoGetCollectionForRead ctx(txn, request.ns);
             Collection* collection = ctx.getCollection();
 
+            // Prevent chunks from being cleaned up during yields - this allows us to only check the
+            // version on initial entry into count.
+            RangePreserver preserver(collection);
+
             PlanExecutor* rawExec;
             Status getExecStatus = getExecutorCount(txn,
                                                     collection,
@@ -114,6 +119,10 @@ namespace mongo {
             AutoGetCollectionForRead ctx(txn, request.ns);
             Collection* collection = ctx.getCollection();
 
+            // Prevent chunks from being cleaned up during yields - this allows us to only check the
+            // version on initial entry into count.
+            RangePreserver preserver(collection);
+
             PlanExecutor* rawExec;
             Status getExecStatus = getExecutorCount(txn,
                                                     collection,
diff --git a/src/mongo/db/commands/geo_near_cmd.cpp b/src/mongo/db/commands/geo_near_cmd.cpp
index cdf8cc9724a..59b2ad8a02b 100644
--- a/src/mongo/db/commands/geo_near_cmd.cpp
+++ b/src/mongo/db/commands/geo_near_cmd.cpp
@@ -33,6 +33,7 @@
 #include "mongo/db/auth/action_set.h"
 #include "mongo/db/auth/action_type.h"
 #include "mongo/db/auth/privilege.h"
+#include "mongo/db/catalog/collection.h"
 #include "mongo/db/catalog/database.h"
 #include "mongo/db/client.h"
 #include "mongo/db/commands.h"
@@ -45,7 +46,7 @@
 #include "mongo/db/jsobj.h"
 #include "mongo/db/query/get_executor.h"
 #include "mongo/db/query/explain.h"
-#include "mongo/db/catalog/collection.h"
+#include "mongo/db/range_preserver.h"
 #include "mongo/platform/unordered_map.h"
 #include "mongo/util/log.h"
 
@@ -181,6 +182,10 @@ namespace mongo {
                 return false;
             }
 
+            // Prevent chunks from being cleaned up during yields - this allows us to only check the
+            // version on initial entry into geoNear.
+            RangePreserver preserver(collection);
+
             PlanExecutor* rawExec;
             if (!getExecutor(txn, collection, cq, PlanExecutor::YIELD_AUTO, &rawExec, 0).isOK()) {
                 errmsg = "can't get query executor";
diff --git a/src/mongo/db/range_preserver.h b/src/mongo/db/range_preserver.h
index dac32c30fd6..9dc40844a9c 100644
--- a/src/mongo/db/range_preserver.h
+++ b/src/mongo/db/range_preserver.h
@@ -49,16 +49,19 @@ namespace mongo {
          * deleted until this object goes out of scope.
          */
         RangePreserver(const Collection* collection) {
-            invariant( collection );
-            // Not a memory leak.  Cached in a static structure by CC's ctor.
-            ClientCursor* cc = new ClientCursor(collection);
+            // Empty collections don't have any data we need to preserve
+            if (collection) {
+                // Not a memory leak.  Cached in a static structure by CC's ctor.
+                ClientCursor* cc = new ClientCursor(collection);
 
-            // Pin keeps the CC from being deleted while it's in scope.  We delete it ourselves.
-            _pin.reset(new ClientCursorPin(collection, cc->cursorid()));
+                // Pin keeps the CC from being deleted while it's in scope.  We delete it ourselves.
+                _pin.reset(new ClientCursorPin(collection, cc->cursorid()));
+            }
         }
 
         ~RangePreserver() {
-            _pin->deleteUnderlying();
+            if (_pin)
+                _pin->deleteUnderlying();
         }
 
     private: