SERVER-58712: Fix update performance on collections with multiple secondary indexes on same key

4 files changed, 86 insertions, 12 deletions

diff --git a/src/mongo/db/exec/multi_plan.cpp b/src/mongo/db/exec/multi_plan.cpp
index 0f2ff61ac66..cff149dd182 100644
--- a/src/mongo/db/exec/multi_plan.cpp
+++ b/src/mongo/db/exec/multi_plan.cpp
@@ -170,14 +170,13 @@ PlanStage::StageState MultiPlanStage::doWork(WorkingSetID* out) {
             .getPlanCache()
             ->remove(plan_cache_key_factory::make<PlanCacheKey>(*_query, collection()));
 
-        _bestPlanIdx = _backupPlanIdx;
-        _backupPlanIdx = kNoSuchPlan;
+        switchToBackupPlan();
         return _candidates[_bestPlanIdx].root->work(out);
     }
 
     if (hasBackupPlan() && PlanStage::ADVANCED == state) {
         LOGV2_DEBUG(20589, 5, "Best plan had a blocking stage, became unblocked");
-        _backupPlanIdx = kNoSuchPlan;
+        removeBackupPlan();
     }
 
     return state;
@@ -277,6 +276,7 @@ Status MultiPlanStage::pickBestPlan(PlanYieldPolicy* yieldPolicy) {
                                      *_query,
                                      std::move(ranking),
                                      _candidates);
+    removeRejectedPlans();
 
     return Status::OK();
 }
@@ -352,6 +352,54 @@ bool MultiPlanStage::workAllPlans(size_t numResults, PlanYieldPolicy* yieldPolic
     return !doneWorking;
 }
 
+void MultiPlanStage::removeRejectedPlans() {
+    // Move the best plan and the backup plan to the front of 'children'.
+    if (_bestPlanIdx != 0) {
+        std::swap(_children[_bestPlanIdx], _children[0]);
+        std::swap(_candidates[_bestPlanIdx], _candidates[0]);
+        if (_backupPlanIdx == 0) {
+            _backupPlanIdx = _bestPlanIdx;
+        }
+        _bestPlanIdx = 0;
+    }
+    size_t startIndex = 1;
+    if (_backupPlanIdx != kNoSuchPlan) {
+        if (_bestPlanIdx != 1) {
+            std::swap(_children[_backupPlanIdx], _children[1]);
+            std::swap(_candidates[_backupPlanIdx], _candidates[1]);
+            _backupPlanIdx = 1;
+        }
+        startIndex = 2;
+    }
+
+    _rejected.reserve(_children.size() - startIndex);
+    for (size_t i = startIndex; i < _children.size(); ++i) {
+        rejectPlan(i);
+    }
+    _children.resize(startIndex);
+}
+
+void MultiPlanStage::switchToBackupPlan() {
+    std::swap(_children[_backupPlanIdx], _children[_bestPlanIdx]);
+    std::swap(_candidates[_backupPlanIdx], _candidates[_bestPlanIdx]);
+    removeBackupPlan();
+}
+
+void MultiPlanStage::rejectPlan(size_t planIdx) {
+    auto rejectedPlan = std::move(_children[planIdx]);
+    if (opCtx() != nullptr) {
+        rejectedPlan->saveState();
+        rejectedPlan->detachFromOperationContext();
+    }
+    _rejected.emplace_back(std::move(rejectedPlan));
+}
+
+void MultiPlanStage::removeBackupPlan() {
+    rejectPlan(_backupPlanIdx);
+    _children.resize(1);
+    _backupPlanIdx = kNoSuchPlan;
+}
+
 bool MultiPlanStage::hasBackupPlan() const {
     return kNoSuchPlan != _backupPlanIdx;
 }
@@ -386,6 +434,9 @@ unique_ptr<PlanStageStats> MultiPlanStage::getStats() {
     for (auto&& child : _children) {
         ret->children.emplace_back(child->getStats());
     }
+    for (auto&& child : _rejected) {
+        ret->children.emplace_back(child->getStats());
+    }
     return ret;
 }
 
