SERVER-16522 listCollections leave cursor open if >1 batch needed

4 files changed, 78 insertions, 34 deletions

diff --git a/src/mongo/SConscript b/src/mongo/SConscript
index 6b8cf504c3c..d4a8ad85c20 100644
--- a/src/mongo/SConscript
+++ b/src/mongo/SConscript
@@ -1105,7 +1105,6 @@ test = env.Install(
                        "s/cluster_ops",
                        "s/cluster_ops_impl",
                        "mocklib",
-                       "db/exec/mock_stage",
                        "$BUILD_DIR/mongo/db/auth/authmocks",
                        "$BUILD_DIR/mongo/db/query/query",
                        "$BUILD_DIR/mongo/db/repl/repl_coordinator_global",
diff --git a/src/mongo/db/commands/list_collections.cpp b/src/mongo/db/commands/list_collections.cpp
index 3115da00e9a..800150706b5 100644
--- a/src/mongo/db/commands/list_collections.cpp
+++ b/src/mongo/db/commands/list_collections.cpp
@@ -31,11 +31,17 @@
 #include "mongo/platform/basic.h"
 
 #include "mongo/db/catalog/collection_catalog_entry.h"
+#include "mongo/db/catalog/cursor_manager.h"
 #include "mongo/db/catalog/database.h"
 #include "mongo/db/catalog/database_catalog_entry.h"
 #include "mongo/db/catalog/database_holder.h"
 #include "mongo/db/client.h"
+#include "mongo/db/clientcursor.h"
 #include "mongo/db/commands.h"
+#include "mongo/db/exec/mock_stage.h"
+#include "mongo/db/exec/working_set.h"
+#include "mongo/db/global_environment_experiment.h"
+#include "mongo/db/query/find_constants.h"
 #include "mongo/db/storage/storage_engine.h"
 
 namespace mongo {
@@ -90,14 +96,10 @@ namespace mongo {
                 matcher.reset( parsed.getValue() );
             }
 
-            // TODO: Handle options specified in the command request object under the "cursor"
-            // field.
+            std::auto_ptr<WorkingSet> ws(new WorkingSet());
+            std::auto_ptr<MockStage> root(new MockStage(ws.get()));
 
-            // TODO: If the full result set does not fit in one batch, allocate a cursor to store
-            // the remainder of the results.
-
-            BSONArrayBuilder arr;
-            for ( list<string>::const_iterator i = names.begin(); i != names.end(); ++i ) {
+            for ( std::list<std::string>::const_iterator i = names.begin(); i != names.end(); ++i ) {
                 string ns = *i;
 
                 StringData collection = nsToCollectionSubstring( ns );
@@ -117,23 +119,64 @@ namespace mongo {
                     continue;
                 }
 
-                arr.append( maybe );
+                WorkingSetID wsId = ws->allocate();
+                WorkingSetMember* member = ws->get(wsId);
+                member->state = WorkingSetMember::OWNED_OBJ;
+                member->keyData.clear();
+                member->loc = RecordId();
+                member->obj = maybe;
+                root->pushBack(*member);
             }
 
-            // TODO: As a temporary compatibility measure for drivers that have not yet been updated
-            // to handle the new cursor response object, we maintain the legacy command format if
-            // the field "cursor" is not present in the request object.  This compatibility layer
-            // should be eventually removed.
-            if ( jsobj["cursor"].type() == mongo::Object ) {
-                const long long cursorId = 0LL;
-                std::string cursorNamespace = str::stream() << dbname << ".$cmd." << name;
-                Command::appendCursorResponseObject( cursorId, cursorNamespace, arr.arr(),
-                                                     &result );
+            std::string cursorNamespace = str::stream() << dbname << ".$cmd." << name;
+
+            PlanExecutor* rawExec;
+            Status makeStatus = PlanExecutor::make(txn,
+                                                   ws.release(),
+                                                   root.release(),
+                                                   cursorNamespace,
+                                                   PlanExecutor::YIELD_MANUAL,
+                                                   &rawExec);
+            std::auto_ptr<PlanExecutor> exec(rawExec);
+            if (!makeStatus.isOK()) {
+                return appendCommandStatus( result, makeStatus );
             }
-            else {
-                result.append( "collections", arr.arr() );
+
+            BSONElement batchSizeElem = jsobj.getFieldDotted("cursor.batchSize");
+            const long long batchSize = batchSizeElem.isNumber()
+                                        ? batchSizeElem.numberLong()
+                                        : -1;
+
+            BSONArrayBuilder firstBatch;
+
+            const int byteLimit = MaxBytesToReturnToClientAtOnce;
+            for (int objCount = 0;
+                 firstBatch.len() < byteLimit && (batchSize == -1 || objCount < batchSize);
+                 objCount++) {
+                BSONObj next;
+                PlanExecutor::ExecState state = exec->getNext(&next, NULL);
+                if ( state == PlanExecutor::IS_EOF ) {
+                    break;
+                }
+                invariant( state == PlanExecutor::ADVANCED );
+                firstBatch.append(next);
             }
 
+            CursorId cursorId = 0LL;
+            if ( !exec->isEOF() ) {
+                exec->saveState();
+                ClientCursor* cursor = new ClientCursor(CursorManager::getGlobalCursorManager(),
+                                                        exec.release());
+                cursorId = cursor->cursorid();
+
+                cursor->setOwnedRecoveryUnit(txn->releaseRecoveryUnit());
+                StorageEngine* storageEngine = getGlobalEnvironment()->getGlobalStorageEngine();
+                txn->setRecoveryUnit(storageEngine->newRecoveryUnit());
+            }
+
+            Command::appendCursorResponseObject( cursorId, cursorNamespace, firstBatch.arr(),
+                                                 &result );
+
             return true;
         }
 
diff --git a/src/mongo/db/exec/SConscript b/src/mongo/db/exec/SConscript
index 0dc22ec7b79..ed4aefd510f 100644
--- a/src/mongo/db/exec/SConscript
+++ b/src/mongo/db/exec/SConscript
@@ -24,16 +24,6 @@ env.CppUnitTest(
 )
 
 env.Library(
-    target = "mock_stage",
-    source = [
-        "mock_stage.cpp",
-    ],
-    LIBDEPS = [
-        "working_set",
-    ],
-)
-
-env.Library(
     target = "scoped_timer",
     source = [
         "scoped_timer.cpp",
@@ -62,6 +52,7 @@ env.Library(
         "keep_mutations.cpp",
         "limit.cpp",
         "merge_sort.cpp",
+        "mock_stage.cpp",
         "multi_iterator.cpp",
         "multi_plan.cpp",
         "near.cpp",
@@ -92,7 +83,6 @@ env.CppUnitTest(
     ],
     LIBDEPS = [
         "exec",
-        "mock_stage",
         "$BUILD_DIR/mongo/serveronly",
         "$BUILD_DIR/mongo/coreserver",
         "$BUILD_DIR/mongo/coredb",
@@ -108,7 +98,6 @@ env.CppUnitTest(
     ],
     LIBDEPS = [
         "exec",
-        "mock_stage",
         "$BUILD_DIR/mongo/serveronly",
         "$BUILD_DIR/mongo/coreserver",
         "$BUILD_DIR/mongo/coredb",
@@ -124,7 +113,6 @@ env.CppUnitTest(
     ],
     LIBDEPS = [
         "exec",
-        "mock_stage",
         "$BUILD_DIR/mongo/serveronly",
         "$BUILD_DIR/mongo/coreserver",
         "$BUILD_DIR/mongo/coredb",
diff --git a/src/mongo/s/commands_public.cpp b/src/mongo/s/commands_public.cpp
index 03b7ea35c0b..6af4251a0bb 100644
--- a/src/mongo/s/commands_public.cpp
+++ b/src/mongo/s/commands_public.cpp
@@ -2703,7 +2703,21 @@ namespace mongo {
                     return false;
                 }
 
-                return passthrough( conf, cmdObj, result );
+                bool retval = passthrough( conf, cmdObj, result );
+
+                BSONObj peekResultObj = result.asTempObj();
+
+                if (peekResultObj["ok"].trueValue() && peekResultObj.hasField("cursor")) {
+                    long long cursorId = peekResultObj["cursor"]["id"].Long();
+                    if (cursorId) {
+                        const string cursorNs = peekResultObj["cursor"]["ns"].String();
+                        cursorCache.storeRef(conf->getPrimary().getConnString(),
+                                             cursorId,
+                                             cursorNs);
+                    }
+                }
+
+                return retval;
             }
         } cmdListCollections;