New upstream version 6.0.20upstream

1 files changed, 467 insertions, 0 deletions

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