ESQL: Implement a MetricsAware interface (elastic#121074)

* ESQL: Implement a MetricsAware interface (elastic#120527)

This implements an interface that export the names of the plan nodes and
functions that need to be counted in the metrics.

Also, the metrics are now counted from within the parser. This should
allow correct accounting for the cases where some nodes can appear both
standalone or part other nodes' children (like Aggregate being a child
of INLINESTATS, so no STATS counting should occur).

The functions counting now also validates that behind a name there is
actually a function registered.

Closes elastic#115992.

(cherry picked from commit a4482d4)

* Drop the HashSet gating when counting commands

The telemetry accounting is no longer done in just one place in the parser,
but split, so that no HashSet is required to discard duplicate accounting of
the same node. This lowers the memory requirements.

Author: bpintea

docs/changelog/121074.yaml

@@ -0,0 +1,5 @@
+pr: 121074
+summary: Implement a `MetricsAware` interface
+area: ES|QL
+type: enhancement
+issues: []

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -74,8 +74,8 @@
 import org.elasticsearch.xpack.esql.plugin.QueryPragmas;
 import org.elasticsearch.xpack.esql.session.Configuration;
 import org.elasticsearch.xpack.esql.session.QueryBuilderResolver;
-import org.elasticsearch.xpack.esql.stats.Metrics;
 import org.elasticsearch.xpack.esql.stats.SearchStats;
+import org.elasticsearch.xpack.esql.telemetry.Metrics;
 import org.elasticsearch.xpack.versionfield.Version;
 import org.junit.Assert;
 

x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/TelemetryIT.java

@@ -20,7 +20,7 @@
 import org.elasticsearch.plugins.PluginsService;
 import org.elasticsearch.telemetry.Measurement;
 import org.elasticsearch.telemetry.TestTelemetryPlugin;
-import org.elasticsearch.xpack.esql.stats.PlanningMetricsManager;
+import org.elasticsearch.xpack.esql.telemetry.PlanTelemetryManager;
 import org.junit.Before;
 
 import java.util.Collection;
@@ -113,6 +113,41 @@ public static Iterable<Object[]> parameters() {
                     Map.ofEntries(Map.entry("TO_IP", 1), Map.entry("TO_STRING", 2)),
                     true
                 ) },
