diff options
| -rw-r--r-- | src/mongo/db/query/get_executor.cpp | 13 | ||||
| -rw-r--r-- | src/mongo/db/query/internal_plans.cpp | 21 | ||||
| -rw-r--r-- | src/mongo/db/query/internal_plans.h | 11 |
3 files changed, 40 insertions, 5 deletions
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp index be3673369a7..84dcfb3e372 100644 --- a/src/mongo/db/query/get_executor.cpp +++ b/src/mongo/db/query/get_executor.cpp @@ -917,11 +917,14 @@ StatusWith<unique_ptr<PlanExecutor>> getExecutorUpdate(OperationContext* opCtx, request->getProj().isEmpty() && hasCollectionDefaultCollation) { LOG(2) << "Using idhack: " << redact(unparsedQuery); - PlanStage* idHackStage = new IDHackStage( - opCtx, collection, unparsedQuery["_id"].wrap(), ws.get(), descriptor); - unique_ptr<UpdateStage> root = make_unique<UpdateStage>( - opCtx, updateStageParams, ws.get(), collection, idHackStage); - return PlanExecutor::make(opCtx, std::move(ws), std::move(root), collection, policy); + // Working set 'ws' is discarded. InternalPlanner::updateWithIdHack() makes its own + // WorkingSet. + return InternalPlanner::updateWithIdHack(opCtx, + collection, + updateStageParams, + descriptor, + unparsedQuery["_id"].wrap(), + policy); } // If we're here then we don't have a parsed query, but we're also not eligible for diff --git a/src/mongo/db/query/internal_plans.cpp b/src/mongo/db/query/internal_plans.cpp index f7c3e96bce1..b09e550b863 100644 --- a/src/mongo/db/query/internal_plans.cpp +++ b/src/mongo/db/query/internal_plans.cpp @@ -36,7 +36,9 @@ #include "mongo/db/exec/delete.h" #include "mongo/db/exec/eof.h" #include "mongo/db/exec/fetch.h" +#include "mongo/db/exec/idhack.h" #include "mongo/db/exec/index_scan.h" +#include "mongo/db/exec/update.h" #include "mongo/stdx/memory.h" namespace mongo { @@ -146,6 +148,25 @@ std::unique_ptr<PlanExecutor> InternalPlanner::deleteWithIndexScan( return std::move(executor.getValue()); } +std::unique_ptr<PlanExecutor> InternalPlanner::updateWithIdHack( + OperationContext* opCtx, + Collection* collection, + const UpdateStageParams& params, + const IndexDescriptor* descriptor, + const BSONObj& key, + PlanExecutor::YieldPolicy yieldPolicy) { + auto ws = stdx::make_unique<WorkingSet>(); + + auto idHackStage = stdx::make_unique<IDHackStage>(opCtx, collection, key, ws.get(), descriptor); + auto root = + stdx::make_unique<UpdateStage>(opCtx, params, ws.get(), collection, idHackStage.release()); + + auto executor = + PlanExecutor::make(opCtx, std::move(ws), std::move(root), collection, yieldPolicy); + invariantOK(executor.getStatus()); + return std::move(executor.getValue()); +} + std::unique_ptr<PlanStage> InternalPlanner::_collectionScan(OperationContext* opCtx, WorkingSet* ws, const Collection* collection, diff --git a/src/mongo/db/query/internal_plans.h b/src/mongo/db/query/internal_plans.h index f9da8e89f70..cc9eb7f0d86 100644 --- a/src/mongo/db/query/internal_plans.h +++ b/src/mongo/db/query/internal_plans.h @@ -41,6 +41,7 @@ class OperationContext; class PlanStage; class WorkingSet; struct DeleteStageParams; +struct UpdateStageParams; /** * The internal planner is a one-stop shop for "off-the-shelf" plans. Most internal procedures @@ -110,6 +111,16 @@ public: PlanExecutor::YieldPolicy yieldPolicy, Direction direction = FORWARD); + /** + * Returns an IDHACK => UPDATE plan. + */ + static std::unique_ptr<PlanExecutor> updateWithIdHack(OperationContext* opCtx, + Collection* collection, + const UpdateStageParams& params, + const IndexDescriptor* descriptor, + const BSONObj& key, + PlanExecutor::YieldPolicy yieldPolicy); + private: /** * Returns a plan stage that can be used for a collection scan. |