Skip to content

Add metric to track queries with script #934

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shivamg2017 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add metric to track queries with script #934

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
88 changes: 88 additions & 0 deletions ...ava/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListener.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,12 @@
import java.util.Map;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import org.apache.lucene.search.BooleanClause;
import org.opensearch.core.action.NotifyOnceListener;
import org.opensearch.index.query.QueryBuilder;
import org.opensearch.index.query.QueryBuilderVisitor;
import org.opensearch.index.query.ScriptQueryBuilder;
import org.opensearch.index.query.functionscore.ScriptScoreQueryBuilder;
import org.opensearch.index.shard.SearchOperationListener;
import org.opensearch.performanceanalyzer.OpenSearchResources;
import org.opensearch.performanceanalyzer.ShardMetricsCollector;
Expand All @@ -23,8 +28,10 @@
import org.opensearch.performanceanalyzer.commons.util.Util;
import org.opensearch.performanceanalyzer.config.PerformanceAnalyzerController;
import org.opensearch.performanceanalyzer.util.Utils;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.search.internal.SearchContext;
import org.opensearch.tasks.Task;
import org.opensearch.telemetry.metrics.Counter;
import org.opensearch.telemetry.metrics.Histogram;
import org.opensearch.telemetry.metrics.MetricsRegistry;
import org.opensearch.telemetry.metrics.tags.Tags;
Expand All @@ -51,8 +58,11 @@ public class RTFPerformanceAnalyzerSearchListener
private final Histogram cpuUtilizationHistogram;
private final Histogram heapUsedHistogram;
private final Histogram searchLatencyHistogram;
private final Counter shardQueryWithScriptCounter;
private final int numProcessors;

public static final String SHARD_QUERY_WITH_SCRIPT_COUNT = "shard_query_with_script_count";

public RTFPerformanceAnalyzerSearchListener(final PerformanceAnalyzerController controller) {
this.controller = controller;
this.cpuUtilizationHistogram =
Expand All @@ -61,6 +71,9 @@ public RTFPerformanceAnalyzerSearchListener(final PerformanceAnalyzerController
createHeapUsedHistogram(OpenSearchResources.INSTANCE.getMetricsRegistry());
this.searchLatencyHistogram =
createSearchLatencyHistogram(OpenSearchResources.INSTANCE.getMetricsRegistry());
this.shardQueryWithScriptCounter =
createShardQueryWithScriptCounter(
OpenSearchResources.INSTANCE.getMetricsRegistry());
this.threadLocal = ThreadLocal.withInitial(HashMap::new);
this.numProcessors = Runtime.getRuntime().availableProcessors();
}
Expand Down Expand Up @@ -103,6 +116,18 @@ private Histogram createSearchLatencyHistogram(MetricsRegistry metricsRegistry)
}
}

private Counter createShardQueryWithScriptCounter(MetricsRegistry metricsRegistry) {
if (metricsRegistry != null) {
return metricsRegistry.createCounter(
SHARD_QUERY_WITH_SCRIPT_COUNT,
"Count of shard query phases where a script was used",
RTFMetrics.MetricUnits.COUNT.toString());
} else {
LOG.debug("MetricsRegistry is null");
return null;
}
}

@Override
public String toString() {
return RTFPerformanceAnalyzerSearchListener.class.getSimpleName();
Expand Down Expand Up @@ -195,6 +220,9 @@ public void queryPhase(SearchContext searchContext, long tookInNanos) {
searchLatencyHistogram.record(
queryTimeInMills, createTags(searchContext, SHARD_QUERY_PHASE, false));

// Emit counter if query contains a script
emitQueryWithScriptMetric(searchContext);

addResourceTrackingCompletionListener(
searchContext, queryStartTime, tookInNanos, SHARD_QUERY_PHASE, false);
}
Expand Down Expand Up @@ -231,6 +259,66 @@ public void failedFetchPhase(SearchContext searchContext) {
searchContext, fetchStartTime, fetchTime, SHARD_FETCH_PHASE, true);
}

/**
* Emits the shard_query_with_script_count counter if the query contains a script. Uses
* QueryBuilderVisitor to recursively walk the query tree and detect ScriptQueryBuilder or
* ScriptScoreQueryBuilder. Also checks for script_fields via hasScriptFields().
*/
private void emitQueryWithScriptMetric(SearchContext searchContext) {
if (shardQueryWithScriptCounter == null) {
return;
}
try {
// Check script_fields
if (searchContext.hasScriptFields()) {
shardQueryWithScriptCounter.add(1, createTags(searchContext));
return;
}

// Walk the query tree for script queries
SearchSourceBuilder source = searchContext.request().source();
if (source != null && source.query() != null) {
ScriptQueryVisitor visitor = new ScriptQueryVisitor();
Copy link
Copy Markdown
Collaborator

@Gaganjuneja Gaganjuneja Apr 2, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What does SearchContext::hasScriptFields method returns?

Copy link
Copy Markdown
Author

@shivamg2017 shivamg2017 Apr 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Returns true when the request has script_fields — scripts that compute new fields in the response (e.g. "script_fields": {"doubled": {"script": ...}}). These are separate from scripts in the query DSL, so we check both: hasScriptFields() for script_fields, and the visitor for script queries inside the query tree

source.query().visit(visitor);
if (visitor.hasScript()) {
shardQueryWithScriptCounter.add(1, createTags(searchContext));
}
}
} catch (Exception e) {
LOG.debug("Error emitting shard_query_with_script_count metric", e);
}
}

/**
* Visitor that recursively walks the query tree to detect script-based queries. Detects
* ScriptQueryBuilder and ScriptScoreQueryBuilder at any nesting level.
*/
private static class ScriptQueryVisitor implements QueryBuilderVisitor {
private boolean scriptDetected = false;

@Override
public void accept(QueryBuilder qb) {
if (scriptDetected) {
return;
}
if (qb instanceof ScriptQueryBuilder || qb instanceof ScriptScoreQueryBuilder) {
Copy link
Copy Markdown
Collaborator

@Gaganjuneja Gaganjuneja Apr 10, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we short circuit if containsScript == true?

Copy link
Copy Markdown
Author

@shivamg2017 shivamg2017 Apr 14, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

scriptDetected = true;
}
}

@Override
public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) {
if (scriptDetected) {
return QueryBuilderVisitor.NO_OP_VISITOR;
}
return this;
}

public boolean hasScript() {
return scriptDetected;
}
}

private void addResourceTrackingCompletionListener(
SearchContext searchContext,
long startTime,
Expand Down
Loading
Loading