SERVER-4680 sendNextBatch logic isn't correct with nonzero batch size (corrected backport)

3 files changed, 22 insertions, 12 deletions

diff --git a/client/parallel.cpp b/client/parallel.cpp
index 3a33eb59488..d7975a6849b 100644
--- a/client/parallel.cpp
+++ b/client/parallel.cpp
@@ -503,7 +503,19 @@ namespace mongo {
                                                             0 , // nToSkip
                                                             _fields.isEmpty() ? 0 : &_fields , // fieldsToReturn
                                                             _options ,
-                                                            _batchSize == 0 ? 0 : _batchSize + _needToSkip // batchSize
+                                                            // NtoReturn is weird.
+                                                            // If zero, it means use default size, so we do that for all cursors
+                                                            // If positive, it's the batch size (we don't want this cursor limiting results), tha
+                                                            // done at a higher level
+                                                            // If negative, it's the batch size, but we don't create a cursor - so we don't want
+                                                            // to create a child cursor either.
+                                                            // Either way, if non-zero, we want to pull back the batch size + the skip amount as
+                                                            // quickly as possible.  Potentially, for a cursor on a single shard or if we keep be
+                                                            // chunks, we can actually add the skip value into the cursor and/or make some assump
+                                                            // return value size ( (batch size + skip amount) / num_servers ).
+                                                            _batchSize == 0 ? 0 :
+                                                                ( _batchSize > 0 ? _batchSize + _needToSkip :
+                                                                                   _batchSize - _needToSkip ) // batchSize
                                                             ) );
 
                 try{
diff --git a/s/cursors.cpp b/s/cursors.cpp
index 12b3d5e7d2f..5957ffcb418 100644
--- a/s/cursors.cpp
+++ b/s/cursors.cpp
@@ -82,7 +82,10 @@ namespace mongo {
         BufBuilder b(32768);
 
         int num = 0;
-        bool sendMore = true;
+
+        // Send more if ntoreturn is 0, or any value > 1 (one is assumed to be a single doc return, with no cursor)
+        bool sendMore = ntoreturn == 0 || ntoreturn > 1;
+        ntoreturn = abs( ntoreturn );
 
         while ( _cursor->more() ) {
             BSONObj o = _cursor->next();
@@ -99,20 +102,15 @@ namespace mongo {
                 break;
             }
 
-            if ( ntoreturn != 0 && ( -1 * num + _totalSent ) == ntoreturn ) {
-                // hard limit - total to send
-                sendMore = false;
-                break;
-            }
-
-            if ( ntoreturn == 0 && _totalSent == 0 && num > 100 ) {
+            if ( ntoreturn == 0 && _totalSent == 0 && num >= 100 ) {
                 // first batch should be max 100 unless batch size specified
                 break;
             }
         }
 
         bool hasMore = sendMore && _cursor->more();
-        LOG(6) << "\t hasMore:" << hasMore << " wouldSendMoreIfHad: " << sendMore << " id:" << getId() << " totalSent: " << _totalSent << endl;
+        LOG(5) << "\t hasMore: " << hasMore << " sendMore: " << sendMore << " cursorMore: " << _cursor->more() << " ntoreturn: " << ntoreturn
+               << " num: " << num << " wouldSendMoreIfHad: " << sendMore << " id:" << getId() << " totalSent: " << _totalSent << endl;
 
         replyToQuery( 0 , r.p() , r.m() , b.buf() , b.len() , num , _totalSent , hasMore ? getId() : 0 );
         _totalSent += num;
diff --git a/s/strategy_shard.cpp b/s/strategy_shard.cpp
index c6b30e7965f..c96a7e1b202 100644
--- a/s/strategy_shard.cpp
+++ b/s/strategy_shard.cpp
@@ -91,10 +91,10 @@ namespace mongo {
             }
 
             ShardedClientCursorPtr cc (new ShardedClientCursor( q , cursor ));
-            if ( ! cc->sendNextBatch( r ) ) {
+            if ( ! cc->sendNextBatch( r, q.ntoreturn ) ) {
                 return;
             }
-            LOG(6) << "storing cursor : " << cc->getId() << endl;
+            LOG(5) << "storing cursor : " << cc->getId() << endl;
             cursorCache.store( cc );
         }