diff options
Diffstat (limited to 'src/mongo/db/query')
120 files changed, 12122 insertions, 857 deletions
diff --git a/src/mongo/db/query/SConscript b/src/mongo/db/query/SConscript index d23c5871aa7..d6da5af72e4 100644 --- a/src/mongo/db/query/SConscript +++ b/src/mongo/db/query/SConscript @@ -1,6 +1,9 @@ # -*- mode: python -*- -Import("env") +Import([ + "env", + "get_option", +]) env = env.Clone() @@ -10,6 +13,8 @@ env.SConscript( "collation", "datetime", 'optimizer', + 'query_stats', + 'query_shape', ], exports=[ 'env' @@ -19,8 +24,11 @@ env.SConscript( env.Library( target='canonical_query', source=[ - "canonical_query.cpp", - "canonical_query_encoder.cpp", + 'canonical_query.cpp', + 'canonical_query_encoder.cpp', + 'parsed_find_command.cpp', + 'query_shape/find_cmd_shape.cpp', + 'query_stats/find_key.cpp', ], LIBDEPS=[ "$BUILD_DIR/mongo/crypto/encrypted_field_config", @@ -53,6 +61,7 @@ env.Library( "query_planner.cpp", "query_settings.cpp", "query_solution.cpp", + "record_id_range.cpp", "stage_types.cpp", ], LIBDEPS=[ @@ -90,6 +99,17 @@ env.Library( ) env.Library( + target='memory_util', + source=[ + 'util/memory_util.cpp', + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/util/processinfo', + '$BUILD_DIR/mongo/util/regex_util', + ], +) + +env.Library( target="query_plan_cache", source=[ "classic_plan_cache.cpp", @@ -101,6 +121,7 @@ env.Library( "$BUILD_DIR/mongo/base", "$BUILD_DIR/mongo/db/exec/sbe/query_sbe", "canonical_query", + "memory_util", ] ) @@ -261,17 +282,18 @@ env.Library( env.Library( target="query_knobs", source=[ - 'plan_cache_size_parameter.cpp', 'query_feature_flags.idl', 'query_knobs.idl', + 'query_stats/query_stats_on_parameter_change.cpp', 'sbe_plan_cache_on_parameter_change.cpp', ], LIBDEPS_PRIVATE=[ - "$BUILD_DIR/mongo/db/service_context", + '$BUILD_DIR/mongo/db/service_context', '$BUILD_DIR/mongo/idl/feature_flag', '$BUILD_DIR/mongo/idl/server_parameter', '$BUILD_DIR/third_party/shim_pcrecpp', - ] + 'memory_util', + ], ) env.Library( @@ -348,7 +370,7 @@ env.Library( '$BUILD_DIR/mongo/db/concurrency/exception_util', '$BUILD_DIR/mongo/db/storage/recovery_unit_base', ], - ) +) env.CppUnitTest( target="db_query_test", @@ -375,10 +397,8 @@ env.CppUnitTest( "interval_test.cpp", "killcursors_request_test.cpp", "lru_key_value_test.cpp", - 'map_reduce_output_format_test.cpp', "parsed_distinct_test.cpp", "plan_cache_indexability_test.cpp", - "plan_cache_size_parameter_test.cpp", "plan_cache_key_info_test.cpp", "plan_cache_test.cpp", "plan_ranker_test.cpp", @@ -392,41 +412,52 @@ env.CppUnitTest( "query_planner_collation_test.cpp", "query_planner_columnar_test.cpp", "query_planner_geo_test.cpp", - "query_planner_pipeline_pushdown_test.cpp", "query_planner_hashed_index_test.cpp", - "query_planner_partialidx_test.cpp", "query_planner_index_test.cpp", "query_planner_operator_test.cpp", "query_planner_options_test.cpp", - "query_planner_tree_test.cpp", + "query_planner_partialidx_test.cpp", + "query_planner_pipeline_pushdown_test.cpp", "query_planner_text_test.cpp", + "query_planner_tree_test.cpp", "query_planner_wildcard_index_test.cpp", "query_request_test.cpp", "query_settings_test.cpp", + 'query_shape/agg_cmd_shape_test.cpp', + 'query_shape/cmd_with_let_shape_test.cpp', + "query_shape/find_cmd_shape_test.cpp", "query_solution_test.cpp", + "record_id_range_test.cpp", "sbe_and_hash_test.cpp", "sbe_and_sorted_test.cpp", + "sbe_shard_filter_test.cpp", "sbe_stage_builder_accumulator_test.cpp", "sbe_stage_builder_lookup_test.cpp", - "sbe_stage_builder_test_fixture.cpp", "sbe_stage_builder_test.cpp", - "sbe_shard_filter_test.cpp", + "sbe_stage_builder_test_fixture.cpp", "shard_filterer_factory_mock.cpp", + "sort_pattern_test.cpp", + "util/deferred_test.cpp", + "util/memory_util_test.cpp", "view_response_formatter_test.cpp", + 'map_reduce_output_format_test.cpp', ], LIBDEPS=[ "$BUILD_DIR/mongo/db/auth/authmocks", "$BUILD_DIR/mongo/db/concurrency/lock_manager", + "$BUILD_DIR/mongo/db/exec/document_value/document_value_test_util", "$BUILD_DIR/mongo/db/exec/sbe/sbe_plan_stage_test", "$BUILD_DIR/mongo/db/multitenancy", "$BUILD_DIR/mongo/db/pipeline/aggregation_request_helper", "$BUILD_DIR/mongo/db/pipeline/document_source_mock", "$BUILD_DIR/mongo/db/query_exec", + '$BUILD_DIR/mongo/db/record_id_helpers', "$BUILD_DIR/mongo/db/repl/replmocks", "$BUILD_DIR/mongo/db/repl/storage_interface_impl", "$BUILD_DIR/mongo/db/service_context_d_test_fixture", "$BUILD_DIR/mongo/db/service_context_test_fixture", "$BUILD_DIR/mongo/dbtests/mocklib", + "$BUILD_DIR/mongo/idl/idl_parser", "$BUILD_DIR/mongo/rpc/rpc", "$BUILD_DIR/mongo/util/clock_source_mock", "collation/collator_factory_mock", @@ -438,6 +469,7 @@ env.CppUnitTest( "query_planner", "query_planner_test_fixture", "query_request", + "query_shape/query_shape", "query_test_service_context", ], ) diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp index d203a6197bf..32ebde51cfb 100644 --- a/src/mongo/db/query/canonical_query.cpp +++ b/src/mongo/db/query/canonical_query.cpp @@ -36,7 +36,6 @@ #include "mongo/crypto/encryption_fields_gen.h" #include "mongo/db/catalog/collection.h" #include "mongo/db/commands/test_commands_enabled.h" -#include "mongo/db/cst/cst_parser.h" #include "mongo/db/jsobj.h" #include "mongo/db/matcher/expression_array.h" #include "mongo/db/namespace_string.h" @@ -45,111 +44,67 @@ #include "mongo/db/query/collation/collator_factory_interface.h" #include "mongo/db/query/fle/server_rewrite.h" #include "mongo/db/query/indexability.h" +#include "mongo/db/query/parsed_find_command.h" #include "mongo/db/query/projection_parser.h" #include "mongo/db/query/query_planner_common.h" #include "mongo/logv2/log.h" namespace mongo { -namespace { - -bool parsingCanProduceNoopMatchNodes(const ExtensionsCallback& extensionsCallback, - MatchExpressionParser::AllowedFeatureSet allowedFeatures) { - return extensionsCallback.hasNoopExtensions() && - (allowedFeatures & MatchExpressionParser::AllowedFeatures::kText || - allowedFeatures & MatchExpressionParser::AllowedFeatures::kJavascript); -} - -} // namespace // static StatusWith<std::unique_ptr<CanonicalQuery>> CanonicalQuery::canonicalize( OperationContext* opCtx, std::unique_ptr<FindCommandRequest> findCommand, bool explain, - const boost::intrusive_ptr<ExpressionContext>& expCtx, + const boost::intrusive_ptr<ExpressionContext>& givenExpCtx, const ExtensionsCallback& extensionsCallback, MatchExpressionParser::AllowedFeatureSet allowedFeatures, const ProjectionPolicies& projectionPolicies, std::vector<std::unique_ptr<InnerPipelineStageInterface>> pipeline) { - tassert(5746107, - "ntoreturn should not be set on the findCommand", - findCommand->getNtoreturn() == boost::none); - auto status = query_request_helper::validateFindCommandRequest(*findCommand); - if (!status.isOK()) { - return status; - } - - std::unique_ptr<CollatorInterface> collator; - if (!findCommand->getCollation().isEmpty()) { - auto statusWithCollator = CollatorFactoryInterface::get(opCtx->getServiceContext()) - ->makeFromBSON(findCommand->getCollation()); - if (!statusWithCollator.isOK()) { - return statusWithCollator.getStatus(); + if (givenExpCtx) { + // Caller provided an ExpressionContext, let's go ahead and use that. + auto swParsedFind = parsed_find_command::parse(givenExpCtx, + std::move(findCommand), + extensionsCallback, + allowedFeatures, + projectionPolicies); + if (!swParsedFind.isOK()) { + return swParsedFind.getStatus(); } - collator = std::move(statusWithCollator.getValue()); - } - - // Make MatchExpression. - boost::intrusive_ptr<ExpressionContext> newExpCtx; - if (!expCtx.get()) { - invariant(findCommand->getNamespaceOrUUID().nss()); - newExpCtx = make_intrusive<ExpressionContext>(opCtx, - std::move(collator), - *findCommand->getNamespaceOrUUID().nss(), - findCommand->getLegacyRuntimeConstants(), - findCommand->getLet()); + return canonicalize(std::move(givenExpCtx), + std::move(swParsedFind.getValue()), + explain, + std::move(pipeline)); } else { - newExpCtx = expCtx; - // A collator can enter through both the FindCommandRequest and ExpressionContext arguments. - // This invariant ensures that both collators are the same because downstream we - // pull the collator from only one of the ExpressionContext carrier. - if (collator.get() && expCtx->getCollator()) { - invariant(CollatorInterface::collatorsMatch(collator.get(), expCtx->getCollator())); + // No ExpressionContext provided, let's call the override that makes one for us. + auto swResults = parsed_find_command::parse( + opCtx, std::move(findCommand), extensionsCallback, allowedFeatures, projectionPolicies); + if (!swResults.isOK()) { + return swResults.getStatus(); } + auto&& [expCtx, parsedFind] = std::move(swResults.getValue()); + return canonicalize(std::move(expCtx), std::move(parsedFind), explain, std::move(pipeline)); } +} + +// static +StatusWith<std::unique_ptr<CanonicalQuery>> CanonicalQuery::canonicalize( + boost::intrusive_ptr<ExpressionContext> expCtx, + std::unique_ptr<ParsedFindCommand> parsedFind, + bool explain, + std::vector<std::unique_ptr<InnerPipelineStageInterface>> pipeline) { // Make the CQ we'll hopefully return. - std::unique_ptr<CanonicalQuery> cq(new CanonicalQuery()); + auto cq = std::make_unique<CanonicalQuery>(); cq->setExplain(explain); - - StatusWithMatchExpression statusWithMatcher = [&]() -> StatusWithMatchExpression { - if (getTestCommandsEnabled() && internalQueryEnableCSTParser.load()) { - try { - return cst::parseToMatchExpression( - findCommand->getFilter(), newExpCtx, extensionsCallback); - } catch (const DBException& ex) { - return ex.toStatus(); - } - } else { - return MatchExpressionParser::parse( - findCommand->getFilter(), newExpCtx, extensionsCallback, allowedFeatures); - } - }(); - if (!statusWithMatcher.isOK()) { - return statusWithMatcher.getStatus(); - } - - // Stop counting expressions after they have been parsed to exclude expressions created - // during optimization and other processing steps. - newExpCtx->stopExpressionCounters(); - - std::unique_ptr<MatchExpression> me = std::move(statusWithMatcher.getValue()); - - Status initStatus = - cq->init(opCtx, - std::move(newExpCtx), - std::move(findCommand), - parsingCanProduceNoopMatchNodes(extensionsCallback, allowedFeatures), - std::move(me), - projectionPolicies, - std::move(pipeline), - true /*optimizeMatchExpression*/ - ); - - if (!initStatus.isOK()) { + if (auto initStatus = cq->init(std::move(expCtx), + std::move(parsedFind), + std::move(pipeline), + true /*optimizeMatchExpression*/); + !initStatus.isOK()) { return initStatus; } - return std::move(cq); + return {std::move(cq)}; } // static @@ -163,67 +118,84 @@ StatusWith<std::unique_ptr<CanonicalQuery>> CanonicalQuery::makeForSubplanner( baseQuery.root()->numChildren() > i); auto root = baseQuery.root()->getChild(i); auto findCommand = std::make_unique<FindCommandRequest>(baseQuery.nss()); - BSONObjBuilder builder; - root->serialize(&builder, true); - findCommand->setFilter(builder.obj()); + findCommand->setFilter(root->serialize()); findCommand->setProjection(baseQuery.getFindCommandRequest().getProjection().getOwned()); findCommand->setSort(baseQuery.getFindCommandRequest().getSort().getOwned()); findCommand->setCollation(baseQuery.getFindCommandRequest().getCollation().getOwned()); - auto status = query_request_helper::validateFindCommandRequest(*findCommand); - if (!status.isOK()) { - return status; - } // Make the CQ we'll hopefully return. - std::unique_ptr<CanonicalQuery> cq(new CanonicalQuery()); + auto cq = std::make_unique<CanonicalQuery>(); cq->setExplain(baseQuery.getExplain()); - + auto swParsedFind = ParsedFindCommand::withExistingFilter( + baseQuery.getExpCtx(), + baseQuery.getCollator() ? baseQuery.getCollator()->clone() : nullptr, + root->shallowClone(), + std::move(findCommand)); + if (!swParsedFind.isOK()) { + return swParsedFind.getStatus(); + } // Note: we do not optimize the MatchExpression representing the branch of the top-level $or // that we are currently examining. This is because repeated invocations of // MatchExpression::optimize() may change the order of predicates in the MatchExpression, due to // new rewrites being unlocked by previous ones. We need to preserve the order of predicates to // allow index tagging to work properly. See SERVER-84013 for more details. - Status initStatus = cq->init(opCtx, - baseQuery.getExpCtx(), - std::move(findCommand), - baseQuery._canHaveNoopMatchNodes, - root->shallowClone(), - ProjectionPolicies::findProjectionPolicies(), + Status initStatus = cq->init(baseQuery.getExpCtx(), + std::move(swParsedFind.getValue()), {} /* an empty pipeline */, false /*optimizeMatchExpression*/); - if (!initStatus.isOK()) { - return initStatus; - } - return std::move(cq); + invariant(initStatus.isOK()); + return {std::move(cq)}; } -Status CanonicalQuery::init(OperationContext* opCtx, - boost::intrusive_ptr<ExpressionContext> expCtx, - std::unique_ptr<FindCommandRequest> findCommand, - bool canHaveNoopMatchNodes, - std::unique_ptr<MatchExpression> root, - const ProjectionPolicies& projectionPolicies, +Status CanonicalQuery::init(boost::intrusive_ptr<ExpressionContext> expCtx, + std::unique_ptr<ParsedFindCommand> parsedFind, std::vector<std::unique_ptr<InnerPipelineStageInterface>> pipeline, bool optimizeMatchExpression) { _expCtx = expCtx; - _findCommand = std::move(findCommand); + _findCommand = std::move(parsedFind->findCommandRequest); + _canHaveNoopMatchNodes = parsedFind->canHaveNoopMatchNodes; - _canHaveNoopMatchNodes = canHaveNoopMatchNodes; _forceClassicEngine = internalQueryForceClassicEngine.load(); - auto validStatus = isValid(root.get(), *_findCommand); - if (!validStatus.isOK()) { - return validStatus.getStatus(); - } - auto unavailableMetadata = validStatus.getValue(); - if (optimizeMatchExpression) { - _root = MatchExpression::normalize(std::move(root)); + _root = MatchExpression::normalize(std::move(parsedFind->filter)); } else { - _root = std::move(root); + _root = std::move(parsedFind->filter); + } + + if (parsedFind->proj) { + // The projection will be optimized only if the query is not compatible with SBE or there's + // no user-specified "let" variable. This is to prevent the user-defined variable being + // optimized out. We will optimize the projection later after we are certain that the query + // is ineligible for SBE. + bool shouldOptimizeProj = !expCtx->sbeCompatible || !_findCommand->getLet(); + if (parsedFind->proj->requiresMatchDetails()) { + // Sadly, in some cases the match details cannot be generated from the unoptimized + // MatchExpression. For example, a rooted-$or of equalities won't work to produce the + // details, but if you optimize that query to an $in, it will work. If we were starting + // from scratch, we may disallow this. But it has already been released as working so we + // will keep it so, and here have to re-parse the projection using the new, normalized + // MatchExpression, before we save this projection for later execution. + _proj.emplace(projection_ast::parseAndAnalyze(expCtx, + _findCommand->getProjection(), + _root.get(), + _findCommand->getFilter(), + *parsedFind->savedProjectionPolicies, + shouldOptimizeProj)); + } else { + _proj.emplace(std::move(*parsedFind->proj)); + if (shouldOptimizeProj) { + _proj->optimize(); + } + } } + if (parsedFind->sort) { + _sortPattern = std::move(parsedFind->sort); + } + _pipeline = std::move(pipeline); + // Perform auto-parameterization only if the query is SBE-compatible and caching is enabled. if (feature_flags::gFeatureFlagSbePlanCache.isEnabledAndIgnoreFCV()) { const bool hasNoTextNodes = !QueryPlannerCommon::hasNode(_root.get(), MatchExpression::TEXT); @@ -240,89 +212,45 @@ Status CanonicalQuery::init(OperationContext* opCtx, } } // The tree must always be valid after normalization. - dassert(isValid(_root.get(), *_findCommand).isOK()); + dassert(parsed_find_command::isValid(_root.get(), *_findCommand).isOK()); if (auto status = isValidNormalized(_root.get()); !status.isOK()) { return status; } - // Validate the projection if there is one. - if (!_findCommand->getProjection().isEmpty()) { - try { - _proj.emplace(projection_ast::parseAndAnalyze(expCtx, - _findCommand->getProjection(), - _root.get(), - _findCommand->getFilter(), - projectionPolicies, - true /* Should optimize? */)); + if (_proj) { + _metadataDeps = _proj->metadataDeps(); - // Fail if any of the projection's dependencies are unavailable. - DepsTracker{unavailableMetadata}.requestMetadata(_proj->metadataDeps()); - } catch (const DBException& e) { - return e.toStatus(); + if (_proj->metadataDeps()[DocumentMetadataFields::kSortKey] && + _findCommand->getSort().isEmpty()) { + return {ErrorCodes::BadValue, "cannot use sortKey $meta projection without a sort"}; } - - _metadataDeps = _proj->metadataDeps(); } - _pipeline = std::move(pipeline); + if (_sortPattern) { + // Be sure to track and add any metadata dependencies from the sort (e.g. text score). + _metadataDeps |= _sortPattern->metadataDeps(parsedFind->unavailableMetadata); - if (_proj && _proj->metadataDeps()[DocumentMetadataFields::kSortKey] && - _findCommand->getSort().isEmpty()) { - return Status(ErrorCodes::BadValue, "cannot use sortKey $meta projection without a sort"); - } - - // If there is a sort, parse it and add any metadata dependencies it induces. - try { - initSortPattern(unavailableMetadata); - } catch (const DBException& ex) { - return ex.toStatus(); + // If the results of this query might have to be merged on a remote node, then that node + // might need the sort key metadata. Request that the plan generates this metadata. + if (_expCtx->needsMerge) { + _metadataDeps.set(DocumentMetadataFields::kSortKey); + } } // If the 'returnKey' option is set, then the plan should produce index key metadata. if (_findCommand->getReturnKey()) { _metadataDeps.set(DocumentMetadataFields::kIndexKey); } - return Status::OK(); } -void CanonicalQuery::initSortPattern(QueryMetadataBitSet unavailableMetadata) { - if (_findCommand->getSort().isEmpty()) { - return; - } - - // A $natural sort is really a hint, and should be handled as such. Furthermore, the downstream - // sort handling code may not expect a $natural sort. - // - // We have already validated that if there is a $natural sort and a hint, that the hint - // also specifies $natural with the same direction. Therefore, it is safe to clear the $natural - // sort and rewrite it as a $natural hint. - if (_findCommand->getSort()[query_request_helper::kNaturalSortField]) { - _findCommand->setHint(_findCommand->getSort().getOwned()); - _findCommand->setSort(BSONObj{}); - } - - if (getTestCommandsEnabled() && internalQueryEnableCSTParser.load()) { - _sortPattern = cst::parseToSortPattern(_findCommand->getSort(), _expCtx); - } else { - _sortPattern = SortPattern{_findCommand->getSort(), _expCtx}; - } - _metadataDeps |= _sortPattern->metadataDeps(unavailableMetadata); - - // If the results of this query might have to be merged on a remote node, then that node might - // need the sort key metadata. Request that the plan generates this metadata. - if (_expCtx->needsMerge) { - _metadataDeps.set(DocumentMetadataFields::kSortKey); - } -} - void CanonicalQuery::setCollator(std::unique_ptr<CollatorInterface> collator) { auto collatorRaw = collator.get(); // We must give the ExpressionContext the same collator. _expCtx->setCollator(std::move(collator)); - // The collator associated with the match expression tree is now invalid, since we have reset - // the collator owned by the ExpressionContext. + // The collator associated with the match expression tree is now invalid, since we have + // reset the collator owned by the ExpressionContext. _root->setCollator(collatorRaw); } @@ -356,138 +284,9 @@ bool CanonicalQuery::isSimpleIdQuery(const BSONObj& query) { return hasID; } -size_t CanonicalQuery::countNodes(const MatchExpression* root, MatchExpression::MatchType type) { - size_t sum = 0; - if (type == root->matchType()) { - sum = 1; - } - for (size_t i = 0; i < root->numChildren(); ++i) { - sum += countNodes(root->getChild(i), type); - } - return sum; -} - -/** - * Does 'root' have a subtree of type 'subtreeType' with a node of type 'childType' inside? - */ -bool hasNodeInSubtree(const MatchExpression* root, - MatchExpression::MatchType childType, - MatchExpression::MatchType subtreeType) { - if (subtreeType == root->matchType()) { - return QueryPlannerCommon::hasNode(root, childType); - } - for (size_t i = 0; i < root->numChildren(); ++i) { - if (hasNodeInSubtree(root->getChild(i), childType, subtreeType)) { - return true; - } - } - return false; -} - -StatusWith<QueryMetadataBitSet> CanonicalQuery::isValid(const MatchExpression* root, - const FindCommandRequest& findCommand) { - QueryMetadataBitSet unavailableMetadata{}; - - // There can only be one TEXT. If there is a TEXT, it cannot appear inside a NOR. - // - // Note that the query grammar (as enforced by the MatchExpression parser) forbids TEXT - // inside of value-expression clauses like NOT, so we don't check those here. - size_t numText = countNodes(root, MatchExpression::TEXT); - if (numText > 1) { - return Status(ErrorCodes::BadValue, "Too many text expressions"); - } else if (1 == numText) { - if (hasNodeInSubtree(root, MatchExpression::TEXT, MatchExpression::NOR)) { - return Status(ErrorCodes::BadValue, "text expression not allowed in nor"); - } - } else { - // Text metadata is not available. - unavailableMetadata.set(DocumentMetadataFields::kTextScore); - } - - // There can only be one NEAR. If there is a NEAR, it must be either the root or the root - // must be an AND and its child must be a NEAR. - size_t numGeoNear = countNodes(root, MatchExpression::GEO_NEAR); - if (numGeoNear > 1) { - return Status(ErrorCodes::BadValue, "Too many geoNear expressions"); - } else if (1 == numGeoNear) { - // Do nothing, we will perform extra checks in CanonicalQuery::isValidNormalized. - } else { - // Geo distance and geo point metadata are unavailable. - unavailableMetadata |= DepsTracker::kAllGeoNearData; - } - - const BSONObj& sortObj = findCommand.getSort(); - BSONElement sortNaturalElt = sortObj["$natural"]; - const BSONObj& hintObj = findCommand.getHint(); - BSONElement hintNaturalElt = hintObj["$natural"]; - - if (sortNaturalElt && sortObj.nFields() != 1) { - return Status(ErrorCodes::BadValue, - str::stream() << "Cannot include '$natural' in compound sort: " << sortObj); - } - - if (hintNaturalElt && hintObj.nFields() != 1) { - return Status(ErrorCodes::BadValue, - str::stream() << "Cannot include '$natural' in compound hint: " << hintObj); - } - - // NEAR cannot have a $natural sort or $natural hint. - if (numGeoNear > 0) { - if (sortNaturalElt) { - return Status(ErrorCodes::BadValue, - "geoNear expression not allowed with $natural sort order"); - } - - if (hintNaturalElt) { - return Status(ErrorCodes::BadValue, - "geoNear expression not allowed with $natural hint"); - } - } - - // TEXT and NEAR cannot both be in the query. - if (numText > 0 && numGeoNear > 0) { - return Status(ErrorCodes::BadValue, "text and geoNear not allowed in same query"); - } - - // TEXT and {$natural: ...} sort order cannot both be in the query. - if (numText > 0 && sortNaturalElt) { - return Status(ErrorCodes::BadValue, "text expression not allowed with $natural sort order"); - } - - // TEXT and hint cannot both be in the query. - if (numText > 0 && !hintObj.isEmpty()) { - return Status(ErrorCodes::BadValue, "text and hint not allowed in same query"); - } - - // TEXT and tailable are incompatible. - if (numText > 0 && findCommand.getTailable()) { - return Status(ErrorCodes::BadValue, "text and tailable cursor not allowed in same query"); - } - - // NEAR and tailable are incompatible. - if (numGeoNear > 0 && findCommand.getTailable()) { - return Status(ErrorCodes::BadValue, - "Tailable cursors and geo $near cannot be used together"); - } - - // $natural sort order must agree with hint. - if (sortNaturalElt) { - if (!hintObj.isEmpty() && !hintNaturalElt) { - return Status(ErrorCodes::BadValue, "index hint not allowed with $natural sort order"); - } - if (hintNaturalElt) { - if (hintNaturalElt.numberInt() != sortNaturalElt.numberInt()) { - return Status(ErrorCodes::BadValue, - "$natural hint must be in the same direction as $natural sort order"); - } - } - } - - return unavailableMetadata; -} - Status CanonicalQuery::isValidNormalized(const MatchExpression* root) { - if (auto numGeoNear = countNodes(root, MatchExpression::GEO_NEAR); numGeoNear > 0) { + if (auto numGeoNear = QueryPlannerCommon::countNodes(root, MatchExpression::GEO_NEAR); + numGeoNear > 0) { tassert(5705300, "Only one geo $near expression is expected", numGeoNear == 1); auto topLevel = false; diff --git a/src/mongo/db/query/canonical_query.h b/src/mongo/db/query/canonical_query.h index 79b49c08463..570b71b6fae 100644 --- a/src/mongo/db/query/canonical_query.h +++ b/src/mongo/db/query/canonical_query.h @@ -37,6 +37,7 @@ #include "mongo/db/matcher/extensions_callback_noop.h" #include "mongo/db/pipeline/inner_pipeline_stage_interface.h" #include "mongo/db/query/collation/collator_interface.h" +#include "mongo/db/query/parsed_find_command.h" #include "mongo/db/query/projection.h" #include "mongo/db/query/projection_policies.h" #include "mongo/db/query/query_request_helper.h" @@ -76,6 +77,15 @@ public: std::vector<std::unique_ptr<InnerPipelineStageInterface>> pipeline = {}); /** + * Creates a CanonicalQuery from a ParsedFindCommand. Uses 'expCtx->opCtx', which must be valid. + */ + static StatusWith<std::unique_ptr<CanonicalQuery>> canonicalize( + boost::intrusive_ptr<ExpressionContext> expCtx, + std::unique_ptr<ParsedFindCommand> parsedFind, + bool explain = false, + std::vector<std::unique_ptr<InnerPipelineStageInterface>> pipeline = {}); + + /** * For testing or for internal clients to use. */ @@ -93,32 +103,22 @@ public: static bool isSimpleIdQuery(const BSONObj& query); /** - * Validates the match expression 'root' as well as the query specified by 'request', checking - * for illegal combinations of operators. Returns a non-OK status if any such illegal - * combination is found. - * - * This method can be called both on normalized and non-normalized 'root'. However, some checks - * can only be performed once the match expressions is normalized. To perform these checks one - * can call 'isValidNormalized()'. - * - * On success, returns a bitset indicating which types of metadata are *unavailable*. For - * example, if 'root' does not contain a $text predicate, then the returned metadata bitset will - * indicate that text score metadata is unavailable. This means that if subsequent - * $meta:"textScore" expressions are found during analysis of the query, we should raise in an - * error. + * Perform validation checks on the normalized 'root' which could not be checked before + * normalization - those should happen in parsed_find_command::isValid(). */ - static StatusWith<QueryMetadataBitSet> isValid(const MatchExpression* root, - const FindCommandRequest& findCommand); + static Status isValidNormalized(const MatchExpression* root); /** - * Perform additional validation checks on the normalized 'root'. + * For internal use only - but public for accessibility for make_unique(). You must go through + * canonicalize to create a CanonicalQuery. */ - static Status isValidNormalized(const MatchExpression* root); + CanonicalQuery() {} const NamespaceString nss() const { invariant(_findCommand->getNamespaceOrUUID().nss()); return *_findCommand->getNamespaceOrUUID().nss(); } + const std::string ns() const { return nss().ns(); } @@ -197,11 +197,6 @@ public: std::string toStringShort() const; /** - * Returns a count of 'type' nodes in expression tree. - */ - static size_t countNodes(const MatchExpression* root, MatchExpression::MatchType type); - - /** * Returns true if this canonical query may have converted extensions such as $where and $text * into no-ops during parsing. This will be the case if it allowed $where and $text in parsing, * but parsed using an ExtensionsCallbackNoop. This does not guarantee that a $where or $text @@ -262,25 +257,18 @@ public: return _pipeline; } -private: - // You must go through canonicalize to create a CanonicalQuery. - CanonicalQuery() {} + void optimizeProjection() { + if (_proj) { + _proj->optimize(); + } + } - Status init(OperationContext* opCtx, - boost::intrusive_ptr<ExpressionContext> expCtx, - std::unique_ptr<FindCommandRequest> findCommand, - bool canHaveNoopMatchNodes, - std::unique_ptr<MatchExpression> root, - const ProjectionPolicies& projectionPolicies, +private: + Status init(boost::intrusive_ptr<ExpressionContext> expCtx, + std::unique_ptr<ParsedFindCommand> parsedFind, std::vector<std::unique_ptr<InnerPipelineStageInterface>> pipeline, bool optimizeMatchExpression); - // Initializes '_sortPattern', adding any metadata dependencies implied by the sort. - // - // Throws a UserException if the sort is illegal, or if any metadata type in - // 'unavailableMetadata' is required. - void initSortPattern(QueryMetadataBitSet unavailableMetadata); - boost::intrusive_ptr<ExpressionContext> _expCtx; std::unique_ptr<FindCommandRequest> _findCommand; diff --git a/src/mongo/db/query/canonical_query_encoder.cpp b/src/mongo/db/query/canonical_query_encoder.cpp index c8a8cd29d48..61ee86dc2c8 100644 --- a/src/mongo/db/query/canonical_query_encoder.cpp +++ b/src/mongo/db/query/canonical_query_encoder.cpp @@ -989,7 +989,7 @@ public: void preVisit(const MatchExpression* expr) { // Encode the type of the node as well as the path (if there is a non-empty path). - _builder->appendStr(encodeMatchType(expr->matchType())); + _builder->appendCStr(encodeMatchType(expr->matchType())); encodeUserString(expr->path(), _builder); // The node encodes itself, and then its children. @@ -1053,7 +1053,7 @@ std::string encodeSBE(const CanonicalQuery& cq) { encodeKeyForAutoParameterizedMatchSBE(cq.root(), &bufBuilder); bufBuilder.appendBuf(proj.objdata(), proj.objsize()); - bufBuilder.appendStr(strBuilderEncoded, false /* includeEndingNull */); + bufBuilder.appendStrBytes(strBuilderEncoded); encodeFindCommandRequest(cq.getFindCommandRequest(), &bufBuilder); diff --git a/src/mongo/db/query/canonical_query_test.cpp b/src/mongo/db/query/canonical_query_test.cpp index 497d2671270..5cba0bcc3f4 100644 --- a/src/mongo/db/query/canonical_query_test.cpp +++ b/src/mongo/db/query/canonical_query_test.cpp @@ -281,9 +281,50 @@ TEST(CanonicalQueryTest, CanonicalizeFromBaseQuery) { MatchExpression* firstClauseExpr = baseCq->root()->getChild(0); auto childCq = assertGet(CanonicalQuery::makeForSubplanner(opCtx.get(), *baseCq, 0)); - BSONObjBuilder expectedFilter; - firstClauseExpr->serialize(&expectedFilter); - ASSERT_BSONOBJ_EQ(childCq->getFindCommandRequest().getFilter(), expectedFilter.obj()); + ASSERT_BSONOBJ_EQ(childCq->getFindCommandRequest().getFilter(), firstClauseExpr->serialize()); + + ASSERT_BSONOBJ_EQ(childCq->getFindCommandRequest().getProjection(), + baseCq->getFindCommandRequest().getProjection()); + ASSERT_BSONOBJ_EQ(childCq->getFindCommandRequest().getSort(), + baseCq->getFindCommandRequest().getSort()); + ASSERT_TRUE(childCq->getExplain()); +} + +TEST(CanonicalQueryTest, CanonicalizeFromBaseQueryWithSpecialFeature) { + // Like the above test, but use $text which is a 'special feature' not always allowed. This is + // meant to reproduce SERVER-XYZ. + QueryTestServiceContext serviceContext; + auto opCtx = serviceContext.makeOperationContext(); + + const bool isExplain = true; + const std::string cmdStr = R"({ + find:'bogusns', + filter: { + $or:[ + {a: 'foo'}, + {$text: {$search: 'bar'}} + ] + }, + projection: {a:1}, + sort: {b:1}, + $db: 'test' + })"; + auto findCommand = query_request_helper::makeFromFindCommandForTests(fromjson(cmdStr)); + auto baseCq = + assertGet(CanonicalQuery::canonicalize(opCtx.get(), + std::move(findCommand), + isExplain, + nullptr, + ExtensionsCallbackNoop(), + MatchExpressionParser::kAllowAllSpecialFeatures)); + + // Note: be sure to use the second child to get $text, since we 'normalize' and sort the + // MatchExpression tree as part of canonicalization. This will put the text search clause + // second. + MatchExpression* secondClauseExpr = baseCq->root()->getChild(1); + auto childCq = assertGet(CanonicalQuery::makeForSubplanner(opCtx.get(), *baseCq, 1)); + + ASSERT_BSONOBJ_EQ(childCq->getFindCommandRequest().getFilter(), secondClauseExpr->serialize()); ASSERT_BSONOBJ_EQ(childCq->getFindCommandRequest().getProjection(), baseCq->getFindCommandRequest().getProjection()); diff --git a/src/mongo/db/query/classic_plan_cache.h b/src/mongo/db/query/classic_plan_cache.h index 8d2e6889bfb..e510b48dd6d 100644 --- a/src/mongo/db/query/classic_plan_cache.h +++ b/src/mongo/db/query/classic_plan_cache.h @@ -235,7 +235,12 @@ using PlanCacheEntry = PlanCacheEntryBase<SolutionCacheData, plan_cache_debug_in using CachedSolution = CachedPlanHolder<SolutionCacheData, plan_cache_debug_info::DebugInfo>; struct BudgetEstimator { - size_t operator()(const std::shared_ptr<const PlanCacheEntry>&) { + /** + * This estimator function is called when an entry is added or removed to LRU cache in order to + * make sure the total plan cache size does not exceed the maximum size. + */ + + size_t operator()(const PlanCacheKey&, const std::shared_ptr<const PlanCacheEntry>&) { return 1; } }; diff --git a/src/mongo/db/query/explain.cpp b/src/mongo/db/query/explain.cpp index 414badb8332..796875bca68 100644 --- a/src/mongo/db/query/explain.cpp +++ b/src/mongo/db/query/explain.cpp @@ -124,7 +124,7 @@ void generatePlannerInfo(PlanExecutor* exec, auto query = exec->getCanonicalQuery(); if (nullptr != query) { BSONObjBuilder parsedQueryBob(plannerBob.subobjStart("parsedQuery")); - query->root()->serialize(&parsedQueryBob); + query->root()->serialize(&parsedQueryBob, {}); parsedQueryBob.doneFast(); if (query->getCollator()) { diff --git a/src/mongo/db/query/find.cpp b/src/mongo/db/query/find.cpp index d2f36155e13..cd8e21bf6c7 100644 --- a/src/mongo/db/query/find.cpp +++ b/src/mongo/db/query/find.cpp @@ -107,19 +107,28 @@ void endQueryOp(OperationContext* opCtx, const CollectionPtr& collection, const PlanExecutor& exec, long long numResults, - CursorId cursorId) { + boost::optional<ClientCursorPin&> cursor, + const BSONObj& cmdObj) { auto curOp = CurOp::get(opCtx); - // Fill out basic CurOp query exec properties. - curOp->debug().nreturned = numResults; - curOp->debug().cursorid = (0 == cursorId ? -1 : cursorId); - curOp->debug().cursorExhausted = (0 == cursorId); + // Fill out basic CurOp query exec properties. More metrics (nreturned and executionTime) + // are collected within collectQueryStatsMongod. + curOp->debug().cursorid = (cursor.has_value() ? cursor->getCursor()->cursorid() : -1); + curOp->debug().cursorExhausted = !cursor.has_value(); + curOp->debug().additiveMetrics.nBatches = 1; // Fill out CurOp based on explain summary statistics. PlanSummaryStats summaryStats; auto&& explainer = exec.getPlanExplainer(); explainer.getSummaryStats(&summaryStats); curOp->debug().setPlanSummaryMetrics(summaryStats); + curOp->setEndOfOpMetrics(numResults); + + if (cursor) { + collectQueryStatsMongod(opCtx, *cursor); + } else { + collectQueryStatsMongod(opCtx, std::move(curOp->debug().queryStatsInfo.key)); + } if (collection) { CollectionQueryInfo::get(collection).notifyOfQuery(opCtx, collection, summaryStats); diff --git a/src/mongo/db/query/find.h b/src/mongo/db/query/find.h index 2b45efcd602..235ebf91ed8 100644 --- a/src/mongo/db/query/find.h +++ b/src/mongo/db/query/find.h @@ -73,6 +73,7 @@ void endQueryOp(OperationContext* opCtx, const CollectionPtr& collection, const PlanExecutor& exec, long long numResults, - CursorId cursorId); + boost::optional<ClientCursorPin&> cursor, + const BSONObj& cmdObj); } // namespace mongo diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp index 5e1d413745e..602be15dbc2 100644 --- a/src/mongo/db/query/get_executor.cpp +++ b/src/mongo/db/query/get_executor.cpp @@ -29,10 +29,13 @@ #define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery +#include "mongo/db/curop.h" #include "mongo/platform/basic.h" #include "mongo/db/query/get_executor.h" +#include "mongo/util/duration.h" +#include "mongo/util/tick_source.h" #include <boost/optional.hpp> #include <limits> #include <memory> @@ -107,6 +110,7 @@ #include "mongo/logv2/log.h" #include "mongo/scripting/engine.h" #include "mongo/util/str.h" +#include "mongo/util/timer.h" namespace mongo { MONGO_FAIL_POINT_DEFINE(includeFakeColumnarIndex); @@ -603,6 +607,8 @@ public: StatusWith<std::unique_ptr<ResultType>> prepare() { const auto& mainColl = getMainCollection(); + + ON_BLOCK_EXIT([&] { CurOp::get(_opCtx)->stopQueryPlanningTimer(); }); if (!mainColl) { LOGV2_DEBUG(20921, 2, @@ -702,10 +708,8 @@ public: "Only one plan is available", "query"_attr = redact(_cq->toStringShort()), "planSummary"_attr = result->getPlanSummary()); - return std::move(result); } - return buildMultiPlan(std::move(solutions)); } @@ -1317,16 +1321,16 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getSlotBasedExe auto&& [roots, solutions] = planningResult->extractResultData(); // In some circumstances (e.g. when have multiple candidate plans or using a cached one), we // might need to execute the plan(s) to pick the best one or to confirm the choice. - if (auto planner = makeRuntimePlannerIfNeeded(opCtx, - collections, - cq.get(), - solutions.size(), - planningResult->decisionWorks(), - planningResult->needsSubplanning(), - yieldPolicy.get(), - plannerParams.options)) { + if (auto runTimePlanner = makeRuntimePlannerIfNeeded(opCtx, + collections, + cq.get(), + solutions.size(), + planningResult->decisionWorks(), + planningResult->needsSubplanning(), + yieldPolicy.get(), + plannerParams.options)) { // Do the runtime planning and pick the best candidate plan. - auto candidates = planner->plan(std::move(solutions), std::move(roots)); + auto candidates = runTimePlanner->plan(std::move(solutions), std::move(roots)); return plan_executor_factory::make(opCtx, std::move(cq), @@ -1397,6 +1401,11 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor( } } + // There's a special case of the projection optimization being skipped when a query has any + // user-defined "let" variable and the query may be run with SBE. Here we make sure the + // projection is optimized for the classic engine. + canonicalQuery->optimizeProjection(); + return getClassicExecutor( opCtx, mainColl, std::move(canonicalQuery), yieldPolicy, plannerParams); } @@ -1409,6 +1418,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutor( PlanYieldPolicy::YieldPolicy yieldPolicy, size_t plannerOptions) { MultipleCollectionAccessor multi{collection}; + return getExecutor(opCtx, multi, std::move(canonicalQuery), @@ -1450,6 +1460,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorFind std::function<void(CanonicalQuery*)> extractAndAttachPipelineStages, bool permitYield, size_t plannerOptions) { + MultipleCollectionAccessor multi{*coll}; return getExecutorFind(opCtx, multi, @@ -1665,6 +1676,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorDele ClassicPrepareExecutionHelper helper{ opCtx, collection, ws.get(), cq.get(), nullptr, defaultPlannerOptions}; auto executionResult = helper.prepare(); + if (!executionResult.isOK()) { return executionResult.getStatus(); } @@ -1852,6 +1864,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorUpda ClassicPrepareExecutionHelper helper{ opCtx, collection, ws.get(), cq.get(), nullptr, defaultPlannerOptions}; auto executionResult = helper.prepare(); + if (!executionResult.isOK()) { return executionResult.getStatus(); } @@ -2125,8 +2138,8 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCoun OperationContext* opCtx = expCtx->opCtx; std::unique_ptr<WorkingSet> ws = std::make_unique<WorkingSet>(); - auto findCommand = std::make_unique<FindCommandRequest>(nss); + findCommand->setFilter(request.getQuery()); auto collation = request.getCollation().value_or(BSONObj()); findCommand->setCollation(collation); @@ -2200,6 +2213,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCoun if (!executionResult.isOK()) { return executionResult.getStatus(); } + auto [root, querySolution] = executionResult.getValue()->extractResultData(); invariant(root); @@ -2208,6 +2222,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> getExecutorCoun expCtx.get(), collection, limit, skip, ws.get(), root.release()); // We must have a tree of stages in order to have a valid plan executor, but the query // solution may be NULL. Takes ownership of all args other than 'collection' and 'opCtx' + return plan_executor_factory::make(std::move(cq), std::move(ws), std::move(root), diff --git a/src/mongo/db/query/indexability.h b/src/mongo/db/query/indexability.h index 48b9e0d91b4..1736d92a89b 100644 --- a/src/mongo/db/query/indexability.h +++ b/src/mongo/db/query/indexability.h @@ -136,6 +136,20 @@ public: } /** + * Returns true if 'me' is ELEM_MATCH_OBJECT and has non-empty path component. + * + * Note: we skip empty path components since they are not allowed in index key patterns. + * Therefore, $elemMatch with an empty path component can never use an index. + * + * Example: {"": {$elemMatch: {a: "hi", b: "bye"}}. + * In this case the predicate cannot use any indexes since the $elemMatch is with an empty path + * component. + */ + static bool isBoundsGeneratingElemMatchObject(const MatchExpression* me) { + return arrayUsesIndexOnChildren(me) && !me->path().empty(); + } + + /** * Returns true if 'me' is a NOT, and the child of the NOT can use * an index on its own field. */ diff --git a/src/mongo/db/query/lru_key_value.h b/src/mongo/db/query/lru_key_value.h index 88186c70923..2786e1c40ab 100644 --- a/src/mongo/db/query/lru_key_value.h +++ b/src/mongo/db/query/lru_key_value.h @@ -28,7 +28,6 @@ */ #pragma once - #include <fmt/format.h> #include <list> #include <memory> @@ -40,30 +39,56 @@ namespace mongo { /** + * 'InsertionEvictionListener' class to use with 'LRUBudgetTracker' that will always noop. + */ +class NoopInsertionEvictionListener { +public: + // Called when a key-value pair is being inserted. Parameters are the key-value pair and its + // estimated size. + template <class K, class V> + void onInsert(const K&, const V&, size_t) {} + + // Called when a key-value pair is being evicted. Parameters are the key-value pair and its + // estimated size. + template <class K, class V> + void onEvict(const K&, const V&, size_t) {} + + // Called when the cache is being cleared. Parameter is the estimated size of the key-value + // pairs in the cache before it was cleared. + void onClear(size_t) {} +}; + +/** * This class tracks a size of entries in 'LRUKeyValue'. * The size can be understood as a number of the entries, an amount of memory they occupied, * or any other value defined by the template parameter 'Estimator'. * The 'Estimator' must be deterministic and always return the same value for the same entry. + * The 'InsertionEvictionListener' will be called on every insertion and eviction as well as when + * the cache is cleared. */ -template <typename V, typename Estimator> +template <class K, class V, typename Estimator, typename InsertionEvictionListener> class LRUBudgetTracker { public: LRUBudgetTracker(size_t maxBudget) : _max(maxBudget), _current(0) {} - void onAdd(const V& v) { - _current += _estimator(v); + void onAdd(const K& k, const V& v) { + size_t budget = _estimator(k, v); + _current += budget; + _listener.onInsert(k, v, budget); } - void onRemove(const V& v) { + void onRemove(const K& k, const V& v) { using namespace fmt::literals; - size_t budget = _estimator(v); + size_t budget = _estimator(k, v); tassert(5968300, "LRU budget underflow: current={}, budget={} "_format(_current, budget), _current >= budget); _current -= budget; + _listener.onEvict(k, v, budget); } void onClear() { + _listener.onClear(_current); _current = 0; } @@ -84,6 +109,7 @@ private: size_t _max; size_t _current; Estimator _estimator; + InsertionEvictionListener _listener; }; /** @@ -91,6 +117,9 @@ private: * policy. The size allowed in the kv-store is controlled by 'LRUBudgetTracker' * set in the constructor. * + * An 'InsertionEvictionListener' may optionally be specified to track the insertion and eviction of + * each key-value pair. + * * Caveat: * This kv-store is NOT thread safe! The client to this utility is responsible * for protecting concurrent access to the LRU store if used in a threaded @@ -102,7 +131,12 @@ private: * TODO: We could move this into the util/ directory and do any cleanup necessary to make it * fully general. */ -template <class K, class V, class BudgetEstimator, class KeyHasher = std::hash<K>> +template <class K, + class V, + class KeyValueBudgetEstimator, + class InsertionEvictionListener = NoopInsertionEvictionListener, + class KeyHasher = std::hash<K>, + class Eq = std::equal_to<K>> class LRUKeyValue { public: LRUKeyValue(size_t maxSize) : _budgetTracker{maxSize} {} @@ -111,13 +145,13 @@ public: clear(); } - typedef std::pair<K, V> KVListEntry; + typedef std::pair<const K*, V> KVListEntry; typedef std::list<KVListEntry> KVList; typedef typename KVList::iterator KVListIt; typedef typename KVList::const_iterator KVListConstIt; - typedef stdx::unordered_map<K, KVListIt, KeyHasher> KVMap; + typedef stdx::unordered_map<K, KVListIt, KeyHasher, Eq> KVMap; typedef typename KVMap::const_iterator KVMapConstIt; // These type declarations are required by the 'Partitioned' utility. @@ -136,14 +170,15 @@ public: KVMapConstIt i = _kvMap.find(key); if (i != _kvMap.end()) { KVListIt found = i->second; - _budgetTracker.onRemove(found->second); + _budgetTracker.onRemove(key, found->second); _kvMap.erase(i); _kvList.erase(found); } - _budgetTracker.onAdd(entry); - _kvList.push_front(std::make_pair(key, std::move(entry))); + _budgetTracker.onAdd(key, entry); + _kvList.push_front(std::make_pair(nullptr, std::move(entry))); _kvMap[key] = _kvList.begin(); + _kvList.begin()->first = &(_kvMap.find(key)->first); return evict(); } @@ -161,10 +196,11 @@ public: KVListIt found = i->second; // Promote the kv-store entry to the front of the list. It is now the most recently used. - _kvList.push_front(std::make_pair(key, std::move(found->second))); + _kvList.push_front(std::make_pair(nullptr, std::move(found->second))); _kvMap.erase(i); _kvList.erase(found); _kvMap[key] = _kvList.begin(); + _kvList.begin()->first = &(_kvMap.find(key)->first); return _kvList.begin(); } @@ -179,7 +215,7 @@ public: return false; } KVListIt found = i->second; - _budgetTracker.onRemove(found->second); + _budgetTracker.onRemove(key, found->second); _kvMap.erase(i); _kvList.erase(found); return true; @@ -193,9 +229,9 @@ public: size_t removeIf(KeyValuePredicate predicate) { size_t removed = 0; for (auto it = _kvList.begin(); it != _kvList.end();) { - if (predicate(it->first, *it->second)) { - _budgetTracker.onRemove(it->second); - _kvMap.erase(it->first); + if (predicate(*it->first, *it->second)) { + _budgetTracker.onRemove(*it->first, it->second); + _kvMap.erase(*it->first); it = _kvList.erase(it); ++removed; } else { @@ -209,9 +245,9 @@ public: * Deletes all entries in the kv-store. */ void clear() { - _budgetTracker.onClear(); _kvList.clear(); _kvMap.clear(); + _budgetTracker.onClear(); } /** @@ -258,8 +294,8 @@ private: while (_budgetTracker.isOverBudget()) { invariant(!_kvList.empty()); - _budgetTracker.onRemove(_kvList.back().second); - _kvMap.erase(_kvList.back().first); + _budgetTracker.onRemove(*_kvList.back().first, _kvList.back().second); + _kvMap.erase(*_kvList.back().first); _kvList.pop_back(); ++nEvicted; @@ -268,13 +304,14 @@ private: return nEvicted; } - LRUBudgetTracker<V, BudgetEstimator> _budgetTracker; + LRUBudgetTracker<K, V, KeyValueBudgetEstimator, InsertionEvictionListener> _budgetTracker; // (K, V) pairs are stored in this std::list. They are sorted in order of use, where the front // is the most recently used and the back is the least recently used. mutable KVList _kvList; // Maps from a key to the corresponding std::list entry. + // TODO: SERVER-73659 LRUKeyValue should track and include the size of _kvMap in overall budget. mutable KVMap _kvMap; }; diff --git a/src/mongo/db/query/lru_key_value_test.cpp b/src/mongo/db/query/lru_key_value_test.cpp index 3ebf47267fb..6dcdfc5a4ea 100644 --- a/src/mongo/db/query/lru_key_value_test.cpp +++ b/src/mongo/db/query/lru_key_value_test.cpp @@ -64,15 +64,15 @@ struct ValueType { struct TrivialBudgetEstimator { static constexpr size_t kSize = 1; - size_t operator()(const ValueType&) { + size_t operator()(const int&, const ValueType&) { return kSize; } - size_t operator()(const std::shared_ptr<int>&) { + size_t operator()(const int&, const std::unique_ptr<int>&) { return kSize; } - size_t operator()(const std::unique_ptr<int>&) { + size_t operator()(const int&, const std::shared_ptr<int>) { return kSize; } }; @@ -87,7 +87,7 @@ struct NonTrivialEntry { }; struct NonTrivialBudgetEstimator { - size_t operator()(const std::shared_ptr<NonTrivialEntry>& value) { + size_t operator()(const int& key, const std::shared_ptr<NonTrivialEntry> value) { return value->budgetSize; } }; @@ -95,8 +95,40 @@ struct NonTrivialBudgetEstimator { using NonTrivialTestSharedPtrValue = LRUKeyValue<size_t, std::shared_ptr<NonTrivialEntry>, NonTrivialBudgetEstimator>; -template <typename Key, typename Value, typename Estimator> -void assertInKVStore(LRUKeyValue<Key, Value, Estimator>& cache, Key key, Value value) { +class NonTrivialInsertionEvictionListener { +public: + NonTrivialInsertionEvictionListener() { + keyTotal = 0; + valueTotal = 0; + budgetTotal = 0; + } + + void onInsert(const int& k, const ValueType& v, size_t budget) { + keyTotal += k; + valueTotal += v.val; + budgetTotal += budget; + } + + void onEvict(const int& k, const ValueType& v, size_t budget) { + keyTotal -= k; + valueTotal -= v.val; + budgetTotal -= budget; + } + + void onClear(size_t budget) { + budgetTotal -= budget; + } + + static size_t keyTotal; + static size_t valueTotal; + static size_t budgetTotal; +}; +size_t NonTrivialInsertionEvictionListener::keyTotal; +size_t NonTrivialInsertionEvictionListener::valueTotal; +size_t NonTrivialInsertionEvictionListener::budgetTotal; + +template <typename Key, typename Value, typename Estimator, typename Listener> +void assertInKVStore(LRUKeyValue<Key, Value, Estimator, Listener>& cache, Key key, Value value) { ASSERT_TRUE(cache.hasKey(key)); auto s = cache.get(key); ASSERT(s.isOK()); @@ -105,8 +137,8 @@ void assertInKVStore(LRUKeyValue<Key, Value, Estimator>& cache, Key key, Value v ASSERT_EQUALS(*(kvItr->second), *value); } -template <typename Key, typename Value, typename Estimator> -void assertNotInKVStore(LRUKeyValue<Key, Value, Estimator>& cache, Key key) { +template <typename Key, typename Value, typename Estimator, typename Listener> +void assertNotInKVStore(LRUKeyValue<Key, Value, Estimator, Listener>& cache, Key key) { ASSERT_FALSE(cache.hasKey(key)); auto s = cache.get(key); ASSERT(!s.isOK()); @@ -295,10 +327,10 @@ TEST(LRUKeyValueTest, IterationTest) { cache.add(2, std::make_shared<int>(2)); auto i = cache.begin(); - ASSERT_EQUALS(i->first, 2); + ASSERT_EQUALS(*i->first, 2); ASSERT_EQUALS(*i->second, 2); ++i; - ASSERT_EQUALS(i->first, 1); + ASSERT_EQUALS(*i->first, 1); ASSERT_EQUALS(*i->second, 1); ++i; ASSERT(i == cache.end()); @@ -356,7 +388,8 @@ TEST(LRUKeyValueTest, UniquePtrKeyValue) { assertNotInKVStore(cacheForEviction, 1); // The entry with key '1' has been Evicted. } -using TestScalarValue = LRUKeyValue<int, ValueType, TrivialBudgetEstimator>; +using TestScalarValue = + LRUKeyValue<int, ValueType, TrivialBudgetEstimator, NonTrivialInsertionEvictionListener>; void assertValueInKVStore(TestScalarValue& cache, int key, ValueType value) { ASSERT_TRUE(cache.hasKey(key)); @@ -373,9 +406,17 @@ TEST(LRUKeyValueTest, ScalarKeyValue) { assertValueInKVStore(cache, 1, ValueType{2}); assertNotInKVStore(cache, 3); + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::keyTotal, 1); + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::valueTotal, 2); + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::budgetTotal, 1); + cache.add(1, ValueType{3}); assertValueInKVStore(cache, 1, ValueType{3}); + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::keyTotal, 1); + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::valueTotal, 3); + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::budgetTotal, 1); + // Test eviction. TestScalarValue cacheForEviction{2}; cacheForEviction.add(1, ValueType{1}); @@ -384,6 +425,18 @@ TEST(LRUKeyValueTest, ScalarKeyValue) { ASSERT_EQUALS(cacheForEviction.size(), static_cast<size_t>(2)); assertNotInKVStore(cacheForEviction, 1); // The entry with key '1' has been Evicted. + + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::keyTotal, 5); + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::valueTotal, 5); + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::budgetTotal, 2); + + // Clear the remaining values. + cacheForEviction.clear(); + + assertNotInKVStore(cacheForEviction, 2); // The entry with key '2' has been Evicted. + assertNotInKVStore(cacheForEviction, 3); // The entry with key '3' has been Evicted. + + ASSERT_EQUALS(NonTrivialInsertionEvictionListener::budgetTotal, 0); } } // namespace diff --git a/src/mongo/db/query/parsed_find_command.cpp b/src/mongo/db/query/parsed_find_command.cpp new file mode 100644 index 00000000000..2ef2e955c06 --- /dev/null +++ b/src/mongo/db/query/parsed_find_command.cpp @@ -0,0 +1,381 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/parsed_find_command.h" + +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery + +#include "mongo/db/cst/cst_parser.h" +#include "mongo/db/query/collation/collator_factory_interface.h" +#include "mongo/db/query/projection_ast_util.h" +#include "mongo/db/query/projection_parser.h" +#include "mongo/db/query/query_planner_common.h" +#include "mongo/db/query/query_request_helper.h" +#include "mongo/logv2/log.h" + +namespace mongo { + +namespace { +/** + * Does 'root' have a subtree of type 'subtreeType' with a node of type 'childType' inside? + */ +bool hasNodeInSubtree(const MatchExpression* root, + MatchExpression::MatchType childType, + MatchExpression::MatchType subtreeType) { + if (subtreeType == root->matchType()) { + return QueryPlannerCommon::hasNode(root, childType); + } + for (size_t i = 0; i < root->numChildren(); ++i) { + if (hasNodeInSubtree(root->getChild(i), childType, subtreeType)) { + return true; + } + } + return false; +} + +bool parsingCanProduceNoopMatchNodes(const ExtensionsCallback& extensionsCallback, + MatchExpressionParser::AllowedFeatureSet allowedFeatures) { + return extensionsCallback.hasNoopExtensions() && + (allowedFeatures & MatchExpressionParser::AllowedFeatures::kText || + allowedFeatures & MatchExpressionParser::AllowedFeatures::kJavascript); +} + +} // namespace + +std::unique_ptr<CollatorInterface> resolveCollator( + OperationContext* opCtx, const std::unique_ptr<FindCommandRequest>& findCommand) { + if (!findCommand->getCollation().isEmpty()) { + return uassertStatusOKWithContext(CollatorFactoryInterface::get(opCtx->getServiceContext()) + ->makeFromBSON(findCommand->getCollation()), + "unable to parse collation"); + } + return nullptr; +} + +/** + * Helper for building 'out.' If there is a projection, parse it and add any metadata dependencies + * it induces. + * + * Throws exceptions if there is an error parsing the projection. + */ +void setProjection(ParsedFindCommand* out, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const std::unique_ptr<FindCommandRequest>& findCommand, + const ProjectionPolicies& policies) { + if (!findCommand->getProjection().isEmpty()) { + out->savedProjectionPolicies.emplace(policies); + out->proj.emplace(projection_ast::parseAndAnalyze(expCtx, + findCommand->getProjection(), + out->filter.get(), + findCommand->getFilter(), + policies)); + + // This will throw if any of the projection's dependencies are unavailable. + DepsTracker{out->unavailableMetadata}.requestMetadata(out->proj->metadataDeps()); + } +} + +/** + * Helper for building 'out.' If there is a sort, parse it and add any metadata dependencies it + * induces. + * + * Throws exceptions if there is an error parsing the sort pattern. + */ +void setSort(ParsedFindCommand* out, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const std::unique_ptr<FindCommandRequest>& findCommand) { + if (!findCommand->getSort().isEmpty()) { + // A $natural sort is really a hint, and should be handled as such. Furthermore, the + // downstream sort handling code may not expect a $natural sort. + // + // We have already validated that if there is a $natural sort and a hint, that the hint + // also specifies $natural with the same direction. Therefore, it is safe to clear the + // $natural sort and rewrite it as a $natural hint. + if (findCommand->getSort()[query_request_helper::kNaturalSortField]) { + findCommand->setHint(findCommand->getSort().getOwned()); + findCommand->setSort(BSONObj{}); + } + if (getTestCommandsEnabled() && internalQueryEnableCSTParser.load()) { + out->sort = cst::parseToSortPattern(findCommand->getSort(), expCtx); + } else { + out->sort.emplace(findCommand->getSort(), expCtx); + } + } +} + +/** + * Helper for building 'out.' If there is a sort, parse it and add any metadata dependencies it + * induces. + */ +Status setSortAndProjection(ParsedFindCommand* out, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const std::unique_ptr<FindCommandRequest>& findCommand, + const ProjectionPolicies& policies) { + try { + setProjection(out, expCtx, findCommand, policies); + setSort(out, expCtx, findCommand); + } catch (const DBException& ex) { + return ex.toStatus(); + } + + return Status::OK(); +} + +/** + * Helper for building 'out.' Sets 'out->filter' and validates that it is well formed. In the + * process, also populates 'out->unavailableMetadata.' + */ +Status setFilter(ParsedFindCommand* out, + std::unique_ptr<MatchExpression> filter, + const std::unique_ptr<FindCommandRequest>& findCommand) { + // Verify the filter follows certain rules like there must be at most one text clause. + auto swMeta = parsed_find_command::isValid(filter.get(), *findCommand); + if (!swMeta.isOK()) { + return swMeta.getStatus(); + } + out->unavailableMetadata = swMeta.getValue(); + out->filter = std::move(filter); + return Status::OK(); +} + + +StatusWith<std::unique_ptr<ParsedFindCommand>> parseWithValidatedCollator( + const boost::intrusive_ptr<ExpressionContext>& expCtx, + std::unique_ptr<FindCommandRequest> findCommand, + const ExtensionsCallback& extensionsCallback, + MatchExpressionParser::AllowedFeatureSet allowedFeatures, + const ProjectionPolicies& projectionPolicies) { + auto out = std::make_unique<ParsedFindCommand>(); + + tassert(5746107, + "ntoreturn should not be set on the findCommand", + findCommand->getNtoreturn() == boost::none); + + if (auto status = query_request_helper::validateFindCommandRequest(*findCommand); + !status.isOK()) { + return status; + } + + // Parse the MatchExpression. + StatusWithMatchExpression statusWithMatcher = [&]() -> StatusWithMatchExpression { + if (getTestCommandsEnabled() && internalQueryEnableCSTParser.load()) { + try { + return cst::parseToMatchExpression( + findCommand->getFilter(), expCtx, extensionsCallback); + } catch (const DBException& ex) { + return ex.toStatus(); + } + } else { + return MatchExpressionParser::parse( + findCommand->getFilter(), expCtx, extensionsCallback, allowedFeatures); + } + }(); + if (!statusWithMatcher.isOK()) { + return statusWithMatcher.getStatus(); + } + + // Stop counting expressions after they have been parsed to exclude expressions created + // during optimization and other processing steps. + expCtx->stopExpressionCounters(); + out->canHaveNoopMatchNodes = + parsingCanProduceNoopMatchNodes(extensionsCallback, allowedFeatures); + + if (auto status = setFilter(out.get(), std::move(statusWithMatcher.getValue()), findCommand); + !status.isOK()) { + return status; + } + + if (auto status = setSortAndProjection(out.get(), expCtx, findCommand, projectionPolicies); + !status.isOK()) { + return status; + } + + out->findCommandRequest = std::move(findCommand); + return {std::move(out)}; +} + +StatusWith<std::unique_ptr<ParsedFindCommand>> ParsedFindCommand::withExistingFilter( + const boost::intrusive_ptr<ExpressionContext>& expCtx, + std::unique_ptr<CollatorInterface> collator, + std::unique_ptr<MatchExpression> filter, + std::unique_ptr<FindCommandRequest> findCommandRequest) { + auto out = std::make_unique<ParsedFindCommand>(); + out->collator = std::move(collator); + if (auto status = setFilter(out.get(), std::move(filter), findCommandRequest); !status.isOK()) { + return status; + } + if (auto status = setSortAndProjection( + out.get(), expCtx, findCommandRequest, ProjectionPolicies::findProjectionPolicies()); + !status.isOK()) { + return status; + } + out->findCommandRequest = std::move(findCommandRequest); + return std::move(out); +} + +namespace parsed_find_command { +StatusWith<QueryMetadataBitSet> isValid(const MatchExpression* root, + const FindCommandRequest& findCommand) { + QueryMetadataBitSet unavailableMetadata{}; + + // There can only be one TEXT. If there is a TEXT, it cannot appear inside a NOR. + // + // Note that the query grammar (as enforced by the MatchExpression parser) forbids TEXT + // inside of value-expression clauses like NOT, so we don't check those here. + size_t numText = QueryPlannerCommon::countNodes(root, MatchExpression::TEXT); + if (numText > 1) { + return Status(ErrorCodes::BadValue, "Too many text expressions"); + } else if (1 == numText) { + if (hasNodeInSubtree(root, MatchExpression::TEXT, MatchExpression::NOR)) { + return Status(ErrorCodes::BadValue, "text expression not allowed in nor"); + } + } else { + // Text metadata is not available. + unavailableMetadata.set(DocumentMetadataFields::kTextScore); + } + + // There can only be one NEAR. If there is a NEAR, it must be either the root or the root + // must be an AND and its child must be a NEAR. + size_t numGeoNear = QueryPlannerCommon::countNodes(root, MatchExpression::GEO_NEAR); + if (numGeoNear > 1) { + return Status(ErrorCodes::BadValue, "Too many geoNear expressions"); + } else if (1 == numGeoNear) { + // Do nothing, we will perform extra checks in CanonicalQuery::isValidNormalized. + } else { + // Geo distance and geo point metadata are unavailable. + unavailableMetadata |= DepsTracker::kAllGeoNearData; + } + + const BSONObj& sortObj = findCommand.getSort(); + BSONElement sortNaturalElt = sortObj["$natural"]; + const BSONObj& hintObj = findCommand.getHint(); + BSONElement hintNaturalElt = hintObj["$natural"]; + + if (sortNaturalElt && sortObj.nFields() != 1) { + return Status(ErrorCodes::BadValue, + str::stream() << "Cannot include '$natural' in compound sort: " << sortObj); + } + + if (hintNaturalElt && hintObj.nFields() != 1) { + return Status(ErrorCodes::BadValue, + str::stream() << "Cannot include '$natural' in compound hint: " << hintObj); + } + + // NEAR cannot have a $natural sort or $natural hint. + if (numGeoNear > 0) { + if (sortNaturalElt) { + return Status(ErrorCodes::BadValue, + "geoNear expression not allowed with $natural sort order"); + } + + if (hintNaturalElt) { + return Status(ErrorCodes::BadValue, + "geoNear expression not allowed with $natural hint"); + } + } + + // TEXT and NEAR cannot both be in the query. + if (numText > 0 && numGeoNear > 0) { + return Status(ErrorCodes::BadValue, "text and geoNear not allowed in same query"); + } + + // TEXT and {$natural: ...} sort order cannot both be in the query. + if (numText > 0 && sortNaturalElt) { + return Status(ErrorCodes::BadValue, "text expression not allowed with $natural sort order"); + } + + // TEXT and hint cannot both be in the query. + if (numText > 0 && !hintObj.isEmpty()) { + return Status(ErrorCodes::BadValue, "text and hint not allowed in same query"); + } + + // TEXT and tailable are incompatible. + if (numText > 0 && findCommand.getTailable()) { + return Status(ErrorCodes::BadValue, "text and tailable cursor not allowed in same query"); + } + + // NEAR and tailable are incompatible. + if (numGeoNear > 0 && findCommand.getTailable()) { + return Status(ErrorCodes::BadValue, + "Tailable cursors and geo $near cannot be used together"); + } + + // $natural sort order must agree with hint. + if (sortNaturalElt) { + if (!hintObj.isEmpty() && !hintNaturalElt) { + return Status(ErrorCodes::BadValue, "index hint not allowed with $natural sort order"); + } + if (hintNaturalElt) { + if (hintNaturalElt.numberInt() != sortNaturalElt.numberInt()) { + return Status(ErrorCodes::BadValue, + "$natural hint must be in the same direction as $natural sort order"); + } + } + } + + return unavailableMetadata; +} + +StatusWith<std::pair<boost::intrusive_ptr<ExpressionContext>, std::unique_ptr<ParsedFindCommand>>> +parse(OperationContext* opCtx, + std::unique_ptr<FindCommandRequest> findCommand, + const ExtensionsCallback& extensionsCallback, + MatchExpressionParser::AllowedFeatureSet allowedFeatures, + const ProjectionPolicies& projectionPolicies) { + // Make the expCtx. + invariant(findCommand->getNamespaceOrUUID().nss().has_value()); + auto expCtx = make_intrusive<ExpressionContext>( + opCtx, *findCommand, resolveCollator(opCtx, findCommand), true /* mayDbProfile */); + auto swResult = parseWithValidatedCollator( + expCtx, std::move(findCommand), extensionsCallback, allowedFeatures, projectionPolicies); + if (!swResult.isOK()) { + return swResult.getStatus(); + } + + return std::pair{std::move(expCtx), std::move(swResult.getValue())}; +} + +StatusWith<std::unique_ptr<ParsedFindCommand>> parse( + const boost::intrusive_ptr<ExpressionContext>& expCtx, + std::unique_ptr<FindCommandRequest> findCommand, + const ExtensionsCallback& extensionsCallback, + MatchExpressionParser::AllowedFeatureSet allowedFeatures, + const ProjectionPolicies& projectionPolicies) { + // A collator can enter through both the FindCommandRequest and ExpressionContext arguments. + // This invariant ensures that both collators are the same because downstream we + // pull the collator from only one of the ExpressionContext carrier. + auto collator = resolveCollator(expCtx->opCtx, findCommand); + if (collator.get() && expCtx->getCollator()) { + invariant(CollatorInterface::collatorsMatch(collator.get(), expCtx->getCollator())); + } + return parseWithValidatedCollator( + expCtx, std::move(findCommand), extensionsCallback, allowedFeatures, projectionPolicies); +} +} // namespace parsed_find_command +} // namespace mongo diff --git a/src/mongo/db/query/parsed_find_command.h b/src/mongo/db/query/parsed_find_command.h new file mode 100644 index 00000000000..f3ec78b9204 --- /dev/null +++ b/src/mongo/db/query/parsed_find_command.h @@ -0,0 +1,127 @@ +/** + * Copyright (C) 2023-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 "mongo/db/matcher/expression.h" +#include "mongo/db/query/find_command_gen.h" +#include "mongo/db/query/projection.h" +#include "mongo/db/query/projection_policies.h" +#include "mongo/db/query/sort_pattern.h" + +namespace mongo { + +/** + * Represents a find command request, but with more fully parsed ASTs for some fields which are + * still raw BSONObj on the FindCommandRequest type. + */ +struct ParsedFindCommand { + ParsedFindCommand() = default; + + /** + * This API adds the ability to construct from a pre-parsed filter. The other arguments will be + * re-parsed again from BSON on the 'findCommandRequest' argument, since we don't have a good + * way of cloning them. + */ + static StatusWith<std::unique_ptr<ParsedFindCommand>> withExistingFilter( + const boost::intrusive_ptr<ExpressionContext>& expCtx, + std::unique_ptr<CollatorInterface> collator, + std::unique_ptr<MatchExpression> filter, + std::unique_ptr<FindCommandRequest> findCommandRequest); + + std::unique_ptr<CollatorInterface> collator; + std::unique_ptr<MatchExpression> filter; + boost::optional<projection_ast::Projection> proj; + boost::optional<SortPattern> sort; + + // Based on parsing the query, which metadata will *not* be available. For example, if there is + // no $text clause, then a text score will not be available. + QueryMetadataBitSet unavailableMetadata; + + // This is saved for an edge case where we need to re-parse a projection later. Only populated + // if there is a non-empty projection. + boost::optional<ProjectionPolicies> savedProjectionPolicies; + + // True if this canonical query may have converted extensions such as $where and $text into + // no-ops during parsing. This will be the case if it allowed $where and $text in parsing, but + // parsed using an ExtensionsCallbackNoop. This does not guarantee that a $where or $text + // existed in the query. Queries with a no-op extension context are special because they can be + // parsed and planned, but they cannot be executed. + bool canHaveNoopMatchNodes; + + // All other parameters to the find command which do not have AST-like types and can be + // appropriately tracked as raw value types like ints. The fields above like 'filter' are all + // still present in their raw form on this FidnCommandRequest, but it is not expected that they + // will be useful other than to keep the original BSON values around in-memory to avoid copying + // large strings and such. + std::unique_ptr<FindCommandRequest> findCommandRequest; +}; + +namespace parsed_find_command { +/** + * Validates the match expression 'root' as well as the query specified by 'request', checking for + * illegal combinations of operators. Returns a non-OK status if any such illegal combination is + * found. + * + * This method can be called both on normalized and non-normalized 'root'. However, some checks can + * only be performed once the match expressions is normalized. To perform these checks one can call + * 'CanonicalQuery::isValidNormalized()'. + * + * On success, returns a bitset indicating which types of metadata are *unavailable*. For example, + * if 'root' does not contain a $text predicate, then the returned metadata bitset will indicate + * that text score metadata is unavailable. This means that if subsequent $meta:"textScore" + * expressions are found during analysis of the query, we should raise in an error. + */ +StatusWith<QueryMetadataBitSet> isValid(const MatchExpression* root, + const FindCommandRequest& findCommand); + +/** + * Parses each big component of the input 'findCommand.' Throws exceptions if failing to parse. + * Comes in one overload which will create an ExpressionContext for the caller, and one overload to + * be used when the caller already has an ExpressionContext. + */ +StatusWith<std::pair<boost::intrusive_ptr<ExpressionContext>, std::unique_ptr<ParsedFindCommand>>> +parse(OperationContext* opCtx, + std::unique_ptr<FindCommandRequest> findCommand, + const ExtensionsCallback& extensionsCallback = ExtensionsCallbackNoop(), + MatchExpressionParser::AllowedFeatureSet allowedFeatures = + MatchExpressionParser::kDefaultSpecialFeatures, + const ProjectionPolicies& projectionPolicies = ProjectionPolicies::findProjectionPolicies()); + +// Overload of the above for when the caller has an available ExpressionContext. +StatusWith<std::unique_ptr<ParsedFindCommand>> parse( + const boost::intrusive_ptr<ExpressionContext>& expCtx, + std::unique_ptr<FindCommandRequest> findCommand, + const ExtensionsCallback& extensionsCallback = ExtensionsCallbackNoop(), + MatchExpressionParser::AllowedFeatureSet allowedFeatures = + MatchExpressionParser::kDefaultSpecialFeatures, + const ProjectionPolicies& projectionPolicies = ProjectionPolicies::findProjectionPolicies()); + +} // namespace parsed_find_command +} // namespace mongo diff --git a/src/mongo/db/query/partitioned_cache.h b/src/mongo/db/query/partitioned_cache.h new file mode 100644 index 00000000000..ca10f731b71 --- /dev/null +++ b/src/mongo/db/query/partitioned_cache.h @@ -0,0 +1,242 @@ +/** + * Copyright (C) 2022-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 "mongo/db/catalog/util/partitioned.h" +#include "mongo/db/commands/server_status_metric.h" +#include "mongo/db/query/lru_key_value.h" +#include "mongo/platform/mutex.h" +#include "mongo/util/container_size_helper.h" + +namespace mongo { + +/** + * A partitioned cache combines a size-bounded map (LRU-based entry eviction) with a partition + * function which allows reducing contention. + */ +template <class KeyType, + class ValueType, + class KeyBudgetEstimator, + class Partitioner, + class InsertionEvictionListener, + class KeyHasher = std::hash<KeyType>, + class Eq = std::equal_to<KeyType>> +class PartitionedCache { +private: + PartitionedCache(const PartitionedCache&) = delete; + PartitionedCache& operator=(const PartitionedCache&) = delete; + +public: + using Lru = LRUKeyValue<KeyType, + ValueType, + KeyBudgetEstimator, + InsertionEvictionListener, + KeyHasher, + Eq>; + using Partition = typename Partitioned<Lru, Partitioner>::OnePartition; + using PartitionId = typename Partitioned<Lru, Partitioner>::PartitionId; + + /** + * Initialize plan cache with the total cache size in bytes and number of partitions. + * + * Important edge cases to consider include: + * + * 1. Adding an entry that is larger than the max partition size to a non-empty partition. + * + * This will evict both entries. This is because entries are evicted from the partition in + * order of least recently used. Thus, the oldest, small entry will be evicted first but the + * partition will still be over budget with the new, too-large entry so it will be evicted as + * well. + * + * 2. Adding a queryStats store entry that is smaller than the overall cache size but larger + * than single partition max size. + * + * It is not possible to write entries to the cache that are larger than a single + * partition's max size, even if it is smaller than the entire cache max size. This is because + * the cache's budget is configured/regulated on the partition level (cacheSize / + * numPartitions). This makes sense as each entry is written to a specific partition, but might + * not be immediately obvious so worthy to highlight. + * + * 3. Too few partitions can cause unnecessary evictions + * + * Every class that implements the PartitionedCache template provides a partitioner() that + * returns the id of the partition to which to write the entry. In existing implementations, + * partitioner() returns the remainder after dividing the entry's key hash by numPartitions. In + * the case where we have only two partitions, every odd key hash will be written to the first + * partition (and vice versa). In this way, it can quickly be the case that one partition + * fills up completely but the partitioner() call keeps returning the already full partition and + * the cache evict old entries from it to put the new one in. At the end of all the write + * operations, the cache is below it's budget (as the second partition is only partially full) + * but we don't have all the entries we expect. It is therefore important to have sufficient + * enough number of partitions so the entries can be more equally dispersed to avoid unnecessary + * evictions. + */ + explicit PartitionedCache(size_t cacheSize, size_t numPartitions) + : _numPartitions(numPartitions) { + invariant(numPartitions > 0); + Lru lru{cacheSize / numPartitions}; + _partitionedCache = + std::make_unique<Partitioned<Lru, Partitioner>>(numPartitions, std::move(lru)); + } + + ~PartitionedCache() = default; + /** + * Inserts the provided <key, value> into the partition associated with that key. Returns the + * number of older entries evicted to fit this new one. + */ + size_t put(const KeyType& key, ValueType value) { + auto partition = _partitionedCache->lockOnePartition(key); + return partition->add(key, std::move(value)); + } + /** + * Inserts the provided <key, value> into the specified partition. Returns the number of older + * entries evicted to fit this new one. + */ + size_t put(const KeyType& key, ValueType value, Partition& partition) { + return partition->add(key, std::move(value)); + } + + StatusWith<ValueType*> lookup(const KeyType& key) const { + auto partition = _partitionedCache->lockOnePartition(key); + auto entry = partition->get(key); + if (!entry.isOK()) { + return {entry.getStatus()}; + } + + return {&entry.getValue()->second}; + } + + /** + * Lookup an entry and also return a lock over the partition. The lock is returned whether + * or not the entry is found. + */ + std::pair<StatusWith<ValueType*>, Partition> getWithPartitionLock(const KeyType& key) const { + auto partition = _partitionedCache->lockOnePartition(key); + auto entry = partition->get(key); + if (!entry.isOK()) { + return std::make_pair(entry.getStatus(), std::move(partition)); + } + + return std::make_pair(StatusWith{&entry.getValue()->second}, std::move(partition)); + } + + /** + * Remove the entry with the 'key' from the cache. If there is no entry for the given key in + * the cache, this call is a no-op. + */ + void remove(const KeyType& key) { + _partitionedCache->erase(key); + } + + /** + * Remove all the entries for keys for which the predicate returns true. Return the number of + * removed entries. + */ + template <typename UnaryPredicate> + size_t removeIf(UnaryPredicate predicate) { + size_t nRemoved = 0; + for (size_t partitionId = 0; partitionId < _numPartitions; ++partitionId) { + auto lockedPartition = _partitionedCache->lockOnePartitionById(partitionId); + nRemoved += lockedPartition->removeIf(predicate); + } + return nRemoved; + } + + /** + * Remove *all* cache entries. + */ + void clear() { + _partitionedCache->clear(); + } + + /** + * Reset total cache size. If the size is set to a smaller value than before, enough entries are + * evicted in order to ensure that the cache fits within the new budget. Returns the number of + * entries evicted. + */ + size_t reset(size_t cacheSize) { + size_t numEvicted = 0; + for (size_t partitionId = 0; partitionId < _numPartitions; ++partitionId) { + auto lockedPartition = _partitionedCache->lockOnePartitionById(partitionId); + numEvicted += lockedPartition->reset(cacheSize / _numPartitions); + } + + return numEvicted; + } + + /** + * Returns the size of the cache. + * Used for testing. + */ + size_t size() const { + return _partitionedCache->size(); + } + + /** + * Returns the number of partitions. + */ + size_t numPartitions() const { + return _numPartitions; + } + + /** + * Invoke `op` for each entry in the cache. Consistency across partitions is not guaranteed. + */ + void forEach(const std::function<void(const KeyType&, const ValueType&)>& op) const { + for (size_t partitionId = 0; partitionId < _numPartitions; ++partitionId) { + auto lockedPartition = _partitionedCache->lockOnePartitionById(partitionId); + + for (auto&& [key, entry] : *lockedPartition) { + op(*key, entry); + } + } + } + + /** + * Allow iterating over partitions. The provided function is called for each partition. The + * argument to the function is another function which can delay acquiring the implicitly locked + * partition until it's needed. + */ + void forEachPartition(const std::function<void(const std::function<Partition()>&)>& op) const { + for (size_t partitionId = 0; partitionId < _numPartitions; ++partitionId) { + op([&]() { return _partitionedCache->lockOnePartitionById(partitionId); }); + } + } + + Partition getPartition(PartitionId partitionId) { + return _partitionedCache->lockOnePartitionById(partitionId); + } + +private: + std::size_t _numPartitions; + std::unique_ptr<Partitioned<Lru, Partitioner>> _partitionedCache; +}; + +} // namespace mongo diff --git a/src/mongo/db/query/plan_cache.h b/src/mongo/db/query/plan_cache.h index dc52b10c4f4..e193309230e 100644 --- a/src/mongo/db/query/plan_cache.h +++ b/src/mongo/db/query/plan_cache.h @@ -31,6 +31,7 @@ #include "mongo/db/catalog/util/partitioned.h" #include "mongo/db/query/lru_key_value.h" +#include "mongo/db/query/partitioned_cache.h" #include "mongo/db/query/plan_cache_callbacks.h" #include "mongo/db/query/plan_cache_debug_info.h" #include "mongo/platform/mutex.h" @@ -281,21 +282,34 @@ private: */ template <class KeyType, class CachedPlanType, - class BudgetEstimator, + class KeyBudgetEstimator, class DebugInfoType, class Partitioner, class KeyHasher = std::hash<KeyType>> -class PlanCacheBase { +class PlanCacheBase + : public PartitionedCache< + KeyType, + // The 'Value' being "std::shared_ptr<const Entry>" is because we allow readers to clone + // cache entries out of the lock, therefore it is illegal to mutate the pieces of a cache + // entry that can be cloned whether you are holding a lock or not. + std::shared_ptr<const PlanCacheEntryBase<CachedPlanType, DebugInfoType>>, + KeyBudgetEstimator, + Partitioner, + NoopInsertionEvictionListener, + KeyHasher> { private: PlanCacheBase(const PlanCacheBase&) = delete; PlanCacheBase& operator=(const PlanCacheBase&) = delete; public: + using Base = + PartitionedCache<KeyType, + std::shared_ptr<const PlanCacheEntryBase<CachedPlanType, DebugInfoType>>, + KeyBudgetEstimator, + Partitioner, + NoopInsertionEvictionListener, + KeyHasher>; using Entry = PlanCacheEntryBase<CachedPlanType, DebugInfoType>; - // The 'Value' being "std::shared_ptr<const Entry>" is because we allow readers to clone cache - // entries out of the lock, therefore it is illegal to mutate the pieces of a cache entry that - // can be cloned whether you are holding a lock or not. - using Lru = LRUKeyValue<KeyType, std::shared_ptr<const Entry>, BudgetEstimator, KeyHasher>; // We have three states for a cache entry to be in. Rather than just 'present' or 'not // present', we use a notion of 'inactive entries' as a way of remembering how performant our @@ -328,11 +342,7 @@ public: * Initialize plan cache with the total cache size in bytes and number of partitions. */ explicit PlanCacheBase(size_t cacheSize, size_t numPartitions = 1) - : _numPartitions(numPartitions) { - invariant(numPartitions > 0); - Lru lru{cacheSize / numPartitions}; - _partitionedCache = std::make_unique<Partitioned<Lru, Partitioner>>(numPartitions, lru); - } + : Base(cacheSize, numPartitions) {} ~PlanCacheBase() = default; @@ -375,7 +385,11 @@ public: }}, why.stats); - auto partition = _partitionedCache->lockOnePartition(key); + auto oldEntryWithPartitionLock = this->getWithPartitionLock(key); + // Can't use reference to structured bindings in a lambda until C++20 so manually + // destructure it here. + auto partitionLock = std::move(oldEntryWithPartitionLock.second); + auto oldEntryWithStatus = std::move(oldEntryWithPartitionLock.first); auto [queryHash, planCacheKey, isNewEntryActive, shouldBeCreated, increasedWorks] = [&]() { if (internalQueryCacheDisableInactiveEntries.load()) { // All entries are always active. @@ -385,32 +399,34 @@ public: true /* shouldBeCreated */, boost::optional<size_t>(boost::none)); } else { - auto oldEntryWithStatus = partition->get(key); tassert(6007020, "LRU store must get value or NoSuchKey error code", oldEntryWithStatus.isOK() || oldEntryWithStatus.getStatus() == ErrorCodes::NoSuchKey); - auto oldEntry = - oldEntryWithStatus.isOK() ? oldEntryWithStatus.getValue()->second : nullptr; + bool hasOldEntry = oldEntryWithStatus.isOK(); const auto newState = getNewEntryState( key, - oldEntry.get(), + // Deference the pointer, then the shared_ptr, and then back to a raw pointer. + hasOldEntry ? &**oldEntryWithStatus.getValue() : nullptr, newWorks, worksGrowthCoefficient.get_value_or(internalQueryCacheWorksGrowthCoefficient), callbacks); // Avoid recomputing the hashes if we've got an old entry to grab them from. - return oldEntry ? std::make_tuple(oldEntry->queryHash, - oldEntry->planCacheKey, - newState.shouldBeActive, - newState.shouldBeCreated, - newState.increasedWorks) - : std::make_tuple(key.queryHash(), - key.planCacheKeyHash(), - newState.shouldBeActive, - newState.shouldBeCreated, - newState.increasedWorks); + auto [queryHash, planCacheKey] = [&]() { + if (hasOldEntry) { + auto&& oldEntry = &**oldEntryWithStatus.getValue(); + return std::make_pair(oldEntry->queryHash, oldEntry->planCacheKey); + } else { + return std::make_pair(key.queryHash(), key.planCacheKeyHash()); + } + }(); + return std::make_tuple(queryHash, + planCacheKey, + newState.shouldBeActive, + newState.shouldBeCreated, + newState.increasedWorks); } }(); @@ -434,7 +450,7 @@ public: increasedWorks ? *increasedWorks : newWorks, callbacks->buildDebugInfo()); - partition->add(key, std::move(newEntry)); + this->put(key, std::move(newEntry), partitionLock); return Status::OK(); } @@ -454,10 +470,7 @@ public: indexFilterKey, now, std::move(debugInfo)); - auto partition = _partitionedCache->lockOnePartition(key); - // We're not interested in the number of evicted entries if the cache store exceeds the - // budget after add(), so we just ignore the return value. - partition->add(key, std::move(entry)); + this->put(key, std::move(entry)); } /** @@ -471,8 +484,8 @@ public: return; } - auto partition = _partitionedCache->lockOnePartition(key); - auto entry = partition->get(key); + auto [entry, partitionLock] = this->getWithPartitionLock(key); + if (!entry.isOK()) { tassert(6007021, "Unexpected error code from LRU store", @@ -480,11 +493,11 @@ public: return; } - auto entryPtr = entry.getValue()->second; + auto entryPtr = *entry.getValue(); if (entryPtr->isActive == true) { std::shared_ptr<Entry> newEntry = entryPtr->clone(); newEntry->isActive = false; - partition->add(key, std::move(newEntry)); + this->put(key, std::move(newEntry), partitionLock); } } @@ -496,29 +509,26 @@ public: * for the query (if there is one). */ GetResult get(const KeyType& key) const { - std::shared_ptr<const Entry> entryPtr; + std::shared_ptr<const Entry> entrySharedPtr; CacheEntryState state; { - auto partition = _partitionedCache->lockOnePartition(key); - auto entry = partition->get(key); + auto [entry, partitionLock] = this->getWithPartitionLock(key); if (!entry.isOK()) { tassert(6007023, "Unexpected error code from LRU store", entry.getStatus() == ErrorCodes::NoSuchKey); return {CacheEntryState::kNotPresent, nullptr}; } - entryPtr = entry.getValue()->second; - state = entryPtr->isActive ? CacheEntryState::kPresentActive - : CacheEntryState::kPresentInactive; + entrySharedPtr = *entry.getValue(); + state = entrySharedPtr->isActive ? CacheEntryState::kPresentActive + : CacheEntryState::kPresentInactive; } - // The purpose of cloning 'entry' after we release the lock is to allow multiple threads to - // clone the same plan cache entry at once. 'entry' cannot be deleted by another thread even - // if the plan cache is being concurrently modified by other threads because we are holding - // a std::shared_ptr to this entry. - tassert(6007024, "LRU store must get a value or an error code", entryPtr); - + // The purpose of cloning 'entry' (in CachedPlanHolder ctor) after we release the lock + // is to allow multiple threads to clone the same plan cache entry at once. 'entry' + // cannot be deleted by another thread even if the plan cache is being concurrently + // modified by other threads because we are holding a std::shared_ptr to this entry. return {state, - std::make_unique<CachedPlanHolder<CachedPlanType, DebugInfoType>>(*entryPtr)}; + std::make_unique<CachedPlanHolder<CachedPlanType, DebugInfoType>>(*entrySharedPtr)}; } /** @@ -537,59 +547,16 @@ public: } /** - * Remove the entry with the 'key' from the cache. If there is no entry for the given key in - * the cache, this call is a no-op. - */ - void remove(const KeyType& key) { - _partitionedCache->erase(key); - } - - /** - * Remove all the entries for keys for which the predicate returns true. Return the number of - * removed entries. - */ - template <typename UnaryPredicate> - size_t removeIf(UnaryPredicate predicate) { - size_t nRemoved = 0; - for (size_t partitionId = 0; partitionId < _numPartitions; ++partitionId) { - auto lockedPartition = _partitionedCache->lockOnePartitionById(partitionId); - nRemoved += lockedPartition->removeIf(predicate); - } - return nRemoved; - } - - /** - * Remove *all* cached plans. Does not clear index information. - */ - void clear() { - _partitionedCache->clear(); - } - - /** - * Reset total cache size. If the size is set to a smaller value than before, enough entries are - * evicted in order to ensure that the cache fits within the new budget. - */ - void reset(size_t cacheSize) { - for (size_t partitionId = 0; partitionId < _numPartitions; ++partitionId) { - auto lockedPartition = _partitionedCache->lockOnePartitionById(partitionId); - lockedPartition->reset(cacheSize / _numPartitions); - } - } - - /** * Returns a copy of a cache entry, looked up by the plan cache key. * * If there is no entry in the cache for the 'query', returns an error Status. */ StatusWith<std::unique_ptr<Entry>> getEntry(const KeyType& key) const { - auto partition = _partitionedCache->lockOnePartition(key); - auto entry = partition->get(key); - if (!entry.isOK()) { - return entry.getStatus(); + auto result = this->lookup(key); + if (!result.isOK()) { + return {result.getStatus()}; } - invariant(entry.getValue()->second); - - return std::unique_ptr<Entry>(entry.getValue()->second->clone()); + return {result.getValue()->get()->clone()}; } /** @@ -598,26 +565,14 @@ public: std::vector<std::unique_ptr<Entry>> getAllEntries() const { std::vector<std::unique_ptr<Entry>> entries; - for (size_t partitionId = 0; partitionId < _numPartitions; ++partitionId) { - auto lockedPartition = _partitionedCache->lockOnePartitionById(partitionId); - - for (auto&& [key, entry] : *lockedPartition) { - entries.emplace_back(entry->clone()); - } - } + this->forEach([&](const KeyType& key, const std::shared_ptr<Entry>& entry) { + entries.emplace_back(entry); + }); return entries; } /** - * Returns the size of the cache. - * Used for testing. - */ - size_t size() const { - return _partitionedCache->size(); - } - - /** * Iterates over the plan cache. For each entry, first filters according to the predicate * function 'cacheKeyFilterFunc', (Note that 'cacheKeyFilterFunc' could be empty, if so, we * don't filter by plan cache key.), then serializes the PlanCacheEntryBase according to @@ -634,20 +589,15 @@ public: std::vector<BSONObj> results; - for (size_t partitionId = 0; partitionId < _numPartitions; ++partitionId) { - auto lockedPartition = _partitionedCache->lockOnePartitionById(partitionId); - - for (auto&& cacheEntry : *lockedPartition) { - if (cacheKeyFilterFunc && !cacheKeyFilterFunc(cacheEntry.first)) { - continue; - } - const auto& entry = cacheEntry.second; - auto serializedEntry = serializationFunc(*entry); - if (filterFunc(serializedEntry)) { - results.push_back(serializedEntry); - } + this->forEach([&](const KeyType& key, const std::shared_ptr<const Entry>& entry) { + if (cacheKeyFilterFunc && !cacheKeyFilterFunc(key)) { + return; } - } + auto serializedEntry = serializationFunc(*entry); + if (filterFunc(serializedEntry)) { + results.push_back(serializedEntry); + } + }); return results; } @@ -735,9 +685,6 @@ private: return res; } - - std::size_t _numPartitions; - std::unique_ptr<Partitioned<Lru, Partitioner>> _partitionedCache; }; } // namespace mongo diff --git a/src/mongo/db/query/plan_cache_size_parameter_test.cpp b/src/mongo/db/query/plan_cache_size_parameter_test.cpp deleted file mode 100644 index 1f0fbf76a27..00000000000 --- a/src/mongo/db/query/plan_cache_size_parameter_test.cpp +++ /dev/null @@ -1,79 +0,0 @@ -/** - * 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. - */ - -#include "mongo/db/query/plan_cache_size_parameter.h" - -#include "mongo/unittest/unittest.h" - -namespace mongo::plan_cache_util { - -bool operator==(const PlanCacheSizeParameter& lhs, const PlanCacheSizeParameter& rhs) { - constexpr double kEpsilon = 1e-10; - return std::abs(lhs.size - rhs.size) < kEpsilon && lhs.units == rhs.units; -} - -TEST(PlanCacheParameterTest, ParseUnitStringPercent) { - ASSERT_TRUE(PlanCacheSizeUnits::kPercent == parseUnitString("%")); -} - -TEST(PlanCacheParameterTest, ParseUnitStringMB) { - ASSERT_TRUE(PlanCacheSizeUnits::kMB == parseUnitString("MB")); - ASSERT_TRUE(PlanCacheSizeUnits::kMB == parseUnitString("mb")); - ASSERT_TRUE(PlanCacheSizeUnits::kMB == parseUnitString("mB")); - ASSERT_TRUE(PlanCacheSizeUnits::kMB == parseUnitString("Mb")); -} - -TEST(PlanCacheParameterTest, ParseUnitStringGB) { - ASSERT_TRUE(PlanCacheSizeUnits::kGB == parseUnitString("GB")); - ASSERT_TRUE(PlanCacheSizeUnits::kGB == parseUnitString("gb")); - ASSERT_TRUE(PlanCacheSizeUnits::kGB == parseUnitString("gB")); - ASSERT_TRUE(PlanCacheSizeUnits::kGB == parseUnitString("Gb")); -} - -TEST(PlanCacheParameterTest, ParseUnitStringIncorrectValue) { - ASSERT_NOT_OK(parseUnitString("").getStatus()); - ASSERT_NOT_OK(parseUnitString(" ").getStatus()); - ASSERT_NOT_OK(parseUnitString("KB").getStatus()); -} - -TEST(PlanCacheParameterTest, ParsePlanCacheSizeParameter) { - ASSERT_TRUE((PlanCacheSizeParameter{10.0, PlanCacheSizeUnits::kPercent}) == - PlanCacheSizeParameter::parse("10%")); - ASSERT_TRUE((PlanCacheSizeParameter{300.0, PlanCacheSizeUnits::kMB}) == - PlanCacheSizeParameter::parse("300MB")); - ASSERT_TRUE((PlanCacheSizeParameter{4.0, PlanCacheSizeUnits::kGB}) == - PlanCacheSizeParameter::parse("4GB")); - ASSERT_TRUE((PlanCacheSizeParameter{5.1, PlanCacheSizeUnits::kPercent}) == - PlanCacheSizeParameter::parse(" 5.1%")); - ASSERT_TRUE((PlanCacheSizeParameter{11.1, PlanCacheSizeUnits::kMB}) == - PlanCacheSizeParameter::parse("11.1 mb")); - ASSERT_TRUE((PlanCacheSizeParameter{12.1, PlanCacheSizeUnits::kGB}) == - PlanCacheSizeParameter::parse(" 12.1 Gb ")); -} -} // namespace mongo::plan_cache_util diff --git a/src/mongo/db/query/plan_enumerator.cpp b/src/mongo/db/query/plan_enumerator.cpp index 4f15905dc6d..9ad6f6b81aa 100644 --- a/src/mongo/db/query/plan_enumerator.cpp +++ b/src/mongo/db/query/plan_enumerator.cpp @@ -35,6 +35,7 @@ #include "mongo/db/query/index_tag.h" #include "mongo/db/query/indexability.h" +#include "mongo/db/query/query_planner_common.h" #include "mongo/logv2/log.h" #include "mongo/util/string_map.h" @@ -58,8 +59,8 @@ std::string getPathPrefix(std::string path) { * is a predicate that is required to use an index. */ bool expressionRequiresIndex(const MatchExpression* node) { - return CanonicalQuery::countNodes(node, MatchExpression::GEO_NEAR) > 0 || - CanonicalQuery::countNodes(node, MatchExpression::TEXT) > 0; + return QueryPlannerCommon::countNodes(node, MatchExpression::GEO_NEAR) > 0 || + QueryPlannerCommon::countNodes(node, MatchExpression::TEXT) > 0; } size_t getPathLength(const MatchExpression* expr) { @@ -261,7 +262,8 @@ PlanEnumerator::PlanEnumerator(const PlanEnumeratorParams& params) _ixisect(params.intersect), _enumerateOrChildrenLockstep(params.enumerateOrChildrenLockstep), _orLimit(params.maxSolutionsPerOr), - _intersectLimit(params.maxIntersectPerAnd) {} + _intersectLimit(params.maxIntersectPerAnd), + _disableOrPushdown(params.disableOrPushdown) {} PlanEnumerator::~PlanEnumerator() { typedef stdx::unordered_map<MemoID, NodeAssignment*> MemoMap; @@ -528,10 +530,14 @@ bool PlanEnumerator::prepMemo(MatchExpression* node, PrepMemoContext context) { // preds to 'indexedPreds'. Adding the mandatory preds directly to 'indexedPreds' would lead // to problems such as pulling a predicate beneath an OR into a set joined by an AND. getIndexedPreds(node, childContext, &indexedPreds); - // Pass in the indexed predicates as outside predicates when prepping the subnodes. + // Pass in the indexed predicates as outside predicates when prepping the subnodes. But if + // match expression optimization is disabled, skip this part: we don't want to do + // OR-pushdown because it relies on the expression being canonicalized. auto childContextCopy = childContext; - for (auto pred : indexedPreds) { - childContextCopy.outsidePreds[pred] = OutsidePredRoute{}; + if (MONGO_likely(!_disableOrPushdown)) { + for (auto pred : indexedPreds) { + childContextCopy.outsidePreds[pred] = OutsidePredRoute{}; + } } if (!prepSubNodes(node, childContextCopy, &subnodes, &mandatorySubnodes)) { return false; @@ -835,6 +841,13 @@ void PlanEnumerator::assignPredicate( MatchExpression* pred, size_t position, OneIndexAssignment* indexAssignment) { + if (MONGO_unlikely(_disableOrPushdown)) { + // If match expression optimization is disabled, we also disable OR-pushdown, + // so we should never get 'outsidePreds' here. + tassert(7059700, + "Tried to do OR-pushdown despite disableMatchExpressionOptimization", + outsidePreds.empty()); + } if (outsidePreds.find(pred) != outsidePreds.end()) { OrPushdownTag::Destination dest; dest.route = outsidePreds.at(pred).route; @@ -1289,6 +1302,8 @@ void PlanEnumerator::getIndexedPreds(MatchExpression* node, std::vector<MatchExpression*>* indexedPreds) { if (Indexability::nodeCanUseIndexOnOwnField(node)) { RelevantTag* rt = static_cast<RelevantTag*>(node->getTag()); + tassert(9074700, "RelevantTag is not assigned to the match expression node", rt != nullptr); + if (context.elemMatchExpr) { // If we're in an $elemMatch context, store the // innermost parent $elemMatch, as well as the @@ -1305,7 +1320,7 @@ void PlanEnumerator::getIndexedPreds(MatchExpression* node, indexedPreds->push_back(node); } else if (Indexability::isBoundsGeneratingNot(node)) { getIndexedPreds(node->getChild(0), context, indexedPreds); - } else if (MatchExpression::ELEM_MATCH_OBJECT == node->matchType()) { + } else if (Indexability::isBoundsGeneratingElemMatchObject(node)) { PrepMemoContext childContext; childContext.elemMatchExpr = node; for (size_t i = 0; i < node->numChildren(); ++i) { diff --git a/src/mongo/db/query/plan_enumerator.h b/src/mongo/db/query/plan_enumerator.h index b82b738c57b..9eabd5b09b9 100644 --- a/src/mongo/db/query/plan_enumerator.h +++ b/src/mongo/db/query/plan_enumerator.h @@ -44,7 +44,8 @@ namespace mongo { struct PlanEnumeratorParams { PlanEnumeratorParams() : maxSolutionsPerOr(internalQueryEnumerationMaxOrSolutions.load()), - maxIntersectPerAnd(internalQueryEnumerationMaxIntersectPerAnd.load()) {} + maxIntersectPerAnd(internalQueryEnumerationMaxIntersectPerAnd.load()), + disableOrPushdown(disableMatchExpressionOptimization.shouldFail()) {} // Do we provide solutions that use more indices than the minimum required to provide // an indexed solution? @@ -69,6 +70,11 @@ struct PlanEnumeratorParams { // all-pairs approach, we could wind up creating a lot of enumeration possibilities for // certain inputs. size_t maxIntersectPerAnd; + + // Whether to disable OR-pushdown optimization. OR-pushdown assumes that the expression has been + // simplified: for example, that single-child $or nodes are unwrapped. To avoid this, when + // the 'disableMatchExpressionOptimization' failpoint is set, we also disable OR-pushdown. + bool disableOrPushdown; }; /** @@ -594,6 +600,9 @@ private: // How many things do we want from each AND? size_t _intersectLimit; + + // Whether we should disable OR-pushdown optimization. + const bool _disableOrPushdown; }; } // namespace mongo diff --git a/src/mongo/db/query/plan_executor_factory.cpp b/src/mongo/db/query/plan_executor_factory.cpp index 0b66c3dafce..386bf38a538 100644 --- a/src/mongo/db/query/plan_executor_factory.cpp +++ b/src/mongo/db/query/plan_executor_factory.cpp @@ -31,6 +31,8 @@ #include "mongo/platform/basic.h" +#include <iostream> + #include "mongo/db/query/plan_executor_factory.h" #include "mongo/db/exec/plan_stage.h" @@ -40,6 +42,7 @@ #include "mongo/db/query/query_planner_params.h" #include "mongo/db/query/util/make_data_structure.h" #include "mongo/logv2/log.h" +#include "mongo/util/duration.h" namespace mongo::plan_executor_factory { @@ -65,6 +68,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make( yieldPolicy); } + StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make( const boost::intrusive_ptr<ExpressionContext>& expCtx, std::unique_ptr<WorkingSet> ws, @@ -74,6 +78,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make( size_t plannerOptions, NamespaceString nss, std::unique_ptr<QuerySolution> qs) { + return make(expCtx->opCtx, std::move(ws), std::move(rt), @@ -98,6 +103,7 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make( NamespaceString nss, PlanYieldPolicy::YieldPolicy yieldPolicy) { dassert(collection); + try { auto execImpl = new PlanExecutorImpl(opCtx, std::move(ws), @@ -128,7 +134,6 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make( NamespaceString nss, std::unique_ptr<PlanYieldPolicySBE> yieldPolicy) { auto&& [rootStage, data] = root; - LOGV2_DEBUG(4822860, 5, "SBE plan", @@ -157,7 +162,6 @@ StatusWith<std::unique_ptr<PlanExecutor, PlanExecutor::Deleter>> make( size_t plannerOptions, NamespaceString nss, std::unique_ptr<PlanYieldPolicySBE> yieldPolicy) { - LOGV2_DEBUG(4822861, 5, "SBE plan", diff --git a/src/mongo/db/query/plan_executor_factory.h b/src/mongo/db/query/plan_executor_factory.h index bf41f169af9..df7184583ec 100644 --- a/src/mongo/db/query/plan_executor_factory.h +++ b/src/mongo/db/query/plan_executor_factory.h @@ -29,6 +29,7 @@ #pragma once +#include "mongo/util/duration.h" #include <queue> #include "mongo/db/exec/sbe/stages/stages.h" diff --git a/src/mongo/db/query/plan_executor_impl.cpp b/src/mongo/db/query/plan_executor_impl.cpp index 03de8319f5c..ae40ecd070e 100644 --- a/src/mongo/db/query/plan_executor_impl.cpp +++ b/src/mongo/db/query/plan_executor_impl.cpp @@ -33,6 +33,7 @@ #include "mongo/db/query/plan_executor_impl.h" +#include "mongo/util/duration.h" #include <memory> #include "mongo/bson/simple_bsonobj_comparator.h" @@ -136,13 +137,6 @@ PlanExecutorImpl::PlanExecutorImpl(OperationContext* opCtx, invariant(!_expCtx || _expCtx->opCtx == _opCtx); invariant(!_cq || !_expCtx || _cq->getExpCtx() == _expCtx); - // If this PlanExecutor is executing a COLLSCAN, keep a pointer directly to the COLLSCAN - // stage. This is used for change streams in order to keep the the latest oplog timestamp - // and post batch resume token up to date as the oplog scan progresses. - if (auto collectionScan = getStageByType(_root.get(), STAGE_COLLSCAN)) { - _collScanStage = static_cast<CollectionScan*>(collectionScan); - } - // If we don't yet have a namespace string, then initialize it from either 'collection' or // '_cq'. if (_nss.isEmpty()) { @@ -174,6 +168,13 @@ PlanExecutorImpl::PlanExecutorImpl(OperationContext* opCtx, _planExplainer->updateEnumeratorExplainInfo( subplanStage->compositeSolution()->_enumeratorExplainInfo); } + + // If this PlanExecutor is executing a COLLSCAN, keep a pointer directly to the COLLSCAN + // stage. This is used for change streams in order to keep the the latest oplog timestamp + // and post batch resume token up to date as the oplog scan progresses. + if (auto collectionScan = getStageByType(_root.get(), STAGE_COLLSCAN)) { + _collScanStage = static_cast<CollectionScan*>(collectionScan); + } } Status PlanExecutorImpl::_pickBestPlan() { diff --git a/src/mongo/db/query/plan_executor_impl.h b/src/mongo/db/query/plan_executor_impl.h index 672cd75f243..ac593cd710c 100644 --- a/src/mongo/db/query/plan_executor_impl.h +++ b/src/mongo/db/query/plan_executor_impl.h @@ -29,6 +29,7 @@ #pragma once +#include "mongo/util/duration.h" #include <boost/optional.hpp> #include <queue> diff --git a/src/mongo/db/query/plan_executor_sbe.cpp b/src/mongo/db/query/plan_executor_sbe.cpp index c8d516cb718..e1909acf95c 100644 --- a/src/mongo/db/query/plan_executor_sbe.cpp +++ b/src/mongo/db/query/plan_executor_sbe.cpp @@ -41,6 +41,7 @@ #include "mongo/db/query/sbe_stage_builder.h" #include "mongo/logv2/log.h" #include "mongo/s/resharding/resume_token_gen.h" +#include "mongo/util/duration.h" namespace mongo { // This failpoint is defined by the classic executor but is also accessed here. diff --git a/src/mongo/db/query/plan_executor_sbe.h b/src/mongo/db/query/plan_executor_sbe.h index 547f0fb0a07..c11d97392ee 100644 --- a/src/mongo/db/query/plan_executor_sbe.h +++ b/src/mongo/db/query/plan_executor_sbe.h @@ -29,6 +29,7 @@ #pragma once +#include "mongo/util/duration.h" #include <queue> #include "mongo/db/exec/sbe/stages/stages.h" diff --git a/src/mongo/db/query/plan_explainer.h b/src/mongo/db/query/plan_explainer.h index b53a10e5655..ffaabecc43a 100644 --- a/src/mongo/db/query/plan_explainer.h +++ b/src/mongo/db/query/plan_explainer.h @@ -29,12 +29,14 @@ #pragma once +#include "mongo/bson/bsonobj.h" #include "mongo/db/exec/plan_stats.h" #include "mongo/db/query/classic_plan_cache.h" #include "mongo/db/query/explain_options.h" #include "mongo/db/query/plan_enumerator_explain_info.h" #include "mongo/db/query/plan_summary_stats.h" #include "mongo/db/query/query_solution.h" +#include "mongo/util/duration.h" namespace mongo { /** diff --git a/src/mongo/db/query/plan_explainer_factory.cpp b/src/mongo/db/query/plan_explainer_factory.cpp index 544ab33fdd2..9307ec349f6 100644 --- a/src/mongo/db/query/plan_explainer_factory.cpp +++ b/src/mongo/db/query/plan_explainer_factory.cpp @@ -34,6 +34,8 @@ #include "mongo/db/exec/plan_cache_util.h" #include "mongo/db/query/plan_explainer_impl.h" #include "mongo/db/query/plan_explainer_sbe.h" +#include "mongo/util/duration.h" +#include <ratio> namespace mongo::plan_explainer_factory { std::unique_ptr<PlanExplainer> make(PlanStage* root) { diff --git a/src/mongo/db/query/plan_explainer_factory.h b/src/mongo/db/query/plan_explainer_factory.h index 5e24a755747..0064df7cb39 100644 --- a/src/mongo/db/query/plan_explainer_factory.h +++ b/src/mongo/db/query/plan_explainer_factory.h @@ -36,6 +36,7 @@ #include "mongo/db/query/plan_explainer.h" #include "mongo/db/query/query_solution.h" #include "mongo/db/query/sbe_plan_ranker.h" +#include "mongo/util/duration.h" namespace mongo::plan_explainer_factory { std::unique_ptr<PlanExplainer> make(PlanStage* root); diff --git a/src/mongo/db/query/plan_explainer_impl.cpp b/src/mongo/db/query/plan_explainer_impl.cpp index c36bc330826..a67c9e98a60 100644 --- a/src/mongo/db/query/plan_explainer_impl.cpp +++ b/src/mongo/db/query/plan_explainer_impl.cpp @@ -655,7 +655,6 @@ boost::optional<double> getWinningPlanScore(PlanStage* root) { void PlanExplainerImpl::getSummaryStats(PlanSummaryStats* statsOut) const { invariant(statsOut); - // We can get some of the fields we need from the common stats stored in the // root stage of the plan tree. const CommonStats* common = _root->getCommonStats(); diff --git a/src/mongo/db/query/plan_explainer_impl.h b/src/mongo/db/query/plan_explainer_impl.h index 73ef81ae825..f73a2ec6250 100644 --- a/src/mongo/db/query/plan_explainer_impl.h +++ b/src/mongo/db/query/plan_explainer_impl.h @@ -29,10 +29,12 @@ #pragma once +#include "mongo/bson/bsonobj.h" #include "mongo/db/exec/plan_stage.h" #include "mongo/db/query/plan_enumerator_explain_info.h" #include "mongo/db/query/plan_explainer.h" #include "mongo/db/query/query_solution.h" +#include "mongo/util/duration.h" namespace mongo { /** @@ -47,7 +49,6 @@ public: PlanExplainerImpl(PlanStage* root, const PlanEnumeratorExplainInfo& explainInfo) : PlanExplainer{explainInfo}, _root{root} {} PlanExplainerImpl(PlanStage* root) : _root{root} {} - const ExplainVersion& getVersion() const final; bool isMultiPlan() const final; std::string getPlanSummary() const final; diff --git a/src/mongo/db/query/plan_explainer_sbe.h b/src/mongo/db/query/plan_explainer_sbe.h index 5dc97f90641..a256e1371c6 100644 --- a/src/mongo/db/query/plan_explainer_sbe.h +++ b/src/mongo/db/query/plan_explainer_sbe.h @@ -35,6 +35,7 @@ #include "mongo/db/query/plan_explainer.h" #include "mongo/db/query/query_solution.h" #include "mongo/db/query/sbe_plan_ranker.h" +#include "mongo/util/duration.h" namespace mongo { /** diff --git a/src/mongo/db/query/plan_summary_stats.h b/src/mongo/db/query/plan_summary_stats.h index ac80b6505f3..5c153ef3e39 100644 --- a/src/mongo/db/query/plan_summary_stats.h +++ b/src/mongo/db/query/plan_summary_stats.h @@ -29,6 +29,7 @@ #pragma once +#include "mongo/util/duration.h" #include <optional> #include <string> diff --git a/src/mongo/db/query/planner_access.cpp b/src/mongo/db/query/planner_access.cpp index 2e8e77c44b5..de294ad661a 100644 --- a/src/mongo/db/query/planner_access.cpp +++ b/src/mongo/db/query/planner_access.cpp @@ -45,14 +45,21 @@ #include "mongo/db/matcher/expression.h" #include "mongo/db/matcher/expression_array.h" #include "mongo/db/matcher/expression_geo.h" +#include "mongo/db/matcher/expression_internal_expr_comparison.h" +#include "mongo/db/matcher/expression_leaf.h" #include "mongo/db/matcher/expression_text.h" +#include "mongo/db/matcher/expression_tree.h" +#include "mongo/db/namespace_string.h" +#include "mongo/db/query/index_bounds.h" #include "mongo/db/query/index_bounds_builder.h" #include "mongo/db/query/index_tag.h" #include "mongo/db/query/indexability.h" +#include "mongo/db/query/parsed_find_command.h" #include "mongo/db/query/planner_wildcard_helpers.h" #include "mongo/db/query/query_knobs_gen.h" #include "mongo/db/query/query_planner.h" #include "mongo/db/query/query_planner_common.h" +#include "mongo/db/query/record_id_range.h" #include "mongo/db/record_id_helpers.h" #include "mongo/logv2/log.h" #include "mongo/util/transitional_tools_do_not_use/vector_spooling.h" @@ -269,7 +276,8 @@ bool compatibleCollator(const QueryPlannerParams& params, void handleRIDRangeMinMax(const CanonicalQuery& query, CollectionScanNode* collScan, const QueryPlannerParams& params, - const CollatorInterface* collator) { + const CollatorInterface* collator, + RecordIdRange& recordRange) { BSONObj minObj = query.getFindCommandRequest().getMin(); BSONObj maxObj = query.getFindCommandRequest().getMax(); if (minObj.isEmpty() && maxObj.isEmpty()) { @@ -289,16 +297,16 @@ void handleRIDRangeMinMax(const CanonicalQuery& query, if (!maxObj.isEmpty() && compatibleCollator(params, collator, maxObj.firstElement())) { // max() is exclusive. // Assumes clustered collection scans are only supported with the forward direction. - collScan->boundInclusion = - CollectionScanParams::ScanBoundInclusion::kIncludeStartRecordOnly; - setLowestRecord(collScan->maxRecord, - IndexBoundsBuilder::objFromElement(maxObj.firstElement(), collator)); + recordRange.maybeNarrowMax( + IndexBoundsBuilder::objFromElement(maxObj.firstElement(), collator), + false /* NOT inclusive*/); } if (!minObj.isEmpty() && compatibleCollator(params, collator, minObj.firstElement())) { // The min() is inclusive as are bounded collection scans by default. - setHighestRecord(collScan->minRecord, - IndexBoundsBuilder::objFromElement(minObj.firstElement(), collator)); + recordRange.maybeNarrowMin( + IndexBoundsBuilder::objFromElement(minObj.firstElement(), collator), + true /* inclusive*/); } } @@ -323,7 +331,8 @@ void handleRIDRangeMinMax(const CanonicalQuery& query, [[nodiscard]] bool handleRIDRangeScan(const MatchExpression* conjunct, CollectionScanNode* collScan, const QueryPlannerParams& params, - const CollatorInterface* collator) { + const CollatorInterface* collator, + RecordIdRange& recordRange) { invariant(params.clusteredInfo); if (conjunct == nullptr) { @@ -334,7 +343,8 @@ void handleRIDRangeMinMax(const CanonicalQuery& query, if (andMatchPtr != nullptr) { bool atLeastOneConjunctCompatibleCollation = false; for (size_t index = 0; index < andMatchPtr->numChildren(); index++) { - if (handleRIDRangeScan(andMatchPtr->getChild(index), collScan, params, collator)) { + if (handleRIDRangeScan( + andMatchPtr->getChild(index), collScan, params, collator, recordRange)) { atLeastOneConjunctCompatibleCollation = true; } } @@ -378,31 +388,35 @@ void handleRIDRangeMinMax(const CanonicalQuery& query, } } + // {min,max}RecordId will bound the range of ids scanned to the highest and lowest present + // in the InMatchExpression, but the filter is still required to filter to _exactly_ the + // requested matches. + // Finally, tighten the collscan bounds with the min/max bounds for the $in. - if (minBound) { - setHighestRecord(collScan->minRecord, *minBound); - } - if (maxBound) { - setLowestRecord(collScan->maxRecord, *maxBound); - } + recordRange.intersectRange(minBound, maxBound); return allEltsCollationCompatible; } - auto match = dynamic_cast<const ComparisonMatchExpression*>(conjunct); + auto match = dynamic_cast<const ComparisonMatchExpressionBase*>(conjunct); if (match == nullptr) { return false; // Not a comparison match expression. } const auto& element = match->getData(); - // Set coarse min/max bounds based on type in case we can't set tight bounds. - BSONObjBuilder minb; - minb.appendMinForType("", element.type()); - setHighestRecord(collScan->minRecord, minb.obj()); + if (!ComparisonMatchExpressionBase::isInternalExprComparison(match->matchType())) { + // Internal comparisons e.g., $_internalExprGt do _not_ carry type bracketing + // semantics (consistent with `$expr{$gt:[a,b]}`). + // For other comparisons which _do_ perform type bracketing, the RecordId bounds + // may be tightened here. + BSONObjBuilder minb; + minb.appendMinForType("", element.type()); + recordRange.maybeNarrowMin(minb.obj(), true /* inclusive */); - BSONObjBuilder maxb; - maxb.appendMaxForType("", element.type()); - setLowestRecord(collScan->maxRecord, maxb.obj()); + BSONObjBuilder maxb; + maxb.appendMaxForType("", element.type()); + recordRange.maybeNarrowMax(maxb.obj(), true /* inclusive */); + } bool compatible = compatibleCollator(params, collator, element); if (!compatible) { @@ -410,19 +424,33 @@ void handleRIDRangeMinMax(const CanonicalQuery& query, } // Even if the collations don't match at this point, it's fine, - // because the bounds exclude values that use it - const auto collated = IndexBoundsBuilder::objFromElement(element, collator); - if (dynamic_cast<const EqualityMatchExpression*>(match)) { - setHighestRecord(collScan->minRecord, collated); - setLowestRecord(collScan->maxRecord, collated); - } else if (dynamic_cast<const LTMatchExpression*>(match) || - dynamic_cast<const LTEMatchExpression*>(match)) { - setLowestRecord(collScan->maxRecord, collated); - } else if (dynamic_cast<const GTMatchExpression*>(match) || - dynamic_cast<const GTEMatchExpression*>(match)) { - setHighestRecord(collScan->minRecord, collated); + // because the bounds exclude values that use it. + const BSONObj collated = IndexBoundsBuilder::objFromElement(element, collator); + using MType = MatchExpression::MatchType; + switch (match->matchType()) { + case MType::EQ: + case MType::INTERNAL_EXPR_EQ: + recordRange.maybeNarrowMin(collated, true /* inclusive */); + recordRange.maybeNarrowMax(collated, true /* inclusive */); + break; + case MType::LT: + case MType::INTERNAL_EXPR_LT: + recordRange.maybeNarrowMax(collated, false /* EXclusive */); + break; + case MType::LTE: + case MType::INTERNAL_EXPR_LTE: + recordRange.maybeNarrowMax(collated, true /* inclusive */); + break; + case MType::GT: + case MType::INTERNAL_EXPR_GT: + recordRange.maybeNarrowMin(collated, false /* EXclusive */); + break; + case MType::GTE: + case MType::INTERNAL_EXPR_GTE: + recordRange.maybeNarrowMin(collated, true /* inclusive */); + break; + default:; } - return true; } @@ -528,11 +556,21 @@ std::unique_ptr<QuerySolutionNode> QueryPlannerAccess::makeCollectionScan( // query is guaranteed to exclude values of the cluster key which are affected by collation. // If so, then even if the query and collection collations differ, the collation difference // won't affect the query results. In that case, we can say hasCompatibleCollation is true. + + RecordIdRange recordRange; + // min/max records may have been set if oplog or change collection. + recordRange.intersectRange(csn->minRecord, csn->maxRecord); bool compatibleCollation = - handleRIDRangeScan(csn->filter.get(), csn.get(), params, queryCollator); + handleRIDRangeScan(csn->filter.get(), csn.get(), params, queryCollator, recordRange); csn->hasCompatibleCollation |= compatibleCollation; - handleRIDRangeMinMax(query, csn.get(), params, queryCollator); + handleRIDRangeMinMax(query, csn.get(), params, queryCollator, recordRange); + + csn->minRecord = recordRange.getMin(); + csn->maxRecord = recordRange.getMax(); + + csn->boundInclusion = CollectionScanParams::makeInclusion(recordRange.isMinInclusive(), + recordRange.isMaxInclusive()); } return csn; diff --git a/src/mongo/db/query/planner_access_test.cpp b/src/mongo/db/query/planner_access_test.cpp index 52df32c871e..d30544b43fa 100644 --- a/src/mongo/db/query/planner_access_test.cpp +++ b/src/mongo/db/query/planner_access_test.cpp @@ -38,9 +38,7 @@ namespace mongo { namespace { BSONObj serializeMatcher(Matcher* matcher) { - BSONObjBuilder builder; - matcher->getMatchExpression()->serialize(&builder); - return builder.obj(); + return matcher->getMatchExpression()->serialize(); } TEST(PlannerAccessTest, PrepareForAccessPlanningSortsEqualNodesByTheirChildren) { diff --git a/src/mongo/db/query/planner_ixselect.cpp b/src/mongo/db/query/planner_ixselect.cpp index c2e06a1027a..b63484717f2 100644 --- a/src/mongo/db/query/planner_ixselect.cpp +++ b/src/mongo/db/query/planner_ixselect.cpp @@ -256,10 +256,9 @@ void QueryPlannerIXSelect::getFields(const MatchExpression* node, if (Indexability::nodeCanUseIndexOnOwnField(node)) { bool supportSparse = Indexability::nodeSupportedBySparseIndex(node); (*out)[prefix + node->path().toString()] = {supportSparse}; - } else if (Indexability::arrayUsesIndexOnChildren(node) && !node->path().empty()) { + } else if (Indexability::isBoundsGeneratingElemMatchObject(node)) { // If the array uses an index on its children, it's something like // {foo : {$elemMatch: {bar: 1}}}, in which case the predicate is really over foo.bar. - // Note we skip empty path components since they are not allowed in index key patterns. prefix += node->path().toString() + "."; for (size_t i = 0; i < node->numChildren(); ++i) { @@ -440,10 +439,13 @@ bool QueryPlannerIXSelect::_compatible(const BSONElement& keyPatternElt, const auto* child = node->getChild(0); const MatchExpression::MatchType childtype = child->matchType(); - // Can't index negations of MOD, REGEX, TYPE_OPERATOR, or ELEM_MATCH_VALUE. + // Can't index negations of MOD, REGEX, TYPE_OPERATOR, or ELEM_MATCH_VALUE; and, as + // above, we can't use a btree-indexed field for geo expressions (or their negations). if (MatchExpression::REGEX == childtype || MatchExpression::MOD == childtype || MatchExpression::TYPE_OPERATOR == childtype || - MatchExpression::ELEM_MATCH_VALUE == childtype) { + MatchExpression::ELEM_MATCH_VALUE == childtype || + MatchExpression::GEO == childtype || MatchExpression::GEO_NEAR == childtype || + MatchExpression::INTERNAL_BUCKET_GEO_WITHIN == childtype) { return false; } diff --git a/src/mongo/db/query/planner_ixselect_test.cpp b/src/mongo/db/query/planner_ixselect_test.cpp index 1df4d714e67..2507e65ce56 100644 --- a/src/mongo/db/query/planner_ixselect_test.cpp +++ b/src/mongo/db/query/planner_ixselect_test.cpp @@ -1333,6 +1333,45 @@ TEST(QueryPlannerIXSelectTest, HashedSparseIndexShouldBeRelevantForExistsTrue) { testRateIndices("{a: {$exists: true}}", "", kSimpleCollator, {entry}, "a", expectedIndices); } +TEST(QueryPlannerIXSelectTest, GeoPredicateCanOnlyUse2dsphereIndex) { + std::vector<IndexEntry> indices; + auto btreeEntry = buildSimpleIndexEntry(BSON("loc" << 1)); + auto twodSphereEntry = buildSimpleIndexEntry(BSON("loc" + << "2dsphere")); + indices.push_back(btreeEntry); + indices.push_back(twodSphereEntry); + std::set<size_t> expectedIndices = {1}; + testRateIndices(R"({loc: {$geoWithin: {$geometry: {type: 'Polygon', + coordinates: [[[0,0],[0,1],[1,0],[0,0]]]}}}})", + "", + kSimpleCollator, + indices, + "loc", + expectedIndices); +} + +TEST(QueryPlannerIXSelectTest, GeoPredicateWithNotCanOnlyUse2dsphereIndex) { + std::vector<IndexEntry> indices; + auto btreeEntry = buildSimpleIndexEntry(BSON("loc" << 1)); + auto twodSphereEntry = buildSimpleIndexEntry(BSON("loc" + << "2dsphere")); + indices.push_back(btreeEntry); + indices.push_back(twodSphereEntry); + // This query gets parsed to {$not: {$and: {$geoWithin: <>}}} and then tags + // the 2dsphere index as relevant. If the $and is optimized away ({$not: {$geoWithin: <>}}), + // that tagging is skipped. + // TODO SERVER-92427: The tagging behavior should be made consistent so that this query has no + // expectedIndices. + std::set<size_t> expectedIndices = {1}; + testRateIndices(R"({loc: {$not: {$geoWithin: {$geometry: {type: 'Polygon', + coordinates: [[[0,0],[0,1],[1,0],[0,0]]]}}}}})", + "", + kSimpleCollator, + indices, + "loc", + expectedIndices); +} + /* * Will compare 'keyPatterns' with 'entries'. As part of comparing, it will sort both of them. */ diff --git a/src/mongo/db/query/projection.cpp b/src/mongo/db/query/projection.cpp index b55fc03602a..2646b39bd8c 100644 --- a/src/mongo/db/query/projection.cpp +++ b/src/mongo/db/query/projection.cpp @@ -227,6 +227,13 @@ void optimizeProjection(ProjectionPathASTNode* root) { Projection::Projection(ProjectionPathASTNode root, ProjectType type) : _root(std::move(root)), _type(type), _deps(analyzeProjection(&_root, type)) {} +void Projection::optimize() { + if (!_projOptimized) { + optimizeProjection(&_root); + _deps = analyzeProjection(&_root, _type); + _projOptimized = true; + } +} namespace { /** diff --git a/src/mongo/db/query/projection.h b/src/mongo/db/query/projection.h index 97b8d1e0d30..95d41caba02 100644 --- a/src/mongo/db/query/projection.h +++ b/src/mongo/db/query/projection.h @@ -142,14 +142,25 @@ public: return _deps.containsElemMatch; } + /** + * Optimizes the projection tree. Additionally, re-computes dependencies in case anything + * changes as in projection {x: {$and: [false, "$b"]}} - which when optimized will no longer + * depend on "b". + */ + void optimize(); + private: ProjectionPathASTNode _root; ProjectType _type; ProjectionDependencies _deps; + bool _projOptimized = false; }; /** - * Walks the projection AST and optimizes each node. + * Walks the projection AST and optimizes each node. Note if you have a 'Projection' instance you + * should prefer to use Projection::optimize() since it will additionally re-compute dependencies in + * case anything changes as in projection {x: {$and: [false, "$b"]}} - which when optimized will no + * longer depend on "b". */ void optimizeProjection(ProjectionPathASTNode* root); diff --git a/src/mongo/db/query/projection_ast.h b/src/mongo/db/query/projection_ast.h index 247a91537f3..1304c68efe5 100644 --- a/src/mongo/db/query/projection_ast.h +++ b/src/mongo/db/query/projection_ast.h @@ -275,7 +275,7 @@ public: ExpressionASTNode(boost::intrusive_ptr<Expression> expr) : _expr(expr) {} ExpressionASTNode(const ExpressionASTNode& other) : ASTNode(other) { BSONObjBuilder bob; - bob << "" << other._expr->serialize(false); + bob << "" << other._expr->serialize(); // TODO SERVER-31003: add a clone() method to Expression. // Temporary stop expression counters while processing the cloned expression. diff --git a/src/mongo/db/query/projection_ast_test.cpp b/src/mongo/db/query/projection_ast_test.cpp index 619bd9a1dde..ea1eb2b683d 100644 --- a/src/mongo/db/query/projection_ast_test.cpp +++ b/src/mongo/db/query/projection_ast_test.cpp @@ -41,6 +41,7 @@ #include "mongo/db/query/projection_ast_util.h" #include "mongo/db/query/projection_parser.h" #include "mongo/db/query/query_planner_test_fixture.h" +#include "mongo/db/query/query_shape/serialization_options.h" namespace { @@ -772,4 +773,62 @@ TEST_F(ProjectionASTTest, ShouldThrowWithPositionalOnExclusion) { DBException, 31395); } + +TEST_F(ProjectionASTTest, TestASTRedaction) { + SerializationOptions options = SerializationOptions::kDebugShapeAndMarkIdentifiers_FOR_TEST; + + auto proj = fromjson("{'a.b': 1}"); + BSONObj output = projection_ast::serialize(*parseWithFindFeaturesEnabled(proj).root(), options); + ASSERT_BSONOBJ_EQ_AUTO( // + R"({"HASH<a>":{"HASH<b>":true},"HASH<_id>":true})", + output); + + proj = fromjson("{'a.b': 0}"); + output = projection_ast::serialize(*parseWithFindFeaturesEnabled(proj).root(), options); + ASSERT_BSONOBJ_EQ_AUTO( // + R"({"HASH<a>":{"HASH<b>":false}})", + output); + + proj = fromjson("{a: 1, b: 1}"); + output = projection_ast::serialize(*parseWithFindFeaturesEnabled(proj).root(), options); + ASSERT_BSONOBJ_EQ_AUTO( // + R"({"HASH<a>":true,"HASH<b>":true,"HASH<_id>":true})", + output); + + // ElemMatch projection + proj = fromjson("{f: {$elemMatch: {foo: 'bar'}}}"); + output = projection_ast::serialize(*parseWithFindFeaturesEnabled(proj).root(), options); + ASSERT_BSONOBJ_EQ_AUTO( // + R"({"HASH<f>":{"$elemMatch":{"HASH<foo>":{"$eq":"?string"}}},"HASH<_id>":true})", + output); + + // Positional projection + proj = fromjson("{'x.$': 1}"); + output = projection_ast::serialize( + *parseWithFindFeaturesEnabled(proj, fromjson("{'x.a': 2}")).root(), {}); + ASSERT_BSONOBJ_EQ_AUTO( // + R"({"x.$":true,"_id":true})", + output); + + // Slice (first form) + proj = fromjson("{a: {$slice: 1}}"); + output = projection_ast::serialize(*parseWithFindFeaturesEnabled(proj).root(), options); + ASSERT_BSONOBJ_EQ_AUTO( // + R"({"HASH<a>":{"$slice":"?number"}})", + output); + + // Slice (second form) + proj = fromjson("{a: {$slice: [1, 3]}}"); + output = projection_ast::serialize(*parseWithFindFeaturesEnabled(proj).root(), options); + ASSERT_BSONOBJ_EQ_AUTO( // + R"({"HASH<a>":{"$slice":["?number","?number"]}})", + output); + + /// $meta projection + proj = fromjson("{foo: {$meta: 'indexKey'}}"); + output = projection_ast::serialize(*parseWithFindFeaturesEnabled(proj).root(), options); + ASSERT_BSONOBJ_EQ_AUTO( // + R"({"HASH<foo>":{"$meta":"indexKey"}})", + output); +} } // namespace diff --git a/src/mongo/db/query/projection_ast_util.cpp b/src/mongo/db/query/projection_ast_util.cpp index e5b4cc1a9c4..23c7b6d9582 100644 --- a/src/mongo/db/query/projection_ast_util.cpp +++ b/src/mongo/db/query/projection_ast_util.cpp @@ -29,28 +29,26 @@ #include "mongo/platform/basic.h" -#include "mongo/db/query/projection_ast_util.h" - #include "mongo/db/query/projection_ast_path_tracking_visitor.h" +#include "mongo/db/query/projection_ast_util.h" +#include "mongo/db/query/query_shape/serialization_options.h" #include "mongo/db/query/tree_walker.h" namespace mongo::projection_ast { namespace { struct BSONVisitorContext { std::stack<BSONObjBuilder> builders; + bool underElemMatch = false; }; class BSONPreVisitor : public ProjectionASTConstVisitor { public: - BSONPreVisitor(PathTrackingVisitorContext<BSONVisitorContext>* context) - : _context(context), _builders(context->data().builders) {} - - virtual void visit(const MatchExpressionASTNode* node) { - static_cast<const MatchExpressionASTNode*>(node)->matchExpression()->serialize( - &_builders.top(), true); - } + using ProjectionASTConstVisitor::visit; + BSONPreVisitor(PathTrackingVisitorContext<BSONVisitorContext>* context, + SerializationOptions options) + : _context(context), _builders(context->data().builders), _options(std::move(options)) {} - virtual void visit(const ProjectionPathASTNode* node) { + void visit(const ProjectionPathASTNode* node) override { if (!node->parent()) { // No root of the tree, thus this node has no field name. _builders.push(BSONObjBuilder()); @@ -59,47 +57,46 @@ public: } } - virtual void visit(const ProjectionPositionalASTNode* node) { - // ProjectionPositional always has the original query's match expression node as its - // child. Serialize as: {"positional.projection.field.$": <original match expression>}. - _context->data().builders.push(_builders.top().subobjStart(getFieldName() + ".$")); - } - - virtual void visit(const ProjectionSliceASTNode* node) { + void visit(const ProjectionSliceASTNode* node) override { BSONObjBuilder sub(_builders.top().subobjStart(getFieldName())); if (node->skip()) { - sub.appendArray("$slice", BSON_ARRAY(*node->skip() << node->limit())); + sub.appendArray("$slice", + BSON_ARRAY(_options.serializeLiteral(*node->skip()) + << _options.serializeLiteral(node->limit()))); } else { - sub.appendNumber("$slice", node->limit()); + _options.appendLiteral(&sub, "$slice", node->limit()); } } - virtual void visit(const ProjectionElemMatchASTNode* node) { - // Defer to the child, match expression node. - } - virtual void visit(const ExpressionASTNode* node) { - node->expression()->serialize(false).addToBsonObj(&_builders.top(), getFieldName()); + void visit(const ExpressionASTNode* node) override { + node->expression()->serialize(_options).addToBsonObj(&_builders.top(), getFieldName()); } - virtual void visit(const BooleanConstantASTNode* node) { + void visit(const BooleanConstantASTNode* node) override { _builders.top().append(getFieldName(), node->value()); } -private: + void visit(const ProjectionPositionalASTNode* node) override = 0; + void visit(const ProjectionElemMatchASTNode* node) override = 0; + void visit(const MatchExpressionASTNode* node) override = 0; + +protected: std::string getFieldName() { - return _context->childPath(); + return _options.serializeFieldPathFromString(_context->childPath()); } PathTrackingVisitorContext<BSONVisitorContext>* _context; std::stack<BSONObjBuilder>& _builders; + SerializationOptions _options; }; class BSONPostVisitor : public ProjectionASTConstVisitor { public: + using ProjectionASTConstVisitor::visit; BSONPostVisitor(BSONVisitorContext* context) : _context(context) {} - virtual void visit(const ProjectionPathASTNode* node) { + void visit(const ProjectionPathASTNode* node) override { // Don't pop the top builder. if (node->parent()) { // Pop the BSONObjBuilder that was added in the pre visitor. @@ -107,25 +104,97 @@ public: } } - virtual void visit(const ProjectionPositionalASTNode* node) { + void visit(const ProjectionSliceASTNode* node) override {} + void visit(const ExpressionASTNode* node) override {} + void visit(const BooleanConstantASTNode* node) override {} + void visit(const MatchExpressionASTNode* node) override {} + + void visit(const ProjectionPositionalASTNode* node) override = 0; + void visit(const ProjectionElemMatchASTNode* node) override = 0; + +protected: + BSONVisitorContext* _context; +}; + +class DebugPreVisitor : public BSONPreVisitor { +public: + using BSONPreVisitor::visit; + DebugPreVisitor(PathTrackingVisitorContext<BSONVisitorContext>* context) + : BSONPreVisitor(context, SerializationOptions{}) {} + + void visit(const ProjectionPositionalASTNode* node) override { + // ProjectionPositional always has the original query's match expression node as its + // child. Serialize as: {"positional.projection.field.$": <original match expression>}. + _context->data().builders.push(_builders.top().subobjStart(getFieldName() + ".$")); + } + + void visit(const ProjectionElemMatchASTNode* node) override { + // Defer to the child, match expression node. + } + + void visit(const MatchExpressionASTNode* node) override { + static_cast<const MatchExpressionASTNode*>(node)->matchExpression()->serialize( + &_builders.top(), {}); + } +}; + +class DebugPostVisitor : public BSONPostVisitor { +public: + using BSONPostVisitor::visit; + DebugPostVisitor(BSONVisitorContext* context) : BSONPostVisitor(context) {} + + void visit(const ProjectionPositionalASTNode* node) override { _context->builders.pop(); } - virtual void visit(const MatchExpressionASTNode* node) {} - virtual void visit(const ProjectionSliceASTNode* node) {} - virtual void visit(const ProjectionElemMatchASTNode* node) {} - virtual void visit(const ExpressionASTNode* node) {} - virtual void visit(const BooleanConstantASTNode* node) {} + void visit(const ProjectionElemMatchASTNode* node) override {} +}; -private: - BSONVisitorContext* _context; +class SerializationPreVisitor : public BSONPreVisitor { +public: + using BSONPreVisitor::visit; + SerializationPreVisitor(PathTrackingVisitorContext<BSONVisitorContext>* context, + const SerializationOptions& options) + : BSONPreVisitor(context, options) {} + + void visit(const ProjectionPositionalASTNode* node) override { + tassert(73488, + "Positional projection should not appear below an $elemMatch projection.", + !_context->data().underElemMatch); + _builders.top().append(getFieldName() + ".$", true); + } + + void visit(const ProjectionElemMatchASTNode* node) override { + // The child match expression node should begin with $elemMatch. + _context->data().underElemMatch = true; + } + + void visit(const MatchExpressionASTNode* node) override { + if (_context->data().underElemMatch) { + static_cast<const MatchExpressionASTNode*>(node)->matchExpression()->serialize( + &_builders.top(), _options); + } + } }; + +class SerializationPostVisitor : public BSONPostVisitor { +public: + using BSONPostVisitor::visit; + SerializationPostVisitor(BSONVisitorContext* context) : BSONPostVisitor(context) {} + + void visit(const ProjectionPositionalASTNode* node) override {} + void visit(const ProjectionElemMatchASTNode* node) override { + _context->underElemMatch = false; + } +}; + } // namespace BSONObj astToDebugBSON(const ASTNode* root) { PathTrackingVisitorContext<BSONVisitorContext> context; - BSONPreVisitor preVisitor{&context}; - BSONPostVisitor postVisitor{&context.data()}; + DebugPreVisitor preVisitor{&context}; + DebugPostVisitor postVisitor{&context.data()}; + PathTrackingWalker walker{&context, {&preVisitor}, {&postVisitor}}; tree_walker::walk<true, projection_ast::ASTNode>(root, &walker); @@ -133,4 +202,15 @@ BSONObj astToDebugBSON(const ASTNode* root) { invariant(context.data().builders.size() == 1); return context.data().builders.top().obj(); } + +BSONObj serialize(const ProjectionPathASTNode& root, const SerializationOptions& options) { + PathTrackingVisitorContext<BSONVisitorContext> context; + SerializationPreVisitor preVisitor{&context, options}; + SerializationPostVisitor postVisitor{&context.data()}; + PathTrackingWalker walker{&context, {&preVisitor}, {&postVisitor}}; + tree_walker::walk<true, projection_ast::ASTNode>(&root, &walker); + + invariant(context.data().builders.size() == 1); + return context.data().builders.top().obj(); +} } // namespace mongo::projection_ast diff --git a/src/mongo/db/query/projection_ast_util.h b/src/mongo/db/query/projection_ast_util.h index af89254a9a1..de102dfbbd7 100644 --- a/src/mongo/db/query/projection_ast_util.h +++ b/src/mongo/db/query/projection_ast_util.h @@ -29,6 +29,7 @@ #pragma once +#include "mongo/db/query/projection.h" #include "mongo/db/query/projection_ast.h" namespace mongo { @@ -37,5 +38,7 @@ namespace projection_ast { * This is intended to be used for debug output, not for serialization. */ BSONObj astToDebugBSON(const ASTNode* root); + +BSONObj serialize(const ProjectionPathASTNode& root, const SerializationOptions& options); } // namespace projection_ast } // namespace mongo diff --git a/src/mongo/db/query/projection_parser.cpp b/src/mongo/db/query/projection_parser.cpp index 85ada560ccc..3495f99f8f2 100644 --- a/src/mongo/db/query/projection_parser.cpp +++ b/src/mongo/db/query/projection_parser.cpp @@ -510,7 +510,16 @@ void parseSubObject(ParseContext* ctx, // It was likely intended to be an expression. Check if it's a valid field path or not to // confirm. try { - FieldPath fp(obj.firstElementFieldNameStringData()); + const auto elementFieldName = obj.firstElementFieldNameStringData(); + if (!hasPositionalOperator(elementFieldName)) { + FieldPath fp(elementFieldName); + } else { + // The 'FieldPath' parser doesn't take positional operators into account, but those + // are valid path projections so trim it off for this validation. + StringData pathWithoutPositionalOperator = + elementFieldName.substr(0, elementFieldName.size() - 2); + FieldPath fp(pathWithoutPositionalOperator); + } } catch (const DBException&) { uasserted(31325, str::stream() diff --git a/src/mongo/db/query/query_feature_flags.idl b/src/mongo/db/query/query_feature_flags.idl index f97cec95cc2..7e5995a022f 100644 --- a/src/mongo/db/query/query_feature_flags.idl +++ b/src/mongo/db/query/query_feature_flags.idl @@ -152,6 +152,12 @@ feature_flags: cpp_varname: gFeatureFlagSbeFull default: false + featureFlagQueryStats: + description: "Feature flag for enabling full queryStats collection." + cpp_varname: gFeatureFlagQueryStats + default: true + version: 6.0 + featureFlagShardedSearchCustomSort: description: "Feature flag to enable user specified sort for sharded $search queries." cpp_varname: gFeatureFlagShardedSearchCustomSort diff --git a/src/mongo/db/query/query_knobs.idl b/src/mongo/db/query/query_knobs.idl index 0748cee3a33..f0f24d3c3dd 100644 --- a/src/mongo/db/query/query_knobs.idl +++ b/src/mongo/db/query/query_knobs.idl @@ -29,8 +29,9 @@ global: cpp_namespace: "mongo" cpp_includes: - - "mongo/db/query/plan_cache_size_parameter.h" - "mongo/db/query/sbe_plan_cache_on_parameter_change.h" + - "mongo/db/query/util/memory_util.h" + - "mongo/db/query/query_stats/query_stats_on_parameter_change.h" - "mongo/platform/atomic_proxy.h" - "mongo/platform/atomic_word.h" @@ -675,7 +676,7 @@ server_parameters: set_at: [ startup, runtime ] cpp_varname: "internalQueryForceClassicEngine" cpp_vartype: AtomicWord<bool> - default: false + default: true internalQueryAppendIdToSetWindowFieldsSort: description: "If true, appends _id to the sort stage generated by desugaring $setWindowFields to @@ -906,6 +907,69 @@ server_parameters: expr: 8 * 1024 * 1024 # 8MB default: 0 + internalQueryStatsRateLimit: + description: "The maximum number of queries per second that are sampled for query stats. + If the rate of queries goes above this number, then rate limiting will kick in, and any + further queries will not be sampled. To sample all queries, this can be set to -1. This can be + set to 0 to turn queryStats off completely." + set_at: [ startup, runtime ] + cpp_varname: "internalQueryStatsRateLimit" + cpp_vartype: AtomicWord<int> + default: 0 + validator: + gte: -1 + on_update: query_stats_util::onQueryStatsSamplingRateUpdate + + internalQueryStatsCacheSize: + description: "The maximum amount of memory that the system will allocate for the query queryStats + cache. This will accept values in either of the following formats: + 1. <number>% indicates a percentage of the physical memory available to the process. E.g.: 15%. + 2. <number>(MB|GB), indicates the amount of memory in MB or GB. E.g.: 1.5GB, 100MB. + The default value is 1%, which means 1% of the physical memory available to the process." + set_at: [ startup, runtime ] + cpp_varname: "internalQueryStatsCacheSize" + cpp_vartype: synchronized_value<std::string> + default: "1%" + on_update: query_stats_util::onQueryStatsStoreSizeUpdate + validator: + callback: query_stats_util::validateQueryStatsStoreSize + + internalQueryStatsErrorsAreCommandFatal: + description: "Whether errors in the $queryStats stage cause the aggregation pipeline to + immediately fail and report the error. Note that this is always the case for debug builds." + set_at: [ startup, runtime ] + cpp_varname: "internalQueryStatsErrorsAreCommandFatal" + cpp_vartype: AtomicWord<bool> + default: false + + internalQueryAggMulticastTimeoutMS: + description: "Timeout in MS for requests to shard servers when aggregations are sent to all shard servers" + set_at: [ startup ] + cpp_vartype: int + cpp_varname: internalQueryAggMulticastTimeoutMS + default: 60000 + validator: + gte: 0 + + internalQueryAggMulticastMaxConcurrency: + description: "Max number of concurrent requests when aggregations are sent to all shard servers" + set_at: startup + cpp_vartype: int + cpp_varname: internalQueryAggMulticastMaxConcurrency + default: 100 + validator: + gte: 1 + + internalQuerySpillingMaxWaitTimeout: + description: "Timeout in MS that the storage engine will block a spilling operation when the + cache is under pressure." + set_at: [ startup, runtime ] + cpp_vartype: AtomicWord<int> + cpp_varname: internalQuerySpillingMaxWaitTimeout + default: 1000 + validator: + gte: 0 + # Note for adding additional query knobs: # # When adding a new query knob, you should consider whether or not you need to add an 'on_update' diff --git a/src/mongo/db/query/query_planner_common.h b/src/mongo/db/query/query_planner_common.h index 6d441155b54..97e94bb0a69 100644 --- a/src/mongo/db/query/query_planner_common.h +++ b/src/mongo/db/query/query_planner_common.h @@ -65,6 +65,20 @@ public: } /** + * Returns a count of 'type' nodes in expression tree. + */ + static size_t countNodes(const MatchExpression* root, MatchExpression::MatchType type) { + size_t sum = 0; + if (type == root->matchType()) { + sum = 1; + } + for (size_t i = 0; i < root->numChildren(); ++i) { + sum += countNodes(root->getChild(i), type); + } + return sum; + } + + /** * Assumes the provided BSONObj is of the form {field1: -+1, ..., field2: -+1} * Returns a BSONObj with the values negated. */ diff --git a/src/mongo/db/query/query_planner_test_lib.cpp b/src/mongo/db/query/query_planner_test_lib.cpp index bec219226b5..849ce6fb6e8 100644 --- a/src/mongo/db/query/query_planner_test_lib.cpp +++ b/src/mongo/db/query/query_planner_test_lib.cpp @@ -1234,7 +1234,10 @@ Status QueryPlannerTestLib::solutionMatches(const BSONObj& testSoln, } BSONObjBuilder bob; - actualGroupNode->groupByExpression->serialize(true).addToBsonObj(&bob, "_id"); + actualGroupNode->groupByExpression + ->serialize(SerializationOptions{ + boost::make_optional(ExplainOptions::Verbosity::kQueryPlanner)}) + .addToBsonObj(&bob, "_id"); auto actualGroupByObj = bob.done(); if (!SimpleBSONObjComparator::kInstance.evaluate(actualGroupByObj == expectedGroupByElem.Obj())) { @@ -1247,7 +1250,10 @@ Status QueryPlannerTestLib::solutionMatches(const BSONObj& testSoln, BSONArrayBuilder actualAccs; for (auto& acc : actualGroupNode->accumulators) { BSONObjBuilder bob; - acc.expr.argument->serialize(true).addToBsonObj(&bob, acc.expr.name); + acc.expr.argument + ->serialize(SerializationOptions{ + boost::make_optional(ExplainOptions::Verbosity::kQueryPlanner)}) + .addToBsonObj(&bob, acc.expr.name); actualAccs.append(BSON(acc.fieldName << bob.done())); } auto expectedAccsObj = expectedGroupObj["accs"].Obj(); diff --git a/src/mongo/db/query/query_planner_tree_test.cpp b/src/mongo/db/query/query_planner_tree_test.cpp index 5b7055dd720..7ea6d55dfb3 100644 --- a/src/mongo/db/query/query_planner_tree_test.cpp +++ b/src/mongo/db/query/query_planner_tree_test.cpp @@ -434,6 +434,39 @@ TEST_F(QueryPlannerTest, RootedOrOfAndDontCollapseDifferentBounds) { "bounds: {c: [[3,3,true,true]], d: [[4,4,true,true]]}}}]}}}}"); } +TEST_F(QueryPlannerTest, DontCrashTryingToPushToSingleChildIndexedOr1) { + FailPointEnableBlock failPoint("disableMatchExpressionOptimization"); + addIndex(BSON("indexed" << 1)); + runQuery( + fromjson("{ $and : [\n" + " { $and : [ { indexed : { $gt : 5 } },\n" + " { unindexed : 42 } ] },\n" + " { $or : [ { indexed: { $lt : 100 } } ] }\n" + " ] }")); + + assertNumSolutions(3U); +} + +TEST_F(QueryPlannerTest, DontCrashTryingToPushToSingleChildIndexedOr2) { + // Test that queries with single-child $and, $or do not crash when match-expression optimization + // is disabled. Normally these single-child nodes are eliminated, so when they are left in place + // it can confuse OR-pushdown optimization. + // + // Originally designed to reproduce SERVER-70597, which would only happen when the + // INDEX_INTERSECTION option is enabled. + FailPointEnableBlock failPoint("disableMatchExpressionOptimization"); + addIndex(BSON("a" << 1 << "b" << 1)); + + params.options |= QueryPlannerParams::INDEX_INTERSECTION; + runQuery( + fromjson("{ $and : [\n" + " { $and : [ { a : 2 } ] },\n" + " { $or : [ { b : 3 } ] }\n" + " ] }")); + + assertNumSolutions(2U); +} + // SERVER-13960: properly handle $or with a mix of exact and inexact predicates. TEST_F(QueryPlannerTest, OrInexactWithExact) { addIndex(BSON("name" << 1)); diff --git a/src/mongo/db/query/query_request_test.cpp b/src/mongo/db/query/query_request_test.cpp index 1493c352b28..52c6dd4aca7 100644 --- a/src/mongo/db/query/query_request_test.cpp +++ b/src/mongo/db/query/query_request_test.cpp @@ -36,6 +36,7 @@ #include "mongo/base/error_codes.h" #include "mongo/db/catalog/collection_catalog.h" #include "mongo/db/catalog/collection_mock.h" +#include "mongo/db/cursor_id.h" #include "mongo/db/dbmessage.h" #include "mongo/db/json.h" #include "mongo/db/namespace_string.h" diff --git a/src/mongo/db/query/query_shape.cpp b/src/mongo/db/query/query_shape.cpp new file mode 100644 index 00000000000..02d4c97d25f --- /dev/null +++ b/src/mongo/db/query/query_shape.cpp @@ -0,0 +1,292 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_shape.h" + +#include "mongo/base/status.h" +#include "mongo/db/query/find_command_gen.h" +#include "mongo/db/query/projection_ast_util.h" +#include "mongo/db/query/query_request_helper.h" +#include "mongo/db/query/query_shape_gen.h" +#include "mongo/db/query/query_stats/query_stats.h" +#include "mongo/db/query/shape_helpers.h" +#include "mongo/db/query/sort_pattern.h" + +namespace mongo::query_shape { + +BSONObj debugPredicateShape(const MatchExpression* predicate) { + SerializationOptions opts; + opts.literalPolicy = LiteralSerializationPolicy::kToDebugTypeString; + return predicate->serialize(opts); +} +BSONObj representativePredicateShape(const MatchExpression* predicate) { + SerializationOptions opts; + opts.literalPolicy = LiteralSerializationPolicy::kToRepresentativeParseableValue; + return predicate->serialize(opts); +} + +BSONObj debugPredicateShape(const MatchExpression* predicate, + std::function<std::string(StringData)> transformIdentifiersCallback) { + SerializationOptions opts; + opts.literalPolicy = LiteralSerializationPolicy::kToDebugTypeString; + opts.transformIdentifiersCallback = transformIdentifiersCallback; + opts.transformIdentifiers = true; + return predicate->serialize(opts); +} + +BSONObj representativePredicateShape( + const MatchExpression* predicate, + std::function<std::string(StringData)> transformIdentifiersCallback) { + SerializationOptions opts; + opts.literalPolicy = LiteralSerializationPolicy::kToRepresentativeParseableValue; + opts.transformIdentifiersCallback = transformIdentifiersCallback; + opts.transformIdentifiers = true; + return predicate->serialize(opts); +} + +BSONObj extractSortShape(const BSONObj& sortSpec, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const SerializationOptions& opts) { + if (sortSpec.isEmpty()) { + return sortSpec; + } + auto natural = sortSpec[query_request_helper::kNaturalSortField]; + + if (!natural) { + return SortPattern{sortSpec, expCtx} + .serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts) + .toBson(); + } + // This '$natural' will fail to parse as a valid SortPattern since it is not a valid field + // path - it is usually considered and converted into a hint. For the query shape, we'll + // keep it unmodified. + BSONObjBuilder bob; + for (auto&& elem : sortSpec) { + if (elem.isABSONObj()) { + // We expect this won't work or parse on the main command path, but for shapification we + // don't really care, just treat it as a literal and don't bother parsing. + opts.appendLiteral( + &bob, opts.serializeFieldPathFromString(elem.fieldNameStringData()), elem); + } else if (elem.fieldNameStringData() == natural.fieldNameStringData()) { + bob.append(elem); + } else { + bob.appendAs(elem, opts.serializeFieldPathFromString(elem.fieldNameStringData())); + } + } + return bob.obj(); +} + +void addShapeLiterals(BSONObjBuilder* bob, + const FindCommandRequest& findCommand, + const SerializationOptions& opts) { + if (auto limit = findCommand.getLimit()) { + opts.appendLiteral( + bob, FindCommandRequest::kLimitFieldName, static_cast<long long>(*limit)); + } + if (auto skip = findCommand.getSkip()) { + opts.appendLiteral(bob, FindCommandRequest::kSkipFieldName, static_cast<long long>(*skip)); + } +} + +static std::vector< + std::pair<StringData, std::function<const OptionalBool(const FindCommandRequest&)>>> + boolArgMap = { + {FindCommandRequest::kSingleBatchFieldName, &FindCommandRequest::getSingleBatch}, + {FindCommandRequest::kAllowDiskUseFieldName, &FindCommandRequest::getAllowDiskUse}, + {FindCommandRequest::kReturnKeyFieldName, &FindCommandRequest::getReturnKey}, + {FindCommandRequest::kShowRecordIdFieldName, &FindCommandRequest::getShowRecordId}, + {FindCommandRequest::kTailableFieldName, &FindCommandRequest::getTailable}, + {FindCommandRequest::kAwaitDataFieldName, &FindCommandRequest::getAwaitData}, +}; +std::vector<std::pair<StringData, std::function<const BSONObj(const FindCommandRequest&)>>> + objArgMap = { + {FindCommandRequest::kCollationFieldName, &FindCommandRequest::getCollation}, + +}; + +void addRemainingFindCommandFields(BSONObjBuilder* bob, + const FindCommandRequest& findCommand, + const SerializationOptions& opts) { + for (auto [fieldName, getterFunction] : boolArgMap) { + auto optBool = getterFunction(findCommand); + optBool.serializeToBSON(fieldName, bob); + } + + if (auto optOplogReplay = findCommand.getOplogReplay()) { + if (optOplogReplay.has_value()) { + opts.appendLiteral( + bob, FindCommandRequest::kOplogReplayFieldName, optOplogReplay.value_or(false)); + } + } + + auto collation = findCommand.getCollation(); + if (!collation.isEmpty()) { + bob->append(FindCommandRequest::kCollationFieldName, collation); + } +} + + +/** + * In a let specification all field names are variable names, and all values are either + * expressions or constants. + */ +BSONObj extractLetSpecShape(BSONObj letSpec, + const SerializationOptions& opts, + const boost::intrusive_ptr<ExpressionContext>& expCtx) { + + BSONObjBuilder bob; + for (BSONElement elem : letSpec) { + auto expr = Expression::parseOperand(expCtx.get(), elem, expCtx->variablesParseState); + auto redactedValue = expr->serialize(opts); + // Note that this will throw on deeply nested let variables. + redactedValue.addToBsonObj(&bob, opts.serializeFieldPathFromString(elem.fieldName())); + } + return bob.obj(); +} + +void appendCmdNs(BSONObjBuilder& bob, + const NamespaceString& nss, + const SerializationOptions& opts) { + BSONObjBuilder nsObj = bob.subobjStart("cmdNs"); + shape_helpers::appendNamespaceShape(nsObj, nss, opts); + nsObj.doneFast(); +} + +BSONObj extractQueryShape(const ParsedFindCommand& findRequest, + const SerializationOptions& opts, + const boost::intrusive_ptr<ExpressionContext>& expCtx) { + const auto& findCmd = *findRequest.findCommandRequest; + BSONObjBuilder bob; + // Serialize the namespace as part of the query shape. + { + auto ns = findCmd.getNamespaceOrUUID(); + if (ns.nss().has_value()) { + appendCmdNs(bob, *ns.nss(), opts); + } else { + BSONObjBuilder cmdNs = bob.subobjStart("cmdNs"); + cmdNs.append("uuid", opts.serializeIdentifier(ns.uuid()->toString())); + cmdNs.append("db", opts.serializeIdentifier(ns.db())); + cmdNs.doneFast(); + } + } + + bob.append("command", "find"); + std::unique_ptr<MatchExpression> filterExpr; + // Filter. + bob.append(FindCommandRequest::kFilterFieldName, findRequest.filter->serialize(opts)); + // Let Spec. + if (auto letSpec = findCmd.getLet()) { + auto redactedObj = extractLetSpecShape(letSpec.get(), opts, expCtx); + auto ownedObj = redactedObj.getOwned(); + bob.append(FindCommandRequest::kLetFieldName, std::move(ownedObj)); + } + + if (findRequest.proj) { + bob.append(FindCommandRequest::kProjectionFieldName, + projection_ast::serialize(*findRequest.proj->root(), opts)); + } + + if (!findCmd.getMax().isEmpty()) { + bob.append(FindCommandRequest::kMaxFieldName, + shape_helpers::extractMinOrMaxShape(findCmd.getMax(), opts)); + } + if (!findCmd.getMin().isEmpty()) { + bob.append(FindCommandRequest::kMinFieldName, + shape_helpers::extractMinOrMaxShape(findCmd.getMin(), opts)); + } + + // Sort. + if (findRequest.sort) { + bob.append( + FindCommandRequest::kSortFieldName, + findRequest.sort + ->serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts) + .toBson()); + } + + // Fields for literal redaction. Adds limit and skip. + addShapeLiterals(&bob, findCmd, opts); + + // Add the fields that require no redaction. + addRemainingFindCommandFields(&bob, findCmd, opts); + + return bob.obj(); +} + +BSONObj extractQueryShape(const AggregateCommandRequest& aggregateCommand, + const Pipeline& pipeline, + const SerializationOptions& opts, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const NamespaceString& nss) { + BSONObjBuilder bob; + + // namespace + appendCmdNs(bob, nss, opts); + bob.append("command", "aggregate"); + + // pipeline + { + BSONArrayBuilder pipelineBab( + bob.subarrayStart(AggregateCommandRequest::kPipelineFieldName)); + auto serializedPipeline = pipeline.serializeToBson(opts); + for (const auto& stage : serializedPipeline) { + pipelineBab.append(stage); + } + pipelineBab.doneFast(); + } + + // explain + if (aggregateCommand.getExplain().has_value()) { + bob.append(AggregateCommandRequest::kExplainFieldName, true); + } + + // allowDiskUse + if (auto param = aggregateCommand.getAllowDiskUse(); param.has_value()) { + bob.append(AggregateCommandRequest::kAllowDiskUseFieldName, param.value_or(false)); + } + + // collation + if (auto param = aggregateCommand.getCollation()) { + bob.append(AggregateCommandRequest::kCollationFieldName, param.get()); + } + + // let + if (auto letSpec = aggregateCommand.getLet()) { + auto redactedObj = extractLetSpecShape(letSpec.get(), opts, expCtx); + auto ownedObj = redactedObj.getOwned(); + bob.append(FindCommandRequest::kLetFieldName, std::move(ownedObj)); + } + return bob.obj(); +} + +QueryShapeHash hash(const BSONObj& queryShape) { + return QueryShapeHash::computeHash(reinterpret_cast<const uint8_t*>(queryShape.objdata()), + queryShape.objsize()); +} +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/README.md b/src/mongo/db/query/query_shape/README.md new file mode 100644 index 00000000000..b3c02d28e69 --- /dev/null +++ b/src/mongo/db/query/query_shape/README.md @@ -0,0 +1,69 @@ +# Query Shape +A query shape is a transformed version of a command with literal values replaced by a "canonical" +BSON Type placeholder. Hence, different instances of a command would be considered to have the same +query shape if they are identical once their literal values are abstracted. + +For example, these two queries would have the same shape: +```js +db.example.findOne({x: 24}); +db.example.findOne({x: 53}); +``` +While these queries would each have a distinct shape: +```js +db.example.findOne({x: 53, y: 1}); +db.example.findOne({x: 53}); +db.example.findOne({x: "string"}); +``` +While different literal _values_ result in the same shape (matching `x` for 23 vs 53), different +BSON _types_ of the literal are considered distinct shapes (matching `x` for 53 vs "string"). + +The concept of a query shape exists not just for the find command, but for many of the CRUD commands +and aggregate. It also includes most (but not all) components of these commands, not just the query +predicate (MatchExpresssion). In these ways, "query" is meant more generally. While some components +included in the query shape are shared across the different types of commands (e.g., the "hint" +field), some are unique. For example, a find command would include a `filter` while an aggregate +command would have a `pipeline`. + +You can see which components are considered part of the query shape or not for each specific shape +type in their respective "shape component" classes, whose purpose is to determine which components +are relevant and should be included for determining the shape for specific type of command. The +structure is as follows: +- [`CmdSpecificShapeComponents`](query_shape.h#L65) + - [`LetShapeComponent`](cmd_with_let_shape.h#L48) + - [`AggCmdShapeComponents`](agg_cmd_shape.h#L82) + - [`FindCmdShapeComponents`](find_cmd_shape.h#L48) + +See more information for the different shapes in their respective classes, structured as follows: +- [`Shape`](query_shape.h) + - [`CmdWithLetShape`](cmd_with_let_shape.h) + - [`AggCmdShape`](agg_cmd_shape.h) + - [`FindCmdShape`](find_cmd_shape.h) + +## Serialization Options +`SerializationOptions` describes the way we serialize literal values. + +There are 3 different serialization options: +- `kUnchanged`: literals are serialized unmodified + - `{x: 5, y: "hello"}` -> `{x: 5, y: "hello"}` +- `kToDebugTypeString`: human readable format, type string of the literal is serialized + - `{x: 5, y: "hello"}` -> `{x: "?number", y: "?string"}` +- `kToRepresentativeParseableValue`: literal serialized to one canonical value for given type, which + must be parseable + - `{x: 5, y: "hello"}` -> `{x: 1, y: "?"}` + - An example of a query which is serialized differently due to the parseable requirement is `{x: + {$regex: "^p.*"}}`. If we serialized the pattern as if it were a normal string we would end up + with `{x: {$regex: "?"}}` however `"?"` is not a valid regex pattern, so this would fail + parsing. Instead we will serialize it this way to maintain parseability, `{x: {$regex: + "\\?"}}`, since `"\\?"` is valid regex. + +See [serialization_options.h](serialization_options.h) for more details. + +When we compute the [query shape hash](query_shape.cpp#L99-107), we use the +`kToRepresentativeParseableValue`, since all literals of the same type will become the same value. +This allows us to group together queries that have the same structure but different literal values +into the same shape, since they will result in the same hash. The term we use to refer to this is +"shapify", as we simplify the queries into their query shape. + +When shapifying, we try to get as close as possible to the original user input, but there are some +stages like `$jsonSchema` and `$setWindowFields` that output "internal" stages that are already +transformed from user input. diff --git a/src/mongo/db/query/query_shape/SConscript b/src/mongo/db/query/query_shape/SConscript new file mode 100644 index 00000000000..d4bddba4934 --- /dev/null +++ b/src/mongo/db/query/query_shape/SConscript @@ -0,0 +1,42 @@ +# -*- mode: python -*- + +Import([ + "env", + "get_option", +]) + +env = env.Clone() + +env.Library( + target='query_shape', source=['query_shape.cpp', 'shape_helpers.cpp'], LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/db/exec/document_value/document_value', + '$BUILD_DIR/mongo/db/pipeline/field_path', + 'query_shape_common', + ], LIBDEPS_PRIVATE=[ + ]) + +env.Library( + target='query_shape_common', source=[ + 'query_shape.idl', + 'serialization_options.cpp', + ], LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/db/exec/document_value/document_value', + '$BUILD_DIR/mongo/db/pipeline/field_path', + ], LIBDEPS_PRIVATE=[ + ]) + +env.CppUnitTest( + target="db_query_query_shape_test", + source=[ + "query_shape_test.cpp", + "query_shape_test.idl", + ], + LIBDEPS=[ + "$BUILD_DIR/mongo/db/auth/authmocks", + "$BUILD_DIR/mongo/db/query/query_test_service_context", + "$BUILD_DIR/mongo/db/service_context_d_test_fixture", + "query_shape", + ], +) diff --git a/src/mongo/db/query/query_shape/agg_cmd_shape.cpp b/src/mongo/db/query/query_shape/agg_cmd_shape.cpp new file mode 100644 index 00000000000..e997150ecc6 --- /dev/null +++ b/src/mongo/db/query/query_shape/agg_cmd_shape.cpp @@ -0,0 +1,125 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_shape/agg_cmd_shape.h" + +#include "mongo/db/query/query_shape/shape_helpers.h" + +namespace mongo::query_shape { + +AggCmdShapeComponents::AggCmdShapeComponents( + const AggregateCommandRequest& aggRequest, + stdx::unordered_set<NamespaceString> involvedNamespaces_, + std::vector<BSONObj> pipeline) + : allowDiskUse(aggRequest.getAllowDiskUse()), + involvedNamespaces(std::move(involvedNamespaces_)), + representativePipeline(std::move(pipeline)) {} + +AggCmdShapeComponents::AggCmdShapeComponents( + OptionalBool allowDiskUse, + stdx::unordered_set<NamespaceString> involvedNamespaces_, + std::vector<BSONObj> pipeline) + : allowDiskUse(allowDiskUse), + involvedNamespaces(std::move(involvedNamespaces_)), + representativePipeline(std::move(pipeline)) {} + +void AggCmdShapeComponents::HashValue(absl::HashState state) const { + state = absl::HashState::combine(std::move(state), allowDiskUse); + for (auto&& shapifiedStage : representativePipeline) { + state = absl::HashState::combine(std::move(state), simpleHash(shapifiedStage)); + } +} + +void AggCmdShape::appendLetCmdSpecificShapeComponents( + BSONObjBuilder& bob, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const SerializationOptions& opts) const { + tassert(7633000, + "We don't support serializing to the unmodified shape here, since we have already " + "shapified and stored the representative query - we've lost the original literals", + opts.literalPolicy != LiteralSerializationPolicy::kUnchanged); + + if (opts == SerializationOptions::kRepresentativeQueryShapeSerializeOptions) { + // We have this copy stored already! + return _components.appendTo(bob); + } else { + // The cached pipeline shape doesn't match the requested options, so we have to + // re-parse the pipeline from the initial request. + expCtx->inMongos = _inMongos; + expCtx->addResolvedNamespaces(_components.involvedNamespaces); + auto reparsed = Pipeline::parse(_components.representativePipeline, expCtx); + auto serializedPipeline = reparsed->serializeToBson(opts); + AggCmdShapeComponents{ + _components.allowDiskUse, _components.involvedNamespaces, serializedPipeline} + .appendTo(bob); + } +} + +void AggCmdShapeComponents::appendTo(BSONObjBuilder& bob) const { + bob.append("command", "aggregate"); + + // pipeline + bob.append(AggregateCommandRequest::kPipelineFieldName, representativePipeline); + + // allowDiskUse + if (allowDiskUse.has_value()) { + bob.append(AggregateCommandRequest::kAllowDiskUseFieldName, bool(allowDiskUse)); + } +} + +// As part of the size, we must track the allocation of elements in the representative +// pipeline, as well as the elements in the unordered set of involved namespaces. +size_t AggCmdShapeComponents::size() const { + return sizeof(AggCmdShapeComponents) + shape_helpers::containerSize(representativePipeline) + + shape_helpers::containerSize(involvedNamespaces); +} + +AggCmdShape::AggCmdShape(const AggregateCommandRequest& aggregateCommand, + NamespaceString origNss, + stdx::unordered_set<NamespaceString> involvedNamespaces_, + const Pipeline& pipeline, + const boost::intrusive_ptr<ExpressionContext>& expCtx) + : CmdWithLetShape(aggregateCommand.getLet(), + expCtx, + _components, + std::move(origNss), + aggregateCommand.getCollation().value_or(BSONObj())), + _components(aggregateCommand, + std::move(involvedNamespaces_), + pipeline.serializeToBson( + SerializationOptions::kRepresentativeQueryShapeSerializeOptions)), + _inMongos(expCtx->inMongos) {} + +size_t AggCmdShape::extraSize() const { + // To account for possible padding, we calculate the extra space with the difference instead of + // using sizeof(bool); + return sizeof(AggCmdShape) - sizeof(CmdWithLetShape) - sizeof(AggCmdShapeComponents); +} + +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/agg_cmd_shape.h b/src/mongo/db/query/query_shape/agg_cmd_shape.h new file mode 100644 index 00000000000..c0ef5a7b06f --- /dev/null +++ b/src/mongo/db/query/query_shape/agg_cmd_shape.h @@ -0,0 +1,103 @@ +/** + * Copyright (C) 2023-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 <boost/intrusive_ptr.hpp> + +#include "mongo/db/pipeline/aggregate_command_gen.h" +#include "mongo/db/pipeline/expression_context.h" +#include "mongo/db/pipeline/pipeline.h" +#include "mongo/db/query/query_shape/cmd_with_let_shape.h" +#include "mongo/db/query/query_shape/query_shape.h" + +namespace mongo::query_shape { + +/** + * A struct representing the aggregate command's specific components that are to be considered part + * of the query shape. + * + * This struct stores the shapified version of the pipeline as a memory optimization. We'll need to + * store the BSON version in either case, since often the parsed version needs that BSON to survive + * as backing memory, so we store the representative pipeline shape so that we are able to parse the + * pipeline again if we need to compute a different shape. + */ +struct AggCmdShapeComponents : public query_shape::CmdSpecificShapeComponents { + AggCmdShapeComponents(const AggregateCommandRequest&, + stdx::unordered_set<NamespaceString> involvedNamespaces, + std::vector<BSONObj> shapifiedPipeline); + + AggCmdShapeComponents(OptionalBool allowDiskUse, + stdx::unordered_set<NamespaceString> involvedNamespaces, + std::vector<BSONObj> shapifiedPipeline); + + size_t size() const final; + + void appendTo(BSONObjBuilder&) const; + + void HashValue(absl::HashState state) const final; + + OptionalBool allowDiskUse; + + stdx::unordered_set<NamespaceString> involvedNamespaces; + + // The representative query shape of the pipeline. + std::vector<BSONObj> representativePipeline; +}; + +/** + * A class representing the query shape of an aggregate command. The components are listed above. + * This class knows how to utilize those components to serialize to BSON with any + * SerializationOptions. Mostly this involves correctly setting up an ExpressionContext to re-parse + * the request if needed. + */ +class AggCmdShape : public CmdWithLetShape { +public: + AggCmdShape(const AggregateCommandRequest&, + NamespaceString origNss, + stdx::unordered_set<NamespaceString> involvedNamespaces, + const Pipeline&, + const boost::intrusive_ptr<ExpressionContext>&); + + void appendLetCmdSpecificShapeComponents(BSONObjBuilder& bob, + const boost::intrusive_ptr<ExpressionContext>&, + const SerializationOptions&) const final; + size_t extraSize() const final override; + +private: + AggCmdShapeComponents _components; + // Flag to denote if the query was run on mongos. Needed to rebuild the "dummy" expression + // context for re-parsing. + bool _inMongos; +}; +static_assert(sizeof(AggCmdShape) <= + sizeof(CmdWithLetShape) + sizeof(AggCmdShapeComponents) + 8 /* bool and padding*/, + "If the class' members have changed, this assert and the extraSize() calculation may " + "need to be updated with a new value."); +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/agg_cmd_shape_test.cpp b/src/mongo/db/query/query_shape/agg_cmd_shape_test.cpp new file mode 100644 index 00000000000..c617391f4e7 --- /dev/null +++ b/src/mongo/db/query/query_shape/agg_cmd_shape_test.cpp @@ -0,0 +1,266 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/bson/json.h" +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/query_shape/agg_cmd_shape.h" +#include "mongo/db/query/query_shape/cmd_with_let_shape.h" +#include "mongo/db/query/query_test_service_context.h" +#include "mongo/unittest/unittest.h" + +namespace mongo::query_shape { + +namespace { +static const NamespaceString kDefaultTestNss = NamespaceString("testDB.testColl"); + +class AggCmdShapeTest : public unittest::Test { +public: + void setUp() final { + _queryTestServiceContext = std::make_unique<QueryTestServiceContext>(); + _operationContext = _queryTestServiceContext->makeOperationContext(); + _expCtx = make_intrusive<ExpressionContextForTest>(); + } + + std::unique_ptr<AggregateCommandRequest> makeAggregateCommandRequest( + std::vector<StringData> stagesJson, + boost::optional<StringData> letJson = boost::none, + boost::optional<StringData> collationJson = boost::none) { + std::vector<BSONObj> pipeline; + for (auto&& stage : stagesJson) { + pipeline.push_back(fromjson(stage.rawData())); + } + + auto aggRequest = + std::make_unique<AggregateCommandRequest>(kDefaultTestNss, std::move(pipeline)); + if (letJson) { + aggRequest->setLet(fromjson(letJson->rawData())); + } + if (collationJson) { + aggRequest->setCollation(fromjson(collationJson->rawData())); + } + return aggRequest; + } + + std::unique_ptr<AggCmdShape> makeShapeFromPipeline( + std::vector<StringData> stagesJson, + boost::optional<StringData> letJson = boost::none, + boost::optional<StringData> collationJson = boost::none) { + + auto aggRequest = makeAggregateCommandRequest( + std::move(stagesJson), std::move(letJson), std::move(collationJson)); + + auto parsedPipeline = Pipeline::parse(aggRequest->getPipeline(), _expCtx); + return std::make_unique<AggCmdShape>(*aggRequest, + kDefaultTestNss, + stdx::unordered_set<NamespaceString>{kDefaultTestNss}, + *parsedPipeline, + _expCtx); + } + std::unique_ptr<AggCmdShapeComponents> makeShapeComponentsFromPipeline( + std::vector<StringData> stagesJson, OptionalBool allowDiskUse = {}) { + auto aggRequest = makeAggregateCommandRequest(std::move(stagesJson)); + + auto parsedPipeline = Pipeline::parse(aggRequest->getPipeline(), _expCtx); + return std::make_unique<AggCmdShapeComponents>( + *aggRequest, + stdx::unordered_set<NamespaceString>{kDefaultTestNss}, + parsedPipeline->serializeToBson( + SerializationOptions::kRepresentativeQueryShapeSerializeOptions)); + } + + std::unique_ptr<QueryTestServiceContext> _queryTestServiceContext; + + ServiceContext::UniqueOperationContext _operationContext; + boost::intrusive_ptr<ExpressionContext> _expCtx; +}; + +TEST_F(AggCmdShapeTest, BasicPipelineShape) { + auto shape = + makeShapeFromPipeline({R"({$match: {x: 3, y: {$lte: 3}}})"_sd, + R"({$group: {_id: "$y", z: {$max: "$z"}, w: {$avg: "$w"}}})"}); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "cmdNs": { + "db": "testDB", + "coll": "testColl" + }, + "command": "aggregate", + "pipeline": [ + { + "$match": { + "$and": [ + { + "x": { + "$eq": "?number" + } + }, + { + "y": { + "$lte": "?number" + } + } + ] + } + }, + { + "$group": { + "_id": "$y", + "z": { + "$max": "$z" + }, + "w": { + "$avg": "$w" + } + } + } + ] + })", + shape->toBson(_operationContext.get(), + SerializationOptions::kDebugQueryShapeSerializeOptions)); +} + +TEST_F(AggCmdShapeTest, IncludesLet) { + auto shape = makeShapeFromPipeline({R"({$match: {x: 3}})"_sd, R"({$limit: 2})"_sd}, + R"({x: 4, y: "str"})"_sd); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "cmdNs": { + "db": "testDB", + "coll": "testColl" + }, + "let": { + "x": "?number", + "y": "?string" + }, + "command": "aggregate", + "pipeline": [ + { + "$match": { + "x": { + "$eq": "?number" + } + } + }, + { + "$limit": "?number" + } + ] + })", + shape->toBson(_operationContext.get(), + SerializationOptions::kDebugQueryShapeSerializeOptions)); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "cmdNs": { + "db": "testDB", + "coll": "testColl" + }, + "let": { + "x": { + "$const": 1 + }, + "y": { + "$const": "?" + } + }, + "command": "aggregate", + "pipeline": [ + { + "$match": { + "x": { + "$eq": 1 + } + } + }, + { + "$limit": 1 + } + ] + })", + shape->toBson(_operationContext.get(), + SerializationOptions::kRepresentativeQueryShapeSerializeOptions)); +} + +TEST_F(AggCmdShapeTest, SizeOfAggCmdShapeComponents) { + auto aggComponents = makeShapeComponentsFromPipeline( + {R"({$match: {x: 3, y: {$lte: 3}}})"_sd, + R"({$group: {_id: "$y", z: {$max: "$z"}, w: {$avg: "$w"}}})"}, + false /*allowDiskUse*/); + + // The sizes of any members of AggCmdShapeComponents are typically accounted for by + // sizeof(AggCmdShapeComponents). The important part of the test here is to ensure that any + // additional memory allocations are also included in the size() operation. In our case, + // we expect additional memory use from the representative pipeline and the involved + // namespaces set. + const auto pipelineSize = shape_helpers::containerSize(aggComponents->representativePipeline); + const auto involvedNamespacesSize = sizeof(kDefaultTestNss) + + kDefaultTestNss.size(); // kDefaultTestNss is the only value in the unordered set. + + ASSERT_EQ(aggComponents->size(), + sizeof(AggCmdShapeComponents) + pipelineSize + involvedNamespacesSize); +} + +TEST_F(AggCmdShapeTest, EquivalentAggCmdShapeComponentSizes) { + auto aggComponentsDiskUseFalse = makeShapeComponentsFromPipeline( + {R"({$match: {x: 3, y: {$lte: 3}}})"_sd, + R"({$group: {_id: "$y", z: {$max: "$z"}, w: {$avg: "$w"}}})"}, + false /*allowDiskUse*/); + auto aggComponentsDiskUseTrue = makeShapeComponentsFromPipeline( + {R"({$match: {x: 3, y: {$lte: 3}}})"_sd, + R"({$group: {_id: "$y", z: {$max: "$z"}, w: {$avg: "$w"}}})"}, + true /*allowDiskUse*/); + ASSERT_EQ(aggComponentsDiskUseFalse->size(), aggComponentsDiskUseTrue->size()); +} + +TEST_F(AggCmdShapeTest, DifferentAggCmdShapeComponentSizes) { + auto smallAggComponents = makeShapeComponentsFromPipeline({R"({$match: {x: 3, y: {$lte: 3}}})"}, + false /*allowDiskUse*/); + auto largeAggComponents = makeShapeComponentsFromPipeline( + {R"({$match: {x: 3, y: {$lte: 3}}})"_sd, + R"({$group: {_id: "$y", z: {$max: "$z"}, w: {$avg: "$w"}}})"}, + false /*allowDiskUse*/); + ASSERT_LT(smallAggComponents->size(), largeAggComponents->size()); +} + +TEST_F(AggCmdShapeTest, SizeOfAggCmdShapeWithAndWithoutLet) { + auto shapeWithoutLet = makeShapeFromPipeline({R"({$match: {x: 3}})"_sd, R"({$limit: 2})"_sd}); + auto shapeWithLet = makeShapeFromPipeline({R"({$match: {x: 3}})"_sd, R"({$limit: 2})"_sd}, + R"({x: 4, y: "str"})"_sd); + ASSERT_LT(shapeWithoutLet->size(), shapeWithLet->size()); +} + +TEST_F(AggCmdShapeTest, SizeOfAggCmdShapeWithAndWithoutCollation) { + auto shapeWithoutCollation = + makeShapeFromPipeline({R"({$match: {x: 3}})"_sd, R"({$limit: 2})"_sd}); + auto shapeWithCollation = makeShapeFromPipeline( + {R"({$match: {x: 3}})"_sd, R"({$limit: 2})"_sd}, boost::none, R"({locale: "en_US"})"_sd); + ASSERT_LT(shapeWithoutCollation->size(), shapeWithCollation->size()); +} +} // namespace +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/cmd_with_let_shape.cpp b/src/mongo/db/query/query_shape/cmd_with_let_shape.cpp new file mode 100644 index 00000000000..2bbb6dfeadc --- /dev/null +++ b/src/mongo/db/query/query_shape/cmd_with_let_shape.cpp @@ -0,0 +1,107 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_shape/cmd_with_let_shape.h" + +namespace mongo::query_shape { + +namespace { +BSONObj extractLetShape(BSONObj letSpec, + const SerializationOptions& opts, + const boost::intrusive_ptr<ExpressionContext>& expCtx) { + if (letSpec.isEmpty()) { + // Fast path for the common case. + return letSpec; + } + + BSONObjBuilder bob; + for (BSONElement elem : letSpec) { + auto expr = Expression::parseOperand(expCtx.get(), elem, expCtx->variablesParseState); + auto redactedValue = expr->serialize(opts); + // Note that this will throw on deeply nested let variables. + redactedValue.addToBsonObj(&bob, opts.serializeFieldPathFromString(elem.fieldName())); + } + return bob.obj(); +} + +auto representativeLetShape(boost::optional<BSONObj> let, + const boost::intrusive_ptr<ExpressionContext>& expCtx) { + return let ? extractLetShape( + *let, SerializationOptions::kRepresentativeQueryShapeSerializeOptions, expCtx) + : BSONObj(); +} +} // namespace + +LetShapeComponent::LetShapeComponent(boost::optional<BSONObj> let, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const CmdSpecificShapeComponents& unownedInnerComponents_) + : shapifiedLet(representativeLetShape(let, expCtx)), + hasLet(bool(let)), + unownedInnerComponents(unownedInnerComponents_) {} + +void LetShapeComponent::HashValue(absl::HashState state) const { + state = absl::HashState::combine( + std::move(state), hasLet, simpleHash(shapifiedLet), unownedInnerComponents); +} + +size_t LetShapeComponent::size() const { + return sizeof(LetShapeComponent) + shapifiedLet.objsize() + unownedInnerComponents.size(); +} + +void LetShapeComponent::addLetBson(BSONObjBuilder& bob, + const SerializationOptions& opts, + const boost::intrusive_ptr<ExpressionContext>& expCtx) const { + if (hasLet) { + auto shapeToAppend = shapifiedLet; + if (opts != SerializationOptions::kRepresentativeQueryShapeSerializeOptions) { + // We have the representative query cached/stored here, but the caller is asking for a + // different format, so we must re-compute. + shapeToAppend = extractLetShape(shapifiedLet, opts, expCtx); + } + bob.append(FindCommandRequest::kLetFieldName, shapeToAppend); + } +} + +void CmdWithLetShape::appendCmdSpecificShapeComponents(BSONObjBuilder& bob, + OperationContext* opCtx, + const SerializationOptions& opts) const { + auto expCtx = + ExpressionContext::makeBlankExpressionContext(opCtx, nssOrUUID, _let.shapifiedLet); + _let.addLetBson(bob, opts, expCtx); + appendLetCmdSpecificShapeComponents(bob, expCtx, opts); +} + +CmdWithLetShape::CmdWithLetShape(boost::optional<BSONObj> let, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const CmdSpecificShapeComponents& unownedInnerComponents, + NamespaceStringOrUUID nssOrUUID, + BSONObj collation) + : Shape(nssOrUUID, collation), _let(let, expCtx, unownedInnerComponents) {} + +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/cmd_with_let_shape.h b/src/mongo/db/query/query_shape/cmd_with_let_shape.h new file mode 100644 index 00000000000..a1c127b9999 --- /dev/null +++ b/src/mongo/db/query/query_shape/cmd_with_let_shape.h @@ -0,0 +1,109 @@ +/** + * Copyright (C) 2023-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 "mongo/bson/bsonobj.h" +#include "mongo/db/pipeline/expression_context.h" +#include "mongo/db/query/parsed_find_command.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_shape/serialization_options.h" + +namespace mongo::query_shape { + +/** + * This struct is bit of a weird one. We want to use it as the shape's _entire_ "specific + * components" (rather than introduce more virtual functions to that interface). So, we track here + * the let component (as the name suggests) but we also keep an unowned reference to the specific + * components of CmdWithLetShape sub-classes. This class doesn't really do all that much with those + * components except track a reference to them and ensure their size is accounted for and their hash + * value is incorporated. + */ +struct LetShapeComponent : public CmdSpecificShapeComponents { + LetShapeComponent(boost::optional<BSONObj> let, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const CmdSpecificShapeComponents& unownedInnerComponents); + + /** + * Hashes to include the shapified let parameters and also the hash of 'unownedInnerComponents'. + */ + void HashValue(absl::HashState state) const final; + + /** + * Includes the size of the let parameters and the size of 'unownedInnerComponents.' + */ + size_t size() const final; + + /** + * Adds _only_ the let params. + */ + void addLetBson(BSONObjBuilder&, + const SerializationOptions&, + const boost::intrusive_ptr<ExpressionContext>&) const; + + BSONObj shapifiedLet; + bool hasLet; + // Tracked so that this can be hash combined correctly. + const CmdSpecificShapeComponents& unownedInnerComponents; +}; + +/** + * The 'let' command argument is semi-generic in that it is supported in a couple commands. However + * it is treated specially since it supports using expressions as the let constants. Using + * expressions induces a library dependency that we don't want in the Shape interface itself. So + * this class handles tracking and adding the 'let' component of the shape for sub-classes. + */ +class CmdWithLetShape : public Shape { +public: + CmdWithLetShape(boost::optional<BSONObj> let, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const CmdSpecificShapeComponents& unownedInnerComponents, + NamespaceStringOrUUID, + BSONObj collation_); + + const CmdSpecificShapeComponents& specificComponents() const final { + return _let; + } + +protected: + void appendCmdSpecificShapeComponents(BSONObjBuilder&, + OperationContext* opCtx, + const SerializationOptions& opts) const final; + virtual void appendLetCmdSpecificShapeComponents( + BSONObjBuilder&, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const SerializationOptions&) const = 0; + + LetShapeComponent _let; +}; +static_assert(sizeof(CmdWithLetShape) == sizeof(Shape) + sizeof(LetShapeComponent), + "If the class' members have changed, this assert and the extraSize() calculation may " + "need to be updated with a new value."); + +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/cmd_with_let_shape_test.cpp b/src/mongo/db/query/query_shape/cmd_with_let_shape_test.cpp new file mode 100644 index 00000000000..21812d97fc7 --- /dev/null +++ b/src/mongo/db/query/query_shape/cmd_with_let_shape_test.cpp @@ -0,0 +1,78 @@ +/** + * Copyright (C) 2024-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. + */ + +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/query_shape/cmd_with_let_shape.h" +#include "mongo/db/service_context_test_fixture.h" +#include "mongo/unittest/unittest.h" + +namespace mongo::query_shape { + +namespace { +static const NamespaceString kDefaultTestNss = NamespaceString("testDB.testColl"); + +class CmdWithLetShapeTest : public unittest::Test {}; + + +struct DummyInnerComponent : public CmdSpecificShapeComponents { + DummyInnerComponent(){}; + void HashValue(absl::HashState state) const {} + size_t size() const final { + return sizeof(*this); + } +}; + +TEST_F(CmdWithLetShapeTest, SizeOfLetShapeComponent) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + auto let = fromjson(R"({x: 4, y: "str"})"); + auto innerComponents = std::make_unique<DummyInnerComponent>(); + auto components = std::make_unique<LetShapeComponent>(let, expCtx, *innerComponents); + + const auto minimumSize = sizeof(CmdSpecificShapeComponents) + sizeof(BSONObj) + sizeof(bool) + + sizeof(void*) /*CmdSpecificShapeComponents&*/ + + static_cast<size_t>(components->shapifiedLet.objsize()) + + components->unownedInnerComponents.size(); + + ASSERT_GTE(components->size(), minimumSize); + ASSERT_LTE(components->size(), minimumSize + 8 /*padding*/); +} + +TEST_F(CmdWithLetShapeTest, SizeOfComponentWithAndWithoutLet) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + auto let = fromjson(R"({x: 4, y: "str"})"); + auto innerComponents = std::make_unique<DummyInnerComponent>(); + auto componentsWithLet = std::make_unique<LetShapeComponent>(let, expCtx, *innerComponents); + auto componentsWithNoLet = + std::make_unique<LetShapeComponent>(boost::none, expCtx, *innerComponents); + + ASSERT_LT(componentsWithNoLet->size(), componentsWithLet->size()); +} + +} // namespace +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/find_cmd_shape.cpp b/src/mongo/db/query/query_shape/find_cmd_shape.cpp new file mode 100644 index 00000000000..2d018de2619 --- /dev/null +++ b/src/mongo/db/query/query_shape/find_cmd_shape.cpp @@ -0,0 +1,227 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_shape/find_cmd_shape.h" + +#include "mongo/db/query/projection_ast_util.h" +#include "mongo/db/query/query_shape/shape_helpers.h" + +namespace mongo::query_shape { +namespace { + +BSONObj projectionShape(const boost::optional<projection_ast::Projection>& proj, + const SerializationOptions& opts = + SerializationOptions::kRepresentativeQueryShapeSerializeOptions) { + return proj ? projection_ast::serialize(*proj->root(), opts) : BSONObj(); +} + +BSONObj sortShape(const boost::optional<SortPattern>& sort, + const SerializationOptions& opts = + SerializationOptions::kRepresentativeQueryShapeSerializeOptions) { + return sort + ? sort->serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts) + .toBson() + : BSONObj(); +} + +void maybeAddWithName(const OptionalBool& optBool, BSONObjBuilder& bob, StringData name) { + if (optBool.has_value()) { + bob.append(name, bool(optBool)); + } +} + +void addRemainingFindCommandFields(const FindCmdShapeComponents& components, BSONObjBuilder& bob) { + maybeAddWithName(components.singleBatch, bob, FindCommandRequest::kSingleBatchFieldName); + maybeAddWithName(components.allowDiskUse, bob, FindCommandRequest::kAllowDiskUseFieldName); + maybeAddWithName(components.returnKey, bob, FindCommandRequest::kReturnKeyFieldName); + maybeAddWithName(components.showRecordId, bob, FindCommandRequest::kShowRecordIdFieldName); + maybeAddWithName(components.tailable, bob, FindCommandRequest::kTailableFieldName); + maybeAddWithName(components.awaitData, bob, FindCommandRequest::kAwaitDataFieldName); + maybeAddWithName(components.oplogReplay, bob, FindCommandRequest::kOplogReplayFieldName); +} + +} // namespace + +FindCmdShapeComponents::FindCmdShapeComponents( + const ParsedFindCommand& request, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const SerializationOptions& opts) + : filter(request.filter->serialize(opts)), + projection(projectionShape(request.proj, opts)), + sort(sortShape(request.sort, opts)), + min(shape_helpers::extractMinOrMaxShape(request.findCommandRequest->getMin(), opts)), + max(shape_helpers::extractMinOrMaxShape(request.findCommandRequest->getMax(), opts)), + singleBatch(request.findCommandRequest->getSingleBatch()), + allowDiskUse(request.findCommandRequest->getAllowDiskUse().has_value() + ? boost::optional<bool>(bool(request.findCommandRequest->getAllowDiskUse())) + : boost::none), + returnKey(request.findCommandRequest->getReturnKey()), + showRecordId(request.findCommandRequest->getShowRecordId()), + tailable(request.findCommandRequest->getTailable()), + awaitData(request.findCommandRequest->getAwaitData()), + oplogReplay(request.findCommandRequest->getOplogReplay()), + hasField(), + serializationOpts(opts) { + hasField.projection = request.proj.has_value(); + hasField.sort = request.sort.has_value(); + hasField.limit = request.findCommandRequest->getLimit().has_value(); + hasField.skip = request.findCommandRequest->getSkip().has_value(); +} + +void FindCmdShapeComponents::appendTo(BSONObjBuilder& bob) const { + + bob.append("command", "find"); + + std::unique_ptr<MatchExpression> filterExpr; + // Filter. + bob.append(FindCommandRequest::kFilterFieldName, filter); + + if (hasField.projection) { + bob.append(FindCommandRequest::kProjectionFieldName, projection); + } + + if (!max.isEmpty()) { + bob.append(FindCommandRequest::kMaxFieldName, max); + } + if (!min.isEmpty()) { + bob.append(FindCommandRequest::kMinFieldName, min); + } + + // Sort. + if (hasField.sort) { + bob.append(FindCommandRequest::kSortFieldName, sort); + } + + // The values here don't matter (assuming we're not using the 'kUnchanged' policy). + tassert(7973601, + "Serialization policy not supported - original values have been discarded", + serializationOpts.literalPolicy != LiteralSerializationPolicy::kUnchanged); + if (hasField.limit) { + serializationOpts.appendLiteral(&bob, FindCommandRequest::kLimitFieldName, 1ll); + } + if (hasField.skip) { + serializationOpts.appendLiteral(&bob, FindCommandRequest::kSkipFieldName, 1ll); + } + + // Add the fields that require no transformation. + addRemainingFindCommandFields(*this, bob); +} + +void FindCmdShapeComponents::HashValue(absl::HashState state) const { + absl::HashState::combine(std::move(state), + simpleHash(filter), + simpleHash(projection), + simpleHash(sort), + simpleHash(min), + simpleHash(max), + singleBatch, + allowDiskUse, + returnKey, + showRecordId, + tailable, + awaitData, + oplogReplay, + hasField); +} + +std::unique_ptr<FindCommandRequest> FindCmdShape::toFindCommandRequest() const { + auto fcr = std::make_unique<FindCommandRequest>(nssOrUUID); + + fcr->setFilter(components.filter); + if (components.hasField.projection) + fcr->setProjection(components.projection); + if (components.hasField.sort) + fcr->setSort(components.sort); + + fcr->setMin(components.min); + fcr->setMax(components.max); + + // Doesn't matter what value to use for limit and skip in the context of a shape. + if (components.hasField.limit) + fcr->setLimit(1ll); + if (components.hasField.skip) + fcr->setSkip(1ll); + + // All the booleans. + if (components.singleBatch.has_value()) + fcr->setSingleBatch(bool(components.singleBatch)); + if (components.allowDiskUse.has_value()) + fcr->setAllowDiskUse(bool(components.allowDiskUse)); + if (components.returnKey.has_value()) + fcr->setReturnKey(bool(components.returnKey)); + if (components.showRecordId.has_value()) + fcr->setShowRecordId(bool(components.showRecordId)); + if (components.tailable.has_value()) + fcr->setTailable(bool(components.tailable)); + if (components.awaitData.has_value()) + fcr->setAwaitData(bool(components.awaitData)); + if (components.oplogReplay.has_value()) + fcr->setOplogReplay(bool(components.oplogReplay)); + + // Common shape components. + if (_let.hasLet) + fcr->setLet(_let.shapifiedLet); + if (!collation.isEmpty()) + fcr->setCollation(collation); + + + return fcr; +} + +FindCmdShape::FindCmdShape(const ParsedFindCommand& findRequest, + const boost::intrusive_ptr<ExpressionContext>& expCtx) + : CmdWithLetShape(findRequest.findCommandRequest->getLet(), + expCtx, + components, + findRequest.findCommandRequest->getNamespaceOrUUID(), + findRequest.findCommandRequest->getCollation()), + components(findRequest, expCtx) {} + +void FindCmdShape::appendLetCmdSpecificShapeComponents( + BSONObjBuilder& bob, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const SerializationOptions& opts) const { + if (opts == SerializationOptions::kRepresentativeQueryShapeSerializeOptions) { + // Fast path: we already have this. + return components.appendTo(bob); + } else { + // Slow path: we need to re-parse from our representative shapes. + auto request = uassertStatusOKWithContext( + parsed_find_command::parse(expCtx, + toFindCommandRequest(), + ExtensionsCallbackNoop(), + MatchExpressionParser::kAllowAllSpecialFeatures), + "Could not re-parse a representative query shape"); + + // This constructor will shapify according to the options. + FindCmdShapeComponents{*request, expCtx, opts}.appendTo(bob); + } +} + +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/find_cmd_shape.h b/src/mongo/db/query/query_shape/find_cmd_shape.h new file mode 100644 index 00000000000..49d70b8ec27 --- /dev/null +++ b/src/mongo/db/query/query_shape/find_cmd_shape.h @@ -0,0 +1,130 @@ +/** + * Copyright (C) 2023-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 "mongo/db/query/query_shape/cmd_with_let_shape.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_shape/shape_helpers.h" + +namespace mongo::query_shape { + +/** + * This struct tracks the components of a find command which are important for the find query shape. + * It attempts to only track those which are _unique_ to a find command - common elements should go + * on some super class. + * + * Data elements which are shapified like 'filter' are stored in their shapified form. By default + * and in most cases this will be the representative query shape form so that it can be re-parsed, + * but as a convenience for serializing it is also supported to construct and serialize this with + * other options. + */ +struct FindCmdShapeComponents : public CmdSpecificShapeComponents { + + FindCmdShapeComponents(const ParsedFindCommand& request, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const SerializationOptions& opts = + SerializationOptions::kRepresentativeQueryShapeSerializeOptions); + + /** + * Appends using the SerializationOptions given in the constructor. + */ + void appendTo(BSONObjBuilder&) const; + + size_t size() const final { + return sizeof(FindCmdShapeComponents) + filter.objsize() + projection.objsize() + + sort.objsize() + min.objsize() + max.objsize(); + } + + BSONObj filter; + BSONObj projection; + BSONObj sort; + BSONObj min; + BSONObj max; + + OptionalBool singleBatch; + OptionalBool allowDiskUse; + OptionalBool returnKey; + OptionalBool showRecordId; + OptionalBool tailable; + OptionalBool awaitData; + OptionalBool oplogReplay; + + // This anonymous struct represents the presence of the member variables as C++ bit fields. + // In doing so, each of these boolean values takes up 1 bit instead of 1 byte. + struct HasField { + HasField() : projection(false), sort(false), limit(false), skip(false) {} + bool projection : 1; + bool sort : 1; + bool limit : 1; + bool skip : 1; + } hasField; + + // We save a copy of the options used when constructed so we know how to properly append things + // like limit and skip - either a 1 or "?number". We could have the caller pass the options + // again during 'appendTo()', but this introduces a risk that the options provided are different + // than the ones we used to compute 'filter' and the other components. + SerializationOptions serializationOpts; + + void HashValue(absl::HashState state) const final; +}; + +class FindCmdShape : public CmdWithLetShape { +public: + FindCmdShape(const ParsedFindCommand& findRequest, + const boost::intrusive_ptr<ExpressionContext>& expCtx); + + /** + * Assembles a parseable FindCommandRequest representing this shape - some of the pieces are + * stored right here in the shape, others are in parent classes. + */ + std::unique_ptr<FindCommandRequest> toFindCommandRequest() const; + + FindCmdShapeComponents components; + +protected: + void appendLetCmdSpecificShapeComponents(BSONObjBuilder& bob, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const SerializationOptions& opts) const final; +}; + +template <typename H> +H AbslHashValue(H h, const FindCmdShapeComponents::HasField& hasField) { + return H::combine( + std::move(h), hasField.projection, hasField.sort, hasField.limit, hasField.skip); +} + +// This assertion is still active on the maintained master branch. On the v6.0 branch, we disable it +// since it is not passing on all toolchains/platforms - notably x86 macOS. The intent of the +// assertion is to prevent accidental additions of data members, which should not happen on this +// branch without first happening on the master branch and passing that assertion. +// static_assert(sizeof(FindCmdShape) == sizeof(CmdWithLetShape) + sizeof(FindCmdShapeComponents), +// "If the class' members have changed, this assert and the extraSize() calculation +// may " "need to be updated with a new value."); +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/find_cmd_shape_test.cpp b/src/mongo/db/query/query_shape/find_cmd_shape_test.cpp new file mode 100644 index 00000000000..0d839a5d3d2 --- /dev/null +++ b/src/mongo/db/query/query_shape/find_cmd_shape_test.cpp @@ -0,0 +1,238 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/query_shape/find_cmd_shape.h" +#include "mongo/db/service_context_test_fixture.h" +#include "mongo/unittest/unittest.h" + +namespace mongo::query_shape { + +namespace { +/** + * TODO this was stolen from another test. Time for a library? + * Simplistic redaction strategy for testing which appends the field name to the prefix "REDACT_". + */ +std::string applyHmacForTest(StringData sd) { + return "REDACT_" + sd.toString(); +} + +static const NamespaceStringOrUUID kDefaultTestNss = + NamespaceStringOrUUID{NamespaceString("testDB.testColl")}; + +struct RequestOptions { + OptionalBool singleBatch = {}; + OptionalBool allowDiskUse = {}; + OptionalBool returnKey = {}; + OptionalBool showRecordId = {}; + OptionalBool tailable = {}; + OptionalBool awaitData = {}; + OptionalBool limit = {}; + OptionalBool skip = {}; +}; +class FindCmdShapeTest : public ServiceContextTest { +public: + void setUp() final { + _expCtx = make_intrusive<ExpressionContextForTest>(); + } + + std::unique_ptr<FindCmdShape> makeShapeFromSort(StringData sortJson) { + auto fcr = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcr->setSort(fromjson(sortJson.rawData())); + auto&& parsedRequest = + uassertStatusOK(::mongo::parsed_find_command::parse(_expCtx, std::move(fcr))); + return std::make_unique<FindCmdShape>(*parsedRequest, _expCtx); + } + + BSONObj sortShape(StringData sortJson) { + auto shape = makeShapeFromSort(sortJson); + return shape->components.sort; + } + + /** + * Returns the shape of the input sort, or boost::none if the input shape was a natural sort + * which got converted into a hint. + */ + boost::optional<BSONObj> maybeRedactedSortShape(StringData sortJson) { + auto shape = makeShapeFromSort(sortJson); + SerializationOptions opts = SerializationOptions::kDebugQueryShapeSerializeOptions; + opts.transformIdentifiers = true; + opts.transformIdentifiersCallback = applyHmacForTest; + auto shapeBson = shape->toBson(_expCtx->opCtx, opts); + if (auto sortElem = shapeBson["sort"]; !sortElem.eoo()) { + return sortElem.Obj().getOwned(); + } + return boost::none; + } + + BSONObj redactedSortShape(StringData sortJson) { + return *maybeRedactedSortShape(sortJson); + } + + boost::intrusive_ptr<ExpressionContext> _expCtx; + + std::unique_ptr<FindCmdShapeComponents> makeShapeComponentsFromFilter( + BSONObj filter, const RequestOptions& requestOptions = {}) { + auto fcr = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcr->setFilter(filter.getOwned()); + fcr->setSingleBatch(requestOptions.singleBatch); + fcr->setAllowDiskUse(requestOptions.allowDiskUse); + fcr->setReturnKey(requestOptions.returnKey); + fcr->setAllowDiskUse(requestOptions.showRecordId); + fcr->setTailable(requestOptions.tailable); + fcr->setAwaitData(requestOptions.awaitData); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(_expCtx, {std::move(fcr)})); + return std::make_unique<FindCmdShapeComponents>(*parsedFind, _expCtx); + } + + std::unique_ptr<FindCmdShape> makeShapeFromFilter(const BSONObj& filter) { + auto fcr = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcr->setFilter(filter.getOwned()); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(_expCtx, {std::move(fcr)})); + return std::make_unique<FindCmdShape>(*parsedFind, _expCtx); + } +}; + +TEST_F(FindCmdShapeTest, NormalSortPattern) { + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({"a.b.c":1,"foo":-1})", + sortShape(R"({"a.b.c": 1, "foo": -1})")); +} + +TEST_F(FindCmdShapeTest, NaturalSortPattern) { + // $natural sorts are interpreted as a hint. Hints are not part of the shape (but should show up + // in the query stats key). + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({})", + sortShape(R"({$natural: 1})")); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({})", + sortShape(R"({$natural: -1})")); +} + +TEST_F(FindCmdShapeTest, NaturalSortPatternWithMeta) { + ASSERT_THROWS_CODE( + sortShape(R"({$natural: 1, x: {$meta: "textScore"}})"), DBException, ErrorCodes::BadValue); +} + +TEST_F(FindCmdShapeTest, MetaPatternWithoutNatural) { + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({"normal":1,"$computed1":{"$meta":"textScore"}})", + sortShape(R"({normal: 1, x: {$meta: "textScore"}})")); +} + +// Here we have one test to ensure that the redaction policy is accepted and applied in the +// query_shape utility, but there are more extensive redaction tests in sort_pattern_test.cpp +TEST_F(FindCmdShapeTest, RespectsRedactionPolicy) { + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({"REDACT_normal":1,"REDACT_y":1})", + redactedSortShape(R"({normal: 1, y: 1})")); + + // No need to redact $natural. Again, this will be interpreted as a hint, but this test is + // interesting to ensure the $-prefix of $natural doesn't confuse us. + ASSERT(!maybeRedactedSortShape(R"({$natural: 1})")); +} + +TEST_F(FindCmdShapeTest, SizeOfShapeComponents) { + auto query = BSON("query" << 1 << "xEquals" << 42); + auto findCmdComponent = makeShapeComponentsFromFilter(query.getOwned()); + const auto querySize = findCmdComponent->filter.objsize(); + + const auto minimumSize = sizeof(FindCmdShapeComponents) + querySize; + ASSERT_GT(findCmdComponent->size(), minimumSize); + ASSERT_LTE(findCmdComponent->size(), + minimumSize + static_cast<size_t>(4 * BSONObj().objsize())); +} + +TEST_F(FindCmdShapeTest, EquivalentShapeComponentsSizes) { + auto query = BSON("query" << 1 << "xEquals" << 42); + // Tailable can not be set together with 'singleBatch' option. + auto mostlyTrueComponent = makeShapeComponentsFromFilter(query.getOwned(), + {/* singleBatch = */ false, + /* allowDiskUse = */ true, + /* returnKey = */ true, + /* showRecordId = */ true, + /* tailable = */ true, + /* awaitData = */ true, + /* limit = */ true, + /* skip = */ true}); + + auto mostlyFalseComponent = makeShapeComponentsFromFilter(query.getOwned(), + {/* singleBatch = */ false, + /* allowDiskUse = */ false, + /* returnKey = */ false, + /* showRecordId = */ false, + /* tailable = */ true, + /* awaitData = */ false, + /* limit = */ false, + /* skip = */ false}); + + ASSERT_EQ(mostlyTrueComponent->size(), mostlyFalseComponent->size()); +} + +TEST_F(FindCmdShapeTest, DifferentShapeComponentsSizes) { + auto smallQuery = BSON("query" << BSONObj()); + auto smallFindCmdComponent = makeShapeComponentsFromFilter(smallQuery.getOwned()); + + auto largeQuery = BSON("query" << 1 << "xEquals" << 42); + auto largeFindCmdComponent = makeShapeComponentsFromFilter(largeQuery.getOwned()); + + ASSERT_LT(smallQuery.objsize(), largeQuery.objsize()); + ASSERT_LT(smallFindCmdComponent->size(), largeFindCmdComponent->size()); +} + +TEST_F(FindCmdShapeTest, SizeOfShapeWithAndWithoutLet) { + auto filter = BSON("query" << 1 << "xEquals" << 42); + auto shapeWithoutLet = makeShapeFromFilter(filter.getOwned()); + + auto fcr = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcr->setFilter(filter.getOwned()); + fcr->setLet(fromjson(R"({x: 4})")); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(_expCtx, {std::move(fcr)})); + auto shapeWithLet = std::make_unique<FindCmdShape>(*parsedFind, _expCtx); + + ASSERT_LT(shapeWithoutLet->size(), shapeWithLet->size()); +} + +TEST_F(FindCmdShapeTest, SizeOfShapeWithAndWithoutCollation) { + auto filter = BSON("query" << 1 << "xEquals" << 42); + auto shapeWithoutCollation = makeShapeFromFilter(filter.getOwned()); + + auto fcr = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcr->setFilter(filter.getOwned()); + fcr->setCollation(fromjson(R"({locale: "en_US"})")); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(_expCtx, {std::move(fcr)})); + auto shapeWithCollation = std::make_unique<FindCmdShape>(*parsedFind, _expCtx); + + ASSERT_LT(shapeWithoutCollation->size(), shapeWithCollation->size()); +} + +} // namespace + +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/query_shape.cpp b/src/mongo/db/query/query_shape/query_shape.cpp new file mode 100644 index 00000000000..2fa0520120e --- /dev/null +++ b/src/mongo/db/query/query_shape/query_shape.cpp @@ -0,0 +1,103 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_shape/query_shape.h" + +#include "mongo/base/status.h" +#include "mongo/crypto/sha256_block.h" +#include "mongo/db/query/find_command_gen.h" +#include "mongo/db/query/projection_ast_util.h" +#include "mongo/db/query/query_request_helper.h" +#include "mongo/db/query/query_shape/agg_cmd_shape.h" +#include "mongo/db/query/query_shape/find_cmd_shape.h" +#include "mongo/db/query/query_shape/query_shape_gen.h" +#include "mongo/db/query/query_shape/shape_helpers.h" +#include "mongo/db/query/sort_pattern.h" + +namespace mongo::query_shape { + +namespace { +void appendCmdNs(BSONObjBuilder& bob, + const NamespaceString& nss, + const SerializationOptions& opts) { + BSONObjBuilder nsObj = bob.subobjStart("cmdNs"); + shape_helpers::appendNamespaceShape(nsObj, nss, opts); + nsObj.doneFast(); +} +} // namespace + +Shape::Shape(NamespaceStringOrUUID nssOrUUID_, BSONObj collation_) + : nssOrUUID(nssOrUUID_), collation(std::move(collation_)) {} + + +BSONObj Shape::toBson(OperationContext* opCtx, const SerializationOptions& opts) const { + BSONObjBuilder bob; + appendCmdNsOrUUID(bob, opts); + if (!collation.isEmpty()) { + // Collation is never shapified. We use find command's collation name definition, but it + // should be the same for all requests. + bob.append(FindCommandRequest::kCollationFieldName, collation); + } + appendCmdSpecificShapeComponents(bob, opCtx, opts); + return bob.obj(); +} + +size_t Shape::size() const { + return sizeof(Shape) + shape_helpers::optionalObjSize(collation) + specificComponents().size() + + extraSize(); +} + +QueryShapeHash Shape::sha256Hash(OperationContext* opCtx) const { + // The Query Shape Hash should use the representative query shape. + auto serialized = + toBson(opCtx, SerializationOptions::kRepresentativeQueryShapeSerializeOptions); + return SHA256Block::computeHash((const uint8_t*)serialized.sharedBuffer().get(), + serialized.objsize()); +} + +void Shape::appendCmdNsOrUUID(BSONObjBuilder& bob, const SerializationOptions& opts) const { + if (nssOrUUID.nss()) { + appendCmdNs(bob, *nssOrUUID.nss(), opts); + } else { + BSONObjBuilder cmdNs = bob.subobjStart("cmdNs"); + cmdNs.append("uuid", opts.serializeIdentifier(nssOrUUID.uuid()->toString())); + cmdNs.append("db", opts.serializeIdentifier(nssOrUUID.db())); + cmdNs.doneFast(); + } +} + +void Shape::appendCmdNs(BSONObjBuilder& bob, + const NamespaceString& nss, + const SerializationOptions& opts) const { + BSONObjBuilder nsObj = bob.subobjStart("cmdNs"); + shape_helpers::appendNamespaceShape(nsObj, nss, opts); + nsObj.doneFast(); +} + +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/query_shape.h b/src/mongo/db/query/query_shape/query_shape.h new file mode 100644 index 00000000000..dc83cbab127 --- /dev/null +++ b/src/mongo/db/query/query_shape/query_shape.h @@ -0,0 +1,165 @@ +/** + * Copyright (C) 2023-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 "mongo/db/matcher/expression.h" +#include "mongo/db/pipeline/aggregate_command_gen.h" +#include "mongo/db/query/find_command_gen.h" +#include "mongo/db/query/query_request_helper.h" +#include "mongo/db/query/query_shape/shape_helpers.h" + +namespace mongo::query_shape { + +/** + * Each type of "query" command likely has different fields/options that are considered important + * for the shape. For example, a find command has a skip and a limit, and an aggregate command has a + * pipeline. This interface is used to allow different sub-commands to diverge in this way but still + * ensure we can appropriately hash them to compare their shapes, and properly account for their + * size. + * + * This struct is split out as a separate inheritence hierarchy from 'Shape' to make it easier to + * ensure each piece is hashed without sub-classes needing to enumerate the parent class's member + * variables. + */ +struct CmdSpecificShapeComponents { + virtual ~CmdSpecificShapeComponents() {} + + /** + * Sub-classes should implement this in a way which includes all shape-relevant state. If two + * shapes should compare equal, they should result in the same hash value. For example for the + * find command - we would include the _shapified_ filter and projection here, but we will not + * include the comment - which is not part of the shape. + */ + virtual void HashValue(absl::HashState state) const = 0; + + /** + * It is important for shape components to accurately report their size, and to make a + * reasonable effort to maintain a minimal size. We use the query shape in memory-constrained + * data structures, so a bigger shape means we can have fewer different shapes stored (for + * example in the query stats store). + * + * We cannot just use sizeof() because there are some variable size data members (like BSON + * objects) which depend on the particular instance. + */ + virtual size_t size() const = 0; + + // Some template boilerplate to allow sub-classes to overload the hash implementation. + template <typename H> + friend H AbslHashValue(H state, const CmdSpecificShapeComponents& value) { + value.HashValue(absl::HashState::Create(&state)); + return std::move(state); + } +}; + +using QueryShapeHash = SHA256Block; + +/** + * A query "shape" is a version of a command with literal values abstracted so that two instances of + * the command may compare/hash equal even if they use slightly different literal values. This + * concept exists not just the find command, but planned for many of the CRUD commands + aggregate. + * It also includes most (but not all) components of these commands, not just the query predicate + * (MatchExpresssion). In these ways, "query" is meant more generally. + * + * A "Query Shape" can vary depending on the command (e.g. find, aggregate, or distinct). This + * abstract struct is the API we must implement for each command which we want to have a "shape" + * concept. + * + * In order to properly account for the size of a query shape, the CmdSpecificShapeComponents should + * include all meaningful memory consumption, and be sure to report it in 'size()'. Subclasses of + * 'Shape' are not expected to have any meaningful memory usage outside of that struct. + */ +class Shape { +public: + virtual ~Shape() {} + + /** + * Sub-classes are expected to implement this as a mechanism for plugging in their command + * specific shape components. + */ + virtual const CmdSpecificShapeComponents& specificComponents() const = 0; + + /** + * Note this may involve re-parsing command BSON and so is not necessarily cheap. + */ + BSONObj toBson(OperationContext*, const SerializationOptions&) const; + + /** + * The Query Shape Hash is defined to be the SHA256 Hash of the representatice query shape. This + * helper computes that. + */ + QueryShapeHash sha256Hash(OperationContext*) const; + + /** + * The size of a query shape is important, since we store these in space-constrained + * environments like the query stats store. + */ + size_t size() const; + + /** + * This should be overriden by a child class if it has members whose sizes are not included in + * specificComponents().size(). + */ + virtual size_t extraSize() const { + return 0; + } + template <typename H> + friend H AbslHashValue(H h, const Shape& shape) { + h = H::combine(std::move(h), shape.nssOrUUID, shape.specificComponents()); + if (!shape.collation.isEmpty()) + h = H::combine(std::move(h), simpleHash(shape.collation)); + return h; + } + + + // Not shapified but it is an identifier so it may be transformed. + NamespaceStringOrUUID nssOrUUID; + + // Never shapified. If it's empty, leave it off. + BSONObj collation; + +protected: + Shape(NamespaceStringOrUUID, BSONObj collation_); + + /** + * Along with the hash implementation, this is the main way that shapes are 'shapified' - + * sub-classes should implement this to add the shapified versions of their literals to an + * object. Depending on 'opts', this may be eligible to be used for output in $queryStats or as + * the object to compute the QueryShapeHash. + */ + virtual void appendCmdSpecificShapeComponents(BSONObjBuilder&, + OperationContext*, + const SerializationOptions& opts) const = 0; + +private: + void appendCmdNsOrUUID(BSONObjBuilder&, const SerializationOptions&) const; + void appendCmdNs(BSONObjBuilder&, const NamespaceString&, const SerializationOptions&) const; +}; + +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/query_shape.idl b/src/mongo/db/query/query_shape/query_shape.idl new file mode 100644 index 00000000000..77e71756467 --- /dev/null +++ b/src/mongo/db/query/query_shape/query_shape.idl @@ -0,0 +1,50 @@ +# Copyright (C) 2023-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. + +global: + cpp_namespace: "mongo::query_shape" + +imports: + - "mongo/idl/basic_types.idl" + + +structs: + CommandNamespace: + description: "Representation of the cmdNs sub-object of the query shape." + fields: + db: + type: string + coll: + type: string + optional: true + uuid: + type: string + optional: true + tenantId: + type: string + optional: true +
\ No newline at end of file diff --git a/src/mongo/db/query/query_shape/query_shape_test.cpp b/src/mongo/db/query/query_shape/query_shape_test.cpp new file mode 100644 index 00000000000..d6185b5c5cb --- /dev/null +++ b/src/mongo/db/query/query_shape/query_shape_test.cpp @@ -0,0 +1,767 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/bson/bsonmisc.h" +#include "mongo/bson/bsonobj.h" +#include "mongo/db/matcher/expression_geo.h" +#include "mongo/db/matcher/extensions_callback_real.h" +#include "mongo/db/matcher/parsed_match_expression_for_test.h" +#include "mongo/db/namespace_string.h" +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_shape/query_shape_test_gen.h" +#include "mongo/db/query/query_shape/serialization_options.h" +#include "mongo/db/query/query_shape/shape_helpers.h" +#include "mongo/db/service_context_test_fixture.h" +#include "mongo/unittest/bson_test_util.h" +#include "mongo/unittest/unittest.h" + +namespace mongo::query_shape { + +namespace { +BSONObj predicateShape(const MatchExpression* expr) { + return expr->serialize(SerializationOptions::kDebugQueryShapeSerializeOptions); +} +BSONObj predicateShape(std::string filterJson) { + return predicateShape(ParsedMatchExpressionForTest(filterJson).get()); +} + +BSONObj predicateShapeRedacted(const MatchExpression* expr) { + return expr->serialize(SerializationOptions::kDebugShapeAndMarkIdentifiers_FOR_TEST); +} +BSONObj predicateShapeRedacted(std::string filterJson) { + return predicateShapeRedacted(ParsedMatchExpressionForTest(filterJson).get()); +} + +// TODO SERVER-87736 There is no 'auto' here, make that more clear. +#define ASSERT_SHAPE_EQ_AUTO(expected, actual) \ + ASSERT_BSONOBJ_EQ_AUTO(expected, predicateShape(actual)) + +#define ASSERT_REDACTED_SHAPE_EQ_AUTO(expected, actual) \ + ASSERT_BSONOBJ_EQ_AUTO(expected, predicateShapeRedacted(actual)) + + +TEST(QueryPredicateShape, Equals) { + ASSERT_SHAPE_EQ_AUTO( // Implicit equals + R"({"a":{"$eq":"?number"}})", + "{a: 5}"); + ASSERT_SHAPE_EQ_AUTO( // Explicit equals + R"({"a":{"$eq":"?number"}})", + "{a: {$eq: 5}}"); + ASSERT_SHAPE_EQ_AUTO( // implicit $and + R"({"$and":[{"a":{"$eq":"?number"}},{"b":{"$eq":"?number"}}]})", + "{a: 5, b: 6}"); + ASSERT_REDACTED_SHAPE_EQ_AUTO( // Implicit equals + R"({"HASH<a>":{"$eq":"?number"}})", + "{a: 5}"); + ASSERT_REDACTED_SHAPE_EQ_AUTO( // Explicit equals + R"({"HASH<a>":{"$eq":"?number"}})", + "{a: {$eq: 5}}"); + ASSERT_REDACTED_SHAPE_EQ_AUTO( // NOLINT + R"({"$and":[{"HASH<a>":{"$eq":"?number"}},{"HASH<b>":{"$eq":"?number"}}]})", + "{a: 5, b: 6}"); + ASSERT_REDACTED_SHAPE_EQ_AUTO( // NOLINT + R"({"HASH<foo>.HASH<$bar>":{"$eq":"?number"}})", + R"({"foo.$bar":0})"); +} + +TEST(QueryPredicateShape, ArraySubTypes) { + ASSERT_SHAPE_EQ_AUTO( // NOLINT + "{a: {$eq: '[]'}}", + "{a: []}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + "{a: {$eq: '?array<?number>'}}", + "{a: [2]}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$eq":"?array<?number>"}})", + "{a: [2, 3]}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$eq":"?array<?object>"}})", + "{a: [{}]}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$eq":"?array<?object>"}})", + "{a: [{}, {}]}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$eq":"?array<?array>"}})", + "{a: [[], [], []]}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$eq":"?array<?array>"}})", + "{a: [[2, 3], ['string'], []]}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$eq":"?array<>"}})", + "{a: [{}, 2]}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$eq":"?array<>"}})", + "{a: [[], 2]}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$eq":"?array<>"}})", + "{a: [[{}, 'string'], 2]}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$eq":"?array<>"}})", + "{a: [[{}, 'string'], 2]}"); +} + +TEST(QueryPredicateShape, Comparisons) { + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({ + "$and": [ + { + "a": { + "$lt": "?number" + } + }, + { + "b": { + "$gt": "?number" + } + }, + { + "c": { + "$gte": "?number" + } + }, + { + "c": { + "$lte": "?number" + } + } + ] + })", + "{a: {$lt: 5}, b: {$gt: 6}, c: {$gte: 3, $lte: 10}}"); +} + +namespace { +void assertShapeIs(std::string filterJson, BSONObj expectedShape) { + ASSERT_BSONOBJ_EQ(expectedShape, predicateShape(filterJson)); +} + +void assertRedactedShapeIs(std::string filterJson, BSONObj expectedShape) { + ASSERT_BSONOBJ_EQ(expectedShape, predicateShapeRedacted(filterJson)); +} +} // namespace + +TEST(QueryPredicateShape, Regex) { + // Note/warning: 'fromjson' will parse $regex into a /regex/, so these tests can't use + // auto-updating BSON assertions. + assertShapeIs("{a: /a+/}", + BSON("a" << BSON("$regex" + << "?string"))); + assertShapeIs("{a: /a+/i}", + BSON("a" << BSON("$regex" + << "?string" + << "$options" + << "?string"))); + assertRedactedShapeIs("{a: /a+/}", + BSON("HASH<a>" << BSON("$regex" + << "?string"))); + assertRedactedShapeIs("{a: /a+/}", + BSON("HASH<a>" << BSON("$regex" + << "?string"))); +} + +TEST(QueryPredicateShape, Mod) { + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$mod":["?number","?number"]}})", + "{a: {$mod: [2, 0]}}"); +} + +TEST(QueryPredicateShape, Exists) { + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$exists":"?bool"}})", + "{a: {$exists: true}}"); +} + +TEST(QueryPredicateShape, In) { + // Any number of children in any order is always the same shape + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$in":"?array<?number>"}})", + "{a: {$in: [1]}}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$in":"?array<>"}})", + "{a: {$in: [1, 4, 'str', /regex/]}}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$in":"?array<>"}})", + "{a: {$in: ['str', /regex/, 1, 4]}}"); +} + +TEST(QueryPredicateShape, BitTestOperators) { + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$bitsAllSet":"?array<?number>"}})", + "{a: {$bitsAllSet: [1, 5]}}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$bitsAllSet":"?array<?number>"}})", + "{a: {$bitsAllSet: 50}}"); + + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$bitsAnySet":"?array<?number>"}})", + "{a: {$bitsAnySet: [1, 5]}}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$bitsAnySet":"?array<?number>"}})", + "{a: {$bitsAnySet: 50}}"); + + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$bitsAllClear":"?array<?number>"}})", + "{a: {$bitsAllClear: [1, 5]}}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$bitsAllClear":"?array<?number>"}})", + "{a: {$bitsAllClear: 50}}"); + + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$bitsAnyClear":"?array<?number>"}})", + "{a: {$bitsAnyClear: [1, 5]}}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$bitsAnyClear":"?array<?number>"}})", + "{a: {$bitsAnyClear: 50}}"); +} + +TEST(QueryPredicateShape, AlwaysBoolean) { + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"$alwaysTrue":"?number"})", + "{$alwaysTrue: 1}"); + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"$alwaysFalse":"?number"})", + "{$alwaysFalse: 1}"); +} + +TEST(QueryPredicateShape, And) { + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({ + "$and": [ + { + "a": { + "$lt": "?number" + } + }, + { + "b": { + "$gte": "?number" + } + }, + { + "c": { + "$lte": "?number" + } + } + ] + })", + "{$and: [{a: {$lt: 5}}, {b: {$gte: 3}}, {c: {$lte: 10}}]}"); +} + +TEST(QueryPredicateShape, Or) { + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({ + "$or": [ + { + "a": { + "$eq": "?number" + } + }, + { + "b": { + "$in": "?array<?number>" + } + }, + { + "c": { + "$gt": "?number" + } + } + ] + })", + "{$or: [{a: 5}, {b: {$in: [1,2,3]}}, {c: {$gt: 10}}]}"); +} + +TEST(QueryPredicateShape, ElemMatch) { + // ElemMatchObjectMatchExpression + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({ + "a": { + "$elemMatch": { + "$and": [ + { + "b": { + "$eq": "?number" + } + }, + { + "c": { + "$exists": "?bool" + } + } + ] + } + } + })", + "{a: {$elemMatch: {b: 5, c: {$exists: true}}}}"); + + // ElemMatchValueMatchExpression + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"a":{"$elemMatch":{"$gt":"?number","$lt":"?number"}}})", + "{a: {$elemMatch: {$gt: 5, $lt: 10}}}"); + + // Nested + ASSERT_REDACTED_SHAPE_EQ_AUTO( // NOLINT + R"({ + "HASH<a>": { + "$elemMatch": { + "$elemMatch": { + "$gt": "?number", + "$lt": "?number" + } + } + } + })", + "{a: {$elemMatch: {$elemMatch: {$gt: 5, $lt: 10}}}}"); +} + +TEST(QueryPredicateShape, InternalBucketGeoWithinMatchExpression) { + auto query = + "{ $_internalBucketGeoWithin: {withinRegion: {$centerSphere: [[0, 0], 10]}, field: " + "\"a\"} " + "}"; + ASSERT_REDACTED_SHAPE_EQ_AUTO( // NOLINT + R"({ + "$_internalBucketGeoWithin": { + "withinRegion": { + "$centerSphere": "?array<>" + }, + "field": "HASH<a>" + } + })", + query); +} + +TEST(QueryPredicateShape, NorMatchExpression) { + ASSERT_REDACTED_SHAPE_EQ_AUTO( // NOLINT + R"({"$nor":[{"HASH<a>":{"$lt":"?number"}},{"HASH<b>":{"$gt":"?number"}}]})", + "{ $nor: [ { a: {$lt: 5} }, { b: {$gt: 4} } ] }"); +} + +TEST(QueryPredicateShape, NotMatchExpression) { + ASSERT_REDACTED_SHAPE_EQ_AUTO( // NOLINT + R"({"HASH<price>":{"$not":{"$gt":"?number"}}})", + "{ price: { $not: { $gt: 1.99 } } }"); + // Test the special case where NotMatchExpression::serialize() reduces to $alwaysFalse. + auto emptyAnd = std::make_unique<AndMatchExpression>(); + const MatchExpression& notExpr = NotMatchExpression(std::move(emptyAnd)); + auto serialized = + notExpr.serialize(SerializationOptions::kDebugShapeAndMarkIdentifiers_FOR_TEST); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({"$alwaysFalse":"?number"})", + serialized); +} + +TEST(QueryPredicateShape, SizeMatchExpression) { + ASSERT_REDACTED_SHAPE_EQ_AUTO( // NOLINT + R"({"HASH<price>":{"$size":"?number"}})", + "{ price: { $size: 2 } }"); +} + +TEST(QueryPredicateShape, TextMatchExpression) { + TextMatchExpressionBase::TextParams params = {"coffee"}; + auto expr = ExtensionsCallbackNoop().createText(params); + auto literalAndFieldRedactOpts = SerializationOptions::kDebugShapeAndMarkIdentifiers_FOR_TEST; + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "$text": { + "$search": "?string", + "$language": "?string", + "$caseSensitive": "?bool", + "$diacriticSensitive": "?bool" + } + })", + expr->serialize(SerializationOptions::kDebugShapeAndMarkIdentifiers_FOR_TEST)); +} + +TEST(QueryPredicateShape, TwoDPtInAnnulusExpression) { + const MatchExpression& expr = TwoDPtInAnnulusExpression({}, {}); + auto literalAndFieldRedactOpts = SerializationOptions::kDebugShapeAndMarkIdentifiers_FOR_TEST; + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({"$TwoDPtInAnnulusExpression":true})", + expr.serialize(SerializationOptions::kDebugShapeAndMarkIdentifiers_FOR_TEST)); +} + +TEST(QueryPredicateShape, WhereMatchExpression) { + ASSERT_SHAPE_EQ_AUTO( // NOLINT + R"({"$where":"?javascript"})", + "{$where: \"some_code()\"}"); +} + +BSONObj queryShapeForOptimizedExprExpression(std::string exprPredicateJson) { + ParsedMatchExpressionForTest expr(exprPredicateJson); + // We need to optimize an $expr expression in order to generate an $_internalExprEq. It's + // not clear we'd want to do optimization before computing the query shape, but we should + // support the computation on any MatchExpression, and this is the easiest way we can create + // this type of MatchExpression node. + auto optimized = MatchExpression::optimize(expr.release()); + return predicateShape(optimized.get()); +} + +TEST(QueryPredicateShape, OptimizedExprPredicates) { + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "$and": [ + { + "a": { + "$_internalExprEq": "?number" + } + }, + { + "$expr": { + "$eq": [ + "$a", + "?number" + ] + } + } + ] + })", + queryShapeForOptimizedExprExpression("{$expr: {$eq: ['$a', 2]}}")); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "$and": [ + { + "a": { + "$_internalExprLt": "?number" + } + }, + { + "$expr": { + "$lt": [ + "$a", + "?number" + ] + } + } + ] + })", + queryShapeForOptimizedExprExpression("{$expr: {$lt: ['$a', 2]}}")); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "$and": [ + { + "a": { + "$_internalExprLte": "?number" + } + }, + { + "$expr": { + "$lte": [ + "$a", + "?number" + ] + } + } + ] + })", + queryShapeForOptimizedExprExpression("{$expr: {$lte: ['$a', 2]}}")); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "$and": [ + { + "a": { + "$_internalExprGt": "?number" + } + }, + { + "$expr": { + "$gt": [ + "$a", + "?number" + ] + } + } + ] + })", + queryShapeForOptimizedExprExpression("{$expr: {$gt: ['$a', 2]}}")); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "$and": [ + { + "a": { + "$_internalExprGte": "?number" + } + }, + { + "$expr": { + "$gte": [ + "$a", + "?number" + ] + } + } + ] + })", + queryShapeForOptimizedExprExpression("{$expr: {$gte: ['$a', 2]}}")); +} + +TEST(QueryShapeIDL, ShapifyIDLStruct) { + SerializationOptions options; + options.transformIdentifiers = true; + options.transformIdentifiersCallback = [](StringData s) -> std::string { + return str::stream() << "HASH<" << s << ">"; + }; + options.literalPolicy = LiteralSerializationPolicy::kToDebugTypeString; + + auto nested = NestedStruct("value", + ExampleEnumEnum::Value1, + "hello", + {1, 2, 3, 4}, + "field.path", + {"field.path.1", "fieldpath2"}, + NamespaceString{"db", "coll"}, + NamespaceString{"db", "coll"}, + 177, + true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "stringField": "value", + "enumField": "EnumValue1", + "stringIntVariantEnum": "hello", + "arrayOfInts": [ + 1, + 2, + 3, + 4 + ], + "fieldpath": "field.path", + "fieldpathList": [ + "field.path.1", + "fieldpath2" + ], + "nss": "db.coll", + "plainNss": "db.coll", + "safeInt64Field": 177, + "boolField": true + })", + nested.toBSON()); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "stringField": "?string", + "enumField": "EnumValue1", + "stringIntVariantEnum": "hello", + "arrayOfInts": "?array<?number>", + "fieldpath": "HASH<field>.HASH<path>", + "fieldpathList": [ + "HASH<field>.HASH<path>.HASH<1>", + "HASH<fieldpath2>" + ], + "nss": "HASH<db.coll>", + "plainNss": "db.coll", + "safeInt64Field": "?number", + "boolField": "?bool" + })", + nested.toBSON(options)); + + + auto parent = ParentStruct(nested, nested); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "nested_shape": { + "stringField": "value", + "enumField": "EnumValue1", + "stringIntVariantEnum": "hello", + "arrayOfInts": [ + 1, + 2, + 3, + 4 + ], + "fieldpath": "field.path", + "fieldpathList": [ + "field.path.1", + "fieldpath2" + ], + "nss": "db.coll", + "plainNss": "db.coll", + "safeInt64Field": 177, + "boolField": true + }, + "nested_no_shape": { + "stringField": "value", + "enumField": "EnumValue1", + "stringIntVariantEnum": "hello", + "arrayOfInts": [ + 1, + 2, + 3, + 4 + ], + "fieldpath": "field.path", + "fieldpathList": [ + "field.path.1", + "fieldpath2" + ], + "nss": "db.coll", + "plainNss": "db.coll", + "safeInt64Field": 177, + "boolField": true + } + })", + parent.toBSON()); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "nested_shape": { + "stringField": "?string", + "enumField": "EnumValue1", + "stringIntVariantEnum": "hello", + "arrayOfInts": "?array<?number>", + "fieldpath": "HASH<field>.HASH<path>", + "fieldpathList": [ + "HASH<field>.HASH<path>.HASH<1>", + "HASH<fieldpath2>" + ], + "nss": "HASH<db.coll>", + "plainNss": "db.coll", + "safeInt64Field": "?number", + "boolField": "?bool" + }, + "nested_no_shape": { + "stringField": "value", + "enumField": "EnumValue1", + "stringIntVariantEnum": "hello", + "arrayOfInts": [ + 1, + 2, + 3, + 4 + ], + "fieldpath": "field.path", + "fieldpathList": [ + "field.path.1", + "fieldpath2" + ], + "nss": "db.coll", + "plainNss": "db.coll", + "safeInt64Field": 177, + "boolField": true + } + })", + parent.toBSON(options)); +} + +} // namespace + +namespace { + +static const NamespaceString kDefaultTestNss = NamespaceString("testDB.testColl"); + +struct DummyShapeSpecificComponents : public query_shape::CmdSpecificShapeComponents { + DummyShapeSpecificComponents(){}; + void HashValue(absl::HashState state) const {} + size_t size() const final { + return sizeof(DummyShapeSpecificComponents); + } +}; + +class DummyShape : public Shape { +public: + DummyShape(NamespaceStringOrUUID nssOrUUID, + BSONObj collation, + DummyShapeSpecificComponents dummyComponents) + : Shape(nssOrUUID, collation) { + components = dummyComponents; + } + + const CmdSpecificShapeComponents& specificComponents() const final { + return components; + } + + void appendCmdSpecificShapeComponents(BSONObjBuilder&, + OperationContext*, + const SerializationOptions& opts) const final {} + DummyShapeSpecificComponents components; +}; + +class DummyShapeWithExtraSize : public Shape { +public: + DummyShapeWithExtraSize(NamespaceStringOrUUID nssOrUUID, + BSONObj collation, + DummyShapeSpecificComponents dummyComponents) + : Shape(nssOrUUID, collation) { + components = dummyComponents; + } + + const CmdSpecificShapeComponents& specificComponents() const final { + return components; + } + + // Random number for testing purposes. + size_t extraSize() const final override { + return 125; + } + void appendCmdSpecificShapeComponents(BSONObjBuilder&, + OperationContext*, + const SerializationOptions& opts) const final {} + + DummyShapeSpecificComponents components; +}; + +class UniversalShapeTest : public ServiceContextTest {}; + +TEST_F(UniversalShapeTest, SizeOfSpecificComponents) { + auto innerComponents = std::make_unique<DummyShapeSpecificComponents>(); + ASSERT_EQ(innerComponents->size(), sizeof(CmdSpecificShapeComponents)); + ASSERT_EQ(innerComponents->size(), sizeof(void*) /*vtable ptr*/); +} + +TEST_F(UniversalShapeTest, SizeOfShape) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + + // Make shape for testing. + auto collation = BSONObj{}; + auto innerComponents = std::make_unique<DummyShapeSpecificComponents>(); + auto shape = std::make_unique<DummyShape>(kDefaultTestNss, collation, *innerComponents); + + ASSERT_EQ(innerComponents->size(), shape->specificComponents().size()); + ASSERT_EQ(shape->size(), + sizeof(NamespaceStringOrUUID) + sizeof(BSONObj) + sizeof(void*) /*vtable ptr*/ + + shape->specificComponents().size() + static_cast<size_t>(collation.objsize())); +} + +TEST_F(UniversalShapeTest, SizeOfShapeWithExtraSize) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + + // Make shape for testing. + auto collation = BSONObj{}; + auto innerComponents = std::make_unique<DummyShapeSpecificComponents>(); + auto shape = std::make_unique<DummyShape>(kDefaultTestNss, collation, *innerComponents); + auto shapeWithExtraSize = + std::make_unique<DummyShapeWithExtraSize>(kDefaultTestNss, collation, *innerComponents); + + ASSERT_EQ(shapeWithExtraSize->size(), shape->size() + shapeWithExtraSize->extraSize()); +} +} // namespace +} // namespace mongo::query_shape diff --git a/src/mongo/db/query/query_shape/query_shape_test.idl b/src/mongo/db/query/query_shape/query_shape_test.idl new file mode 100644 index 00000000000..06efb7ed1ef --- /dev/null +++ b/src/mongo/db/query/query_shape/query_shape_test.idl @@ -0,0 +1,91 @@ +# Copyright (C) 2023-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. +# + +global: + cpp_namespace: "mongo" + +imports: + - "mongo/idl/basic_types.idl" + +enums: + ExampleEnum: + description: "" + type: string + values: + Value1: "EnumValue1" + Value2: "EnumValue2" + +structs: + NestedStruct: + query_shape_component: true + strict: true + description: "" + fields: + stringField: + query_shape: literal + type: string + enumField: + query_shape: parameter + type: ExampleEnum + stringIntVariantEnum: + query_shape: parameter + type: + variant: [string, int] + arrayOfInts: + query_shape: literal + type: array<int> + fieldpath: + query_shape: anonymize + type: string + fieldpathList: + query_shape: anonymize + type: array<string> + nss: + query_shape: custom + type: namespacestring + plainNss: + query_shape: parameter + type: namespacestring + safeInt64Field: + query_shape: literal + type: safeInt64 + boolField: + query_shape: literal + type: bool + + ParentStruct: + query_shape_component: true + strict: true + description: "" + fields: + nested_shape: + query_shape: literal + type: NestedStruct + nested_no_shape: + query_shape: parameter + type: NestedStruct diff --git a/src/mongo/db/query/query_shape/serialization_options.cpp b/src/mongo/db/query/query_shape/serialization_options.cpp new file mode 100644 index 00000000000..e6008f8579b --- /dev/null +++ b/src/mongo/db/query/query_shape/serialization_options.cpp @@ -0,0 +1,515 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "serialization_options.h" +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery + +#include "mongo/db/query/query_shape/serialization_options.h" + +#include <boost/optional.hpp> +#include <string> + +#include "mongo/base/string_data.h" +#include "mongo/bson/timestamp.h" +#include "mongo/db/exec/document_value/document.h" +#include "mongo/db/exec/document_value/value.h" +#include "mongo/db/query/explain_options.h" +#include "mongo/logv2/log.h" +#include "mongo/util/assert_util.h" + +namespace mongo { + +namespace { + +// We'll pre-declare all of these strings so that we can avoid the allocations when we reference +// them later. +static constexpr StringData kUndefinedTypeString = "?undefined"_sd; +static constexpr StringData kStringTypeString = "?string"_sd; +static constexpr StringData kNumberTypeString = "?number"_sd; +static constexpr StringData kMinKeyTypeString = "?minKey"_sd; +static constexpr StringData kObjectTypeString = "?object"_sd; +static constexpr StringData kArrayTypeString = "?array"_sd; +static constexpr StringData kBinDataTypeString = "?binData"_sd; +static constexpr StringData kObjectIdTypeString = "?objectId"_sd; +static constexpr StringData kBoolTypeString = "?bool"_sd; +static constexpr StringData kDateTypeString = "?date"_sd; +static constexpr StringData kNullTypeString = "?null"_sd; +static constexpr StringData kRegexTypeString = "?regex"_sd; +static constexpr StringData kDbPointerTypeString = "?dbPointer"_sd; +static constexpr StringData kJavascriptTypeString = "?javascript"_sd; +static constexpr StringData kJavascriptWithScopeTypeString = "?javascriptWithScope"_sd; +static constexpr StringData kTimestampTypeString = "?timestamp"_sd; +static constexpr StringData kMaxKeyTypeString = "?maxKey"_sd; + +static const StringMap<StringData> kArrayTypeStringConstants{ + {kUndefinedTypeString.rawData(), "?array<?undefined>"_sd}, + {kStringTypeString.rawData(), "?array<?string>"_sd}, + {kNumberTypeString.rawData(), "?array<?number>"_sd}, + {kMinKeyTypeString.rawData(), "?array<?minKey>"_sd}, + {kObjectTypeString.rawData(), "?array<?object>"_sd}, + {kArrayTypeString.rawData(), "?array<?array>"_sd}, + {kBinDataTypeString.rawData(), "?array<?binData>"_sd}, + {kObjectIdTypeString.rawData(), "?array<?objectId>"_sd}, + {kBoolTypeString.rawData(), "?array<?bool>"_sd}, + {kDateTypeString.rawData(), "?array<?date>"_sd}, + {kNullTypeString.rawData(), "?array<?null>"_sd}, + {kRegexTypeString.rawData(), "?array<?regex>"_sd}, + {kDbPointerTypeString.rawData(), "?array<?dbPointer>"_sd}, + {kJavascriptTypeString.rawData(), "?array<?javascript>"_sd}, + {kJavascriptWithScopeTypeString.rawData(), "?array<?javascriptWithScope>"_sd}, + {kTimestampTypeString.rawData(), "?array<?timestamp>"_sd}, + {kMaxKeyTypeString.rawData(), "?array<?maxKey>"_sd}, +}; + +static constexpr auto kRepresentativeString = "?"_sd; +static constexpr auto kRepresentativeNumber = 1; +static const auto kRepresentativeObject = BSON("?" + << "?"); +static const auto kRepresentativeArray = BSONArray(); +static constexpr auto kRepresentativeBinData = BSONBinData(); +static const auto kRepresentativeObjectId = OID::max(); +static constexpr auto kRepresentativeBool = true; +static const auto kRepresentativeDate = Date_t::fromMillisSinceEpoch(0); +static const auto kRepresentativeRegex = BSONRegEx("/\?/"); +static const auto kRepresentativeDbPointer = BSONDBRef("?.?", OID::max()); +static const auto kRepresentativeJavascript = BSONCode("return ?;"); +static const auto kRepresentativeJavascriptWithScope = BSONCodeWScope("return ?;", BSONObj()); +static const auto kRepresentativeTimestamp = Timestamp::min(); + +/** + * A default redaction strategy that generates easy to check results for testing purposes. + */ +std::string applyHmacForTest(StringData s) { + // Avoid ending in a parenthesis since the results will occur in a raw string where the )" + // sequence will accidentally terminate the string. + return str::stream() << "HASH<" << s << ">"; +} + +/** + * Computes a debug string meant to represent "any value of type t", where "t" is the type of the + * provided argument. For example "?number" for any number (int, double, etc.). + */ +StringData debugTypeString(BSONType t) { + // This is tightly coupled with 'canonicalizeBSONType' and therefore also with + // sorting/comparison semantics. + switch (t) { + case EOO: + case Undefined: + return kUndefinedTypeString; + case Symbol: + case String: + return kStringTypeString; + case NumberInt: + case NumberLong: + case NumberDouble: + case NumberDecimal: + return kNumberTypeString; + case MinKey: + return kMinKeyTypeString; + case Object: + return kObjectTypeString; + case Array: + // This case should only happen if we have an array within an array. + return kArrayTypeString; + case BinData: + return kBinDataTypeString; + case jstOID: + return kObjectIdTypeString; + case Bool: + return kBoolTypeString; + case Date: + return kDateTypeString; + case jstNULL: + return kNullTypeString; + case RegEx: + return kRegexTypeString; + case DBRef: + return kDbPointerTypeString; + case Code: + return kJavascriptTypeString; + case CodeWScope: + return kJavascriptWithScopeTypeString; + case bsonTimestamp: + return kTimestampTypeString; + case MaxKey: + return kMaxKeyTypeString; + default: + MONGO_UNREACHABLE_TASSERT(7539806); + } +} + +/** + * Returns an arbitrary value of the same type as the one given. For any number, this will be the + * number 1. For any boolean this will be true. + * TODO if you need a different value to make sure it will parse, you should not use this API. + */ +ImplicitValue defaultLiteralOfType(BSONType t) { + // This is tightly coupled with 'canonicalizeBSONType' and therefore also with + // sorting/comparison semantics. + switch (t) { + case EOO: + case Undefined: + return BSONUndefined; + case Symbol: + case String: + return kRepresentativeString; + case NumberInt: + case NumberLong: + case NumberDouble: + case NumberDecimal: + return kRepresentativeNumber; + case MinKey: + return MINKEY; + case Object: + return kRepresentativeObject; + case Array: + // This case should only happen if we have an array within an array. + return kRepresentativeArray; + case BinData: + return kRepresentativeBinData; + case jstOID: + return kRepresentativeObjectId; + case Bool: + return kRepresentativeBool; + case Date: + return kRepresentativeDate; + case jstNULL: + return BSONNULL; + case RegEx: + return kRepresentativeRegex; + case DBRef: + return kRepresentativeDbPointer; + case Code: + return kRepresentativeJavascript; + case CodeWScope: + return kRepresentativeJavascriptWithScope; + case bsonTimestamp: + return kRepresentativeTimestamp; + case MaxKey: + return MAXKEY; + default: + MONGO_UNREACHABLE_TASSERT(7539803); + } +} + +/** + * A struct representing the sub-type information for an array. + */ +struct ArraySubtypeInfo { + /** + * Whether the values of an array are all the same BSON type or not (mixed). + */ + enum class NTypes { kEmpty, kOneType, kMixed }; + ArraySubtypeInfo(NTypes nTypes_) : nTypes(nTypes_) {} + ArraySubtypeInfo(BSONType oneType) : nTypes(NTypes::kOneType), singleType(oneType) {} + + NTypes nTypes; + boost::optional<BSONType> singleType = boost::none; +}; + +template <typename ValueType> +using GetTypeFn = std::function<BSONType(ValueType)>; + +static GetTypeFn<BSONElement> getBSONElementType = [](const BSONElement& e) { return e.type(); }; +static GetTypeFn<Value> getValueType = [](const Value& v) { return v.getType(); }; + +/** + * Scans 'arrayOfValues' to see if all values are of the same type or not. Returns this info in a + * struct - see the struct definition for how it is represented. + * + * Templated algorithm to handle both iterators of BSONElements or iterators of Values. + * 'getTypeCallback' is provided to abstract away the different '.type()' vs '.getType()' APIs. + */ +template <typename ArrayType, typename ValueType> +ArraySubtypeInfo determineArraySubType(const ArrayType& arrayOfValues, + GetTypeFn<ValueType> getTypeCallback) { + boost::optional<BSONType> firstType = boost::none; + for (auto&& v : arrayOfValues) { + if (!firstType) { + firstType.emplace(getTypeCallback(v)); + } else if (*firstType != getTypeCallback(v)) { + return {ArraySubtypeInfo::NTypes::kMixed}; + } + } + return firstType ? ArraySubtypeInfo{*firstType} + : ArraySubtypeInfo{ArraySubtypeInfo::NTypes::kEmpty}; +} + +ArraySubtypeInfo determineArraySubType(const BSONObj& arrayAsObj) { + return determineArraySubType<BSONObj, BSONElement>(arrayAsObj, getBSONElementType); +} +ArraySubtypeInfo determineArraySubType(const std::vector<Value>& values) { + return determineArraySubType<std::vector<Value>, Value>(values, getValueType); +} + +template <typename ValueType> +StringData debugTypeString( + const ValueType& v, + GetTypeFn<ValueType> getTypeCallback, + std::function<ArraySubtypeInfo(ValueType)> determineArraySubTypeCallback) { + if (getTypeCallback(v) == BSONType::Array) { + // Iterating the array as .Obj(), as if it were a BSONObj (with field names '0', '1', etc.) + // is faster than converting the whole thing to an array which would force a copy. + auto typeInfo = determineArraySubTypeCallback(v); + switch (typeInfo.nTypes) { + case ArraySubtypeInfo::NTypes::kEmpty: + return "[]"_sd; + case ArraySubtypeInfo::NTypes::kOneType: + return kArrayTypeStringConstants.at(debugTypeString(*typeInfo.singleType)); + case ArraySubtypeInfo::NTypes::kMixed: + return "?array<>"; + default: + MONGO_UNREACHABLE_TASSERT(7539801); + } + } + return debugTypeString(getTypeCallback(v)); +} + +template <typename ValueType> +ImplicitValue defaultLiteralOfType( + const ValueType& v, + GetTypeFn<ValueType> getTypeCallback, + std::function<ArraySubtypeInfo(ValueType)> determineArraySubTypeCallback) { + if (getTypeCallback(v) == BSONType::Array) { + auto typeInfo = determineArraySubTypeCallback(v); + switch (typeInfo.nTypes) { + case ArraySubtypeInfo::NTypes::kEmpty: + return BSONArray(); + case ArraySubtypeInfo::NTypes::kOneType: + return std::vector<Value>{defaultLiteralOfType(*typeInfo.singleType)}; + case ArraySubtypeInfo::NTypes::kMixed: + // We don't care which types, we'll use a number and a string as the canonical + // mixed type array regardless. This is to ensure we don't get 2^N possibilities + // for mixed type scenarios - we wish to collapse all "mixed type" arrays to one + // canonical mix. The choice of int and string is mostly arbitrary - hopefully + // somewhat comprehensible at a glance. + return std::vector<Value>{Value(2), Value("or more types"_sd)}; + default: + MONGO_UNREACHABLE_TASSERT(7539805); + } + } + return defaultLiteralOfType(getTypeCallback(v)); +} + +ArraySubtypeInfo getSubTypeFromBSONElemArray(BSONElement arrayElem) { + // Iterating the array as .Obj(), as if it were a BSONObj (with field names '0', '1', etc.) + // is faster than converting the whole thing to an array which would force a copy. + return determineArraySubType(arrayElem.Obj()); +} +ArraySubtypeInfo getSubTypeFromValueArray(const Value& arrayVal) { + return determineArraySubType(arrayVal.getArray()); +} + +void appendDefaultOfNonArrayType(BSONObjBuilder* bob, StringData name, const BSONElement& e) { + switch (e.type()) { + case EOO: + case Undefined: + bob->appendUndefined(name); + return; + case Symbol: + case String: + bob->append(name, kRepresentativeString); + return; + case NumberInt: + case NumberLong: + case NumberDouble: + case NumberDecimal: + bob->append(name, kRepresentativeNumber); + return; + case MinKey: + bob->appendMinKey(name); + return; + case Object: + bob->append(name, kRepresentativeObject); + return; + case Array: + // This case is more complicated and callers should use a more generic helper. + MONGO_UNREACHABLE_TASSERT(8094100); + case BinData: + bob->append(name, kRepresentativeBinData); + return; + case jstOID: + bob->append(name, kRepresentativeObjectId); + return; + case Bool: + bob->append(name, kRepresentativeBool); + return; + case Date: + bob->append(name, kRepresentativeDate); + return; + case jstNULL: + bob->appendNull(name); + return; + case RegEx: + bob->append(name, kRepresentativeRegex); + return; + case DBRef: + bob->append(name, kRepresentativeDbPointer); + return; + case Code: + bob->append(name, kRepresentativeJavascript); + return; + case CodeWScope: + bob->append(name, kRepresentativeJavascriptWithScope); + return; + case bsonTimestamp: + bob->append(name, kRepresentativeTimestamp); + return; + case MaxKey: + bob->appendMaxKey(name); + return; + default: + MONGO_UNREACHABLE_TASSERT(8094101); + }; +} +} // namespace + +const SerializationOptions SerializationOptions::kRepresentativeQueryShapeSerializeOptions = + SerializationOptions{LiteralSerializationPolicy::kToRepresentativeParseableValue}; + +const SerializationOptions SerializationOptions::kDebugQueryShapeSerializeOptions = + SerializationOptions{LiteralSerializationPolicy::kToDebugTypeString}; + +SerializationOptions::SerializationOptions(LiteralSerializationPolicy policy) + : literalPolicy(policy) {} +SerializationOptions::SerializationOptions( + boost::optional<ExplainOptions::Verbosity> explainVerbosity) + : verbosity(explainVerbosity) {} + +SerializationOptions::SerializationOptions(LiteralSerializationPolicy policy, + bool transformIdentifiers, + TokenizeIdentifierFunc transformIdentifiersCallbackFn) + : literalPolicy(policy), + transformIdentifiers(transformIdentifiers), + transformIdentifiersCallback(transformIdentifiersCallbackFn) {} + +const SerializationOptions SerializationOptions::kMarkIdentifiers_FOR_TEST{ + LiteralSerializationPolicy::kUnchanged, true, applyHmacForTest}; + +const SerializationOptions SerializationOptions::kDebugShapeAndMarkIdentifiers_FOR_TEST{ + LiteralSerializationPolicy::kToDebugTypeString, true, applyHmacForTest}; + +// Overloads for BSONElem and Value. +StringData debugTypeString(BSONElement e) { + return debugTypeString<BSONElement>(e, getBSONElementType, getSubTypeFromBSONElemArray); +} +StringData debugTypeString(const Value& v) { + return debugTypeString<Value>(v, getValueType, getSubTypeFromValueArray); +} + +// Overloads for BSONElem and Value. +ImplicitValue defaultLiteralOfType(const Value& v) { + return defaultLiteralOfType<Value>(v, getValueType, getSubTypeFromValueArray); +} +ImplicitValue defaultLiteralOfType(BSONElement e) { + return defaultLiteralOfType<BSONElement>(e, getBSONElementType, getSubTypeFromBSONElemArray); +} + +void SerializationOptions::appendLiteral(BSONObjBuilder* bob, const BSONElement& e) const { + appendLiteral(bob, e.fieldNameStringData(), e); +} +void SerializationOptions::appendLiteral(BSONObjBuilder* bob, + StringData name, + const BSONElement& e) const { + // The first two cases are particularly performance sensitive. We could answer everything here + // with the code inside the 'kToDebugTypeString' branch, but there are some relatively easy ways + // to accomplish the first two policy cases (in the common cases), so we'll special case those + // in order to avoid constructing a temporary Value. + switch (literalPolicy) { + case LiteralSerializationPolicy::kUnchanged: + bob->appendAs(e, name); + return; + case LiteralSerializationPolicy::kToRepresentativeParseableValue: { + if (e.type() != BSONType::Array) { + appendDefaultOfNonArrayType(bob, name, e); + return; + } + // If it's an array we'll default to the slow but general codepath below. + [[fallthrough]]; + } + case LiteralSerializationPolicy::kToDebugTypeString: { + // Performance isn't as sensitive here. + return serializeLiteral(e).addToBsonObj(bob, name); + } + default: + MONGO_UNREACHABLE_TASSERT(8094102); + } +} + +void SerializationOptions::appendLiteral(BSONObjBuilder* bob, + StringData fieldName, + const ImplicitValue& v, + const boost::optional<Value>& representativeValue) const { + serializeLiteral(v, representativeValue).addToBsonObj(bob, fieldName); +} + +Value SerializationOptions::serializeLiteral( + const BSONElement& e, const boost::optional<Value>& representativeValue) const { + switch (literalPolicy) { + case LiteralSerializationPolicy::kUnchanged: + return Value(e); + case LiteralSerializationPolicy::kToDebugTypeString: + return Value(debugTypeString(e)); + case LiteralSerializationPolicy::kToRepresentativeParseableValue: + return representativeValue.value_or(defaultLiteralOfType(e)); + default: + MONGO_UNREACHABLE_TASSERT(7539802); + } +} + +Value SerializationOptions::serializeLiteral( + const ImplicitValue& v, const boost::optional<Value>& representativeValue) const { + switch (literalPolicy) { + case LiteralSerializationPolicy::kUnchanged: + return v; + case LiteralSerializationPolicy::kToDebugTypeString: + return Value(debugTypeString(v)); + case LiteralSerializationPolicy::kToRepresentativeParseableValue: + return representativeValue.value_or(defaultLiteralOfType(v)); + default: + MONGO_UNREACHABLE_TASSERT(7539804); + } +} + +std::string SerializationOptions::serializeFieldPathFromString(StringData path) const { + if (transformIdentifiers) { + try { + return serializeFieldPath(FieldPath(path, false)); + } catch (DBException& ex) { + LOGV2_DEBUG(7549808, + 1, + "Failed to convert a path string to a FieldPath", + "pathString"_attr = path, + "failure"_attr = ex.toStatus()); + return serializeFieldPath("invalidFieldPathPlaceholder"); + } + } + return path.toString(); +} +} // namespace mongo diff --git a/src/mongo/db/query/query_shape/serialization_options.h b/src/mongo/db/query/query_shape/serialization_options.h new file mode 100644 index 00000000000..226da7689d3 --- /dev/null +++ b/src/mongo/db/query/query_shape/serialization_options.h @@ -0,0 +1,236 @@ +/** + * Copyright (C) 2023-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 "mongo/base/string_data.h" +#include "mongo/bson/bsonobj.h" +#include "mongo/bson/bsonobjbuilder.h" +#include "mongo/db/exec/document_value/document.h" +#include "mongo/db/exec/document_value/value.h" +#include "mongo/db/pipeline/field_path.h" +#include "mongo/db/query/explain_options.h" +#include "mongo/util/assert_util.h" +#include <boost/optional.hpp> +#include <string> + +namespace mongo { +namespace { +// Should never be called, throw to ensure we catch this in tests. +std::string defaultHmacStrategy(StringData s) { + MONGO_UNREACHABLE_TASSERT(7332410); +} +} // namespace + +/** + * A policy enum for how to serialize literal values. + */ +enum class LiteralSerializationPolicy { + // The default way to serialize. Just serialize whatever literals were given if they are still + // available, or whatever you parsed them to. This is expected to be able to parse again, since + // it worked the first time. + kUnchanged, + // Serialize any literal value as "?number" or similar. For example "?bool" for any boolean. Use + // 'debugTypeString()' helper. + kToDebugTypeString, + // Serialize any literal value to one canonical value of the given type, with the constraint + // that the chosen representative value should be parseable in this context. There are some + // default implementations that will usually work (e.g. using the number 1 almost always works + // for numbers), but serializers should be careful to think about and test this if their parsers + // reject certain values. + kToRepresentativeParseableValue, +}; + +/** + * A struct with options for how you want to serialize a match or aggregation expression. + */ +struct SerializationOptions { + using TokenizeIdentifierFunc = std::function<std::string(StringData)>; + + // The default serialization options for a query shape. No need to redact identifiers for the + // this purpose. We may do that on the $queryStats read path. + static const SerializationOptions kRepresentativeQueryShapeSerializeOptions; + static const SerializationOptions kDebugQueryShapeSerializeOptions; + static const SerializationOptions kMarkIdentifiers_FOR_TEST; + static const SerializationOptions kDebugShapeAndMarkIdentifiers_FOR_TEST; + + SerializationOptions() = default; + SerializationOptions(LiteralSerializationPolicy policy); + SerializationOptions(boost::optional<ExplainOptions::Verbosity> explain); + SerializationOptions(LiteralSerializationPolicy policy, + bool transformIdentifiers, + TokenizeIdentifierFunc transformIdentifiersCallbackFn); + + /** + * Checks if this SerializationOptions represents the same options as another + * SerializationOptions. Note it cannot compare whether the two 'transformIdentifiersCallback's + * are the same - the language purposefully leaves the comparison operator undefined. + */ + bool operator==(const SerializationOptions& other) const { + return this->transformIdentifiers == other.transformIdentifiers && + // You cannot well determine std::function equivalence in C++, so this is the best we'll + // do. + (this->transformIdentifiersCallback == nullptr) == + (other.transformIdentifiersCallback == nullptr) && + this->literalPolicy == other.literalPolicy && this->verbosity == other.verbosity; + } + bool operator!=(const SerializationOptions& other) const { + return !(*this == other); + } + + // Helper function for removing identifiable information (like collection/db names). + // Note: serializeFieldPath/serializeFieldPathFromString should be used for field + // names. + std::string serializeIdentifier(StringData str) const { + if (transformIdentifiers) { + return transformIdentifiersCallback(str); + } + return str.toString(); + } + + std::string serializeFieldPath(FieldPath path) const { + if (transformIdentifiers) { + std::stringstream hmaced; + for (size_t i = 0; i < path.getPathLength(); ++i) { + if (i > 0) { + hmaced << "."; + } + hmaced << transformIdentifiersCallback(path.getFieldName(i)); + } + return hmaced.str(); + } + return path.fullPath(); + } + + std::string serializeFieldPathWithPrefix(FieldPath path) const { + return "$" + serializeFieldPath(path); + } + + std::string serializeFieldPathFromString(StringData path) const; + + std::vector<std::string> serializeFieldPathFromString( + const std::vector<std::string>& paths) const { + std::vector<std::string> result; + result.reserve(paths.size()); + for (auto& p : paths) { + result.push_back(serializeFieldPathFromString(p)); + } + return result; + } + + // Helper functions for applying hmac to BSONObj. Does not take into account anything to do with + // MQL semantics, removes all field names and literals in the passed in obj. + void addHmacedArrayToBuilder(BSONArrayBuilder* bab, std::vector<BSONElement> array) const { + for (const auto& elem : array) { + if (elem.type() == BSONType::Object) { + BSONObjBuilder subObj(bab->subobjStart()); + addHmacedObjToBuilder(&subObj, elem.Obj()); + subObj.done(); + } else if (elem.type() == BSONType::Array) { + BSONArrayBuilder subArr(bab->subarrayStart()); + addHmacedArrayToBuilder(&subArr, elem.Array()); + subArr.done(); + } else { + *bab << serializeLiteral(elem); + } + } + } + + void addHmacedObjToBuilder(BSONObjBuilder* bob, BSONObj objToHmac) const { + for (const auto& elem : objToHmac) { + auto fieldName = serializeFieldPath(elem.fieldName()); + if (elem.type() == BSONType::Object) { + BSONObjBuilder subObj(bob->subobjStart(fieldName)); + addHmacedObjToBuilder(&subObj, elem.Obj()); + subObj.done(); + } else if (elem.type() == BSONType::Array) { + BSONArrayBuilder subArr(bob->subarrayStart(fieldName)); + addHmacedArrayToBuilder(&subArr, elem.Array()); + subArr.done(); + } else { + appendLiteral(bob, fieldName, elem); + } + } + } + + /** + * Helper method to call 'serializeLiteral()' on 'e' and append the resulting value to 'bob' + * using the same name as 'e'. + */ + void appendLiteral(BSONObjBuilder* bob, const BSONElement& e) const; + void appendLiteral(BSONObjBuilder* bob, StringData name, const BSONElement& e) const; + /** + * Helper method to call 'serializeLiteral()' on 'v' and append the result to 'bob' using field + * name 'fieldName'. + */ + void appendLiteral(BSONObjBuilder* bob, + StringData fieldName, + const ImplicitValue& v, + const boost::optional<Value>& representativeValue = boost::none) const; + + /** + * Depending on the configured 'literalPolicy', serializeLiteral will return the appropriate + * value for adding literals to serialization output: + * - If 'literalPolicy' is 'kUnchanged', returns the input value unmodified. + * - If it is 'kToDebugTypeString', computes and returns the type string as a string Value. + * - If it is 'kToRepresentativeValue', it returns an arbitrary value of the same type as the + * one given. For any number, this will be the number 1. For any boolean this will be true. + * If the 'representativeValue' parameter if it is not none, returns it (regardless of type). + * + * Example usage: BSON("myArg" << options.serializeLiteral(_myArg)); + */ + Value serializeLiteral(const BSONElement& e, + const boost::optional<Value>& representativeValue = boost::none) const; + Value serializeLiteral(const ImplicitValue& v, + const boost::optional<Value>& representativeValue = boost::none) const; + + // 'literalPolicy' is an independent option to serialize in a general format with the aim of + // similar "shaped" queries serializing to the same object. For example, if set to + // 'kToDebugTypeString', then the serialization of {a: {$gt: 2}} should result in {a: {$gt: + // '?number'}}, as will the serialization of {a: {$gt: 3}}. + // + // "Literal" here is meant to stand in contrast to expression arguments, as in the $gt + // expressions in {$and: [{a: {$gt: 3}}, {b: {$gt: 4}}]}. There the only literals are 3 and 4, + // so the serialization expected for 'kToDebugTypeString' would be {$and: [{a: {$gt: + // '?number'}}, {b: {$lt: '?number'}}]}. + LiteralSerializationPolicy literalPolicy = LiteralSerializationPolicy::kUnchanged; + + // If true the caller must set transformIdentifiersCallback. 'transformIdentifiers' if set along + // with a strategy the redaction strategy will be called on any personal identifiable + // information (e.g., field paths/names, collection names) encountered before serializing them. + bool transformIdentifiers = false; + std::function<std::string(StringData)> transformIdentifiersCallback = defaultHmacStrategy; + + // For aggregation indicate whether we should use the more verbose serialization format. + boost::optional<ExplainOptions::Verbosity> verbosity = boost::none; + + // If set to true, serializes each stage and expression as needed for query analysis. + bool serializeForQueryAnalysis = false; +}; + +} // namespace mongo diff --git a/src/mongo/db/query/query_shape/shape_helpers.cpp b/src/mongo/db/query/query_shape/shape_helpers.cpp new file mode 100644 index 00000000000..8eea475ab78 --- /dev/null +++ b/src/mongo/db/query/query_shape/shape_helpers.cpp @@ -0,0 +1,108 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_shape/shape_helpers.h" + +#include "mongo/db/query/query_shape/query_shape_gen.h" + +namespace mongo::shape_helpers { + +static constexpr StringData hintSpecialField = "$hint"_sd; +// A "Flat" object is one with only top-level fields. We won't descend recursively to shapify any +// sub-objects. +BSONObj shapifyFlatObj(BSONObj obj, const SerializationOptions& opts, bool valuesAreLiterals) { + if (obj.isEmpty()) { + // fast-path for the common case. + return obj; + } + + BSONObjBuilder bob; + for (BSONElement elem : obj) { + if (hintSpecialField.compare(elem.fieldNameStringData()) == 0) { + if (elem.type() == BSONType::String) { + bob.append(hintSpecialField, opts.serializeFieldPathFromString(elem.String())); + } else if (elem.type() == BSONType::Object) { + opts.appendLiteral(&bob, hintSpecialField, elem.Obj()); + } else { + // SERVER-85500: $hint syntax will not be validated if the collection does not + // exist, so we should accept a value that is neither string nor object here. + opts.appendLiteral(&bob, hintSpecialField, elem); + } + continue; + } + + // $natural doesn't need to be redacted. + if (elem.fieldNameStringData().compare(query_request_helper::kNaturalSortField) == 0) { + bob.append(elem); + continue; + } + + if (valuesAreLiterals) { + opts.appendLiteral(&bob, opts.serializeFieldPathFromString(elem.fieldName()), elem); + } else { + bob.appendAs(elem, opts.serializeFieldPathFromString(elem.fieldName())); + } + } + return bob.obj(); +} + +BSONObj extractHintShape(BSONObj hintObj, const SerializationOptions& opts) { + return shapifyFlatObj(hintObj, opts, /* valuesAreLiterals = */ false); +} + +BSONObj extractMinOrMaxShape(BSONObj obj, const SerializationOptions& opts) { + return shapifyFlatObj(obj, opts, /* valuesAreLiterals = */ true); +} + +void appendNamespaceShape(BSONObjBuilder& bob, + const NamespaceString& nss, + const SerializationOptions& opts) { + bob.append("db", opts.serializeIdentifier(nss.db())); + bob.append("coll", opts.serializeIdentifier(nss.coll())); +} + +NamespaceStringOrUUID parseNamespaceShape(BSONElement cmdNsElt) { + tassert(7632900, "cmdNs must be an object.", cmdNsElt.type() == BSONType::Object); + auto cmdNs = query_shape::CommandNamespace::parse("cmdNs"_sd, cmdNsElt.embeddedObject()); + + if (cmdNs.getColl().has_value()) { + tassert(7632903, + "Exactly one of 'uuid' and 'coll' can be defined.", + !cmdNs.getUuid().has_value()); + return NamespaceString(cmdNs.getDb(), cmdNs.getColl().value()); + } else { + tassert(7632904, + "Exactly one of 'uuid' and 'coll' can be defined.", + !cmdNs.getColl().has_value()); + UUID uuid = uassertStatusOK(UUID::parse(cmdNs.getUuid().value().toString())); + return NamespaceStringOrUUID(cmdNs.getDb().toString(), uuid); + } +} + +} // namespace mongo::shape_helpers diff --git a/src/mongo/db/query/query_shape/shape_helpers.h b/src/mongo/db/query/query_shape/shape_helpers.h new file mode 100644 index 00000000000..4d0fadb4a47 --- /dev/null +++ b/src/mongo/db/query/query_shape/shape_helpers.h @@ -0,0 +1,101 @@ +/** + * Copyright (C) 2023-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 "mongo/bson/simple_bsonobj_comparator.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_shape/serialization_options.h" + +namespace mongo::shape_helpers { + +int64_t inline optionalObjSize(boost::optional<BSONObj> optionalObj) { + if (!optionalObj) + return 0; + return optionalObj->objsize(); +} + +template <typename T> +int64_t optionalSize(boost::optional<T> optionalVal) { + if (!optionalVal) + return 0; + return optionalVal->size(); +} + +template <typename T> +std::function<size_t(size_t, const T&)> sizeAccumulatorFunc() { + MONGO_UNREACHABLE; // Don't know how to compute the size of this template type. +}; + +template <> +inline std::function<size_t(size_t, const BSONObj&)> sizeAccumulatorFunc<BSONObj>() { + return [](size_t total, const BSONObj& obj) { + return total + sizeof(BSONObj) + static_cast<size_t>(obj.objsize()); + }; +} + +template <> +inline std::function<size_t(size_t, const NamespaceString&)> +sizeAccumulatorFunc<NamespaceString>() { + return [](size_t total, const NamespaceString& nss) { + // For each element, we have to track the size of the + // nss as well as the size allocated by the nss. It would be + // ideal to be able to ask the underlying namespace string for + // its capacity, but it's not something we have access to. + // Further, namespace strings appear to shrink to fit (i.e + // resize to correct size), so it may not be necessary. Should + // we also try to consider short string optimization? At the + // very least, the current approach gives us a good upper bound + // memory usage (assuming shrink to fit). + return total + sizeof(nss) + nss.size(); + }; +} + +template <typename Container> +size_t containerSize(const Container& container) { + return std::accumulate(container.begin(), + container.end(), + 0, + sizeAccumulatorFunc<typename Container::value_type>()); +} + +/** + * Serializes the given 'hintObj' in accordance with the options. Assumes the hint is correct and + * contains field names. It is possible that this hint doesn't actually represent an index, but we + * can't detect that here. + */ +BSONObj extractHintShape(BSONObj hintObj, const SerializationOptions& opts); +BSONObj extractMinOrMaxShape(BSONObj obj, const SerializationOptions& opts); + +NamespaceStringOrUUID parseNamespaceShape(BSONElement cmdNsElt); +void appendNamespaceShape(BSONObjBuilder& bob, + const NamespaceString& nss, + const SerializationOptions& opts); + +} // namespace mongo::shape_helpers diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp index 669777b61ae..191218b2f4d 100644 --- a/src/mongo/db/query/query_solution.cpp +++ b/src/mongo/db/query/query_solution.cpp @@ -1566,11 +1566,11 @@ void GroupNode::appendToString(str::stream* ss, int indent) const { if (idx > 0) { *ss << ", "; } - *ss << "{" << groupName << ": " << exprObj->serialize(false).toString() << "}"; + *ss << "{" << groupName << ": " << exprObj->serialize().toString() << "}"; ++idx; } } else { - *ss << "{_id: " << groupByExpression->serialize(false).toString() << "}"; + *ss << "{_id: " << groupByExpression->serialize().toString() << "}"; } *ss << '\n'; addIndent(ss, indent + 1); @@ -1581,7 +1581,11 @@ void GroupNode::appendToString(str::stream* ss, int indent) const { } auto& acc = accumulators[idx]; *ss << "{" << acc.fieldName << ": {" << acc.expr.name << ": " - << acc.expr.argument->serialize(true).toString() << "}}"; + << acc.expr.argument + ->serialize(SerializationOptions{ + boost::make_optional(ExplainOptions::Verbosity::kQueryPlanner)}) + .toString() + << "}}"; } *ss << "]" << '\n'; addCommon(ss, indent); diff --git a/src/mongo/db/query/query_stats/README.md b/src/mongo/db/query/query_stats/README.md new file mode 100644 index 00000000000..6f2667fbfd9 --- /dev/null +++ b/src/mongo/db/query/query_stats/README.md @@ -0,0 +1,200 @@ +# Query Stats +This directory is the home of the infrastructure related to recording runtime query statistics for +the database. It is not to be confused with `src/mongo/db/query/stats/` which is the home of the +logic for computing and maintaining statistics about a collection or index's data distribution - for +use by the query planner. + +The system will collect metrics for each query execution, and the results will be aggregated in a +structure called the [`QueryStatsStore`](#querystatsstore) upon completion of each successful +execution. Metrics will be aggregated according to an abstracted version of the query known as the +query stats key and will be collected on any mongod or mongos process for which they are configured, +including primaries and secondaries. + +## QueryStatsStore +At the center of everything here is the [`QueryStatsStore`](query_stats.h#93-97), which is a +partitioned hash table that maps the hash of a [Query Stats Key](#glossary) (also known as the +_Query Stats Store Key_) to some metrics about how often each one occurs. + +### Computing the Query Stats Store Key +A query stats store key contains various dimensions that distinctify a specific query. One main +attribute to the query stats store key, is the query shape (`query_shape::Shape`). For example, if +the client does this: +```js +db.example.findOne({x: 24}); +db.example.findOne({x: 53}); +``` +then the `QueryStatsStore` should contain an entry for a single query shape which would record 2 +executions and some related statistics (see [`QueryStatsEntry`](query_stats_entry.h) for details). + +For more information on query shape, see the [query_shape](../query_shape/README.md) directory. + +The query stats store has _more_ dimensions (i.e. more granularity) to group incoming queries than +just the query shape. For example, these queries would all three have the same shape but the first +would have a different query stats store entry from the other two: +```js +db.example.find({x: 55}); +db.example.find({x: 55}).batchSize(2); +db.example.find({x: 55}).batchSize(3); +``` +There are two distinct query stats store entries here - both the examples which include the batch +size will be treated separately from the example which does not specify a batch size. + +The dimensions considered will depend on the command, but can generally be found in the +[`KeyGenerator`](key_generator.h) interface, which will generate the query stats store keys by which +we accumulate statistics. As one example, you can find the +[`FindKey`](find_key.h) which will include all the things tracked in the +`FindCmdQueryStatsStoreKeyComponents` (including `batchSize` shown in this example). + +### Query Stats Store Cache Size +The size of the`QueryStatsStore` can be set by the server parameter +[`internalQueryStatsCacheSize`](#server-parameters), and the partitions will be created based off +that. See [`queryStatsStoreManagerRegisterer`](query_stats.cpp#L138-L154) for more details about how +the number of partitions and their size is determined; Each partition is an LRU cache, therefore, if +adding a new entry to the partition makes it go over its size limit, the least recently used entries +will be evicted to drop below the max size. Eviction will be tracked in the new [server status +metrics](#server-status-metrics) for queryStats. + +## Metric Collection +At a high level, when a query is run and collection of query stats is enabled, during planning we +call [`registerRequest`]((query_stats.h#L195-L198)) in which the query stats store key will be +generated based on the query's shape and the various other dimensions. The key will always be serialized +and stored on the `opDebug`, and also on the cursor in the case that there are `getMore`s, so that we can +continue to aggregate the operation's metrics. Once the query execution is fully complete, +[`writeQueryStats`](query_stats.h#L200-216) will be called and will either retrieve the entry for +the key from the store if it exists and update it, or create a new one and add it to the store. See +more details in the [comments](query_stats.h#L158-L216). + +### Rate Limiting +Whether or not query stats will be recorded for a specific query execution depends on a Rate +Limiter, which limits the number of recordings per second based on the server parameter +[internalQueryStatsRateLimit](#server-parameters). The goal of the rate limiter is to minimize +impact to overall system performance through restricting excessive traffic. If a query is run but +the rate limit has been reached, the query will still execute as expected but query stats will not +be updated in the query stats store. Our rate limiter uses the sliding window algorithm; see details +[here](rate_limiting.h#82-87). + +## Metric Retrieval +To retrieve the stats gathered in the `QueryStatsStore`, there is a new aggregation stage, +`$queryStats`. This stage must be the first in a pipeline and it must be run against the admin +database. The structure of the command is as follows (note `aggregate: 1` reflecting there is no collection): +```js +db.adminCommand({ + aggregate: 1, + pipeline: [{ + $queryStats: { + tranformIdentifiers: { + algorithm: "hmac-sha-256", + hmacKey: BinData(8, "87c4082f169d3fef0eef34dc8e23458cbb457c3sf3n2") /* bindata + subtype 8 - a new type for sensitive data */, + } + } + }] +}) +``` +`transformIdentifiers` is optional. If not present, we will generate the regular Query Stats Key. If +present: +- `algorithm` is required and the only currently supported option is "hmac-sha-256". +- `hmacKey` is required +- We will generate the [One-way Tokenized](#glossary) Query Stats Key by applying the "hmac-sha-256" + to the names of any field, collection, or database. Application Name field is not transformed. + +The query stats store will output one document for each query stats key, which is structured in the +following way: +```js +{ + key: {/* Query Stats Key */}, + asOf: ISODate(/* … */), + metrics: { + execCount: 0, + firstSeenTimestamp: ISODate(/* … */), + latestSeenTimestamp: ISODate(/* … */), + docsReturned: {sum: 0, max: 0, min: 0, sumOfSquares: 0}, + firstResponseExecMicros: {sum: 0, max: 0, min: 0, sumOfSquares: 0}, + totalExecMicros: {sum: 0, max: 0, min: 0, sumOfSquares: 0}, + lastExecutionMicros: 0, + } +} +``` +- `key`: Query Stats Key. +- `asOf`: UTC time when $queryStats read this entry from the store. This will not return the same + UTC time for each result. The data structure used for the store is partitioned, and each partition + will be read at a snapshot individually. You may see up to the number of partitions in unique + timestamps returned by one $queryStats cursor. +- `metrics`: the metrics collected; these may be flawed due to: + - Server restarts, which will reset metrics. + - LRU eviction, which will reset metrics. + - Rate limiting, which will skew metrics. +- `metrics.execCount`: Number of recorded observations of this query. +- `metrics.firstSeenTimestamp`: UTC time taken at query completion (including getMores) for the + first recording of this query stats store entry. +- `metrics.lastSeenTimestamp`: UTC time taken at query completion (including getMores) for the + latest recording of this query stats store entry. +- `metrics.docsReturned`: Various broken down metrics for the number of documents returned by + observation of this query. +- `metrics.firstResponseExecMicros`: Estimated time spent computing and returning the first batch. +- `metrics.totalExecMicros`: Estimated time spent computing and returning all batches, which is the + same as the above for single-batch queries. +- `metrics.lastExecutionMicros`: Estimated time spent processing the latest query (akin to + "totalExecMicros", not "firstResponseExecMicros"). + +#### Permissions +`$queryStats` is restricted by two privilege actions: +- `queryStatsRead` privilege allows running `$queryStats` without passing the `transformIdentifiers` + options. +- `queryStatsReadTransformed` allows running `$queryStats` with `transformIdentifiers` set. These +two privileges are included in the clusterMonitor role in Atlas. + +### Server Parameters +- `internalQueryStatsCacheSize`: + * Max query stats store size, specified as a string like "4MB" or "1%". Defaults to 1% of the + machine's total memory. + * Query stats store is a LRU cache structure with partitions, so we may be under the cap due to + implementation. + +- `internalQueryStatsRateLimit`: + * The rate limit is an integer which imposes a maximum number of recordings per second. Default is + 0 which has the effect of disabling query stats collection. Setting the parameter to -1 means + there will be no rate limit. + +- `logComponentVerbosity.queryStats`: + * Controls the logging behavior for query stats. See [Logging](#logging) for details. + +### Logging +Setting `logComponentVerbosity.queryStats` will do the following for each level: +* Level 0 (default): Nothing will be logged. +* Level 1 or higher: Invocations of $queryStats will be logged if and only if the algorithm is + "hmac-sha-256". The specification of the $queryStats stage is logged, with any provided hmac key + redacted. +* Level 2 or higher: Nothing extra, reserved for future use. +* Level 3 or higher: All results of any "hmac-sha-256" $queryStats invocation are logged. Each + result will be its own entry and there will be one final entry that says "we finished". +* Levels 4 and 5 do nothing extra. + +### Server Status Metrics +The following will be added to the `serverStatus.metrics`: +```js +queryStats: { + numEvicted: NumberLong(0), + numHmacApplicationErrors: NumberLong(0), + numQueryStatsStoreWriteErrors: NumberLong(0), + numRateLimitedRequests: NumberLong(0), + queryStatsStoreSizeEstimateBytes: NumberLong(0) +} +``` + +# Glossary +**Query Execution**: This term implies the overall execution of what a client would consider one +query, but which may or may not involve one or more getMore commands to iterate a cursor. For +example, a find command and two getMore commands on the returned cursor is one query execution. An +aggregate command which returns everything in one batch is also one query execution. + +**One-way Tokenized Object**: A one-way tokenized object has an HMAC hashing function applied to +particular sensitive elements/pieces of an object. It is "one-way" because it is never meant to be +undone. This allows us to detect when two queries are using the same identifiers, but never to +reveal what those identifiers were. + +**Query Shape**: [Query Shape](../query_shape/README.md) + +**Query Stats Key**: Also known as the _Query Stats Store Key_, this is the collection of attributes +championed by the query shape which identifies one grouping of metrics. The $queryStats stage will +output one document per query stats key - output in the "key" field. diff --git a/src/mongo/db/query/query_stats/SConscript b/src/mongo/db/query/query_stats/SConscript new file mode 100644 index 00000000000..f9f3a8b1c2e --- /dev/null +++ b/src/mongo/db/query/query_stats/SConscript @@ -0,0 +1,121 @@ +# -*- mode: python -*- + +Import([ + "env", + "get_option", +]) + +env = env.Clone() + +env.Library( + target='rate_limiting', + source=[ + 'rate_limiting.cpp', + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/util/clock_sources', + ], +) + +env.Library(target='query_stats_parse', source=['transform_algorithm.idl'], LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/idl/idl_parser', +]) + +env.Library( + target='query_stats', + source=[ + '$BUILD_DIR/mongo/db/curop.cpp', + 'key.cpp', + 'query_stats.cpp', + 'query_stats_entry.cpp' + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/bson/mutable/mutable_bson', + '$BUILD_DIR/mongo/db/commands', + '$BUILD_DIR/mongo/db/concurrency/lock_manager', + '$BUILD_DIR/mongo/db/exec/document_value/document_value', + '$BUILD_DIR/mongo/db/generic_cursor', + '$BUILD_DIR/mongo/db/profile_filter', + '$BUILD_DIR/mongo/db/query/command_request_response', + '$BUILD_DIR/mongo/db/query/memory_util', + '$BUILD_DIR/mongo/db/query/query_knobs', + '$BUILD_DIR/mongo/db/query/query_shape/query_shape', + '$BUILD_DIR/mongo/db/server_options', + '$BUILD_DIR/mongo/db/service_context', + '$BUILD_DIR/mongo/db/stats/counters', + '$BUILD_DIR/mongo/db/stats/timer_stats', + '$BUILD_DIR/mongo/db/storage/storage_engine_parameters', + '$BUILD_DIR/mongo/rpc/client_metadata', + '$BUILD_DIR/mongo/transport/service_executor', + '$BUILD_DIR/mongo/util/diagnostic_info' if get_option('use-diagnostic-latches') == 'on' else [], + '$BUILD_DIR/mongo/util/fail_point', + '$BUILD_DIR/mongo/util/net/network', + '$BUILD_DIR/mongo/util/processinfo', + '$BUILD_DIR/mongo/util/progress_meter', + 'query_stats_parse', + 'rate_limiting', + ], + LIBDEPS_PRIVATE=[ + '$BUILD_DIR/mongo/db/auth/auth', + '$BUILD_DIR/mongo/db/auth/user_acquisition_stats', + '$BUILD_DIR/mongo/db/exec/projection_executor', + '$BUILD_DIR/mongo/db/prepare_conflict_tracker', + '$BUILD_DIR/mongo/db/stats/resource_consumption_metrics', + ], +) + +env.CppUnitTest( + target="db_query_query_stats_test", + source=[ + "agg_key_test.cpp", + "find_key_test.cpp", + "key_test.cpp", + "query_stats_test.cpp", + "query_stats_store_test.cpp", + "rate_limiting_test.cpp", + ], + LIBDEPS=[ + "$BUILD_DIR/mongo/db/auth/authmocks", + "$BUILD_DIR/mongo/db/query/query_shape/query_shape", + "$BUILD_DIR/mongo/db/query/query_test_service_context", + "$BUILD_DIR/mongo/db/service_context_d_test_fixture", + "query_stats", + "rate_limiting", + ], +) + +env.Benchmark( + target='rate_limiting_bm', + source=[ + 'rate_limiting_bm.cpp', + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/unittest/unittest', + '$BUILD_DIR/mongo/util/processinfo', + 'rate_limiting', + ], +) + +env.Benchmark( + target='shapifying_bm', + source=[ + 'shapifying_bm.cpp', + ], + LIBDEPS=[ + '$BUILD_DIR/mongo/base', + '$BUILD_DIR/mongo/db/auth/auth', + '$BUILD_DIR/mongo/db/pipeline/pipeline', + '$BUILD_DIR/mongo/db/query/canonical_query', + '$BUILD_DIR/mongo/db/query/query_shape/query_shape', + '$BUILD_DIR/mongo/db/query/query_test_service_context', + '$BUILD_DIR/mongo/db/service_context', + '$BUILD_DIR/mongo/rpc/client_metadata', + '$BUILD_DIR/mongo/unittest/unittest', + '$BUILD_DIR/mongo/util/processinfo', + 'query_stats', + ], +) diff --git a/src/mongo/db/query/query_stats/agg_key.cpp b/src/mongo/db/query/query_stats/agg_key.cpp new file mode 100644 index 00000000000..1d53418d371 --- /dev/null +++ b/src/mongo/db/query/query_stats/agg_key.cpp @@ -0,0 +1,174 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_stats/agg_key.h" + +#include "mongo/db/query/explain_options.h" +#include <absl/container/node_hash_set.h> +#include <boost/cstdint.hpp> +#include <functional> +#include <initializer_list> +#include <memory> +#include <numeric> +#include <vector> + +#include <boost/move/utility_core.hpp> +#include <boost/optional/optional.hpp> +#include <boost/smart_ptr/intrusive_ptr.hpp> + +#include "mongo/crypto/fle_field_schema_gen.h" +#include "mongo/db/pipeline/exchange_spec_gen.h" +#include "mongo/db/pipeline/pipeline.h" +#include "mongo/db/query/query_shape/agg_cmd_shape.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_shape/serialization_options.h" +#include "mongo/db/query/query_shape/shape_helpers.h" +#include "mongo/idl/basic_types_gen.h" +#include "mongo/util/assert_util.h" + +namespace mongo::query_stats { + +AggCmdComponents::AggCmdComponents(const AggregateCommandRequest& request_, + stdx::unordered_set<NamespaceString> involvedNamespaces_) + : involvedNamespaces(std::move(involvedNamespaces_)), + _bypassDocumentValidation(request_.getBypassDocumentValidation().value_or(false)), + _verbosity(request_.getExplain()), + _hasField() { + _hasField.batchSize = request_.getCursor().getBatchSize().has_value(); + _hasField.bypassDocumentValidation = request_.getBypassDocumentValidation().has_value(); + _hasField.explain = request_.getExplain().has_value(); + _hasField.passthroughToShard = request_.getPassthroughToShard().has_value(); +} + + +void AggCmdComponents::HashValue(absl::HashState state) const { + // The hashing for verbosity in this branch needed to be different because the compiler was + // complaining about the different wrappers. This is not important since this computation is + // only used locally in memory on a single machine, and the query shape is still stable. + auto verbosity = + _hasField.explain ? std::string(ExplainOptions::verbosityString(_verbosity.value())) : ""; + state = absl::HashState::combine(std::move(state), + _bypassDocumentValidation, + _hasField.batchSize, + _hasField.bypassDocumentValidation, + verbosity, + _hasField.explain, + _hasField.passthroughToShard); + // We don't need to add 'involvedNamespaces' here since they are already tracked/duplicated in + // the Pipeline component of the query shape. We just expose them here for ease of + // analysis/querying. +} + +void AggCmdComponents::appendTo(BSONObjBuilder& bob, const SerializationOptions& opts) const { + + // otherNss + if (!involvedNamespaces.empty()) { + BSONArrayBuilder otherNss = bob.subarrayStart(kOtherNssFieldName); + for (const auto& nss : involvedNamespaces) { + BSONObjBuilder otherNsEntryBob = otherNss.subobjStart(); + shape_helpers::appendNamespaceShape(otherNsEntryBob, nss, opts); + otherNsEntryBob.doneFast(); + } + otherNss.doneFast(); + } + + // bypassDocumentValidation + if (_hasField.bypassDocumentValidation) { + bob.append(AggregateCommandRequest::kBypassDocumentValidationFieldName, + _bypassDocumentValidation); + } + + // We don't store the specified batch size values since they don't matter. + // Provide an arbitrary literal long here. + + tassert(78429, + "Serialization policy not supported - original values have been discarded", + opts.literalPolicy != LiteralSerializationPolicy::kUnchanged); + + if (_hasField.batchSize) { + // cursor + BSONObjBuilder cursorInfo = bob.subobjStart(AggregateCommandRequest::kCursorFieldName); + opts.appendLiteral(&cursorInfo, SimpleCursorOptions::kBatchSizeFieldName, 0ll); + cursorInfo.doneFast(); + } + + if (_hasField.explain) { + // The verbosity can be explicitly set by using the .explain() command, but when using the + // flag {explain: true} it is set to 'queryPlanner'. + bob.append(AggregateCommandRequest::kExplainFieldName, + ExplainOptions::verbosityString(_verbosity.value())); + } + + // The values here don't matter (assuming we're not using the 'kUnchanged' policy). + tassert(8949601, + "Serialization policy not supported - original values have been discarded", + opts.literalPolicy != LiteralSerializationPolicy::kUnchanged); + if (_hasField.passthroughToShard) { + BSONObjBuilder passthroughToShardInfo = + bob.subobjStart(AggregateCommandRequest::kPassthroughToShardFieldName); + static const PassthroughToShardOptions representativePassthroughOptions = []() { + PassthroughToShardOptions passthroughOpts; + // The value doesn't matter since we will only use this for shapified output. + passthroughOpts.setShard("?"); + return passthroughOpts; + }(); + representativePassthroughOptions.serialize(&passthroughToShardInfo, opts); + passthroughToShardInfo.doneFast(); + } +} + +size_t AggCmdComponents::size() const { + return sizeof(AggCmdComponents) + + std::accumulate(involvedNamespaces.begin(), + involvedNamespaces.end(), + 0, + [](int64_t total, const auto& nss) { return total + nss.size(); }); +} + +void AggKey::appendCommandSpecificComponents(BSONObjBuilder& bob, + const SerializationOptions& opts) const { + return _components.appendTo(bob, opts); +} + +AggKey::AggKey(AggregateCommandRequest request, + const Pipeline& pipeline, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + stdx::unordered_set<NamespaceString> involvedNamespaces, + const NamespaceString& origNss, + query_shape::CollectionType collectionType) + : Key(expCtx->opCtx, + std::make_unique<query_shape::AggCmdShape>( + request, origNss, involvedNamespaces, pipeline, expCtx), + request.getHint(), + request.getReadConcern(), + request.getMaxTimeMS().has_value(), + collectionType), + _components(request, std::move(involvedNamespaces)) {} + +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/agg_key.h b/src/mongo/db/query/query_stats/agg_key.h new file mode 100644 index 00000000000..38b80e28006 --- /dev/null +++ b/src/mongo/db/query/query_stats/agg_key.h @@ -0,0 +1,129 @@ +/** + * Copyright (C) 2023-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 <cstdint> +#include <utility> + +#include <absl/container/node_hash_map.h> +#include <boost/move/utility_core.hpp> +#include <boost/none.hpp> +#include <boost/optional/optional.hpp> +#include <boost/smart_ptr/intrusive_ptr.hpp> + +#include "mongo/base/string_data.h" +#include "mongo/bson/bsonobj.h" +#include "mongo/bson/bsonobjbuilder.h" +#include "mongo/db/collection_type.h" +#include "mongo/db/namespace_string.h" +#include "mongo/db/operation_context.h" +#include "mongo/db/pipeline/aggregate_command_gen.h" +#include "mongo/db/pipeline/expression_context.h" +#include "mongo/db/pipeline/pipeline.h" +#include "mongo/db/pipeline/variables.h" +#include "mongo/db/query/explain_options.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_stats/key.h" + +namespace mongo::query_stats { + +/** + * Struct representing the aggregate command's unique arguments which should be included in the + * query stats key. + */ +struct AggCmdComponents : public SpecificKeyComponents { + static constexpr StringData kOtherNssFieldName = "otherNss"_sd; + + AggCmdComponents(const AggregateCommandRequest&, + stdx::unordered_set<NamespaceString> involvedNamespaces); + + void HashValue(absl::HashState state) const final; + + void appendTo(BSONObjBuilder& bob, const SerializationOptions& opts) const; + + size_t size() const; + + stdx::unordered_set<NamespaceString> involvedNamespaces; + bool _bypassDocumentValidation; + const boost::optional<mongo::ExplainOptions::Verbosity> _verbosity; + + // This anonymous struct represents the presence of the member variables as C++ bit fields. + // In doing so, each of these boolean values takes up 1 bit instead of 1 byte. + struct HasField { + HasField() : batchSize(false), bypassDocumentValidation(false), explain(false) {} + bool batchSize : 1; + bool bypassDocumentValidation : 1; + bool explain : 1; + bool passthroughToShard : 1; + } _hasField; +}; + +/** + * Handles shapification for AggregateCommandRequests. Requires a pre-parsed pipeline in order to + * avoid parsing the raw pipeline multiple times, but users should be sure to provide a + * non-optimized pipeline. + */ +class AggKey final : public Key { +public: + AggKey(AggregateCommandRequest request, + const Pipeline& pipeline, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + stdx::unordered_set<NamespaceString> involvedNamespaces, + const NamespaceString& origNss, + query_shape::CollectionType collectionType = query_shape::CollectionType::kUnknown); + + const SpecificKeyComponents& specificComponents() const final { + return _components; + } + + // The default implementation of hashing for smart pointers is not a good one for our purposes. + // Here we overload them to actually take the hash of the object, rather than hashing the + // pointer itself. + template <typename H> + friend H AbslHashValue(H h, const std::unique_ptr<const AggKey>& key) { + return H::combine(std::move(h), *key); + } + template <typename H> + friend H AbslHashValue(H h, const std::shared_ptr<const AggKey>& key) { + return H::combine(std::move(h), *key); + } + + +protected: + void appendCommandSpecificComponents(BSONObjBuilder& bob, + const SerializationOptions& opts) const final override; + +private: + const AggCmdComponents _components; +}; +static_assert( + sizeof(AggKey) == sizeof(Key) + sizeof(AggCmdComponents), + "If the class' members have changed, this assert may need to be updated with a new value."); +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/agg_key_test.cpp b/src/mongo/db/query/query_stats/agg_key_test.cpp new file mode 100644 index 00000000000..35d0ae20d86 --- /dev/null +++ b/src/mongo/db/query/query_stats/agg_key_test.cpp @@ -0,0 +1,204 @@ +/** + * Copyright (C) 2024-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. + */ + +#include <boost/smart_ptr/intrusive_ptr.hpp> + +#include "mongo/db/pipeline/expression_context.h" +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/pipeline/pipeline.h" +#include "mongo/db/query/query_shape/agg_cmd_shape.h" +#include "mongo/db/query/query_stats/agg_key.h" +#include "mongo/db/service_context_test_fixture.h" +#include "mongo/idl/basic_types.h" +#include "mongo/unittest/unittest.h" +#include "mongo/util/intrusive_counter.h" + +namespace mongo::query_stats { + +namespace { + +static const NamespaceString kDefaultTestNss = NamespaceString("testDB.testColl"); + +static constexpr auto collectionType = query_shape::CollectionType::kCollection; + +class AggKeyTest : public ServiceContextTest { +public: + static std::unique_ptr<const Key> makeAggKeyFromRawPipeline( + const std::vector<BSONObj>& rawPipeline) { + auto expCtx = make_intrusive<ExpressionContextForTest>(kDefaultTestNss); + AggregateCommandRequest acr(kDefaultTestNss); + acr.setPipeline(rawPipeline); + auto pipeline = Pipeline::parse(rawPipeline, expCtx); + return std::make_unique<AggKey>(acr, + *pipeline, + expCtx, + pipeline->getInvolvedCollections(), + acr.getNamespace(), + collectionType); + } + size_t namespaceSize(stdx::unordered_set<NamespaceString> involvedNamespaces) { + return std::accumulate(involvedNamespaces.begin(), + involvedNamespaces.end(), + 0, + [](int64_t total, const auto& nss) { return total + nss.size(); }); + } +}; + +TEST_F(AggKeyTest, SizeOfAggCmdComponents) { + auto expCtx = make_intrusive<ExpressionContextForTest>(kDefaultTestNss); + auto rawPipeline = {fromjson(R"({ + $match: { + foo: { $in: ["a", "b"] }, + bar: { $gte: { $date: "2022-01-01T00:00:00Z" } } + } + })")}; + + AggregateCommandRequest acr(kDefaultTestNss); + acr.setPipeline(rawPipeline); + auto pipeline = Pipeline::parse(rawPipeline, expCtx); + auto namespaces = pipeline->getInvolvedCollections(); + auto aggComponents = std::make_unique<AggCmdComponents>(acr, namespaces); + + const auto minimumSize = sizeof(SpecificKeyComponents) + + sizeof(stdx::unordered_set<NamespaceString>) + 2 /*size for bool and HasField*/ + + sizeof(boost::optional<mongo::ExplainOptions::Verbosity>) + namespaceSize(namespaces); + ASSERT_GTE(aggComponents->size(), minimumSize); + ASSERT_LTE(aggComponents->size(), minimumSize + 8 /*padding*/); +} + +TEST_F(AggKeyTest, EquivalentAggCmdComponentSizes) { + auto expCtx = make_intrusive<ExpressionContextForTest>(kDefaultTestNss); + auto rawPipeline = {fromjson(R"({ + $match: { + foo: { $in: ["a", "b"] }, + bar: { $gte: { $date: "2022-01-01T00:00:00Z" } } + } + })")}; + + // Set different values in the command request. + AggregateCommandRequest acrBypassTrue(kDefaultTestNss); + acrBypassTrue.setPipeline(rawPipeline); + acrBypassTrue.setBypassDocumentValidation(true); + auto pipeline = Pipeline::parse(rawPipeline, expCtx); + auto namespaces = pipeline->getInvolvedCollections(); + auto aggComponentsBypassTrue = std::make_unique<AggCmdComponents>(acrBypassTrue, namespaces); + + + AggregateCommandRequest acrBypassFalse(kDefaultTestNss); + acrBypassFalse.setPipeline(rawPipeline); + acrBypassFalse.setBypassDocumentValidation(false); + auto aggComponentsBypassFalse = std::make_unique<AggCmdComponents>(acrBypassFalse, namespaces); + + ASSERT_EQ(aggComponentsBypassTrue->size(), aggComponentsBypassFalse->size()); +} + +TEST_F(AggKeyTest, DifferentAggCmdComponentSizes) { + auto expCtx = make_intrusive<ExpressionContextForTest>(kDefaultTestNss); + auto rawPipeline = {fromjson(R"({ + $match: { + foo: { $in: ["a", "b"] }, + bar: { $gte: { $date: "2022-01-01T00:00:00Z" } } + } + })")}; + + AggregateCommandRequest acr(kDefaultTestNss); + acr.setPipeline(rawPipeline); + // Manually creating different namespaces for testing purposes. + const auto namespaceStringOne = NamespaceString("testDB.testColl1"); + const auto namespaceStringTwo = NamespaceString("testDB.testColl2"); + + stdx::unordered_set<NamespaceString> smallNamespaces; + smallNamespaces.insert(namespaceStringOne); + + stdx::unordered_set<NamespaceString> largeNamespaces; + largeNamespaces.insert(namespaceStringOne); + largeNamespaces.insert(namespaceStringTwo); + + auto smallAggComponents = std::make_unique<AggCmdComponents>(acr, smallNamespaces); + auto largeAggComponents = std::make_unique<AggCmdComponents>(acr, largeNamespaces); + + ASSERT_LT(namespaceSize(smallNamespaces), namespaceSize(largeNamespaces)); + ASSERT_LT(smallAggComponents->size(), largeAggComponents->size()); +} + +// Testing item in opCtx that should impact key size. +TEST_F(AggKeyTest, SizeOfAggKeyWithAndWithoutWriteConcern) { + auto rawPipeline = {fromjson(R"({ + $match: { + foo: { $in: ["a", "b"] }, + bar: { $gte: { $date: "2022-01-01T00:00:00Z" } } + } + })")}; + + auto keyWithoutComment = makeAggKeyFromRawPipeline(rawPipeline); + + auto expCtx = make_intrusive<ExpressionContextForTest>(kDefaultTestNss); + AggregateCommandRequest acrWithComment(kDefaultTestNss); + acrWithComment.setPipeline(rawPipeline); + expCtx->opCtx->setComment(BSON("comment" + << " foo")); + auto pipelineWithComment = Pipeline::parse(rawPipeline, expCtx); + auto keyWithComment = std::make_unique<AggKey>(acrWithComment, + *pipelineWithComment, + expCtx, + pipelineWithComment->getInvolvedCollections(), + acrWithComment.getNamespace(), + collectionType); + + ASSERT_LT(keyWithoutComment->size(), keyWithComment->size()); +} + +// Testing item in command request that should impact key size. +TEST_F(AggKeyTest, SizeOfAggKeyWithAndWithoutReadConcern) { + auto rawPipeline = {fromjson(R"({ + $match: { + foo: { $in: ["a", "b"] }, + bar: { $gte: { $date: "2022-01-01T00:00:00Z" } } + } + })")}; + + auto keyWithoutReadConcern = makeAggKeyFromRawPipeline(rawPipeline); + + auto expCtx = make_intrusive<ExpressionContextForTest>(kDefaultTestNss); + AggregateCommandRequest acrWithReadConcern(kDefaultTestNss); + acrWithReadConcern.setPipeline(rawPipeline); + acrWithReadConcern.setReadConcern(fromjson(R"({level: "local"})")); + auto pipelineWithReadConcern = Pipeline::parse(rawPipeline, expCtx); + auto keyWithReadConcern = + std::make_unique<AggKey>(acrWithReadConcern, + *pipelineWithReadConcern, + expCtx, + pipelineWithReadConcern->getInvolvedCollections(), + acrWithReadConcern.getNamespace(), + collectionType); + + ASSERT_LT(keyWithoutReadConcern->size(), keyWithReadConcern->size()); +} +} // namespace +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/aggregate_key_generator.cpp b/src/mongo/db/query/query_stats/aggregate_key_generator.cpp new file mode 100644 index 00000000000..f175df296f2 --- /dev/null +++ b/src/mongo/db/query/query_stats/aggregate_key_generator.cpp @@ -0,0 +1,185 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_stats/aggregate_key_generator.h" + +#include "mongo/db/pipeline/pipeline.h" +#include "mongo/db/query/query_shape.h" +#include "mongo/db/query/serialization_options.h" +#include "mongo/db/query/shape_helpers.h" + +namespace mongo::query_stats { + +BSONObj AggregateKeyGenerator::generate( + OperationContext* opCtx, + boost::optional<SerializationOptions::TokenizeIdentifierFunc> hmacPolicy) const { + // TODO SERVER-76087 We will likely want to set a flag here to stop $search from calling out + // to mongot. + auto expCtx = makeDummyExpCtx(opCtx); + SerializationOptions opts{LiteralSerializationPolicy::kToDebugTypeString}; + if (hmacPolicy) { + opts.transformIdentifiers = true; + opts.transformIdentifiersCallback = *hmacPolicy; + opts.includePath = true; + opts.verbosity = boost::none; + } + + return makeQueryStatsKey(opts, expCtx); +} + +void AggregateKeyGenerator::appendCommandSpecificComponents( + BSONObjBuilder& bob, const SerializationOptions& opts) const { + // cursor + if (auto param = _request.getCursor().getBatchSize()) { + BSONObjBuilder cursorInfo = bob.subobjStart(AggregateCommandRequest::kCursorFieldName); + opts.appendLiteral(&cursorInfo, + SimpleCursorOptions::kBatchSizeFieldName, + static_cast<long long>(param.get())); + cursorInfo.doneFast(); + } + + // maxTimeMS + if (auto param = _request.getMaxTimeMS()) { + opts.appendLiteral(&bob, + AggregateCommandRequest::kMaxTimeMSFieldName, + static_cast<long long>(param.get())); + } + + // bypassDocumentValidation + if (auto param = _request.getBypassDocumentValidation()) { + opts.appendLiteral( + &bob, AggregateCommandRequest::kBypassDocumentValidationFieldName, bool(param.get())); + } + + // otherNss + if (!_involvedNamespaces.empty()) { + BSONArrayBuilder otherNss = bob.subarrayStart(kOtherNssFieldName); + for (const auto& nss : _involvedNamespaces) { + BSONObjBuilder otherNsEntryBob = otherNss.subobjStart(); + shape_helpers::appendNamespaceShape(otherNsEntryBob, nss, opts); + otherNsEntryBob.doneFast(); + } + otherNss.doneFast(); + } +} + +BSONObj AggregateKeyGenerator::makeQueryStatsKey( + const SerializationOptions& opts, const boost::intrusive_ptr<ExpressionContext>& expCtx) const { + auto pipeline = Pipeline::parse(_request.getPipeline(), expCtx); + return _makeQueryStatsKeyHelper(opts, expCtx, *pipeline); +} + +BSONObj AggregateKeyGenerator::_makeQueryStatsKeyHelper( + const SerializationOptions& opts, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + const Pipeline& pipeline) const { + return generateWithQueryShape( + query_shape::extractQueryShape(_request, pipeline, opts, expCtx, _origNss), opts); +} + +namespace { + +int64_t sum(const std::initializer_list<int64_t>& sizes) { + return std::accumulate(sizes.begin(), sizes.end(), 0, std::plus{}); +} + +int64_t size(const std::vector<BSONObj>& objects) { + return std::accumulate(objects.begin(), objects.end(), 0, [](int64_t total, const auto& obj) { + // Include the 'sizeof' to account for the variable number in the vector. + return total + sizeof(BSONObj) + obj.objsize(); + }); +} + +int64_t size(const boost::optional<PassthroughToShardOptions>& passthroughToShardOpts) { + if (!passthroughToShardOpts) { + return 0; + } + return passthroughToShardOpts->getShard().size(); +} + +int64_t size(const boost::optional<ExchangeSpec>& exchange) { + if (!exchange) { + return 0; + } + return sum( + {exchange->getKey().objsize(), + (exchange->getBoundaries() ? size(exchange->getBoundaries().get()) : 0), + (exchange->getConsumerIds() ? 4 * static_cast<int64_t>(exchange->getConsumerIds()->size()) + : 0)}); +} + +int64_t size(const boost::optional<EncryptionInformation>& encryptInfo) { + if (!encryptInfo) { + return 0; + } + tasserted(7659700, + "Unexpected encryption information - not expecting to collect query shape stats on " + "encrypted querys"); +} + +int64_t size(const StringData& str) { + return str.size(); +} + +int64_t size(const boost::optional<BSONObj>& obj) { + return optionalObjSize(obj); +} + +// variadic base case. +template <typename T> +int64_t sumOfSizes(const T& t) { + return size(t); +} + +// variadic recursive case. Making the compiler expand the pluses everywhere to give us good +// formatting at the call site. sumOfSizes(x, y, z) rather than size(x) + size(y) + size(z). +template <typename T, typename... Args> +int64_t sumOfSizes(const T& t, const Args&... args) { + return size(t) + sumOfSizes(args...); +} + +int64_t aggRequestSize(const AggregateCommandRequest& request) { + return sumOfSizes(request.getPipeline(), + request.getLet(), + request.getUnwrappedReadPref(), + request.getExchange(), + request.getPassthroughToShard(), + request.getEncryptionInformation(), + request.getDbName()); +} + +} // namespace + +int64_t AggregateKeyGenerator::doGetSize() const { + return sum({sizeof(*this), + static_cast<int64_t>(_origNss.size()), + optionalObjSize(_initialQueryStatsKey), + aggRequestSize(_request)}); +} +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/aggregated_metric.h b/src/mongo/db/query/query_stats/aggregated_metric.h new file mode 100644 index 00000000000..2e933a17a98 --- /dev/null +++ b/src/mongo/db/query/query_stats/aggregated_metric.h @@ -0,0 +1,79 @@ +/** + * Copyright (C) 2023-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 <algorithm> +#include <cstdint> + +#include "mongo/base/string_data.h" +#include "mongo/bson/bsonobjbuilder.h" +#include "mongo/platform/decimal128.h" + +namespace mongo::query_stats { + +/** + * An aggregated metric stores a compressed view of data. It balances the loss of information + * with the reduction in required storage. + */ +struct AggregatedMetric { + + /** + * Aggregate an observed value into the metric. + */ + void aggregate(uint64_t val) { + sum += val; + max = std::max(val, max); + min = std::min(val, min); + sumOfSquares = sumOfSquares.add(Decimal128(val).multiply(Decimal128(val))); + } + + void appendTo(BSONObjBuilder& builder, const StringData& fieldName) const { + BSONObjBuilder metricsBuilder = builder.subobjStart(fieldName); + metricsBuilder.append("sum", (long long)sum); + metricsBuilder.append("max", (long long)max); + metricsBuilder.append("min", (long long)min); + metricsBuilder.append("sumOfSquares", sumOfSquares); + metricsBuilder.done(); + } + + uint64_t sum = 0; + // Default to the _signed_ maximum (which fits in unsigned range) because we cast to + // BSONNumeric when serializing. + uint64_t min = (uint64_t)std::numeric_limits<int64_t>::max; + uint64_t max = 0; + + /** + * The sum of squares along with (an externally stored) count will allow us to compute the + * variance/stddev. + */ + Decimal128 sumOfSquares{}; +}; + +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/find_key.cpp b/src/mongo/db/query/query_stats/find_key.cpp new file mode 100644 index 00000000000..437c75aecaa --- /dev/null +++ b/src/mongo/db/query/query_stats/find_key.cpp @@ -0,0 +1,69 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_stats/find_key.h" + +namespace mongo::query_stats { + +void FindCmdComponents::appendTo(BSONObjBuilder& bob, const SerializationOptions& opts) const { + + if (_hasField.allowPartialResults) { + bob.append(FindCommandRequest::kAllowPartialResultsFieldName, _allowPartialResults); + } + + // Fields for literal redaction. Adds batchSize, and noCursorTimeOut. + + if (_hasField.noCursorTimeout) { + bob.append(FindCommandRequest::kNoCursorTimeoutFieldName, _noCursorTimeout); + } + + // We don't store the specified batch size value since it doesn't matter. + // Provide an arbitrary literal long here. + tassert(7973602, + "Serialization policy not supported - original values have been discarded", + opts.literalPolicy != LiteralSerializationPolicy::kUnchanged); + + if (_hasField.batchSize) { + opts.appendLiteral(&bob, FindCommandRequest::kBatchSizeFieldName, 0ll); + } +} + +std::unique_ptr<FindCommandRequest> FindKey::reparse(OperationContext* opCtx) const { + auto fcr = + static_cast<const query_shape::FindCmdShape*>(universalComponents()._queryShape.get()) + ->toFindCommandRequest(); + if (_components._hasField.allowPartialResults) + fcr->setAllowPartialResults(_components._allowPartialResults); + if (_components._hasField.noCursorTimeout) + fcr->setNoCursorTimeout(_components._noCursorTimeout); + if (_components._hasField.batchSize) + fcr->setBatchSize(1ll); + return fcr; +} +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/find_key.h b/src/mongo/db/query/query_stats/find_key.h new file mode 100644 index 00000000000..8578a77e573 --- /dev/null +++ b/src/mongo/db/query/query_stats/find_key.h @@ -0,0 +1,152 @@ +/** + * Copyright (C) 2023-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 <memory> + +#include "mongo/db/namespace_string.h" +#include "mongo/db/operation_context.h" +#include "mongo/db/query/query_shape/find_cmd_shape.h" +#include "mongo/db/query/query_stats/key.h" + +namespace mongo::query_stats { + +struct FindCmdComponents : public SpecificKeyComponents { + FindCmdComponents(const FindCommandRequest* findCmd) + : _allowPartialResults(findCmd->getAllowPartialResults().value_or(false)), + _noCursorTimeout(findCmd->getNoCursorTimeout().value_or(false)), + _hasField() { + _hasField.batchSize = findCmd->getBatchSize().has_value(); + _hasField.allowPartialResults = findCmd->getAllowPartialResults().has_value(); + _hasField.noCursorTimeout = findCmd->getNoCursorTimeout().has_value(); + } + + std::size_t size() const { + return sizeof(FindCmdComponents); + } + + void HashValue(absl::HashState state) const final { + absl::HashState::combine( + std::move(state), _hasField, _allowPartialResults, _noCursorTimeout); + } + + void appendTo(BSONObjBuilder& bob, const SerializationOptions& opts) const; + + // Avoid using boost::optional here because it creates extra padding at the beginning of the + // struct. Since each QueryStatsEntry can have its own FindKey, it's better to + // minimize the struct's size as much as possible. + + // Preserved literal. + bool _allowPartialResults; + bool _noCursorTimeout; + + // This anonymous struct represents the presence of the member variables as C++ bit fields. + // In doing so, each of these boolean values takes up 1 bit instead of 1 byte. + struct HasField { + HasField() : batchSize(false), allowPartialResults(false), noCursorTimeout(false) {} + bool batchSize : 1; + bool allowPartialResults : 1; + bool noCursorTimeout : 1; + bool operator==(const HasField& other) const { + return batchSize == other.batchSize && + allowPartialResults == other.allowPartialResults && + noCursorTimeout == other.noCursorTimeout; + } + + } _hasField; + + template <typename H> + friend H AbslHashValue(H h, const HasField& hasField) { + return H::combine(std::move(h), + hasField.batchSize, + hasField.noCursorTimeout, + hasField.allowPartialResults); + } +}; + +// This static assert checks to ensure that the struct's size is changed thoughtfully. If adding +// or otherwise changing the members, this assert may be updated with care. +static_assert( + // Expecting two bytes for allowPartialResults and noCursorTimeout, and another + // byte for _hasField. For alignment reasons (alignment is 8 bytes here), this means the trailer + // will bring up the total bytecount to a multiple of 8. + sizeof(FindCmdComponents) <= sizeof(SpecificKeyComponents) + 8, + "Size of FindCmdComponents is too large! " + "Make sure that the struct has been align- and padding-optimized. " + "If the struct's members have changed, this assert may need to be updated with a new " + "value."); + +class FindKey final : public Key { +public: + FindKey(const boost::intrusive_ptr<ExpressionContext>& expCtx, + const ParsedFindCommand& request, + query_shape::CollectionType collectionType = query_shape::CollectionType::kUnknown) + : Key(expCtx->opCtx, + std::make_unique<query_shape::FindCmdShape>(request, expCtx), + request.findCommandRequest->getHint(), + request.findCommandRequest->getReadConcern(), + request.findCommandRequest->getMaxTimeMS().has_value(), + collectionType), + _components(request.findCommandRequest.get()) {} + + // The default implementation of hashing for smart pointers is not a good one for our purposes. + // Here we overload them to actually take the hash of the object, rather than hashing the + // pointer itself. + template <typename H> + friend H AbslHashValue(H h, const std::unique_ptr<const FindKey>& key) { + return H::combine(std::move(h), *key); + } + template <typename H> + friend H AbslHashValue(H h, const std::shared_ptr<const FindKey>& key) { + return H::combine(std::move(h), *key); + } + + const SpecificKeyComponents& specificComponents() const { + return _components; + } + +private: + void appendCommandSpecificComponents(BSONObjBuilder& bob, + const SerializationOptions& opts) const final { + _components.appendTo(bob, opts); + } + + std::unique_ptr<FindCommandRequest> reparse(OperationContext* opCtx) const; + + FindCmdComponents _components; +}; + +// This static assert checks to ensure that the struct's size is changed thoughtfully. If adding +// or otherwise changing the members, this assert may be updated with care. +static_assert(sizeof(FindKey) == sizeof(Key) + sizeof(FindCmdComponents), + "If the class' members have changed, this assert may need to be updated with a new " + "value and the size calcuation will need to be changed."); + +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/find_key_test.cpp b/src/mongo/db/query/query_stats/find_key_test.cpp new file mode 100644 index 00000000000..6c34ba5a606 --- /dev/null +++ b/src/mongo/db/query/query_stats/find_key_test.cpp @@ -0,0 +1,133 @@ +/** + * Copyright (C) 2022-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. + */ + +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/parsed_find_command.h" +#include "mongo/db/query/query_stats/find_key.h" +#include "mongo/db/service_context_test_fixture.h" +#include "mongo/unittest/unittest.h" + +namespace mongo::query_stats { + +namespace { +static const NamespaceString kDefaultTestNss = NamespaceString("testDB.testColl"); + +static constexpr auto collectionType = query_shape::CollectionType::kCollection; + +class FindKeyTest : public ServiceContextTest { +public: + static std::unique_ptr<const Key> makeFindKeyFromQuery(const BSONObj& filter) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + auto fcr = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcr->setFilter(filter.getOwned()); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(expCtx, {std::move(fcr)})); + return std::make_unique<FindKey>(expCtx, *parsedFind, collectionType); + } +}; + +TEST_F(FindKeyTest, SizeOfFindCmdComponents) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + + auto fcr = std::make_unique<FindCommandRequest>(kDefaultTestNss); + auto query = BSON("query" << 1 << "xEquals" << 42); + fcr->setFilter(query.getOwned()); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(expCtx, {std::move(fcr)})); + auto findComponents = std::make_unique<FindCmdComponents>(parsedFind->findCommandRequest.get()); + + ASSERT_GTE(findComponents->size(), sizeof(SpecificKeyComponents) + 3 /*bools and HasField*/); + ASSERT_LTE(findComponents->size(), + sizeof(SpecificKeyComponents) + 8 /*bools, HasField, and padding*/); +} + +TEST_F(FindKeyTest, EquivalentFindCmdComponentsSizes) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + auto query = BSON("query" << 1 << "xEquals" << 42); + + // Set different fields in the find commands. + auto fcrCursorTimeout = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcrCursorTimeout->setFilter(query.getOwned()); + fcrCursorTimeout->setNoCursorTimeout(true); + auto parsedFindCursorTimeout = + uassertStatusOK(parsed_find_command::parse(expCtx, {std::move(fcrCursorTimeout)})); + auto findComponentsCursorTimeout = + std::make_unique<FindCmdComponents>(parsedFindCursorTimeout->findCommandRequest.get()); + + auto fcrAllowPartial = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcrAllowPartial->setFilter(query.getOwned()); + fcrAllowPartial->setAllowPartialResults(true); + auto parsedFindAllowPartial = + uassertStatusOK(parsed_find_command::parse(expCtx, {std::move(fcrAllowPartial)})); + auto findComponentsAllowPartial = + std::make_unique<FindCmdComponents>(parsedFindAllowPartial->findCommandRequest.get()); + + ASSERT_EQ(findComponentsCursorTimeout->size(), findComponentsAllowPartial->size()); +} + +// Testing item from opCtx that should impact key size. +TEST_F(FindKeyTest, SizeOfFindKeyWithAndWithoutComment) { + auto query = BSON("query" << 1 << "xEquals" << 42); + + auto keyWithoutComment = makeFindKeyFromQuery(query); + + auto opCtx = makeOperationContext(); + auto fcrWithComment = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcrWithComment->setFilter(query.getOwned()); + opCtx->setComment(BSON("comment" + << " foo")); + auto expCtxWithComment = make_intrusive<ExpressionContext>( + opCtx.get(), *fcrWithComment, nullptr, true /* mayDbProfile*/); + auto parsedFindWithComment = + uassertStatusOK(parsed_find_command::parse(expCtxWithComment, {std::move(fcrWithComment)})); + auto keyWithComment = std::make_unique<query_stats::FindKey>( + expCtxWithComment, *parsedFindWithComment, collectionType); + + ASSERT_LT(keyWithoutComment->size(), keyWithComment->size()); +} + +// Testing item from command request that should impact key size. +TEST_F(FindKeyTest, SizeOfFindKeyWithAndWithoutReadConcern) { + auto query = BSON("query" << 1 << "xEquals" << 42); + + auto keyWithoutReadConcern = makeFindKeyFromQuery(query); + + auto expCtxWithReadConcern = make_intrusive<ExpressionContextForTest>(); + auto fcrWithReadConcern = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcrWithReadConcern->setFilter(query.getOwned()); + fcrWithReadConcern->setReadConcern(fromjson(R"({level: "local"})")); + auto parsedFindWithReadConcern = uassertStatusOK( + parsed_find_command::parse(expCtxWithReadConcern, {std::move(fcrWithReadConcern)})); + auto keyWithReadConcern = std::make_unique<query_stats::FindKey>( + expCtxWithReadConcern, *parsedFindWithReadConcern, collectionType); + + ASSERT_LT(keyWithoutReadConcern->size(), keyWithReadConcern->size()); +} + + +} // namespace +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/key.cpp b/src/mongo/db/query/query_stats/key.cpp new file mode 100644 index 00000000000..f282ef21a2a --- /dev/null +++ b/src/mongo/db/query/query_stats/key.cpp @@ -0,0 +1,223 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/query_stats/key.h" + +#include "mongo/db/query/query_stats/query_stats_helpers.h" +#include "mongo/rpc/metadata/client_metadata.h" + +namespace mongo::query_stats { + +namespace { + +BSONObj scrubHighCardinalityFields(const ClientMetadata* clientMetadata) { + if (!clientMetadata) { + return BSONObj(); + } + return clientMetadata->documentWithoutMongosInfo(); +} + +BSONObj shapifyReadPreference(boost::optional<BSONObj> readPreference) { + if (!readPreference) { + return BSONObj(); + } + + BSONObjBuilder builder; + for (const auto& elem : *readPreference) { + if (elem.fieldNameStringData() != "tags"_sd) { + builder.append(elem); + continue; + } + + // Sort the $readPreference tags so that different orderings still map to one query stats + // store key. + BSONObjSet sortedTags = SimpleBSONObjComparator::kInstance.makeBSONObjSet(); + for (const auto& tag : elem.Array()) { + sortedTags.insert(tag.Obj()); + } + + BSONArrayBuilder arrBuilder(builder.subarrayStart("tags"_sd)); + for (const auto& tag : sortedTags) { + arrBuilder.append(tag); + } + } + return builder.obj(); +} + +} // namespace + +UniversalKeyComponents::UniversalKeyComponents(std::unique_ptr<query_shape::Shape> queryShape, + const ClientMetadata* clientMetadata, + boost::optional<BSONObj> commentObj, + boost::optional<BSONObj> hint, + boost::optional<BSONObj> readPreference, + boost::optional<BSONObj> writeConcern, + boost::optional<BSONObj> readConcern, + std::unique_ptr<APIParameters> apiParams, + query_shape::CollectionType collectionType, + bool maxTimeMS) + : _clientMetaData(scrubHighCardinalityFields(clientMetadata)), + _commentObj(commentObj.value_or(BSONObj()).getOwned()), + _hintObj(hint.value_or(BSONObj()).getOwned()), + _writeConcern(writeConcern.value_or(BSONObj()).getOwned()), + _shapifiedReadPreference(shapifyReadPreference(readPreference)), + _shapifiedReadConcern(shapifyReadConcern(readConcern.value_or(BSONObj()))), + _comment(commentObj ? _commentObj.firstElement() : BSONElement()), + _queryShape(std::move(queryShape)), + _apiParams(std::move(apiParams)), + _clientMetaDataHash(clientMetadata ? clientMetadata->hashWithoutMongosInfo() + : simpleHash(BSONObj())), + _collectionType(collectionType), + _hasField() { + _hasField.clientMetaData = bool(clientMetadata); + _hasField.comment = bool(commentObj); + _hasField.hint = bool(hint); + _hasField.readPreference = bool(readPreference); + _hasField.writeConcern = bool(writeConcern); + _hasField.readConcern = bool(readConcern); + _hasField.maxTimeMS = maxTimeMS; + tassert(7973600, "shape must not be null", _queryShape); +} + +BSONObj UniversalKeyComponents::shapifyReadConcern(const BSONObj& readConcern, + const SerializationOptions& opts) { + // Read concern should not be considered a literal. + // afterClusterTime is distinct for every operation with causal consistency enabled. We + // normalize it in order not to blow out the queryStats store cache. + if (readConcern["afterClusterTime"].eoo() && readConcern["atClusterTime"].eoo()) { + return readConcern.copy(); + } else { + BSONObjBuilder bob; + + if (auto levelElem = readConcern["level"]) { + bob.append(levelElem); + } + if (auto afterClusterTime = readConcern["afterClusterTime"]) { + opts.appendLiteral(&bob, "afterClusterTime", afterClusterTime); + } + if (auto atClusterTime = readConcern["atClusterTime"]) { + opts.appendLiteral(&bob, "atClusterTime", atClusterTime); + } + return bob.obj(); + } +} + +size_t UniversalKeyComponents::size() const { + return sizeof(*this) + _queryShape->size() + + (_apiParams ? sizeof(*_apiParams) + shape_helpers::optionalSize(_apiParams->getAPIVersion()) + : 0) + + _hintObj.objsize() + (_hasField.clientMetaData ? _clientMetaData.objsize() : 0) + + _commentObj.objsize() + + (_hasField.readPreference ? _shapifiedReadPreference.objsize() : 0) + + (_hasField.readConcern ? _shapifiedReadConcern.objsize() : 0) + + (_hasField.writeConcern ? _writeConcern.objsize() : 0); +} + +void UniversalKeyComponents::appendTo(BSONObjBuilder& bob, const SerializationOptions& opts) const { + if (_hasField.comment) { + opts.appendLiteral(&bob, "comment", _comment); + } + + if (_hasField.readConcern) { + auto readConcernToAppend = _shapifiedReadConcern; + if (opts != SerializationOptions::kRepresentativeQueryShapeSerializeOptions) { + // The options aren't the same as the first time we shapified, so re-computation is + // necessary (e.g. use "?timestamp" instead of the representative Timestamp(0, 0)). + readConcernToAppend = shapifyReadConcern(_shapifiedReadConcern, opts); + } + bob.append("readConcern", readConcernToAppend); + } + + if (const auto& apiVersion = _apiParams->getAPIVersion()) { + bob.append("apiVersion", apiVersion.value()); + } + + if (const auto& apiStrict = _apiParams->getAPIStrict()) { + bob.append("apiStrict", apiStrict.value()); + } + + if (const auto& apiDeprecationErrors = _apiParams->getAPIDeprecationErrors()) { + bob.append("apiDeprecationErrors", apiDeprecationErrors.value()); + } + + if (_hasField.readPreference) { + bob.append("$readPreference", _shapifiedReadPreference); + } + + if (_hasField.writeConcern) { + bob.append("writeConcern", _writeConcern); + } + + if (_hasField.clientMetaData) { + bob.append("client", _clientMetaData); + } + if (_collectionType > query_shape::CollectionType::kUnknown) { + bob.append("collectionType", toStringData(_collectionType)); + } + if (!_hintObj.isEmpty()) { + bob.append("hint", shape_helpers::extractHintShape(_hintObj, opts)); + } + if (_hasField.maxTimeMS) { + opts.appendLiteral(&bob, "maxTimeMS", 0ll); + } +} +Key::Key(OperationContext* opCtx, + std::unique_ptr<query_shape::Shape> queryShape, + boost::optional<BSONObj> hint, + boost::optional<BSONObj> readConcern, + bool maxTimeMS, + query_shape::CollectionType collectionType) + : _universalComponents( + std::move(queryShape), + ClientMetadata::get(opCtx->getClient()), + opCtx->getCommentOwnedCopy(), + hint, + ReadPreferenceSetting::get(opCtx).usedDefaultReadPrefValue() + ? boost::none + : boost::make_optional(ReadPreferenceSetting::get(opCtx).toInnerBSON()), + opCtx->getWriteConcern().isImplicitDefaultWriteConcern() + ? boost::none + : boost::make_optional(opCtx->getWriteConcern().toBSON()), + readConcern, + std::make_unique<APIParameters>(APIParameters::get(opCtx)), + collectionType, + maxTimeMS) {} + +BSONObj Key::toBson(OperationContext* opCtx, const SerializationOptions& opts) const { + BSONObjBuilder bob; + + // We'll take care of appending this one outside of the appendTo() call below since it needs + // an OperationContext in some re-parsing cases. The rest is simpler. + bob.append("queryShape", _universalComponents._queryShape->toBson(opCtx, opts)); + + _universalComponents.appendTo(bob, opts); + appendCommandSpecificComponents(bob, opts); + return bob.obj(); +} +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/key.h b/src/mongo/db/query/query_stats/key.h new file mode 100644 index 00000000000..bb83fca82ed --- /dev/null +++ b/src/mongo/db/query/query_stats/key.h @@ -0,0 +1,304 @@ +/** + * Copyright (C) 2023-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 <memory> + +#include "mongo/bson/bsonobj.h" +#include "mongo/db/api_parameters.h" +#include "mongo/db/collection_type.h" +#include "mongo/db/pipeline/expression_context.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_shape/serialization_options.h" +#include "mongo/db/query/query_shape/shape_helpers.h" +#include "mongo/db/query/query_stats/transform_algorithm_gen.h" +#include "mongo/rpc/metadata/client_metadata.h" + +namespace mongo::query_stats { + +/** + * A struct holding pieces of the command request that are a component of the query stats store key + * and are options/arguments to all supported query stats commands. + * + * This struct (and the SpecificKeyComponents) are split out as a separate inheritence hierarchy to + * make it easier to ensure each piece is hashed without sub-classes needing to enumerate the parent + * class's member variables. + */ +struct UniversalKeyComponents { + UniversalKeyComponents(std::unique_ptr<query_shape::Shape> queryShape, + const ClientMetadata* clientMetadata, + boost::optional<BSONObj> commentObj, + boost::optional<BSONObj> hint, + boost::optional<BSONObj> readPreference, + boost::optional<BSONObj> writeConcern, + boost::optional<BSONObj> readConcern, + std::unique_ptr<APIParameters> apiParams, + query_shape::CollectionType collectionType, + bool maxTimeMS); + /** + * Returns a copy of the read concern object. If there is an "afterClusterTime" or + * "atClusterTime" component, the timestamp is shapified according to 'opts'. + */ + static BSONObj shapifyReadConcern( + const BSONObj& readConcern, + const SerializationOptions& opts = + SerializationOptions::kRepresentativeQueryShapeSerializeOptions); + + size_t size() const; + + void appendTo(BSONObjBuilder& bob, const SerializationOptions& opts) const; + + // Avoid using boost::optional here because it creates extra padding at the beginning of the + // struct. Since each QueryStatsEntry has its own Key subclass, it's better to minimize + // the struct's size as much as possible. + + BSONObj _clientMetaData; // Preserve this value. + BSONObj _commentObj; // Shapify this value. + BSONObj _hintObj; // Preserve this value. + BSONObj _writeConcern; // Preserve this value. + + // Preserved literal except value of 'tags' field is sorted. + BSONObj _shapifiedReadPreference; + // Preserved literal except 'afterClusterTime' and 'atClusterTime' are shapified. + BSONObj _shapifiedReadConcern; + + // Separate the possibly-enormous BSONObj from the remaining members + + BSONElement _comment; + + std::unique_ptr<query_shape::Shape> _queryShape; + std::unique_ptr<APIParameters> _apiParams; // Preserve this value in the query shape. + + // Simple hash of the client metadata object. This value is stored separately because it is + // cached on the client to avoid re-computing on every operation. If no client metadata is + // present, this will be the hash of an empty BSON object (otherwise known as 0). + const unsigned long _clientMetaDataHash; + + // This value is not known when run a query is run on mongos over an unsharded collection, so it + // is not set through that code path. + query_shape::CollectionType _collectionType; + + // This anonymous struct represents the presence of the member variables as C++ bit fields. + // In doing so, each of these boolean values takes up 1 bit instead of 1 byte. + struct HasField { + HasField() + : clientMetaData(false), + comment(false), + hint(false), + readPreference(false), + writeConcern(false), + readConcern(false), + maxTimeMS(false) {} + + bool clientMetaData : 1; + bool comment : 1; + bool hint : 1; + bool readPreference : 1; + bool writeConcern : 1; + bool readConcern : 1; + bool maxTimeMS : 1; + } _hasField; +}; + +/** + * A base class for sub-classes to derive from to expose the hashing ability for all of their + * sub-components. + * + * This struct (and the UniversalKeyComponents) are split out as a separate inheritence hierarchy to + * make it easier to ensure each piece is hashed without sub-classes needing to enumerate the parent + * class's member variables. + */ +struct SpecificKeyComponents { + virtual ~SpecificKeyComponents() {} + + virtual void HashValue(absl::HashState state) const = 0; + + /** + * Sub-classes should implement this to report how much memory is used. This is important to do + * carefully since we are under a budget in the query stats store and use this to do the + * accounting. Implementers should include sizeof(*derivedThis) and be sure to also include the + * size of any owned pointer-like objects such as BSONObj or NamespaceString which are + * indirectly using memory elsehwhere. + * + * We cannot just use sizeof() because there are some variable size data members (like BSON + * objects) which depend on the particular instance. + */ + virtual size_t size() const = 0; +}; + +template <typename H> +H AbslHashValue(H state, const SpecificKeyComponents& value) { + value.HashValue(absl::HashState::Create(&state)); + return std::move(state); +} + +template <typename H> +H AbslHashValue(H h, const UniversalKeyComponents& components) { + return H::combine(std::move(h), + *components._queryShape, + components._clientMetaDataHash, + // Note we use the comment's type in the hash function. + components._comment.type(), + simpleHash(components._hintObj), + simpleHash(components._shapifiedReadPreference), + simpleHash(components._writeConcern), + simpleHash(components._shapifiedReadConcern), + components._apiParams ? APIParameters::Hash{}(*components._apiParams) : 0, + components._collectionType, + components._hasField); +} + +template <typename H> +H AbslHashValue(H h, const UniversalKeyComponents::HasField& hasField) { + return H::combine(std::move(h), + hasField.clientMetaData, + hasField.comment, + hasField.hint, + hasField.readPreference, + hasField.writeConcern, + hasField.readConcern, + hasField.maxTimeMS); +} + + +// This static assert checks to ensure that the struct's size is changed thoughtfully. If adding +// or otherwise changing the members, this assert may be updated with care. +static_assert( + sizeof(UniversalKeyComponents) <= sizeof(query_shape::Shape) + 6 * sizeof(BSONObj) + + sizeof(BSONElement) + sizeof(std::unique_ptr<APIParameters>) + + sizeof(query_shape::CollectionType) + sizeof(query_shape::QueryShapeHash) + + sizeof(int64_t), + "Size of Key is too large! " + "Make sure that the struct has been align- and padding-optimized. " + "If the struct's members have changed, this assert may need to be updated with a new value."); + +/** + * An abstract base class representing a query stats store key for a given request. All query stats + * store entries should include some common elements, tracked in `_universalComponents`. For + * example, everything tracked must have a `query_shape::Shape`. + * + * Subclasses can add more components to include as discriminating factors in which entries should + * be tracked separately. For example, two find commands which are identical except in their read + * concern should be tracked differently. Maybe they will have quite different performance + * characteristics or help us determine when the read concern was changed by the client. + * + * The interface to do this is to split out the state/memory for these components as a separate + * struct which can indpendently hash itself and compute its size (both of which are important for + * the query stats store). Subclasses of Key itself should not have any meaningfully sized + * state other than the 'specificComponents().' + */ +class Key { +public: + virtual ~Key() = default; + + /** + * All Keys will share these characteristics as part of their query stats store key. + * Returns an unowned reference so the caller must ensure the result does not outlive this + * Key instance. + */ + const auto& universalComponents() const { + return _universalComponents; + } + + /** + * Different commands will have different components they want to be included in the query stats + * store key. This interface allows them to do so and easily have those components incorporated + * into this key generation and hashing. + */ + virtual const SpecificKeyComponents& specificComponents() const = 0; + + /** + * Materializes the query stats store key. Not expected to be used on ingestion, since we should + * store this object and its components directly in their native C++ data structures - we can + * use the absl::Hash<query_stats::Key>{}() API to look them up. Instead, this may be useful to + * display the key (as it is used for $queryStats) or perhaps one day persist it to storage. + */ + BSONObj toBson(OperationContext* opCtx, const SerializationOptions& opts) const; + + /** + * Convenience function. + */ + query_shape::QueryShapeHash getQueryShapeHash(OperationContext* opCtx) const { + // TODO (future ticket?) should we cache this somewhere else? + return _universalComponents._queryShape->sha256Hash(opCtx); + } + + size_t size() const { + return sizeof(Key) + specificComponents().size() + _universalComponents.size(); + } + + template <typename H> + friend H AbslHashValue(H h, const Key& key) { + return H::combine(std::move(h), key._universalComponents, key.specificComponents()); + } + + // The default implementation of hashing for smart pointers is not a good one for our purposes. + // Here we overload them to actually take the hash of the object, rather than hashing the + // pointer itself. + template <typename H> + friend H AbslHashValue(H h, const std::unique_ptr<const Key>& key) { + return H::combine(std::move(h), *key); + } + template <typename H> + friend H AbslHashValue(H h, const std::shared_ptr<const Key>& key) { + return H::combine(std::move(h), *key); + } + +protected: + /** + * Sub-classes can use this to instantiate a 'real' Key. 'queryShape' must not be null, + * but is tracked as a pointer since it is a virtual class and we want to own it here. + */ + Key(OperationContext* opCtx, + std::unique_ptr<query_shape::Shape> queryShape, + boost::optional<BSONObj> hint, + boost::optional<BSONObj> readConcern, + bool maxTimeMS, + query_shape::CollectionType collectionType = query_shape::CollectionType::kUnknown); + + /** + * With a given BSONObjBuilder, append the command-specific components of the query stats key. + * + * You may be wondering why this API is here rather than as a virtual method on + * CmdSpecificComponents - and that would be because many implementations can involve a re-parse + * of the request if it needs to serialize with different serialization options. This re-parsing + * process often needs the context of things tracked in _universalComponents, which is hard to + * access from the specific components. + */ + virtual void appendCommandSpecificComponents(BSONObjBuilder& bob, + const SerializationOptions& opts) const = 0; + +private: + UniversalKeyComponents _universalComponents; +}; +static_assert( + sizeof(Key) == sizeof(void*) /*vtable ptr*/ + sizeof(UniversalKeyComponents), + "If the class' members have changed, this assert may need to be updated with a new value."); +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/key_test.cpp b/src/mongo/db/query/query_stats/key_test.cpp new file mode 100644 index 00000000000..69359d08bda --- /dev/null +++ b/src/mongo/db/query/query_stats/key_test.cpp @@ -0,0 +1,177 @@ +/** + * Copyright (C) 2024-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. + */ + +#include "mongo/bson/bsonelement.h" +#include "mongo/db/collection_type.h" +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/parsed_find_command.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_stats/find_key.h" +#include "mongo/db/query/query_stats/key.h" +#include "mongo/db/service_context_test_fixture.h" +#include "mongo/unittest/unittest.h" + +namespace mongo::query_stats { + +namespace { +static const NamespaceString kDefaultTestNss = NamespaceString("testDB.testColl"); + + +struct DummyShapeSpecificComponents : public query_shape::CmdSpecificShapeComponents { + DummyShapeSpecificComponents(){}; + void HashValue(absl::HashState state) const {} + size_t size() const final { + return sizeof(DummyShapeSpecificComponents); + } +}; + +class DummyShape : public query_shape::Shape { +public: + DummyShape(NamespaceStringOrUUID nssOrUUID, + BSONObj collation, + DummyShapeSpecificComponents dummyComponents) + : Shape(nssOrUUID, collation) { + components = dummyComponents; + } + + const query_shape::CmdSpecificShapeComponents& specificComponents() const final { + return components; + } + + void appendCmdSpecificShapeComponents(BSONObjBuilder&, + OperationContext*, + const SerializationOptions& opts) const final {} + DummyShapeSpecificComponents components; +}; + +struct DummyKeyComponents : public SpecificKeyComponents { + DummyKeyComponents(){}; + + void HashValue(absl::HashState state) const {} + size_t size() const { + return sizeof(DummyKeyComponents); + } +}; + +class DummyKey : public Key { +public: + DummyKey(OperationContext* opCtx, + std::unique_ptr<query_shape::Shape> queryShape, + boost::optional<BSONObj> hint, + boost::optional<BSONObj> readConcern, + bool maxTimeMS, + query_shape::CollectionType collectionType, + DummyKeyComponents dummyComponents) + : Key(opCtx, std::move(queryShape), hint, readConcern, maxTimeMS, collectionType) { + components = dummyComponents; + } + const SpecificKeyComponents& specificComponents() const { + return components; + }; + void appendCommandSpecificComponents(BSONObjBuilder& bob, + const SerializationOptions& opts) const {}; + DummyKeyComponents components; +}; +class UniversalKeyTest : public ServiceContextTest {}; + +TEST_F(UniversalKeyTest, SizeOfUniversalComponents) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + + // Make shape for testing. + auto collation = BSONObj{}; + auto innerComponents = std::make_unique<DummyShapeSpecificComponents>(); + auto shape = std::make_unique<DummyShape>(kDefaultTestNss, collation, *innerComponents); + + // Gather sizes and create universalComponents. + const auto shapeSize = shape->size(); + auto clientMetadata = ClientMetadata::get(expCtx->opCtx->getClient()); + + auto clientMetadataSize = clientMetadata ? clientMetadata->documentWithoutMongosInfo().objsize() + : BSONObj().objsize(); + + auto apiParams = std::make_unique<APIParameters>(APIParameters::get(expCtx->opCtx)); + const auto apiParamsSize = static_cast<size_t>( + apiParams ? sizeof(*apiParams) + shape_helpers::optionalSize(apiParams->getAPIVersion()) + : 0); + auto universalComponents = + std::make_unique<UniversalKeyComponents>(std::move(shape), + clientMetadata, + BSONObj(), + BSONObj(), + BSONObj(), + BSONObj(), + BSONObj(), + std::move(apiParams), + query_shape::CollectionType::kUnknown, + true); + + const auto minimumUniversalKeyComponentSize = sizeof(std::unique_ptr<query_shape::Shape>) + + (6 * sizeof(BSONObj)) + sizeof(std::unique_ptr<APIParameters>) + sizeof(BSONElement) + + sizeof(query_shape::CollectionType) + sizeof(unsigned long) + 1 /*HasField*/; + ASSERT_GTE(sizeof(UniversalKeyComponents), minimumUniversalKeyComponentSize); + ASSERT_LTE(sizeof(UniversalKeyComponents), minimumUniversalKeyComponentSize + 8 /*padding*/); + + ASSERT_GT(universalComponents->size(), + sizeof(UniversalKeyComponents) + shapeSize + clientMetadataSize + apiParamsSize); + ASSERT_LTE(universalComponents->size(), + sizeof(UniversalKeyComponents) + shapeSize + clientMetadataSize + + (5 * static_cast<size_t>(BSONObj().objsize())) + apiParamsSize); +} + +TEST_F(UniversalKeyTest, SizeOfSpecificComponents) { + auto innerComponents = std::make_unique<DummyShapeSpecificComponents>(); + auto keyComponents = std::make_unique<DummyKeyComponents>(); + + ASSERT_EQ(keyComponents->size(), sizeof(SpecificKeyComponents)); + ASSERT_EQ(sizeof(SpecificKeyComponents), sizeof(void*) /*vtable ptr*/); +} + +TEST_F(UniversalKeyTest, SizeOfKey) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + + auto collation = BSONObj{}; + auto innerComponents = std::make_unique<DummyShapeSpecificComponents>(); + auto shape = std::make_unique<DummyShape>(kDefaultTestNss, collation, *innerComponents); + + auto keyComponents = std::make_unique<DummyKeyComponents>(); + + auto key = std::make_unique<DummyKey>(expCtx->opCtx, + std::move(shape), + BSONObj(), + BSONObj(), + false, + query_shape::CollectionType::kUnknown, + *keyComponents); + ASSERT_EQ(innerComponents->size(), key->specificComponents().size()); + ASSERT_EQ(sizeof(Key), sizeof(UniversalKeyComponents) + sizeof(void*)); + ASSERT_EQ(key->size(), + sizeof(Key) + key->universalComponents().size() + key->specificComponents().size()); +} +} // namespace +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/query_stats.cpp b/src/mongo/db/query/query_stats/query_stats.cpp new file mode 100644 index 00000000000..a8bd49e0533 --- /dev/null +++ b/src/mongo/db/query/query_stats/query_stats.cpp @@ -0,0 +1,467 @@ +/** + * Copyright (C) 2022-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. + */ + +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQueryStats + +#include "mongo/db/query/query_stats/query_stats.h" + +#include "mongo/crypto/hash_block.h" +#include "mongo/db/concurrency/d_concurrency.h" +#include "mongo/db/concurrency/locker.h" +#include "mongo/db/curop.h" +#include "mongo/db/exec/projection_executor_builder.h" +#include "mongo/db/namespace_string.h" +#include "mongo/db/pipeline/aggregate_command_gen.h" +#include "mongo/db/pipeline/process_interface/stub_mongo_process_interface.h" +#include "mongo/db/query/find_command_gen.h" +#include "mongo/db/query/plan_explainer.h" +#include "mongo/db/query/projection_ast_util.h" +#include "mongo/db/query/projection_parser.h" +#include "mongo/db/query/query_feature_flags_gen.h" +#include "mongo/db/query/query_planner_params.h" +#include "mongo/db/query/query_request_helper.h" +#include "mongo/db/query/query_shape/serialization_options.h" +#include "mongo/db/query/query_stats/query_stats_on_parameter_change.h" +#include "mongo/db/query/sort_pattern.h" +#include "mongo/logv2/log.h" +#include "mongo/rpc/metadata/client_metadata.h" +#include "mongo/util/assert_util.h" +#include "mongo/util/debug_util.h" +#include "mongo/util/processinfo.h" +#include "mongo/util/system_clock_source.h" +#include <optional> + +namespace mongo::query_stats { + +Counter64 queryStatsStoreSizeEstimateBytesMetric; +ServerStatusMetricField<Counter64> displaySizeEstimateMetric( + "queryStats.queryStatsStoreSizeEstimateBytes", &queryStatsStoreSizeEstimateBytesMetric); + + +const Decorable<ServiceContext>::Decoration<std::unique_ptr<QueryStatsStoreManager>> + QueryStatsStoreManager::get = + ServiceContext::declareDecoration<std::unique_ptr<QueryStatsStoreManager>>(); + +const Decorable<ServiceContext>::Decoration<std::unique_ptr<RateLimiting>> + QueryStatsStoreManager::getRateLimiter = + ServiceContext::declareDecoration<std::unique_ptr<RateLimiting>>(); + + +namespace { + +Counter64 queryStatsEvictedMetric; +ServerStatusMetricField<Counter64> displayEvictedMetric("queryStats.numEvicted", + &queryStatsEvictedMetric); +Counter64 queryStatsRateLimitedRequestsMetric; +ServerStatusMetricField<Counter64> displayRateLimitMetric("queryStats.numRateLimitedRequests", + &queryStatsRateLimitedRequestsMetric); +Counter64 queryStatsStoreWriteErrorsMetric; +ServerStatusMetricField<Counter64> displayWriteErrorsMetric( + "queryStats.numQueryStatsStoreWriteErrors", &queryStatsStoreWriteErrorsMetric); + +/** + * Indicates whether or not query stats is enabled via the feature flag. + */ +bool isQueryStatsFeatureEnabled() { + // We need to call isVersionInitialized() first because this could run during startup while the + // FCV is still uninitialized. + if (serverGlobalParams.featureCompatibility.isVersionInitialized()) { + return feature_flags::gFeatureFlagQueryStats.isEnabled( + serverGlobalParams.featureCompatibility); + } + // (Generic FCV reference): This reference is needed to ensure we correctly initialize query + // stats during startup. + return feature_flags::gFeatureFlagQueryStats.isEnabledOnVersion( + multiversion::GenericFCV::kLatest); +} + +/** + * Cap the queryStats store size. + */ +size_t capQueryStatsStoreSize(size_t requestedSize) { + size_t cappedStoreSize = memory_util::capMemorySize( + requestedSize /*requestedSizeBytes*/, 1 /*maximumSizeGB*/, 25 /*percentTotalSystemMemory*/); + // If capped size is less than requested size, the queryStats store has been capped at its + // upper limit. + if (cappedStoreSize < requestedSize) { + LOGV2_DEBUG(7106502, + 1, + "The queryStats store size has been capped", + "cappedSize"_attr = cappedStoreSize); + } + return cappedStoreSize; +} + +/** + * Get the queryStats store size based on the query job's value. + */ +size_t getQueryStatsStoreSize() { + auto status = memory_util::MemorySize::parse(internalQueryStatsCacheSize.get()); + uassertStatusOK(status); + size_t requestedSize = memory_util::convertToSizeInBytes(status.getValue()); + return capQueryStatsStoreSize(requestedSize); +} + +void assertConfigurationAllowed() { + uassert(ErrorCodes::QueryFeatureNotAllowed, + "Cannot configure queryStats store. The feature flag is not enabled. Please restart " + "and specify the feature flag, or upgrade the feature compatibility version to one " + "where it is enabled by default.", + isQueryStatsFeatureEnabled()); +} + +class QueryStatsOnParamChangeUpdaterImpl final : public query_stats_util::OnParamChangeUpdater { +public: + void updateCacheSize(ServiceContext* serviceCtx, memory_util::MemorySize memSize) final { + assertConfigurationAllowed(); + auto requestedSize = memory_util::convertToSizeInBytes(memSize); + auto cappedSize = capQueryStatsStoreSize(requestedSize); + auto& queryStatsStoreManager = QueryStatsStoreManager::get(serviceCtx); + size_t numEvicted = queryStatsStoreManager->resetSize(cappedSize); + queryStatsEvictedMetric.increment(numEvicted); + } + + void updateSamplingRate(ServiceContext* serviceCtx, int samplingRate) { + assertConfigurationAllowed(); + QueryStatsStoreManager::getRateLimiter(serviceCtx).get()->setSamplingRate(samplingRate); + } +}; + +ServiceContext::ConstructorActionRegisterer queryStatsStoreManagerRegisterer{ + "QueryStatsStoreManagerRegisterer", [](ServiceContext* serviceCtx) { + // Note: it is possible that this is called before FCV is properly set up. The feature flags + // can only be specified at startup, but the feature compatibility version may change at + // runtime. If the feature compatibility version upgrades at runtime, the feature may now be + // enabled by default, even if the flag was not specified. To allow for this possibility, we + // will always configure a query stats store of the size currently specified by + // 'internalQueryStatsCacheSize', but we will prevent changing its shape or rate limit at + // runtime unless the feature flag is enabled (at whatever current FCV when the + // configuration setParameter command is run). + + query_stats_util::queryStatsStoreOnParamChangeUpdater(serviceCtx) = + std::make_unique<QueryStatsOnParamChangeUpdaterImpl>(); + size_t size = getQueryStatsStoreSize(); + auto&& globalQueryStatsStoreManager = QueryStatsStoreManager::get(serviceCtx); + // Initially the queryStats store used the same number of partitions as the plan cache, that + // is the number of cpu cores. However, with performance investigation we found that when + // the size of the partitions was too large, it took too long to copy out and read one + // partition. We are now capping each partition at 16MB (the largest size a query shape can + // be. If that gives us fewer partitions than we have cores, we set it to match the + // number of cores. The size needs to be cast to a double since we want to round up the + // number of partitions, and therefore need to avoid int division. + size_t numPartitions = std::ceil(double(size) / (16 * 1024 * 1024)); + auto numLogicalCores = ProcessInfo::getNumCores(); + if (numPartitions < numLogicalCores) { + numPartitions = numLogicalCores; + } + + globalQueryStatsStoreManager = + std::make_unique<QueryStatsStoreManager>(size, numPartitions); + auto configuredSamplingRate = internalQueryStatsRateLimit.load(); + QueryStatsStoreManager::getRateLimiter(serviceCtx) = std::make_unique<RateLimiting>( + configuredSamplingRate < 0 ? INT_MAX : configuredSamplingRate, Seconds{1}); + }}; + +/** + * Top-level checks for whether queryStats collection is enabled. If this returns false, we must + * go no further. + */ +bool isQueryStatsEnabled(const ServiceContext* serviceCtx) { + // During initialization, FCV may not yet be setup but queries could be run. We can't + // check whether queryStats should be enabled without FCV, so default to not recording + // those queries. + return isQueryStatsFeatureEnabled() && + QueryStatsStoreManager::get(serviceCtx)->getMaxSize() > 0; +} + +/** + * Internal check for whether we should collect metrics. This checks the rate limiting + * configuration for a global on/off decision and, if enabled, delegates to the rate limiter. + */ +bool shouldCollect(const ServiceContext* serviceCtx) { + // Cannot collect queryStats if sampling rate is not greater than 0. Note that we do not + // increment queryStatsRateLimitedRequestsMetric here since queryStats is entirely disabled. + auto samplingRate = QueryStatsStoreManager::getRateLimiter(serviceCtx)->getSamplingRate(); + if (samplingRate <= 0) { + LOGV2_DEBUG(8473001, + 5, + "sampling rate is <= 0, skipping this request", + "samplingRate"_attr = samplingRate); + return false; + } + // Check if rate limiting allows us to collect queryStats for this request. + if (samplingRate < INT_MAX && + !QueryStatsStoreManager::getRateLimiter(serviceCtx)->handleRequestSlidingWindow()) { + queryStatsRateLimitedRequestsMetric.increment(); + LOGV2_DEBUG(8473002, + 5, + "rate limited this request", + "samplingRate"_attr = samplingRate, + "totalLimited"_attr = queryStatsRateLimitedRequestsMetric.get()); + return false; + } + return true; +} + +void updateStatistics(const QueryStatsStore::Partition& proofOfLock, + QueryStatsEntry& toUpdate, + const uint64_t queryExecMicros, + const uint64_t firstResponseExecMicros, + const uint64_t docsReturned) { + toUpdate.latestSeenTimestamp = Date_t::now(); + toUpdate.lastExecutionMicros = queryExecMicros; + toUpdate.execCount++; + toUpdate.totalExecMicros.aggregate(queryExecMicros); + toUpdate.firstResponseExecMicros.aggregate(firstResponseExecMicros); + toUpdate.docsReturned.aggregate(docsReturned); +} + +} // namespace + +void registerRequest(OperationContext* opCtx, + const NamespaceString& collection, + std::function<std::unique_ptr<Key>(void)> makeKey, + bool willNeverExhaust) { + if (!isQueryStatsEnabled(opCtx->getServiceContext())) { + LOGV2_DEBUG(8473000, + 5, + "not collecting query stats for this request since it is disabled", + "featureEnabled"_attr = isQueryStatsFeatureEnabled()); + return; + } + + // Queries against metadata collections should never appear in queryStats data. + if (collection.isFLE2StateCollection()) { + return; + } + + // Don't record queries from internal clients. + if (opCtx->getClient()->session() && + (opCtx->getClient()->session()->getTags() & transport::Session::kInternalClient)) { + return; + } + + auto& opDebug = CurOp::get(opCtx)->debug(); + + if (opDebug.queryStatsInfo.wasRateLimited) { + LOGV2_DEBUG( + 8288900, + 4, + "Query stats request was previously rate limited. We expect this is a query on a view"); + return; + } + + if (!shouldCollect(opCtx->getServiceContext())) { + opDebug.queryStatsInfo.wasRateLimited = true; + return; + } + + if (opDebug.queryStatsInfo.key) { + // A find() request may have already registered the shapifier. Ie, it's a find command over + // a non-physical collection, eg view, which is implemented by generating an agg pipeline. + LOGV2_DEBUG(7198700, + 2, + "Query stats request shapifier already registered", + "collection"_attr = collection); + return; + } + + opDebug.queryStatsInfo.willNeverExhaust = willNeverExhaust; + // There are a few cases where a query shape can be larger than the original query. For example, + // {$exists: false} in the input query serializes to {$not: {$exists: true}. In rare cases where + // an input query has thousands of clauses, the cumulative bloat that shapification adds results + // in a BSON object that exceeds the 16 MB memory limit. In these cases, we want to exclude the + // original query from queryStats metrics collection and let it execute normally. + try { + opDebug.queryStatsInfo.key = makeKey(); + } catch (const DBException& ex) { + queryStatsStoreWriteErrorsMetric.increment(); + + const auto status = ex.toStatus(); + if (status.code() == ErrorCodes::BSONObjectTooLarge) { + LOGV2_DEBUG(7979400, + 2, + "Query Stats shapification has exceeded the 16 MB memory limit. Metrics " + "will not be collected"); + return; + } + + const auto& cmdObj = CurOp::get(opCtx)->opDescription(); + LOGV2_DEBUG(9423100, + 2, + "Error encountered when creating the Query Stats store key. Metrics will not " + "be collected for this command", + "status"_attr = status, + "command"_attr = cmdObj); + if (kDebugBuild || internalQueryStatsErrorsAreCommandFatal.load()) { + // uassert rather than tassert so that we avoid creating fatal failures on queries that + // were going to fail anyway, but trigger the error here first. A query that ONLY fails + // when query stats is enabled will still be surfaced by the uassert. + // Note that in the former case, these queries will fail with a different error code + // than they would have otherwise. Since this block is only applicable in test + // environments, this is fine. We make this tradeoff because it is desirable to have + // real bugs clearly surfaced as query stats issues. + uasserted(9423101, + str::stream() << "Failed to create query stats store key. Status: " << status + << " Command: " << cmdObj); + } + + return; + } + opDebug.queryStatsInfo.keyHash = absl::Hash<query_stats::Key>{}(*opDebug.queryStatsInfo.key); + // TODO look up this query shape (sub-component of query stats store key) in some new shared + // data structure that the query settings component could share. See if the query SHAPE hash has + // been computed before. If so, record the query shape hash on the opDebug. If not, compute the + // hash and store it there so we can avoid re-doing this for each request. +} + +QueryStatsStore& getQueryStatsStore(OperationContext* opCtx) { + uassert(ErrorCodes::QueryFeatureNotAllowed, + "Query stats is not enabled without the feature flag on and a cache size greater than " + "0 bytes", + isQueryStatsEnabled(opCtx->getServiceContext())); + return QueryStatsStoreManager::get(opCtx->getServiceContext())->getQueryStatsStore(); +} + +void writeQueryStats(OperationContext* opCtx, + boost::optional<size_t> queryStatsKeyHash, + std::unique_ptr<Key> key, + const uint64_t queryExecMicros, + const uint64_t firstResponseExecMicros, + const uint64_t docsReturned, + bool willNeverExhaust) { + // Generally we expect a 'key' to write query stats. However, for a change stream query, we + // expect it has no 'key' after its first writeQueryStats(), but it must have a + // 'queryStatsKeyHash' for its entry to be updated. + // TODO SERVER-89058 Modify comment to include tailable cursors. + if (!key && !(willNeverExhaust && queryStatsKeyHash)) { + return; + } + + // It's possible that query stats was enabled in registerRequest but has been disabled since + // (e.g., by FCV downgrade or setting the store size to 0). Rather than calling + // getQueryStatsStore (which would trigger a uassert if queryStats is disabled), we return and + // log a message if query stats is disabled, and otherwise grab the query stats store directly. + if (!isQueryStatsEnabled(opCtx->getServiceContext())) { + LOGV2_DEBUG(8456700, + 2, + "Query stats was enabled when the command started but is now disabled. " + "Metrics will not be collected.", + "queryStatsKeyHash"_attr = queryStatsKeyHash); + return; + } + auto&& queryStatsStore = + QueryStatsStoreManager::get(opCtx->getServiceContext())->getQueryStatsStore(); + if (key) { + dassert(absl::Hash<query_stats::Key>{}(*key) == queryStatsKeyHash, + "Expecting query stats key to hash to the given hash. Is the OpCtx state being " + "incorrectly re-used?"); + } + auto&& [statusWithMetrics, partitionLock] = + queryStatsStore.getWithPartitionLock(*queryStatsKeyHash); + if (statusWithMetrics.isOK()) { + // Found an existing entry! Just update the metrics and we're done. + return updateStatistics(partitionLock, + *statusWithMetrics.getValue(), + queryExecMicros, + firstResponseExecMicros, + docsReturned); + } + + // It is possible a cursor that lives forever has no key associated with it and its entry may + // have been evicted. + if (willNeverExhaust && !key) { + return; + } + + // Otherwise we didn't find an existing entry. Try to create one. + tassert(7315200, + "key cannot be null when writing a new entry to the queryStats store", + key != nullptr); + size_t numEvicted = + queryStatsStore.put(*queryStatsKeyHash, QueryStatsEntry(std::move(key)), partitionLock); + queryStatsEvictedMetric.increment(numEvicted); + auto newMetrics = partitionLock->get(*queryStatsKeyHash); + if (!newMetrics.isOK()) { + // This can happen if the budget is immediately exceeded. Specifically if the there is + // not enough room for a single new entry if the number of partitions is too high + // relative to the size. + queryStatsStoreWriteErrorsMetric.increment(); + LOGV2_DEBUG(7560900, + 0, + "Failed to store queryStats entry.", + "status"_attr = newMetrics.getStatus(), + "queryStatsKeyHash"_attr = queryStatsKeyHash); + return; + } + + return updateStatistics(partitionLock, + newMetrics.getValue()->second, + queryExecMicros, + firstResponseExecMicros, + docsReturned); +} + +void writeQueryStatsOnCursorDisposeOrKill(OperationContext* opCtx, + boost::optional<size_t> queryStatsKeyHash, + std::unique_ptr<Key> key, + bool willNeverExhaust, + const uint64_t queryExecMicros, + const uint64_t firstResponseExecMicros, + const uint64_t docsReturned) { + // It is discouraged but technically possible for a user to enable queryStats on the mongods of + // a replica set. In this case, a cursor will be created for each mongod. However, the + // queryStatsKey is behind a unique_ptr on CurOp. The ClientCursor constructor std::moves the + // queryStatsKey so it uniquely owns it (and also makes the queryStatsKey on CurOp now a + // nullptr) and copies over the queryStatsKeyHash as the latter is a cheap copy. + // In the case of sharded $search, two cursors will be created per mongod. In this way, + // two cursors are part of the same thread/operation, and therefore share a OpCtx/CurOp/OpDebug. + // The first cursor that is created will own the queryStatsKey and have a copy of the + // queryStatsKeyHash. On the other hand, the second one will only have a copy of the hash since + // the queryStatsKey will be null on CurOp from being std::move'd in the first cursor + // construction call. To not trip the tassert in writeQueryStats and because all cursors are + // guaranteed to have a copy of the hash, we check that the cursor has a key + if (key && opCtx) { + query_stats::writeQueryStats(opCtx, + queryStatsKeyHash, + std::move(key), + queryExecMicros, + firstResponseExecMicros, + docsReturned, + willNeverExhaust); + } else if (willNeverExhaust && opCtx) { + // Since we already recorded information about the possible getMores associated with a + // cursor that never ends, the only information left to record is about the kill/dispose + // cursor operation. This operation is not timed and does not have any metrics associated + // with it. + query_stats::writeQueryStats(opCtx, queryStatsKeyHash, nullptr, 0, 0, 0, willNeverExhaust); + } +} + +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/query_stats.h b/src/mongo/db/query/query_stats/query_stats.h new file mode 100644 index 00000000000..fc96a8be179 --- /dev/null +++ b/src/mongo/db/query/query_stats/query_stats.h @@ -0,0 +1,211 @@ +/** + * Copyright (C) 2022-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 "mongo/base/status.h" +#include "mongo/bson/bsonobj.h" +#include "mongo/db/concurrency/d_concurrency.h" +#include "mongo/db/curop.h" +#include "mongo/db/namespace_string.h" +#include "mongo/db/query/partitioned_cache.h" +#include "mongo/db/query/plan_explainer.h" +#include "mongo/db/query/query_stats/key.h" +#include "mongo/db/query/query_stats/query_stats_entry.h" +#include "mongo/db/query/query_stats/rate_limiting.h" +#include "mongo/db/service_context.h" +#include "mongo/db/views/view.h" +#include <cstdint> +#include <memory> + +namespace mongo::query_stats { + +extern Counter64 queryStatsStoreSizeEstimateBytesMetric; + +struct QueryStatsPartitioner { + // The partitioning function for use with the 'Partitioned' utility. + std::size_t operator()(const std::size_t hash, const std::size_t nPartitions) const { + return hash % nPartitions; + } +}; + +struct QueryStatsStoreEntryBudgetor { + size_t operator()(const std::size_t hash, const QueryStatsEntry& value) { + return sizeof(decltype(value)) + sizeof(decltype(hash)) + value.key->size(); + } +}; + +/* + * 'QueryStatsStore insertion and eviction listener implementation. This class adjusts the + * 'queryStatsStoreSize' serverStatus metric when entries are inserted or evicted. + */ +struct QueryStatsStoreInsertionEvictionListener { + void onInsert(const std::size_t&, const QueryStatsEntry&, size_t estimatedSize) { + queryStatsStoreSizeEstimateBytesMetric.increment(estimatedSize); + } + + void onEvict(const std::size_t&, const QueryStatsEntry&, size_t estimatedSize) { + queryStatsStoreSizeEstimateBytesMetric.decrement(estimatedSize); + } + + void onClear(size_t estimatedSize) { + queryStatsStoreSizeEstimateBytesMetric.decrement(estimatedSize); + } +}; +using QueryStatsStore = PartitionedCache<std::size_t, + QueryStatsEntry, + QueryStatsStoreEntryBudgetor, + QueryStatsPartitioner, + QueryStatsStoreInsertionEvictionListener>; + +/** + * A manager for the queryStats store allows a "pointer swap" on the queryStats store itself. The + * usage patterns are as follows: + * + * - Updating the queryStats store uses the `getQueryStatsStore()` method. The queryStats store + * instance is obtained, entries are looked up and mutated, or created anew. + * - The queryStats store is "reset". This involves atomically allocating a new instance, once + * there are no more updaters (readers of the store "pointer"), and returning the existing + * instance. + */ +class QueryStatsStoreManager { +public: + // The query stats store can be configured using these objects on a per-ServiceContext level. + // This is essentially global, but can be manipulated by unit tests. + static const ServiceContext::Decoration<std::unique_ptr<QueryStatsStoreManager>> get; + static const ServiceContext::Decoration<std::unique_ptr<RateLimiting>> getRateLimiter; + + template <typename... QueryStatsStoreArgs> + QueryStatsStoreManager(size_t cacheSize, size_t numPartitions) + : _queryStatsStore(std::make_unique<QueryStatsStore>(cacheSize, numPartitions)), + _maxSize(cacheSize) {} + + /** + * Acquire the instance of the queryStats store. + */ + QueryStatsStore& getQueryStatsStore() { + return *_queryStatsStore; + } + + size_t getMaxSize() { + return _maxSize.load(); + } + + /** + * Resize the queryStats store and return the number of evicted + * entries. + */ + size_t resetSize(size_t cacheSize) { + _maxSize.store(cacheSize); + return _queryStatsStore->reset(cacheSize); + } + +private: + std::unique_ptr<QueryStatsStore> _queryStatsStore; + + /** + * Max size of the queryStats store. Tracked here to avoid having to recompute after it's + * divided up into partitions. + */ + AtomicWord<size_t> _maxSize; +}; + +/** + * Acquire a reference to the global queryStats store. + */ +QueryStatsStore& getQueryStatsStore(OperationContext* opCtx); + +/** + * Registers a request for query stats collection. The function may decide not to collect anything, + * so this should be called for all requests. The decision is made based on the feature flag and + * query stats rate limiting. + * + * The originating command/query does not persist through the end of query execution due to + * optimizations made to the original query and the expiration of OpCtx across getMores. In order + * to pair the query stats metrics that are collected at the end of execution with the original + * query, it is necessary to store the original query during planning and persist it through + * getMores. + * + * During planning, registerRequest is called to serialize the query stats key and save it to + * OpDebug. If a query's execution is complete within the original operation, + * collectQueryStatsMongod/collectQueryStatsMongos will call writeQueryStats() and pass along the + * query stats key to be saved in the query stats store alongside metrics collected. + * + * However, OpDebug does not persist through cursor iteration, so if a query's execution will span + * more than one request/operation, it's necessary to save the query stats context to the cursor + * upon cursor registration. In these cases, collectQueryStatsMongod/collectQueryStatsMongos will + * aggregate each operation's metrics within the cursor. Once the request is eventually complete, + * the cursor calls writeQueryStats() on its destruction. + * + * Notes: + * - It's important to call registerRequest with the original request, before canonicalizing or + * optimizing it, in order to preserve the user's input for the query shape. + * - Calling this affects internal state. It should be called exactly once for each request for + * which query stats may be collected. + * - The std::function argument to construct an abstracted Key is provided to break + * library cycles so this library does not need to know how to parse everything. It is done as a + * deferred construction callback to ensure that this feature does not impact performance if + * collecting stats is not needed due to the feature being disabled or the request being rate + * limited. + */ +void registerRequest(OperationContext* opCtx, + const NamespaceString& collection, + std::function<std::unique_ptr<Key>(void)> makeKey, + bool willNeverExhaust = false); + +/** + * Writes query stats to the query stats store for the operation identified by `queryStatsKeyHash`. + * + * Direct calls to writeQueryStats in new code should be avoided in favor of calling existing + * functions: + * - collectQueryStatsMongod/collectQueryStatsMongos in the case of requests that span one + * operation + * - writeQueryStatsOnCursorDisposeOrKill() in the case of requests that span + * multiple operations (via getMore) + */ +void writeQueryStats(OperationContext* opCtx, + boost::optional<size_t> queryStatsKeyHash, + std::unique_ptr<Key> key, + uint64_t queryExecMicros, + uint64_t firstResponseExecMicros, + uint64_t docsReturned, + bool willNeverExhaust = false); + +/** + * Called from ClientCursor::dispose/ClusterClientCursorImpl::kill to set up and writeQueryStats() + * at the end of life of a cursor. + */ +void writeQueryStatsOnCursorDisposeOrKill(OperationContext* opCtx, + boost::optional<size_t> queryStatsKeyHash, + std::unique_ptr<Key> key, + bool willNeverExhaust, + uint64_t queryExecMicros, + uint64_t firstResponseExecMicros, + uint64_t docsReturned); +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/plan_cache_size_parameter.cpp b/src/mongo/db/query/query_stats/query_stats_entry.cpp index 46e42efafbf..f69f0a6ee2a 100644 --- a/src/mongo/db/query/plan_cache_size_parameter.cpp +++ b/src/mongo/db/query/query_stats/query_stats_entry.cpp @@ -1,5 +1,5 @@ /** - * Copyright (C) 2021-present MongoDB, Inc. + * Copyright (C) 2023-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, @@ -27,48 +27,28 @@ * it in the license file. */ -#include "mongo/db/query/plan_cache_size_parameter.h" +#include "mongo/db/query/query_stats/query_stats_entry.h" -#include <pcrecpp.h> +#include <boost/optional.hpp> -#include "mongo/db/query/query_knobs_gen.h" +#include "mongo/crypto/hash_block.h" +#include "mongo/crypto/sha256_block.h" -namespace mongo::plan_cache_util { +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery -StatusWith<PlanCacheSizeUnits> parseUnitString(const std::string& strUnit) { - if (strUnit.empty()) { - return Status(ErrorCodes::Error{6007010}, "Unit value cannot be empty"); - } +namespace mongo::query_stats { - if (strUnit[0] == '%') { - return PlanCacheSizeUnits::kPercent; - } else if (strUnit[0] == 'M' || strUnit[0] == 'm') { - return PlanCacheSizeUnits::kMB; - } else if (strUnit[0] == 'G' || strUnit[0] == 'g') { - return PlanCacheSizeUnits::kGB; - } - - return Status(ErrorCodes::Error{6007011}, "Incorrect unit value"); +BSONObj QueryStatsEntry::toBSON() const { + BSONObjBuilder builder{sizeof(QueryStatsEntry) + 100}; + builder.append("lastExecutionMicros", (long long)lastExecutionMicros); + builder.append("execCount", (long long)execCount); + totalExecMicros.appendTo(builder, "totalExecMicros"); + firstResponseExecMicros.appendTo(builder, "firstResponseExecMicros"); + docsReturned.appendTo(builder, "docsReturned"); + builder.append("firstSeenTimestamp", firstSeenTimestamp); + builder.append("latestSeenTimestamp", latestSeenTimestamp); + return builder.obj(); } -StatusWith<PlanCacheSizeParameter> PlanCacheSizeParameter::parse(const std::string& str) { - pcrecpp::RE_Options opt; - opt.set_caseless(true); - // Looks for a floating point number with followed by a unit suffix (MB, GB, %). - pcrecpp::RE re("\\s*(\\d+\\.?\\d*)\\s*(MB|GB|%)\\s*", opt); - - double size{}; - std::string strUnit{}; - if (!re.FullMatch(str, &size, &strUnit)) { - return {ErrorCodes::Error{6007012}, "Unable to parse plan cache size string"}; - } - - auto statusWithUnit = parseUnitString(strUnit); - if (!statusWithUnit.isOK()) { - return statusWithUnit.getStatus(); - } - - return PlanCacheSizeParameter{size, statusWithUnit.getValue()}; -} -} // namespace mongo::plan_cache_util +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/query_stats_entry.h b/src/mongo/db/query/query_stats/query_stats_entry.h new file mode 100644 index 00000000000..6b61a6a6dcf --- /dev/null +++ b/src/mongo/db/query/query_stats/query_stats_entry.h @@ -0,0 +1,95 @@ +/** + * Copyright (C) 2023-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 <algorithm> +#include <cstdint> +#include <memory> + +#include "mongo/db/commands/server_status_metric.h" +#include "mongo/db/query/query_stats/aggregated_metric.h" +#include "mongo/db/query/query_stats/key.h" +#include "mongo/db/query/query_stats/transform_algorithm_gen.h" +#include "mongo/util/time_support.h" + +namespace mongo::query_stats { + +/** + * The value stored in the query stats store. It contains a Key representing this "kind" of + * query, and some metrics about that shape. This class is responsible for knowing its size and + * updating our server status metrics about the size of the query stats store accordingly. At the + * time of this writing, the LRUCache utility does not easily expose its size in a way we could use + * as server status metrics. + */ +struct QueryStatsEntry { + QueryStatsEntry(std::unique_ptr<const Key> key_) + : firstSeenTimestamp(Date_t::now()), key(std::move(key_)) {} + + BSONObj toBSON() const; + + /** + * Timestamp for when this query shape was added to the store. Set on construction. + */ + const Date_t firstSeenTimestamp; + + /** + * Timestamp for when the latest time this query shape was seen. + */ + Date_t latestSeenTimestamp; + + /** + * Last execution time in microseconds. + */ + uint64_t lastExecutionMicros = 0; + + /** + * Number of query executions. + */ + uint64_t execCount = 0; + + /** + * Aggregates the total time for execution including getMore requests. + */ + AggregatedMetric totalExecMicros; + + /** + * Aggregates the time for execution for first batch only. + */ + AggregatedMetric firstResponseExecMicros; + + AggregatedMetric docsReturned; + + /** + * The Key that can generate the query stats key for this request. + */ + std::shared_ptr<const Key> key; +}; + +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/query_stats_helpers.h b/src/mongo/db/query/query_stats/query_stats_helpers.h new file mode 100644 index 00000000000..6d53cc8d4ce --- /dev/null +++ b/src/mongo/db/query/query_stats/query_stats_helpers.h @@ -0,0 +1,52 @@ +/** + * Copyright (C) 2023-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 <absl/hash/hash.h> +#include <boost/optional.hpp> + +#include "mongo/bson/bsonobj.h" +#include "mongo/bson/simple_bsonobj_comparator.h" +#include "mongo/db/query/query_shape/shape_helpers.h" + +namespace mongo::query_stats { + +/** + * An abseil compatible hash function for BSONObjects. Note that this hasher ignores any collation + * and uses the "simple" comparisons. This is fine and correct for query stats, but this is + * intentionally placed within the 'query_stats' namespace to avoid polluting the whole codebase + * with this helper which could cause an accidental bug where we ignore the request's collation. + */ +template <typename H> +H AbslHashValue(H h, const BSONObj& obj) { + return H::combine(std::move(h), simpleHash(obj)); +} + +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/query_stats_on_parameter_change.cpp b/src/mongo/db/query/query_stats/query_stats_on_parameter_change.cpp new file mode 100644 index 00000000000..a8b7df9fccb --- /dev/null +++ b/src/mongo/db/query/query_stats/query_stats_on_parameter_change.cpp @@ -0,0 +1,97 @@ +/** + * Copyright (C) 2022-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. + */ + + +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery + +#include "mongo/db/query/query_stats/query_stats_on_parameter_change.h" + +#include "mongo/base/status.h" +#include "mongo/db/concurrency/d_concurrency.h" +#include "mongo/db/query/partitioned_cache.h" +#include "mongo/db/query/query_knobs_gen.h" +#include "mongo/db/query/util/memory_util.h" +#include "mongo/db/service_context.h" +#include "mongo/logv2/log.h" + +namespace mongo::query_stats_util { + +namespace { +/** + * Given the current 'Client', returns a pointer to the 'ServiceContext' and an interface for + * updating the queryStats store. + */ +std::pair<ServiceContext*, OnParamChangeUpdater*> getUpdater(const Client& client) { + auto serviceCtx = client.getServiceContext(); + tassert(7106500, "ServiceContext must be non null", serviceCtx); + + auto updater = queryStatsStoreOnParamChangeUpdater(serviceCtx).get(); + tassert(7106501, "queryStats store size updater must be non null", updater); + return {serviceCtx, updater}; +} +} // namespace + + +Status onQueryStatsStoreSizeUpdate(const std::string& str) { + auto newSize = memory_util::MemorySize::parse(str); + if (!newSize.isOK()) { + return newSize.getStatus(); + } + + // The client is nullptr if the parameter is supplied from the command line. In this case, we + // ignore the update event, the parameter will be processed when initializing the service + // context. + if (auto client = Client::getCurrent()) { + auto&& [serviceCtx, updater] = getUpdater(*client); + updater->updateCacheSize(serviceCtx, newSize.getValue()); + } + + return Status::OK(); +} + +Status validateQueryStatsStoreSize(const std::string& str) { + return memory_util::MemorySize::parse(str).getStatus(); +} + +Status onQueryStatsSamplingRateUpdate(int samplingRate) { + // The client is nullptr if the parameter is supplied from the command line. In this case, we + // ignore the update event, the parameter will be processed when initializing the service + // context. + if (auto client = Client::getCurrent()) { + auto&& [serviceCtx, updater] = getUpdater(*client); + updater->updateSamplingRate(serviceCtx, samplingRate < 0 ? INT_MAX : samplingRate); + } + + return Status::OK(); +} + +const Decorable<ServiceContext>::Decoration<std::unique_ptr<OnParamChangeUpdater>> + queryStatsStoreOnParamChangeUpdater = + ServiceContext::declareDecoration<std::unique_ptr<OnParamChangeUpdater>>(); +} // namespace mongo::query_stats_util diff --git a/src/mongo/db/query/query_stats/query_stats_on_parameter_change.h b/src/mongo/db/query/query_stats/query_stats_on_parameter_change.h new file mode 100644 index 00000000000..2a824961b34 --- /dev/null +++ b/src/mongo/db/query/query_stats/query_stats_on_parameter_change.h @@ -0,0 +1,76 @@ +/** + * Copyright (C) 2022-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 "mongo/base/status.h" +#include "mongo/db/concurrency/d_concurrency.h" +#include "mongo/db/query/partitioned_cache.h" +#include "mongo/db/query/util/memory_util.h" + + +namespace mongo::query_stats_util { + +Status onQueryStatsStoreSizeUpdate(const std::string& str); + + +Status validateQueryStatsStoreSize(const std::string& str); + +Status onQueryStatsSamplingRateUpdate(int samplingRate); + +/** + * An interface used to modify the queryStats store when query setParameters are modified. This is + * done via an interface decorating the 'ServiceContext' in order to avoid a link-time dependency of + * the query knobs library on the queryStats code. + */ +class OnParamChangeUpdater { +public: + virtual ~OnParamChangeUpdater() = default; + + /** + * Resizes the queryStats store decorating 'serviceCtx' to the new size given by 'memSize'. If + * the new size is smaller than the old, cache entries are evicted in order to ensure the + * cache fits within the new size bound. + */ + virtual void updateCacheSize(ServiceContext* serviceCtx, memory_util::MemorySize memSize) = 0; + + /** + * Updates the sampling rate for the queryStats rate limiter. + */ + virtual void updateSamplingRate(ServiceContext* serviceCtx, int samplingRate) = 0; +}; + +/** + * Decorated accessor to the 'OnParamChangeUpdater' stored in 'ServiceContext'. Again, this is done + * via a decoration and interface to avoid a link-time dependency from the query knobs library on + * the queryStats code. + */ +extern const Decorable<ServiceContext>::Decoration<std::unique_ptr<OnParamChangeUpdater>> + queryStatsStoreOnParamChangeUpdater; +} // namespace mongo::query_stats_util diff --git a/src/mongo/db/query/query_stats/query_stats_store_test.cpp b/src/mongo/db/query/query_stats/query_stats_store_test.cpp new file mode 100644 index 00000000000..74965c4ece9 --- /dev/null +++ b/src/mongo/db/query/query_stats/query_stats_store_test.cpp @@ -0,0 +1,1427 @@ +/** + * Copyright (C) 2022-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. + */ + +#include "mongo/bson/simple_bsonobj_comparator.h" +#include "mongo/db/catalog/rename_collection.h" +#include "mongo/db/collection_type.h" +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/query_feature_flags_gen.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_shape/serialization_options.h" +#include "mongo/db/query/query_stats/agg_key.h" +#include "mongo/db/query/query_stats/find_key.h" +#include "mongo/db/query/query_stats/key.h" +#include "mongo/db/query/query_stats/query_stats.h" +#include "mongo/db/service_context_test_fixture.h" +#include "mongo/idl/server_parameter_test_util.h" +#include "mongo/unittest/unittest.h" + +namespace mongo::query_stats { + +int countAllEntries(const QueryStatsStore& store) { + int numKeys = 0; + store.forEach([&](auto&& key, auto&& entry) { numKeys++; }); + return numKeys; +} + +static const NamespaceStringOrUUID kDefaultTestNss = NamespaceString("testDB.testColl"); +class QueryStatsStoreTest : public ServiceContextTest { +public: + static std::unique_ptr<const Key> makeFindKeyFromQuery(BSONObj filter) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + auto fcr = std::make_unique<FindCommandRequest>(kDefaultTestNss); + fcr->setFilter(filter.getOwned()); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(expCtx, std::move(fcr))); + return std::make_unique<FindKey>(expCtx, *parsedFind, collectionType); + } + + static constexpr auto collectionType = query_shape::CollectionType::kCollection; + BSONObj makeQueryStatsKeyFindRequest(const FindCommandRequest& fcr, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + bool applyHmac) { + auto fcrCopy = std::make_unique<FindCommandRequest>(fcr); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(expCtx, std::move(fcrCopy))); + FindKey findKey(expCtx, *parsedFind, collectionType); + SerializationOptions opts = SerializationOptions::kDebugShapeAndMarkIdentifiers_FOR_TEST; + if (!applyHmac) { + opts.transformIdentifiers = false; + opts.transformIdentifiersCallback = defaultHmacStrategy; + } + return findKey.toBson(expCtx->opCtx, opts); + } + + BSONObj makeQueryStatsKeyAggregateRequest(AggregateCommandRequest acr, + const Pipeline& pipeline, + const boost::intrusive_ptr<ExpressionContext>& expCtx, + LiteralSerializationPolicy literalPolicy, + bool applyHmac = false) { + auto aggKey = std::make_unique<AggKey>(acr, + pipeline, + expCtx, + pipeline.getInvolvedCollections(), + acr.getNamespace(), + collectionType); + + // SerializationOptions opts{.literalPolicy = literalPolicy}; + SerializationOptions opts = SerializationOptions::kMarkIdentifiers_FOR_TEST; + opts.literalPolicy = literalPolicy; + if (!applyHmac) { + opts.transformIdentifiers = false; + opts.transformIdentifiersCallback = defaultHmacStrategy; + } + return aggKey->toBson(expCtx->opCtx, opts); + } +}; + +TEST_F(QueryStatsStoreTest, BasicUsage) { + QueryStatsStore queryStatsStore{5000000, 1000}; + + auto getMetrics = [&](BSONObj query) { + auto key = makeFindKeyFromQuery(query); + auto lookupResult = queryStatsStore.lookup(absl::Hash<query_stats::Key>{}(*key)); + ASSERT_OK(lookupResult); + return *lookupResult.getValue(); + }; + + auto collectMetrics = [&](BSONObj query) { + auto key = makeFindKeyFromQuery(query); + auto lookupHash = absl::Hash<query_stats::Key>{}(*key); + auto lookupResult = queryStatsStore.lookup(lookupHash); + if (!lookupResult.isOK()) { + queryStatsStore.put(lookupHash, QueryStatsEntry{std::move(key)}); + lookupResult = queryStatsStore.lookup(lookupHash); + } + auto metrics = lookupResult.getValue(); + metrics->execCount += 1; + metrics->lastExecutionMicros += 123456; + }; + + auto query1 = BSON("query" << 1 << "xEquals" << 42); + // same value, different instance (tests hashing & equality) + auto query1x = BSON("query" << 1 << "xEquals" << 42); + auto query2 = BSON("query" << 2 << "yEquals" << 43); + + collectMetrics(query1); + collectMetrics(query1); + collectMetrics(query1x); + collectMetrics(query2); + + ASSERT_EQ(getMetrics(query1).execCount, 3); + ASSERT_EQ(getMetrics(query1x).execCount, 3); + ASSERT_EQ(getMetrics(query2).execCount, 1); + + auto collectMetricsWithLock = [&](BSONObj& filter) { + auto key = makeFindKeyFromQuery(filter); + auto [lookupResult, lock] = + queryStatsStore.getWithPartitionLock(absl::Hash<query_stats::Key>{}(*key)); + ASSERT_OK(lookupResult); + auto& metrics = *lookupResult.getValue(); + metrics.execCount += 1; + metrics.lastExecutionMicros += 123456; + }; + + collectMetricsWithLock(query1x); + collectMetricsWithLock(query2); + + ASSERT_EQ(getMetrics(query1).execCount, 4); + ASSERT_EQ(getMetrics(query1x).execCount, 4); + ASSERT_EQ(getMetrics(query2).execCount, 2); + + ASSERT_EQ(2, countAllEntries(queryStatsStore)); +} + +TEST_F(QueryStatsStoreTest, EvictionTest) { + // This creates a queryStats store with a single partition to specifically test the eviction + // behavior with very large queries. + // Add an entry that is smaller than the max partition size. + auto query = BSON("query" << 1 << "xEquals" << 42); + auto key = makeFindKeyFromQuery(query); + + const size_t cacheSize = key->size() + sizeof(QueryStatsEntry) + 100; + const auto numPartitions = 1; + QueryStatsStore queryStatsStore{cacheSize, numPartitions}; + + auto hash = absl::Hash<query_stats::Key>{}(*key); + queryStatsStore.put(hash, QueryStatsEntry{std::move(key)}); + ASSERT_EQ(countAllEntries(queryStatsStore), 1); + + // We'll do this again later so save this as a helper function. + auto addLargeEntry = [&](auto& queryStatsStore) { + // Add an entry that is larger than the max partition size to the non-empty partition. This + // should evict both entries, the first small entry written to the partition and the current + // too large entry we wish to write to the partition. The reason is because entries are + // evicted from the partition in order of least recently used. Thus, the small entry will be + // evicted first but the partition will still be over budget so the final, too large entry + // will also be evicted. + auto opCtx = makeOperationContext(); + auto fcr = std::make_unique<FindCommandRequest>( + NamespaceStringOrUUID(NamespaceString("testDB.testColl"))); + fcr->setLet(BSON("var" << 2)); + fcr->setFilter(fromjson("{$expr: [{$eq: ['$a', '$$var']}]}")); + fcr->setProjection(fromjson("{varIs: '$$var'}")); + fcr->setLimit(5); + fcr->setSkip(2); + fcr->setBatchSize(25); + fcr->setMaxTimeMS(1000); + fcr->setNoCursorTimeout(false); + opCtx->setComment(BSON("comment" + << " foo bar baz")); + fcr->setSingleBatch(false); + fcr->setAllowDiskUse(false); + fcr->setAllowPartialResults(true); + fcr->setAllowDiskUse(false); + fcr->setShowRecordId(true); + fcr->setHint(BSON("z" << 1 << "c" << 1)); + fcr->setMax(BSON("z" << 25)); + fcr->setMin(BSON("z" << 80)); + fcr->setSort(BSON("sortVal" << 1 << "otherSort" << -1)); + auto&& [expCtx, parsedFind] = + uassertStatusOK(parsed_find_command::parse(opCtx.get(), std::move(fcr))); + + key = std::make_unique<query_stats::FindKey>(expCtx, *parsedFind, collectionType); + auto lookupHash = absl::Hash<query_stats::Key>{}(*key); + QueryStatsEntry testMetrics{std::move(key)}; + queryStatsStore.put(lookupHash, testMetrics); + }; + + addLargeEntry(queryStatsStore); + ASSERT_EQ(countAllEntries(queryStatsStore), 0); + + // This creates a queryStats store where each partition has a max size of 500 bytes. + QueryStatsStore queryStatsStoreTwo{/*cacheSize*/ cacheSize * 3, /*numPartitions*/ 3}; + // Adding a queryStats store entry that is smaller than the overal cache size but larger + // than a single partition max size, will cause an eviction. testMetrics is larger than 500 + // bytes and thus over budget for the partitions of this cache. + addLargeEntry(queryStatsStoreTwo); + ASSERT_EQ(countAllEntries(queryStatsStoreTwo), 0); +} + +TEST_F(QueryStatsStoreTest, GenerateMaxBsonSizeQueryShape) { + const NamespaceString nss = NamespaceString("testDB.testColl"); + FindCommandRequest fcr((NamespaceStringOrUUID(nss))); + // This creates a query that is just below the 16 MB memory limit. + int limit = 225500; + BSONObjBuilder bob; + BSONArrayBuilder andBob(bob.subarrayStart("$and")); + for (int i = 1; i <= limit; i++) { + BSONObjBuilder childrenBob; + childrenBob.append("x", BSON("$lt" << i << "$gte" << i)); + andBob.append(childrenBob.obj()); + } + andBob.doneFast(); + fcr.setFilter(bob.obj()); + auto fcrCopy = std::make_unique<FindCommandRequest>(fcr); + auto opCtx = makeOperationContext(); + auto parsedFindPair = + uassertStatusOK(parsed_find_command::parse(opCtx.get(), std::move(fcrCopy))); + + auto&& globalQueryStatsStoreManager = QueryStatsStoreManager::get(opCtx->getServiceContext()); + globalQueryStatsStoreManager = std::make_unique<QueryStatsStoreManager>(500000, 1000); + + // The shapification process will bloat the input query over the 16 MB memory limit. Assert that + // calling registerRequest() doesn't throw and that the opDebug isn't registered with a key hash + // (thus metrics won't be tracked for this query). + ASSERT_DOES_NOT_THROW(query_stats::registerRequest(opCtx.get(), nss, [&]() { + return std::make_unique<query_stats::FindKey>( + parsedFindPair.first, *parsedFindPair.second, query_shape::CollectionType::kCollection); + })); + auto& opDebug = CurOp::get(*opCtx)->debug(); + ASSERT_FALSE(opDebug.queryStatsInfo.keyHash.has_value()); +} + +TEST_F(QueryStatsStoreTest, CorrectlyRedactsFindCommandRequestAllFields) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + FindCommandRequest fcr(kDefaultTestNss); + + fcr.setFilter(BSON("a" << 1)); + + auto key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "find", + "filter": { + "HASH<a>": { + "$eq": "?number" + } + } + }, + "collectionType": "collection" + })", + key); + + // Add sort. + fcr.setSort(BSON("sortVal" << 1 << "otherSort" << -1)); + key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "find", + "filter": { + "HASH<a>": { + "$eq": "?number" + } + }, + "sort": { + "HASH<sortVal>": 1, + "HASH<otherSort>": -1 + } + }, + "collectionType": "collection" + })", + key); + + // Add inclusion projection. + fcr.setProjection(BSON("e" << true << "f" << true)); + key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "find", + "filter": { + "HASH<a>": { + "$eq": "?number" + } + }, + "projection": { + "HASH<e>": true, + "HASH<f>": true, + "HASH<_id>": true + }, + "sort": { + "HASH<sortVal>": 1, + "HASH<otherSort>": -1 + } + }, + "collectionType": "collection" + })", + key); + + // Add let. + fcr.setLet(BSON("var1" << 1 << "var2" + << "const1")); + key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "let": { + "HASH<var1>": "?number", + "HASH<var2>": "?string" + }, + "command": "find", + "filter": { + "HASH<a>": { + "$eq": "?number" + } + }, + "projection": { + "HASH<e>": true, + "HASH<f>": true, + "HASH<_id>": true + }, + "sort": { + "HASH<sortVal>": 1, + "HASH<otherSort>": -1 + } + }, + "collectionType": "collection" + })", + key); + + // Add hinting fields. + fcr.setHint(BSON("z" << 1 << "c" << 1)); + fcr.setMax(BSON("z" << 25)); + fcr.setMin(BSON("z" << 80)); + key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "let": { + "HASH<var1>": "?number", + "HASH<var2>": "?string" + }, + "command": "find", + "filter": { + "HASH<a>": { + "$eq": "?number" + } + }, + "projection": { + "HASH<e>": true, + "HASH<f>": true, + "HASH<_id>": true + }, + "max": { + "HASH<z>": "?number" + }, + "min": { + "HASH<z>": "?number" + }, + "sort": { + "HASH<sortVal>": 1, + "HASH<otherSort>": -1 + } + }, + "collectionType": "collection", + "hint": { + "HASH<z>": 1, + "HASH<c>": 1 + } + })", + key); + + // Add the literal redaction fields. + fcr.setLimit(5); + fcr.setSkip(2); + fcr.setBatchSize(25); + fcr.setMaxTimeMS(1000); + fcr.setNoCursorTimeout(false); + + key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "let": { + "HASH<var1>": "?number", + "HASH<var2>": "?string" + }, + "command": "find", + "filter": { + "HASH<a>": { + "$eq": "?number" + } + }, + "projection": { + "HASH<e>": true, + "HASH<f>": true, + "HASH<_id>": true + }, + "max": { + "HASH<z>": "?number" + }, + "min": { + "HASH<z>": "?number" + }, + "sort": { + "HASH<sortVal>": 1, + "HASH<otherSort>": -1 + }, + "limit": "?number", + "skip": "?number" + }, + "collectionType": "collection", + "hint": { + "HASH<z>": 1, + "HASH<c>": 1 + }, + "maxTimeMS": "?number", + "noCursorTimeout": false, + "batchSize": "?number" + })", + key); + + // Add the fields that shouldn't be hmacApplied. + fcr.setSingleBatch(true); + fcr.setAllowDiskUse(false); + fcr.setAllowPartialResults(true); + fcr.setAllowDiskUse(false); + fcr.setShowRecordId(true); + auto readPreference = BSON("mode" + << "nearest" + << "tags" + << BSON_ARRAY(BSON("some" + << "tag") + << BSON("some" + << "other tag"))); + ReadPreferenceSetting::get(expCtx->opCtx) = + uassertStatusOK(ReadPreferenceSetting::fromInnerBSON(readPreference)); + key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "let": { + "HASH<var1>": "?number", + "HASH<var2>": "?string" + }, + "command": "find", + "filter": { + "HASH<a>": { + "$eq": "?number" + } + }, + "projection": { + "HASH<e>": true, + "HASH<f>": true, + "HASH<_id>": true + }, + "max": { + "HASH<z>": "?number" + }, + "min": { + "HASH<z>": "?number" + }, + "sort": { + "HASH<sortVal>": 1, + "HASH<otherSort>": -1 + }, + "limit": "?number", + "skip": "?number", + "singleBatch": true, + "allowDiskUse": false, + "showRecordId": true + }, + "$readPreference": { + "mode": "nearest", + "tags": [ { "some": "other tag" }, { "some": "tag" } ], + "hedge": { "enabled": true } + }, + "collectionType": "collection", + "hint": { + "HASH<z>": 1, + "HASH<c>": 1 + }, + "maxTimeMS": "?number", + "allowPartialResults": true, + "noCursorTimeout": false, + "batchSize": "?number" + })", + key); + + fcr.setAllowPartialResults(false); + key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + // Make sure that a false allowPartialResults is also accurately captured. + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "let": { + "HASH<var1>": "?number", + "HASH<var2>": "?string" + }, + "command": "find", + "filter": { + "HASH<a>": { + "$eq": "?number" + } + }, + "projection": { + "HASH<e>": true, + "HASH<f>": true, + "HASH<_id>": true + }, + "max": { + "HASH<z>": "?number" + }, + "min": { + "HASH<z>": "?number" + }, + "sort": { + "HASH<sortVal>": 1, + "HASH<otherSort>": -1 + }, + "limit": "?number", + "skip": "?number", + "singleBatch": true, + "allowDiskUse": false, + "showRecordId": true + }, + "$readPreference": { + "mode": "nearest", + "tags": [ { "some": "other tag" }, { "some": "tag" } ], + "hedge": { "enabled": true } + }, + "collectionType": "collection", + "hint": { + "HASH<z>": 1, + "HASH<c>": 1 + }, + "maxTimeMS": "?number", + "allowPartialResults": false, + "noCursorTimeout": false, + "batchSize": "?number" + })", + key); +} + +TEST_F(QueryStatsStoreTest, CorrectlyRedactsTailableFindCommandRequest) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + + FindCommandRequest fcr(NamespaceStringOrUUID(NamespaceString("testDB.testColl"))); + fcr.setAwaitData(true); + fcr.setTailable(true); + fcr.setSort(BSON("$natural" << 1)); + auto key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "find", + "filter": {}, + "tailable": true, + "awaitData": true + }, + "collectionType": "collection", + "hint": { + "$natural": 1 + } + })", + key); +} + +TEST_F(QueryStatsStoreTest, CorrectlyRedactsFindCommandRequestEmptyFields) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + FindCommandRequest fcr(NamespaceStringOrUUID(NamespaceString("testDB.testColl"))); + fcr.setFilter(BSONObj()); + fcr.setSort(BSONObj()); + fcr.setProjection(BSONObj()); + + auto hmacApplied = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "find", + "filter": {} + }, + "collectionType": "collection" + })", + hmacApplied); // NOLINT (test auto-update) +} + +TEST_F(QueryStatsStoreTest, CorrectlyRedactsHintsWithOptions) { + auto expCtx = make_intrusive<ExpressionContextForTest>(); + FindCommandRequest fcr(NamespaceStringOrUUID(NamespaceString("testDB.testColl"))); + + fcr.setFilter(BSON("b" << 1)); + fcr.setHint(BSON("z" << 1 << "c" << 1)); + fcr.setMax(BSON("z" << 25)); + fcr.setMin(BSON("z" << 80)); + + auto key = makeQueryStatsKeyFindRequest(fcr, expCtx, false); + + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "testDB", + "coll": "testColl" + }, + "command": "find", + "filter": { + "b": { + "$eq": "?number" + } + }, + "max": { + "z": "?number" + }, + "min": { + "z": "?number" + } + }, + "collectionType": "collection", + "hint": { + "z": 1, + "c": 1 + } + })", + key); + // Test with a string hint. Note that this is the internal representation of the string hint + // generated at parse time. + fcr.setHint(BSON("$hint" + << "z")); + + key = makeQueryStatsKeyFindRequest(fcr, expCtx, false); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "testDB", + "coll": "testColl" + }, + "command": "find", + "filter": { + "b": { + "$eq": "?number" + } + }, + "max": { + "z": "?number" + }, + "min": { + "z": "?number" + } + }, + "collectionType": "collection", + "hint": { + "$hint": "z" + } + })", + key); + + fcr.setHint(BSON("z" << 1 << "c" << 1)); + key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "find", + "filter": { + "HASH<b>": { + "$eq": "?number" + } + }, + "max": { + "HASH<z>": "?number" + }, + "min": { + "HASH<z>": "?number" + } + }, + "collectionType": "collection", + "hint": { + "HASH<z>": 1, + "HASH<c>": 1 + } + })", + key); + + // Test that $natural comes through unmodified. + fcr.setHint(BSON("$natural" << -1)); + key = makeQueryStatsKeyFindRequest(fcr, expCtx, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "find", + "filter": { + "HASH<b>": { + "$eq": "?number" + } + }, + "max": { + "HASH<z>": "?number" + }, + "min": { + "HASH<z>": "?number" + } + }, + "collectionType": "collection", + "hint": { + "$natural": -1 + } + })", + key); +} + +TEST_F(QueryStatsStoreTest, DefinesLetVariables) { + // Test that the expression context we use to apply hmac will understand the 'let' part of + // the find command while parsing the other pieces of the command. + + // Note that this ExpressionContext will not have the let variables defined - we expect the + // 'makeQueryStatsKey' call to do that. + auto opCtx = makeOperationContext(); + auto fcr = std::make_unique<FindCommandRequest>(NamespaceString("testDB.testColl")); + fcr->setLet(BSON("var" << 2)); + fcr->setFilter(fromjson("{$expr: [{$eq: ['$a', '$$var']}]}")); + fcr->setProjection(fromjson("{varIs: '$$var'}")); + + auto expCtx = make_intrusive<ExpressionContextForTest>(opCtx.get()); + expCtx->variables.seedVariablesWithLetParameters(expCtx.get(), *fcr->getLet()); + auto hmacApplied = makeQueryStatsKeyFindRequest(*fcr, expCtx, false); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "testDB", + "coll": "testColl" + }, + "let": { + "var": "?number" + }, + "command": "find", + "filter": { + "$expr": [ + { + "$eq": [ + "$a", + "$$var" + ] + } + ] + }, + "projection": { + "varIs": "$$var", + "_id": true + } + }, + "collectionType": "collection" + })", + hmacApplied); + + hmacApplied = makeQueryStatsKeyFindRequest(*fcr, expCtx, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "let": { + "HASH<var>": "?number" + }, + "command": "find", + "filter": { + "$expr": [ + { + "$eq": [ + "$HASH<a>", + "$$HASH<var>" + ] + } + ] + }, + "projection": { + "HASH<varIs>": "$$HASH<var>", + "HASH<_id>": true + } + }, + "collectionType": "collection" + })", + hmacApplied); +} + +TEST_F(QueryStatsStoreTest, CorrectlyTokenizesAggregateCommandRequestAllFieldsSimplePipeline) { + auto expCtx = make_intrusive<ExpressionContextForTest>(*kDefaultTestNss.nss()); + AggregateCommandRequest acr(*kDefaultTestNss.nss()); + auto matchStage = fromjson(R"({ + $match: { + foo: { $in: ["a", "b"] }, + bar: { $gte: { $date: "2022-01-01T00:00:00Z" } } + } + })"); + auto unwindStage = fromjson("{$unwind: '$x'}"); + auto groupStage = fromjson(R"({ + $group: { + _id: "$_id", + c: { $first: "$d.e" }, + f: { $sum: 1 } + } + })"); + auto limitStage = fromjson("{$limit: 10}"); + auto outStage = fromjson(R"({$out: 'outColl'})"); + auto rawPipeline = {matchStage, unwindStage, groupStage, limitStage, outStage}; + acr.setPipeline(rawPipeline); + auto pipeline = Pipeline::parse(rawPipeline, expCtx); + + auto shapified = makeQueryStatsKeyAggregateRequest( + acr, *pipeline, expCtx, LiteralSerializationPolicy::kToDebugTypeString, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "aggregate", + "pipeline": [ + { + "$match": { + "$and": [ + { + "HASH<foo>": { + "$in": "?array<?string>" + } + }, + { + "HASH<bar>": { + "$gte": "?date" + } + } + ] + } + }, + { + "$unwind": { + "path": "$HASH<x>" + } + }, + { + "$group": { + "_id": "$HASH<_id>", + "HASH<c>": { + "$first": "$HASH<d>.HASH<e>" + }, + "HASH<f>": { + "$sum": "?number" + } + } + }, + { + "$limit": "?number" + }, + { + "$out": { + "coll": "HASH<outColl>", + "db": "HASH<testDB>" + } + } + ] + }, + "collectionType": "collection" + })", + shapified); + + // Add the fields that shouldn't be abstracted. + acr.setAllowDiskUse(false); + acr.setHint(BSON("z" << 1 << "c" << 1)); + acr.setCollation(BSON("locale" + << "simple")); + shapified = makeQueryStatsKeyAggregateRequest( + acr, *pipeline, expCtx, LiteralSerializationPolicy::kToDebugTypeString, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "collation": { + "locale": "simple" + }, + "command": "aggregate", + "pipeline": [ + { + "$match": { + "$and": [ + { + "HASH<foo>": { + "$in": "?array<?string>" + } + }, + { + "HASH<bar>": { + "$gte": "?date" + } + } + ] + } + }, + { + "$unwind": { + "path": "$HASH<x>" + } + }, + { + "$group": { + "_id": "$HASH<_id>", + "HASH<c>": { + "$first": "$HASH<d>.HASH<e>" + }, + "HASH<f>": { + "$sum": "?number" + } + } + }, + { + "$limit": "?number" + }, + { + "$out": { + "coll": "HASH<outColl>", + "db": "HASH<testDB>" + } + } + ], + "allowDiskUse": false + }, + "collectionType": "collection", + "hint": { + "HASH<z>": 1, + "HASH<c>": 1 + } + })", + shapified); + + // Add let. + acr.setLet(BSON("var1" << BSON("$literal" + << "$foo") + << "var2" + << "bar")); + shapified = makeQueryStatsKeyAggregateRequest( + acr, *pipeline, expCtx, LiteralSerializationPolicy::kToDebugTypeString, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "collation": { + "locale": "simple" + }, + "let": { + "HASH<var1>": "?string", + "HASH<var2>": "?string" + }, + "command": "aggregate", + "pipeline": [ + { + "$match": { + "$and": [ + { + "HASH<foo>": { + "$in": "?array<?string>" + } + }, + { + "HASH<bar>": { + "$gte": "?date" + } + } + ] + } + }, + { + "$unwind": { + "path": "$HASH<x>" + } + }, + { + "$group": { + "_id": "$HASH<_id>", + "HASH<c>": { + "$first": "$HASH<d>.HASH<e>" + }, + "HASH<f>": { + "$sum": "?number" + } + } + }, + { + "$limit": "?number" + }, + { + "$out": { + "coll": "HASH<outColl>", + "db": "HASH<testDB>" + } + } + ], + "allowDiskUse": false + }, + "collectionType": "collection", + "hint": { + "HASH<z>": 1, + "HASH<c>": 1 + } + })", + shapified); + + // Add the fields that should be abstracted. + auto cursorOptions = SimpleCursorOptions(); + cursorOptions.setBatchSize(10); + acr.setCursor(cursorOptions); + acr.setMaxTimeMS(500); + acr.setBypassDocumentValidation(true); + expCtx->opCtx->setComment(BSON("comment" + << "note to self")); + shapified = makeQueryStatsKeyAggregateRequest( + acr, *pipeline, expCtx, LiteralSerializationPolicy::kToDebugTypeString, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "collation": { + "locale": "simple" + }, + "let": { + "HASH<var1>": "?string", + "HASH<var2>": "?string" + }, + "command": "aggregate", + "pipeline": [ + { + "$match": { + "$and": [ + { + "HASH<foo>": { + "$in": "?array<?string>" + } + }, + { + "HASH<bar>": { + "$gte": "?date" + } + } + ] + } + }, + { + "$unwind": { + "path": "$HASH<x>" + } + }, + { + "$group": { + "_id": "$HASH<_id>", + "HASH<c>": { + "$first": "$HASH<d>.HASH<e>" + }, + "HASH<f>": { + "$sum": "?number" + } + } + }, + { + "$limit": "?number" + }, + { + "$out": { + "coll": "HASH<outColl>", + "db": "HASH<testDB>" + } + } + ], + "allowDiskUse": false + }, + "comment": "?string", + "collectionType": "collection", + "hint": { + "HASH<z>": 1, + "HASH<c>": 1 + }, + "maxTimeMS": "?number", + "bypassDocumentValidation": true, + "cursor": { + "batchSize": "?number" + } + })", + shapified); + + // Test again but with the representative query shape. + shapified = makeQueryStatsKeyAggregateRequest( + acr, *pipeline, expCtx, LiteralSerializationPolicy::kToRepresentativeParseableValue, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "collation": { + "locale": "simple" + }, + "let": { + "HASH<var1>": { + "$const": "?" + }, + "HASH<var2>": { + "$const": "?" + } + }, + "command": "aggregate", + "pipeline": [ + { + "$match": { + "$and": [ + { + "HASH<foo>": { + "$in": [ + "?" + ] + } + }, + { + "HASH<bar>": { + "$gte": {"$date":"1970-01-01T00:00:00.000Z"} + } + } + ] + } + }, + { + "$unwind": { + "path": "$HASH<x>" + } + }, + { + "$group": { + "_id": "$HASH<_id>", + "HASH<c>": { + "$first": "$HASH<d>.HASH<e>" + }, + "HASH<f>": { + "$sum": { + "$const": 1 + } + } + } + }, + { + "$limit": 1 + }, + { + "$out": { + "coll": "HASH<outColl>", + "db": "HASH<testDB>" + } + } + ], + "allowDiskUse": false + }, + "comment": "?", + "collectionType": "collection", + "hint": { + "HASH<z>": 1, + "HASH<c>": 1 + }, + "maxTimeMS": 1, + "bypassDocumentValidation": true, + "cursor": { + "batchSize": 1 + } + })", + shapified); +} + +TEST_F(QueryStatsStoreTest, CorrectlyTokenizesAggregateCommandRequestEmptyFields) { + auto expCtx = make_intrusive<ExpressionContextForTest>(*kDefaultTestNss.nss()); + AggregateCommandRequest acr(*kDefaultTestNss.nss()); + acr.setPipeline({}); + auto pipeline = Pipeline::parse({}, expCtx); + + auto shapified = makeQueryStatsKeyAggregateRequest( + acr, *pipeline, expCtx, LiteralSerializationPolicy::kToDebugTypeString, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "aggregate", + "pipeline": [] + }, + "collectionType": "collection" + })", + shapified); // NOLINT (test auto-update) + + // Test again with the representative query shape. + shapified = makeQueryStatsKeyAggregateRequest( + acr, *pipeline, expCtx, LiteralSerializationPolicy::kToRepresentativeParseableValue, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "aggregate", + "pipeline": [] + }, + "collectionType": "collection" + })", + shapified); // NOLINT (test auto-update) +} + +TEST_F(QueryStatsStoreTest, + CorrectlyTokenizesAggregateCommandRequestPipelineWithSecondaryNamespaces) { + auto expCtx = make_intrusive<ExpressionContextForTest>(*kDefaultTestNss.nss()); + auto nsToUnionWith = NamespaceString(expCtx->ns.db(), "otherColl"); + expCtx->addResolvedNamespaces({nsToUnionWith}); + + AggregateCommandRequest acr(*kDefaultTestNss.nss()); + auto unionWithStage = fromjson(R"({ + $unionWith: { + coll: "otherColl", + pipeline: [{$match: {val: "foo"}}] + } + })"); + auto sortStage = fromjson("{$sort: {age: 1}}"); + auto rawPipeline = {unionWithStage, sortStage}; + acr.setPipeline(rawPipeline); + auto pipeline = Pipeline::parse(rawPipeline, expCtx); + + auto shapified = makeQueryStatsKeyAggregateRequest( + acr, *pipeline, expCtx, LiteralSerializationPolicy::kToDebugTypeString, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "aggregate", + "pipeline": [ + { + "$unionWith": { + "coll": "HASH<otherColl>", + "pipeline": [ + { + "$match": { + "HASH<val>": { + "$eq": "?string" + } + } + } + ] + } + }, + { + "$sort": { + "HASH<age>": 1 + } + } + ] + }, + "collectionType": "collection", + "otherNss": [ + { + "db": "HASH<testDB>", + "coll": "HASH<otherColl>" + } + ] + })", + shapified); + + // Do the same thing with the representative query shape. + shapified = makeQueryStatsKeyAggregateRequest( + acr, *pipeline, expCtx, LiteralSerializationPolicy::kToRepresentativeParseableValue, true); + ASSERT_BSONOBJ_EQ_AUTO( // NOLINT + R"({ + "queryShape": { + "cmdNs": { + "db": "HASH<testDB>", + "coll": "HASH<testColl>" + }, + "command": "aggregate", + "pipeline": [ + { + "$unionWith": { + "coll": "HASH<otherColl>", + "pipeline": [ + { + "$match": { + "HASH<val>": { + "$eq": "?" + } + } + } + ] + } + }, + { + "$sort": { + "HASH<age>": 1 + } + } + ] + }, + "collectionType": "collection", + "otherNss": [ + { + "db": "HASH<testDB>", + "coll": "HASH<otherColl>" + } + ] + })", + shapified); +} + +BSONObj toBSON(AggregatedMetric am) { + BSONObjBuilder builder; + am.appendTo(builder, "m"); + return builder.obj(); +} + +TEST_F(QueryStatsStoreTest, SumOfSquaresOverflowTest) { + // Ensure sumOfSquares is initialized correctly. + AggregatedMetric aggMetric; + auto res = toBSON(aggMetric).getObjectField("m").getField("sumOfSquares").Decimal(); + + ASSERT_EQ(res, Decimal128()); + + // Aggregating with the maximum int value does not overflow the sumOfSquares field. + auto maxVal = std::numeric_limits<uint64_t>::max(); + aggMetric.aggregate(maxVal); + res = toBSON(aggMetric).getObjectField("m").getField("sumOfSquares").Decimal(); + + ASSERT_EQ(res, Decimal128(maxVal).power(Decimal128(2.0))); +} +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/query_stats_test.cpp b/src/mongo/db/query/query_stats/query_stats_test.cpp new file mode 100644 index 00000000000..4b9462e3e25 --- /dev/null +++ b/src/mongo/db/query/query_stats/query_stats_test.cpp @@ -0,0 +1,223 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/bson/bsonobj.h" +#include "mongo/db/collection_type.h" +#include "mongo/db/namespace_string.h" +#include "mongo/db/operation_context.h" +#include "mongo/db/pipeline/expression_context.h" +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/parsed_find_command.h" +#include "mongo/db/query/query_stats/find_key.h" +#include "mongo/db/query/query_stats/query_stats.h" +#include "mongo/db/service_context_test_fixture.h" +#include "mongo/idl/server_parameter_test_util.h" +#include "mongo/unittest/unittest.h" +#include "mongo/util/assert_util.h" + +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQueryStats + +namespace mongo::query_stats { +class QueryStatsTest : public ServiceContextTest {}; + +TEST_F(QueryStatsTest, TwoRegisterRequestsWithSameOpCtxRateLimitedFirstCall) { + // This test simulates what happens with queries over views where two calls to registerRequest() + // can be made with the same opCtx. + + // Make query for query stats. + const NamespaceString nss = NamespaceString("testDB.testColl"); + FindCommandRequest fcr((NamespaceStringOrUUID(nss))); + fcr.setFilter(BSONObj()); + + auto fcrCopy = std::make_unique<FindCommandRequest>(fcr); + auto opCtx = makeOperationContext(); + auto expCtx = make_intrusive<ExpressionContextForTest>(); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(expCtx, {std::move(fcrCopy)})); + + auto& opDebug = CurOp::get(*opCtx)->debug(); + ASSERT_EQ(opDebug.queryStatsInfo.wasRateLimited, false); + + // First call to registerRequest() should be rate limited. + QueryStatsStoreManager::getRateLimiter(opCtx->getServiceContext()) = + std::make_unique<RateLimiting>(0, Seconds{1}); + ASSERT_DOES_NOT_THROW(query_stats::registerRequest(opCtx.get(), nss, [&]() { + return std::make_unique<query_stats::FindKey>( + expCtx, *parsedFind, query_shape::CollectionType::kCollection); + })); + + // Since the query was rate limited, no key should have been created. + ASSERT(opDebug.queryStatsInfo.key == nullptr); + ASSERT_EQ(opDebug.queryStatsInfo.wasRateLimited, true); + + // Second call should not be rate limited. + QueryStatsStoreManager::getRateLimiter(opCtx->getServiceContext()) + .get() + ->setSamplingRate(INT_MAX); + + ASSERT_DOES_NOT_THROW(query_stats::registerRequest(opCtx.get(), nss, [&]() { + return std::make_unique<query_stats::FindKey>( + expCtx, *parsedFind, query_shape::CollectionType::kCollection); + })); + + // queryStatsKey should not be created for previously rate limited query. + ASSERT(opDebug.queryStatsInfo.key == nullptr); + ASSERT_EQ(opDebug.queryStatsInfo.wasRateLimited, true); + ASSERT_FALSE(opDebug.queryStatsInfo.keyHash.has_value()); +} + +TEST_F(QueryStatsTest, TwoRegisterRequestsWithSameOpCtxDisabledBetween) { + // This test simulates an observed bug where an opCtx is used for two requests, and between the + // first and the second the query stats store is emptied/disabled. + + // Make query for query stats. + const NamespaceString nss = NamespaceString("testDB.testColl"); + FindCommandRequest fcr((NamespaceStringOrUUID(nss))); + fcr.setFilter(BSONObj()); + + auto serviceCtx = getServiceContext(); + auto opCtx = makeOperationContext(); + + auto& opDebug = CurOp::get(*opCtx)->debug(); + ASSERT(opDebug.queryStatsInfo.key == nullptr); + ASSERT_FALSE(opDebug.queryStatsInfo.keyHash.has_value()); + QueryStatsStoreManager::get(serviceCtx) = + std::make_unique<QueryStatsStoreManager>(16 * 1024 * 1024, 1); + + QueryStatsStoreManager::getRateLimiter(serviceCtx) = + std::make_unique<RateLimiting>(-1, Seconds{1}); + + { + auto fcrCopy = std::make_unique<FindCommandRequest>(fcr); + auto expCtx = make_intrusive<ExpressionContext>( + opCtx.get(), *fcrCopy, nullptr, true /* mayDbProfile*/); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(expCtx, {std::move(fcrCopy)})); + ASSERT_DOES_NOT_THROW(query_stats::registerRequest(opCtx.get(), nss, [&]() { + return std::make_unique<query_stats::FindKey>( + expCtx, *parsedFind, query_shape::CollectionType::kCollection); + })); + + ASSERT(opDebug.queryStatsInfo.key != nullptr); + ASSERT(opDebug.queryStatsInfo.keyHash.has_value()); + + ASSERT_DOES_NOT_THROW(query_stats::writeQueryStats(opCtx.get(), + opDebug.queryStatsInfo.keyHash, + std::move(opDebug.queryStatsInfo.key), + 0 /*queryExecMicros*/, + 0 /*firstResponseExecMicros*/, + 0 /*docsReturned*/)); + } + + // Second call should see that query stats are now disabled. + { + // To reproduce SERVER-84730 we need to clear out the query stats store so that writing the + // stats at the end will attempt to insert a new entry. + QueryStatsStoreManager::get(serviceCtx)->resetSize(0); + + auto fcrCopy = std::make_unique<FindCommandRequest>(fcr); + fcrCopy->setFilter(BSON("x" << 1)); + auto expCtx = make_intrusive<ExpressionContext>( + opCtx.get(), *fcrCopy, nullptr, true /* mayDbProfile*/); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(expCtx, {std::move(fcrCopy)})); + + ASSERT_DOES_NOT_THROW(query_stats::registerRequest(opCtx.get(), nss, [&]() { + return std::make_unique<query_stats::FindKey>( + expCtx, *parsedFind, query_shape::CollectionType::kCollection); + })); + + // queryStatsKey should not be created since we have a size budget of 0. + ASSERT(opDebug.queryStatsInfo.key == nullptr); + // This is not a rate limit, but rather a lack of space rendering it entirely disabled. + ASSERT_FALSE(opDebug.queryStatsInfo.wasRateLimited); + + // Interestingly, we purposefully leave the hash value around on the OperationContext after + // the previous operation finishes. This is because we think it may have value in being + // logged in the future, even after query stats have been written. Excepting obscure + // internal use-cases, most OperationContexts will die shortly after the query stats are + // written, so this isn't expected to be a large issue. + ASSERT(opDebug.queryStatsInfo.keyHash.has_value()); + + QueryStatsStoreManager::get(serviceCtx)->resetSize(16 * 1024 * 1024); + // SERVER-84730 this assertion used to throw since there is no key, but there is a hash. + ASSERT_DOES_NOT_THROW(query_stats::writeQueryStats(opCtx.get(), + opDebug.queryStatsInfo.keyHash, + std::move(opDebug.queryStatsInfo.key), + 0 /*queryExecMicros*/, + 0 /*firstResponseExecMicros*/, + 0 /*docsReturned*/)); + } +} + +TEST_F(QueryStatsTest, RegisterRequestAbsorbsErrors) { + const NamespaceString nss = NamespaceString("testDB.testColl"); + + auto opCtx = makeOperationContext(); + auto& opDebug = CurOp::get(*opCtx)->debug(); + + QueryStatsStoreManager::getRateLimiter(getServiceContext()) = + std::make_unique<RateLimiting>(-1, Seconds{1}); + + // First case - don't treat errors as fatal. + internalQueryStatsErrorsAreCommandFatal.store(false); + + // Skip these checks for debug builds because errors are always fatal in that environment. + if (!kDebugBuild) { + ASSERT_DOES_NOT_THROW(query_stats::registerRequest(opCtx.get(), nss, [&]() { + uasserted(ErrorCodes::BSONObjectTooLarge, "size error"); + return nullptr; + })); + + opDebug.queryStatsInfo = OpDebug::QueryStatsInfo{}; + ASSERT_DOES_NOT_THROW(query_stats::registerRequest(opCtx.get(), nss, [&]() { + uasserted(ErrorCodes::BadValue, "fake error"); + return nullptr; + })); + } + + // Now make sure that errors are propagated when the knob is set. + internalQueryStatsErrorsAreCommandFatal.store(true); + + // We shouldn't propagate 'BSONObjectTooLarge' errors under any circumstances. + opDebug.queryStatsInfo = OpDebug::QueryStatsInfo{}; + ASSERT_DOES_NOT_THROW(query_stats::registerRequest(opCtx.get(), nss, [&]() { + uasserted(ErrorCodes::BSONObjectTooLarge, "size error"); + return nullptr; + })); + + // This should hit our assertion. + opDebug.queryStatsInfo = OpDebug::QueryStatsInfo{}; + ASSERT_THROWS(query_stats::registerRequest(opCtx.get(), + nss, + [&]() { + uasserted(ErrorCodes::BadValue, "fake error"); + return nullptr; + }), + DBException); +} + +} // namespace mongo::query_stats diff --git a/src/mongo/db/query/query_stats/rate_limiting.cpp b/src/mongo/db/query/query_stats/rate_limiting.cpp new file mode 100644 index 00000000000..aa8ca645bf1 --- /dev/null +++ b/src/mongo/db/query/query_stats/rate_limiting.cpp @@ -0,0 +1,96 @@ +/** + * Copyright (C) 2022-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. + */ + +#include "rate_limiting.h" +#include "mongo/stdx/mutex.h" +#include "mongo/util/clock_source.h" + +namespace mongo { +RateLimiting::RateLimiting(RequestCount samplingRate, + Milliseconds timePeriod, + ClockSource* clockSource) + : _clockSource(clockSource != nullptr ? clockSource : SystemClockSource::get()), + _samplingRate(samplingRate), + _timePeriod(timePeriod), + _windowStart(_clockSource->now()), + _prevCount(0), + _currentCount(0) {} + +Date_t RateLimiting::tickWindow() { + Date_t currentTime = _clockSource->now(); + + // Elapsed time since window start exceeds the time period. Start a new window. + if (currentTime - _windowStart > _timePeriod) { + _windowStart = currentTime; + _prevCount = _currentCount; + _currentCount = 0; + } + return currentTime; +} + +bool RateLimiting::handleRequestFixedWindow() { + stdx::unique_lock windowLock{_windowMutex}; + tickWindow(); + + if (_currentCount < _samplingRate.load()) { + _currentCount += 1; + return true; + } + return false; +} + +bool RateLimiting::handleRequestSlidingWindow() { + stdx::unique_lock windowLock{_windowMutex}; + + Date_t currentTime = tickWindow(); + auto windowStart = _windowStart; + auto prevCount = _prevCount; + + // Sliding window is implemented over fixed size time periods/blocks as follows. Instead of + // making the decision to limit the rate using only the current time period, we look to the rate + // of the previous period to predicate the rate of the current. This smooths the "sampling" of + // the events by predicting a constant rate and limiting accordingly. + + // Percentage of time remaining in current window. + double percentRemainingOfCurrentWindow = + ((double)(_timePeriod.count() - (currentTime - windowStart).count())) / _timePeriod.count(); + // Estimate the number of requests remaining in the current period. We assume the requests in + // the previous time block occurred at a constant rate. We multiply the total number of requests + // in the previous period by the percentage of time remaining in the current period. + double estimatedRemaining = prevCount * percentRemainingOfCurrentWindow; + // Add this estimate to the requests we know have taken place within the current time block. + double estimatedCount = _currentCount + estimatedRemaining; + + if (estimatedCount < _samplingRate.load()) { + _currentCount += 1; + return true; + } + return false; +} +} // namespace mongo diff --git a/src/mongo/db/query/query_stats/rate_limiting.h b/src/mongo/db/query/query_stats/rate_limiting.h new file mode 100644 index 00000000000..66e38d7119b --- /dev/null +++ b/src/mongo/db/query/query_stats/rate_limiting.h @@ -0,0 +1,126 @@ +/** + * Copyright (C) 2022-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 "mongo/util/clock_source.h" +#include "mongo/util/concurrency/mutex.h" +#include "mongo/util/system_clock_source.h" + +namespace mongo { + +/** + * Rate limiting is used to put a bound on the number of requests to a certain resource over a fixed + * time window. This implementation is approximate in the sense that it may permit the bound to + * exceeded. The bound is approximate as a trade off to reduce contention on internal resources. + */ +class RateLimiting { + using RequestCount = uint32_t; + +public: + /* + * Constructor for a rate limiter. Specify the number of requests you want to take place, as + * well as the time period in milliseconds. + */ + RateLimiting(RequestCount samplingRate, + Milliseconds timePeriod = Seconds{1}, + ClockSource* clockSource = nullptr); + + /* + * Getter for the sampling rate. + */ + RequestCount getSamplingRate() { + return _samplingRate.load(); + } + + /* + * Setter for the sampling rate. + */ + void setSamplingRate(RequestCount samplingRate) { + _samplingRate.store(samplingRate); + } + + /* + * A simple method for rate limiting. Returns false if we have reached the request limit for the + * current time window; otherwise, returns true and adds the request to the count for the + * current window. If we have passed the end of the previous window, the slate is wiped clean. + */ + bool handleRequestFixedWindow(); + + /* + * A method that ensures a more steady rate of requests. Rather than only looking at the current + * time block, this method simulates a sliding window to estimate how many requests occurred in + * the last full time period. Like the above, returns whether the request should be handled, and + * resets the window if enough time has passed. + */ + bool handleRequestSlidingWindow(); + +private: + /* + * Resets the current window if it has ended. Returns the current time. This must be called in + * the beginning of each handleRequest...() method. + */ + Date_t tickWindow(); + + /* + * Clock source used to track time. + */ + ClockSource* const _clockSource; + + /* + * Sampling rate is the bound on the number of requests we want to admit per window. + */ + AtomicWord<RequestCount> _samplingRate; + + /* + * Time period is the window size in ms. + */ + const Milliseconds _timePeriod; + + /* + * Window start. + */ + Date_t _windowStart; + + /* + * Count of requests handled in the previous window. + */ + RequestCount _prevCount; + + /* + * Count of requests handled in the current window. + */ + RequestCount _currentCount; + + /* + * Mutex used when reading/writing the window. + */ + SimpleMutex _windowMutex; +}; +} // namespace mongo diff --git a/src/mongo/db/query/query_stats/rate_limiting_bm.cpp b/src/mongo/db/query/query_stats/rate_limiting_bm.cpp new file mode 100644 index 00000000000..06308e6b0d8 --- /dev/null +++ b/src/mongo/db/query/query_stats/rate_limiting_bm.cpp @@ -0,0 +1,144 @@ +/** + * Copyright (C) 2023-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. + */ + + +#include <benchmark/benchmark.h> +#include <climits> +#include <memory> + +#include "mongo/bson/json.h" +#include "mongo/db/matcher/expression_leaf.h" +#include "mongo/db/matcher/expression_parser.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_stats/rate_limiting.h" +#include "mongo/util/duration.h" +#include "mongo/util/processinfo.h" +#include "mongo/util/time_support.h" + +namespace mongo { +namespace { + +// Local testing determined that these parameter values drove the most lock contention, which is +// what we want to capture in this benchmark. +constexpr long long rateLimitedWorkTimeMicros = 5; +constexpr long long consistentWorkTimeMicros = 10; + +constexpr long long numThreads = 256; + +// Rate limit some fraction of the overall work for a request with a sliding window. +int requestWithSlidingWindow(RateLimiting& limit) { + if (limit.handleRequestSlidingWindow()) { + sleepmicros(rateLimitedWorkTimeMicros); + } + sleepmicros(consistentWorkTimeMicros); + return 0; +} + +// Represent a request that bypasses the rate limiter. +int requestUnlimited() { + constexpr long long totalTime = rateLimitedWorkTimeMicros + consistentWorkTimeMicros; + sleepmicros(totalTime); + return 0; +} + +// Represent a request without the rate limited work. +int requestDeactivated() { + sleepmicros(consistentWorkTimeMicros); + return 0; +} + +// Benchmark sliding window rate limiting. +void BM_SlidingWindow(benchmark::State& state) { + // The rate limiter needs a clock source passed in. + static std::unique_ptr<ClockSource> clockSource; + static std::unique_ptr<RateLimiting> rateLimit; + + // Initialize the rate limiter only on the first thread to start up. + if (state.thread_index == 0) { + clockSource = std::make_unique<SystemClockSource>(); + rateLimit = + std::make_unique<RateLimiting>(state.range(0), Milliseconds(1), clockSource.get()); + } + + // Run the benchmark. + for (auto keepRunning : state) { + benchmark::DoNotOptimize(requestWithSlidingWindow(*rateLimit)); + } + + // Clean up the rate limiter when the benchmark is done. + if (state.thread_index == 0) { + rateLimit.reset(); + clockSource.reset(); + } +} + +// "Control" benchmark that does not rate limit requests. In other words, the extra work is always +// done for every request. This benchmark can be thought of as the "goal" performance for the peak, +// or the highest rate limit in BM_SlidingWindow, to compare against. +void BM_Unlimited(benchmark::State& state) { + for (auto keepRunning : state) { + benchmark::DoNotOptimize(requestUnlimited()); + } +} +// Another control benchmark, where the extra work is never done for any request. This can be +// thought of as the goal performance for when rate limit equals 0. +void BM_Deactivated(benchmark::State& state) { + for (auto keepRunning : state) { + benchmark::DoNotOptimize(requestDeactivated()); + } +} + +// Google microbenchmarks report time T (in nanoseconds) spent per operation. But at Mongo we are +// interested in total opereations performed per second. The former can easily be converted to the +// latter by diving 10^6 by T. Use this benchmark to determine the natural throughput of the +// operation. This can be compared to the rate limited benchmarks (BM_SlidingWindow) to determine +// the overhead of rate limiting. Looking at the percentage change in throughput between the control +// benchmarks and the rate limited benchmark, will indicate how much overhead is due to lock +// contention. +BENCHMARK(BM_Unlimited)->Threads(numThreads); + +BENCHMARK(BM_Deactivated)->Threads(numThreads); + +// Local testing has confirmed that the higher the rate limit, the worse the throughput. This makes +// sense as putting a higher upper bound on number of requests allowed in a given time period, means +// longer wait times for the lock. +BENCHMARK(BM_SlidingWindow) + ->ArgName("rate limit") + ->Arg(0) + ->Arg(64) + ->Arg(128) + ->Arg(256) + ->Arg(512) + ->Arg(1024) + ->Arg(2048) + ->Arg(4816) + ->Threads(numThreads); + +} // namespace +} // namespace mongo diff --git a/src/mongo/db/query/query_stats/rate_limiting_test.cpp b/src/mongo/db/query/query_stats/rate_limiting_test.cpp new file mode 100644 index 00000000000..380636a2a20 --- /dev/null +++ b/src/mongo/db/query/query_stats/rate_limiting_test.cpp @@ -0,0 +1,77 @@ +/** + * Copyright (C) 2022-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. + */ + +#include "mongo/db/query/query_stats/rate_limiting.h" +#include "mongo/unittest/unittest.h" +#include "mongo/util/time_support.h" + +namespace mongo { +TEST(RateLimitingTest, FixedWindowSucceeds) { + auto rl = RateLimiting(1); + ASSERT_TRUE(rl.handleRequestFixedWindow()); +} + +TEST(RateLimitingTest, SlidingWindowSucceeds) { + auto rl = RateLimiting(1); + ASSERT_TRUE(rl.handleRequestSlidingWindow()); +} + +TEST(RateLimitingTest, FixedWindowFails) { + auto rl = RateLimiting(0); + ASSERT_FALSE(rl.handleRequestFixedWindow()); +} + +TEST(RateLimitingTest, SlidingWindowFails) { + auto rl = RateLimiting(0); + ASSERT_FALSE(rl.handleRequestSlidingWindow()); +} + +TEST(RateLimitingTest, FixedWindowSucceedsThenFails) { + auto rl = RateLimiting(1, Hours{1}); + ASSERT_TRUE(rl.handleRequestFixedWindow()); + ASSERT_FALSE(rl.handleRequestFixedWindow()); + ASSERT_FALSE(rl.handleRequestFixedWindow()); +} + +TEST(RateLimitingTest, SlidingWindowSucceedsThenFails) { + auto rl = RateLimiting(1, Hours{1}); + ASSERT_TRUE(rl.handleRequestSlidingWindow()); + ASSERT_FALSE(rl.handleRequestSlidingWindow()); + ASSERT_FALSE(rl.handleRequestSlidingWindow()); +} + +TEST(RateLimitingTest, FixedWindowPermitsRequestAfterWindowExpires) { + auto rl = RateLimiting(1, Milliseconds{10}); + ASSERT_TRUE(rl.handleRequestFixedWindow()); + ASSERT_FALSE(rl.handleRequestFixedWindow()); + sleepmillis(11); + ASSERT_TRUE(rl.handleRequestFixedWindow()); +} + +} // namespace mongo diff --git a/src/mongo/db/query/query_stats/shapifying_bm.cpp b/src/mongo/db/query/query_stats/shapifying_bm.cpp new file mode 100644 index 00000000000..fd7f605c855 --- /dev/null +++ b/src/mongo/db/query/query_stats/shapifying_bm.cpp @@ -0,0 +1,142 @@ +/** + * Copyright (C) 2023-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. + */ + + +#include <benchmark/benchmark.h> +#include <climits> +#include <memory> + +#include "mongo/bson/json.h" +#include "mongo/db/concurrency/locker_noop_client_observer.h" +#include "mongo/db/matcher/expression_leaf.h" +#include "mongo/db/matcher/expression_parser.h" +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/query_shape/query_shape.h" +#include "mongo/db/query/query_stats/find_key.h" +#include "mongo/db/query/query_stats/query_stats.h" +#include "mongo/db/query/query_stats/rate_limiting.h" +#include "mongo/idl/server_parameter_test_util.h" +#include "mongo/rpc/metadata/client_metadata.h" +#include "mongo/util/duration.h" +#include "mongo/util/processinfo.h" +#include "mongo/util/testing_proctor.h" +#include "mongo/util/time_support.h" + +namespace mongo { +namespace { + +static const NamespaceStringOrUUID kDefaultTestNss = + NamespaceStringOrUUID{NamespaceString("testDB.testColl")}; + +static constexpr auto kCollectionType = query_shape::CollectionType::kCollection; + +// This is a snapshot of the client metadata generated from our IDHACK genny workload. The +// specifics aren't so important, but it chosen in an attempt to be indicative of the size/shape +// of this kind of thing "in the wild". +const auto kMetadataWrapper = fromjson(R"({metadata: { + "application" : { + "name" : "Genny" + }, + "driver" : { + "name" : "mongoc / mongocxx", + "version" : "1.23.2 / 3.7.0" + }, + "os" : { + "type" : "Linux", + "name" : "Ubuntu", + "version" : "22.04", + "architecture" : "aarch64" + }, + "platform" : "cfg=0x03215e88e9 posix=200809 stdc=201710 CC=GCC 11.3.0 CFLAGS=\"-fPIC\" LDFLAGS=\"\"" + }})"); +auto kMockClientMetadataElem = kMetadataWrapper["metadata"]; + +auto makeFindKey(const boost::intrusive_ptr<ExpressionContext>& expCtx, + const ParsedFindCommand& parsedFind) { + return std::make_unique<const query_stats::FindKey>(expCtx, parsedFind, kCollectionType); +} + +int shapifyAndHashRequest(const boost::intrusive_ptr<ExpressionContext>& expCtx, + const ParsedFindCommand& parsedFind) { + auto key = makeFindKey(expCtx, parsedFind); + [[maybe_unused]] auto hash = absl::Hash<query_stats::Key>{}(*key); + return 0; +} + +// Benchmark the performance of computing and hashing the query stats key for an IDHACK query. +void BM_ShapfiyIDHack(benchmark::State& state) { + auto serviceCtx = ServiceContext::make(); + serviceCtx->registerClientObserver(std::make_unique<LockerNoopClientObserver>()); + + auto client = serviceCtx->makeClient("query_test"); + + auto opCtx = client->makeOperationContext(); + auto expCtx = make_intrusive<ExpressionContextForTest>(opCtx.get()); + auto fcr = std::make_unique<FindCommandRequest>(expCtx->ns); + fcr->setFilter(fromjson("{_id: 4}")); + ClientMetadata::setFromMetadata(opCtx->getClient(), kMockClientMetadataElem); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(expCtx, std::move(fcr))); + + // Run the benchmark. + for (auto keepRunning : state) { + benchmark::DoNotOptimize(shapifyAndHashRequest(expCtx, *parsedFind)); + } +} + +// Benchmark computing the query stats key and its hash for a mildly complex query predicate. +void BM_ShapfiyMildlyComplex(benchmark::State& state) { + auto serviceCtx = ServiceContext::make(); + serviceCtx->registerClientObserver(std::make_unique<LockerNoopClientObserver>()); + + auto client = serviceCtx->makeClient("query_test"); + + auto opCtx = client->makeOperationContext(); + auto expCtx = make_intrusive<ExpressionContextForTest>(opCtx.get()); + auto fcr = std::make_unique<FindCommandRequest>(expCtx->ns); + fcr->setFilter(fromjson(R"({ + clientId: {$nin: ["432345", "4386945", "111111"]}, + nEmployees: {$gte: 4, $lt: 20}, + deactivated: false, + region: "US", + yearlySpend: {$lte: 1000} + })")); + ClientMetadata::setFromMetadata(opCtx->getClient(), kMockClientMetadataElem); + auto parsedFind = uassertStatusOK(parsed_find_command::parse(expCtx, std::move(fcr))); + + // Run the benchmark. + for (auto keepRunning : state) { + benchmark::DoNotOptimize(shapifyAndHashRequest(expCtx, *parsedFind)); + } +} + +BENCHMARK(BM_ShapfiyIDHack)->Threads(1); +BENCHMARK(BM_ShapfiyMildlyComplex)->Threads(1); + +} // namespace +} // namespace mongo diff --git a/src/mongo/db/query/query_stats/transform_algorithm.idl b/src/mongo/db/query/query_stats/transform_algorithm.idl new file mode 100644 index 00000000000..cd0a5ba43db --- /dev/null +++ b/src/mongo/db/query/query_stats/transform_algorithm.idl @@ -0,0 +1,37 @@ +# Copyright (C) 2023-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. +# +global: + cpp_namespace: "mongo" + +enums: + TransformAlgorithm: + description: "The type of algorithm to be used for the transformIdentifiers field of $queryStats." + type: string + values: + kHmacSha256: "hmac-sha-256" + kNone: "none" diff --git a/src/mongo/db/query/record_id_bound.h b/src/mongo/db/query/record_id_bound.h index 99400ae938d..d6a37617a4f 100644 --- a/src/mongo/db/query/record_id_bound.h +++ b/src/mongo/db/query/record_id_bound.h @@ -29,6 +29,7 @@ #pragma once +#include <boost/operators.hpp> #include <boost/optional.hpp> #include <fmt/format.h> #include <ostream> @@ -44,7 +45,7 @@ namespace mongo { /** * A RecordId bound for a collection scan, with an optional BSON representation for pretty printing. */ -class RecordIdBound { +class RecordIdBound : boost::totally_ordered<RecordIdBound> { public: RecordIdBound() = default; @@ -81,6 +82,14 @@ public: return _recordId.compare(rhs._recordId); } + bool operator==(const RecordIdBound& rhs) const { + return compare(rhs) == 0; + } + + bool operator<(const RecordIdBound& rhs) const { + return compare(rhs) < 0; + } + private: RecordId _recordId; boost::optional<BSONObj> _bson; diff --git a/src/mongo/db/query/record_id_range.cpp b/src/mongo/db/query/record_id_range.cpp new file mode 100644 index 00000000000..88720183ef4 --- /dev/null +++ b/src/mongo/db/query/record_id_range.cpp @@ -0,0 +1,109 @@ +/** + * Copyright (C) 2023-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. + */ + +#include <boost/optional.hpp> + +#include "mongo/db/query/record_id_range.h" + +namespace mongo { +void RecordIdRange::maybeNarrowMin(const BSONObj& newMin, bool inclusive) { + maybeNarrowMin(RecordIdBound(record_id_helpers::keyForObj(newMin), newMin), inclusive); +} + +void RecordIdRange::maybeNarrowMin(const RecordIdBound& newMin, bool inclusive) { + if (_min) { + auto cmp = _min->compare(newMin); + // The range only needs updating if: + // * There's no existing _min + // * The provided value is greater than the current _min + // * The value == _min, but is _not_ inclusive, but the existing value is + + if (cmp > 0) { + // Current min is strictly greater than the provided value (and existing value has been + // initialised), nothing to do. + return; + } + + if (cmp == 0) { + // Inclusivity moving true -> false narrows the range. + _minInclusive = _minInclusive && inclusive; + return; + } + } + _min = newMin; + // The bound value changed, so the previous value of _minInclusive is irrelevant. + _minInclusive = inclusive; +} + +void RecordIdRange::maybeNarrowMax(const BSONObj& newMax, bool inclusive) { + maybeNarrowMax(RecordIdBound(record_id_helpers::keyForObj(newMax), newMax), inclusive); +} + +void RecordIdRange::maybeNarrowMax(const RecordIdBound& newMax, bool inclusive) { + if (_max) { + auto cmp = _max->compare(newMax); + // The range only needs updating if: + // * There's no existing _max + // * The provided value is less than the current _max + // * The value == _max, but is _not_ inclusive, but the existing value is + + if (cmp < 0) { + // Current max is strictly less than the provided value (and existing value has been + // initialised), nothing to do. + return; + } + + if (cmp == 0) { + // Inclusivity moving true -> false narrows the range. + _maxInclusive = _maxInclusive && inclusive; + return; + } + } + _max = newMax; + // The bound value changed, so the previous value of _maxInclusive is irrelevant. + _maxInclusive = inclusive; +} + +void RecordIdRange::intersectRange(const RecordIdRange& other) { + intersectRange(other._min, other._max, other._minInclusive, other._maxInclusive); +} + +void RecordIdRange::intersectRange(const boost::optional<RecordIdBound>& min, + const boost::optional<RecordIdBound>& max, + bool minInclusive, + bool maxInclusive) { + if (min) { + maybeNarrowMin(*min, minInclusive); + } + if (max) { + maybeNarrowMax(*max, maxInclusive); + } +} + +} // namespace mongo diff --git a/src/mongo/db/query/record_id_range.h b/src/mongo/db/query/record_id_range.h new file mode 100644 index 00000000000..d25e17356c1 --- /dev/null +++ b/src/mongo/db/query/record_id_range.h @@ -0,0 +1,117 @@ +/** + * Copyright (C) 2023-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 <boost/optional.hpp> + +#include "mongo/bson/bsonobj.h" +#include "mongo/db/query/record_id_bound.h" + +namespace mongo { + +class RecordIdRange { +public: + /** + * If the provided value @p newMin is greater than the existing min, + * update the lower bound to equal @p newMin + * + * @return true if range was adjusted + */ + void maybeNarrowMin(const BSONObj& newMin, bool inclusive); + void maybeNarrowMin(const RecordIdBound& newMin, bool inclusive); + + /** + * If the provided value @p newMax is less than the existing max, + * update the upper bound to equal @p newMax + * + * @return true if range was adjusted + */ + void maybeNarrowMax(const BSONObj& newMax, bool inclusive); + void maybeNarrowMax(const RecordIdBound& newMax, bool inclusive); + + /** + * Update this range to the intersection of this range + * and @p other. This may update both, one of, or neither of + * min and max. + * + * Results in a range which is either unchanged, or made + * narrower (possibly becoming an empty range). + */ + void intersectRange(const RecordIdRange& other); + /** + * Overload of intersectRange taking the components of a RecordIdRange, + * for convenience when the other range is not handled as a RecordIdRange. + */ + void intersectRange(const boost::optional<RecordIdBound>& min, + const boost::optional<RecordIdBound>& max, + bool minInclusive = true, + bool maxInclusive = true); + + bool isEmpty() const; + + + const auto& getMin() const { + return _min; + } + + const auto& getMax() const { + return _max; + } + + bool isMinInclusive() const { + return _minInclusive; + } + + bool isMaxInclusive() const { + return _maxInclusive; + } + + +private: + // If present, this parameter sets the start point of a forward scan or the end point of a + // reverse scan. + boost::optional<RecordIdBound> _min; + + // If present, this parameter sets the start point of a reverse scan or the end point of a + // forward scan. + boost::optional<RecordIdBound> _max; + + // TODO: investigate folding this into RecordIdBound; many other usages pair RecordIdBound + // with ScanBoundInclusion to convey this information + // If min is present, this indicates whether the range is inclusive or exclusive of the + // set min value + bool _minInclusive = true; + // If max is present, this indicates whether the range is inclusive or exclusive of the + // set max value + bool _maxInclusive = true; +}; + +} // namespace mongo diff --git a/src/mongo/db/query/record_id_range_test.cpp b/src/mongo/db/query/record_id_range_test.cpp new file mode 100644 index 00000000000..5db866a73d3 --- /dev/null +++ b/src/mongo/db/query/record_id_range_test.cpp @@ -0,0 +1,119 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/base/error_extra_info.h" +#include "mongo/bson/bsonmisc.h" +#include "mongo/db/query/record_id_range.h" + +#include "mongo/bson/bsonmisc.h" +#include "mongo/unittest/unittest.h" + +namespace { + +using namespace mongo; + +template <class BoundsCheck, class NarrowingCheck, class InclusivityCheck> +void testRange(const BoundsCheck& getBoundOptional, + const NarrowingCheck& maybeNarrowBound, + const InclusivityCheck& checkInclusivity, + int initialValue, + int narrowerValue, + int widerValue) { + ASSERT_FALSE(getBoundOptional()); + + auto assertValueEq = [&](auto value) { + auto bson = BSON("value" << value); + auto recordId = record_id_helpers::keyForObj(bson); + ASSERT_EQ(recordId, getBoundOptional()->recordId()); + }; + + // narrow from unset + maybeNarrowBound(BSON("value" << initialValue), true /* inclusive */); + ASSERT_TRUE(getBoundOptional()); + assertValueEq(initialValue); + ASSERT_TRUE(checkInclusivity()); + + // narrow by removing inclusivity of the bound + maybeNarrowBound(BSON("value" << initialValue), false /* not inclusive */); + ASSERT_TRUE(getBoundOptional()); + assertValueEq(initialValue); + ASSERT_FALSE(checkInclusivity()); + + // cannot widen by re-adding inclusivity + maybeNarrowBound(BSON("value" << initialValue), true /* inclusive */); + ASSERT_TRUE(getBoundOptional()); + assertValueEq(initialValue); + ASSERT_FALSE(checkInclusivity()); + + // cannot widen by setting a wider bound + maybeNarrowBound(BSON("value" << widerValue), true /* inclusive */); + ASSERT_TRUE(getBoundOptional()); + assertValueEq(initialValue); + ASSERT_FALSE(checkInclusivity()); + + // cannot widen by setting a wider bound, regardless of inclusivity + maybeNarrowBound(BSON("value" << widerValue), false /* not inclusive */); + ASSERT_TRUE(getBoundOptional()); + assertValueEq(initialValue); + ASSERT_FALSE(checkInclusivity()); + + // narrow to a non-inclusive bound at a narrower value + maybeNarrowBound(BSON("value" << narrowerValue), false /* not inclusive */); + ASSERT_TRUE(getBoundOptional()); + assertValueEq(narrowerValue); + ASSERT_FALSE(checkInclusivity()); +} + +TEST(RecordIdRangeTest, NarrowMin) { + RecordIdRange range; + + testRange([&] { return range.getMin(); }, + [&](const BSONObj& newVal, bool inclusive) { + return range.maybeNarrowMin(newVal, inclusive); + }, + [&] { return range.isMinInclusive(); }, + 10, + 11, + 9); +} + +TEST(RecordIdRangeTest, NarrowMax) { + RecordIdRange range; + + testRange([&] { return range.getMax(); }, + [&](const BSONObj& newVal, bool inclusive) { + return range.maybeNarrowMax(newVal, inclusive); + }, + [&] { return range.isMaxInclusive(); }, + 10, + 9, + 11); +} + +} // namespace diff --git a/src/mongo/db/query/sbe_cached_solution_planner.cpp b/src/mongo/db/query/sbe_cached_solution_planner.cpp index c594c70f22e..cd7f23016eb 100644 --- a/src/mongo/db/query/sbe_cached_solution_planner.cpp +++ b/src/mongo/db/query/sbe_cached_solution_planner.cpp @@ -74,6 +74,7 @@ CandidatePlans CachedSolutionPlanner::plan( std::move(roots[0].first), std::move(roots[0].second), maxReadsBeforeReplan); + auto explainer = plan_explainer_factory::make( candidate.root.get(), &candidate.data, diff --git a/src/mongo/db/query/sbe_plan_cache.cpp b/src/mongo/db/query/sbe_plan_cache.cpp index c8944057639..70d864e7074 100644 --- a/src/mongo/db/query/sbe_plan_cache.cpp +++ b/src/mongo/db/query/sbe_plan_cache.cpp @@ -31,7 +31,7 @@ #include "mongo/db/query/sbe_plan_cache.h" -#include "mongo/db/query/plan_cache_size_parameter.h" +#include "mongo/db/query/util/memory_util.h" #include "mongo/db/server_options.h" #include "mongo/logv2/log.h" #include "mongo/util/processinfo.h" @@ -42,71 +42,23 @@ namespace { const auto sbePlanCacheDecoration = ServiceContext::declareDecoration<std::unique_ptr<sbe::PlanCache>>(); -size_t convertToSizeInBytes(const plan_cache_util::PlanCacheSizeParameter& param) { - constexpr size_t kBytesInMB = 1024 * 1024; - constexpr size_t kMBytesInGB = 1024; - - double sizeInMB = param.size; - - switch (param.units) { - case plan_cache_util::PlanCacheSizeUnits::kPercent: - sizeInMB *= ProcessInfo::getMemSizeMB() / 100.0; - break; - case plan_cache_util::PlanCacheSizeUnits::kMB: - break; - case plan_cache_util::PlanCacheSizeUnits::kGB: - sizeInMB *= kMBytesInGB; - break; - } - - return static_cast<size_t>(sizeInMB * kBytesInMB); -} - -/** - * Sets upper size limit on the PlanCache size to 500GB or 25% of the system's memory, whichever is - * smaller. - */ -size_t capPlanCacheSize(size_t planCacheSize) { - constexpr size_t kBytesInGB = 1024 * 1024 * 1024; - - // Maximum size of the plan cache expressed in bytes. - constexpr size_t kMaximumPlanCacheSize = 500 * kBytesInGB; - - // Maximum size of the plan cache expressed as a share of the memory available to the process. - const plan_cache_util::PlanCacheSizeParameter limitToProcessSize{ - 25, plan_cache_util::PlanCacheSizeUnits::kPercent}; - const size_t limitToProcessSizeInBytes = convertToSizeInBytes(limitToProcessSize); - - // The size will be capped by the minimum of the two values defined above. - const size_t maxPlanCacheSize = std::min(kMaximumPlanCacheSize, limitToProcessSizeInBytes); - - if (planCacheSize > maxPlanCacheSize) { - planCacheSize = maxPlanCacheSize; - LOGV2_DEBUG(6007000, - 1, - "The plan cache size has been capped", - "maxPlanCacheSize"_attr = maxPlanCacheSize); - } - - return planCacheSize; -} - -size_t getPlanCacheSizeInBytes(const plan_cache_util::PlanCacheSizeParameter& param) { - size_t planCacheSize = convertToSizeInBytes(param); - uassert(5968001, - "Cache size must be at least 1KB * number of cores", - planCacheSize >= 1024 * ProcessInfo::getNumCores()); - return capPlanCacheSize(planCacheSize); -} class PlanCacheOnParamChangeUpdaterImpl final : public plan_cache_util::OnParamChangeUpdater { public: - void updateCacheSize(ServiceContext* serviceCtx, - plan_cache_util::PlanCacheSizeParameter parameter) final { + void updateCacheSize(ServiceContext* serviceCtx, memory_util::MemorySize memSize) final { if (feature_flags::gFeatureFlagSbePlanCache.isEnabledAndIgnoreFCV()) { - auto size = getPlanCacheSizeInBytes(parameter); + auto newSizeBytes = memory_util::getRequestedMemSizeInBytes(memSize); + auto cappedCacheSize = memory_util::capMemorySize(newSizeBytes /*requestedSizeBytes*/, + 500 /*maximumSizeGB*/, + 25 /*percentTotalSystemMemory*/); + if (cappedCacheSize < newSizeBytes) { + LOGV2_DEBUG(6007001, + 1, + "The plan cache size has been capped", + "cappedSize"_attr = cappedCacheSize); + } auto& globalPlanCache = sbePlanCacheDecoration(serviceCtx); - globalPlanCache->reset(size); + globalPlanCache->reset(cappedCacheSize); } } @@ -124,12 +76,21 @@ ServiceContext::ConstructorActionRegisterer planCacheRegisterer{ std::make_unique<PlanCacheOnParamChangeUpdaterImpl>(); if (feature_flags::gFeatureFlagSbePlanCache.isEnabledAndIgnoreFCV()) { - auto status = plan_cache_util::PlanCacheSizeParameter::parse(planCacheSize.get()); + auto status = memory_util::MemorySize::parse(planCacheSize.get()); uassertStatusOK(status); - - auto size = getPlanCacheSizeInBytes(status.getValue()); + auto size = memory_util::getRequestedMemSizeInBytes(status.getValue()); + auto cappedCacheSize = memory_util::capMemorySize(size /*requestedSizeBytes*/, + 500 /*maximumSizeGB*/, + 25 /*percentTotalSystemMemory*/); + if (cappedCacheSize < size) { + LOGV2_DEBUG(6007000, + 1, + "The plan cache size has been capped", + "cappedSize"_attr = cappedCacheSize); + } auto& globalPlanCache = sbePlanCacheDecoration(serviceCtx); - globalPlanCache = std::make_unique<sbe::PlanCache>(size, ProcessInfo::getNumCores()); + globalPlanCache = + std::make_unique<sbe::PlanCache>(cappedCacheSize, ProcessInfo::getNumCores()); } }}; diff --git a/src/mongo/db/query/sbe_plan_cache.h b/src/mongo/db/query/sbe_plan_cache.h index 6e7853fa817..090df251bdf 100644 --- a/src/mongo/db/query/sbe_plan_cache.h +++ b/src/mongo/db/query/sbe_plan_cache.h @@ -190,7 +190,14 @@ struct CachedSbePlan { using PlanCacheEntry = PlanCacheEntryBase<CachedSbePlan, plan_cache_debug_info::DebugInfoSBE>; struct BudgetEstimator { - size_t operator()(const std::shared_ptr<const PlanCacheEntry>& entry) { + /** + * This estimator function is called when an entry is added or removed to LRU cache in order to + * make sure the total plan cache size does not exceed the maximum size. + */ + size_t operator()(const sbe::PlanCacheKey& key, + const std::shared_ptr<const PlanCacheEntry>& entry) { + // TODO: SERVER-73649 include size of underlying query shape and size of int_32 key hash in + // total size estimation. return entry->estimatedEntrySizeBytes; } }; diff --git a/src/mongo/db/query/sbe_plan_cache_on_parameter_change.cpp b/src/mongo/db/query/sbe_plan_cache_on_parameter_change.cpp index 4279a142e4b..9f6cbef6d4f 100644 --- a/src/mongo/db/query/sbe_plan_cache_on_parameter_change.cpp +++ b/src/mongo/db/query/sbe_plan_cache_on_parameter_change.cpp @@ -56,7 +56,7 @@ Status clearSbeCacheOnParameterChangeHelper() { } Status onPlanCacheSizeUpdate(const std::string& str) { - auto newSize = PlanCacheSizeParameter::parse(str); + auto newSize = memory_util::MemorySize::parse(str); if (!newSize.isOK()) { return newSize.getStatus(); } diff --git a/src/mongo/db/query/sbe_plan_cache_on_parameter_change.h b/src/mongo/db/query/sbe_plan_cache_on_parameter_change.h index e126dda9f48..d0f8a026a05 100644 --- a/src/mongo/db/query/sbe_plan_cache_on_parameter_change.h +++ b/src/mongo/db/query/sbe_plan_cache_on_parameter_change.h @@ -32,7 +32,7 @@ #include <string> #include "mongo/base/status.h" -#include "mongo/db/query/plan_cache_size_parameter.h" +#include "mongo/db/query/util/memory_util.h" #include "mongo/db/service_context.h" namespace mongo::plan_cache_util { @@ -70,11 +70,11 @@ public: virtual ~OnParamChangeUpdater() = default; /** - * Resizes the SBE plan cache decorating 'serviceCtx' to the new size given by 'parameter'. If + * Resizes the SBE plan cache decorating 'serviceCtx' to the new size given by 'memSize'. If * the new cache size is smaller than the old, cache entries are evicted in order to ensure the * cache fits within the new size bound. */ - virtual void updateCacheSize(ServiceContext* serviceCtx, PlanCacheSizeParameter parameter) = 0; + virtual void updateCacheSize(ServiceContext* serviceCtx, memory_util::MemorySize memSize) = 0; /** * Deletes all plans from the SBE plan cache decorating 'serviceCtx'. diff --git a/src/mongo/db/query/sbe_stage_builder_expression.cpp b/src/mongo/db/query/sbe_stage_builder_expression.cpp index 89a8c2d68ff..8ecc4961435 100644 --- a/src/mongo/db/query/sbe_stage_builder_expression.cpp +++ b/src/mongo/db/query/sbe_stage_builder_expression.cpp @@ -2122,12 +2122,52 @@ public: std::vector<EvalExprStagePair> branches; branches.reserve(numChildren); + auto childStageCount = 0; for (size_t i = 0; i < numChildren; ++i) { auto [expr, stage] = _context->popFrame(); + if (stage.stage.get() != nullptr) { + childStageCount++; + } branches.emplace_back(std::move(expr), std::move(stage)); } std::reverse(branches.begin(), branches.end()); + // If there is no separate child stage branch, then we can implement $ifNull as a simple + // projection of SBE if expression, instead of with union stages. + if (childStageCount == 0) { + auto stage = _context->extractCurrentEvalStage(); + + std::vector<sbe::value::SlotId> slots; + slots.reserve(branches.size()); + sbe::value::SlotMap<std::unique_ptr<sbe::EExpression>> projects; + for (auto& branch : branches) { + if (branch.first.getSlot()) { + slots.push_back(*branch.first.getSlot()); + } else { + auto slot = _context->state.slotId(); + slots.push_back(slot); + projects.emplace(slot, branch.first.extractExpr()); + } + } + if (!projects.empty()) { + stage = makeProject(std::move(stage), std::move(projects), _context->planNodeId); + } + + auto expr = sbe::makeE<sbe::EVariable>(slots[slots.size() - 1]); + for (int i = slots.size() - 2; i >= 0; i--) { + auto thenExpr = sbe::makeE<sbe::EVariable>(slots[i]); + auto condExpr = makeNot(generateNullOrMissing(thenExpr->clone())); + expr = + sbe::makeE<sbe::EIf>(std::move(condExpr), std::move(thenExpr), std::move(expr)); + } + + auto outSlot = _context->state.slotId(); + stage = makeProject(std::move(stage), _context->planNodeId, outSlot, std::move(expr)); + + _context->pushExpr(outSlot, std::move(stage)); + return; + } + // Prepare to create limit-1/union with N branches (where N is the number of operands). Each // branch will be evaluated from left to right until one of the branches produces a value. auto branchFn = [](EvalExpr evalExpr, diff --git a/src/mongo/db/query/sort_pattern.cpp b/src/mongo/db/query/sort_pattern.cpp index fcd3cd177e1..5a444cd31f2 100644 --- a/src/mongo/db/query/sort_pattern.cpp +++ b/src/mongo/db/query/sort_pattern.cpp @@ -112,13 +112,13 @@ QueryMetadataBitSet SortPattern::metadataDeps(QueryMetadataBitSet unavailableMet return depsTracker.metadataDeps(); } -Document SortPattern::serialize(SortKeySerialization serializationMode) const { +Document SortPattern::serialize(SortKeySerialization serializationMode, + const SerializationOptions& options) const { MutableDocument keyObj; const size_t n = _sortPattern.size(); for (size_t i = 0; i < n; ++i) { if (_sortPattern[i].fieldPath) { - // Append a named integer based on whether the sort is ascending/descending. - keyObj.setField(_sortPattern[i].fieldPath->fullPath(), + keyObj.setField(options.serializeFieldPath(*_sortPattern[i].fieldPath), Value(_sortPattern[i].isAscending ? 1 : -1)); } else { // Sorting by an expression, use a made up field name. @@ -127,7 +127,12 @@ Document SortPattern::serialize(SortKeySerialization serializationMode) const { case SortKeySerialization::kForExplain: case SortKeySerialization::kForPipelineSerialization: { const bool isExplain = (serializationMode == SortKeySerialization::kForExplain); - keyObj[computedFieldName] = _sortPattern[i].expression->serialize(isExplain); + auto opts = SerializationOptions{}; + if (isExplain) { + opts.verbosity = + boost::make_optional(ExplainOptions::Verbosity::kQueryPlanner); + } + keyObj[computedFieldName] = _sortPattern[i].expression->serialize(opts); break; } case SortKeySerialization::kForSortKeyMerging: { diff --git a/src/mongo/db/query/sort_pattern.h b/src/mongo/db/query/sort_pattern.h index b659ed0124e..c709a43eac8 100644 --- a/src/mongo/db/query/sort_pattern.h +++ b/src/mongo/db/query/sort_pattern.h @@ -33,6 +33,7 @@ #include "mongo/db/exec/document_value/document.h" #include "mongo/db/pipeline/document_path_support.h" #include "mongo/db/pipeline/expression.h" +#include "mongo/db/query/query_shape/serialization_options.h" namespace mongo { class SortPattern { @@ -72,7 +73,8 @@ public: /** * Write out a Document whose contents are the sort key pattern. */ - Document serialize(SortKeySerialization) const; + Document serialize(SortKeySerialization serializationMode, + const SerializationOptions& options = {}) const; /** * Serializes the document to BSON, only keeping the paths specified in the sort pattern. diff --git a/src/mongo/db/query/sort_pattern_test.cpp b/src/mongo/db/query/sort_pattern_test.cpp new file mode 100644 index 00000000000..2abe1bffa81 --- /dev/null +++ b/src/mongo/db/query/sort_pattern_test.cpp @@ -0,0 +1,97 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "query_shape/serialization_options.h" + +#include "mongo/db/exec/document_value/document.h" +#include "mongo/db/exec/document_value/document_value_test_util.h" +#include "mongo/db/exec/document_value/value.h" +#include "mongo/db/pipeline/expression_context_for_test.h" +#include "mongo/db/query/sort_pattern.h" +#include "mongo/unittest/unittest.h" +namespace mongo { +namespace { + +auto getExpCtx() { + auto nss = NamespaceString("db", "coll"); + return boost::intrusive_ptr<ExpressionContextForTest>{new ExpressionContextForTest(nss)}; +} + +TEST(SerializeSortPatternTest, SerializeAndRedactFieldName) { + auto expCtx = getExpCtx(); + auto sortPattern = SortPattern(fromjson("{val: 1}"), expCtx); + SerializationOptions opts = SerializationOptions::kMarkIdentifiers_FOR_TEST; + + // Most basic sort pattern, confirm that field name gets redacted. + ASSERT_DOCUMENT_EQ_AUTO( // NOLINT + R"({"HASH<val>":1})", + sortPattern.serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts)); + + // Confirm that multiple sort fields get redacted. + sortPattern = SortPattern(fromjson("{val: 1, test: -1, third: -1}"), expCtx); + ASSERT_DOCUMENT_EQ_AUTO( // NOLINT + R"({"HASH<val>":1,"HASH<test>":-1,"HASH<third>":-1})", + sortPattern.serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts)); + + // Test sort pattern that contains an expression. + sortPattern = SortPattern(fromjson("{val: 1, test: {$meta: \"randVal\"}}"), expCtx); + ASSERT_DOCUMENT_EQ_AUTO( // NOLINT + R"({"HASH<val>":1,"$computed1":{"$meta":"randVal"}})", + sortPattern.serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts)); + + // Sorting by only an expression results in a made up field name in serialization and therefore + // doesn't get redacted. + sortPattern = SortPattern(fromjson("{val: {$meta: \"textScore\"}}"), expCtx); + ASSERT_DOCUMENT_EQ_AUTO( // NOLINT + R"({"$computed0":{"$meta":"textScore"}})", + sortPattern.serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts)); + + sortPattern = SortPattern(fromjson("{'a.b.c': 1}"), expCtx); + ASSERT_DOCUMENT_EQ_AUTO( // NOLINT + R"({"HASH<a>.HASH<b>.HASH<c>":1})", + sortPattern.serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts)); +} + +TEST(SerializeSortPatternTest, SerializeNoRedaction) { + auto expCtx = getExpCtx(); + auto sortPattern = SortPattern(fromjson("{val: 1}"), expCtx); + SerializationOptions opts = {}; + opts.transformIdentifiers = false; + ASSERT_DOCUMENT_EQ_AUTO( // NOLINT + R"({"val":1})", + sortPattern.serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts)); + + // Call serialize() with no options. + ASSERT_DOCUMENT_EQ_AUTO( // NOLINT + R"({"val":1})", + sortPattern.serialize(SortPattern::SortKeySerialization::kForPipelineSerialization)); +} + +} // namespace +} // namespace mongo diff --git a/src/mongo/db/query/util/deferred.h b/src/mongo/db/query/util/deferred.h new file mode 100644 index 00000000000..a2609bb6b49 --- /dev/null +++ b/src/mongo/db/query/util/deferred.h @@ -0,0 +1,118 @@ +/** + * Copyright (C) 2023-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> + +namespace mongo { + +/** + * A template class that provides a way to defer the initialization of an object until its value is + * actually required. This is also commonly referred to as lazy initialization. + * + * Dangers: + * - This implementation is currently not thread safe, and it shouldn't be used in multi-threaded + * fashion. + * - Be careful about using this for lazy initialization of data members and capturing the 'this' + * variable. Code like this will result in buggy/unsafe move constructors, which would have a + * dangling reference to the moved-from type: + * + * class MyType { + * int x; + * // !!! Dangling 'this' when moved !!! + * Deferred<int> xSquared{[this]() { return this->x * this-> x; }; + * }; + * Instead, it is better to do something like this: + * class MyType { + * int xSquared() const { + * return *_xSquared.get(_x); + * } + * + * int _x; + * Deferred<int, int> _xSquared{[](int x) { return x * x; }; + * }; + * - As a similar danger, the value is only computed once. if you initialize it with arguments like + * the above 'xSquared()' implementation, then be cogniscent that the value will never change. If + * '_x' changes, '_xSquared' will not. + * + * A Deferred class can be constructed with either an initial value (eager initialization) or a + * function that will generate the value when needed. + */ +template <typename T, typename... Args> +class Deferred { +public: + /** + * Instantiates a Deffered<T> with the given data - no callbacks or lazy initialization. + */ + Deferred(T data) : _data(data) {} + + /** + * Stores a function to compute a T later. Please note the warnings described in this class + * comment. + */ + Deferred(std::function<T(Args&&...)> initializer) : _initializer(std::move(initializer)) {} + + /** + * Returns a pointer to the managed object. Initializes the object if it hasn't done so already. + */ + T& get(Args&&... args) const { + if (_initializer) { + _data = _initializer(std::forward<Args>(args)...); + _initializer = nullptr; + } + return _data; + } + + /** + * Dereferences the pointer to the managed object. Note this is only a valid shortcut if there + * are no arguments to '_initializer'. + */ + T* operator->() const { + return &get(); + } + + /** + * Returns a referenced to the managed object. Initializes the object if it hasn't done so + * already. Note this is only a valid shortcut if there are no arguments to '_initializer'. + */ + const T& operator*() const { + return get(); + } + + bool isInitialized() const { + return _initializer ? false : true; + } + +private: + mutable T _data; + mutable std::function<T(Args&&...)> _initializer; +}; + +} // namespace mongo diff --git a/src/mongo/db/query/util/deferred_test.cpp b/src/mongo/db/query/util/deferred_test.cpp new file mode 100644 index 00000000000..de256394787 --- /dev/null +++ b/src/mongo/db/query/util/deferred_test.cpp @@ -0,0 +1,98 @@ +/** + * Copyright (C) 2023-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. + */ + +#include "mongo/db/query/util/deferred.h" + +#include "mongo/unittest/unittest.h" + +namespace mongo { +using std::string; +using namespace std::string_literals; + + +TEST(DeferredTest, EagerInitialization) { + Deferred<string> eager{"someString"}; + ASSERT_TRUE(eager.isInitialized()); + ASSERT_EQ(eager.get(), "someString"s); + ASSERT_EQ(*eager, "someString"s); +} + +TEST(DeferredTest, DeferredInitialization) { + size_t initializationCount = 0; + Deferred<string> deferred{[&]() { + initializationCount++; + return "someString"s; + }}; + ASSERT_FALSE(deferred.isInitialized()); + + // Ensure the deferred object wasn't initialized on creation. + ASSERT_EQ(initializationCount, 0); + + // Ensure that the deferred object is initialized on pointer dereferences. + ASSERT_FALSE(deferred->empty()); + ASSERT_TRUE(deferred.isInitialized()); + + ASSERT_EQ(initializationCount, 1); + + // Ensure that the content of the deferred object is equal to its raw counterpart, while also + // verifing that it is initialized at most once. + ASSERT_EQ(deferred.get(), "someString"s); + ASSERT_EQ(initializationCount, 1); +} + +TEST(DeferredTest, DeferredInitializationWithOneArgument) { + size_t initializationCount = 0; + Deferred<string, const string&> deferred{[&](const string& input) { + initializationCount++; + return "{" + input + "}"; + }}; + + // Ensure the deferred object wasn't initialized on creation. + ASSERT_EQ(initializationCount, 0); + + // Ensure that the content of the deferred object is equal to its raw counterpart, while also + // verifing that it is initialized at most once. + ASSERT_EQ(deferred.get("more curlies"), "{more curlies}"s); + ASSERT_EQ(initializationCount, 1); + + // Note that the value is cached, so it's not really valid to call it with a different argument. + ASSERT_EQ(deferred.get("merganser"), "{more curlies}"s); + ASSERT_EQ(initializationCount, 1); +} + +TEST(DeferredTest, DeferredInitializationWithTwoArgs) { + Deferred<string, const string&, const string&> deferred{ + [&](const auto& input, const auto& prefix) { return prefix + input; }}; + + ASSERT_EQ(deferred.get("cowbell", "more "), "more cowbell"s); + ASSERT_EQ(deferred.get("cowbell", "more "), "more cowbell"s); + ASSERT_EQ(deferred.get("cowbell", "less?"), "more cowbell"s); + ASSERT_EQ(deferred.get("tests", "better"), "more cowbell"s); +} +} // namespace mongo diff --git a/src/mongo/db/query/util/memory_util.cpp b/src/mongo/db/query/util/memory_util.cpp new file mode 100644 index 00000000000..8a206deb2b7 --- /dev/null +++ b/src/mongo/db/query/util/memory_util.cpp @@ -0,0 +1,128 @@ +/** + * 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. + */ + +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery + +#include "mongo/db/query/util/memory_util.h" + +#include <cstddef> +#include <pcrecpp.h> + +#include "mongo/logv2/log.h" +#include "mongo/util/processinfo.h" + + +namespace mongo::memory_util { + +StatusWith<MemoryUnits> parseUnitString(const std::string& strUnit) { + if (strUnit.empty()) { + return Status(ErrorCodes::Error{6007010}, "Unit value cannot be empty"); + } + + if (strUnit[0] == '%') { + return MemoryUnits::kPercent; + } else if (strUnit[0] == 'M' || strUnit[0] == 'm') { + return MemoryUnits::kMB; + } else if (strUnit[0] == 'G' || strUnit[0] == 'g') { + return MemoryUnits::kGB; + } + + return Status(ErrorCodes::Error{6007011}, "Incorrect unit value"); +} + +StatusWith<MemorySize> MemorySize::parse(const std::string& str) { + pcrecpp::RE_Options opt; + opt.set_caseless(true); + // Looks for a floating point number with followed by a unit suffix (MB, GB, %). + pcrecpp::RE re("\\s*(\\d+\\.?\\d*)\\s*(MB|GB|%)\\s*", opt); + + double size{}; + std::string strUnit{}; + if (!re.FullMatch(str, &size, &strUnit)) { + return {ErrorCodes::Error{6007012}, "Unable to parse memory size string"}; + } + + auto statusWithUnit = parseUnitString(strUnit); + if (!statusWithUnit.isOK()) { + return statusWithUnit.getStatus(); + } + return MemorySize{size, statusWithUnit.getValue()}; +} + +size_t convertToSizeInBytes(const MemorySize& memSize) { + constexpr size_t kBytesInMB = 1024 * 1024; + constexpr size_t kMBytesInGB = 1024; + + double sizeInMB = memSize.size; + + switch (memSize.units) { + case MemoryUnits::kPercent: + sizeInMB *= ProcessInfo::getMemSizeMB() / 100.0; + break; + case MemoryUnits::kMB: + break; + case MemoryUnits::kGB: + sizeInMB *= kMBytesInGB; + break; + } + + return static_cast<size_t>(sizeInMB * kBytesInMB); +} + +size_t getRequestedMemSizeInBytes(const MemorySize& memSize) { + size_t planCacheSize = convertToSizeInBytes(memSize); + uassert(5968001, + "Cache size must be at least 1KB * number of cores", + planCacheSize >= 1024 * ProcessInfo::getNumCores()); + return planCacheSize; +} + +/** + * Sets upper limit on a storage structure's size. Either that structure's maximumSize or to + * percentage of the total system's memory (both known at call site), whichever is smaller. + */ +size_t capMemorySize(size_t requestedSizeBytes, + size_t maximumSizeGB, + double percentTotalSystemMemory) { + constexpr size_t kBytesInGB = 1024 * 1024 * 1024; + // Express maximum size in bytes. + const size_t maximumSizeBytes = maximumSizeGB * kBytesInGB; + const memory_util::MemorySize limitToProcessSize{percentTotalSystemMemory, + memory_util::MemoryUnits::kPercent}; + const size_t limitToProcessSizeInBytes = convertToSizeInBytes(limitToProcessSize); + + // The size will be capped by the minimum of the two values defined above. + const size_t upperLimit = std::min(maximumSizeBytes, limitToProcessSizeInBytes); + + if (requestedSizeBytes > upperLimit) { + requestedSizeBytes = upperLimit; + } + return requestedSizeBytes; +} +} // namespace mongo::memory_util diff --git a/src/mongo/db/query/plan_cache_size_parameter.h b/src/mongo/db/query/util/memory_util.h index 322a1fff564..345780b4c84 100644 --- a/src/mongo/db/query/plan_cache_size_parameter.h +++ b/src/mongo/db/query/util/memory_util.h @@ -31,29 +31,36 @@ #include <string> +#include "mongo/base/error_codes.h" #include "mongo/base/status_with.h" -namespace mongo::plan_cache_util { +namespace mongo::memory_util { /** - * Defines units of planCacheSize parameter. + * Defines units of memory. */ -enum class PlanCacheSizeUnits { +enum class MemoryUnits { kPercent, kMB, kGB, }; -StatusWith<PlanCacheSizeUnits> parseUnitString(const std::string& strUnit); - /** - * Represents parsed planCacheSize parameter. + * Represents parsed memory size parameter. */ -struct PlanCacheSizeParameter { - static StatusWith<PlanCacheSizeParameter> parse(const std::string& str); +struct MemorySize { + static StatusWith<MemorySize> parse(const std::string& str); const double size; - const PlanCacheSizeUnits units; + const MemoryUnits units; }; -} // namespace mongo::plan_cache_util +StatusWith<MemoryUnits> parseUnitString(const std::string& strUnit); +size_t convertToSizeInBytes(const MemorySize& memSize); +size_t capMemorySize(size_t requestedSizeBytes, + size_t maximumSizeGB, + double percentTotalSystemMemory); +size_t getRequestedMemSizeInBytes(const MemorySize& memSize); + + +} // namespace mongo::memory_util diff --git a/src/mongo/db/query/util/memory_util_test.cpp b/src/mongo/db/query/util/memory_util_test.cpp new file mode 100644 index 00000000000..78f7b3098d6 --- /dev/null +++ b/src/mongo/db/query/util/memory_util_test.cpp @@ -0,0 +1,73 @@ +/** + * 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. + */ + +#include "mongo/db/query/util/memory_util.h" + +#include "mongo/unittest/unittest.h" + +namespace mongo::memory_util { + +bool operator==(const MemorySize& lhs, const MemorySize& rhs) { + constexpr double kEpsilon = 1e-10; + return std::abs(lhs.size - rhs.size) < kEpsilon && lhs.units == rhs.units; +} + +TEST(MemorySizeTest, ParseUnitStringPercent) { + ASSERT_TRUE(MemoryUnits::kPercent == parseUnitString("%")); +} + +TEST(MemorySizeTest, ParseUnitStringMB) { + ASSERT_TRUE(MemoryUnits::kMB == parseUnitString("MB")); + ASSERT_TRUE(MemoryUnits::kMB == parseUnitString("mb")); + ASSERT_TRUE(MemoryUnits::kMB == parseUnitString("mB")); + ASSERT_TRUE(MemoryUnits::kMB == parseUnitString("Mb")); +} + +TEST(MemorySizeTest, ParseUnitStringGB) { + ASSERT_TRUE(MemoryUnits::kGB == parseUnitString("GB")); + ASSERT_TRUE(MemoryUnits::kGB == parseUnitString("gb")); + ASSERT_TRUE(MemoryUnits::kGB == parseUnitString("gB")); + ASSERT_TRUE(MemoryUnits::kGB == parseUnitString("Gb")); +} + +TEST(MemorySizeTest, ParseUnitStringIncorrectValue) { + ASSERT_NOT_OK(parseUnitString("").getStatus()); + ASSERT_NOT_OK(parseUnitString(" ").getStatus()); + ASSERT_NOT_OK(parseUnitString("KB").getStatus()); +} + +TEST(MemorySizeTest, ParseMemorySize) { + ASSERT_TRUE((MemorySize{10.0, MemoryUnits::kPercent}) == MemorySize::parse("10%")); + ASSERT_TRUE((MemorySize{300.0, MemoryUnits::kMB}) == MemorySize::parse("300MB")); + ASSERT_TRUE((MemorySize{4.0, MemoryUnits::kGB}) == MemorySize::parse("4GB")); + ASSERT_TRUE((MemorySize{5.1, MemoryUnits::kPercent}) == MemorySize::parse(" 5.1%")); + ASSERT_TRUE((MemorySize{11.1, MemoryUnits::kMB}) == MemorySize::parse("11.1 mb")); + ASSERT_TRUE((MemorySize{12.1, MemoryUnits::kGB}) == MemorySize::parse(" 12.1 Gb ")); +} +} // namespace mongo::memory_util diff --git a/src/mongo/db/query/view_response_formatter.cpp b/src/mongo/db/query/view_response_formatter.cpp index 76c8457c8b5..397f5e6769a 100644 --- a/src/mongo/db/query/view_response_formatter.cpp +++ b/src/mongo/db/query/view_response_formatter.cpp @@ -28,6 +28,8 @@ */ #include "mongo/platform/basic.h" +#include "mongo/util/assert_util.h" +#include "mongo/util/str.h" #include "mongo/db/query/view_response_formatter.h" @@ -57,7 +59,11 @@ Status ViewResponseFormatter::appendAsCountResponse(BSONObjBuilder* resultBuilde } else { invariant(cursorFirstBatch.size() == 1); auto countObj = cursorFirstBatch.back(); - resultBuilder->append(kCountField, countObj["count"].Int()); + auto countElem = countObj["count"]; + tassert(9384400, + str::stream() << "the 'count' should be of number type, but found " << countElem, + countElem.isNumber()); + resultBuilder->appendAs(countElem, kCountField); } resultBuilder->append(kOkField, 1); return Status::OK(); diff --git a/src/mongo/db/query/view_response_formatter_test.cpp b/src/mongo/db/query/view_response_formatter_test.cpp index dc86c5c9fc7..c0d710d838e 100644 --- a/src/mongo/db/query/view_response_formatter_test.cpp +++ b/src/mongo/db/query/view_response_formatter_test.cpp @@ -51,6 +51,16 @@ TEST(ViewResponseFormatter, FormatInitialCountResponseSuccessfully) { ASSERT_BSONOBJ_EQ(fromjson("{'n': 7, ok: 1}"), builder.obj()); } +TEST(ViewResponseFormatter, FormatInitialCountResponseWithNumberLong) { + CursorResponse cr( + testNss, testCursor, {BSON("count" << std::numeric_limits<long long>::max())}); + ViewResponseFormatter formatter(cr.toBSON(CursorResponse::ResponseType::InitialResponse)); + BSONObjBuilder builder; + ASSERT_OK(formatter.appendAsCountResponse(&builder)); + ASSERT_BSONOBJ_EQ(BSON("n" << std::numeric_limits<long long>::max() << "ok" << 1), + builder.obj()); +} + TEST(ViewResponseFormatter, FormatSubsequentCountResponseSuccessfully) { CursorResponse cr(testNss, testCursor, {BSON("count" << 7)}); ViewResponseFormatter formatter(cr.toBSON(CursorResponse::ResponseType::SubsequentResponse)); @@ -59,6 +69,16 @@ TEST(ViewResponseFormatter, FormatSubsequentCountResponseSuccessfully) { ASSERT_BSONOBJ_EQ(fromjson("{'n': 7, ok: 1}"), builder.obj()); } +TEST(ViewResponseFormatter, FormatSubsequentCountResponseWithLong) { + CursorResponse cr( + testNss, testCursor, {BSON("count" << std::numeric_limits<long long>::max())}); + ViewResponseFormatter formatter(cr.toBSON(CursorResponse::ResponseType::SubsequentResponse)); + BSONObjBuilder builder; + ASSERT_OK(formatter.appendAsCountResponse(&builder)); + ASSERT_BSONOBJ_EQ(BSON("n" << std::numeric_limits<long long>::max() << "ok" << 1), + builder.obj()); +} + TEST(ViewResponseFormatter, FormatEmptyInitialCountResponseSuccessfully) { CursorResponse cr(testNss, testCursor, {}); ViewResponseFormatter formatter(cr.toBSON(CursorResponse::ResponseType::InitialResponse)); |