From 1d991b8d67336ebc1bc219c9084617bf2a35d934 Mon Sep 17 00:00:00 2001 From: Siddhant Deshmukh Date: Thu, 16 Jan 2025 15:19:30 -0800 Subject: [PATCH] Model changes for hashcode and id (#191) Signed-off-by: Siddhant Deshmukh --- .../core/listener/QueryInsightsListener.java | 2 +- .../grouper/MinMaxHeapQueryGrouper.java | 2 +- .../insights/rules/model/Attribute.java | 4 +- .../rules/model/SearchQueryRecord.java | 43 +++++++++++++-- .../insights/QueryInsightsTestUtils.java | 14 ++--- .../grouper/MinMaxHeapQueryGrouperTests.java | 4 +- .../top_queries/TopQueriesResponseTests.java | 55 +++++++++++++++++-- .../rules/model/SearchQueryRecordTests.java | 10 ++-- 8 files changed, 105 insertions(+), 29 deletions(-) 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 84e8463e..e2c6255f 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 @@ -288,7 +288,7 @@ private void constructSearchQueryRecord(final SearchPhaseContext context, final // Add hashcode attribute when grouping is enabled if (queryInsightsService.isGroupingEnabled()) { String hashcode = queryShapeGenerator.getShapeHashCodeAsString(queryShape); - attributes.put(Attribute.ID, hashcode); + attributes.put(Attribute.QUERY_GROUP_HASHCODE, hashcode); } } diff --git a/src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java b/src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java index 56fc0eb9..33bbecb6 100644 --- a/src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java +++ b/src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java @@ -271,7 +271,7 @@ int numberOfTopGroups() { private String getGroupingId(final SearchQueryRecord searchQueryRecord) { switch (groupingType) { case SIMILARITY: - return searchQueryRecord.getAttributes().get(Attribute.ID).toString(); + return searchQueryRecord.getAttributes().get(Attribute.QUERY_GROUP_HASHCODE).toString(); case NONE: throw new IllegalArgumentException("Should not try to group queries if grouping type is NONE"); default: diff --git a/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java b/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java index c24a97f9..5cf233a3 100644 --- a/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java +++ b/src/main/java/org/opensearch/plugin/insights/rules/model/Attribute.java @@ -58,9 +58,9 @@ public enum Attribute { */ LABELS, /** - * Query Group hashcode or query hashcode representing a unique identifier for the query/group + * Query Group hashcode */ - ID, + QUERY_GROUP_HASHCODE, /** * Grouping type of the query record (none, similarity) */ 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 1fcf8b0f..93113d08 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 @@ -15,6 +15,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.UUID; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; import org.opensearch.Version; @@ -41,6 +42,8 @@ public class SearchQueryRecord implements ToXContentObject, Writeable { private final long timestamp; private final Map measurements; private final Map attributes; + private final String id; + /** * Timestamp */ @@ -93,12 +96,16 @@ public class SearchQueryRecord implements ToXContentObject, Writeable { * Grouping type of the query record (none, similarity) */ public static final String GROUP_BY = "group_by"; - /** - * Query Group hashcode or query hashcode representing a unique identifier for the query/group + * UUID */ public static final String ID = "id"; + /** + * Query Group hashcode + */ + public static final String QUERY_GROUP_HASHCODE = "query_group_hashcode"; + public static final String MEASUREMENTS = "measurements"; private String groupingId; @@ -111,6 +118,7 @@ public class SearchQueryRecord implements ToXContentObject, Writeable { */ public SearchQueryRecord(final StreamInput in) throws IOException, ClassCastException { this.timestamp = in.readLong(); + this.id = in.readString(); if (in.getVersion().onOrAfter(Version.V_2_17_0)) { measurements = new LinkedHashMap<>(); in.readOrderedMap(MetricType::readFromStream, Measurement::readFromStream) @@ -137,12 +145,30 @@ public SearchQueryRecord(final StreamInput in) throws IOException, ClassCastExce * @param attributes A list of Attributes associated with this query */ public SearchQueryRecord(final long timestamp, Map measurements, final Map attributes) { + this(timestamp, measurements, attributes, UUID.randomUUID().toString()); + } + + /** + * Constructor of SearchQueryRecord + * + * @param timestamp The timestamp of the query. + * @param measurements A list of Measurement associated with this query + * @param attributes A list of Attributes associated with this query + * @param id unique id for a search query record + */ + public SearchQueryRecord( + final long timestamp, + Map measurements, + final Map attributes, + String id + ) { if (measurements == null) { throw new IllegalArgumentException("Measurements cannot be null"); } this.measurements = measurements; this.attributes = attributes; this.timestamp = timestamp; + this.id = id; } /** @@ -156,6 +182,7 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc long timestamp = 0L; Map measurements = new HashMap<>(); Map attributes = new HashMap<>(); + String id = null; parser.nextToken(); XContentParserUtils.ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); @@ -167,6 +194,9 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc case TIMESTAMP: timestamp = parser.longValue(); break; + case ID: + id = parser.text(); + break; case LATENCY: case CPU: case MEMORY: @@ -179,8 +209,8 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc case GROUP_BY: attributes.put(Attribute.GROUP_BY, parser.text()); break; - case ID: - attributes.put(Attribute.ID, parser.text()); + case QUERY_GROUP_HASHCODE: + attributes.put(Attribute.QUERY_GROUP_HASHCODE, parser.text()); break; case SOURCE: XContentParserUtils.ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser); @@ -264,7 +294,7 @@ public static SearchQueryRecord fromXContent(XContentParser parser) throws IOExc log.error("Error when parsing through search hit", e); } } - return new SearchQueryRecord(timestamp, measurements, attributes); + return new SearchQueryRecord(timestamp, measurements, attributes, id); } /** @@ -337,6 +367,8 @@ public void addAttribute(final Attribute attribute, final Object value) { public XContentBuilder toXContent(final XContentBuilder builder, final Params params) throws IOException { builder.startObject(); builder.field("timestamp", timestamp); + builder.field("id", id); + for (Map.Entry entry : attributes.entrySet()) { builder.field(entry.getKey().toString(), entry.getValue()); } @@ -358,6 +390,7 @@ public XContentBuilder toXContent(final XContentBuilder builder, final Params pa @Override public void writeTo(final StreamOutput out) throws IOException { out.writeLong(timestamp); + out.writeString(id); if (out.getVersion().onOrAfter(Version.V_2_17_0)) { out.writeMap( measurements, diff --git a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java b/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java index fe9fd515..84f4e646 100644 --- a/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java +++ b/src/test/java/org/opensearch/plugin/insights/QueryInsightsTestUtils.java @@ -139,7 +139,7 @@ public static List generateQueryInsightRecords( 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.ID, Objects.hashCode(i)); + attributes.put(Attribute.QUERY_GROUP_HASHCODE, Objects.hashCode(i)); attributes.put(Attribute.GROUP_BY, GroupingType.NONE); attributes.put( Attribute.TASK_RESOURCE_USAGES, @@ -200,13 +200,13 @@ public static List> generateMultipleQueryInsightsRecords public static void populateSameQueryHashcodes(List searchQueryRecords) { for (SearchQueryRecord record : searchQueryRecords) { - record.getAttributes().put(Attribute.ID, 1); + record.getAttributes().put(Attribute.QUERY_GROUP_HASHCODE, 1); } } public static void populateHashcode(List searchQueryRecords, int hash) { for (SearchQueryRecord record : searchQueryRecords) { - record.getAttributes().put(Attribute.ID, hash); + record.getAttributes().put(Attribute.QUERY_GROUP_HASHCODE, hash); } } @@ -223,7 +223,7 @@ public static TopQueries createRandomTopQueries() { return new TopQueries(node, records); } - public static TopQueries createFixedTopQueries() { + public static TopQueries createFixedTopQueries(String id) { DiscoveryNode node = new DiscoveryNode( "node_for_top_queries_test", buildNewFakeTransportAddress(), @@ -232,12 +232,12 @@ public static TopQueries createFixedTopQueries() { VersionUtils.randomVersion(random()) ); List records = new ArrayList<>(); - records.add(createFixedSearchQueryRecord()); + records.add(createFixedSearchQueryRecord(id)); return new TopQueries(node, records); } - public static SearchQueryRecord createFixedSearchQueryRecord() { + public static SearchQueryRecord createFixedSearchQueryRecord(String id) { long timestamp = 1706574180000L; Map measurements = Map.of(MetricType.LATENCY, new Measurement(1L)); @@ -256,7 +256,7 @@ public static SearchQueryRecord createFixedSearchQueryRecord() { ) ); - return new SearchQueryRecord(timestamp, measurements, attributes); + return new SearchQueryRecord(timestamp, measurements, attributes, id); } public static void compareJson(ToXContent param1, ToXContent param2) throws IOException { diff --git a/src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouperTests.java b/src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouperTests.java index a87d1035..d5ca7bff 100644 --- a/src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouperTests.java +++ b/src/test/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouperTests.java @@ -44,7 +44,7 @@ public void testWithAllDifferentHashcodes() { Set hashcodeSet = new HashSet<>(); for (SearchQueryRecord record : records) { groupedRecord = minMaxHeapQueryGrouper.add(record); - int hashcode = (int) groupedRecord.getAttributes().get(Attribute.ID); + int hashcode = (int) groupedRecord.getAttributes().get(Attribute.QUERY_GROUP_HASHCODE); hashcodeSet.add(hashcode); } assertEquals(numOfRecords, hashcodeSet.size()); @@ -58,7 +58,7 @@ public void testWithAllSameHashcodes() { Set hashcodeSet = new HashSet<>(); for (SearchQueryRecord record : records) { groupedRecord = minMaxHeapQueryGrouper.add(record); - int hashcode = (int) groupedRecord.getAttributes().get(Attribute.ID); + int hashcode = (int) groupedRecord.getAttributes().get(Attribute.QUERY_GROUP_HASHCODE); hashcodeSet.add(hashcode); } assertEquals(1, hashcodeSet.size()); diff --git a/src/test/java/org/opensearch/plugin/insights/rules/action/top_queries/TopQueriesResponseTests.java b/src/test/java/org/opensearch/plugin/insights/rules/action/top_queries/TopQueriesResponseTests.java index e8b5ca9c..fdfa5fb2 100644 --- a/src/test/java/org/opensearch/plugin/insights/rules/action/top_queries/TopQueriesResponseTests.java +++ b/src/test/java/org/opensearch/plugin/insights/rules/action/top_queries/TopQueriesResponseTests.java @@ -40,18 +40,61 @@ public void testSerialize() throws Exception { } public void testToXContent() throws IOException { - char[] expectedXcontent = - "{\"top_queries\":[{\"timestamp\":1706574180000,\"node_id\":\"node_for_top_queries_test\",\"phase_latency_map\":{\"expand\":1,\"query\":10,\"fetch\":1},\"task_resource_usages\":[{\"action\":\"action\",\"taskId\":2,\"parentTaskId\":1,\"nodeId\":\"id\",\"taskResourceUsage\":{\"cpu_time_in_nanos\":1000,\"memory_in_bytes\":2000}},{\"action\":\"action2\",\"taskId\":3,\"parentTaskId\":1,\"nodeId\":\"id2\",\"taskResourceUsage\":{\"cpu_time_in_nanos\":2000,\"memory_in_bytes\":1000}}],\"search_type\":\"query_then_fetch\",\"measurements\":{\"latency\":{\"number\":1,\"count\":1,\"aggregationType\":\"NONE\"}}}]}" - .toCharArray(); - TopQueries topQueries = QueryInsightsTestUtils.createFixedTopQueries(); + String id = "sample_id"; + + char[] expectedXContent = ("{" + + "\"top_queries\":[{" + + "\"timestamp\":1706574180000," + + "\"node_id\":\"node_for_top_queries_test\"," + + "\"phase_latency_map\":{" + + "\"expand\":1," + + "\"query\":10," + + "\"fetch\":1" + + "}," + + "\"task_resource_usages\":[{" + + "\"action\":\"action\"," + + "\"taskId\":2," + + "\"parentTaskId\":1," + + "\"nodeId\":\"id\"," + + "\"taskResourceUsage\":{" + + "\"cpu_time_in_nanos\":1000," + + "\"memory_in_bytes\":2000" + + "}" + + "},{" + + "\"action\":\"action2\"," + + "\"taskId\":3," + + "\"parentTaskId\":1," + + "\"nodeId\":\"id2\"," + + "\"taskResourceUsage\":{" + + "\"cpu_time_in_nanos\":2000," + + "\"memory_in_bytes\":1000" + + "}" + + "}]," + + "\"search_type\":\"query_then_fetch\"," + + "\"measurements\":{" + + "\"latency\":{" + + "\"number\":1," + + "\"count\":1," + + "\"aggregationType\":\"NONE\"" + + "}" + + "}," + + "\"id\":\"" + + id + + "\"" + + "}]" + + "}").toCharArray(); + + TopQueries topQueries = QueryInsightsTestUtils.createFixedTopQueries(id); ClusterName clusterName = new ClusterName("test-cluster"); TopQueriesResponse response = new TopQueriesResponse(clusterName, List.of(topQueries), new ArrayList<>(), 10, MetricType.LATENCY); + XContentBuilder builder = MediaTypeRegistry.contentBuilder(MediaTypeRegistry.JSON); char[] xContent = BytesReference.bytes(response.toXContent(builder, ToXContent.EMPTY_PARAMS)).utf8ToString().toCharArray(); - Arrays.sort(expectedXcontent); + + Arrays.sort(expectedXContent); Arrays.sort(xContent); - assertEquals(Arrays.hashCode(expectedXcontent), Arrays.hashCode(xContent)); + assertEquals(Arrays.hashCode(expectedXContent), Arrays.hashCode(xContent)); } /** diff --git a/src/test/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecordTests.java b/src/test/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecordTests.java index 17b442f5..eefce983 100644 --- a/src/test/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecordTests.java +++ b/src/test/java/org/opensearch/plugin/insights/rules/model/SearchQueryRecordTests.java @@ -45,19 +45,19 @@ public void testAllMetricTypes() { } public void testCompare() { - SearchQueryRecord record1 = QueryInsightsTestUtils.createFixedSearchQueryRecord(); - SearchQueryRecord record2 = QueryInsightsTestUtils.createFixedSearchQueryRecord(); + SearchQueryRecord record1 = QueryInsightsTestUtils.createFixedSearchQueryRecord("id"); + SearchQueryRecord record2 = QueryInsightsTestUtils.createFixedSearchQueryRecord("id"); assertEquals(0, SearchQueryRecord.compare(record1, record2, MetricType.LATENCY)); } public void testEqual() { - SearchQueryRecord record1 = QueryInsightsTestUtils.createFixedSearchQueryRecord(); - SearchQueryRecord record2 = QueryInsightsTestUtils.createFixedSearchQueryRecord(); + SearchQueryRecord record1 = QueryInsightsTestUtils.createFixedSearchQueryRecord("id"); + SearchQueryRecord record2 = QueryInsightsTestUtils.createFixedSearchQueryRecord("id"); assertEquals(record1, record2); } public void testFromXContent() { - SearchQueryRecord record = QueryInsightsTestUtils.createFixedSearchQueryRecord(); + SearchQueryRecord record = QueryInsightsTestUtils.createFixedSearchQueryRecord("id"); try (XContentParser recordParser = createParser(JsonXContent.jsonXContent, record.toString())) { SearchQueryRecord parsedRecord = SearchQueryRecord.fromXContent(recordParser); QueryInsightsTestUtils.checkRecordsEquals(List.of(record), List.of(parsedRecord));