Make search query counters dynamic to support all query types (#12601) (#13397)

deshsidd · web-flow · commit 1b8cba309d01 · 2024-04-25T17:13:45.000-07:00
Signed-off-by: Siddhant Deshmukh &lt;deshsid@amazon.com&gt;
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -24,6 +24,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Add an individual setting of rate limiter for segment replication ([#12959](https://github.com/opensearch-project/OpenSearch/pull/12959))
 - [Streaming Indexing] Ensure support of the new transport by security plugin ([#13174](https://github.com/opensearch-project/OpenSearch/pull/13174))
 - Add cluster setting to dynamically configure the buckets for filter rewrite optimization. ([#13179](https://github.com/opensearch-project/OpenSearch/pull/13179))
+- Make search query counters dynamic to support all query types ([#12601](https://github.com/opensearch-project/OpenSearch/pull/12601))
 
 ### Dependencies
 - Bump `org.apache.commons:commons-configuration2` from 2.10.0 to 2.10.1 ([#12896](https://github.com/opensearch-project/OpenSearch/pull/12896))
diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java
@@ -8,39 +8,14 @@
 
 package org.opensearch.action.search;
 
-import org.opensearch.index.query.BoolQueryBuilder;
-import org.opensearch.index.query.BoostingQueryBuilder;
-import org.opensearch.index.query.ConstantScoreQueryBuilder;
-import org.opensearch.index.query.DisMaxQueryBuilder;
-import org.opensearch.index.query.DistanceFeatureQueryBuilder;
-import org.opensearch.index.query.ExistsQueryBuilder;
-import org.opensearch.index.query.FieldMaskingSpanQueryBuilder;
-import org.opensearch.index.query.FuzzyQueryBuilder;
-import org.opensearch.index.query.GeoBoundingBoxQueryBuilder;
-import org.opensearch.index.query.GeoDistanceQueryBuilder;
-import org.opensearch.index.query.GeoPolygonQueryBuilder;
-import org.opensearch.index.query.GeoShapeQueryBuilder;
-import org.opensearch.index.query.IntervalQueryBuilder;
-import org.opensearch.index.query.MatchAllQueryBuilder;
-import org.opensearch.index.query.MatchPhraseQueryBuilder;
-import org.opensearch.index.query.MatchQueryBuilder;
-import org.opensearch.index.query.MultiMatchQueryBuilder;
-import org.opensearch.index.query.PrefixQueryBuilder;
 import org.opensearch.index.query.QueryBuilder;
-import org.opensearch.index.query.QueryStringQueryBuilder;
-import org.opensearch.index.query.RangeQueryBuilder;
-import org.opensearch.index.query.RegexpQueryBuilder;
-import org.opensearch.index.query.ScriptQueryBuilder;
-import org.opensearch.index.query.SimpleQueryStringBuilder;
-import org.opensearch.index.query.TermQueryBuilder;
-import org.opensearch.index.query.WildcardQueryBuilder;
-import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder;
 import org.opensearch.telemetry.metrics.Counter;
 import org.opensearch.telemetry.metrics.MetricsRegistry;
 import org.opensearch.telemetry.metrics.tags.Tags;
 
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * Class contains all the Counters related to search query types.
@@ -49,244 +24,47 @@ final class SearchQueryCounters {
     private static final String LEVEL_TAG = "level";
     private static final String UNIT = "1";
     private final MetricsRegistry metricsRegistry;
-
-    // Counters related to Query types
     public final Counter aggCounter;
-    public final Counter boolCounter;
-    public final Counter boostingCounter;
-    public final Counter constantScoreCounter;
-    public final Counter disMaxCounter;
-    public final Counter distanceFeatureCounter;
-    public final Counter existsCounter;
-    public final Counter fieldMaskingSpanCounter;
-    public final Counter functionScoreCounter;
-    public final Counter fuzzyCounter;
-    public final Counter geoBoundingBoxCounter;
-    public final Counter geoDistanceCounter;
-    public final Counter geoPolygonCounter;
-    public final Counter geoShapeCounter;
-    public final Counter intervalCounter;
-    public final Counter matchCounter;
-    public final Counter matchallCounter;
-    public final Counter matchPhrasePrefixCounter;
-    public final Counter multiMatchCounter;
     public final Counter otherQueryCounter;
-    public final Counter prefixCounter;
-    public final Counter queryStringCounter;
-    public final Counter rangeCounter;
-    public final Counter regexpCounter;
-    public final Counter scriptCounter;
-    public final Counter simpleQueryStringCounter;
     public final Counter sortCounter;
-    public final Counter skippedCounter;
-    public final Counter termCounter;
-    public final Counter totalCounter;
-    public final Counter wildcardCounter;
-    public final Counter numberOfInputFieldsCounter;
     private final Map<Class<? extends QueryBuilder>, Counter> queryHandlers;
+    public final ConcurrentHashMap<String, Counter> nameToQueryTypeCounters;
 
     public SearchQueryCounters(MetricsRegistry metricsRegistry) {
         this.metricsRegistry = metricsRegistry;
+        this.nameToQueryTypeCounters = new ConcurrentHashMap<>();
         this.aggCounter = metricsRegistry.createCounter(
             "search.query.type.agg.count",
             "Counter for the number of top level agg search queries",
             UNIT
         );
-        this.boolCounter = metricsRegistry.createCounter(
-            "search.query.type.bool.count",
-            "Counter for the number of top level and nested bool search queries",
-            UNIT
-        );
-        this.boostingCounter = metricsRegistry.createCounter(
-            "search.query.type.boost.count",
-            "Counter for the number of top level and nested boost search queries",
-            UNIT
-        );
-        this.constantScoreCounter = metricsRegistry.createCounter(
-            "search.query.type.counstantscore.count",
-            "Counter for the number of top level and nested constant score search queries",
-            UNIT
-        );
-        this.disMaxCounter = metricsRegistry.createCounter(
-            "search.query.type.dismax.count",
-            "Counter for the number of top level and nested disjuntion max search queries",
-            UNIT
-        );
-        this.distanceFeatureCounter = metricsRegistry.createCounter(
-            "search.query.type.distancefeature.count",
-            "Counter for the number of top level and nested distance feature search queries",
-            UNIT
-        );
-        this.existsCounter = metricsRegistry.createCounter(
-            "search.query.type.exists.count",
-            "Counter for the number of top level and nested exists search queries",
-            UNIT
-        );
-        this.fieldMaskingSpanCounter = metricsRegistry.createCounter(
-            "search.query.type.fieldmaskingspan.count",
-            "Counter for the number of top level and nested field masking span search queries",
-            UNIT
-        );
-        this.functionScoreCounter = metricsRegistry.createCounter(
-            "search.query.type.functionscore.count",
-            "Counter for the number of top level and nested function score search queries",
-            UNIT
-        );
-        this.fuzzyCounter = metricsRegistry.createCounter(
-            "search.query.type.fuzzy.count",
-            "Counter for the number of top level and nested fuzzy search queries",
-            UNIT
-        );
-        this.geoBoundingBoxCounter = metricsRegistry.createCounter(
-            "search.query.type.geoboundingbox.count",
-            "Counter for the number of top level and nested geo bounding box queries",
-            UNIT
-        );
-        this.geoDistanceCounter = metricsRegistry.createCounter(
-            "search.query.type.geodistance.count",
-            "Counter for the number of top level and nested geo distance queries",
-            UNIT
-        );
-        this.geoPolygonCounter = metricsRegistry.createCounter(
-            "search.query.type.geopolygon.count",
-            "Counter for the number of top level and nested geo polygon queries",
-            UNIT
-        );
-        this.geoShapeCounter = metricsRegistry.createCounter(
-            "search.query.type.geoshape.count",
-            "Counter for the number of top level and nested geo shape queries",
-            UNIT
-        );
-        this.intervalCounter = metricsRegistry.createCounter(
-            "search.query.type.interval.count",
-            "Counter for the number of top level and nested interval queries",
-            UNIT
-        );
-        this.matchCounter = metricsRegistry.createCounter(
-            "search.query.type.match.count",
-            "Counter for the number of top level and nested match search queries",
-            UNIT
-        );
-        this.matchallCounter = metricsRegistry.createCounter(
-            "search.query.type.matchall.count",
-            "Counter for the number of top level and nested match all search queries",
-            UNIT
-        );
-        this.matchPhrasePrefixCounter = metricsRegistry.createCounter(
-            "search.query.type.matchphrase.count",
-            "Counter for the number of top level and nested match phrase prefix search queries",
-            UNIT
-        );
-        this.multiMatchCounter = metricsRegistry.createCounter(
-            "search.query.type.multimatch.count",
-            "Counter for the number of top level and nested multi match search queries",
-            UNIT
-        );
         this.otherQueryCounter = metricsRegistry.createCounter(
             "search.query.type.other.count",
             "Counter for the number of top level and nested search queries that do not match any other categories",
             UNIT
         );
-        this.prefixCounter = metricsRegistry.createCounter(
-            "search.query.type.prefix.count",
-            "Counter for the number of top level and nested search queries that match prefix queries",
-            UNIT
-        );
-        this.queryStringCounter = metricsRegistry.createCounter(
-            "search.query.type.querystringquery.count",
-            "Counter for the number of top level and nested queryStringQuery search queries",
-            UNIT
-        );
-        this.rangeCounter = metricsRegistry.createCounter(
-            "search.query.type.range.count",
-            "Counter for the number of top level and nested range search queries",
-            UNIT
-        );
-        this.regexpCounter = metricsRegistry.createCounter(
-            "search.query.type.regex.count",
-            "Counter for the number of top level and nested regex search queries",
-            UNIT
-        );
-        this.scriptCounter = metricsRegistry.createCounter(
-            "search.query.type.script.count",
-            "Counter for the number of top level and nested script search queries",
-            UNIT
-        );
-        this.simpleQueryStringCounter = metricsRegistry.createCounter(
-            "search.query.type.simplequerystring.count",
-            "Counter for the number of top level and nested script simple query string search queries",
-            UNIT
-        );
-        this.skippedCounter = metricsRegistry.createCounter(
-            "search.query.type.skipped.count",
-            "Counter for the number queries skipped due to error",
-            UNIT
-        );
         this.sortCounter = metricsRegistry.createCounter(
             "search.query.type.sort.count",
             "Counter for the number of top level sort search queries",
             UNIT
         );
-        this.termCounter = metricsRegistry.createCounter(
-            "search.query.type.term.count",
-            "Counter for the number of top level and nested term search queries",
-            UNIT
-        );
-        this.totalCounter = metricsRegistry.createCounter(
-            "search.query.type.total.count",
-            "Counter for the number of top level and nested search queries",
-            UNIT
-        );
-        this.wildcardCounter = metricsRegistry.createCounter(
-            "search.query.type.wildcard.count",
-            "Counter for the number of top level and nested wildcard search queries",
-            UNIT
-        );
-        this.numberOfInputFieldsCounter = metricsRegistry.createCounter(
-            "search.query.type.numberofinputfields.count",
-            "Counter for the number of input fields in the search queries",
-            UNIT
-        );
         this.queryHandlers = new HashMap<>();
-        initializeQueryHandlers();
+
     }
 
     public void incrementCounter(QueryBuilder queryBuilder, int level) {
-        Counter counter = queryHandlers.get(queryBuilder.getClass());
-        if (counter != null) {
-            counter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        } else {
-            otherQueryCounter.add(1, Tags.create().addTag(LEVEL_TAG, level));
-        }
-    }
+        String uniqueQueryCounterName = queryBuilder.getName();
 
-    private void initializeQueryHandlers() {
+        Counter counter = nameToQueryTypeCounters.computeIfAbsent(uniqueQueryCounterName, k -> createQueryCounter(k));
+        counter.add(1, Tags.create().addTag(LEVEL_TAG, level));
+    }
 
-        queryHandlers.put(BoolQueryBuilder.class, boolCounter);
-        queryHandlers.put(FunctionScoreQueryBuilder.class, functionScoreCounter);
-        queryHandlers.put(MatchQueryBuilder.class, matchCounter);
-        queryHandlers.put(MatchPhraseQueryBuilder.class, matchPhrasePrefixCounter);
-        queryHandlers.put(MultiMatchQueryBuilder.class, multiMatchCounter);
-        queryHandlers.put(QueryStringQueryBuilder.class, queryStringCounter);
-        queryHandlers.put(RangeQueryBuilder.class, rangeCounter);
-        queryHandlers.put(RegexpQueryBuilder.class, regexpCounter);
-        queryHandlers.put(TermQueryBuilder.class, termCounter);
-        queryHandlers.put(WildcardQueryBuilder.class, wildcardCounter);
-        queryHandlers.put(BoostingQueryBuilder.class, boostingCounter);
-        queryHandlers.put(ConstantScoreQueryBuilder.class, constantScoreCounter);
-        queryHandlers.put(DisMaxQueryBuilder.class, disMaxCounter);
-        queryHandlers.put(DistanceFeatureQueryBuilder.class, distanceFeatureCounter);
-        queryHandlers.put(ExistsQueryBuilder.class, existsCounter);
-        queryHandlers.put(FieldMaskingSpanQueryBuilder.class, fieldMaskingSpanCounter);
-        queryHandlers.put(FuzzyQueryBuilder.class, fuzzyCounter);
-        queryHandlers.put(GeoBoundingBoxQueryBuilder.class, geoBoundingBoxCounter);
-        queryHandlers.put(GeoDistanceQueryBuilder.class, geoDistanceCounter);
-        queryHandlers.put(GeoPolygonQueryBuilder.class, geoPolygonCounter);
-        queryHandlers.put(GeoShapeQueryBuilder.class, geoShapeCounter);
-        queryHandlers.put(IntervalQueryBuilder.class, intervalCounter);
-        queryHandlers.put(MatchAllQueryBuilder.class, matchallCounter);
-        queryHandlers.put(PrefixQueryBuilder.class, prefixCounter);
-        queryHandlers.put(ScriptQueryBuilder.class, scriptCounter);
-        queryHandlers.put(SimpleQueryStringBuilder.class, simpleQueryStringCounter);
+    private Counter createQueryCounter(String counterName) {
+        Counter counter = metricsRegistry.createCounter(
+            "search.query.type." + counterName + ".count",
+            "Counter for the number of top level and nested " + counterName + " search queries",
+            UNIT
+        );
+        return counter;
     }
 }
diff --git a/server/src/test/java/org/opensearch/action/search/SearchQueryCategorizerTests.java b/server/src/test/java/org/opensearch/action/search/SearchQueryCategorizerTests.java
