SERVER-16523 listIndexes leave cursor open if >1 batch needed

2 files changed, 79 insertions, 16 deletions

diff --git a/src/mongo/db/commands/list_indexes.cpp b/src/mongo/db/commands/list_indexes.cpp
index 3403ede8b1e..3bd6d639604 100644
--- a/src/mongo/db/commands/list_indexes.cpp
+++ b/src/mongo/db/commands/list_indexes.cpp
@@ -32,9 +32,15 @@
 
 #include "mongo/db/catalog/collection.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/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 {
@@ -112,27 +118,70 @@ namespace mongo {
             vector<string> indexNames;
             cce->getAllIndexes( txn, &indexNames );
 
-            // 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 ( size_t i = 0; i < indexNames.size(); i++ ) {
-                arr.append( cce->getIndexSpec( txn, indexNames[i] ) );
+                BSONObj indexSpec = cce->getIndexSpec( txn, indexNames[i] );
+
+                WorkingSetID wsId = ws->allocate();
+                WorkingSetMember* member = ws->get(wsId);
+                member->state = WorkingSetMember::OWNED_OBJ;
+                member->keyData.clear();
+                member->loc = RecordId();
+                member->obj = indexSpec;
+                root->pushBack(*member);
+            }
+
+            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 );
             }
 
-            if ( cmdObj["cursor"].type() == mongo::Object ) {
-                const long long cursorId = 0LL;
-                std::string cursorNamespace = str::stream()
-                    << dbname << ".$cmd." << name << "." << ns.coll();
-                Command::appendCursorResponseObject( cursorId, cursorNamespace, arr.arr(),
-                                                     &result );
-            } else {
-                result.append( "indexes", arr.arr() );
+            BSONElement batchSizeElem = cmdObj.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/s/commands_public.cpp b/src/mongo/s/commands_public.cpp
index 6af4251a0bb..98606f7b91c 100644
--- a/src/mongo/s/commands_public.cpp
+++ b/src/mongo/s/commands_public.cpp
@@ -2745,7 +2745,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;
             }
         } cmdListIndexes;