+            new Object[] {
+                new Test(
+                    // Using the `::` cast operator and a function alias
+                    """
+                        ROW host = "1.1.1.1"
+                        | EVAL ip = host::ip::string, y = to_str(host)
+                        """,
+                    Map.ofEntries(Map.entry("ROW", 1), Map.entry("EVAL", 1)),
+                    Map.ofEntries(Map.entry("TO_IP", 1), Map.entry("TO_STRING", 2)),
+                    true
+                ) },
+            new Object[] {
+                new Test(
+                    // Using the `::` cast operator and a function alias
+                    """
+                        FROM idx
+                        | EVAL ip = host::ip::string, y = to_str(host)
+                        """,
+                    Map.ofEntries(Map.entry("FROM", 1), Map.entry("EVAL", 1)),
+                    Map.ofEntries(Map.entry("TO_IP", 1), Map.entry("TO_STRING", 2)),
+                    true
+                ) },
+            new Object[] {
+                new Test(
+                    """
+                        FROM idx
+                        | EVAL y = to_str(host)
+                        | LOOKUP JOIN lookup_idx ON host
+                        """,
+                    Build.current().isSnapshot()
+                        ? Map.ofEntries(Map.entry("FROM", 1), Map.entry("EVAL", 1), Map.entry("LOOKUP JOIN", 1))
+                        : Collections.emptyMap(),
+                    Build.current().isSnapshot() ? Map.ofEntries(Map.entry("TO_STRING", 1)) : Collections.emptyMap(),
+                    Build.current().isSnapshot()
+                ) },
             new Object[] {
                 new Test(
                     "METRICS idx | LIMIT 10",
@@ -123,9 +158,7 @@ public static Iterable<Object[]> parameters() {
             new Object[] {
                 new Test(
                     "METRICS idx max(id) BY host | LIMIT 10",
-                    Build.current().isSnapshot()
-                        ? Map.ofEntries(Map.entry("METRICS", 1), Map.entry("LIMIT", 1), Map.entry("FROM TS", 1))
-                        : Collections.emptyMap(),
+                    Build.current().isSnapshot() ? Map.ofEntries(Map.entry("METRICS", 1), Map.entry("LIMIT", 1)) : Collections.emptyMap(),
                     Build.current().isSnapshot() ? Map.ofEntries(Map.entry("MAX", 1)) : Collections.emptyMap(),
                     Build.current().isSnapshot()
                 ) }
@@ -138,7 +171,7 @@ public static Iterable<Object[]> parameters() {
             // | EVAL ip = to_ip(host), x = to_string(host), y = to_string(host)
             // | INLINESTATS max(id)
             // """,
-            // Build.current().isSnapshot() ? Map.of("FROM", 1, "EVAL", 1, "INLINESTATS", 1, "STATS", 1) : Collections.emptyMap(),
+            // Build.current().isSnapshot() ? Map.of("FROM", 1, "EVAL", 1, "INLINESTATS", 1) : Collections.emptyMap(),
             // Build.current().isSnapshot()
             // ? Map.ofEntries(Map.entry("MAX", 1), Map.entry("TO_IP", 1), Map.entry("TO_STRING", 2))
             // : Collections.emptyMap(),
@@ -186,19 +219,19 @@ private static void testQuery(
                 client(dataNode.getName()).execute(EsqlQueryAction.INSTANCE, request, ActionListener.running(() -> {
                     try {
                         // test total commands used
-                        final List<Measurement> commandMeasurementsAll = measurements(plugin, PlanningMetricsManager.FEATURE_METRICS_ALL);
+                        final List<Measurement> commandMeasurementsAll = measurements(plugin, PlanTelemetryManager.FEATURE_METRICS_ALL);
                         assertAllUsages(expectedCommands, commandMeasurementsAll, iteration, success);
 
                         // test num of queries using a command
-                        final List<Measurement> commandMeasurements = measurements(plugin, PlanningMetricsManager.FEATURE_METRICS);
+                        final List<Measurement> commandMeasurements = measurements(plugin, PlanTelemetryManager.FEATURE_METRICS);
                         assertUsageInQuery(expectedCommands, commandMeasurements, iteration, success);
 
                         // test total functions used
-                        final List<Measurement> functionMeasurementsAll = measurements(plugin, PlanningMetricsManager.FUNCTION_METRICS_ALL);
+                        final List<Measurement> functionMeasurementsAll = measurements(plugin, PlanTelemetryManager.FUNCTION_METRICS_ALL);
                         assertAllUsages(expectedFunctions, functionMeasurementsAll, iteration, success);
 
                         // test number of queries using a function
-                        final List<Measurement> functionMeasurements = measurements(plugin, PlanningMetricsManager.FUNCTION_METRICS);
+                        final List<Measurement> functionMeasurements = measurements(plugin, PlanTelemetryManager.FUNCTION_METRICS);
                         assertUsageInQuery(expectedFunctions, functionMeasurements, iteration, success);
                     } finally {
                         latch.countDown();
@@ -216,8 +249,8 @@ private static void assertAllUsages(Map<String, Integer> expected, List<Measurem
         Set<String> found = featureNames(metrics);
         assertThat(found, is(expected.keySet()));
         for (Measurement metric : metrics) {
-            assertThat(metric.attributes().get(PlanningMetricsManager.SUCCESS), is(success));
-            String featureName = (String) metric.attributes().get(PlanningMetricsManager.FEATURE_NAME);
+            assertThat(metric.attributes().get(PlanTelemetryManager.SUCCESS), is(success));
+            String featureName = (String) metric.attributes().get(PlanTelemetryManager.FEATURE_NAME);
             assertThat(metric.getLong(), is(iteration * expected.get(featureName)));
         }
     }
@@ -227,7 +260,7 @@ private static void assertUsageInQuery(Map<String, Integer> expected, List<Measu
         functionsFound = featureNames(found);
         assertThat(functionsFound, is(expected.keySet()));
         for (Measurement measurement : found) {
-            assertThat(measurement.attributes().get(PlanningMetricsManager.SUCCESS), is(success));
+            assertThat(measurement.attributes().get(PlanTelemetryManager.SUCCESS), is(success));
             assertThat(measurement.getLong(), is(iteration));
         }
     }
@@ -238,7 +271,7 @@ private static List<Measurement> measurements(TestTelemetryPlugin plugin, String
 
     private static Set<String> featureNames(List<Measurement> functionMeasurements) {
         return functionMeasurements.stream()
-            .map(x -> x.attributes().get(PlanningMetricsManager.FEATURE_NAME))
+            .map(x -> x.attributes().get(PlanTelemetryManager.FEATURE_NAME))
             .map(String.class::cast)
             .collect(Collectors.toSet());
     }
@@ -268,6 +301,19 @@ private static void loadData(String nodeName) {
         }
 
         client().admin().indices().prepareRefresh("idx").get();
+
+        assertAcked(
+            client().admin()
+                .indices()
+                .prepareCreate("lookup_idx")
+                .setSettings(
+                    Settings.builder()
+                        .put("index.routing.allocation.require._name", nodeName)
+                        .put("index.mode", "lookup")
+                        .put(IndexMetadata.SETTING_NUMBER_OF_SHARDS, 1)
+                )
+                .setMapping("ip", "type=ip", "host", "type=keyword")
+        );
     }
 
     private DiscoveryNode randomDataNode() {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Analyzer.java

@@ -92,7 +92,7 @@
 import org.elasticsearch.xpack.esql.rule.Rule;
 import org.elasticsearch.xpack.esql.rule.RuleExecutor;
 import org.elasticsearch.xpack.esql.session.Configuration;
-import org.elasticsearch.xpack.esql.stats.FeatureMetric;
+import org.elasticsearch.xpack.esql.telemetry.FeatureMetric;
 import org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter;
 
 import java.time.Duration;
@@ -133,7 +133,7 @@
 import static org.elasticsearch.xpack.esql.core.type.DataType.TIME_DURATION;
 import static org.elasticsearch.xpack.esql.core.type.DataType.VERSION;
 import static org.elasticsearch.xpack.esql.core.type.DataType.isTemporalAmount;
-import static org.elasticsearch.xpack.esql.stats.FeatureMetric.LIMIT;
+import static org.elasticsearch.xpack.esql.telemetry.FeatureMetric.LIMIT;
 import static org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter.maybeParseTemporalAmount;
 
 /**
@@ -220,7 +220,7 @@ private LogicalPlan resolveIndex(UnresolvedRelation plan, IndexResolution indexR
                         plan.metadataFields(),
                         plan.indexMode(),
                         indexResolutionMessage,
-                        plan.commandName()
+                        plan.telemetryLabel()
                     );
             }
             IndexPattern table = plan.indexPattern();
@@ -233,7 +233,7 @@ private LogicalPlan resolveIndex(UnresolvedRelation plan, IndexResolution indexR
                     plan.metadataFields(),
                     plan.indexMode(),
                     "invalid [" + table + "] resolution to [" + indexResolution + "]",
-                    plan.commandName()
+                    plan.telemetryLabel()
                 );
             }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/analysis/Verifier.java

@@ -32,8 +32,8 @@
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.logical.Lookup;
 import org.elasticsearch.xpack.esql.plan.logical.Project;
-import org.elasticsearch.xpack.esql.stats.FeatureMetric;
-import org.elasticsearch.xpack.esql.stats.Metrics;
+import org.elasticsearch.xpack.esql.telemetry.FeatureMetric;
+import org.elasticsearch.xpack.esql.telemetry.Metrics;
 
 import java.util.ArrayList;
 import java.util.BitSet;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/capabilities/TelemetryAware.java

@@ -0,0 +1,23 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.esql.capabilities;
+
+import java.util.Locale;
+
+/**
+ * Interface for plan nodes that need to be accounted in the statistics
+ */
+public interface TelemetryAware {
+
+    /**
+     * @return the label reported in the telemetry data. Only needs to be overwritten if the label doesn't match the class name.
+     */
+    default String telemetryLabel() {
+        return getClass().getSimpleName().toUpperCase(Locale.ROOT);
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/execution/PlanExecutor.java

@@ -26,10 +26,10 @@
 import org.elasticsearch.xpack.esql.session.IndexResolver;
 import org.elasticsearch.xpack.esql.session.QueryBuilderResolver;
 import org.elasticsearch.xpack.esql.session.Result;
-import org.elasticsearch.xpack.esql.stats.Metrics;
-import org.elasticsearch.xpack.esql.stats.PlanningMetrics;
-import org.elasticsearch.xpack.esql.stats.PlanningMetricsManager;
-import org.elasticsearch.xpack.esql.stats.QueryMetric;
+import org.elasticsearch.xpack.esql.telemetry.Metrics;
+import org.elasticsearch.xpack.esql.telemetry.PlanTelemetry;
+import org.elasticsearch.xpack.esql.telemetry.PlanTelemetryManager;
+import org.elasticsearch.xpack.esql.telemetry.QueryMetric;
 
 import static org.elasticsearch.action.ActionListener.wrap;
 
@@ -41,7 +41,7 @@ public class PlanExecutor {
     private final Mapper mapper;
     private final Metrics metrics;
     private final Verifier verifier;
-    private final PlanningMetricsManager planningMetricsManager;
+    private final PlanTelemetryManager planTelemetryManager;
 
     public PlanExecutor(IndexResolver indexResolver, MeterRegistry meterRegistry, XPackLicenseState licenseState) {
         this.indexResolver = indexResolver;
@@ -50,7 +50,7 @@ public PlanExecutor(IndexResolver indexResolver, MeterRegistry meterRegistry, XP
         this.mapper = new Mapper();
         this.metrics = new Metrics(functionRegistry);
         this.verifier = new Verifier(metrics, licenseState);
-        this.planningMetricsManager = new PlanningMetricsManager(meterRegistry);
+        this.planTelemetryManager = new PlanTelemetryManager(meterRegistry);
     }
 
     public void esql(
@@ -65,7 +65,7 @@ public void esql(
         QueryBuilderResolver queryBuilderResolver,
         ActionListener<Result> listener
     ) {
-        final PlanningMetrics planningMetrics = new PlanningMetrics();
+        final PlanTelemetry planTelemetry = new PlanTelemetry(functionRegistry);
         final var session = new EsqlSession(
             sessionId,
             cfg,
@@ -76,20 +76,20 @@ public void esql(
             new LogicalPlanOptimizer(new LogicalOptimizerContext(cfg, foldContext)),
             mapper,
             verifier,
-            planningMetrics,
+            planTelemetry,
             indicesExpressionGrouper,
             queryBuilderResolver
         );
         QueryMetric clientId = QueryMetric.fromString("rest");
         metrics.total(clientId);
 
         ActionListener<Result> executeListener = wrap(x -> {
-            planningMetricsManager.publish(planningMetrics, true);
+            planTelemetryManager.publish(planTelemetry, true);
             listener.onResponse(x);
         }, ex -> {
             // TODO when we decide if we will differentiate Kibana from REST, this String value will likely come from the request
             metrics.failed(clientId);
-            planningMetricsManager.publish(planningMetrics, false);
+            planTelemetryManager.publish(planTelemetry, false);
             listener.onFailure(ex);
         });
         // Wrap it in a listener so that if we have any exceptions during execution, the listener picks it up

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java

@@ -224,6 +224,7 @@ public class EsqlFunctionRegistry {
     // it has with the alias name associated to the FunctionDefinition instance
     private final Map<String, FunctionDefinition> defs = new LinkedHashMap<>();
     private final Map<String, String> aliases = new HashMap<>();
+    private final Map<Class<? extends Function>, String> names = new HashMap<>();
 
     private SnapshotFunctionRegistry snapshotRegistry = null;
 
@@ -258,6 +259,12 @@ public boolean functionExists(String functionName) {
         return defs.containsKey(functionName);
     }
 
+    public String functionName(Class<? extends Function> clazz) {
+        String name = names.get(clazz);
+        Check.notNull(name, "Cannot find function by class {}", clazz);
+        return name;
+    }
+
     public Collection<FunctionDefinition> listFunctions() {
         // It is worth double checking if we need this copy. These are immutable anyway.
         return defs.values();
@@ -758,6 +765,14 @@ void register(FunctionDefinition... functions) {
                 }
                 aliases.put(alias, f.name());
             }
+            Check.isTrue(
+                names.containsKey(f.clazz()) == false,
+                "function type [{}} is registered twice with names [{}] and [{}]",
+                f.clazz(),
+                names.get(f.clazz()),
+                f.name()
+            );
+            names.put(f.clazz(), f.name());
         }
         // sort the temporary map by key name and add it to the global map of functions
         defs.putAll(

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/AstBuilder.java

@@ -8,7 +8,7 @@
 package org.elasticsearch.xpack.esql.parser;
 
 public class AstBuilder extends LogicalPlanBuilder {
-    public AstBuilder(QueryParams params) {
-        super(params);
+    public AstBuilder(ParsingContext context) {
+        super(context);
     }
 }

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlParser.java

@@ -18,7 +18,9 @@
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 import org.elasticsearch.xpack.esql.core.util.StringUtils;
+import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
+import org.elasticsearch.xpack.esql.telemetry.PlanTelemetry;
 
 import java.util.BitSet;
 import java.util.function.BiFunction;
@@ -52,20 +54,27 @@ public void setEsqlConfig(EsqlConfig config) {
         this.config = config;
     }
 
+    // testing utility
     public LogicalPlan createStatement(String query) {
         return createStatement(query, new QueryParams());
     }
 
+    // testing utility
     public LogicalPlan createStatement(String query, QueryParams params) {
+        return createStatement(query, params, new PlanTelemetry(new EsqlFunctionRegistry()));
+    }
+
+    public LogicalPlan createStatement(String query, QueryParams params, PlanTelemetry metrics) {
         if (log.isDebugEnabled()) {
             log.debug("Parsing as statement: {}", query);
         }
-        return invokeParser(query, params, EsqlBaseParser::singleStatement, AstBuilder::plan);
+        return invokeParser(query, params, metrics, EsqlBaseParser::singleStatement, AstBuilder::plan);
     }
 
     private <T> T invokeParser(
         String query,
         QueryParams params,
+        PlanTelemetry metrics,
         Function<EsqlBaseParser, ParserRuleContext> parseFunction,
         BiFunction<AstBuilder, ParserRuleContext, T> result
     ) {
@@ -99,7 +108,7 @@ private <T> T invokeParser(
                 log.trace("Parse tree: {}", tree.toStringTree());
             }
 
-            return result.apply(new AstBuilder(params), tree);
+            return result.apply(new AstBuilder(new ExpressionBuilder.ParsingContext(params, metrics)), tree);
         } catch (StackOverflowError e) {
             throw new ParsingException("ESQL statement is too large, causing stack overflow when generating the parsing tree: [{}]", query);
         }