diff --git a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java index c54f1add..7c4f212c 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java +++ b/src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java @@ -138,6 +138,10 @@ public void onRequestFailure(final SearchPhaseContext context, final SearchReque constructSearchQueryRecord(context, searchRequestContext); } + private boolean shouldCollect(MetricType metricType) { + return queryInsightsService.isSearchQueryMetricsFeatureEnabled() || queryInsightsService.isCollectionEnabled(metricType); + } + private void constructSearchQueryRecord(final SearchPhaseContext context, final SearchRequestContext searchRequestContext) { SearchTask searchTask = context.getTask(); List tasksResourceUsages = searchRequestContext.getPhaseResourceUsage(); @@ -154,19 +158,19 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final final SearchRequest request = context.getRequest(); try { Map measurements = new HashMap<>(); - if (queryInsightsService.isCollectionEnabled(MetricType.LATENCY)) { + if (shouldCollect(MetricType.LATENCY)) { measurements.put( MetricType.LATENCY, TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - searchRequestContext.getAbsoluteStartNanos()) ); } - if (queryInsightsService.isCollectionEnabled(MetricType.CPU)) { + if (shouldCollect(MetricType.CPU)) { measurements.put( MetricType.CPU, tasksResourceUsages.stream().map(a -> a.getTaskResourceUsage().getCpuTimeInNanos()).mapToLong(Long::longValue).sum() ); } - if (queryInsightsService.isCollectionEnabled(MetricType.MEMORY)) { + if (shouldCollect(MetricType.MEMORY)) { measurements.put( MetricType.MEMORY, tasksResourceUsages.stream().map(a -> a.getTaskResourceUsage().getMemoryInBytes()).mapToLong(Long::longValue).sum() diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java b/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java index 6a36c9aa..b467cbb5 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/QueryInsightsService.java @@ -123,17 +123,19 @@ public QueryInsightsService( * @param record the record to ingest */ public boolean addRecord(final SearchQueryRecord record) { - boolean shouldAdd = false; - for (Map.Entry entry : topQueriesServices.entrySet()) { - if (!enableCollect.get(entry.getKey())) { - continue; - } - List currentSnapshot = entry.getValue().getTopQueriesCurrentSnapshot(); - // skip add to top N queries store if the incoming record is smaller than the Nth record - if (currentSnapshot.size() < entry.getValue().getTopNSize() - || SearchQueryRecord.compare(record, currentSnapshot.get(0), entry.getKey()) > 0) { - shouldAdd = true; - break; + boolean shouldAdd = searchQueryMetricsEnabled; + if (!shouldAdd) { + for (Map.Entry entry : topQueriesServices.entrySet()) { + if (!enableCollect.get(entry.getKey())) { + continue; + } + List currentSnapshot = entry.getValue().getTopQueriesCurrentSnapshot(); + // skip add to top N queries store if the incoming record is smaller than the Nth record + if (currentSnapshot.size() < entry.getValue().getTopNSize() + || SearchQueryRecord.compare(record, currentSnapshot.get(0), entry.getKey()) > 0) { + shouldAdd = true; + break; + } } } if (shouldAdd) { diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryAggregationCategorizer.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryAggregationCategorizer.java index e72e5dec..256ecd6f 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryAggregationCategorizer.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryAggregationCategorizer.java @@ -38,7 +38,10 @@ public SearchQueryAggregationCategorizer(SearchQueryCounters searchQueryCounters * @param aggregatorFactories input aggregations * @param measurements latency, cpu, memory measurements */ - public void incrementSearchQueryAggregationCounters(Collection aggregatorFactories, Map measurements) { + public void incrementSearchQueryAggregationCounters( + Collection aggregatorFactories, + Map measurements + ) { for (AggregationBuilder aggregationBuilder : aggregatorFactories) { incrementCountersRecursively(aggregationBuilder, measurements); } diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java index ccf0e0f5..902e8cb6 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/categorizer/SearchQueryCounters.java @@ -103,6 +103,7 @@ public SearchQueryCounters(MetricsRegistry metricsRegistry) { * Increment counter * @param queryBuilder query builder * @param level level of query builder, 0 being highest level + * @param measurements metrics measurements */ public void incrementCounter(QueryBuilder queryBuilder, int level, Map measurements) { String uniqueQueryCounterName = queryBuilder.getName(); @@ -116,6 +117,7 @@ public void incrementCounter(QueryBuilder queryBuilder, int level, Map measurements) { aggCounter.add(value, tags); @@ -126,6 +128,7 @@ public void incrementAggCounter(double value, Tags tags, Map * Increment sort counter * @param value value to increment * @param tags tags + * @param measurements metrics measurements */ public void incrementSortCounter(double value, Tags tags, Map measurements) { sortCounter.add(value, tags); diff --git a/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java b/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java index 4b8d7a41..630e273c 100644 --- a/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java +++ b/src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java @@ -133,9 +133,9 @@ public void writeTo(final StreamOutput out) throws IOException { out.writeLong(timestamp); out.writeMap(measurements, (stream, metricType) -> MetricType.writeTo(out, metricType), StreamOutput::writeGenericValue); out.writeMap( - attributes, - (stream, attribute) -> Attribute.writeTo(out, attribute), - (stream, attributeValue) -> Attribute.writeValueTo(out, attributeValue) + attributes, + (stream, attribute) -> Attribute.writeTo(out, attribute), + (stream, attributeValue) -> Attribute.writeValueTo(out, attributeValue) ); } diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java index 443ca8b8..5bf54599 100644 --- a/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java +++ b/src/test/java/org/opensearch/plugin/insights/QueryInsightsPluginTests.java @@ -8,7 +8,6 @@ package org.opensearch.plugin.insights; -import org.junit.Before; import org.opensearch.action.ActionRequest; import org.opensearch.client.Client; import org.opensearch.cluster.service.ClusterService; @@ -30,6 +29,7 @@ import org.opensearch.threadpool.ExecutorBuilder; import org.opensearch.threadpool.ScalingExecutorBuilder; import org.opensearch.threadpool.ThreadPool; +import org.junit.Before; import java.util.Arrays; import java.util.List; diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java index a96533d9..f4590186 100644 --- a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java +++ b/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java @@ -37,8 +37,6 @@ import static java.util.Collections.emptyMap; import static java.util.Collections.emptySet; -import static org.junit.Assert.assertEquals; -import static org.junit.Assert.assertNull; import static org.opensearch.common.xcontent.XContentFactory.jsonBuilder; import static org.opensearch.test.OpenSearchTestCase.buildNewFakeTransportAddress; import static org.opensearch.test.OpenSearchTestCase.random; @@ -47,6 +45,8 @@ import static org.opensearch.test.OpenSearchTestCase.randomIntBetween; import static org.opensearch.test.OpenSearchTestCase.randomLong; import static org.opensearch.test.OpenSearchTestCase.randomLongBetween; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertNull; final public class QueryInsightsTestUtils { @@ -209,8 +209,8 @@ public static boolean checkRecordsEquals(List records1, List< return false; } else if (value instanceof Map && !Maps.deepEquals((Map) value, (Map) attributes2.get(attribute))) { - return false; - } + return false; + } } } return true; diff --git a/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java b/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java index 6db7ffea..558ddeec 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java @@ -8,17 +8,12 @@ package org.opensearch.plugin.insights.core.listener; -import org.junit.Before; -import org.mockito.ArgumentCaptor; -import org.mockito.MockitoAnnotations; -import org.mockito.MockitoSession; import org.opensearch.action.search.SearchPhaseContext; import org.opensearch.action.search.SearchRequest; import org.opensearch.action.search.SearchRequestContext; import org.opensearch.action.search.SearchTask; import org.opensearch.action.search.SearchType; import org.opensearch.action.support.replication.ClusterStateCreationUtils; -import org.opensearch.client.Client; import org.opensearch.cluster.ClusterState; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.collect.Tuple; @@ -37,12 +32,11 @@ import org.opensearch.search.aggregations.support.ValueType; import org.opensearch.search.builder.SearchSourceBuilder; import org.opensearch.tasks.Task; -import org.opensearch.telemetry.metrics.MetricsRegistry; -import org.opensearch.telemetry.metrics.tags.Tags; import org.opensearch.test.ClusterServiceUtils; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; +import org.junit.Before; import java.util.ArrayList; import java.util.Collections; @@ -54,11 +48,11 @@ import java.util.concurrent.Phaser; import java.util.concurrent.TimeUnit; +import org.mockito.ArgumentCaptor; + import static org.mockito.ArgumentMatchers.any; -import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; -import static org.mockito.Mockito.spy; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -115,7 +109,7 @@ public void testOnRequestEnd() throws InterruptedException { Collections.singletonMap(Task.X_OPAQUE_ID, "userLabel") ); - String[] indices = new String[]{"index-1", "index-2"}; + String[] indices = new String[] { "index-1", "index-2" }; Map phaseLatencyMap = new HashMap<>(); phaseLatencyMap.put("expand", 0L); @@ -164,7 +158,7 @@ public void testConcurrentOnRequestEnd() throws InterruptedException { Collections.singletonMap(Task.X_OPAQUE_ID, "userLabel") ); - String[] indices = new String[]{"index-1", "index-2"}; + String[] indices = new String[] { "index-1", "index-2" }; Map phaseLatencyMap = new HashMap<>(); phaseLatencyMap.put("expand", 0L); @@ -252,7 +246,6 @@ public void testTopNFeatureEnabledDisabled() { verify(queryInsightsService3, never()).checkAndStopQueryInsights(); verify(queryInsightsService3, never()).checkAndRestartQueryInsights(); - // Test case 4: Only CPU enabled initially, enable latency and verify expected behavior QueryInsightsService queryInsightsService4 = mock(QueryInsightsService.class); QueryInsightsListener queryInsightsListener4 = new QueryInsightsListener(clusterService, queryInsightsService4); diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java index cd7a7a72..2ba03ec1 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/service/QueryInsightsServiceTests.java @@ -8,7 +8,6 @@ package org.opensearch.plugin.insights.core.service; -import org.junit.Before; import org.opensearch.client.Client; import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Settings; @@ -19,12 +18,11 @@ import org.opensearch.telemetry.metrics.noop.NoopMetricsRegistry; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.ThreadPool; +import org.junit.Before; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.spy; -import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.verifyNoMoreInteractions; /** * Unit Tests for {@link QueryInsightsService}. diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/TopQueriesServiceTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/TopQueriesServiceTests.java index 396e1822..8478fe16 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/TopQueriesServiceTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/service/TopQueriesServiceTests.java @@ -75,22 +75,16 @@ public void testSmallNSize() { } public void testValidateTopNSize() { - assertThrows(IllegalArgumentException.class, () -> { - topQueriesService.validateTopNSize(QueryInsightsSettings.MAX_N_SIZE + 1); - }); + assertThrows(IllegalArgumentException.class, () -> { topQueriesService.validateTopNSize(QueryInsightsSettings.MAX_N_SIZE + 1); }); } public void testValidateNegativeTopNSize() { - assertThrows(IllegalArgumentException.class, () -> { - topQueriesService.validateTopNSize(-1); - }); + assertThrows(IllegalArgumentException.class, () -> { topQueriesService.validateTopNSize(-1); }); } public void testGetTopQueriesWhenNotEnabled() { topQueriesService.setEnabled(false); - assertThrows(IllegalArgumentException.class, () -> { - topQueriesService.getTopQueriesRecords(false); - }); + assertThrows(IllegalArgumentException.class, () -> { topQueriesService.getTopQueriesRecords(false); }); } public void testValidateWindowSize() { @@ -100,12 +94,8 @@ public void testValidateWindowSize() { assertThrows(IllegalArgumentException.class, () -> { topQueriesService.validateWindowSize(new TimeValue(QueryInsightsSettings.MIN_WINDOW_SIZE.getSeconds() - 1, TimeUnit.SECONDS)); }); - assertThrows(IllegalArgumentException.class, () -> { - topQueriesService.validateWindowSize(new TimeValue(2, TimeUnit.DAYS)); - }); - assertThrows(IllegalArgumentException.class, () -> { - topQueriesService.validateWindowSize(new TimeValue(7, TimeUnit.MINUTES)); - }); + assertThrows(IllegalArgumentException.class, () -> { topQueriesService.validateWindowSize(new TimeValue(2, TimeUnit.DAYS)); }); + assertThrows(IllegalArgumentException.class, () -> { topQueriesService.validateWindowSize(new TimeValue(7, TimeUnit.MINUTES)); }); } private static void runUntilTimeoutOrFinish(DeterministicTaskQueue deterministicTaskQueue, long duration) { diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java index 7ce2499e..3a7f87a4 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/service/categorizor/SearchQueryCategorizerTests.java @@ -8,9 +8,6 @@ package org.opensearch.plugin.insights.core.service.categorizor; -import org.junit.After; -import org.mockito.invocation.InvocationOnMock; -import org.mockito.stubbing.Answer; import org.opensearch.index.query.BoolQueryBuilder; import org.opensearch.index.query.BoostingQueryBuilder; import org.opensearch.index.query.MatchNoneQueryBuilder; @@ -37,6 +34,7 @@ import org.opensearch.telemetry.metrics.MetricsRegistry; import org.opensearch.telemetry.metrics.tags.Tags; import org.opensearch.test.OpenSearchTestCase; +import org.junit.After; import org.junit.Before; import java.util.Arrays; @@ -44,14 +42,16 @@ import java.util.Map; import org.mockito.ArgumentCaptor; +import org.mockito.invocation.InvocationOnMock; +import org.mockito.stubbing.Answer; +import static org.opensearch.plugin.insights.QueryInsightsTestUtils.generateQueryInsightRecords; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; -import static org.opensearch.plugin.insights.QueryInsightsTestUtils.generateQueryInsightRecords; public final class SearchQueryCategorizerTests extends OpenSearchTestCase { @@ -71,19 +71,18 @@ public void setup() { invocation -> mock(Counter.class) ); - when(metricsRegistry.createHistogram(any(String.class), any(String.class), any(String.class))) - .thenAnswer(new Answer() { - @Override - public Histogram answer(InvocationOnMock invocation) throws Throwable { - // Extract arguments to identify which histogram is being created - String name = invocation.getArgument(0); - // Create a mock histogram - Histogram histogram = mock(Histogram.class); - // Store histogram in map for lookup - histogramMap.put(name, histogram); - return histogram; - } - }); + when(metricsRegistry.createHistogram(any(String.class), any(String.class), any(String.class))).thenAnswer(new Answer() { + @Override + public Histogram answer(InvocationOnMock invocation) throws Throwable { + // Extract arguments to identify which histogram is being created + String name = invocation.getArgument(0); + // Create a mock histogram + Histogram histogram = mock(Histogram.class); + // Store histogram in map for lookup + histogramMap.put(name, histogram); + return histogram; + } + }); searchQueryCategorizer = SearchQueryCategorizer.getInstance(metricsRegistry); } @@ -327,7 +326,6 @@ private void verifyMeasurementHistogramsIncremented(SearchQueryRecord record, in Histogram queryTypeCpuHistogram = histogramMap.get("search.query.type.cpu.histogram"); Histogram queryTypeMemoryHistogram = histogramMap.get("search.query.type.memory.histogram"); - verify(queryTypeLatencyHistogram, times(times)).record(eq(expectedLatency), any(Tags.class)); verify(queryTypeCpuHistogram, times(times)).record(eq(expectedCpu), any(Tags.class)); verify(queryTypeMemoryHistogram, times(times)).record(eq(expectedMemory), any(Tags.class)); diff --git a/src/yamlRestTest/java/org/opensearch/plugin/insights/QueryInsightsClientYamlTestSuiteIT.java b/src/yamlRestTest/java/org/opensearch/plugin/insights/QueryInsightsClientYamlTestSuiteIT.java index 2860acb3..488e73d0 100644 --- a/src/yamlRestTest/java/org/opensearch/plugin/insights/QueryInsightsClientYamlTestSuiteIT.java +++ b/src/yamlRestTest/java/org/opensearch/plugin/insights/QueryInsightsClientYamlTestSuiteIT.java @@ -9,10 +9,10 @@ import com.carrotsearch.randomizedtesting.annotations.Name; import com.carrotsearch.randomizedtesting.annotations.ParametersFactory; + import org.opensearch.test.rest.yaml.ClientYamlTestCandidate; import org.opensearch.test.rest.yaml.OpenSearchClientYamlSuiteTestCase; - public class QueryInsightsClientYamlTestSuiteIT extends OpenSearchClientYamlSuiteTestCase { public QueryInsightsClientYamlTestSuiteIT(@Name("yaml") ClientYamlTestCandidate testCandidate) {