Skip to content

Commit

Permalink
Integrate metrics framework, refactor code and update tests
Browse files Browse the repository at this point in the history 
Signed-off-by: Siddhant Deshmukh <[email protected]>
  • Loading branch information
@deshsidd
deshsidd committed Oct 9, 2023
1 parent cbc1aeb commit 4e79d74
Show file tree
Hide file tree
Showing 4 changed files with 44 additions and 32 deletions.
61 changes: 32 additions & 29 deletions server/src/main/java/org/opensearch/action/search/SearchQueryCategorizor.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -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));
}
}

Expand All @@ -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);
}

}
3 changes: 3 additions & 0 deletions server/src/main/java/org/opensearch/action/search/SearchQueryCounters.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down
10 changes: 8 additions & 2 deletions server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -173,6 +174,8 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,

private final SearchRequestStats searchRequestStats;

private final MetricsRegistry metricsRegistry;

@Inject
public TransportSearchAction(
NodeClient client,
Expand All @@ -187,7 +190,8 @@ public TransportSearchAction(
IndexNameExpressionResolver indexNameExpressionResolver,
NamedWriteableRegistry namedWriteableRegistry,
SearchPipelineService searchPipelineService,
SearchRequestStats searchRequestStats
SearchRequestStats searchRequestStats,
MetricsRegistry metricsRegistry
) {
super(SearchAction.NAME, transportService, actionFilters, (Writeable.Reader<SearchRequest>) SearchRequest::new);
this.client = client;
Expand All @@ -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) {
Expand Down Expand Up @@ -435,7 +440,8 @@ private void executeRequest(
return;
}

SearchQueryCategorizor.categorize(searchRequest.source());
SearchQueryCategorizor searchQueryCategorizor = new SearchQueryCategorizor(metricsRegistry);
searchQueryCategorizor.categorize(searchRequest.source());

ActionListener<SearchSourceBuilder> rewriteListener = ActionListener.wrap(source -> {
if (source != searchRequest.source()) {
Expand Down
2 changes: 1 addition & 1 deletion server/src/test/java/org/opensearch/index/query/QueryBuilderVisitorTests.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ public void testNoOpsVisitor() {

List<QueryBuilder> 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);
Expand Down

0 comments on commit 4e79d74

Please sign in to comment.