SERVER-18772 Add ntoreturn to LiteParsedQuery

13 files changed, 200 insertions, 49 deletions

diff --git a/src/mongo/client/remote_command_runner_impl.cpp b/src/mongo/client/remote_command_runner_impl.cpp
index f025ca271fc..720e0ecb3c0 100644
--- a/src/mongo/client/remote_command_runner_impl.cpp
+++ b/src/mongo/client/remote_command_runner_impl.cpp
@@ -115,6 +115,10 @@ Status runDownconvertedFindCommand(DBClientConnection* conn,
 
     auto& lpq = lpqStatus.getValue();
 
+    // We are downconverting a find command, and find command can only have ntoreturn
+    // if it was generated by mongos.
+    invariant(!lpq->getNToReturn());
+
     Query query(lpq->getFilter());
     if (!lpq->getSort().isEmpty()) {
         query.sort(lpq->getSort());
diff --git a/src/mongo/db/commands/find_cmd.cpp b/src/mongo/db/commands/find_cmd.cpp
index 9ed570e8df5..775518e402a 100644
--- a/src/mongo/db/commands/find_cmd.cpp
+++ b/src/mongo/db/commands/find_cmd.cpp
@@ -216,7 +216,7 @@ public:
         }
 
         // Fill out curop information.
-        long long ntoreturn = lpq->getBatchSize().value_or(0);
+        long long ntoreturn = lpq->getNToReturn().value_or(0);
         beginQueryOp(txn, nss, cmdObj, ntoreturn, lpq->getSkip().value_or(0));
 
         // 1b) Finish the parsing step by using the LiteParsedQuery to create a CanonicalQuery.
diff --git a/src/mongo/db/exec/multi_plan.cpp b/src/mongo/db/exec/multi_plan.cpp
index b93af367f89..26200beb931 100644
--- a/src/mongo/db/exec/multi_plan.cpp
+++ b/src/mongo/db/exec/multi_plan.cpp
@@ -203,10 +203,10 @@ size_t MultiPlanStage::getTrialPeriodNumToReturn(const CanonicalQuery& query) {
     // Determine the number of results which we will produce during the plan
     // ranking phase before stopping.
     size_t numResults = static_cast<size_t>(internalQueryPlanEvaluationMaxResults);
-    if (query.getParsed().getLimit()) {
+    if (query.getParsed().getNToReturn()) {
+        numResults = std::min(static_cast<size_t>(*query.getParsed().getNToReturn()), numResults);
+    } else if (query.getParsed().getLimit()) {
         numResults = std::min(static_cast<size_t>(*query.getParsed().getLimit()), numResults);
-    } else if (!query.getParsed().isFromFindCommand() && query.getParsed().getBatchSize()) {
-        numResults = std::min(static_cast<size_t>(*query.getParsed().getBatchSize()), numResults);
     }
 
     return numResults;
diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp
index 7b780cccd6d..427c2083f2d 100644
--- a/src/mongo/db/query/canonical_query.cpp
+++ b/src/mongo/db/query/canonical_query.cpp
@@ -571,7 +571,11 @@ std::string CanonicalQuery::toString() const {
     }
 
     if (_pq->getSkip()) {
-        ss << " skip=" << *_pq->getSkip() << "\n";
+        ss << " skip=" << *_pq->getSkip();
+    }
+
+    if (_pq->getNToReturn()) {
+        ss << " ntoreturn=" << *_pq->getNToReturn() << '\n';
     }
 
     // The expression tree puts an endl on for us.
@@ -598,6 +602,10 @@ std::string CanonicalQuery::toStringShort() const {
         ss << " skip: " << *_pq->getSkip();
     }
 
+    if (_pq->getNToReturn()) {
+        ss << " ntoreturn=" << *_pq->getNToReturn();
+    }
+
     return ss;
 }
 
diff --git a/src/mongo/db/query/find.cpp b/src/mongo/db/query/find.cpp
index 46dcaa014d6..def9c310d5b 100644
--- a/src/mongo/db/query/find.cpp
+++ b/src/mongo/db/query/find.cpp
@@ -91,7 +91,7 @@ bool shouldSaveCursor(OperationContext* txn,
         return false;
     }
 
-    if (!pq.isFromFindCommand() && pq.getBatchSize() && *pq.getBatchSize() == 1) {
+    if (pq.getNToReturn().value_or(0) == 1) {
         return false;
     }
 
@@ -574,7 +574,7 @@ std::string runQuery(OperationContext* txn,
 
         if (FindCommon::enoughForFirstBatch(pq, numResults, bb.len())) {
             LOG(5) << "Enough for first batch, wantMore=" << pq.wantMore()
-                   << " batchSize=" << pq.getBatchSize().value_or(0) << " numResults=" << numResults
+                   << " ntoreturn=" << pq.getNToReturn().value_or(0) << " numResults=" << numResults
                    << endl;
             break;
         }
diff --git a/src/mongo/db/query/find_common.cpp b/src/mongo/db/query/find_common.cpp
index 8d6a718313e..07e5256f1ad 100644
--- a/src/mongo/db/query/find_common.cpp
+++ b/src/mongo/db/query/find_common.cpp
@@ -37,7 +37,7 @@ namespace mongo {
 bool FindCommon::enoughForFirstBatch(const LiteParsedQuery& pq,
                                      long long numDocs,
                                      int bytesBuffered) {
-    if (!pq.getBatchSize()) {
+    if (!pq.getEffectiveBatchSize()) {
         // If there is no batch size, we stop generating additional results as soon as we have
         // either 101 documents or at least 1MB of data.
         return (bytesBuffered > 1024 * 1024) || numDocs >= LiteParsedQuery::kDefaultBatchSize;
@@ -45,11 +45,15 @@ bool FindCommon::enoughForFirstBatch(const LiteParsedQuery& pq,
 
     // If there is a batch size, we add results until either satisfying this batch size or exceeding
     // the 4MB size threshold.
-    return numDocs >= *pq.getBatchSize() || bytesBuffered > kMaxBytesToReturnToClientAtOnce;
+    return numDocs >= pq.getEffectiveBatchSize().value() ||
+        bytesBuffered > kMaxBytesToReturnToClientAtOnce;
 }
 
-bool FindCommon::enoughForGetMore(long long ntoreturn, long long numDocs, int bytesBuffered) {
-    return (ntoreturn && numDocs >= ntoreturn) || (bytesBuffered > kMaxBytesToReturnToClientAtOnce);
+bool FindCommon::enoughForGetMore(long long effectiveBatchSize,
+                                  long long numDocs,
+                                  int bytesBuffered) {
+    return (effectiveBatchSize && numDocs >= effectiveBatchSize) ||
+        (bytesBuffered > kMaxBytesToReturnToClientAtOnce);
 }
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/find_common.h b/src/mongo/db/query/find_common.h
index df69ce56fb3..a16c3187ffc 100644
--- a/src/mongo/db/query/find_common.h
+++ b/src/mongo/db/query/find_common.h
@@ -52,12 +52,14 @@ public:
     /**
      * Returns true if enough results have been prepared to stop adding more to a getMore batch.
      *
-     * An 'ntoreturn' value of zero is interpreted as the absence of a batchSize; in this case,
-     * returns true only once the size threshold is exceeded. If 'ntoreturn' is positive, returns
-     * true once either are added until we have either satisfied the batch size or exceeded the size
-     * threshold.
+     * An 'effectiveBatchSize' value of zero is interpreted as the absence of a batchSize;
+     * in this case, returns true only once the size threshold is exceeded. If 'effectiveBatchSize'
+     * is positive, returns true once either are added until we have either satisfied the batch size
+     * or exceeded the size threshold.
      */
-    static bool enoughForGetMore(long long ntoreturn, long long numDocs, int bytesBuffered);
+    static bool enoughForGetMore(long long effectiveBatchSize,
+                                 long long numDocs,
+                                 int bytesBuffered);
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/lite_parsed_query.cpp b/src/mongo/db/query/lite_parsed_query.cpp
index 306a9856bc6..f114f899966 100644
--- a/src/mongo/db/query/lite_parsed_query.cpp
+++ b/src/mongo/db/query/lite_parsed_query.cpp
@@ -78,6 +78,7 @@ const char kHintField[] = "hint";
 const char kSkipField[] = "skip";
 const char kLimitField[] = "limit";
 const char kBatchSizeField[] = "batchSize";
+const char kNToReturnField[] = "ntoreturn";
 const char kSingleBatchField[] = "singleBatch";
 const char kCommentField[] = "comment";
 const char kMaxScanField[] = "maxScan";
@@ -107,7 +108,6 @@ StatusWith<unique_ptr<LiteParsedQuery>> LiteParsedQuery::makeFromFindCommand(Nam
                                                                              const BSONObj& cmdObj,
                                                                              bool isExplain) {
     unique_ptr<LiteParsedQuery> pq(new LiteParsedQuery(std::move(nss)));
-    pq->_fromCommand = true;
     pq->_explain = isExplain;
 
     // Parse the command BSON by looping through one element at a time.
@@ -207,6 +207,20 @@ StatusWith<unique_ptr<LiteParsedQuery>> LiteParsedQuery::makeFromFindCommand(Nam
             }
 
             pq->_batchSize = batchSize;
+        } else if (str::equals(fieldName, kNToReturnField)) {
+            if (!el.isNumber()) {
+                str::stream ss;
+                ss << "Failed to parse: " << cmdObj.toString() << ". "
+                   << "'ntoreturn' field must be numeric.";
+                return Status(ErrorCodes::FailedToParse, ss);
+            }
+
+            long long ntoreturn = el.numberLong();
+            if (ntoreturn < 0) {
+                return Status(ErrorCodes::BadValue, "ntoreturn value must be non-negative");
+            }
+
+            pq->_ntoreturn = ntoreturn;
         } else if (str::equals(fieldName, kSingleBatchField)) {
             Status status = checkFieldType(el, Bool);
             if (!status.isOK()) {
@@ -411,6 +425,7 @@ std::unique_ptr<LiteParsedQuery> LiteParsedQuery::makeAsFindCmd(
     boost::optional<long long> skip,
     boost::optional<long long> limit,
     boost::optional<long long> batchSize,
+    boost::optional<long long> ntoreturn,
     bool wantMore,
     bool isExplain,
     const std::string& comment,
@@ -429,7 +444,10 @@ std::unique_ptr<LiteParsedQuery> LiteParsedQuery::makeAsFindCmd(
     bool isAwaitData,
     bool isPartial) {
     unique_ptr<LiteParsedQuery> pq(new LiteParsedQuery(std::move(nss)));
-    pq->_fromCommand = true;
+    // ntoreturn and batchSize or limit are mutually exclusive.
+    if (batchSize || limit) {
+        invariant(!ntoreturn);
+    }
 
     pq->_filter = filter;
     pq->_proj = projection;
@@ -439,6 +457,7 @@ std::unique_ptr<LiteParsedQuery> LiteParsedQuery::makeAsFindCmd(
     pq->_skip = skip;
     pq->_limit = limit;
     pq->_batchSize = batchSize;
+    pq->_ntoreturn = ntoreturn;
     pq->_wantMore = wantMore;
 
     pq->_explain = isExplain;
@@ -492,6 +511,10 @@ void LiteParsedQuery::asFindCommand(BSONObjBuilder* cmdBuilder) const {
         cmdBuilder->append(kSkipField, *_skip);
     }
 
+    if (_ntoreturn) {
+        cmdBuilder->append(kNToReturnField, *_ntoreturn);
+    }
+
     if (_limit) {
         cmdBuilder->append(kLimitField, *_limit);
     }
@@ -622,6 +645,11 @@ Status LiteParsedQuery::validate() const {
         }
     }
 
+    if ((_limit || _batchSize) && _ntoreturn) {
+        return Status(ErrorCodes::BadValue,
+                      "'limit' or 'batchSize' fields can not be set with 'ntoreturn' field.");
+    }
+
     return Status::OK();
 }
 
@@ -746,31 +774,30 @@ Status LiteParsedQuery::init(int ntoskip,
     _proj = proj.getOwned();
 
     if (ntoskip) {
+        if (ntoskip < 0) {
+            str::stream ss;
+            ss << "Skip value must be positive, but received: " << ntoskip << ". ";
+            return Status(ErrorCodes::BadValue, ss);
+        }
         _skip = ntoskip;
     }
 
     if (ntoreturn) {
-        _batchSize = ntoreturn;
+        if (ntoreturn < 0) {
+            if (ntoreturn == std::numeric_limits<int>::min()) {
+                // ntoreturn is negative but can't be negated.
+                return Status(ErrorCodes::BadValue, "bad ntoreturn value in query");
+            }
+            _ntoreturn = -ntoreturn;
+            _wantMore = false;
+        } else {
+            _ntoreturn = ntoreturn;
+        }
     }
 
     // Initialize flags passed as 'queryOptions' bit vector.
     initFromInt(queryOptions);
 
-    if (_skip && *_skip < 0) {
-        return Status(ErrorCodes::BadValue, "bad skip value in query");
-    }
-
-    if (_batchSize && *_batchSize < 0) {
-        if (*_batchSize == std::numeric_limits<int>::min()) {
-            // _batchSize is negative but can't be negated.
-            return Status(ErrorCodes::BadValue, "bad limit value in query");
-        }
-
-        // A negative number indicates that the cursor should be closed after the first batch.
-        _wantMore = false;
-        _batchSize = -*_batchSize;
-    }
-
     if (fromQueryMessage) {
         BSONElement queryField = queryObj["query"];
         if (!queryField.isABSONObj()) {
@@ -954,4 +981,8 @@ Status LiteParsedQuery::validateFindCmd() {
     return validate();
 }
 
+boost::optional<long long> LiteParsedQuery::getEffectiveBatchSize() const {
+    return _batchSize ? _batchSize : _ntoreturn;
+}
+
 }  // namespace mongo
diff --git a/src/mongo/db/query/lite_parsed_query.h b/src/mongo/db/query/lite_parsed_query.h
index 5cc841922c0..ab1f8ab3aeb 100644
--- a/src/mongo/db/query/lite_parsed_query.h
+++ b/src/mongo/db/query/lite_parsed_query.h
@@ -92,6 +92,7 @@ public:
         boost::optional<long long> skip = boost::none,
         boost::optional<long long> limit = boost::none,
         boost::optional<long long> batchSize = boost::none,
+        boost::optional<long long> ntoreturn = boost::none,
         bool wantMore = true,
         bool isExplain = false,
         const std::string& comment = "",
@@ -194,13 +195,20 @@ public:
     boost::optional<long long> getBatchSize() const {
         return _batchSize;
     }
+    boost::optional<long long> getNToReturn() const {
+        return _ntoreturn;
+    }
+
+    /**
+     * Returns batchSize or ntoreturn value if either is set. If neither is set,
+     * returns boost::none.
+     */
+    boost::optional<long long> getEffectiveBatchSize() const;
+
     bool wantMore() const {
         return _wantMore;
     }
 
-    bool isFromFindCommand() const {
-        return _fromCommand;
-    }
     bool isExplain() const {
         return _explain;
     }
@@ -349,7 +357,11 @@ private:
     // allowed.
     boost::optional<long long> _batchSize;
 
-    bool _fromCommand = false;
+    // Set only when parsed from an OP_QUERY find message. The value is computed by driver or shell
+    // and is set to be a min of batchSize and limit provided by user. LPQ can have set either
+    // ntoreturn or batchSize / limit.
+    boost::optional<long long> _ntoreturn;
+
     bool _explain = false;
 
     std::string _comment;
diff --git a/src/mongo/db/query/lite_parsed_query_test.cpp b/src/mongo/db/query/lite_parsed_query_test.cpp
index 05909fd4388..1af3279e438 100644
--- a/src/mongo/db/query/lite_parsed_query_test.cpp
+++ b/src/mongo/db/query/lite_parsed_query_test.cpp
@@ -107,7 +107,7 @@ TEST(LiteParsedQueryTest, NumToReturn) {
                                                                              false,     // snapshot
                                                                              false)));  // explain
 
-    ASSERT_EQUALS(6, *lpq->getBatchSize());
+    ASSERT_EQUALS(6, *lpq->getNToReturn());
     ASSERT(lpq->wantMore());
 }
 
@@ -125,7 +125,7 @@ TEST(LiteParsedQueryTest, NumToReturnNegative) {
                                                                              false,     // snapshot
                                                                              false)));  // explain
 
-    ASSERT_EQUALS(6, *lpq->getBatchSize());
+    ASSERT_EQUALS(6, *lpq->getNToReturn());
     ASSERT(!lpq->wantMore());
 }
 
@@ -257,7 +257,6 @@ TEST(LiteParsedQueryTest, ForbidMetaSortOnFieldWithoutMetaProject) {
 
 TEST(LiteParsedQueryTest, MakeAsFindCmdDefaultArgs) {
     auto lpq = LiteParsedQuery::makeAsFindCmd(NamespaceString("test.ns"));
-    ASSERT_TRUE(lpq->isFromFindCommand());
 
     ASSERT_EQUALS("test.ns", lpq->ns());
 
@@ -269,6 +268,7 @@ TEST(LiteParsedQueryTest, MakeAsFindCmdDefaultArgs) {
     ASSERT_FALSE(lpq->getSkip());
     ASSERT_FALSE(lpq->getLimit());
     ASSERT_FALSE(lpq->getBatchSize());
+    ASSERT_FALSE(lpq->getNToReturn());
     ASSERT_TRUE(lpq->wantMore());
 
     ASSERT_FALSE(lpq->isExplain());
@@ -300,6 +300,7 @@ TEST(LiteParsedQueryTest, MakeFindCmdAllArgs) {
                                               4,
                                               5,
                                               6,
+                                              boost::none,
                                               false,
                                               true,
                                               "this is a comment",
@@ -317,7 +318,6 @@ TEST(LiteParsedQueryTest, MakeFindCmdAllArgs) {
                                               true,
                                               true,
                                               true);
-    ASSERT_TRUE(lpq->isFromFindCommand());
 
     ASSERT_EQUALS("test.ns", lpq->ns());
 
@@ -329,7 +329,71 @@ TEST(LiteParsedQueryTest, MakeFindCmdAllArgs) {
     ASSERT_EQ(4, *lpq->getSkip());
     ASSERT_EQ(5, *lpq->getLimit());
     ASSERT_EQ(6, *lpq->getBatchSize());
+    ASSERT_FALSE(lpq->getNToReturn());
     ASSERT_FALSE(lpq->wantMore());
+    ASSERT_EQ(6, *lpq->getEffectiveBatchSize());
+
+    ASSERT_TRUE(lpq->isExplain());
+    ASSERT_EQ("this is a comment", lpq->getComment());
+    ASSERT_EQUALS(7, lpq->getMaxScan());
+    ASSERT_EQUALS(8, lpq->getMaxTimeMS());
+
+    ASSERT_EQUALS(BSON("e" << 1), lpq->getMin());
+    ASSERT_EQUALS(BSON("f" << 1), lpq->getMax());
+
+    ASSERT_TRUE(lpq->returnKey());
+    ASSERT_TRUE(lpq->showRecordId());
+    ASSERT_TRUE(lpq->isSnapshot());
+    ASSERT_TRUE(lpq->hasReadPref());
+    ASSERT_TRUE(lpq->isTailable());
+    ASSERT_TRUE(lpq->isSlaveOk());
+    ASSERT_TRUE(lpq->isOplogReplay());
+    ASSERT_TRUE(lpq->isNoCursorTimeout());
+    ASSERT_TRUE(lpq->isAwaitData());
+    ASSERT_TRUE(lpq->isPartial());
+}
+
+TEST(LiteParsedQueryTest, MakeAsFindCmdNToReturn) {
+    auto lpq = LiteParsedQuery::makeAsFindCmd(NamespaceString("test.ns"),
+                                              BSON("a" << 1),
+                                              BSON("b" << 1),
+                                              BSON("c" << 1),
+                                              BSON("d" << 1),
+                                              4,
+                                              boost::none,
+                                              boost::none,
+                                              5,
+                                              false,
+                                              true,
+                                              "this is a comment",
+                                              7,
+                                              8,
+                                              BSON("e" << 1),
+                                              BSON("f" << 1),
+                                              true,
+                                              true,
+                                              true,
+                                              true,
+                                              true,
+                                              true,
+                                              true,
+                                              true,
+                                              true,
+                                              true);
+
+    ASSERT_EQUALS("test.ns", lpq->ns());
+
+    ASSERT_EQUALS(BSON("a" << 1), lpq->getFilter());
+    ASSERT_EQUALS(BSON("b" << 1), lpq->getProj());
+    ASSERT_EQUALS(BSON("c" << 1), lpq->getSort());
+    ASSERT_EQUALS(BSON("d" << 1), lpq->getHint());
+
+    ASSERT_EQ(4, *lpq->getSkip());
+    ASSERT_FALSE(lpq->getLimit());
+    ASSERT_FALSE(lpq->getBatchSize());
+    ASSERT_EQ(5, *lpq->getNToReturn());
+    ASSERT_FALSE(lpq->wantMore());
+    ASSERT_EQ(5, *lpq->getEffectiveBatchSize());
 
     ASSERT_TRUE(lpq->isExplain());
     ASSERT_EQ("this is a comment", lpq->getComment());
@@ -1013,7 +1077,7 @@ TEST(LiteParsedQueryTest, ParseCommandIsFromFindCommand) {
     unique_ptr<LiteParsedQuery> lpq(
         assertGet(LiteParsedQuery::makeFromFindCommand(nss, cmdObj, isExplain)));
 
-    ASSERT(lpq->isFromFindCommand());
+    ASSERT_FALSE(lpq->getNToReturn());
 }
 
 TEST(LiteParsedQueryTest, ParseCommandNotFromFindCommand) {
@@ -1030,7 +1094,7 @@ TEST(LiteParsedQueryTest, ParseCommandNotFromFindCommand) {
                                                  BSONObj(),
                                                  false,     // snapshot
                                                  false)));  // explain
-    ASSERT(!lpq->isFromFindCommand());
+    ASSERT_TRUE(lpq->getNToReturn());
 }
 
 TEST(LiteParsedQueryTest, ParseCommandAwaitDataButNotTailable) {
@@ -1068,7 +1132,7 @@ TEST(LiteParsedQueryTest, DefaultQueryParametersCorrect) {
     ASSERT_FALSE(lpq->getLimit());
 
     ASSERT_EQUALS(true, lpq->wantMore());
-    ASSERT_EQUALS(true, lpq->isFromFindCommand());
+    ASSERT_FALSE(lpq->getNToReturn());
     ASSERT_EQUALS(false, lpq->isExplain());
     ASSERT_EQUALS(0, lpq->getMaxScan());
     ASSERT_EQUALS(0, lpq->getMaxTimeMS());
diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 01dddc60fba..ae932223137 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -523,14 +523,14 @@ QuerySolutionNode* QueryPlannerAnalysis::analyzeSort(const CanonicalQuery& query
         // We have a true limit. The limit can be combined with the SORT stage.
         sort->limit =
             static_cast<size_t>(*lpq.getLimit()) + static_cast<size_t>(lpq.getSkip().value_or(0));
-    } else if (!lpq.isFromFindCommand() && lpq.getBatchSize()) {
+    } else if (lpq.getNToReturn()) {
         // We have an ntoreturn specified by an OP_QUERY style find. This is used
         // by clients to mean both batchSize and limit.
         //
         // Overflow here would be bad and could cause a nonsense limit. Cast
         // skip and limit values to unsigned ints to make sure that the
         // sum is never stored as signed. (See SERVER-13537).
-        sort->limit = static_cast<size_t>(*lpq.getBatchSize()) +
+        sort->limit = static_cast<size_t>(*lpq.getNToReturn()) +
             static_cast<size_t>(lpq.getSkip().value_or(0));
 
         // This is a SORT with a limit. The wire protocol has a single quantity
@@ -792,11 +792,11 @@ QuerySolution* QueryPlannerAnalysis::analyzeDataAccess(const CanonicalQuery& que
             limit->limit = *lpq.getLimit();
             limit->children.push_back(solnRoot);
             solnRoot = limit;
-        } else if (!lpq.isFromFindCommand() && lpq.getBatchSize() && !lpq.wantMore()) {
+        } else if (lpq.getNToReturn() && !lpq.wantMore()) {
             // We have a "legacy limit", i.e. a negative ntoreturn value from an OP_QUERY style
             // find.
             LimitNode* limit = new LimitNode();
-            limit->limit = *lpq.getBatchSize();
+            limit->limit = *lpq.getNToReturn();
             limit->children.push_back(solnRoot);
             solnRoot = limit;
         }
diff --git a/src/mongo/db/query/query_planner_test.cpp b/src/mongo/db/query/query_planner_test.cpp
index 0d85ccae1ad..7d4a3ab1e5e 100644
--- a/src/mongo/db/query/query_planner_test.cpp
+++ b/src/mongo/db/query/query_planner_test.cpp
@@ -3956,4 +3956,28 @@ TEST(BadInputTest, TagAccordingToCache) {
     ASSERT_NOT_OK(s);
 }
 
+// A query run as a find command with a sort and ntoreturn should generate a plan implementing
+// the 'ntoreturn hack'.
+TEST_F(QueryPlannerTest, NToReturnHackWithFindCommand) {
+    params.options |= QueryPlannerParams::SPLIT_LIMITED_SORT;
+
+    runQueryAsCommand(fromjson("{find: 'testns', sort: {a:1}, ntoreturn:3}"));
+
+    assertNumSolutions(1U);
+    assertSolutionExists(
+        "{or: {nodes: ["
+        "{sort: {limit:3, pattern: {a:1}, node: {cscan: {dir:1}}}}, "
+        "{sort: {limit:0, pattern: {a:1}, node: {cscan: {dir:1}}}}"
+        "]}}");
+}
+
+TEST_F(QueryPlannerTest, NToReturnHackWithSingleBatch) {
+    params.options |= QueryPlannerParams::SPLIT_LIMITED_SORT;
+
+    runQueryAsCommand(fromjson("{find: 'testns', sort: {a:1}, ntoreturn:3, singleBatch:true}"));
+
+    assertNumSolutions(1U);
+    assertSolutionExists("{sort: {pattern: {a:1}, limit:3, node: {cscan: {dir:1, filter: {}}}}}");
+}
+
 }  // namespace
diff --git a/src/mongo/s/query/cluster_find.cpp b/src/mongo/s/query/cluster_find.cpp
index d3f03681a1a..5e0a9957323 100644
--- a/src/mongo/s/query/cluster_find.cpp
+++ b/src/mongo/s/query/cluster_find.cpp
@@ -74,6 +74,7 @@ std::unique_ptr<LiteParsedQuery> transformQueryForShards(const LiteParsedQuery&
                                           boost::none,  // Don't forward skip.
                                           newLimit,
                                           lpq.getBatchSize(),
+                                          lpq.getNToReturn(),
                                           lpq.wantMore(),
                                           lpq.isExplain(),
                                           lpq.getComment(),
@@ -117,6 +118,7 @@ StatusWith<CursorId> runQueryWithoutRetrying(OperationContext* txn,
     }
 
     ClusterClientCursorParams params(query.nss());
+    params.limit = query.getParsed().getLimit();
     params.batchSize = query.getParsed().getBatchSize();
     params.limit = query.getParsed().getLimit();
     params.sort = query.getParsed().getSort();