SERVER-75369 Fix the overflow of total possible enumeration count in LockstepOr enumeration strategy

(cherry picked from commit 6bb0a27e70d7484f4dc9532c35dc1f0c7be873a3)

3 files changed, 37 insertions, 6 deletions

diff --git a/src/mongo/db/query/plan_enumerator.cpp b/src/mongo/db/query/plan_enumerator.cpp
index a11bc9b415a..96fc70fb99c 100644
--- a/src/mongo/db/query/plan_enumerator.cpp
+++ b/src/mongo/db/query/plan_enumerator.cpp
@@ -1689,8 +1689,8 @@ bool PlanEnumerator::LockstepOrAssignment::allIdentical() const {
     return true;
 }
 
-bool PlanEnumerator::LockstepOrAssignment::shouldResetBeforeProceeding(
-    size_t totalEnumerated) const {
+bool PlanEnumerator::LockstepOrAssignment::shouldResetBeforeProceeding(size_t totalEnumerated,
+                                                                       size_t orLimit) const {
     if (totalEnumerated == 0 || !exhaustedLockstepIteration) {
         return false;
     }
@@ -1700,7 +1700,12 @@ bool PlanEnumerator::LockstepOrAssignment::shouldResetBeforeProceeding(
         if (!subnode.maxIterCount) {
             return false;  // Haven't yet looped over this child entirely, not ready yet.
         }
-        totalPossibleEnumerations *= subnode.maxIterCount.get();
+        totalPossibleEnumerations *= subnode.maxIterCount.value();
+        // If 'totalPossibleEnumerations' reaches the limit, we can just shortcut it. Otherwise,
+        // 'totalPossibleEnumerations' could overflow if we have a large $or.
+        if (totalPossibleEnumerations >= orLimit) {
+            return false;
+        }
     }
 
     // If we're able to compute a total number expected enumerations, we must have already cycled
@@ -1737,7 +1742,7 @@ bool PlanEnumerator::_nextMemoForLockstepOrAssignment(
         }
         // Edge case: if every child has only one option available, we are already finished
         // enumerating.
-        if (assignment->shouldResetBeforeProceeding(assignment->totalEnumerated)) {
+        if (assignment->shouldResetBeforeProceeding(assignment->totalEnumerated, _orLimit)) {
             assignment->exhaustedLockstepIteration = false;
             return true;  // We're back at the beginning, no need to reset.
         }
@@ -1776,7 +1781,7 @@ bool PlanEnumerator::_nextMemoForLockstepOrAssignment(
     // This special ordering is tricky to reset. Because it iterates the sub nodes in such a
     // unique order, it can be difficult to know when it has actually finished iterating. Our
     // strategy is just to compute a total and go back to the beginning once we hit that total.
-    if (!assignment->shouldResetBeforeProceeding(assignment->totalEnumerated)) {
+    if (!assignment->shouldResetBeforeProceeding(assignment->totalEnumerated, _orLimit)) {
         return false;
     }
     // Reset!
diff --git a/src/mongo/db/query/plan_enumerator.h b/src/mongo/db/query/plan_enumerator.h
index 60344f0c9ee..b82b738c57b 100644
--- a/src/mongo/db/query/plan_enumerator.h
+++ b/src/mongo/db/query/plan_enumerator.h
@@ -244,7 +244,7 @@ private:
          * Returns true if 'totalEnumerated' matches the total number of expected plans for this
          * assignment.
          */
-        bool shouldResetBeforeProceeding(size_t totalEnumerated) const;
+        bool shouldResetBeforeProceeding(size_t totalEnumerated, size_t orLimit) const;
 
         /**
          * Returns true if each sub node is at the same iterationCount.
diff --git a/src/mongo/db/query/query_planner_tree_test.cpp b/src/mongo/db/query/query_planner_tree_test.cpp
index 9d403989792..853c02388b0 100644
--- a/src/mongo/db/query/query_planner_tree_test.cpp
+++ b/src/mongo/db/query/query_planner_tree_test.cpp
@@ -2464,6 +2464,32 @@ TEST_F(QueryPlannerTest, LockstepOrEnumerationSanityCheckTwoChildrenTwoIndexesEa
         "{ixscan: {pattern: {a: 1, c: 1}}}}}}}");
 }
 
+TEST_F(QueryPlannerTest, TotalPossibleLockstepOrEnumerationReachesTheOrLimit) {
+    params.options =
+        QueryPlannerParams::NO_TABLE_SCAN | QueryPlannerParams::ENUMERATE_OR_CHILDREN_LOCKSTEP;
+    addIndex(BSON("a" << 1 << "b" << 1));
+    addIndex(BSON("a" << 1 << "c" << 1));
+
+    BSONArrayBuilder orBuilder;
+    // This max number has a value of 65 in order to potentillay triger any overflow of the possible
+    // enumeration count, because each predicate in $or has two possible indexes, allowing for 2^65
+    // possible enumerations.
+    const int maxPredicates = 65;
+    for (int i = 0; i < maxPredicates; i++) {
+        orBuilder.append(BSON("b" << i << "c" << i));
+    }
+
+    auto cmd = BSON("find"
+                    << "testns"
+                    << "filter" << BSON("a" << 1 << "$or" << orBuilder.arr()));
+
+    // Ensure that the query runs fine.
+    runQueryAsCommand(cmd);
+
+    // internalQueryMaxOrSolutions.load() + 2.
+    assertNumSolutions(12U);
+}
+
 // Test that we enumerate the expected plans with the special parameter set. In this test we have
 // two branches of an $or, each with one possible indexed solution.
 TEST_F(QueryPlannerTest, LockstepOrEnumerationSanityCheckTwoChildrenOneIndexEach) {