SERVER-63102 Introduce internalQueryPlanEvaluationWorksSbe

The 'internalQueryPlanEvaluationWorks' parameter now only
affects the classic engine. The newly added parameter has
similar behavior, but applies only to queries using SBE.

(cherry picked from commit a94caa502cf94fa6c8fcfea7283d7eaf3bd55ad5)

9 files changed, 80 insertions, 32 deletions

diff --git a/jstests/noPassthrough/explain_reports_correct_trial_period_statistics.js b/jstests/noPassthrough/explain_reports_correct_trial_period_statistics.js
index 05289cfde86..eafce2c9a62 100644
--- a/jstests/noPassthrough/explain_reports_correct_trial_period_statistics.js
+++ b/jstests/noPassthrough/explain_reports_correct_trial_period_statistics.js
@@ -18,6 +18,7 @@ assert.commandWorked(coll.createIndex({b: 1}));
 
 // Configure the server such that the trial period should end after doing 10 reads from storage.
 assert.commandWorked(db.adminCommand({setParameter: 1, internalQueryPlanEvaluationWorks: 10}));
+assert.commandWorked(db.adminCommand({setParameter: 1, internalQueryPlanEvaluationWorksSbe: 10}));
 
 assert.commandWorked(coll.insert(Array.from({length: 20}, (v, i) => {
     return {a: 1, b: 1, c: i};
diff --git a/jstests/noPassthrough/query_knobs_validation.js b/jstests/noPassthrough/query_knobs_validation.js
index 21376673f84..37c6fdabc92 100644
--- a/jstests/noPassthrough/query_knobs_validation.js
+++ b/jstests/noPassthrough/query_knobs_validation.js
@@ -12,6 +12,7 @@ const conn = MongoRunner.runMongod();
 const testDB = conn.getDB("admin");
 const expectedParamDefaults = {
     internalQueryPlanEvaluationWorks: 10000,
+    internalQueryPlanEvaluationWorksSbe: 10000,
     internalQueryPlanEvaluationCollFraction: 0.3,
     internalQueryPlanEvaluationCollFractionSbe: 0.0,
     internalQueryPlanEvaluationMaxResults: 101,
@@ -85,9 +86,11 @@ assert.eq(getParamRes["internalPipelineLengthLimit"],
 // Verify that the default values are set as expected when the server starts up.
 assertDefaultParameterValues();
 
-assertSetParameterSucceeds("internalQueryPlanEvaluationWorks", 11);
-assertSetParameterFails("internalQueryPlanEvaluationWorks", 0);
-assertSetParameterFails("internalQueryPlanEvaluationWorks", -1);
+for (let paramName of ["internalQueryPlanEvaluationWorks", "internalQueryPlanEvaluationWorksSbe"]) {
+    assertSetParameterSucceeds(paramName, 11);
+    assertSetParameterFails(paramName, 0);
+    assertSetParameterFails(paramName, -1);
+}
 
 for (let paramName of ["internalQueryPlanEvaluationCollFraction",
                        "internalQueryPlanEvaluationCollFractionSbe"]) {
diff --git a/jstests/noPassthrough/sbe_multiplanner_trial_termination.js b/jstests/noPassthrough/sbe_multiplanner_trial_termination.js
index 8918df3bdb1..75739587a9f 100644
--- a/jstests/noPassthrough/sbe_multiplanner_trial_termination.js
+++ b/jstests/noPassthrough/sbe_multiplanner_trial_termination.js
@@ -12,6 +12,7 @@ const collName = "sbe_multiplanner_coll";
 const collFracKnob = "internalQueryPlanEvaluationCollFraction";
 const collFracKnobSbe = "internalQueryPlanEvaluationCollFractionSbe";
 const worksKnob = "internalQueryPlanEvaluationWorks";
+const worksKnobSbe = "internalQueryPlanEvaluationWorksSbe";
 
 const defaultCollFrac = 0.3;
 const trialLengthFromCollFrac = defaultCollFrac * numDocs;
@@ -46,7 +47,7 @@ for (let i = 0; i < numDocs; ++i) {
 // collection. Since the classic multiplanner takes either the works limit or 30% of the collection
 // size -- whichever is larger -- this should cause the trial period to run for about 0.3 * numDocs
 // work cycles.
-const getParamRes = assert.commandWorked(db.adminCommand({getParameter: 1, [collFracKnob]: 1}));
+let getParamRes = assert.commandWorked(db.adminCommand({getParameter: 1, [collFracKnob]: 1}));
 assert.eq(getParamRes[collFracKnob], defaultCollFrac);
 assert.commandWorked(db.adminCommand({setParameter: 1, [worksKnob]: trialLengthFromWorksKnob}));
 
@@ -62,22 +63,42 @@ for (let plan of allPlans) {
     assert.eq(executionStages.works, trialLengthFromCollFrac, plan);
 }
 
-// Verifies that for each SBE plan in the 'allPlans' array, the number of storage reads done by the
-// plan is equal to 'expectedNumReads'.
-function verifySbeNumReads(allPlans, expectedNumReads) {
+// For each SBE plan in the 'allPlans' array, verifies the number of storage reads
+// done by the plan with respect to 'expectedNumReads' by calling 'assertionFn(actualNumReads,
+// expectedNumReads)'.
+//
+// By default, 'assertionFn' is 'assert.eq()' -- meaning that the number of storage cursor reads
+// done by each candidate SBE plan is checked exactly -- but the caller can pass a different
+// assertion function to override this behavior.
+function verifySbeNumReads(allPlans, expectedNumReads, assertionFn = assert.eq) {
     for (let plan of allPlans) {
         // Infer the number of reads (SBE's equivalent of work units) as the sum of keys and
         // documents examined.
         assert(plan.hasOwnProperty("totalKeysExamined"), plan);
         assert(plan.hasOwnProperty("totalDocsExamined"), plan);
         const numReads = plan.totalKeysExamined + plan.totalDocsExamined;
-        assert.eq(numReads, expectedNumReads, plan);
+        assertionFn(numReads, expectedNumReads, plan);
     }
 }
 
-// Allow the query to use SBE. This time, the trial period should terminate based on the works knob.
+// Verify the default values of the SBE-specific knobs.
+getParamRes = assert.commandWorked(db.adminCommand({getParameter: 1, [collFracKnobSbe]: 1}));
+assert.eq(getParamRes[collFracKnobSbe], 0.0);
+getParamRes = assert.commandWorked(db.adminCommand({getParameter: 1, [worksKnobSbe]: 1}));
+assert.gt(getParamRes[worksKnobSbe], numDocs);
+
+// Allow the query to use SBE. Since we haven't modified any SBE-specific knobs yet, we expect the
+// length of the trial period to be determined by the default value of the SBE works knob. Since the
+// default value of SBE's works knob exceeds the size of the collection, we expect the number of
+// reads to exceed the collection size as well. By construction of the test, this also means that
+// the trial period length exceeds both 'trialLengthFromCollFrac' and 'trialLengthFromWorksKnob'.
 assert.commandWorked(db.adminCommand({setParameter: 1, internalQueryForceClassicEngine: false}));
 allPlans = getAllPlansExecution("2");
+verifySbeNumReads(allPlans, numDocs, assert.gt);
+
+// Setting the SBE works knob lower will reduce the length of the trial period.
+assert.commandWorked(db.adminCommand({setParameter: 1, [worksKnobSbe]: trialLengthFromWorksKnob}));
+allPlans = getAllPlansExecution("2");
 verifySbeNumReads(allPlans, trialLengthFromWorksKnob);
 
 // If the SBE "collection fraction" knob is set to the same value as the equivalent knob for the
diff --git a/src/mongo/db/exec/multi_plan.cpp b/src/mongo/db/exec/multi_plan.cpp
index 892d05cb5ca..96aaf4ca52d 100644
--- a/src/mongo/db/exec/multi_plan.cpp
+++ b/src/mongo/db/exec/multi_plan.cpp
@@ -163,8 +163,11 @@ Status MultiPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) {
     // make sense.
     auto optTimer = getOptTimer();
 
-    size_t numWorks = trial_period::getTrialPeriodMaxWorks(
-        opCtx(), collection(), internalQueryPlanEvaluationCollFraction.load());
+    const size_t numWorks =
+        trial_period::getTrialPeriodMaxWorks(opCtx(),
+                                             collection(),
+                                             internalQueryPlanEvaluationWorks.load(),
+                                             internalQueryPlanEvaluationCollFraction.load());
     size_t numResults = trial_period::getTrialPeriodNumToReturn(*_query);
 
     try {
diff --git a/src/mongo/db/exec/trial_period_utils.cpp b/src/mongo/db/exec/trial_period_utils.cpp
index 2472885c7ed..b158f297b47 100644
--- a/src/mongo/db/exec/trial_period_utils.cpp
+++ b/src/mongo/db/exec/trial_period_utils.cpp
@@ -36,13 +36,12 @@
 namespace mongo::trial_period {
 size_t getTrialPeriodMaxWorks(OperationContext* opCtx,
                               const CollectionPtr& collection,
+                              int maxWorksParam,
                               double collFraction) {
-    // Run each plan some number of times. This number is at least as great as
-    // 'internalQueryPlanEvaluationWorks', but may be larger for big collections.
-    size_t numWorks = internalQueryPlanEvaluationWorks.load();
+    size_t numWorks = static_cast<size_t>(maxWorksParam);
     if (collection) {
-        numWorks = std::max(static_cast<size_t>(internalQueryPlanEvaluationWorks.load()),
-                            static_cast<size_t>(collFraction * collection->numRecords(opCtx)));
+        numWorks =
+            std::max(numWorks, static_cast<size_t>(collFraction * collection->numRecords(opCtx)));
     }
 
     return numWorks;
diff --git a/src/mongo/db/exec/trial_period_utils.h b/src/mongo/db/exec/trial_period_utils.h
index 53a9c91a889..f8e4d4a725f 100644
--- a/src/mongo/db/exec/trial_period_utils.h
+++ b/src/mongo/db/exec/trial_period_utils.h
@@ -39,11 +39,14 @@ namespace trial_period {
 /**
  * Returns the number of times that we are willing to work a plan during a trial period.
  *
- * Calculated based on a fixed query knob and the size of the collection multiplied by
- * 'collFraction'.
+ * Calculated with the following formula, where "|collection|" denotes the approximate number of
+ * documents in the collection:
+ *
+ *   max(maxWorksParam, collFraction * |collection|)
  */
 size_t getTrialPeriodMaxWorks(OperationContext* opCtx,
                               const CollectionPtr& collection,
+                              int maxWorksParam,
                               double collFraction);
 
 /**
diff --git a/src/mongo/db/query/query_knobs.idl b/src/mongo/db/query/query_knobs.idl
index cbaa8310d2d..fea0273c3fe 100644
--- a/src/mongo/db/query/query_knobs.idl
+++ b/src/mongo/db/query/query_knobs.idl
@@ -39,7 +39,10 @@ server_parameters:
   # multi-plan ranking
   #
   internalQueryPlanEvaluationWorks:
-    description: "Max number of times we call work() on plans before comparing them, for small collections."
+    description: "For small collections, the max number of times we call work() on plans during the
+    runtime plan selection trial period. Applies only to the classic execution engine. The complete
+    formula for calculating the maximum trial period works depends on
+    'internalQueryPlanEvaluationCollFraction' as well as this parameter."
     set_at: [ startup, runtime ]
     cpp_varname: "internalQueryPlanEvaluationWorks"
     cpp_vartype: AtomicWord<int>
@@ -47,6 +50,19 @@ server_parameters:
     validator:
       gt: 0
 
+  internalQueryPlanEvaluationWorksSbe:
+    description: "The maximum number of individual storage cursor reads performed by any candidate
+    plan during the runtime plan selection trial period. Applies only to queries using the SBE
+    execution engine. This is the analog of the 'internalQueryPlanEvaluationWorks' knob above but
+    for SBE. When 'internalQueryPlanEvaluationCollFractionSbe' has its default value of 0, this
+    parameter acts as a hard limit on trial period length regardless of collection size."
+    set_at: [ startup, runtime ]
+    cpp_varname: "internalQueryPlanEvaluationWorksSbe"
+    cpp_vartype: AtomicWord<int>
+    default: 10000
+    validator:
+      gt: 0
+
   internalQueryPlanEvaluationCollFraction:
     description: "For large collections, the ceiling for the number times we work() candidate plans
     is taken as this fraction of the collection size. Applies only to the classic execution engine."
diff --git a/src/mongo/db/query/sbe_multi_planner.cpp b/src/mongo/db/query/sbe_multi_planner.cpp
index 8f7735e0a39..e8ddbaa1232 100644
--- a/src/mongo/db/query/sbe_multi_planner.cpp
+++ b/src/mongo/db/query/sbe_multi_planner.cpp
@@ -49,8 +49,10 @@ CandidatePlans MultiPlanner::plan(
     auto candidates = collectExecutionStats(
         std::move(solutions),
         std::move(roots),
-        trial_period::getTrialPeriodMaxWorks(
-            _opCtx, _collection, internalQueryPlanEvaluationCollFractionSbe.load()));
+        trial_period::getTrialPeriodMaxWorks(_opCtx,
+                                             _collection,
+                                             internalQueryPlanEvaluationWorksSbe.load(),
+                                             internalQueryPlanEvaluationCollFractionSbe.load()));
     auto decision = uassertStatusOK(mongo::plan_ranker::pickBestPlan<PlanStageStats>(candidates));
     return finalizeExecutionPlans(std::move(decision), std::move(candidates));
 }
diff --git a/src/mongo/dbtests/plan_ranking.cpp b/src/mongo/dbtests/plan_ranking.cpp
index 076e0a80abd..a503185daf2 100644
--- a/src/mongo/dbtests/plan_ranking.cpp
+++ b/src/mongo/dbtests/plan_ranking.cpp
@@ -211,17 +211,13 @@ public:
     }
 
     void run() {
-        // We get the number of works done during the trial period in order to make sure that there
-        // are more documents in the collection than works done in the trial period. This ensures
-        // neither of the plans reach EOF or produce results.
-        size_t numWorks = trial_period::getTrialPeriodMaxWorks(opCtx(), nullptr, 0);
-        size_t smallNumber = 10;
-        // The following condition must be met in order for the following test to work. Specifically
-        // this condition guarantees that the score of the plan using the index on d will score
-        // higher than the the plan using the index on a.
-        ASSERT(smallNumber < numWorks);
-        for (size_t i = 0; i < numWorks * 2; ++i) {
-            insert(BSON("a" << static_cast<int>(i >= ((numWorks * 2) - smallNumber)) << "d"
+        const size_t numDocs = 1000;
+        const size_t smallNumber = 10;
+
+        // Insert 'numDocs' documents. A number of documents given by 'smallNumber' should have
+        // a==1, while all other docs have a==0.
+        for (size_t i = 0; i < numDocs; ++i) {
+            insert(BSON("a" << static_cast<int>(i >= (numDocs - smallNumber)) << "d"
                             << static_cast<int>(i)));
         }
 
@@ -230,6 +226,10 @@ public:
         addIndex(BSON("a" << 1));
         addIndex(BSON("d" << 1));
 
+        // Run a query where we expect the most efficient plan to fail due to exhausting the
+        // blocking sort memory limit during multi-planning. We expect this error to be swallowed
+        // and the less efficient plan using index {d: 1} to be chosen instead.
+        //
         // Query: find({a: 1}).sort({d: 1})
         auto findCommand = std::make_unique<FindCommandRequest>(nss);
         findCommand->setFilter(BSON("a" << 1));