Add wlm resiliency orchestrator (query group service) (#15925) (#16225)

* cancellation related * Update CHANGELOG.md * add better cancellation reason * Update DefaultTaskCancellationTests.java * refactor * refactor * Update DefaultTaskCancellation.java * Update DefaultTaskCancellation.java * Update DefaultTaskCancellation.java * Update DefaultTaskSelectionStrategy.java * refactor * refactor node level threshold * use query group task * code clean up and refactorings * add unit tests and fix existing ones * uncomment the test case * update CHANGELOG * fix imports * add queryGroupService * refactor and add UTs for new constructs * fix javadocs * remove code clutter * change annotation version and task selection strategy * rename a util class * remove wrappers from resource type * apply spotless * address comments * add rename changes * address comments * initial changes * refactor changes and logical bug fix * add chanegs * address comments * temp changes * add UTs * add changelog * add task completion listener hook * add remaining pieces to make the feature functional * extend stats and fix bugs * fix bugs and add logic to make SBP work with wlm * address comments * fix bugs and SBP ITs * add missed applyCluster state change * address comments * decouple queryGroupService and cancellationService * replace StateApplier with StateListener interface * fix precommit errors --------- Signed-off-by: Kiran Prakash <[email protected]> Signed-off-by: Kaushal Kumar <[email protected]> Co-authored-by: Kiran Prakash <[email protected]>
opensearch-project · Oct 8, 2024 · 3b4e11d · 3b4e11d
1 parent 1b1c90f
commit 3b4e11d
Show file tree

Hide file tree

Showing 22 changed files with 1,410 additions and 117 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased 2.x]
 ### Added
