diff options
25 files changed, 219 insertions, 79 deletions
diff --git a/src/mongo/db/commands/index_filter_commands_test.cpp b/src/mongo/db/commands/index_filter_commands_test.cpp index a503804ec64..b1b7e7b7e29 100644 --- a/src/mongo/db/commands/index_filter_commands_test.cpp +++ b/src/mongo/db/commands/index_filter_commands_test.cpp @@ -139,7 +139,7 @@ void addQueryShapeToPlanCache(OperationContext* opCtx, ASSERT_OK(statusWithCQ.getStatus()); std::unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); - QuerySolution qs; + QuerySolution qs{QueryPlannerParams::Options::DEFAULT}; qs.cacheData.reset(new SolutionCacheData()); qs.cacheData->tree.reset(new PlanCacheIndexTree()); std::vector<QuerySolution*> solns; diff --git a/src/mongo/db/dbhelpers.cpp b/src/mongo/db/dbhelpers.cpp index cafe10dd59c..94d5ef3e3d3 100644 --- a/src/mongo/db/dbhelpers.cpp +++ b/src/mongo/db/dbhelpers.cpp @@ -107,6 +107,7 @@ RecordId Helpers::findOne(OperationContext* opCtx, unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue()); size_t options = requireIndex ? QueryPlannerParams::NO_TABLE_SCAN : QueryPlannerParams::DEFAULT; + options = options | QueryPlannerParams::OMIT_REPL_STATE_PERMITS_READS_CHECK; auto exec = uassertStatusOK(getExecutor( opCtx, &collection, std::move(cq), PlanYieldPolicy::YieldPolicy::NO_YIELD, options)); diff --git a/src/mongo/db/dbhelpers.h b/src/mongo/db/dbhelpers.h index 9ac3164f6e0..54fe7635ce4 100644 --- a/src/mongo/db/dbhelpers.h +++ b/src/mongo/db/dbhelpers.h @@ -48,29 +48,31 @@ class QueryRequest; * all helpers assume locking is handled above them */ struct Helpers { - - /* fetch a single object from collection ns that matches query. - set your db SavedContext first. - - @param query - the query to perform. note this is the low level portion of query so - "orderby : ..." won't work. - - @param requireIndex if true, assert if no index for the query. a way to guard against - writing a slow query. - - @return true if object found - */ + /** + * Executes the given match expression ('query') and returns true if there is at least one + * one matching document. The first found matching document is returned via the 'result' output + * parameter. + * + * If 'requireIndex' is true, then this forces the query system to choose an indexed plan. An + * exception is thrown if no 'requireIndex' is set to true but no indexed plan exists. + * + * Performs the read successfully regardless of a replica set node's state, meaning that the + * node does not need to be primary or secondary. + */ static bool findOne(OperationContext* opCtx, const CollectionPtr& collection, const BSONObj& query, BSONObj& result, bool requireIndex = false); + /** + * Similar to the 'findOne()' overload above, except returns the RecordId of the first matching + * document, or a null RecordId if no such document exists. + */ static RecordId findOne(OperationContext* opCtx, const CollectionPtr& collection, const BSONObj& query, bool requireIndex); - static RecordId findOne(OperationContext* opCtx, const CollectionPtr& collection, std::unique_ptr<QueryRequest> qr, diff --git a/src/mongo/db/exec/sbe/parser/parser.cpp b/src/mongo/db/exec/sbe/parser/parser.cpp index 49bd40b3cbc..0c17e710f3d 100644 --- a/src/mongo/db/exec/sbe/parser/parser.cpp +++ b/src/mongo/db/exec/sbe/parser/parser.cpp @@ -574,7 +574,8 @@ void Parser::walkScan(AstQuery& ast) { boost::none, forward, nullptr, - getCurrentPlanNodeId()); + getCurrentPlanNodeId(), + LockAcquisitionCallback{}); } void Parser::walkParallelScan(AstQuery& ast) { @@ -642,7 +643,8 @@ void Parser::walkSeek(AstQuery& ast) { lookupSlot(ast.nodes[0]->identifier), true /* forward */, nullptr, - getCurrentPlanNodeId()); + getCurrentPlanNodeId(), + LockAcquisitionCallback{}); } void Parser::walkIndexScan(AstQuery& ast) { @@ -692,7 +694,8 @@ void Parser::walkIndexScan(AstQuery& ast) { boost::none, boost::none, nullptr, - getCurrentPlanNodeId()); + getCurrentPlanNodeId(), + LockAcquisitionCallback{}); } void Parser::walkIndexSeek(AstQuery& ast) { @@ -742,7 +745,8 @@ void Parser::walkIndexSeek(AstQuery& ast) { lookupSlot(ast.nodes[0]->identifier), lookupSlot(ast.nodes[1]->identifier), nullptr, - getCurrentPlanNodeId()); + getCurrentPlanNodeId(), + LockAcquisitionCallback{}); } void Parser::walkProject(AstQuery& ast) { diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.cpp b/src/mongo/db/exec/sbe/stages/ix_scan.cpp index a1d4fb3fc97..2fe95dc3980 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.cpp +++ b/src/mongo/db/exec/sbe/stages/ix_scan.cpp @@ -36,7 +36,6 @@ #include "mongo/db/exec/sbe/values/bson.h" #include "mongo/db/exec/trial_run_tracker.h" #include "mongo/db/index/index_access_method.h" -#include "mongo/db/repl/replication_coordinator.h" namespace mongo::sbe { IndexScanStage::IndexScanStage(const NamespaceStringOrUUID& name, @@ -49,7 +48,8 @@ IndexScanStage::IndexScanStage(const NamespaceStringOrUUID& name, boost::optional<value::SlotId> seekKeySlotLow, boost::optional<value::SlotId> seekKeySlotHigh, PlanYieldPolicy* yieldPolicy, - PlanNodeId nodeId) + PlanNodeId nodeId, + LockAcquisitionCallback lockAcquisitionCallback) : PlanStage(seekKeySlotLow ? "ixseek"_sd : "ixscan"_sd, yieldPolicy, nodeId), _name(name), _indexName(indexName), @@ -59,7 +59,8 @@ IndexScanStage::IndexScanStage(const NamespaceStringOrUUID& name, _indexKeysToInclude(indexKeysToInclude), _vars(std::move(vars)), _seekKeySlotLow(seekKeySlotLow), - _seekKeySlotHigh(seekKeySlotHigh) { + _seekKeySlotHigh(seekKeySlotHigh), + _lockAcquisitionCallback(std::move(lockAcquisitionCallback)) { // The valid state is when both boundaries, or none is set, or only low key is set. invariant((_seekKeySlotLow && _seekKeySlotHigh) || (!_seekKeySlotLow && !_seekKeySlotHigh) || (_seekKeySlotLow && !_seekKeySlotHigh)); @@ -78,7 +79,8 @@ std::unique_ptr<PlanStage> IndexScanStage::clone() const { _seekKeySlotLow, _seekKeySlotHigh, _yieldPolicy, - _commonStats.nodeId); + _commonStats.nodeId, + _lockAcquisitionCallback); } void IndexScanStage::prepare(CompileCtx& ctx) { @@ -138,9 +140,9 @@ void IndexScanStage::doRestoreState() { } _coll.emplace(_opCtx, _name); - - uassertStatusOK(repl::ReplicationCoordinator::get(_opCtx)->checkCanServeReadsFor( - _opCtx, _coll->getNss(), true)); + if (_lockAcquisitionCallback) { + _lockAcquisitionCallback(_opCtx, *_coll); + } if (_cursor) { _cursor->restore(); @@ -175,9 +177,9 @@ void IndexScanStage::open(bool reOpen) { invariant(!_cursor); invariant(!_coll); _coll.emplace(_opCtx, _name); - - uassertStatusOK(repl::ReplicationCoordinator::get(_opCtx)->checkCanServeReadsFor( - _opCtx, _coll->getNss(), true)); + if (_lockAcquisitionCallback) { + _lockAcquisitionCallback(_opCtx, *_coll); + } } else { invariant(_cursor); invariant(_coll); diff --git a/src/mongo/db/exec/sbe/stages/ix_scan.h b/src/mongo/db/exec/sbe/stages/ix_scan.h index 6023957a3d2..c3ec85ebb05 100644 --- a/src/mongo/db/exec/sbe/stages/ix_scan.h +++ b/src/mongo/db/exec/sbe/stages/ix_scan.h @@ -31,12 +31,12 @@ #include "mongo/bson/ordering.h" #include "mongo/db/db_raii.h" +#include "mongo/db/exec/sbe/stages/lock_acquisition_callback.h" #include "mongo/db/exec/sbe/stages/stages.h" #include "mongo/db/storage/record_store.h" #include "mongo/db/storage/sorted_data_interface.h" namespace mongo::sbe { - /** * A stage that iterates the entries of a collection index, starting from a bound specified by the * value in 'seekKeySlotLow' and ending (via IS_EOF) with the 'seekKeySlotHigh' bound. (An @@ -69,7 +69,8 @@ public: boost::optional<value::SlotId> seekKeySlotLow, boost::optional<value::SlotId> seekKeySlotHigh, PlanYieldPolicy* yieldPolicy, - PlanNodeId nodeId); + PlanNodeId nodeId, + LockAcquisitionCallback lockAcquisitionCallback); std::unique_ptr<PlanStage> clone() const final; @@ -102,6 +103,8 @@ private: const boost::optional<value::SlotId> _seekKeySlotLow; const boost::optional<value::SlotId> _seekKeySlotHigh; + LockAcquisitionCallback _lockAcquisitionCallback; + std::unique_ptr<value::ViewOfValueAccessor> _recordAccessor; std::unique_ptr<value::ViewOfValueAccessor> _recordIdAccessor; diff --git a/src/mongo/db/exec/sbe/stages/lock_acquisition_callback.h b/src/mongo/db/exec/sbe/stages/lock_acquisition_callback.h new file mode 100644 index 00000000000..6d5a22ae460 --- /dev/null +++ b/src/mongo/db/exec/sbe/stages/lock_acquisition_callback.h @@ -0,0 +1,45 @@ +/** + * Copyright (C) 2021-present MongoDB, Inc. + * + * This program is free software: you can redistribute it and/or modify + * it under the terms of the Server Side Public License, version 1, + * as published by MongoDB, Inc. + * + * This program is distributed in the hope that it will be useful, + * but WITHOUT ANY WARRANTY; without even the implied warranty of + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the + * Server Side Public License for more details. + * + * You should have received a copy of the Server Side Public License + * along with this program. If not, see + * <http://www.mongodb.com/licensing/server-side-public-license>. + * + * As a special exception, the copyright holders give permission to link the + * code of portions of this program with the OpenSSL library under certain + * conditions as described in each individual source file and distribute + * linked combinations including the program with the OpenSSL library. You + * must comply with the Server Side Public License in all respects for + * all of the code used other than as permitted herein. If you modify file(s) + * with this exception, you may extend this exception to your version of the + * file(s), but you are not obligated to do so. If you do not wish to do so, + * delete this exception statement from your version. If you delete this + * exception statement from all source files in the program, then also delete + * it in the license file. + */ + +#pragma once + +#include <functional> + +#include "mongo/db/db_raii.h" +#include "mongo/db/operation_context.h" + +namespace mongo::sbe { +/** + * A callback which gets called whenever a stage which accesses the storage engine (e.g. "scan", + * "seek", or "ixscan") obtains or re-obtains its AutoGet*. + */ +using LockAcquisitionCallback = + std::function<void(OperationContext*, const AutoGetCollectionForReadMaybeLockFree&)>; + +} // namespace mongo::sbe diff --git a/src/mongo/db/exec/sbe/stages/scan.cpp b/src/mongo/db/exec/sbe/stages/scan.cpp index 2b7b62240d9..69d35aa62bd 100644 --- a/src/mongo/db/exec/sbe/stages/scan.cpp +++ b/src/mongo/db/exec/sbe/stages/scan.cpp @@ -33,7 +33,6 @@ #include "mongo/db/exec/sbe/expressions/expression.h" #include "mongo/db/exec/trial_run_tracker.h" -#include "mongo/db/repl/replication_coordinator.h" #include "mongo/util/str.h" namespace mongo { @@ -47,6 +46,7 @@ ScanStage::ScanStage(const NamespaceStringOrUUID& name, bool forward, PlanYieldPolicy* yieldPolicy, PlanNodeId nodeId, + LockAcquisitionCallback lockAcquisitionCallback, ScanOpenCallback openCallback) : PlanStage(seekKeySlot ? "seek"_sd : "scan"_sd, yieldPolicy, nodeId), _name(name), @@ -56,6 +56,7 @@ ScanStage::ScanStage(const NamespaceStringOrUUID& name, _vars(std::move(vars)), _seekKeySlot(seekKeySlot), _forward(forward), + _lockAcquisitionCallback(std::move(lockAcquisitionCallback)), _openCallback(openCallback) { invariant(_fields.size() == _vars.size()); invariant(!_seekKeySlot || _forward); @@ -71,6 +72,7 @@ std::unique_ptr<PlanStage> ScanStage::clone() const { _forward, _yieldPolicy, _commonStats.nodeId, + _lockAcquisitionCallback, _openCallback); } @@ -130,9 +132,9 @@ void ScanStage::doRestoreState() { } _coll.emplace(_opCtx, _name); - - uassertStatusOK(repl::ReplicationCoordinator::get(_opCtx)->checkCanServeReadsFor( - _opCtx, _coll->getNss(), true)); + if (_lockAcquisitionCallback) { + _lockAcquisitionCallback(_opCtx, *_coll); + } if (_cursor) { const bool couldRestore = _cursor->restore(); @@ -170,9 +172,9 @@ void ScanStage::open(bool reOpen) { invariant(!_cursor); invariant(!_coll); _coll.emplace(_opCtx, _name); - - uassertStatusOK(repl::ReplicationCoordinator::get(_opCtx)->checkCanServeReadsFor( - _opCtx, _coll->getNss(), true)); + if (_lockAcquisitionCallback) { + _lockAcquisitionCallback(_opCtx, *_coll); + } } else { invariant(_cursor); invariant(_coll); diff --git a/src/mongo/db/exec/sbe/stages/scan.h b/src/mongo/db/exec/sbe/stages/scan.h index 71a3a8438e3..f4653f8f0c3 100644 --- a/src/mongo/db/exec/sbe/stages/scan.h +++ b/src/mongo/db/exec/sbe/stages/scan.h @@ -29,7 +29,7 @@ #pragma once -#include "mongo/db/db_raii.h" +#include "mongo/db/exec/sbe/stages/lock_acquisition_callback.h" #include "mongo/db/exec/sbe/stages/stages.h" #include "mongo/db/exec/sbe/values/bson.h" #include "mongo/db/storage/record_store.h" @@ -49,6 +49,7 @@ public: bool forward, PlanYieldPolicy* yieldPolicy, PlanNodeId nodeId, + LockAcquisitionCallback lockAcquisitionCallback, ScanOpenCallback openCallback = {}); std::unique_ptr<PlanStage> clone() const final; @@ -84,6 +85,7 @@ private: // run is complete, this pointer is reset to nullptr. TrialRunTracker* _tracker{nullptr}; + LockAcquisitionCallback _lockAcquisitionCallback; ScanOpenCallback _openCallback; std::unique_ptr<value::ViewOfValueAccessor> _recordAccessor; diff --git a/src/mongo/db/query/classic_stage_builder_test.cpp b/src/mongo/db/query/classic_stage_builder_test.cpp index 5d912b81d3d..64f00ebf4c2 100644 --- a/src/mongo/db/query/classic_stage_builder_test.cpp +++ b/src/mongo/db/query/classic_stage_builder_test.cpp @@ -59,7 +59,7 @@ public: * Converts a 'QuerySolutionNode' to a 'QuerySolution'. */ std::unique_ptr<QuerySolution> makeQuerySolution(std::unique_ptr<QuerySolutionNode> root) { - auto querySoln = std::make_unique<QuerySolution>(); + auto querySoln = std::make_unique<QuerySolution>(QueryPlannerParams::Options::DEFAULT); querySoln->setRoot(std::move(root)); return querySoln; } diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp index 884bbc9c311..444fc33f74f 100644 --- a/src/mongo/db/query/get_executor.cpp +++ b/src/mongo/db/query/get_executor.cpp @@ -550,7 +550,7 @@ public: "namespace"_attr = _cq->ns(), "canonicalQuery"_attr = redact(_cq->toStringShort())); - auto solution = std::make_unique<QuerySolution>(); + auto solution = std::make_unique<QuerySolution>(_plannerOptions); solution->setRoot(std::make_unique<EofNode>()); auto root = buildExecutableTree(*solution); diff --git a/src/mongo/db/query/plan_cache_test.cpp b/src/mongo/db/query/plan_cache_test.cpp index c10a1d2a2cb..e465979baaa 100644 --- a/src/mongo/db/query/plan_cache_test.cpp +++ b/src/mongo/db/query/plan_cache_test.cpp @@ -286,7 +286,7 @@ std::pair<CoreIndexInfo, std::unique_ptr<WildcardProjection>> makeWildcardUpdate */ struct GenerateQuerySolution { QuerySolution* operator()() const { - unique_ptr<QuerySolution> qs(new QuerySolution()); + unique_ptr<QuerySolution> qs(new QuerySolution(QueryPlannerParams::Options::DEFAULT)); qs->cacheData.reset(new SolutionCacheData()); qs->cacheData->solnType = SolutionCacheData::COLLSCAN_SOLN; qs->cacheData->tree.reset(new PlanCacheIndexTree()); @@ -349,7 +349,8 @@ void assertShouldNotCacheQuery(const char* queryStr) { } std::unique_ptr<QuerySolution> getQuerySolutionForCaching() { - std::unique_ptr<QuerySolution> qs = std::make_unique<QuerySolution>(); + std::unique_ptr<QuerySolution> qs = + std::make_unique<QuerySolution>(QueryPlannerParams::Options::DEFAULT); qs->cacheData = std::make_unique<SolutionCacheData>(); qs->cacheData->tree = std::make_unique<PlanCacheIndexTree>(); return qs; @@ -1136,7 +1137,7 @@ protected: // Create a CachedSolution the long way.. // QuerySolution -> PlanCacheEntry -> CachedSolution - QuerySolution qs; + QuerySolution qs{QueryPlannerParams::Options::DEFAULT}; qs.cacheData = soln.cacheData->clone(); std::vector<QuerySolution*> solutions; solutions.push_back(&qs); diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp index 6c19619ba1c..f61f7fc3264 100644 --- a/src/mongo/db/query/planner_analysis.cpp +++ b/src/mongo/db/query/planner_analysis.cpp @@ -912,7 +912,7 @@ std::unique_ptr<QuerySolution> QueryPlannerAnalysis::analyzeDataAccess( const CanonicalQuery& query, const QueryPlannerParams& params, std::unique_ptr<QuerySolutionNode> solnRoot) { - auto soln = std::make_unique<QuerySolution>(); + auto soln = std::make_unique<QuerySolution>(params.options); soln->indexFilterApplied = params.indexFiltersApplied; solnRoot->computeProperties(); diff --git a/src/mongo/db/query/query_planner.cpp b/src/mongo/db/query/query_planner.cpp index 2d30ca830c2..bfa853b37cc 100644 --- a/src/mongo/db/query/query_planner.cpp +++ b/src/mongo/db/query/query_planner.cpp @@ -179,6 +179,9 @@ string optionString(size_t options) { case QueryPlannerParams::ENUMERATE_OR_CHILDREN_LOCKSTEP: ss << "ENUMERATE_OR_CHILDREN_LOCKSTEP "; break; + case QueryPlannerParams::OMIT_REPL_STATE_PERMITS_READS_CHECK: + ss << "OMIT_REPL_STATE_PERMITS_READS_CHECK"; + break; case QueryPlannerParams::DEFAULT: MONGO_UNREACHABLE; break; diff --git a/src/mongo/db/query/query_planner_params.h b/src/mongo/db/query/query_planner_params.h index 626062fe2ec..d40eb397ff8 100644 --- a/src/mongo/db/query/query_planner_params.h +++ b/src/mongo/db/query/query_planner_params.h @@ -118,6 +118,13 @@ struct QueryPlannerParams { // is thought to be helpful in general, but particularly in cases where all children of the // $or use the same fields and have the same indexes available, as in this example. ENUMERATE_OR_CHILDREN_LOCKSTEP = 1 << 12, + + // Instructs the planner to produce a plan which will *not* check at runtime that the node's + // replica set member state allows reads. Typically, replica set members will only serve + // reads to clients if thet are in parimary or secondary state. Client reads are disallowed + // in other states, e.g. during initial sync. Internal operations, on the other hand, can + // use this flag to exempt themselves from this repl set note state requirement. + OMIT_REPL_STATE_PERMITS_READS_CHECK = 1 << 13, }; // See Options enum above. diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h index 4d3c272f7fa..6445b7d8f2f 100644 --- a/src/mongo/db/query/query_solution.h +++ b/src/mongo/db/query/query_solution.h @@ -326,7 +326,7 @@ struct QuerySolutionNodeWithSortSet : public QuerySolutionNode { */ class QuerySolution { public: - QuerySolution() = default; + explicit QuerySolution(size_t plannerOptions) : plannerOptions(plannerOptions) {} /** * Return true if this solution tree contains a node of the given 'type'. @@ -362,6 +362,18 @@ public: */ void setRoot(std::unique_ptr<QuerySolutionNode> root); + /** + * Returns true if the execution plan which is constructed from this QuerySolution should check + * that the node is eligible to serve reads prior to actually performing any reads. + */ + bool shouldCheckCanServeReads() const { + return !(plannerOptions & QueryPlannerParams::OMIT_REPL_STATE_PERMITS_READS_CHECK); + } + + // A bit vector of flags which clients to the QueryPlanner pass to control which plans are + // generated and their properties. + const size_t plannerOptions; + // There are two known scenarios in which a query solution might potentially block: // // Sort stage: diff --git a/src/mongo/db/query/query_solution_test.cpp b/src/mongo/db/query/query_solution_test.cpp index 8763b208fcd..71470e22c50 100644 --- a/src/mongo/db/query/query_solution_test.cpp +++ b/src/mongo/db/query/query_solution_test.cpp @@ -1061,7 +1061,7 @@ TEST(QuerySolutionTest, NodeIdsAssignedInPostOrderFashionStartingFromOne) { ASSERT_EQ(orNode->children[0]->nodeId(), 0u); ASSERT_EQ(orNode->children[1]->nodeId(), 0u); - auto querySolution = std::make_unique<QuerySolution>(); + auto querySolution = std::make_unique<QuerySolution>(QueryPlannerParams::Options::DEFAULT); querySolution->setRoot(std::move(orNode)); auto root = querySolution->root(); diff --git a/src/mongo/db/query/sbe_stage_builder.cpp b/src/mongo/db/query/sbe_stage_builder.cpp index 32ac09989f4..787e5746639 100644 --- a/src/mongo/db/query/sbe_stage_builder.cpp +++ b/src/mongo/db/query/sbe_stage_builder.cpp @@ -114,6 +114,18 @@ const QuerySolutionNode* getNodeByType(const QuerySolutionNode* root, StageType return nullptr; } + +sbe::LockAcquisitionCallback makeLockAcquisitionCallback(bool checkNodeCanServeReads) { + if (!checkNodeCanServeReads) { + return {}; + } + + return [](OperationContext* opCtx, const AutoGetCollectionForReadMaybeLockFree& coll) { + uassertStatusOK(repl::ReplicationCoordinator::get(opCtx)->checkCanServeReadsFor( + opCtx, coll.getNss(), true)); + }; +} + } // namespace SlotBasedStageBuilder::SlotBasedStageBuilder(OperationContext* opCtx, @@ -125,7 +137,8 @@ SlotBasedStageBuilder::SlotBasedStageBuilder(OperationContext* opCtx, : StageBuilder(opCtx, collection, cq, solution), _yieldPolicy(yieldPolicy), _data(makeRuntimeEnvironment(_opCtx, &_slotIdGenerator)), - _shardFiltererFactory(shardFiltererFactory) { + _shardFiltererFactory(shardFiltererFactory), + _lockAcquisitionCallback(makeLockAcquisitionCallback(solution.shouldCheckCanServeReads())) { // SERVER-52803: In the future if we need to gather more information from the QuerySolutionNode // tree, rather than doing one-off scans for each piece of information, we should add a formal // analysis pass here. @@ -184,7 +197,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder &_frameIdGenerator, _yieldPolicy, _data.env, - reqs.getIsTailableCollScanResumeBranch()); + reqs.getIsTailableCollScanResumeBranch(), + _lockAcquisitionCallback); if (reqs.has(kReturnKey)) { // Assign the 'returnKeySlot' to be the empty object. @@ -245,8 +259,14 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder // Index scans cannot produce an oplogTsSlot, so assert that the caller doesn't need it. invariant(!reqs.has(kOplogTs)); - return generateIndexScan( - _opCtx, _collection, ixn, reqs, &_slotIdGenerator, &_spoolIdGenerator, _yieldPolicy); + return generateIndexScan(_opCtx, + _collection, + ixn, + reqs, + &_slotIdGenerator, + &_spoolIdGenerator, + _yieldPolicy, + _lockAcquisitionCallback); } std::tuple<sbe::value::SlotId, sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> @@ -267,7 +287,8 @@ SlotBasedStageBuilder::makeLoopJoinForFetch(std::unique_ptr<sbe::PlanStage> inpu seekKeySlot, true, nullptr, - planNodeId); + planNodeId, + _lockAcquisitionCallback); // Get the recordIdSlot from the outer side (e.g., IXSCAN) and feed it to the inner side, // limiting the result set to 1 row. @@ -826,7 +847,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> SlotBasedStageBuilder boost::none, // recordSlot &_slotIdGenerator, _yieldPolicy, - root->nodeId()); + root->nodeId(), + _lockAcquisitionCallback); indexScanList.push_back(std::move(ixscan)); ixscanOutputSlots.push_back(sbe::makeSV(recordIdSlot)); } diff --git a/src/mongo/db/query/sbe_stage_builder.h b/src/mongo/db/query/sbe_stage_builder.h index 61b728f5d9b..2dd50579b76 100644 --- a/src/mongo/db/query/sbe_stage_builder.h +++ b/src/mongo/db/query/sbe_stage_builder.h @@ -30,6 +30,7 @@ #pragma once #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/lock_acquisition_callback.h" #include "mongo/db/exec/sbe/values/slot.h" #include "mongo/db/exec/sbe/values/value.h" #include "mongo/db/exec/trial_period_utils.h" @@ -337,5 +338,9 @@ private: // A factory to construct shard filters. ShardFiltererFactoryInterface* _shardFiltererFactory; + + // A callback that should be installed on "scan" and "ixscan" nodes. It will get invoked when + // these data access stages acquire their AutoGet*. + const sbe::LockAcquisitionCallback _lockAcquisitionCallback; }; } // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp b/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp index ef882219f37..832597adfc5 100644 --- a/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp +++ b/src/mongo/db/query/sbe_stage_builder_coll_scan.cpp @@ -122,7 +122,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateOptimizedOplo sbe::value::FrameIdGenerator* frameIdGenerator, PlanYieldPolicy* yieldPolicy, sbe::RuntimeEnvironment* env, - bool isTailableResumeBranch) { + bool isTailableResumeBranch, + sbe::LockAcquisitionCallback lockAcquisitionCallback) { invariant(collection->ns().isOplog()); // The minTs and maxTs optimizations are not compatible with resumeAfterRecordId and can only // be done for a forward scan. @@ -174,6 +175,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateOptimizedOplo true /* forward */, yieldPolicy, csn->nodeId(), + std::move(lockAcquisitionCallback), makeOpenCallbackIfNeeded(collection, csn)); // Start the scan from the seekRecordId if we can use the oplogStartHack. @@ -276,7 +278,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateOptimizedOplo seekRecordIdSlot, true /* forward */, yieldPolicy, - csn->nodeId()), + csn->nodeId(), + std::move(lockAcquisitionCallback)), sbe::makeSV(), sbe::makeSV(*seekRecordIdSlot), nullptr, @@ -311,7 +314,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateGenericCollSc sbe::value::FrameIdGenerator* frameIdGenerator, PlanYieldPolicy* yieldPolicy, sbe::RuntimeEnvironment* env, - bool isTailableResumeBranch) { + bool isTailableResumeBranch, + sbe::LockAcquisitionCallback lockAcquisitionCallback) { const auto forward = csn->direction == CollectionScanParams::FORWARD; invariant(!csn->shouldTrackLatestOplogTimestamp || collection->ns().isOplog()); @@ -345,6 +349,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateGenericCollSc forward, yieldPolicy, csn->nodeId(), + lockAcquisitionCallback, makeOpenCallbackIfNeeded(collection, csn)); // Check if the scan should be started after the provided resume RecordId and construct a nested @@ -376,8 +381,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateGenericCollSc seekSlot, forward, yieldPolicy, - csn->nodeId()), - + csn->nodeId(), + lockAcquisitionCallback), sbe::makeSV(seekSlot), sbe::makeSV(seekSlot), nullptr, @@ -461,7 +466,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateCollScan( sbe::value::FrameIdGenerator* frameIdGenerator, PlanYieldPolicy* yieldPolicy, sbe::RuntimeEnvironment* env, - bool isTailableResumeBranch) { + bool isTailableResumeBranch, + sbe::LockAcquisitionCallback lockAcquisitionCallback) { if (csn->minTs || csn->maxTs) { return generateOptimizedOplogScan(opCtx, collection, @@ -470,7 +476,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateCollScan( frameIdGenerator, yieldPolicy, env, - isTailableResumeBranch); + isTailableResumeBranch, + std::move(lockAcquisitionCallback)); } else { return generateGenericCollScan(opCtx, collection, @@ -479,7 +486,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateCollScan( frameIdGenerator, yieldPolicy, env, - isTailableResumeBranch); + isTailableResumeBranch, + std::move(lockAcquisitionCallback)); } } } // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_coll_scan.h b/src/mongo/db/query/sbe_stage_builder_coll_scan.h index 1a1683d986f..4d3a73f18fe 100644 --- a/src/mongo/db/query/sbe_stage_builder_coll_scan.h +++ b/src/mongo/db/query/sbe_stage_builder_coll_scan.h @@ -30,6 +30,7 @@ #pragma once #include "mongo/db/exec/sbe/expressions/expression.h" +#include "mongo/db/exec/sbe/stages/lock_acquisition_callback.h" #include "mongo/db/exec/sbe/stages/stages.h" #include "mongo/db/exec/sbe/values/value.h" #include "mongo/db/query/query_solution.h" @@ -58,6 +59,7 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateCollScan( sbe::value::FrameIdGenerator* frameIdGenerator, PlanYieldPolicy* yieldPolicy, sbe::RuntimeEnvironment* env, - bool isTailableResumeBranch); + bool isTailableResumeBranch, + sbe::LockAcquisitionCallback lockAcquisitionCallback); } // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_index_scan.cpp b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp index 3512a24296f..3a460cf03f8 100644 --- a/src/mongo/db/query/sbe_stage_builder_index_scan.cpp +++ b/src/mongo/db/query/sbe_stage_builder_index_scan.cpp @@ -277,7 +277,8 @@ generateOptimizedMultiIntervalIndexScan( sbe::value::SlotVector indexKeySlots, sbe::value::SlotIdGenerator* slotIdGenerator, PlanYieldPolicy* yieldPolicy, - PlanNodeId planNodeId) { + PlanNodeId planNodeId, + sbe::LockAcquisitionCallback lockAcquisitionCallback) { using namespace std::literals; auto recordIdSlot = slotIdGenerator->generate(); @@ -344,7 +345,8 @@ generateOptimizedMultiIntervalIndexScan( lowKeySlot, highKeySlot, yieldPolicy, - planNodeId); + planNodeId, + std::move(lockAcquisitionCallback)); // Finally, get the keys from the outer side and feed them to the inner side (ixscan). return {recordIdSlot, @@ -397,7 +399,8 @@ makeRecursiveBranchForGenericIndexScan(const CollectionPtr& collection, sbe::value::SlotVector savedIndexKeySlots, sbe::value::SlotIdGenerator* slotIdGenerator, PlanYieldPolicy* yieldPolicy, - PlanNodeId planNodeId) { + PlanNodeId planNodeId, + sbe::LockAcquisitionCallback lockAcquisitionCallback) { auto resultSlot = slotIdGenerator->generate(); auto recordIdSlot = slotIdGenerator->generate(); @@ -425,7 +428,8 @@ makeRecursiveBranchForGenericIndexScan(const CollectionPtr& collection, lowKeySlot, boost::none, yieldPolicy, - planNodeId); + planNodeId, + std::move(lockAcquisitionCallback)); // Get the low key from the outer side and feed it to the inner side (ixscan). auto nlj = sbe::makeS<sbe::LoopJoinStage>(std::move(project), @@ -522,7 +526,8 @@ generateGenericMultiIntervalIndexScan(const CollectionPtr& collection, sbe::value::SlotVector indexKeySlots, sbe::value::SlotIdGenerator* slotIdGenerator, sbe::value::SpoolIdGenerator* spoolIdGenerator, - PlanYieldPolicy* yieldPolicy) { + PlanYieldPolicy* yieldPolicy, + sbe::LockAcquisitionCallback lockAcquisitionCallback) { using namespace std::literals; @@ -572,7 +577,8 @@ generateGenericMultiIntervalIndexScan(const CollectionPtr& collection, savedIndexKeySlots, slotIdGenerator, yieldPolicy, - ixn->nodeId()); + ixn->nodeId(), + std::move(lockAcquisitionCallback)); // Construct a union stage from the two branches. auto makeSlotVector = [](sbe::value::SlotId headSlot, const sbe::value::SlotVector& varSlots) { @@ -624,7 +630,8 @@ std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateSingleInt boost::optional<sbe::value::SlotId> recordSlot, sbe::value::SlotIdGenerator* slotIdGenerator, PlanYieldPolicy* yieldPolicy, - PlanNodeId planNodeId) { + PlanNodeId planNodeId, + sbe::LockAcquisitionCallback lockAcquisitionCallback) { auto recordIdSlot = slotIdGenerator->generate(); auto lowKeySlot = slotIdGenerator->generate(); auto highKeySlot = slotIdGenerator->generate(); @@ -656,7 +663,8 @@ std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateSingleInt lowKeySlot, highKeySlot, yieldPolicy, - planNodeId); + planNodeId, + std::move(lockAcquisitionCallback)); // Finally, get the keys from the outer side and feed them to the inner side. return {recordIdSlot, @@ -675,7 +683,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScan( PlanStageReqs reqs, sbe::value::SlotIdGenerator* slotIdGenerator, sbe::value::SpoolIdGenerator* spoolIdGenerator, - PlanYieldPolicy* yieldPolicy) { + PlanYieldPolicy* yieldPolicy, + sbe::LockAcquisitionCallback lockAcquisitionCallback) { uassert(4822864, "Index scans with a filter are not supported in SBE", !ixn->filter); auto descriptor = @@ -737,7 +746,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScan( boost::none, // recordSlot slotIdGenerator, yieldPolicy, - ixn->nodeId()); + ixn->nodeId(), + std::move(lockAcquisitionCallback)); outputs.set(PlanStageSlots::kRecordId, recordIdSlot); } else if (intervals.size() > 1) { @@ -753,7 +763,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScan( indexKeySlots, slotIdGenerator, yieldPolicy, - ixn->nodeId()); + ixn->nodeId(), + std::move(lockAcquisitionCallback)); outputs.set(PlanStageSlots::kRecordId, recordIdSlot); } else { @@ -768,7 +779,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScan( indexKeySlots, slotIdGenerator, spoolIdGenerator, - yieldPolicy); + yieldPolicy, + std::move(lockAcquisitionCallback)); outputs.set(PlanStageSlots::kRecordId, recordIdSlot); } diff --git a/src/mongo/db/query/sbe_stage_builder_index_scan.h b/src/mongo/db/query/sbe_stage_builder_index_scan.h index ab54dfa02cf..bdfa8661291 100644 --- a/src/mongo/db/query/sbe_stage_builder_index_scan.h +++ b/src/mongo/db/query/sbe_stage_builder_index_scan.h @@ -29,6 +29,7 @@ #pragma once +#include "mongo/db/exec/sbe/stages/lock_acquisition_callback.h" #include "mongo/db/exec/sbe/stages/stages.h" #include "mongo/db/exec/sbe/values/value.h" #include "mongo/db/query/query_solution.h" @@ -55,7 +56,8 @@ std::pair<std::unique_ptr<sbe::PlanStage>, PlanStageSlots> generateIndexScan( PlanStageReqs reqs, sbe::value::SlotIdGenerator* slotIdGenerator, sbe::value::SpoolIdGenerator* spoolIdGenerator, - PlanYieldPolicy* yieldPolicy); + PlanYieldPolicy* yieldPolicy, + sbe::LockAcquisitionCallback lockAcquisitionCallback); /** * Constructs the most simple version of an index scan from the single interval index bounds. The @@ -86,6 +88,7 @@ std::pair<sbe::value::SlotId, std::unique_ptr<sbe::PlanStage>> generateSingleInt boost::optional<sbe::value::SlotId> recordSlot, sbe::value::SlotIdGenerator* slotIdGenerator, PlanYieldPolicy* yieldPolicy, - PlanNodeId nodeId); + PlanNodeId nodeId, + sbe::LockAcquisitionCallback lockAcquisitionCallback); } // namespace mongo::stage_builder diff --git a/src/mongo/db/query/sbe_stage_builder_test_fixture.cpp b/src/mongo/db/query/sbe_stage_builder_test_fixture.cpp index efb6c08639f..d832a3e30b6 100644 --- a/src/mongo/db/query/sbe_stage_builder_test_fixture.cpp +++ b/src/mongo/db/query/sbe_stage_builder_test_fixture.cpp @@ -39,7 +39,7 @@ namespace mongo { std::unique_ptr<QuerySolution> SbeStageBuilderTestFixture::makeQuerySolution( std::unique_ptr<QuerySolutionNode> root) { - auto querySoln = std::make_unique<QuerySolution>(); + auto querySoln = std::make_unique<QuerySolution>(QueryPlannerParams::Options::DEFAULT); querySoln->setRoot(std::move(root)); return querySoln; } diff --git a/src/mongo/dbtests/query_stage_multiplan.cpp b/src/mongo/dbtests/query_stage_multiplan.cpp index 3275bcaed17..7a2ee9aef72 100644 --- a/src/mongo/dbtests/query_stage_multiplan.cpp +++ b/src/mongo/dbtests/query_stage_multiplan.cpp @@ -74,7 +74,7 @@ using std::vector; static const NamespaceString nss("unittests.QueryStageMultiPlan"); std::unique_ptr<QuerySolution> createQuerySolution() { - auto soln = std::make_unique<QuerySolution>(); + auto soln = std::make_unique<QuerySolution>(QueryPlannerParams::Options::DEFAULT); soln->cacheData = std::make_unique<SolutionCacheData>(); soln->cacheData->solnType = SolutionCacheData::COLLSCAN_SOLN; soln->cacheData->tree = std::make_unique<PlanCacheIndexTree>(); @@ -492,8 +492,12 @@ TEST_F(QueryStageMultiPlanTest, MPSExplainAllPlans) { std::make_unique<MultiPlanStage>(_expCtx.get(), ctx.getCollection(), cq.get()); // Put each plan into the MultiPlanStage. Takes ownership of 'firstPlan' and 'secondPlan'. - mps->addPlan(std::make_unique<QuerySolution>(), std::move(firstPlan), ws.get()); - mps->addPlan(std::make_unique<QuerySolution>(), std::move(secondPlan), ws.get()); + mps->addPlan(std::make_unique<QuerySolution>(QueryPlannerParams::Options::DEFAULT), + std::move(firstPlan), + ws.get()); + mps->addPlan(std::make_unique<QuerySolution>(QueryPlannerParams::Options::DEFAULT), + std::move(secondPlan), + ws.get()); // Making a PlanExecutor chooses the best plan. auto exec = |