SERVER-36871 $sample can loop infinitely on orphaned data

3 files changed, 13 insertions, 0 deletions

diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp
index 23bbd3e5428..c703301205a 100644
--- a/src/mongo/db/query/plan_executor_impl.cpp
+++ b/src/mongo/db/query/plan_executor_impl.cpp
@@ -47,6 +47,7 @@
 #include "mongo/db/exec/plan_stage.h"
 #include "mongo/db/exec/plan_stats.h"
 #include "mongo/db/exec/subplan.h"
+#include "mongo/db/exec/trial_stage.h"
 #include "mongo/db/exec/working_set.h"
 #include "mongo/db/exec/working_set_common.h"
 #include "mongo/db/query/find_common.h"
@@ -273,6 +274,14 @@ Status PlanExecutorImpl::_pickBestPlan() {
         return cachedPlan->pickBestPlan(_yieldPolicy.get());
     }
 
+    // Finally, we might have an explicit TrialPhase. This specifies exactly two candidate plans,
+    // one of which is to be evaluated. If it fails the trial, then the backup plan is adopted.
+    foundStage = getStageByType(_root.get(), STAGE_TRIAL);
+    if (foundStage) {
+        TrialStage* trialStage = static_cast<TrialStage*>(foundStage);
+        return trialStage->pickBestPlan(_yieldPolicy.get());
+    }
+
     // Either we chose a plan, or no plan selection was required. In both cases,
     // our work has been successfully completed.
     return Status::OK();
diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp
index c0bc1958371..3060a06115d 100644
--- a/src/mongo/db/query/stage_builder.cpp
+++ b/src/mongo/db/query/stage_builder.cpp
@@ -376,6 +376,7 @@ PlanStage* buildStages(OperationContext* opCtx,
         case STAGE_SUBPLAN:
         case STAGE_TEXT_MATCH:
         case STAGE_TEXT_OR:
+        case STAGE_TRIAL:
         case STAGE_UNKNOWN:
         case STAGE_UPDATE: {
             mongoutils::str::stream ss;
diff --git a/src/mongo/db/query/stage_types.h b/src/mongo/db/query/stage_types.h
index bb2ce7bd0ad..d99d212124d 100644
--- a/src/mongo/db/query/stage_types.h
+++ b/src/mongo/db/query/stage_types.h
@@ -95,6 +95,9 @@ enum StageType {
     STAGE_TEXT_OR,
     STAGE_TEXT_MATCH,
 
+    // Stage for choosing between two alternate plans based on an initial trial period.
+    STAGE_TRIAL,
+
     STAGE_UNKNOWN,
 
     STAGE_UPDATE,