Skip to content

Commit 4fd6b0f

Browse files
committed
Add data models for health stats api
Signed-off-by: Chenyang Ji <cyji@amazon.com>
1 parent 290c701 commit 4fd6b0f

15 files changed

+723
-23
lines changed

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

+20
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
package org.opensearch.plugin.insights.core.service;
1010

1111
import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.DEFAULT_GROUPING_TYPE;
12+
import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.QUERY_INSIGHTS_EXECUTOR;
1213
import static org.opensearch.plugin.insights.settings.QueryInsightsSettings.getExporterSettings;
1314

1415
import java.io.IOException;
@@ -33,6 +34,8 @@
3334
import org.opensearch.plugin.insights.rules.model.GroupingType;
3435
import org.opensearch.plugin.insights.rules.model.MetricType;
3536
import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
37+
import org.opensearch.plugin.insights.rules.model.healthStats.QueryInsightsHealthStats;
38+
import org.opensearch.plugin.insights.rules.model.healthStats.TopQueriesHealthStats;
3639
import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
3740
import org.opensearch.telemetry.metrics.MetricsRegistry;
3841
import org.opensearch.threadpool.Scheduler;
@@ -439,4 +442,21 @@ protected void doClose() throws IOException {
439442
queryInsightsExporterFactory.closeAllExporters();
440443
queryInsightsReaderFactory.closeAllReaders();
441444
}
445+
446+
/**
447+
* Get health stats for query insights services
448+
*
449+
* @return QueryInsightsHealthStats
450+
*/
451+
public QueryInsightsHealthStats getHealthStats() {
452+
Map<MetricType, TopQueriesHealthStats> topQueriesHealthStatsMap = new HashMap<>();
453+
for (Map.Entry<MetricType, TopQueriesService> entry : topQueriesServices.entrySet()) {
454+
topQueriesHealthStatsMap.put(entry.getKey(), entry.getValue().getHealthStats());
455+
}
456+
return new QueryInsightsHealthStats(
457+
threadPool.info(QUERY_INSIGHTS_EXECUTOR),
458+
this.queryRecordsQueue.size(),
459+
topQueriesHealthStatsMap
460+
);
461+
}
442462
}

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

+10
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,7 @@
4949
import org.opensearch.plugin.insights.rules.model.GroupingType;
5050
import org.opensearch.plugin.insights.rules.model.MetricType;
5151
import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
52+
import org.opensearch.plugin.insights.rules.model.healthStats.TopQueriesHealthStats;
5253
import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
5354
import org.opensearch.threadpool.ThreadPool;
5455

@@ -516,4 +517,13 @@ private void drain() {
516517
topQueriesHistorySnapshot.set(new ArrayList<>());
517518
topQueriesCurrentSnapshot.set(new ArrayList<>());
518519
}
520+
521+
/**
522+
* Get top queries service health stats
523+
*
524+
* @return TopQueriesHealthStats
525+
*/
526+
public TopQueriesHealthStats getHealthStats() {
527+
return new TopQueriesHealthStats(this.topQueriesStore.size(), this.queryGrouper.getHealthStats());
528+
}
519529
}

src/main/java/org/opensearch/plugin/insights/core/service/grouper/MinMaxHeapQueryGrouper.java

+34-21
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import org.opensearch.plugin.insights.rules.model.GroupingType;
2121
import org.opensearch.plugin.insights.rules.model.MetricType;
2222
import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
23+
import org.opensearch.plugin.insights.rules.model.healthStats.QueryGrouperHealthStats;
2324
import org.opensearch.plugin.insights.settings.QueryInsightsSettings;
2425

2526
/**
@@ -39,12 +40,12 @@ public class MinMaxHeapQueryGrouper implements QueryGrouper {
3940
/**
4041
* Metric type for the current grouping service
4142
*/
42-
private MetricType metricType;
43+
private final MetricType metricType;
4344

