SERVER-26334 fix mongos handling of ntoreturn:1

2 files changed, 6 insertions, 5 deletions

diff --git a/src/mongo/db/query/find.cpp b/src/mongo/db/query/find.cpp
index 14645bd55e1..bbce97fc50b 100644
--- a/src/mongo/db/query/find.cpp
+++ b/src/mongo/db/query/find.cpp
@@ -89,11 +89,7 @@ bool shouldSaveCursor(OperationContext* txn,
     }
 
     const QueryRequest& qr = exec->getCanonicalQuery()->getQueryRequest();
-    if (!qr.wantMore() && !qr.isTailable()) {
-        return false;
-    }
-
-    if (qr.getNToReturn().value_or(0) == 1) {
+    if (!qr.wantMore()) {
         return false;
     }
 
diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp
index b0aa759e853..af2c46b550c 100644
--- a/src/mongo/db/query/query_request.cpp
+++ b/src/mongo/db/query/query_request.cpp
@@ -734,6 +734,11 @@ Status QueryRequest::init(int ntoskip,
         }
     }
 
+    // An ntoreturn of 1 is special because it also means to return at most one batch.
+    if (_ntoreturn.value_or(0) == 1) {
+        _wantMore = false;
+    }
+
     // Initialize flags passed as 'queryOptions' bit vector.
     initFromInt(queryOptions);