SERVER-18955 Have ShardedClientCursor set the batch size before all calls

comment update

Signed-off-by: Randolph Tan <randolph@10gen.com>
(cherry picked from commit 52840b1595449155acecb1a7401a332cbdb53fad)

Conflicts:
	src/mongo/s/cursors.cpp
	src/mongo/s/cursors.h

(cherry picked from commit 61349d3c00c5a6a4fecfcd12f528574d25ba1bf1)

3 files changed, 50 insertions, 25 deletions

diff --git a/jstests/sharding/sharded_limit_batchsize.js b/jstests/sharding/sharded_limit_batchsize.js
index e5a8c62acda..4c488499d6c 100644
--- a/jstests/sharding/sharded_limit_batchsize.js
+++ b/jstests/sharding/sharded_limit_batchsize.js
@@ -2,11 +2,29 @@
 // of limit and batchSize with sort return the correct results, and do not issue
 // unnecessary getmores (see SERVER-14299).
 
+
+/**
+ * This function checks that the log on the mongoD never shows the mongoS asking for all documents
+ * to be returned.  This occurs if ntoreturn:0 is present in the log.
+ */
+function checkLogForBatching(log) {
+	var haveGetMore = false;
+	log.forEach(function(logline) {
+		if (logline.search("getmore") != 0) {
+			haveGetMore = true;
+			//All ntoreturn should be the batch size, not 0
+			assert.eq(logline.search("ntoreturn:0"), -1)
+		}
+	});
+	assert.eq(haveGetMore, true);
+}
+
 /**
  * Test the correctness of queries with sort and batchSize on a sharded cluster,
  * running the queries against collection 'coll'.
  */
 function testBatchSize(coll) {
+	//Roll the cursor over the second batch and make sure it's correctly sized
     assert.eq(20, coll.find().sort({x: 1}).batchSize(3).itcount());
     assert.eq(15, coll.find().sort({x: 1}).batchSize(3).skip(5).itcount());
 }
@@ -88,10 +106,13 @@ for (var i=1; i<=10; ++i) {
 //
 
 jsTest.log("Running batchSize tests against sharded collection.");
-testBatchSize(shardedCol, st.shard0);
+st.shard0.adminCommand({setParameter: 1, logLevel : 1});
+testBatchSize(shardedCol);
+st.shard0.adminCommand({setParameter: 1, logLevel : 0});
+checkLogForBatching(st.shard0.adminCommand({ getLog: 'global' }).log);
 
 jsTest.log("Running batchSize tests against non-sharded collection.");
-testBatchSize(unshardedCol, st.shard0);
+testBatchSize(unshardedCol);
 
 //
 // Run tests for limit. These should *not* issue getmores. We confirm this
diff --git a/src/mongo/s/cursors.cpp b/src/mongo/s/cursors.cpp
index e1287677351..8bafd53f4b4 100644
--- a/src/mongo/s/cursors.cpp
+++ b/src/mongo/s/cursors.cpp
@@ -113,7 +113,7 @@ namespace mongo {
         return hasMore;
     }
 
-    bool ShardedClientCursor::sendNextBatch( Request& r , int ntoreturn ,
+    bool ShardedClientCursor::sendNextBatch( Request& r , int batchSize ,
             BufBuilder& buffer, int& docCount ) {
         uassert( 10191 ,  "cursor already done" , ! _done );
 
@@ -123,44 +123,48 @@ namespace mongo {
 
         docCount = 0;
 
-        // If ntoreturn is negative, it means that we should send up to -ntoreturn results
-        // back to the client, and that we should only send a *single batch*. An ntoreturn of
+        // If batchSize is negative, it means that we should send up to -batchSize results
+        // back to the client, and that we should only send a *single batch*. An batchSize of
         // 1 is also a special case which means "return up to 1 result in a single batch" (so
-        // that +1 actually has the same meaning of -1). For all other values of ntoreturn, we
+        // that +1 actually has the same meaning of -1). For all other values of batchSize, we
         // may have to return multiple batches.
-        const bool sendMoreBatches = ntoreturn == 0 || ntoreturn > 1;
-        ntoreturn = abs( ntoreturn );
+        const bool sendMoreBatches = batchSize == 0 || batchSize > 1;
+        batchSize = abs( batchSize );
+
+        // Set the initial batch size to 101, just like mongoD.
+        if (batchSize == 0 && _totalSent == 0)
+            batchSize = 101;
+
+        // Set batch size to batchSize requested by the current operation unconditionally.  This is
+        // necessary because if the loop exited due to docCount == batchSize then setBatchSize(0) was
+        // called, so the next _cusor->more() will be called with a batch size of 0 if the cursor
+        // buffer was drained the previous run.  Unconditionally setting the batch size ensures that
+        // we don't ask for a batch size of zero as a side effect.
+        _cursor->setBatchSize(batchSize);
 
         bool cursorHasMore = true;
         while ( ( cursorHasMore = _cursor->more() ) ) {
             BSONObj o = _cursor->next();
 
             buffer.appendBuf( (void*)o.objdata() , o.objsize() );
-            docCount++;
+            ++docCount;
+
             // Ensure that the next batch will never wind up requesting more docs from the shard
-            // than are remaining to satisfy the initial ntoreturn.
-            if (ntoreturn != 0) {
-                _cursor->setBatchSize(ntoreturn - docCount);
+            // than are remaining to satisfy the initial batchSize.
+            if (batchSize != 0) {
+                if (docCount == batchSize)
+                    break;
+                _cursor->setBatchSize(batchSize - docCount);
             }
 
             if ( buffer.len() > maxSize ) {
                 break;
             }
-
-            if ( docCount == ntoreturn ) {
-                // soft limit aka batch size
-                break;
-            }
-
-            if ( ntoreturn == 0 && _totalSent == 0 && docCount >= 100 ) {
-                // first batch should be max 100 unless batch size specified
-                break;
-            }
         }
 
         // We need to request another batch if the following two conditions hold:
         //
-        //  1. ntoreturn is positive and not equal to 1 (see the comment above). This condition
+        //  1. batchSize is positive and not equal to 1 (see the comment above). This condition
         //  is stored in 'sendMoreBatches'.
         //
         //  2. The last call to _cursor->more() was true (i.e. we never explicitly got a false
@@ -180,7 +184,7 @@ namespace mongo {
         LOG(5) << "\t hasMoreBatches: " << hasMoreBatches
                << " sendMoreBatches: " << sendMoreBatches
                << " cursorHasMore: " << cursorHasMore
-               << " ntoreturn: " << ntoreturn
+               << " batchSize: " << batchSize
                << " num: " << docCount
                << " id:" << getId()
                << " totalSent: " << _totalSent << endl;
diff --git a/src/mongo/s/cursors.h b/src/mongo/s/cursors.h
index 73503bced7b..11aa864e672 100644
--- a/src/mongo/s/cursors.h
+++ b/src/mongo/s/cursors.h
@@ -73,7 +73,7 @@ namespace mongo {
          *
          * @return true if this is not the final batch.
          */
-        bool sendNextBatch( Request& r, int ntoreturn, BufBuilder& buffer, int& docCount );
+        bool sendNextBatch( Request& r, int batchSize, BufBuilder& buffer, int& docCount );
 
         void accessed();
         /** @return idle time in ms */