4445
/**
4546
* Aggregation type for the current grouping service
4647
*/
47-
private AggregationType aggregationType;
48+
private final AggregationType aggregationType;
4849
/**
4950
* Map storing groupingId to Tuple containing Aggregate search query record and boolean.
5051
* SearchQueryRecord: Aggregate search query record to store the aggregate of a metric type based on the aggregation type..
@@ -53,18 +54,18 @@ public class MinMaxHeapQueryGrouper implements QueryGrouper {
5354
* boolean: True if the aggregate record is in the Top N queries priority query (min heap) and False if the aggregate
5455
* record is in the Max Heap
5556
*/
56-
private ConcurrentHashMap<String, Tuple<SearchQueryRecord, Boolean>> groupIdToAggSearchQueryRecord;
57+
private final ConcurrentHashMap<String, Tuple<SearchQueryRecord, Boolean>> groupIdToAggSearchQueryRecord;
5758
/**
5859
* Min heap to keep track of the Top N query groups and is passed from TopQueriesService as the topQueriesStore
5960
*/
60-
private PriorityBlockingQueue<SearchQueryRecord> minHeapTopQueriesStore;
61+
private final PriorityBlockingQueue<SearchQueryRecord> minHeapTopQueriesStore;
6162
/**
6263
* The Max heap is an overflow data structure used to manage records that exceed the capacity of the Min heap.
6364
* It stores all records not included in the Top N query results. When the aggregate measurement for one of these
6465
* records is updated and it now qualifies as part of the Top N, the record is moved from the Max heap to the Min heap,
6566
* and the records are rearranged accordingly.
6667
*/
67-
private PriorityBlockingQueue<SearchQueryRecord> maxHeapQueryStore;
68+
private final PriorityBlockingQueue<SearchQueryRecord> maxHeapQueryStore;
6869

