SERVER-4150 add exhaust support

2 files changed, 32 insertions, 1 deletions

diff --git a/src/mongo/db/ops/query.cpp b/src/mongo/db/ops/query.cpp
index 404978dfdf2..9c1fae4133a 100644
--- a/src/mongo/db/ops/query.cpp
+++ b/src/mongo/db/ops/query.cpp
@@ -1196,6 +1196,7 @@ namespace mongo {
             {
                 QueryResponseBuilder queryResponseBuilder( pq, cursor, oldPlan );
                 long long cursorid = 0;
+                const char * exhaust = 0;
                 OpTime slaveReadTill;
                 ClientCursor::CleanupPointer ccPointer;
                 ccPointer.reset( new ClientCursor( QueryOption_NoCursorTimeout, cursor, ns ) );
@@ -1276,6 +1277,10 @@ namespace mongo {
                         DEV tlog() << "query has no more but tailable, cursorid: "
                             << cursorid << endl;
                     }
+                    if( queryOptions & QueryOption_Exhaust ) {
+                        exhaust = ns;
+                        curop.debug().exhaust = true;
+                    }
                     ccPointer.release();
                     // undo unlimited timeout
                 }
@@ -1289,7 +1294,7 @@ namespace mongo {
                 qr->setOperation(opReply);
                 qr->startingFrom = 0;
                 qr->nReturned = nReturned;
-                return 0;
+                return exhaust;
             }
         } catch ( const UserException &u ) {
             if ( retries != 0 || u.getCode() != ScanAndOrderMemoryLimitExceededAssertionCode ) {
diff --git a/src/mongo/dbtests/querytests.cpp b/src/mongo/dbtests/querytests.cpp
index 55b5ca9fe34..92611686589 100644
--- a/src/mongo/dbtests/querytests.cpp
+++ b/src/mongo/dbtests/querytests.cpp
@@ -34,6 +34,8 @@
 
 namespace mongo {
     extern int __findingStartInitialTimeout;
+    void assembleRequest( const string &ns, BSONObj query, int nToReturn, int nToSkip,
+                         const BSONObj *fieldsToReturn, int queryOptions, Message &toSend );
 }
 
 namespace QueryTests {
@@ -1224,6 +1226,29 @@ namespace QueryTests {
         }
     };
     
+    class Exhaust : public CollectionBase {
+    public:
+        Exhaust() : CollectionBase( "exhaust" ) {}
+        void run() {
+            BSONObj info;
+            ASSERT( client().runCommand( "unittests",
+                                        BSON( "create" << "querytests.exhaust" <<
+                                             "capped" << true << "size" << 8192 ), info ) );
+            client().insert( ns(), BSON( "ts" << 0 ) );
+            Message message;
+            assembleRequest( ns(), BSON( "ts" << GTE << 0 ), 0, 0, 0,
+                            QueryOption_OplogReplay | QueryOption_CursorTailable |
+                            QueryOption_Exhaust,
+                            message );
+            DbMessage dbMessage( message );
+            QueryMessage queryMessage( dbMessage );
+            Message result;
+            const char *exhaust = runQuery( message, queryMessage, *cc().curop(), result );
+            ASSERT( exhaust );
+            ASSERT_EQUALS( string( ns() ), exhaust );
+        }
+    };
+    
     namespace parsedtests {
         class basic1 {
         public:
@@ -1495,6 +1520,7 @@ namespace QueryTests {
             add< FindingStartPartiallyFull >();
             add< FindingStartStale >();
             add< WhatsMyUri >();
+            add< Exhaust >();
 
             add< parsedtests::basic1 >();