diff --git a/src/mongo/db/exec/multi_plan.h b/src/mongo/db/exec/multi_plan.h
index 14d8ea413af..12ab61659db 100644
--- a/src/mongo/db/exec/multi_plan.h
+++ b/src/mongo/db/exec/multi_plan.h
@@ -75,7 +75,6 @@ public:
 
     std::unique_ptr<PlanStageStats> getStats() final;
 
-
     const SpecificStats* getSpecificStats() const final;
 
     boost::optional<double> getCandidateScore(size_t candidateIdx) const;
@@ -163,6 +162,19 @@ private:
      */
     void tryYield(PlanYieldPolicy* yieldPolicy);
 
+    /**
+     * Deletes all children, except for best and backup plans.
+     *
+     * This is necessary to release any resources that rejected plans might have.
+     * For example, if multi-update can be done by scanning several indexes,
+     * it will be slowed down by rejected index scans because of index cursors
+     * that need to be reopeneed after every update.
+     */
+    void removeRejectedPlans();
+    void rejectPlan(size_t planIdx);
+    void switchToBackupPlan();
+    void removeBackupPlan();
+
     static const int kNoSuchPlan = -1;
 
     // Describes the cases in which we should write an entry for the winning plan to the plan cache.
@@ -178,6 +190,9 @@ private:
     // one-to-one with _candidates.
     std::vector<plan_ranker::CandidatePlan> _candidates;
 
+    // Rejected plans in saved and detached state.
+    std::vector<std::unique_ptr<PlanStage>> _rejected;
+
     // index into _candidates, of the winner of the plan competition
     // uses -1 / kNoSuchPlan when best plan is not (yet) known
     int _bestPlanIdx;
diff --git a/src/mongo/db/query/sbe_multi_planner.cpp b/src/mongo/db/query/sbe_multi_planner.cpp
index 35737b0502b..4fb0601563e 100644
--- a/src/mongo/db/query/sbe_multi_planner.cpp
+++ b/src/mongo/db/query/sbe_multi_planner.cpp
@@ -31,7 +31,6 @@
 
 #include "mongo/db/query/sbe_multi_planner.h"
 
-#include "mongo/db/exec/multi_plan.h"
 #include "mongo/db/exec/sbe/expressions/expression.h"
 #include "mongo/db/exec/sbe/values/bson.h"
 #include "mongo/db/query/collection_query_info.h"
diff --git a/src/mongo/dbtests/query_stage_multiplan.cpp b/src/mongo/dbtests/query_stage_multiplan.cpp
index c8f202d26fb..bafb77956db 100644
--- a/src/mongo/dbtests/query_stage_multiplan.cpp
+++ b/src/mongo/dbtests/query_stage_multiplan.cpp
@@ -180,6 +180,11 @@ unique_ptr<PlanStage> getCollScanPlan(ExpressionContext* expCtx,
     return root;
 }
 
