From 4e79d74854fa1669944a487b48c0eb81a97fe86c Mon Sep 17 00:00:00 2001 From: Siddhant Deshmukh Date: Mon, 9 Oct 2023 01:43:38 -0700 Subject: [PATCH] Integrate metrics framework, refactor code and update tests Signed-off-by: Siddhant Deshmukh --- .../action/search/SearchQueryCategorizor.java | 61 ++++++++++--------- .../action/search/SearchQueryCounters.java | 3 + .../action/search/TransportSearchAction.java | 10 ++- .../index/query/QueryBuilderVisitorTests.java | 2 +- 4 files changed, 44 insertions(+), 32 deletions(-) diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java index d1995f3da01e6..5d50d42f4e7a3 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java +++ b/server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java @@ -11,9 +11,7 @@ import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.apache.lucene.search.BooleanClause; -import org.opensearch.identity.IdentityService; import org.opensearch.index.query.BoolQueryBuilder; -import org.opensearch.index.query.IntervalsSourceProvider; import org.opensearch.index.query.MatchPhraseQueryBuilder; import org.opensearch.index.query.MatchQueryBuilder; import org.opensearch.index.query.MultiMatchQueryBuilder; @@ -27,56 +25,54 @@ import org.opensearch.index.query.WildcardQueryBuilder; import org.opensearch.index.query.functionscore.FunctionScoreQueryBuilder; import org.opensearch.search.builder.SearchSourceBuilder; -import org.opensearch.telemetry.tracing.attributes.Attributes; +import org.opensearch.telemetry.metrics.MetricsRegistry; +import org.opensearch.telemetry.metrics.tags.Tags; public class SearchQueryCategorizor { private static final Logger log = LogManager.getLogger(SearchQueryCategorizor.class); - public static SearchQueryCounters searchQueryCounters = new SearchQueryCounters(); // What metrics registry to use here? + public static SearchQueryCounters searchQueryCounters; - public static void categorize(SearchSourceBuilder source) { + public SearchQueryCategorizor(MetricsRegistry metricsRegistry) { + searchQueryCounters = new SearchQueryCounters(metricsRegistry); + } + + public void categorize(SearchSourceBuilder source) { QueryBuilder topLevelQueryBuilder = source.query(); - // Get and log the query shape - QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); - topLevelQueryBuilder.visit(shapeVisitor, 0); - String queryShapeJson = shapeVisitor.prettyPrintTree(" "); - log.debug("Query shape : " + queryShapeJson); + logQueryShape(topLevelQueryBuilder); + + incrementQueryCounters(topLevelQueryBuilder); + } + private static void incrementQueryCounters(QueryBuilder topLevelQueryBuilder) { // Increment the query counters using Metric Framework QueryBuilderVisitor queryBuilderVisitor = new QueryBuilderVisitor() { @Override public void accept(QueryBuilder qb, int level) { - // This method will be called for every QueryBuilder node in the tree. - // The tree referred to here is the tree of querybuilders for the incoming search - // query with the topLevelQueryBuilder as the root. - - // Increment counter for current QueryBuilder using Metric Framework. - if (qb instanceof AggregationQ) { - searchQueryCounters.aggCounter.add(1, Attributes.create().addAttribute("level", level)); - } else if (qb instanceof BoolQueryBuilder) { - searchQueryCounters.boolCounter.add(1, Attributes.create().addAttribute("level", level)); + if (qb instanceof BoolQueryBuilder) { + searchQueryCounters.boolCounter.add(1, Tags.create().addTag("level", level)); } else if (qb instanceof FunctionScoreQueryBuilder) { - searchQueryCounters.functionScoreCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.functionScoreCounter.add(1, Tags.create().addTag("level", level)); } else if (qb instanceof MatchQueryBuilder) { - searchQueryCounters.matchCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.matchCounter.add(1, Tags.create().addTag("level", level)); } else if (qb instanceof MatchPhraseQueryBuilder) { - searchQueryCounters.matchPhrasePrefixCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.matchPhrasePrefixCounter.add(1, Tags.create().addTag("level", level)); } else if (qb instanceof MultiMatchQueryBuilder) { - searchQueryCounters.multiMatchCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.multiMatchCounter.add(1, Tags.create().addTag("level", level)); } else if (qb instanceof QueryStringQueryBuilder) { - searchQueryCounters.queryStringQueryCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.queryStringQueryCounter.add(1, Tags.create().addTag("level", level)); } else if (qb instanceof RangeQueryBuilder) { - searchQueryCounters.rangeCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.rangeCounter.add(1, Tags.create().addTag("level", level)); } else if (qb instanceof RegexpQueryBuilder) { - searchQueryCounters.regexCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.regexCounter.add(1, Tags.create().addTag("level", level)); } else if (qb instanceof TermQueryBuilder) { - searchQueryCounters.termCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.termCounter.add(1, Tags.create().addTag("level", level)); } else if (qb instanceof WildcardQueryBuilder) { - searchQueryCounters.wildcardCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.wildcardCounter.add(1, Tags.create().addTag("level", level)); } else { - searchQueryCounters.otherQueryCounter.add(1, Attributes.create().addAttribute("level", level)); + searchQueryCounters.otherQueryCounter.add(1, Tags.create().addTag("level", level)); } } @@ -88,4 +84,11 @@ public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { topLevelQueryBuilder.visit(queryBuilderVisitor, 0); } + private static void logQueryShape(QueryBuilder topLevelQueryBuilder) { + QueryShapeVisitor shapeVisitor = new QueryShapeVisitor(); + topLevelQueryBuilder.visit(shapeVisitor, 0); + String queryShapeJson = shapeVisitor.prettyPrintTree(" "); + log.debug("Query shape : " + queryShapeJson); + } + } diff --git a/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java b/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java index 6e4fdb8d21ba5..43195a9a1afa2 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java +++ b/server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java @@ -8,6 +8,9 @@ package org.opensearch.action.search; +import org.opensearch.telemetry.metrics.Counter; +import org.opensearch.telemetry.metrics.MetricsRegistry; + public class SearchQueryCounters { private final MetricsRegistry metricsRegistry; diff --git a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java index 0e06bf3f2dd77..15e1642d5d374 100644 --- a/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java +++ b/server/src/main/java/org/opensearch/action/search/TransportSearchAction.java @@ -97,6 +97,7 @@ import org.opensearch.transport.RemoteTransportException; import org.opensearch.transport.Transport; import org.opensearch.transport.TransportService; +import org.opensearch.telemetry.metrics.MetricsRegistry; import java.util.ArrayList; import java.util.Arrays; @@ -173,6 +174,8 @@ public class TransportSearchAction extends HandledTransportAction) SearchRequest::new); this.client = client; @@ -205,6 +209,7 @@ public TransportSearchAction( this.isRequestStatsEnabled = clusterService.getClusterSettings().get(SEARCH_REQUEST_STATS_ENABLED); clusterService.getClusterSettings().addSettingsUpdateConsumer(SEARCH_REQUEST_STATS_ENABLED, this::setIsRequestStatsEnabled); this.searchRequestStats = searchRequestStats; + this.metricsRegistry = metricsRegistry; } private void setIsRequestStatsEnabled(boolean isRequestStatsEnabled) { @@ -435,7 +440,8 @@ private void executeRequest( return; } - SearchQueryCategorizor.categorize(searchRequest.source()); + SearchQueryCategorizor searchQueryCategorizor = new SearchQueryCategorizor(metricsRegistry); + searchQueryCategorizor.categorize(searchRequest.source()); ActionListener rewriteListener = ActionListener.wrap(source -> { if (source != searchRequest.source()) { diff --git a/server/src/test/java/org/opensearch/index/query/QueryBuilderVisitorTests.java b/server/src/test/java/org/opensearch/index/query/QueryBuilderVisitorTests.java index 7849d3985ca59..0fd627fe642b7 100644 --- a/server/src/test/java/org/opensearch/index/query/QueryBuilderVisitorTests.java +++ b/server/src/test/java/org/opensearch/index/query/QueryBuilderVisitorTests.java @@ -21,7 +21,7 @@ public void testNoOpsVisitor() { List visitedQueries = new ArrayList<>(); QueryBuilderVisitor qbv = createTestVisitor(visitedQueries); - boolQueryBuilder.visit(qbv); + boolQueryBuilder.visit(qbv, 0); QueryBuilderVisitor subQbv = qbv.getChildVisitor(BooleanClause.Occur.MUST_NOT); assertEquals(0, visitedQueries.size()); assertEquals(qbv, subQbv);