diff options
Diffstat (limited to 'src/mongo/db/query/query_stats')
28 files changed, 5474 insertions, 0 deletions
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/query_stats/query_stats_entry.cpp b/src/mongo/db/query/query_stats/query_stats_entry.cpp new file mode 100644 index 00000000000..f69f0a6ee2a --- /dev/null +++ b/src/mongo/db/query/query_stats/query_stats_entry.cpp @@ -0,0 +1,54 @@ +/** + * 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/query_stats_entry.h" + +#include <boost/optional.hpp> + +#include "mongo/crypto/hash_block.h" +#include "mongo/crypto/sha256_block.h" + +#define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kQuery + +namespace mongo::query_stats { + +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(); +} + + +} // 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" |