sync bug fixes from core to the plugin repo

Signed-off-by: Chenyang Ji <[email protected]>

.editorconfig

@@ -0,0 +1,33 @@
+# EditorConfig: http://editorconfig.org/
+
+root = true
+
+[*]
+charset = utf-8
+trim_trailing_whitespace = true
+insert_final_newline = true
+indent_style = space
+
+[*.gradle]
+indent_size = 2
+
+[*.groovy]
+indent_size = 4
+
+[*.java]
+indent_size = 4
+
+[*.json]
+indent_size = 2
+
+[*.py]
+indent_size = 2
+
+[*.sh]
+indent_size = 2
+
+[*.{yml,yaml}]
+indent_size = 2
+
+[*.{xsd,xml}]
+indent_size = 4

src/main/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListener.java

@@ -52,7 +52,7 @@ public final class QueryInsightsListener extends SearchRequestOperationsListener
     /**
      * Constructor for QueryInsightsListener
      *
-     * @param clusterService The Node's cluster service.
+     * @param clusterService       The Node's cluster service.
      * @param queryInsightsService The topQueriesByLatencyService associated with this listener
      */
     @Inject
@@ -91,7 +91,7 @@ public QueryInsightsListener(final ClusterService clusterService, final QueryIns
      * and query insights services.
      *
      * @param metricType {@link MetricType}
-     * @param enabled boolean
+     * @param enabled    boolean
      */
     public void setEnableTopQueries(final MetricType metricType, final boolean enabled) {
         boolean isAllMetricsDisabled = !queryInsightsService.isEnabled();
@@ -119,16 +119,20 @@ public boolean isEnabled() {
     }
 
     @Override
-    public void onPhaseStart(SearchPhaseContext context) {}
+    public void onPhaseStart(SearchPhaseContext context) {
+    }
 
     @Override
-    public void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {}
+    public void onPhaseEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+    }
 
     @Override
-    public void onPhaseFailure(SearchPhaseContext context, Throwable cause) {}
+    public void onPhaseFailure(SearchPhaseContext context, Throwable cause) {
+    }
 
     @Override
-    public void onRequestStart(SearchRequestContext searchRequestContext) {}
+    public void onRequestStart(SearchRequestContext searchRequestContext) {
+    }
 
     @Override
     public void onRequestEnd(final SearchPhaseContext context, final SearchRequestContext searchRequestContext) {
@@ -143,6 +147,9 @@ public void onRequestFailure(final SearchPhaseContext context, final SearchReque
     private void constructSearchQueryRecord(final SearchPhaseContext context, final SearchRequestContext searchRequestContext) {
         SearchTask searchTask = context.getTask();
         List<TaskResourceInfo> tasksResourceUsages = searchRequestContext.getPhaseResourceUsage();
+        if (clusterService.getTaskResourceTrackingService() != null) {
+            clusterService.getTaskResourceTrackingService().refreshResourceStats(searchTask);
+        }
         tasksResourceUsages.add(
             new TaskResourceInfo(
                 searchTask.getAction(),

src/main/java/org/opensearch/plugin/insights/core/service/TopQueriesService.java

@@ -136,24 +136,23 @@ public int getTopNSize() {
      * @param size the wanted top N size
      */
     public void validateTopNSize(final int size) {
-        if (size > QueryInsightsSettings.MAX_N_SIZE) {
+        if (size < 1 || size > QueryInsightsSettings.MAX_N_SIZE) {
             throw new IllegalArgumentException(
                 "Top N size setting for ["
                     + metricType
                     + "]"
-                    + " should be smaller than max top N size ["
+                    + " should be between 1 and "
                     + QueryInsightsSettings.MAX_N_SIZE
-                    + "was ("
+                    + ", was ("
                     + size
-                    + " > "
-                    + QueryInsightsSettings.MAX_N_SIZE
                     + ")"
             );
         }
     }
 
     /**
      * Set enable flag for the service
+     *
      * @param enabled boolean
      */
     public void setEnabled(final boolean enabled) {
@@ -251,7 +250,7 @@ public void validateExporterConfig(Settings settings) {
     /**
      * Get all top queries records that are in the current top n queries store
      * Optionally include top N records from the last window.
-     *
+     * <p>
      * By default, return the records in sorted order.
      *
      * @param includeLastWindow if the top N queries from the last window should be included

src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java

@@ -8,11 +8,18 @@
 
 package org.opensearch.plugin.insights.rules.model;
 
+import org.apache.lucene.util.ArrayUtil;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
+import org.opensearch.core.common.io.stream.Writeable;
+import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
 import java.util.Locale;
+import java.util.Map;
 
 /**
  * Valid attributes for a search query record
@@ -65,14 +72,78 @@ static Attribute readFromStream(final StreamInput in) throws IOException {
     /**
      * Write Attribute to a StreamOutput
      *
-     * @param out the StreamOutput to write
+     * @param out       the StreamOutput to write
      * @param attribute the Attribute to write
      * @throws IOException IOException
      */
     static void writeTo(final StreamOutput out, final Attribute attribute) throws IOException {
         out.writeString(attribute.toString());
     }
 
+    /**
+     * Write Attribute value to a StreamOutput
+     *
+     * @param out            the StreamOutput to write
+     * @param attributeValue the Attribute value to write
+     */
+    @SuppressWarnings("unchecked")
+    public static void writeValueTo(StreamOutput out, Object attributeValue) throws IOException {
+        if (attributeValue instanceof List) {
+            out.writeList((List<? extends Writeable>) attributeValue);
+        } else {
+            out.writeGenericValue(attributeValue);
+        }
+    }
+
+    /**
+     * Read attribute value from the input stream given the Attribute type
+     *
+     * @param in        the {@link StreamInput} input to read
+     * @param attribute attribute type to differentiate between Source and others
+     * @return parse value
+     * @throws IOException IOException
+     */
+    public static Object readAttributeValue(StreamInput in, Attribute attribute) throws IOException {
+        if (attribute == Attribute.TASK_RESOURCE_USAGES) {
+            return in.readList(TaskResourceInfo::readFromStream);
+        } else {
+            return in.readGenericValue();
+        }
+    }
+
+    /**
+     * Read attribute map from the input stream
+     *
+     * @param in the {@link StreamInput} to read
+     * @return parsed attribute map
+     * @throws IOException IOException
+     */
+    public static Map<Attribute, Object> readAttributeMap(StreamInput in) throws IOException {
+        int size = readArraySize(in);
+        if (size == 0) {
+            return Collections.emptyMap();
+        }
+        Map<Attribute, Object> map = new HashMap<>(size);
+
+        for (int i = 0; i < size; i++) {
+            Attribute key = readFromStream(in);
+            Object value = readAttributeValue(in, key);
+            map.put(key, value);
+        }
+        return map;
+    }
+
+    private static int readArraySize(StreamInput in) throws IOException {
+        final int arraySize = in.readVInt();
+        if (arraySize > ArrayUtil.MAX_ARRAY_LENGTH) {
+            throw new IllegalStateException("array length must be <= to " + ArrayUtil.MAX_ARRAY_LENGTH + " but was: " + arraySize);
+        }
+        if (arraySize < 0) {
+            throw new NegativeArraySizeException("array size must be positive but was: " + arraySize);
+        }
+        return arraySize;
+    }
+
     @Override
     public String toString() {
         return this.name().toLowerCase(Locale.ROOT);

src/main/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecord.java

@@ -43,7 +43,7 @@ public SearchQueryRecord(final StreamInput in) throws IOException, ClassCastExce
         measurements = new HashMap<>();
         in.readMap(MetricType::readFromStream, StreamInput::readGenericValue)
             .forEach(((metricType, o) -> measurements.put(metricType, metricType.parseValue(o))));
-        this.attributes = in.readMap(Attribute::readFromStream, StreamInput::readGenericValue);
+        this.attributes = Attribute.readAttributeMap(in);
     }
 
     /**
@@ -132,7 +132,11 @@ public XContentBuilder toXContent(final XContentBuilder builder, final ToXConten
     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), StreamOutput::writeGenericValue);
+        out.writeMap(
+                attributes,
+                (stream, attribute) -> Attribute.writeTo(out, attribute),
+                (stream, attributeValue) -> Attribute.writeValueTo(out, attributeValue)
+        );
     }
 
     /**

src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java

@@ -12,6 +12,8 @@
 import org.opensearch.cluster.node.DiscoveryNode;
 import org.opensearch.common.settings.ClusterSettings;
 import org.opensearch.common.util.Maps;
+import org.opensearch.core.tasks.resourcetracker.TaskResourceInfo;
+import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage;
 import org.opensearch.core.xcontent.ToXContent;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.plugin.insights.rules.action.top_queries.TopQueries;
@@ -46,7 +48,8 @@
 
 final public class QueryInsightsTestUtils {
 
-    public QueryInsightsTestUtils() {}
+    public QueryInsightsTestUtils() {
+    }
 
     public static List<SearchQueryRecord> generateQueryInsightRecords(int count) {
         return generateQueryInsightRecords(count, count, System.currentTimeMillis(), 0);
@@ -80,6 +83,25 @@ public static List<SearchQueryRecord> generateQueryInsightRecords(int lower, int
             attributes.put(Attribute.TOTAL_SHARDS, randomIntBetween(1, 100));
             attributes.put(Attribute.INDICES, randomArray(1, 3, Object[]::new, () -> randomAlphaOfLengthBetween(5, 10)));
             attributes.put(Attribute.PHASE_LATENCY_MAP, phaseLatencyMap);
+            attributes.put(
+                Attribute.TASK_RESOURCE_USAGES,
+                List.of(
+                    new TaskResourceInfo(
+                        randomAlphaOfLengthBetween(5, 10),
+                        randomLongBetween(1, 1000),
+                        randomLongBetween(1, 1000),
+                        randomAlphaOfLengthBetween(5, 10),
+                        new TaskResourceUsage(randomLongBetween(1, 1000), randomLongBetween(1, 1000))
+                    ),
+                    new TaskResourceInfo(
+                        randomAlphaOfLengthBetween(5, 10),
+                        randomLongBetween(1, 1000),
+                        randomLongBetween(1, 1000),
+                        randomAlphaOfLengthBetween(5, 10),
+                        new TaskResourceUsage(randomLongBetween(1, 1000), randomLongBetween(1, 1000))
+                    )
+                )
+            );
 
             records.add(new SearchQueryRecord(timestamp, measurements, attributes));
             timestamp += interval;
@@ -163,8 +185,8 @@ public static boolean checkRecordsEquals(List<SearchQueryRecord> records1, List<
                     return false;
                 } else if (value instanceof Map
                     && !Maps.deepEquals((Map<Object, Object>) value, (Map<Object, Object>) attributes2.get(attribute))) {
-                        return false;
-                    }
+                    return false;
+                }
             }
         }
         return true;

src/test/java/org/opensearch/plugin/insights/core/listener/QueryInsightsListenerTests.java

@@ -32,6 +32,7 @@
 import org.opensearch.search.aggregations.support.ValueType;
 import org.opensearch.search.builder.SearchSourceBuilder;
 import org.opensearch.tasks.Task;
+import org.opensearch.tasks.TaskResourceTrackingService;
 import org.opensearch.test.ClusterServiceUtils;
 import org.opensearch.test.OpenSearchTestCase;
 import org.opensearch.threadpool.TestThreadPool;
@@ -65,6 +66,7 @@ public class QueryInsightsListenerTests extends OpenSearchTestCase {
     private final SearchRequest searchRequest = mock(SearchRequest.class);
     private final QueryInsightsService queryInsightsService = mock(QueryInsightsService.class);
     private final TopQueriesService topQueriesService = mock(TopQueriesService.class);
+    private final TaskResourceTrackingService taskResourceTrackingService = mock(TaskResourceTrackingService.class);
     private final ThreadPool threadPool = new TestThreadPool("QueryInsightsThreadPool");
     private ClusterService clusterService;
 
@@ -77,6 +79,7 @@ public void setup() {
         ClusterState state = ClusterStateCreationUtils.stateWithActivePrimary("test", true, 1 + randomInt(3), randomInt(2));
         clusterService = ClusterServiceUtils.createClusterService(threadPool, state.getNodes().getLocalNode(), clusterSettings);
         ClusterServiceUtils.setState(clusterService, state);
+        clusterService.setTaskResourceTrackingService(taskResourceTrackingService);
         when(queryInsightsService.isCollectionEnabled(MetricType.LATENCY)).thenReturn(true);
         when(queryInsightsService.getTopQueriesService(MetricType.LATENCY)).thenReturn(topQueriesService);
 
@@ -108,7 +111,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<String, Long> phaseLatencyMap = new HashMap<>();
         phaseLatencyMap.put("expand", 0L);
@@ -139,6 +142,7 @@ public void testOnRequestEnd() throws InterruptedException {
         assertEquals(searchSourceBuilder.toString(), generatedRecord.getAttributes().get(Attribute.SOURCE));
         Map<String, String> labels = (Map<String, String>) generatedRecord.getAttributes().get(Attribute.LABELS);
         assertEquals("userLabel", labels.get(Task.X_OPAQUE_ID));
+        verify(taskResourceTrackingService, times(1)).refreshResourceStats(task);
     }
 
     public void testConcurrentOnRequestEnd() throws InterruptedException {
@@ -157,7 +161,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<String, Long> phaseLatencyMap = new HashMap<>();
         phaseLatencyMap.put("expand", 0L);
@@ -200,6 +204,7 @@ public void testConcurrentOnRequestEnd() throws InterruptedException {
         countDownLatch.await();
 
         verify(queryInsightsService, times(numRequests)).addRecord(any());
+        verify(taskResourceTrackingService, times(numRequests)).refreshResourceStats(task);
     }
 
     public void testSetEnabled() {

src/test/java/org/opensearch/plugin/insights/core/service/TopQueriesServiceTests.java

@@ -75,12 +75,22 @@ 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);
+        });
     }
 
     public void testGetTopQueriesWhenNotEnabled() {
         topQueriesService.setEnabled(false);
-        assertThrows(IllegalArgumentException.class, () -> { topQueriesService.getTopQueriesRecords(false); });
+        assertThrows(IllegalArgumentException.class, () -> {
+            topQueriesService.getTopQueriesRecords(false);
+        });
     }
 
     public void testValidateWindowSize() {
@@ -90,8 +100,12 @@ 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) {