6970
/**
7071
* Top N size based on the configuration set
@@ -80,11 +81,11 @@ public class MinMaxHeapQueryGrouper implements QueryGrouper {
8081
private int maxGroups;
8182

8283
public MinMaxHeapQueryGrouper(
83-
MetricType metricType,
84-
GroupingType groupingType,
85-
AggregationType aggregationType,
86-
PriorityBlockingQueue<SearchQueryRecord> topQueriesStore,
87-
int topNSize
84+
final MetricType metricType,
85+
final GroupingType groupingType,
86+
final AggregationType aggregationType,
87+
final PriorityBlockingQueue<SearchQueryRecord> topQueriesStore,
88+
final int topNSize
8889
) {
8990
this.groupingType = groupingType;
9091
this.metricType = metricType;
@@ -103,7 +104,7 @@ public MinMaxHeapQueryGrouper(
103104
* @return return the search query record that represents the group
104105
*/
105106
@Override
106-
public SearchQueryRecord add(SearchQueryRecord searchQueryRecord) {
107+
public SearchQueryRecord add(final SearchQueryRecord searchQueryRecord) {
107108
if (groupingType == GroupingType.NONE) {
108109
throw new IllegalArgumentException("Do not use addQueryToGroup when GroupingType is None");
109110
}
@@ -120,8 +121,7 @@ public SearchQueryRecord add(SearchQueryRecord searchQueryRecord) {
120121
// Add to min PQ and promote to max
121122
// If max PQ is empty return else try to promote record from max to min
122123
if (!groupIdToAggSearchQueryRecord.containsKey(groupId)) {
123-
boolean maxGroupsLimitReached = checkMaxGroupsLimitReached(groupId);
124-
if (maxGroupsLimitReached) {
124+
if (checkMaxGroupsLimitReached(groupId)) {
125125
return null;
126126
}
127127
aggregateSearchQueryRecord = searchQueryRecord;
@@ -158,7 +158,7 @@ public void drain() {
158158
* @return grouping type changed
159159
*/
160160
@Override
161-
public boolean setGroupingType(GroupingType newGroupingType) {
161+
public boolean setGroupingType(final GroupingType newGroupingType) {
162162
if (this.groupingType != newGroupingType) {
163163
this.groupingType = newGroupingType;
164164
drain();
@@ -183,7 +183,7 @@ public GroupingType getGroupingType() {
183183
* @return max groups changed
184184
*/
185185
@Override
186-
public boolean setMaxGroups(int maxGroups) {
186+
public boolean setMaxGroups(final int maxGroups) {
187187
if (this.maxGroups != maxGroups) {
188188
this.maxGroups = maxGroups;
189189
drain();
@@ -197,17 +197,21 @@ public boolean setMaxGroups(int maxGroups) {
197197
* @param newSize new size
198198
*/
199199
@Override
200-
public void updateTopNSize(int newSize) {
200+
public void updateTopNSize(final int newSize) {
201201
this.topNSize = newSize;
202202
}
203203

204-
private void addToMinPQ(SearchQueryRecord searchQueryRecord, String groupId) {
204+
private void addToMinPQ(final SearchQueryRecord searchQueryRecord, final String groupId) {
205205
minHeapTopQueriesStore.add(searchQueryRecord);
206206
groupIdToAggSearchQueryRecord.put(groupId, new Tuple<>(searchQueryRecord, true));
207207
overflow();
208208
}
209209

210-
private void addAndPromote(SearchQueryRecord searchQueryRecord, SearchQueryRecord aggregateSearchQueryRecord, String groupId) {
210+
private void addAndPromote(
211+
final SearchQueryRecord searchQueryRecord,
212+
final SearchQueryRecord aggregateSearchQueryRecord,
213+
final String groupId
214+
) {
211215
Number measurementToAdd = searchQueryRecord.getMeasurement(metricType);
212216
aggregateSearchQueryRecord.addMeasurement(metricType, measurementToAdd);
213217
addToMinPQ(aggregateSearchQueryRecord, groupId);
@@ -228,7 +232,7 @@ private void overflow() {
228232
}
229233
}
230234

231-
private boolean checkMaxGroupsLimitReached(String groupId) {
235+
private boolean checkMaxGroupsLimitReached(final String groupId) {
232236
if (maxGroups <= maxHeapQueryStore.size() && minHeapTopQueriesStore.size() >= topNSize) {
233237
log.warn(
234238
"Exceeded [{}] setting threshold which is set at {}. Discarding new group with id {}.",
@@ -259,11 +263,11 @@ int numberOfTopGroups() {
259263
}
260264

261265
/**
262-
* Get groupingId. This should be query hashcode for SIMILARITY grouping and user_id for USER_ID grouping.
266+
* Get groupingId. This should be the query hashcode for SIMILARITY grouping and user_id for USER_ID grouping.
263267
* @param searchQueryRecord record
264268
* @return Grouping Id
265269
*/
266-
private String getGroupingId(SearchQueryRecord searchQueryRecord) {
270+
private String getGroupingId(final SearchQueryRecord searchQueryRecord) {
267271
switch (groupingType) {
268272
case SIMILARITY:
269273
return searchQueryRecord.getAttributes().get(Attribute.QUERY_HASHCODE).toString();
@@ -273,4 +277,13 @@ private String getGroupingId(SearchQueryRecord searchQueryRecord) {
273277
throw new IllegalArgumentException("The following grouping type is not supported : " + groupingType);
274278
}
275279
}
280+
281+
/**
282+
* Get health stats of the MinMaxHeapQueryGrouperService
283+
*
284+
* @return QueryGrouperHealthStats
285+
*/
286+
public QueryGrouperHealthStats getHealthStats() {
287+
return new QueryGrouperHealthStats(this.groupIdToAggSearchQueryRecord.size(), this.maxHeapQueryStore.size());
288+
}
276289
}

src/main/java/org/opensearch/plugin/insights/core/service/grouper/QueryGrouper.java

+8
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010

1111
import org.opensearch.plugin.insights.rules.model.GroupingType;
1212
import org.opensearch.plugin.insights.rules.model.SearchQueryRecord;
13+
import org.opensearch.plugin.insights.rules.model.healthStats.QueryGrouperHealthStats;
1314

1415
/**
1516
* Interface for grouping search queries based on grouping type for the metric type.
@@ -57,4 +58,11 @@ public interface QueryGrouper {
5758
* @param topNSize the new top N size
5859
*/
5960
void updateTopNSize(int topNSize);
61+
62+
/**
63+
* Get health stats of the QueryGrouperService
64+
*
65+
* @return QueryGrouperHealthStats
66+
*/
67+
QueryGrouperHealthStats getHealthStats();
6068
}

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -62,7 +62,7 @@ public static MetricType fromString(final String metricType) {
6262
* @param metricType the MetricType to write
6363
* @throws IOException IOException
6464
*/
65-
static void writeTo(final StreamOutput out, final MetricType metricType) throws IOException {
65+
public static void writeTo(final StreamOutput out, final MetricType metricType) throws IOException {
6666
out.writeString(metricType.toString());
6767
}
6868

Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
/*
2+
* SPDX-License-Identifier: Apache-2.0
3+
*
4+
* The OpenSearch Contributors require contributions made to
5+
* this file be licensed under the Apache-2.0 license or a
6+
* compatible open source license.
7+
*/
8+
9+
package org.opensearch.plugin.insights.rules.model.healthStats;
10+
11+
import java.io.IOException;
12+
import org.opensearch.core.common.io.stream.StreamInput;
13+
import org.opensearch.core.common.io.stream.StreamOutput;
14+
import org.opensearch.core.common.io.stream.Writeable;
15+
import org.opensearch.core.xcontent.ToXContentFragment;
16+
import org.opensearch.core.xcontent.XContentBuilder;
17+
18+
/**
19+
* Represents the health statistics of the query grouper.
20+
*/
21+
public class QueryGrouperHealthStats implements ToXContentFragment, Writeable {
22+
private final int queryGroupCount;
23+
private final int queryGroupHeapSize;
24+
private static final String QUERY_GROUP_COUNT = "QueryGroupCount";
25+
private static final String QUERY_GROUP_HEAP_SIZE = "QueryGroupHeapSize";
26+
27+
/**
28+
* Constructor to read QueryGrouperHealthStats from a StreamInput.
29+
*
30+
* @param in the StreamInput to read the QueryGrouperHealthStats from
31+
* @throws IOException IOException
32+
*/
33+
public QueryGrouperHealthStats(final StreamInput in) throws IOException {
34+
this.queryGroupCount = in.readInt();
35+
this.queryGroupHeapSize = in.readInt();
36+
}
37+
38+
/**
39+
* Constructor of QueryGrouperHealthStats
40+
*
41+
* @param queryGroupCount Number of groups in the grouper
42+
* @param queryGroupHeapSize Heap size of the grouper
43+
*/
44+
public QueryGrouperHealthStats(final int queryGroupCount, final int queryGroupHeapSize) {
45+
this.queryGroupCount = queryGroupCount;
46+
this.queryGroupHeapSize = queryGroupHeapSize;
47+
}
48+
49+
/**
50+
* Write QueryGrouperHealthStats Object to output stream
51+
* @param out streamOutput
52+
* @throws IOException IOException
53+
*/
54+
@Override
55+
public void writeTo(StreamOutput out) throws IOException {
56+
out.writeInt(queryGroupCount);
57+
out.writeInt(queryGroupHeapSize);
58+
}
59+
60+
/**
61+
* Write QueryGrouperHealthStats object to XContent
62+
*
63+
* @param builder XContentBuilder
64+
* @param params Parameters
65+
* @return XContentBuilder
66+
* @throws IOException IOException
67+
*/
68+
@Override
69+
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
70+
builder.field(QUERY_GROUP_COUNT, queryGroupCount);
71+
builder.field(QUERY_GROUP_HEAP_SIZE, queryGroupHeapSize);
72+
return builder;
73+
}
74+
75+
/**
76+
* Gets the number of query groups.
77+
*
78+
* @return the query group count
79+
*/
80+
public int getQueryGroupCount() {
81+
return queryGroupCount;
82+
}
83+
84+
/**
85+
* Gets the query group heap size.
86+
*
87+
* @return the query group heap size
88+
*/
89+
public int getQueryGroupHeapSize() {
90+
return queryGroupHeapSize;
91+
}
92+
}

0 commit comments

Comments
 (0)