88
99package org .opensearch .wlm ;
1010
11+ import org .apache .logging .log4j .LogManager ;
12+ import org .apache .logging .log4j .Logger ;
13+ import org .opensearch .action .search .SearchShardTask ;
14+ import org .opensearch .cluster .ClusterChangedEvent ;
15+ import org .opensearch .cluster .ClusterStateListener ;
16+ import org .opensearch .cluster .metadata .Metadata ;
17+ import org .opensearch .cluster .metadata .QueryGroup ;
18+ import org .opensearch .cluster .service .ClusterService ;
19+ import org .opensearch .common .lifecycle .AbstractLifecycleComponent ;
1120import org .opensearch .core .concurrency .OpenSearchRejectedExecutionException ;
21+ import org .opensearch .monitor .jvm .JvmStats ;
22+ import org .opensearch .monitor .process .ProcessProbe ;
23+ import org .opensearch .search .backpressure .trackers .NodeDuressTrackers ;
24+ import org .opensearch .search .backpressure .trackers .NodeDuressTrackers .NodeDuressTracker ;
25+ import org .opensearch .tasks .Task ;
26+ import org .opensearch .tasks .TaskResourceTrackingService ;
27+ import org .opensearch .threadpool .Scheduler ;
28+ import org .opensearch .threadpool .ThreadPool ;
29+ import org .opensearch .wlm .cancellation .QueryGroupTaskCancellationService ;
1230import org .opensearch .wlm .stats .QueryGroupState ;
1331import org .opensearch .wlm .stats .QueryGroupStats ;
1432import org .opensearch .wlm .stats .QueryGroupStats .QueryGroupStatsHolder ;
1533
34+ import java .io .IOException ;
1635import java .util .HashMap ;
36+ import java .util .HashSet ;
1737import java .util .Map ;
38+ import java .util .Optional ;
39+ import java .util .Set ;
40+
41+ import static org .opensearch .wlm .tracker .QueryGroupResourceUsageTrackerService .TRACKED_RESOURCES ;
1842
1943/**
2044 * As of now this is a stub and main implementation PR will be raised soon.Coming PR will collate these changes with core QueryGroupService changes
2145 */
22- public class QueryGroupService {
23- // This map does not need to be concurrent since we will process the cluster state change serially and update
24- // this map with new additions and deletions of entries. QueryGroupState is thread safe
25- private final Map <String , QueryGroupState > queryGroupStateMap ;
46+ public class QueryGroupService extends AbstractLifecycleComponent
47+ implements
48+ ClusterStateListener ,
49+ TaskResourceTrackingService .TaskCompletionListener {
50+
51+ private static final Logger logger = LogManager .getLogger (QueryGroupService .class );
52+
53+ private final QueryGroupTaskCancellationService taskCancellationService ;
54+ private volatile Scheduler .Cancellable scheduledFuture ;
55+ private final ThreadPool threadPool ;
56+ private final ClusterService clusterService ;
57+ private final WorkloadManagementSettings workloadManagementSettings ;
58+ private Set <QueryGroup > activeQueryGroups ;
59+ private final Set <QueryGroup > deletedQueryGroups ;
60+ private final NodeDuressTrackers nodeDuressTrackers ;
61+ private final QueryGroupsStateAccessor queryGroupsStateAccessor ;
62+
63+ public QueryGroupService (
64+ QueryGroupTaskCancellationService taskCancellationService ,
65+ ClusterService clusterService ,
66+ ThreadPool threadPool ,
67+ WorkloadManagementSettings workloadManagementSettings ,
68+ QueryGroupsStateAccessor queryGroupsStateAccessor
69+ ) {
70+
71+ this (
72+ taskCancellationService ,
73+ clusterService ,
74+ threadPool ,
75+ workloadManagementSettings ,
76+ new NodeDuressTrackers (
77+ Map .of (
78+ ResourceType .CPU ,
79+ new NodeDuressTracker (
80+ () -> workloadManagementSettings .getNodeLevelCpuCancellationThreshold () < ProcessProbe .getInstance ()
81+ .getProcessCpuPercent () / 100.0 ,
82+ workloadManagementSettings ::getDuressStreak
83+ ),
84+ ResourceType .MEMORY ,
85+ new NodeDuressTracker (
86+ () -> workloadManagementSettings .getNodeLevelMemoryCancellationThreshold () <= JvmStats .jvmStats ()
87+ .getMem ()
88+ .getHeapUsedPercent () / 100.0 ,
89+ workloadManagementSettings ::getDuressStreak
90+ )
91+ )
92+ ),
93+ queryGroupsStateAccessor ,
94+ new HashSet <>(),
95+ new HashSet <>()
96+ );
97+ }
98+
99+ public QueryGroupService (
100+ QueryGroupTaskCancellationService taskCancellationService ,
101+ ClusterService clusterService ,
102+ ThreadPool threadPool ,
103+ WorkloadManagementSettings workloadManagementSettings ,
104+ NodeDuressTrackers nodeDuressTrackers ,
105+ QueryGroupsStateAccessor queryGroupsStateAccessor ,
106+ Set <QueryGroup > activeQueryGroups ,
107+ Set <QueryGroup > deletedQueryGroups
108+ ) {
109+ this .taskCancellationService = taskCancellationService ;
110+ this .clusterService = clusterService ;
111+ this .threadPool = threadPool ;
112+ this .workloadManagementSettings = workloadManagementSettings ;
113+ this .nodeDuressTrackers = nodeDuressTrackers ;
114+ this .activeQueryGroups = activeQueryGroups ;
115+ this .deletedQueryGroups = deletedQueryGroups ;
116+ this .queryGroupsStateAccessor = queryGroupsStateAccessor ;
117+ activeQueryGroups .forEach (queryGroup -> this .queryGroupsStateAccessor .addNewQueryGroup (queryGroup .get_id ()));
118+ this .queryGroupsStateAccessor .addNewQueryGroup (QueryGroupTask .DEFAULT_QUERY_GROUP_ID_SUPPLIER .get ());
119+ this .clusterService .addListener (this );
120+ }
121+
122+ /**
123+ * run at regular interval
124+ */
125+ void doRun () {
126+ if (workloadManagementSettings .getWlmMode () == WlmMode .DISABLED ) {
127+ return ;
128+ }
129+ taskCancellationService .cancelTasks (nodeDuressTrackers ::isNodeInDuress , activeQueryGroups , deletedQueryGroups );
130+ taskCancellationService .pruneDeletedQueryGroups (deletedQueryGroups );
131+ }
132+
133+ /**
134+ * {@link AbstractLifecycleComponent} lifecycle method
135+ */
136+ @ Override
137+ protected void doStart () {
138+ scheduledFuture = threadPool .scheduleWithFixedDelay (() -> {
139+ try {
140+ doRun ();
141+ } catch (Exception e ) {
142+ logger .debug ("Exception occurred in Query Sandbox service" , e );
143+ }
144+ }, this .workloadManagementSettings .getQueryGroupServiceRunInterval (), ThreadPool .Names .GENERIC );
145+ }
26146
27- public QueryGroupService () {
28- this (new HashMap <>());
147+ @ Override
148+ protected void doStop () {
149+ if (scheduledFuture != null ) {
150+ scheduledFuture .cancel ();
151+ }
29152 }
30153
31- public QueryGroupService (Map <String , QueryGroupState > queryGroupStateMap ) {
32- this .queryGroupStateMap = queryGroupStateMap ;
154+ @ Override
155+ protected void doClose () throws IOException {}
156+
157+ @ Override
158+ public void clusterChanged (ClusterChangedEvent event ) {
159+ // Retrieve the current and previous cluster states
160+ Metadata previousMetadata = event .previousState ().metadata ();
161+ Metadata currentMetadata = event .state ().metadata ();
162+
163+ // Extract the query groups from both the current and previous cluster states
164+ Map <String , QueryGroup > previousQueryGroups = previousMetadata .queryGroups ();
165+ Map <String , QueryGroup > currentQueryGroups = currentMetadata .queryGroups ();
166+
167+ // Detect new query groups added in the current cluster state
168+ for (String queryGroupName : currentQueryGroups .keySet ()) {
169+ if (!previousQueryGroups .containsKey (queryGroupName )) {
170+ // New query group detected
171+ QueryGroup newQueryGroup = currentQueryGroups .get (queryGroupName );
172+ // Perform any necessary actions with the new query group
173+ queryGroupsStateAccessor .addNewQueryGroup (newQueryGroup .get_id ());
174+ }
175+ }
176+
177+ // Detect query groups deleted in the current cluster state
178+ for (String queryGroupName : previousQueryGroups .keySet ()) {
179+ if (!currentQueryGroups .containsKey (queryGroupName )) {
180+ // Query group deleted
181+ QueryGroup deletedQueryGroup = previousQueryGroups .get (queryGroupName );
182+ // Perform any necessary actions with the deleted query group
183+ this .deletedQueryGroups .add (deletedQueryGroup );
184+ queryGroupsStateAccessor .removeQueryGroup (deletedQueryGroup .get_id ());
185+ }
186+ }
187+ this .activeQueryGroups = new HashSet <>(currentMetadata .queryGroups ().values ());
33188 }
34189
35190 /**
36191 * updates the failure stats for the query group
192+ *
37193 * @param queryGroupId query group identifier
38194 */
39195 public void incrementFailuresFor (final String queryGroupId ) {
40- QueryGroupState queryGroupState = queryGroupStateMap . get (queryGroupId );
196+ QueryGroupState queryGroupState = queryGroupsStateAccessor . getQueryGroupState (queryGroupId );
41197 // This can happen if the request failed for a deleted query group
42198 // or new queryGroup is being created and has not been acknowledged yet
43199 if (queryGroupState == null ) {
@@ -47,12 +203,11 @@ public void incrementFailuresFor(final String queryGroupId) {
47203 }
48204
49205 /**
50- *
51206 * @return node level query group stats
52207 */
53208 public QueryGroupStats nodeStats () {
54209 final Map <String , QueryGroupStatsHolder > statsHolderMap = new HashMap <>();
55- for (Map .Entry <String , QueryGroupState > queryGroupsState : queryGroupStateMap .entrySet ()) {
210+ for (Map .Entry <String , QueryGroupState > queryGroupsState : queryGroupsStateAccessor . getQueryGroupStateMap () .entrySet ()) {
56211 final String queryGroupId = queryGroupsState .getKey ();
57212 final QueryGroupState currentState = queryGroupsState .getValue ();
58213
@@ -63,18 +218,113 @@ public QueryGroupStats nodeStats() {
63218 }
64219
65220 /**
66- *
67221 * @param queryGroupId query group identifier
68222 */
69223 public void rejectIfNeeded (String queryGroupId ) {
70- if (queryGroupId == null ) return ;
71- boolean reject = false ;
72- final StringBuilder reason = new StringBuilder ();
73- // TODO: At this point this is dummy and we need to decide whether to cancel the request based on last
74- // reported resource usage for the queryGroup. We also need to increment the rejection count here for the
75- // query group
76- if (reject ) {
77- throw new OpenSearchRejectedExecutionException ("QueryGroup " + queryGroupId + " is already contended." + reason .toString ());
224+ if (workloadManagementSettings .getWlmMode () != WlmMode .ENABLED ) {
225+ return ;
226+ }
227+
228+ if (queryGroupId == null || queryGroupId .equals (QueryGroupTask .DEFAULT_QUERY_GROUP_ID_SUPPLIER .get ())) return ;
229+ QueryGroupState queryGroupState = queryGroupsStateAccessor .getQueryGroupState (queryGroupId );
230+
231+ // This can happen if the request failed for a deleted query group
232+ // or new queryGroup is being created and has not been acknowledged yet or invalid query group id
233+ if (queryGroupState == null ) {
234+ return ;
235+ }
236+
237+ // rejections will not happen for SOFT mode QueryGroups
238+ Optional <QueryGroup > optionalQueryGroup = activeQueryGroups .stream ().filter (x -> x .get_id ().equals (queryGroupId )).findFirst ();
239+
240+ if (optionalQueryGroup .isPresent () && optionalQueryGroup .get ().getResiliencyMode () == MutableQueryGroupFragment .ResiliencyMode .SOFT )
241+ return ;
242+
243+ optionalQueryGroup .ifPresent (queryGroup -> {
244+ boolean reject = false ;
245+ final StringBuilder reason = new StringBuilder ();
246+ for (ResourceType resourceType : TRACKED_RESOURCES ) {
247+ if (queryGroup .getResourceLimits ().containsKey (resourceType )) {
248+ final double threshold = getNormalisedRejectionThreshold (
249+ queryGroup .getResourceLimits ().get (resourceType ),
250+ resourceType
251+ );
252+ final double lastRecordedUsage = queryGroupState .getResourceState ().get (resourceType ).getLastRecordedUsage ();
253+ if (threshold < lastRecordedUsage ) {
254+ reject = true ;
255+ reason .append (resourceType )
256+ .append (" limit is breaching for ENFORCED type QueryGroup: (" )
257+ .append (threshold )
258+ .append (" < " )
259+ .append (lastRecordedUsage )
260+ .append ("). " );
261+ queryGroupState .getResourceState ().get (resourceType ).rejections .inc ();
262+ // should not double count even if both the resource limits are breaching
263+ break ;
264+ }
265+ }
266+ }
267+ if (reject ) {
268+ queryGroupState .totalRejections .inc ();
269+ throw new OpenSearchRejectedExecutionException (
270+ "QueryGroup " + queryGroupId + " is already contended. " + reason .toString ()
271+ );
272+ }
273+ });
274+ }
275+
276+ private double getNormalisedRejectionThreshold (double limit , ResourceType resourceType ) {
277+ if (resourceType == ResourceType .CPU ) {
278+ return limit * workloadManagementSettings .getNodeLevelCpuRejectionThreshold ();
279+ } else if (resourceType == ResourceType .MEMORY ) {
280+ return limit * workloadManagementSettings .getNodeLevelMemoryRejectionThreshold ();
281+ }
282+ throw new IllegalArgumentException (resourceType + " is not supported in WLM yet" );
283+ }
284+
285+ public Set <QueryGroup > getActiveQueryGroups () {
286+ return activeQueryGroups ;
287+ }
288+
289+ public Set <QueryGroup > getDeletedQueryGroups () {
290+ return deletedQueryGroups ;
291+ }
292+
293+ /**
294+ * This method determines whether the task should be accounted by SBP if both features co-exist
295+ * @param t QueryGroupTask
296+ * @return whether or not SBP handle it
297+ */
298+ public boolean shouldSBPHandle (Task t ) {
299+ QueryGroupTask task = (QueryGroupTask ) t ;
300+ boolean isInvalidQueryGroupTask = true ;
301+ if (!task .getQueryGroupId ().equals (QueryGroupTask .DEFAULT_QUERY_GROUP_ID_SUPPLIER .get ())) {
302+ isInvalidQueryGroupTask = activeQueryGroups .stream ()
303+ .noneMatch (queryGroup -> queryGroup .get_id ().equals (task .getQueryGroupId ()));
304+ }
305+ return workloadManagementSettings .getWlmMode () != WlmMode .ENABLED || isInvalidQueryGroupTask ;
306+ }
307+
308+ @ Override
309+ public void onTaskCompleted (Task task ) {
310+ if (!(task instanceof QueryGroupTask )) {
311+ return ;
312+ }
313+ final QueryGroupTask queryGroupTask = (QueryGroupTask ) task ;
314+ String queryGroupId = queryGroupTask .getQueryGroupId ();
315+
316+ // set the default queryGroupId if not existing in the active query groups
317+ String finalQueryGroupId = queryGroupId ;
318+ boolean exists = activeQueryGroups .stream ().anyMatch (queryGroup -> queryGroup .get_id ().equals (finalQueryGroupId ));
319+
320+ if (!exists ) {
321+ queryGroupId = QueryGroupTask .DEFAULT_QUERY_GROUP_ID_SUPPLIER .get ();
322+ }
323+
324+ if (task instanceof SearchShardTask ) {
325+ queryGroupsStateAccessor .getQueryGroupState (queryGroupId ).shardCompletions .inc ();
326+ } else {
327+ queryGroupsStateAccessor .getQueryGroupState (queryGroupId ).completions .inc ();
78328 }
79329 }
80330}
0 commit comments