SERVER-57037 Improve precision of operator counters

14 files changed, 104 insertions, 33 deletions

diff --git a/src/mongo/db/catalog/collection_impl.cpp b/src/mongo/db/catalog/collection_impl.cpp
index a15a507b98c..a70e9b6ec84 100644
--- a/src/mongo/db/catalog/collection_impl.cpp
+++ b/src/mongo/db/catalog/collection_impl.cpp
@@ -647,8 +647,10 @@ Collection::Validator CollectionImpl::parseValidator(
             ValidationLevelEnum::moderate)
         allowedFeatures &= ~MatchExpressionParser::AllowedFeatures::kEncryptKeywords;
 
+    expCtx->startExpressionCounters();
     auto statusWithMatcher =
         MatchExpressionParser::parse(validator, expCtx, ExtensionsCallbackNoop(), allowedFeatures);
+    expCtx->stopExpressionCounters();
 
     if (!statusWithMatcher.isOK()) {
         return {
diff --git a/src/mongo/db/ops/parsed_delete.cpp b/src/mongo/db/ops/parsed_delete.cpp
index 5028ccde159..6075a550e2f 100644
--- a/src/mongo/db/ops/parsed_delete.cpp
+++ b/src/mongo/db/ops/parsed_delete.cpp
@@ -80,6 +80,7 @@ Status ParsedDelete::parseRequest() {
         return Status::OK();
     }
 
+    _expCtx->startExpressionCounters();
     return parseQueryToCQ();
 }
 
diff --git a/src/mongo/db/ops/parsed_update.cpp b/src/mongo/db/ops/parsed_update.cpp
index 2d4df260b1a..ac281c89101 100644
--- a/src/mongo/db/ops/parsed_update.cpp
+++ b/src/mongo/db/ops/parsed_update.cpp
@@ -39,7 +39,8 @@ namespace mongo {
 
 ParsedUpdate::ParsedUpdate(OperationContext* opCtx,
                            const UpdateRequest* request,
-                           const ExtensionsCallback& extensionsCallback)
+                           const ExtensionsCallback& extensionsCallback,
+                           bool forgoOpCounterIncrements)
     : _opCtx(opCtx),
       _request(request),
       _expCtx(make_intrusive<ExpressionContext>(
@@ -53,7 +54,11 @@ ParsedUpdate::ParsedUpdate(OperationContext* opCtx,
           request->explain())),
       _driver(_expCtx),
       _canonicalQuery(),
-      _extensionsCallback(extensionsCallback) {}
+      _extensionsCallback(extensionsCallback) {
+    if (forgoOpCounterIncrements) {
+        _expCtx->enabledCounters = false;
+    }
+}
 
 Status ParsedUpdate::parseRequest() {
     // It is invalid to request that the UpdateStage return the prior or newly-updated version
@@ -158,6 +163,7 @@ Status ParsedUpdate::parseQueryToCQ() {
         findCommand->setLet(*letParams);
     }
 
+    _expCtx->startExpressionCounters();
     auto statusWithCQ = CanonicalQuery::canonicalize(_opCtx,
                                                      std::move(findCommand),
                                                      static_cast<bool>(_request->explain()),
diff --git a/src/mongo/db/ops/parsed_update.h b/src/mongo/db/ops/parsed_update.h
index 76440516ecb..40e3a279712 100644
--- a/src/mongo/db/ops/parsed_update.h
+++ b/src/mongo/db/ops/parsed_update.h
@@ -68,7 +68,8 @@ public:
      */
     ParsedUpdate(OperationContext* opCtx,
                  const UpdateRequest* request,
-                 const ExtensionsCallback& extensionsCallback);
+                 const ExtensionsCallback& extensionsCallback,
+                 bool forgoOpCounterIncrements = false);
 
     /**
      * Parses the update request to a canonical query and an update driver. On success, the
diff --git a/src/mongo/db/ops/write_ops_exec.cpp b/src/mongo/db/ops/write_ops_exec.cpp
index eaa1a6aea7e..f43dcf8581d 100644
--- a/src/mongo/db/ops/write_ops_exec.cpp
+++ b/src/mongo/db/ops/write_ops_exec.cpp
@@ -721,7 +721,8 @@ static SingleWriteResult performSingleUpdateOp(OperationContext* opCtx,
                                                const NamespaceString& ns,
                                                UpdateRequest* updateRequest,
                                                OperationSource source,
-                                               bool* containsDotsAndDollarsField) {
+                                               bool* containsDotsAndDollarsField,
+                                               bool forgoOpCounterIncrements) {
     CurOpFailpointHelpers::waitWhileFailPointEnabled(
         &hangDuringBatchUpdate,
         opCtx,
@@ -810,7 +811,7 @@ static SingleWriteResult performSingleUpdateOp(OperationContext* opCtx,
     }
 
     const ExtensionsCallbackReal extensionsCallback(opCtx, &updateRequest->getNamespaceString());
-    ParsedUpdate parsedUpdate(opCtx, updateRequest, extensionsCallback);
+    ParsedUpdate parsedUpdate(opCtx, updateRequest, extensionsCallback, forgoOpCounterIncrements);
     uassertStatusOK(parsedUpdate.parseRequest());
 
     CurOpFailpointHelpers::waitWhileFailPointEnabled(
@@ -883,7 +884,8 @@ static SingleWriteResult performSingleUpdateOpWithDupKeyRetry(
     const write_ops::UpdateOpEntry& op,
     LegacyRuntimeConstants runtimeConstants,
     const boost::optional<BSONObj>& letParams,
-    OperationSource source) {
+    OperationSource source,
+    bool forgoOpCounterIncrements) {
     globalOpCounters.gotUpdate();
     ServerWriteConcernMetrics::get(opCtx)->recordWriteConcernForUpdate(opCtx->getWriteConcern());
     auto& curOp = *CurOp::get(opCtx);
@@ -920,8 +922,12 @@ static SingleWriteResult performSingleUpdateOpWithDupKeyRetry(
 
         try {
             bool containsDotsAndDollarsField = false;
-            const auto ret =
-                performSingleUpdateOp(opCtx, ns, &request, source, &containsDotsAndDollarsField);
+            const auto ret = performSingleUpdateOp(opCtx,
+                                                   ns,
+                                                   &request,
+                                                   source,
+                                                   &containsDotsAndDollarsField,
+                                                   forgoOpCounterIncrements);
 
             if (containsDotsAndDollarsField) {
                 // If it's an upsert, increment 'inserts' metric, otherwise increment 'updates'.
@@ -986,6 +992,9 @@ WriteResult performUpdates(OperationContext* opCtx,
     const auto& runtimeConstants =
         wholeOp.getLegacyRuntimeConstants().value_or(Variables::generateRuntimeConstants(opCtx));
 
+    // Increment operator counters only during the fisrt single update operation in a batch of
+    // updates.
+    bool forgoOpCounterIncrements = false;
     for (auto&& singleOp : wholeOp.getUpdates()) {
         const auto stmtId = getStmtIdForWriteOp(opCtx, wholeOp, stmtIdIndex++);
         if (opCtx->getTxnNumber()) {
@@ -1024,8 +1033,16 @@ WriteResult performUpdates(OperationContext* opCtx,
                 ? *wholeOp.getStmtIds()
                 : std::vector<StmtId>{stmtId};
 
-            out.results.emplace_back(performSingleUpdateOpWithDupKeyRetry(
-                opCtx, ns, stmtIds, singleOp, runtimeConstants, wholeOp.getLet(), source));
+            out.results.emplace_back(
+                performSingleUpdateOpWithDupKeyRetry(opCtx,
+                                                     ns,
+                                                     stmtIds,
+                                                     singleOp,
+                                                     runtimeConstants,
+                                                     wholeOp.getLet(),
+                                                     source,
+                                                     forgoOpCounterIncrements));
+            forgoOpCounterIncrements = true;
             lastOpFixer.finishedOpSuccessfully();
         } catch (const DBException& ex) {
             out.canContinue = handleError(
diff --git a/src/mongo/db/pipeline/document_source_lookup.cpp b/src/mongo/db/pipeline/document_source_lookup.cpp
index be4dd4bdee3..eefa1e13626 100644
--- a/src/mongo/db/pipeline/document_source_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup.cpp
@@ -977,8 +977,10 @@ void DocumentSourceLookUp::resolveLetVariables(const Document& localDoc, Variabl
 
 void DocumentSourceLookUp::initializeResolvedIntrospectionPipeline() {
     _variables.copyToExpCtx(_variablesParseState, _fromExpCtx.get());
+    _fromExpCtx->startExpressionCounters();
     _resolvedIntrospectionPipeline =
         Pipeline::parse(_resolvedPipeline, _fromExpCtx, lookupPipeValidator);
+    _fromExpCtx->stopExpressionCounters();
 }
 
 void DocumentSourceLookUp::recordPlanSummaryStats(const Pipeline& pipeline) {
diff --git a/src/mongo/db/pipeline/expression.cpp b/src/mongo/db/pipeline/expression.cpp
index b04276b82b9..17c1d061a2b 100644
--- a/src/mongo/db/pipeline/expression.cpp
+++ b/src/mongo/db/pipeline/expression.cpp
@@ -181,7 +181,7 @@ void Expression::registerExpression(
     parserMap[key] =
         ParserRegistration{parser, allowedWithApiStrict, allowedWithClientType, requiredMinVersion};
     // Add this expression to the global map of operator counters for expressions.
-    operatorCountersExpressions.addExpressionCounter(key);
+    operatorCountersAggExpressions.addAggExpressionCounter(key);
 }
 
 intrusive_ptr<Expression> Expression::parseExpression(ExpressionContext* const expCtx,
@@ -219,8 +219,8 @@ intrusive_ptr<Expression> Expression::parseExpression(ExpressionContext* const e
             expCtx->opCtx, opName, entry.allowedWithApiStrict, entry.allowedWithClientType);
     }
 
-    // Increment the global counter for this expression.
-    operatorCountersExpressions.incrementExpressionCounter(opName);
+    // Increment the counter for this expression in the current context.
+    expCtx->incrementAggExprCounter(opName);
     return entry.parser(expCtx, obj.firstElement(), vps);
 }
 
diff --git a/src/mongo/db/pipeline/expression_context.cpp b/src/mongo/db/pipeline/expression_context.cpp
index 89cc44bf094..cb92a6e9ef9 100644
--- a/src/mongo/db/pipeline/expression_context.cpp
+++ b/src/mongo/db/pipeline/expression_context.cpp
@@ -226,21 +226,28 @@ intrusive_ptr<ExpressionContext> ExpressionContext::copyWith(
 }
 
 void ExpressionContext::startExpressionCounters() {
-    if (!_expressionCounters) {
+    if (enabledCounters && !_expressionCounters) {
         _expressionCounters = boost::make_optional<ExpressionCounters>({});
     }
 }
 
 void ExpressionContext::incrementMatchExprCounter(StringData name) {
-    if (_expressionCounters) {
+    if (enabledCounters && _expressionCounters) {
         ++_expressionCounters.get().matchExprCountersMap[name];
     }
 }
 
+void ExpressionContext::incrementAggExprCounter(StringData name) {
+    if (enabledCounters && _expressionCounters) {
+        ++_expressionCounters.get().aggExprCountersMap[name];
+    }
+}
+
 void ExpressionContext::stopExpressionCounters() {
-    if (_expressionCounters) {
+    if (enabledCounters && _expressionCounters) {
         operatorCountersMatchExpressions.mergeCounters(
             _expressionCounters.get().matchExprCountersMap);
+        operatorCountersAggExpressions.mergeCounters(_expressionCounters.get().aggExprCountersMap);
     }
     _expressionCounters = boost::none;
 }
diff --git a/src/mongo/db/pipeline/expression_context.h b/src/mongo/db/pipeline/expression_context.h
index 88340b8b846..be915aa98f1 100644
--- a/src/mongo/db/pipeline/expression_context.h
+++ b/src/mongo/db/pipeline/expression_context.h
@@ -108,6 +108,15 @@ public:
     };
 
     /**
+     * The structure ExpressionCounters encapsulates counters for match, aggregate, and other
+     * expression types as seen in end-user queries.
+     */
+    struct ExpressionCounters {
+        StringMap<uint64_t> aggExprCountersMap;
+        StringMap<uint64_t> matchExprCountersMap;
+    };
+
+    /**
      * Constructs an ExpressionContext to be used for Pipeline parsing and evaluation.
      * 'resolvedNamespaces' maps collection names (not full namespaces) to ResolvedNamespaces.
      */
@@ -344,11 +353,20 @@ public:
     void incrementMatchExprCounter(StringData name);
 
     /**
+     * Increment the counter for the aggregate expression with a given name.
+     */
+    void incrementAggExprCounter(StringData name);
+
+    /**
      * Merge expression counters from the current expression context into the global maps
      * and stop counting.
      */
     void stopExpressionCounters();
 
+    bool expressionCountersAreActive() {
+        return _expressionCounters.is_initialized();
+    }
+
     // The explain verbosity requested by the user, or boost::none if no explain was requested.
     boost::optional<ExplainOptions::Verbosity> explain;
 
@@ -429,6 +447,11 @@ public:
     // When non-empty, contains the unmodified user provided aggregation command.
     BSONObj originalAggregateCommand;
 
+    // True if the expression context is the original one for a given pipeline.
+    // False if another context is created for the same pipeline. Used to disable duplicate
+    // expression counting.
+    bool enabledCounters = true;
+
 protected:
     static const int kInterruptCheckPeriod = 128;
 
diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp
index fda5cdd9300..d6531198ff2 100644
--- a/src/mongo/db/query/canonical_query.cpp
+++ b/src/mongo/db/query/canonical_query.cpp
@@ -127,7 +127,7 @@ StatusWith<std::unique_ptr<CanonicalQuery>> CanonicalQuery::canonicalize(
         return statusWithMatcher.getStatus();
     }
 
-    // Stop counting match expressions after they have been parsed to exclude expressions created
+    // Stop counting expressions after they have been parsed to exclude expressions created
     // during optimization and other processing steps.
     newExpCtx->stopExpressionCounters();
 
diff --git a/src/mongo/db/query/projection_ast.h b/src/mongo/db/query/projection_ast.h
index ca5d3145da9..247a91537f3 100644
--- a/src/mongo/db/query/projection_ast.h
+++ b/src/mongo/db/query/projection_ast.h
@@ -278,10 +278,17 @@ public:
         bob << "" << other._expr->serialize(false);
 
         // TODO SERVER-31003: add a clone() method to Expression.
-        boost::intrusive_ptr<Expression> clonedExpr =
-            Expression::parseOperand(other._expr->getExpressionContext(),
-                                     bob.obj().firstElement(),
-                                     other._expr->getExpressionContext()->variablesParseState);
+        // Temporary stop expression counters while processing the cloned expression.
+        auto otherCtx = other._expr->getExpressionContext();
+        auto activeCounting = otherCtx->expressionCountersAreActive();
+        if (activeCounting) {
+            otherCtx->enabledCounters = false;
+        }
+        boost::intrusive_ptr<Expression> clonedExpr = Expression::parseOperand(
+            otherCtx, bob.obj().firstElement(), otherCtx->variablesParseState);
+        if (activeCounting) {
+            otherCtx->enabledCounters = true;
+        }
         _expr = clonedExpr;
     }
 
diff --git a/src/mongo/db/stats/counters.cpp b/src/mongo/db/stats/counters.cpp
index c1d7b6f7ec8..6d57b3d8b6a 100644
--- a/src/mongo/db/stats/counters.cpp
+++ b/src/mongo/db/stats/counters.cpp
@@ -313,6 +313,6 @@ NetworkCounter networkCounter;
 AuthCounter authCounter;
 AggStageCounters aggStageCounters;
 DotsAndDollarsFieldsCounters dotsAndDollarsFieldsCounters;
-OperatorCountersExpressions operatorCountersExpressions;
+OperatorCountersAggExpressions operatorCountersAggExpressions;
 OperatorCountersMatchExpressions operatorCountersMatchExpressions;
 }  // namespace mongo
diff --git a/src/mongo/db/stats/counters.h b/src/mongo/db/stats/counters.h
index bec806d7005..acde64ba160 100644
--- a/src/mongo/db/stats/counters.h
+++ b/src/mongo/db/stats/counters.h
@@ -336,33 +336,36 @@ public:
 
 extern DotsAndDollarsFieldsCounters dotsAndDollarsFieldsCounters;
 
-class OperatorCountersExpressions {
+class OperatorCountersAggExpressions {
 private:
-    struct ExprCounter {
-        ExprCounter(StringData name) : metric("operatorCounters.expressions." + name, &counter) {}
+    struct AggExprCounter {
+        AggExprCounter(StringData name)
+            : metric("operatorCounters.expressions." + name, &counter) {}
 
         Counter64 counter;
         ServerStatusMetricField<Counter64> metric;
     };
 
 public:
-    void addExpressionCounter(StringData name) {
-        operatorCountersExpressionMap[name] = std::make_unique<ExprCounter>(name);
+    void addAggExpressionCounter(StringData name) {
+        operatorCountersAggExpressionMap[name] = std::make_unique<AggExprCounter>(name);
     }
 
-    void incrementExpressionCounter(StringData name) {
-        if (auto it = operatorCountersExpressionMap.find(name);
-            it != operatorCountersExpressionMap.end()) {
-            it->second->counter.increment(1);
+    void mergeCounters(StringMap<uint64_t>& toMerge) {
+        for (auto&& [name, cnt] : toMerge) {
+            if (auto it = operatorCountersAggExpressionMap.find(name);
+                it != operatorCountersAggExpressionMap.end()) {
+                it->second->counter.increment(cnt);
+            }
         }
     }
 
 private:
     // Map of aggregation expressions to the number of occurrences in aggregation pipelines.
-    StringMap<std::unique_ptr<ExprCounter>> operatorCountersExpressionMap = {};
+    StringMap<std::unique_ptr<AggExprCounter>> operatorCountersAggExpressionMap = {};
 };
 
-extern OperatorCountersExpressions operatorCountersExpressions;
+extern OperatorCountersAggExpressions operatorCountersAggExpressions;
 
 /**
  * Global counters for match expressions.
diff --git a/src/mongo/db/update/pipeline_executor.cpp b/src/mongo/db/update/pipeline_executor.cpp
index ac6df9ca086..710803d6b04 100644
--- a/src/mongo/db/update/pipeline_executor.cpp
+++ b/src/mongo/db/update/pipeline_executor.cpp
@@ -72,7 +72,9 @@ PipelineExecutor::PipelineExecutor(const boost::intrusive_ptr<ExpressionContext>
     }
 
     _expCtx->setResolvedNamespaces(resolvedNamespaces);
+    _expCtx->startExpressionCounters();
     _pipeline = Pipeline::parse(pipeline, _expCtx);
+    _expCtx->stopExpressionCounters();
 
     // Validate the update pipeline.
     for (auto&& stage : _pipeline->getSources()) {