New upstream version 6.0.20upstream

1 files changed, 117 insertions, 40 deletions

diff --git a/src/mongo/s/query/cluster_aggregate.cpp b/src/mongo/s/query/cluster_aggregate.cpp
index 3d6e9b5c2af..6374bcfd494 100644
--- a/src/mongo/s/query/cluster_aggregate.cpp
+++ b/src/mongo/s/query/cluster_aggregate.cpp
@@ -27,6 +27,7 @@
  *    it in the license file.
  */
 
+#include "mongo/s/chunk_manager.h"
 #define MONGO_LOGV2_DEFAULT_COMPONENT ::mongo::logv2::LogComponent::kCommand
 
 #include "mongo/platform/basic.h"
@@ -56,6 +57,9 @@
 #include "mongo/db/query/explain_common.h"
 #include "mongo/db/query/find_common.h"
 #include "mongo/db/query/fle/server_rewrite.h"
+#include "mongo/db/query/query_stats/agg_key.h"
+#include "mongo/db/query/query_stats/key.h"
+#include "mongo/db/query/query_stats/query_stats.h"
 #include "mongo/db/timeseries/timeseries_options.h"
 #include "mongo/db/views/resolved_view.h"
 #include "mongo/db/views/view.h"
@@ -95,7 +99,7 @@ namespace {
 // definition. It's okay that this is incorrect, we will repopulate the real namespace map on the
 // mongod. Note that this function must be called before forwarding an aggregation command on an
 // unsharded collection, in order to verify that the involved namespaces are allowed to be sharded.
-auto resolveInvolvedNamespaces(stdx::unordered_set<NamespaceString> involvedNamespaces) {
+auto resolveInvolvedNamespaces(const stdx::unordered_set<NamespaceString>& involvedNamespaces) {
     StringMap<ExpressionContext::ResolvedNamespace> resolvedNamespaces;
     for (auto&& nss : involvedNamespaces) {
         resolvedNamespaces.try_emplace(nss.coll(), nss, std::vector<BSONObj>{});
@@ -258,6 +262,68 @@ std::vector<BSONObj> rebuildPipelineWithTimeSeriesGranularity(const std::vector<
     return newPipeline;
 }
 
+/**
+ * Builds an expCtx with which to parse the request's pipeline, then parses the pipeline and
+ * registers the pre-optimized pipeline with query stats collection.
+ */
+std::unique_ptr<Pipeline, PipelineDeleter> parsePipelineAndRegisterQueryStats(
+    OperationContext* opCtx,
+    const stdx::unordered_set<NamespaceString>& involvedNamespaces,
+    const NamespaceString& executionNss,
+    AggregateCommandRequest& request,
+    const boost::optional<ChunkManager>& cm,
+    const LiteParsedPipeline& liteParsedPipeline,
+    bool hasChangeStream,
+    bool shouldDoFLERewrite) {
+    // Populate the collection UUID and the appropriate collation to use.
+    auto [collationObj, uuid] = [&]() -> std::pair<BSONObj, boost::optional<UUID>> {
+        // If this is a change stream, take the user-defined collation if one exists, or an
+        // empty BSONObj otherwise. Change streams never inherit the collection's default
+        // collation, and since collectionless aggregations generally run on the 'admin'
+        // database, the standard logic would attempt to resolve its non-existent UUID and
+        // collation by sending a specious 'listCollections' command to the config servers.
+        if (hasChangeStream) {
+            return {request.getCollation().value_or(BSONObj()), boost::none};
+        }
+
+        return cluster_aggregation_planner::getCollationAndUUID(
+            opCtx, cm, executionNss, request.getCollation().value_or(BSONObj()));
+    }();
+
+    // Build an ExpressionContext for the pipeline. This instantiates an appropriate collator,
+    // resolves all involved namespaces, and creates a shared MongoProcessInterface for use by the
+    // pipeline's stages.
+    boost::intrusive_ptr<ExpressionContext> expCtx =
+        makeExpressionContext(opCtx,
+                              request,
+                              collationObj,
+                              uuid,
+                              resolveInvolvedNamespaces(involvedNamespaces),
+                              hasChangeStream);
+
+    // A pipeline with $changeStreamSplitLargeEvent requires the use of resume token format v2,
+    // since the 'fragmentNum' field only exists in this version and later.
+    if (hasChangeStream && liteParsedPipeline.endsWithChangeStreamSplitLargeEvent()) {
+        expCtx->changeStreamTokenVersion = 2;
+    }
+
+    // Parse and optimize the full pipeline.
+    auto pipeline = Pipeline::parse(request.getPipeline(), expCtx);
+
+    // Skip query stats recording for queryable encryption queries.
+    if (!shouldDoFLERewrite) {
+        query_stats::registerRequest(
+            opCtx,
+            executionNss,
+            [&]() {
+                return std::make_unique<query_stats::AggKey>(
+                    request, *pipeline, expCtx, involvedNamespaces, executionNss);
+            },
+            hasChangeStream);
+    }
+    return pipeline;
+}
+
 }  // namespace
 
 Status ClusterAggregate::runAggregate(OperationContext* opCtx,
@@ -351,39 +417,15 @@ Status ClusterAggregate::runAggregate(OperationContext* opCtx,
 
     boost::intrusive_ptr<ExpressionContext> expCtx;
     const auto pipelineBuilder = [&]() {
-        // Populate the collection UUID and the appropriate collation to use.
-        auto [collationObj, uuid] = [&]() -> std::pair<BSONObj, boost::optional<UUID>> {
-            // If this is a change stream, take the user-defined collation if one exists, or an
-            // empty BSONObj otherwise. Change streams never inherit the collection's default
-            // collation, and since collectionless aggregations generally run on the 'admin'
-            // database, the standard logic would attempt to resolve its non-existent UUID and
-            // collation by sending a specious 'listCollections' command to the config servers.
-            if (hasChangeStream) {
-                return {request.getCollation().value_or(BSONObj()), boost::none};
-            }
-
-            return cluster_aggregation_planner::getCollationAndUUID(
-                opCtx, cm, namespaces.executionNss, request.getCollation().value_or(BSONObj()));
-        }();
-
-        // Build an ExpressionContext for the pipeline. This instantiates an appropriate collator,
-        // resolves all involved namespaces, and creates a shared MongoProcessInterface for use by
-        // the pipeline's stages.
-        expCtx = makeExpressionContext(opCtx,
-                                       request,
-                                       collationObj,
-                                       uuid,
-                                       resolveInvolvedNamespaces(involvedNamespaces),
-                                       hasChangeStream);
-
-        // A pipeline with $changeStreamSplitLargeEvent requires the use of resume token format v2,
-        // since the 'fragmentNum' field only exists in this version and later.
-        if (hasChangeStream && liteParsedPipeline.endsWithChangeStreamSplitLargeEvent()) {
-            expCtx->changeStreamTokenVersion = 2;
-        }
-
-        // Parse and optimize the full pipeline.
-        auto pipeline = Pipeline::parse(request.getPipeline(), expCtx);
+        auto pipeline = parsePipelineAndRegisterQueryStats(opCtx,
+                                                           involvedNamespaces,
+                                                           namespaces.executionNss,
+                                                           request,
+                                                           cm,
+                                                           liteParsedPipeline,
+                                                           hasChangeStream,
+                                                           shouldDoFLERewrite);
+        expCtx = pipeline->getContext();
 
         // If the aggregate command supports encrypted collections, do rewrites of the pipeline to
         // support querying against encrypted fields.
@@ -429,15 +471,48 @@ Status ClusterAggregate::runAggregate(OperationContext* opCtx,
                 cluster_aggregation_planner::AggregationTargeter::TargetingPolicy::kMongosRequired);
 
     if (!expCtx) {
-        // When the AggregationTargeter chooses a "passthrough" policy, it does not call the
-        // 'pipelineBuilder' function, so we never get an expression context. Because this is a
-        // passthrough, we only need a bare minimum expression context anyway.
+        // When the AggregationTargeter chooses a "passthrough" or "specific shard only" policy, it
+        // does not call the 'pipelineBuilder' function, so we've yet to construct an expression
+        // context or register query stats. Because this is a passthrough, we only need a bare
+        // minimum expression context on mongos.
         invariant(targeter.policy ==
                       cluster_aggregation_planner::AggregationTargeter::kPassthrough ||
                   targeter.policy ==
                       cluster_aggregation_planner::AggregationTargeter::kSpecificShardOnly);
+
         expCtx = make_intrusive<ExpressionContext>(
             opCtx, nullptr, namespaces.executionNss, boost::none, request.getLet());
+        expCtx->addResolvedNamespaces(involvedNamespaces);
+
+
+        // We might need 'inMongos' temporarily set to true for query stats parsing, but we don't
+        // want to modify the value of 'expCtx' for future code execution so we will set it back to
+        // its original value.
+        ON_BLOCK_EXIT([&expCtx, originalInMongosVal = expCtx->inMongos]() {
+            expCtx->inMongos = originalInMongosVal;
+        });
+
+        // In order to parse a change stream request for query stats, 'inMongos' needs
+        // to be set to true.
+        if (hasChangeStream) {
+            expCtx->inMongos = true;
+        }
+
+        // Skip query stats recording for queryable encryption queries.
+        if (!shouldDoFLERewrite) {
+            // We want to hold off parsing the pipeline until it's clear we must. Because of that,
+            // we wait to parse the pipeline until this callback is invoked within
+            // query_stats::registerRequest.
+            query_stats::registerRequest(
+                opCtx,
+                namespaces.executionNss,
+                [&]() {
+                    auto pipeline = Pipeline::parse(request.getPipeline(), expCtx);
+                    return std::make_unique<query_stats::AggKey>(
+                        request, *pipeline, expCtx, involvedNamespaces, namespaces.executionNss);
+                },
+                hasChangeStream);
+        }
     }
 
     if (request.getExplain()) {
@@ -465,10 +540,11 @@ Status ClusterAggregate::runAggregate(OperationContext* opCtx,
                 // If this is an explain write the explain output and return.
                 auto expCtx = targeter.pipeline->getContext();
                 if (expCtx->explain) {
+                    auto opts = SerializationOptions{};
+                    opts.verbosity = boost::make_optional(ExplainOptions::Verbosity::kQueryPlanner);
                     *result << "splitPipeline" << BSONNULL << "mongos"
                             << Document{{"host", getHostNameCachedAndPort()},
-                                        {"stages",
-                                         targeter.pipeline->writeExplainOps(*expCtx->explain)}};
+                                        {"stages", targeter.pipeline->writeExplainOps(opts)}};
                     return Status::OK();
                 }
 
@@ -537,11 +613,12 @@ Status ClusterAggregate::runAggregate(OperationContext* opCtx,
         updateHostsTargetedMetrics(opCtx, namespaces.executionNss, cm, involvedNamespaces);
         // Report usage statistics for each stage in the pipeline.
         liteParsedPipeline.tickGlobalStageCounters();
-
         // Add 'command' object to explain output.
         if (expCtx->explain) {
             explain_common::appendIfRoom(
                 aggregation_request_helper::serializeToCommandObj(request), "command", result);
+            collectQueryStatsMongos(opCtx,
+                                    std::move(CurOp::get(opCtx)->debug().queryStatsInfo.key));
         }
     }
     return status;