Drop the HashSet gating when counting commands

This distributes the telemetry counting in the parser so that no HashSet is required to correctly account the commands.
bpintea · Jan 28, 2025 · c36b45c · c36b45c
1 parent bd4e845
commit c36b45c
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 16 deletions.
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/LogicalPlanBuilder.java
@@ -103,9 +103,6 @@ public LogicalPlanBuilder(ParsingContext context) {
 
     protected LogicalPlan plan(ParseTree ctx) {
         LogicalPlan p = ParserUtils.typedParsing(this, ctx, LogicalPlan.class);
-        if (p instanceof TelemetryAware ma) {
-            this.context.telemetry().command(ma);
-        }
         var errors = this.context.params().parsingErrors();
         if (errors.hasNext() == false) {
             return p;
@@ -130,7 +127,9 @@ protected List<LogicalPlan> plans(List<? extends ParserRuleContext> ctxs) {
 
     @Override
     public LogicalPlan visitSingleStatement(EsqlBaseParser.SingleStatementContext ctx) {
-        return plan(ctx.query());
+        var plan = plan(ctx.query());
+        telemetryCount(plan);
+        return plan;
     }
 
     @Override
@@ -145,13 +144,21 @@ public LogicalPlan visitCompositeQuery(EsqlBaseParser.CompositeQueryContext ctx)
         }
         try {
             LogicalPlan input = plan(ctx.query());
+            telemetryCount(input);
             PlanFactory makePlan = typedParsing(this, ctx.processingCommand(), PlanFactory.class);
             return makePlan.apply(input);
         } finally {
             queryDepth--;
         }
     }
 
+    private LogicalPlan telemetryCount(LogicalPlan node) {
+        if (node instanceof TelemetryAware ma) {
+            this.context.telemetry().command(ma);
+        }
+        return node;
+    }
+
     @Override
     public PlanFactory visitEvalCommand(EsqlBaseParser.EvalCommandContext ctx) {
         return p -> new Eval(source(ctx), p, visitFields(ctx.fields()));

diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/telemetry/PlanTelemetry.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/telemetry/PlanTelemetry.java
@@ -8,24 +8,19 @@
 package org.elasticsearch.xpack.esql.telemetry;
 
 import org.elasticsearch.xpack.esql.capabilities.TelemetryAware;
-import org.elasticsearch.xpack.esql.core.QlIllegalArgumentException;
 import org.elasticsearch.xpack.esql.core.expression.function.Function;
+import org.elasticsearch.xpack.esql.core.util.Check;
 import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
 
 import java.util.HashMap;
-import java.util.HashSet;
 import java.util.Locale;
 import java.util.Map;
-import java.util.Set;
-
-import static org.elasticsearch.common.Strings.format;
 
 /**
  * This class is responsible for collecting metrics related to ES|QL planning.
  */
 public class PlanTelemetry {
     private final EsqlFunctionRegistry functionRegistry;
-    private final Set<TelemetryAware> telemetryAwares = new HashSet<>();
     private final Map<String, Integer> commands = new HashMap<>();
     private final Map<String, Integer> functions = new HashMap<>();
 
@@ -38,12 +33,8 @@ private void add(Map<String, Integer> map, String key) {
     }
 
     public void command(TelemetryAware command) {
-        if (telemetryAwares.add(command)) {
-            if (command.telemetryLabel() == null) {
-                throw new QlIllegalArgumentException(format("TelemetryAware [{}] has no metric name", command));
-            }
-            add(commands, command.telemetryLabel());
-        }
+        Check.notNull(command.telemetryLabel(), "TelemetryAware [{}] has no telemetry label", command);
+        add(commands, command.telemetryLabel());
     }
 
     public void function(String name) {