+- [Workload Management] Add orchestrator for wlm resiliency (QueryGroupService) ([#15925](https://github.com/opensearch-project/OpenSearch/pull/15925))
 - [Offline Nodes] Adds offline-tasks library containing various interfaces to be used for Offline Background Tasks. ([#13574](https://github.com/opensearch-project/OpenSearch/pull/13574))
 - Add support for async deletion in S3BlobContainer ([#15621](https://github.com/opensearch-project/OpenSearch/pull/15621))
 - [Workload Management] QueryGroup resource cancellation framework changes ([#15651](https://github.com/opensearch-project/OpenSearch/pull/15651))

diff --git a/...src/internalClusterTest/java/org/opensearch/search/backpressure/SearchBackpressureIT.java b/...src/internalClusterTest/java/org/opensearch/search/backpressure/SearchBackpressureIT.java
@@ -39,6 +39,7 @@
 import org.opensearch.test.ParameterizedStaticSettingsOpenSearchIntegTestCase;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.TransportService;
+import org.opensearch.wlm.QueryGroupTask;
 import org.hamcrest.MatcherAssert;
 import org.junit.After;
 import org.junit.Before;
@@ -411,6 +412,7 @@ protected void doExecute(Task task, TestRequest request, ActionListener<TestResp
             threadPool.executor(ThreadPool.Names.SEARCH).execute(() -> {
                 try {
                     CancellableTask cancellableTask = (CancellableTask) task;
+                    ((QueryGroupTask) task).setQueryGroupId(threadPool.getThreadContext());
                     long startTime = System.nanoTime();
 
                     // Doing a busy-wait until task cancellation or timeout.

diff --git a/server/src/main/java/org/opensearch/action/ActionModule.java b/server/src/main/java/org/opensearch/action/ActionModule.java
@@ -469,6 +469,7 @@
 import org.opensearch.tasks.Task;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.usage.UsageService;
+import org.opensearch.wlm.QueryGroupTask;
 
 import java.util.ArrayList;
 import java.util.Collections;
@@ -552,7 +553,10 @@ public ActionModule(
         destructiveOperations = new DestructiveOperations(settings, clusterSettings);
         Set<RestHeaderDefinition> headers = Stream.concat(
             actionPlugins.stream().flatMap(p -> p.getRestHeaders().stream()),
-            Stream.of(new RestHeaderDefinition(Task.X_OPAQUE_ID, false))
+            Stream.of(
+                new RestHeaderDefinition(Task.X_OPAQUE_ID, false),
+                new RestHeaderDefinition(QueryGroupTask.QUERY_GROUP_ID_HEADER, false)
+            )
         ).collect(Collectors.toSet());
         UnaryOperator<RestHandler> restWrapper = null;
         for (ActionPlugin plugin : actionPlugins) {

diff --git a/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java b/server/src/main/java/org/opensearch/common/settings/ClusterSettings.java
@@ -797,6 +797,9 @@ public void apply(Settings value, Settings current, Settings previous) {
                 WorkloadManagementSettings.NODE_LEVEL_CPU_CANCELLATION_THRESHOLD,
                 WorkloadManagementSettings.NODE_LEVEL_MEMORY_REJECTION_THRESHOLD,
                 WorkloadManagementSettings.NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD,
+                WorkloadManagementSettings.WLM_MODE_SETTING,
+                WorkloadManagementSettings.QUERYGROUP_SERVICE_RUN_INTERVAL_SETTING,
+                WorkloadManagementSettings.QUERYGROUP_SERVICE_DURESS_STREAK_SETTING,
 
                 SearchService.CLUSTER_ALLOW_DERIVED_FIELD_SETTING,
 

diff --git a/server/src/main/java/org/opensearch/node/Node.java b/server/src/main/java/org/opensearch/node/Node.java
@@ -269,8 +269,13 @@
 import org.opensearch.usage.UsageService;
 import org.opensearch.watcher.ResourceWatcherService;
 import org.opensearch.wlm.QueryGroupService;
+import org.opensearch.wlm.QueryGroupsStateAccessor;
+import org.opensearch.wlm.WorkloadManagementSettings;
 import org.opensearch.wlm.WorkloadManagementTransportInterceptor;
+import org.opensearch.wlm.cancellation.MaximumResourceTaskSelectionStrategy;
+import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService;
 import org.opensearch.wlm.listeners.QueryGroupRequestOperationListener;
+import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService;
 
 import javax.net.ssl.SNIHostName;
 
@@ -1019,8 +1024,30 @@ protected Node(
             List<IdentityAwarePlugin> identityAwarePlugins = pluginsService.filterPlugins(IdentityAwarePlugin.class);
             identityService.initializeIdentityAwarePlugins(identityAwarePlugins);
 
-            final QueryGroupService queryGroupService = new QueryGroupService(); // We will need to replace this with actual instance of the
-                                                                                 // queryGroupService
+            final QueryGroupResourceUsageTrackerService queryGroupResourceUsageTrackerService = new QueryGroupResourceUsageTrackerService(
+                taskResourceTrackingService
+            );
+            final WorkloadManagementSettings workloadManagementSettings = new WorkloadManagementSettings(
+                settings,
+                settingsModule.getClusterSettings()
+            );
+
+            final QueryGroupsStateAccessor queryGroupsStateAccessor = new QueryGroupsStateAccessor();
+
+            final QueryGroupService queryGroupService = new QueryGroupService(
+                new QueryGroupTaskCancellationService(
+                    workloadManagementSettings,
+                    new MaximumResourceTaskSelectionStrategy(),
+                    queryGroupResourceUsageTrackerService,
+                    queryGroupsStateAccessor
+                ),
+                clusterService,
+                threadPool,
+                workloadManagementSettings,
+                queryGroupsStateAccessor
+            );
+            taskResourceTrackingService.addTaskCompletionListener(queryGroupService);
+
             final QueryGroupRequestOperationListener queryGroupRequestOperationListener = new QueryGroupRequestOperationListener(
                 queryGroupService,
                 threadPool
@@ -1086,7 +1113,7 @@ protected Node(
 
             WorkloadManagementTransportInterceptor workloadManagementTransportInterceptor = new WorkloadManagementTransportInterceptor(
                 threadPool,
-                new QueryGroupService() // We will need to replace this with actual implementation
+                queryGroupService
             );
 
             final Collection<SecureSettingsFactory> secureSettingsFactories = pluginsService.filterPlugins(Plugin.class)
@@ -1180,7 +1207,8 @@ protected Node(
                 searchBackpressureSettings,
                 taskResourceTrackingService,
                 threadPool,
-                transportService.getTaskManager()
+                transportService.getTaskManager(),
+                queryGroupService
             );
 
             final SegmentReplicationStatsTracker segmentReplicationStatsTracker = new SegmentReplicationStatsTracker(indicesService);
@@ -1392,6 +1420,7 @@ protected Node(
                 b.bind(IndexingPressureService.class).toInstance(indexingPressureService);
                 b.bind(TaskResourceTrackingService.class).toInstance(taskResourceTrackingService);
                 b.bind(SearchBackpressureService.class).toInstance(searchBackpressureService);
+                b.bind(QueryGroupService.class).toInstance(queryGroupService);
                 b.bind(AdmissionControlService.class).toInstance(admissionControlService);
                 b.bind(UsageService.class).toInstance(usageService);
                 b.bind(AggregationUsageService.class).toInstance(searchModule.getValuesSourceRegistry().getUsageService());
@@ -1583,6 +1612,7 @@ public Node start() throws NodeValidationException {
         nodeService.getMonitorService().start();
         nodeService.getSearchBackpressureService().start();
         nodeService.getTaskCancellationMonitoringService().start();
+        injector.getInstance(QueryGroupService.class).start();
 
         final ClusterService clusterService = injector.getInstance(ClusterService.class);
 
@@ -1756,6 +1786,7 @@ private Node stop() {
         injector.getInstance(FsHealthService.class).stop();
         injector.getInstance(NodeResourceUsageTracker.class).stop();
         injector.getInstance(ResourceUsageCollectorService.class).stop();
+        injector.getInstance(QueryGroupService.class).stop();
         nodeService.getMonitorService().stop();
         nodeService.getSearchBackpressureService().stop();
         injector.getInstance(GatewayService.class).stop();

diff --git a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java
@@ -42,6 +42,7 @@
 import org.opensearch.tasks.TaskResourceTrackingService.TaskCompletionListener;
 import org.opensearch.threadpool.Scheduler;
 import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.wlm.QueryGroupService;
 import org.opensearch.wlm.ResourceType;
 
 import java.io.IOException;
@@ -86,12 +87,14 @@ public class SearchBackpressureService extends AbstractLifecycleComponent implem
 
     private final Map<Class<? extends SearchBackpressureTask>, SearchBackpressureState> searchBackpressureStates;
     private final TaskManager taskManager;
+    private final QueryGroupService queryGroupService;
 
     public SearchBackpressureService(
         SearchBackpressureSettings settings,
         TaskResourceTrackingService taskResourceTrackingService,
         ThreadPool threadPool,
-        TaskManager taskManager
+        TaskManager taskManager,
+        QueryGroupService queryGroupService
     ) {
         this(settings, taskResourceTrackingService, threadPool, System::nanoTime, new NodeDuressTrackers(new EnumMap<>(ResourceType.class) {
             {
@@ -131,7 +134,8 @@ public SearchBackpressureService(
                 settings.getClusterSettings(),
                 SearchShardTaskSettings.SETTING_HEAP_MOVING_AVERAGE_WINDOW_SIZE
             ),
-            taskManager
+            taskManager,
+            queryGroupService
         );
     }
 
@@ -143,14 +147,16 @@ public SearchBackpressureService(
         NodeDuressTrackers nodeDuressTrackers,
         TaskResourceUsageTrackers searchTaskTrackers,
         TaskResourceUsageTrackers searchShardTaskTrackers,
-        TaskManager taskManager
+        TaskManager taskManager,
+        QueryGroupService queryGroupService
     ) {
         this.settings = settings;
         this.taskResourceTrackingService = taskResourceTrackingService;
         this.taskResourceTrackingService.addTaskCompletionListener(this);
         this.threadPool = threadPool;
         this.nodeDuressTrackers = nodeDuressTrackers;
         this.taskManager = taskManager;
+        this.queryGroupService = queryGroupService;
 
         this.searchBackpressureStates = Map.of(
             SearchTask.class,
@@ -346,6 +352,7 @@ <T extends CancellableTask & SearchBackpressureTask> List<CancellableTask> getTa
             .stream()
             .filter(type::isInstance)
             .map(type::cast)
+            .filter(queryGroupService::shouldSBPHandle)
             .collect(Collectors.toUnmodifiableList());
     }