From 2e20f7615c340c8c8faa2777dfd734e0dcabb5e2 Mon Sep 17 00:00:00 2001
From: Zixuan Zhuang <zixuan.zhuang@mongodb.com>
Date: Thu, 17 Aug 2023 00:54:38 +0000
Subject: [PATCH] SERVER-78558: Establish mongot cursor before SBE builder

---
 src/mongo/db/pipeline/search_helper.h | 18 ++++++++++++++++++
 src/mongo/db/query/query_planner.cpp  |  5 ++++-
 src/mongo/db/query/query_solution.cpp |  9 ++++++++-
 src/mongo/db/query/query_solution.h   | 27 ++++++++++++++++++++++++++-
 4 files changed, 56 insertions(+), 3 deletions(-)

diff --git a/src/mongo/db/pipeline/search_helper.h b/src/mongo/db/pipeline/search_helper.h
index 062d3236ed1..8c52385799f 100644
--- a/src/mongo/db/pipeline/search_helper.h
+++ b/src/mongo/db/pipeline/search_helper.h
@@ -139,6 +139,24 @@ public:
     virtual bool isSearchMetaStage(DocumentSource* stage) {
         return false;
     }
+
+    /**
+     * Gets the information for the search QSN from DocumentSourceSearch.
+     * The results are returned as a tuple of the format:
+     * <limit, mongotDocsRequested, searchQuery, taskExecutor, intermediateResultsProtocolVersion>
+     */
+    virtual std::unique_ptr<SearchNode> getSearchNode(DocumentSource* stage) {
+        return nullptr;
+    }
+
+    /**
+     * Executes the initial $search query to get the cursor id and first batch for building the
+     * $search SBE plan.
+     */
+    virtual std::pair<CursorResponse, CursorResponse> establishSearchQueryCursors(
+        const boost::intrusive_ptr<ExpressionContext>& expCtx, const SearchNode* searchNode) {
+        return {CursorResponse(), CursorResponse()};
+    }
 };
 
 /**
diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp
index 64c3710bc16..13a25b37042 100644
--- a/src/mongo/db/query/query_planner.cpp
+++ b/src/mongo/db/query/query_planner.cpp
@@ -556,8 +556,11 @@ StatusWith<std::unique_ptr<QuerySolution>> tryToBuildSearchQuerySolution(
                 "Pushing down $search into SBE but featureFlagSearchInSbe is disabled.",
                 feature_flags::gFeatureFlagSearchInSbe.isEnabledAndIgnoreFCVUnsafe());
 
+        auto searchNode =
+            getSearchHelpers(query.getOpCtx()->getServiceContext())->getSearchNode(stage);
+
         auto querySoln = std::make_unique<QuerySolution>();
-        querySoln->setRoot(std::make_unique<SearchNode>(isSearchMeta));
+        querySoln->setRoot(std::move(searchNode));
         return std::move(querySoln);
     }
 
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index c3e2ebd04ff..35202e2b65d 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -1819,7 +1819,8 @@ void SentinelNode::appendToString(str::stream* ss, int indent) const {
 }
 
 std::unique_ptr<QuerySolutionNode> SearchNode::clone() const {
-    return std::make_unique<SearchNode>(isSearchMeta);
+    return std::make_unique<SearchNode>(
+        isSearchMeta, searchQuery, limit, intermediateResultsProtocolVersion);
 }
 
 void SearchNode::appendToString(str::stream* ss, int indent) const {
@@ -1827,6 +1828,12 @@ void SearchNode::appendToString(str::stream* ss, int indent) const {
     *ss << "SEARCH\n";
     addIndent(ss, indent + 1);
     *ss << "isSearchMeta = " << isSearchMeta << '\n';
+    addIndent(ss, indent + 1);
+    *ss << "searchQuery = " << searchQuery << '\n';
+    if (limit) {
+        addIndent(ss, indent + 1);
+        *ss << "limit = " << limit << '\n';
+    }
 }
 
 /**
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index b2061180f31..cb8ae5dd1bf 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -1733,7 +1733,16 @@ struct SentinelNode : public QuerySolutionNode {
 };
 
 struct SearchNode : public QuerySolutionNode {
-    explicit SearchNode(bool isSearchMeta) : isSearchMeta(isSearchMeta) {
+    SearchNode() = default;
+
+    SearchNode(bool isSearchMeta,
+               BSONObj searchQuery,
+               boost::optional<long long> limit,
+               boost::optional<int> intermediateResultsProtocolVersion)
+        : isSearchMeta(isSearchMeta),
+          searchQuery(searchQuery),
+          limit(limit),
+          intermediateResultsProtocolVersion(intermediateResultsProtocolVersion) {
         // TODO SERVER-78565: Support $search in SBE plan cache
         eligibleForPlanCache = false;
     }
@@ -1766,6 +1775,22 @@ struct SearchNode : public QuerySolutionNode {
      * True for $searchMeta, False for $search query.
      */
     bool isSearchMeta;
+
+    const BSONObj searchQuery;
+
+    /**
+     * This will populate the docsRequested field of the cursorOptions document sent as part of the
+     * command to mongot in the case where the query has an extractable limit that can guide the
+     * number of documents that mongot returns to mongod.
+     */
+    boost::optional<long long> limit;
+
+    /**
+     * Protocol version if it must be communicated via the search request.
+     * If we are in a sharded environment but are targeting unsharded collection we may have a
+     * protocol version even though it should not be sent to mongot.
+     */
+    boost::optional<int> intermediateResultsProtocolVersion;
 };
 
 /**