diff --git a/src/main/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListener.java b/src/main/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListener.java index 6d7b0532..30abd662 100644 --- a/src/main/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListener.java +++ b/src/main/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListener.java @@ -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; @@ -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; @@ -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 = @@ -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(); } @@ -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(); @@ -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); } @@ -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(); + 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) { + 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, diff --git a/src/test/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListenerTests.java b/src/test/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListenerTests.java index 7d5c726c..b28698d0 100644 --- a/src/test/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListenerTests.java +++ b/src/test/java/org/opensearch/performanceanalyzer/listener/RTFPerformanceAnalyzerSearchListenerTests.java @@ -13,6 +13,7 @@ import org.junit.Before; import org.junit.BeforeClass; import org.junit.Test; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.Mockito; import org.opensearch.action.search.SearchShardTask; @@ -20,15 +21,22 @@ import org.opensearch.core.index.Index; import org.opensearch.core.index.shard.ShardId; import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; +import org.opensearch.index.query.BoolQueryBuilder; +import org.opensearch.index.query.QueryBuilders; +import org.opensearch.index.query.ScriptQueryBuilder; +import org.opensearch.index.query.functionscore.ScriptScoreQueryBuilder; import org.opensearch.performanceanalyzer.OpenSearchResources; import org.opensearch.performanceanalyzer.ShardMetricsCollector; import org.opensearch.performanceanalyzer.commons.metrics.RTFMetrics; import org.opensearch.performanceanalyzer.commons.util.Util; import org.opensearch.performanceanalyzer.config.PerformanceAnalyzerController; import org.opensearch.performanceanalyzer.util.Utils; +import org.opensearch.script.Script; +import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.search.internal.SearchContext; import org.opensearch.search.internal.ShardSearchRequest; 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; @@ -48,6 +56,7 @@ public class RTFPerformanceAnalyzerSearchListenerTests { @Mock private Histogram searchLatencyHistogram; @Mock private Histogram shardMetricsCpuHistogram; @Mock private Histogram shardMetricsHeapHistogram; + @Mock private Counter shardQueryWithScriptCounter; @Mock private Index index; @Mock private TaskResourceUsage taskResourceUsage; @@ -89,6 +98,10 @@ public void init() { } return null; }); + Mockito.when( + metricsRegistry.createCounter( + Mockito.anyString(), Mockito.anyString(), Mockito.anyString())) + .thenReturn(shardQueryWithScriptCounter); searchListener = new RTFPerformanceAnalyzerSearchListener(controller); assertEquals( RTFPerformanceAnalyzerSearchListener.class.getSimpleName(), @@ -203,6 +216,308 @@ public void testTaskCompletionListener() { Mockito.verify(shardHeap).record(Mockito.anyDouble(), Mockito.any(Tags.class)); } + // ==================== Script Query Detection Tests ==================== + + @Test + public void testQueryPhaseEmitsScriptCounterForTopLevelScriptQuery() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(new ScriptQueryBuilder(new Script("doc['field'].value > 5"))); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter).add(Mockito.eq(1L), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseEmitsScriptCounterForScriptScoreQuery() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query( + new ScriptScoreQueryBuilder( + QueryBuilders.matchAllQuery(), new Script("_score * doc['boost'].value"))); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter).add(Mockito.eq(1L), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseEmitsScriptCounterForScriptFields() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(true); + + // No script in query tree — but script_fields present + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(QueryBuilders.matchAllQuery()); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter).add(Mockito.eq(1L), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseDoesNotEmitScriptCounterForPlainQuery() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(QueryBuilders.termQuery("status", "active")); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter, Mockito.never()) + .add(Mockito.anyLong(), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseDetectsScriptQueryNestedInBoolMust() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + // bool { must: [term, script_query] } + BoolQueryBuilder boolQuery = + QueryBuilders.boolQuery() + .must(QueryBuilders.termQuery("status", "active")) + .must(new ScriptQueryBuilder(new Script("doc['age'].value > 21"))); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(boolQuery); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter).add(Mockito.eq(1L), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseDetectsScriptQueryNestedInBoolFilter() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + // bool { filter: [script_query] } + BoolQueryBuilder boolQuery = + QueryBuilders.boolQuery() + .filter(new ScriptQueryBuilder(new Script("doc['score'].value > 0"))); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(boolQuery); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter).add(Mockito.eq(1L), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseDetectsScriptQueryNestedInBoolShould() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + // bool { should: [match_all, script_query] } + BoolQueryBuilder boolQuery = + QueryBuilders.boolQuery() + .should(QueryBuilders.matchAllQuery()) + .should(new ScriptQueryBuilder(new Script("true"))); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(boolQuery); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter).add(Mockito.eq(1L), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseDetectsScriptQueryDeeplyNested() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + // bool { must: [ bool { filter: [ bool { must_not: [script_query] } ] } ] } + BoolQueryBuilder innermost = + QueryBuilders.boolQuery() + .mustNot(new ScriptQueryBuilder(new Script("doc['x'].value < 0"))); + BoolQueryBuilder middle = QueryBuilders.boolQuery().filter(innermost); + BoolQueryBuilder outer = QueryBuilders.boolQuery().must(middle); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(outer); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter).add(Mockito.eq(1L), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseNoScriptInDeeplyNestedBool() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + // Deep nesting with no scripts at all + BoolQueryBuilder inner = QueryBuilders.boolQuery().must(QueryBuilders.termQuery("a", "1")); + BoolQueryBuilder middle = + QueryBuilders.boolQuery().filter(inner).should(QueryBuilders.matchAllQuery()); + BoolQueryBuilder outer = + QueryBuilders.boolQuery().must(middle).mustNot(QueryBuilders.termQuery("b", "2")); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(outer); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter, Mockito.never()) + .add(Mockito.anyLong(), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseScriptFieldsTakesPriorityOverQueryTree() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(true); + + // Script in both script_fields AND query tree — counter should fire once (short-circuits on + // script_fields) + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(new ScriptQueryBuilder(new Script("true"))); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + // Should be called exactly once — from the script_fields check, not the query tree + Mockito.verify(shardQueryWithScriptCounter, Mockito.times(1)) + .add(Mockito.eq(1L), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseNullSourceDoesNotEmitScriptCounter() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + Mockito.when(shardSearchRequest.source()).thenReturn(null); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter, Mockito.never()) + .add(Mockito.anyLong(), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseNullQueryInSourceDoesNotEmitScriptCounter() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + SearchSourceBuilder source = new SearchSourceBuilder(); + // No query set — source.query() returns null + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter, Mockito.never()) + .add(Mockito.anyLong(), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseScriptDetectionHandlesExceptionGracefully() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()) + .thenThrow(new RuntimeException("simulated failure")); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + Mockito.verify(shardQueryWithScriptCounter, Mockito.never()) + .add(Mockito.anyLong(), Mockito.any(Tags.class)); + } + + @Test + public void testQueryPhaseScriptCounterEmitsCorrectTags() { + initializeValidSearchContext(true); + setupIndexMocks(); + Mockito.when(controller.getCollectorsRunModeValue()) + .thenReturn(Util.CollectorMode.TELEMETRY.getValue()); + Mockito.when(searchContext.hasScriptFields()).thenReturn(false); + + SearchSourceBuilder source = new SearchSourceBuilder(); + source.query(new ScriptQueryBuilder(new Script("true"))); + Mockito.when(shardSearchRequest.source()).thenReturn(source); + + searchListener.preQueryPhase(searchContext); + searchListener.queryPhase(searchContext, 0L); + + ArgumentCaptor tagsCaptor = ArgumentCaptor.forClass(Tags.class); + Mockito.verify(shardQueryWithScriptCounter).add(Mockito.eq(1L), tagsCaptor.capture()); + + String tagsString = tagsCaptor.getValue().toString(); + assertTrue("Tags should contain index name", tagsString.contains("myTestIndex")); + assertTrue("Tags should contain index UUID", tagsString.contains("abc-def")); + } + + private void setupIndexMocks() { + Mockito.when(shardId.getIndex()).thenReturn(index); + Mockito.when(index.getName()).thenReturn("myTestIndex"); + Mockito.when(index.getUUID()).thenReturn("abc-def"); + } + private void initializeValidSearchContext(boolean isValid) { if (isValid) { Mockito.when(searchContext.request()).thenReturn(shardSearchRequest);