+const PlanStage* getBestPlanRoot(const MultiPlanStage* mps) {
+    auto bestPlanIdx = mps->bestPlanIdx();
+    return bestPlanIdx ? mps->getChildren()[bestPlanIdx.get()].get() : nullptr;
+}
+
 std::unique_ptr<MultiPlanStage> runMultiPlanner(ExpressionContext* expCtx,
                                                 const NamespaceString& nss,
                                                 const CollectionPtr& coll,
@@ -189,6 +194,7 @@ std::unique_ptr<MultiPlanStage> runMultiPlanner(ExpressionContext* expCtx,
     // at least).
     unique_ptr<WorkingSet> sharedWs(new WorkingSet());
     unique_ptr<PlanStage> ixScanRoot = getIxScanPlan(expCtx, coll, sharedWs.get(), desiredFooValue);
+    const auto* ixScanRootPtr = ixScanRoot.get();
 
     // Plan 1: CollScan.
     BSONObj filterObj = BSON("foo" << desiredFooValue);
@@ -207,7 +213,7 @@ std::unique_ptr<MultiPlanStage> runMultiPlanner(ExpressionContext* expCtx,
     NoopYieldPolicy yieldPolicy(expCtx->opCtx->getServiceContext()->getFastClockSource());
     ASSERT_OK(mps->pickBestPlan(&yieldPolicy));
     ASSERT(mps->bestPlanChosen());
-    ASSERT_EQUALS(0, *mps->bestPlanIdx());
+    ASSERT_EQUALS(getBestPlanRoot(mps.get()), ixScanRootPtr);
 
     return mps;
 }
@@ -240,6 +246,8 @@ TEST_F(QueryStageMultiPlanTest, MPSCollectionScanVsHighlySelectiveIXScan) {
     unique_ptr<WorkingSet> sharedWs(new WorkingSet());
     unique_ptr<PlanStage> ixScanRoot = getIxScanPlan(_expCtx.get(), coll, sharedWs.get(), 7);
 
+    const auto* ixScanRootPtr = ixScanRoot.get();
+
     // Plan 1: CollScan with matcher.
     BSONObj filterObj = BSON("foo" << 7);
     unique_ptr<MatchExpression> filter = makeMatchExpressionFromFilter(_expCtx.get(), filterObj);
@@ -254,11 +262,7 @@ TEST_F(QueryStageMultiPlanTest, MPSCollectionScanVsHighlySelectiveIXScan) {
     mps->addPlan(createQuerySolution(), std::move(ixScanRoot), sharedWs.get());
     mps->addPlan(createQuerySolution(), std::move(collScanRoot), sharedWs.get());
 
-    // Plan 0 aka the first plan aka the index scan should be the best.
-    NoopYieldPolicy yieldPolicy(_clock);
-    ASSERT_OK(mps->pickBestPlan(&yieldPolicy));
-    ASSERT(mps->bestPlanChosen());
-    ASSERT_EQUALS(0, *mps->bestPlanIdx());
+    const auto* mpsPtr = mps.get();
 
     // Takes ownership of arguments other than 'collection'.
     auto statusWithPlanExecutor =
@@ -271,6 +275,9 @@ TEST_F(QueryStageMultiPlanTest, MPSCollectionScanVsHighlySelectiveIXScan) {
     ASSERT_OK(statusWithPlanExecutor.getStatus());
     auto exec = std::move(statusWithPlanExecutor.getValue());
 
+    ASSERT_TRUE(mpsPtr->bestPlanChosen());
+    ASSERT_EQUALS(mpsPtr->getChildren()[mpsPtr->bestPlanIdx().get()].get(), ixScanRootPtr);
+
     // Get all our results out.
     int results = 0;
     BSONObj obj;
@@ -482,6 +489,8 @@ TEST_F(QueryStageMultiPlanTest, MPSExplainAllPlans) {
 
     auto ws = std::make_unique<WorkingSet>();
     auto firstPlan = std::make_unique<MockStage>(_expCtx.get(), ws.get());
+    const auto* firstPlanPtr = firstPlan.get();
+
     auto secondPlan = std::make_unique<MockStage>(_expCtx.get(), ws.get());
 
     for (int i = 0; i < nDocs; ++i) {
@@ -517,7 +526,7 @@ TEST_F(QueryStageMultiPlanTest, MPSExplainAllPlans) {
     auto root = static_cast<MultiPlanStage*>(execImpl->getRootStage());
     ASSERT_TRUE(root->bestPlanChosen());
     // The first candidate plan should have won.
-    ASSERT_EQ(*root->bestPlanIdx(), 0);
+    ASSERT_EQ(getBestPlanRoot(root), firstPlanPtr);
 
     BSONObjBuilder bob;
     Explain::explainStages(exec.get(),