From b74cc5c736f477bdbabc4c164eead549357a9811 Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Mon, 7 Oct 2024 17:21:20 -0700 Subject: [PATCH] Add wlm resiliency orchestrator (query group service) (#15925) * cancellation related Signed-off-by: Kiran Prakash * Update CHANGELOG.md Signed-off-by: Kiran Prakash * add better cancellation reason Signed-off-by: Kiran Prakash * Update DefaultTaskCancellationTests.java Signed-off-by: Kiran Prakash * refactor Signed-off-by: Kiran Prakash * refactor Signed-off-by: Kiran Prakash * Update DefaultTaskCancellation.java Signed-off-by: Kiran Prakash * Update DefaultTaskCancellation.java Signed-off-by: Kiran Prakash * Update DefaultTaskCancellation.java Signed-off-by: Kiran Prakash * Update DefaultTaskSelectionStrategy.java Signed-off-by: Kiran Prakash * refactor Signed-off-by: Kiran Prakash * refactor node level threshold Signed-off-by: Kiran Prakash * use query group task Signed-off-by: Kaushal Kumar * code clean up and refactorings Signed-off-by: Kaushal Kumar * add unit tests and fix existing ones Signed-off-by: Kaushal Kumar * uncomment the test case Signed-off-by: Kaushal Kumar * update CHANGELOG Signed-off-by: Kaushal Kumar * fix imports Signed-off-by: Kaushal Kumar * add queryGroupService Signed-off-by: Kaushal Kumar * refactor and add UTs for new constructs Signed-off-by: Kaushal Kumar * fix javadocs Signed-off-by: Kaushal Kumar * remove code clutter Signed-off-by: Kaushal Kumar * change annotation version and task selection strategy Signed-off-by: Kaushal Kumar * rename a util class Signed-off-by: Kaushal Kumar * remove wrappers from resource type Signed-off-by: Kaushal Kumar * apply spotless Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * add rename changes Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * initial changes Signed-off-by: Kaushal Kumar * refactor changes and logical bug fix Signed-off-by: Kaushal Kumar * add chanegs Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * temp changes Signed-off-by: Kaushal Kumar * add UTs Signed-off-by: Kaushal Kumar * add changelog Signed-off-by: Kaushal Kumar * add task completion listener hook Signed-off-by: Kaushal Kumar * add remaining pieces to make the feature functional Signed-off-by: Kaushal Kumar * extend stats and fix bugs Signed-off-by: Kaushal Kumar * fix bugs and add logic to make SBP work with wlm Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * fix bugs and SBP ITs Signed-off-by: Kaushal Kumar * add missed applyCluster state change Signed-off-by: Kaushal Kumar * address comments Signed-off-by: Kaushal Kumar * decouple queryGroupService and cancellationService Signed-off-by: Kaushal Kumar * replace StateApplier with StateListener interface Signed-off-by: Kaushal Kumar * fix precommit errors Signed-off-by: Kaushal Kumar --------- Signed-off-by: Kiran Prakash Signed-off-by: Kaushal Kumar Co-authored-by: Kiran Prakash --- CHANGELOG.md | 1 + .../backpressure/SearchBackpressureIT.java | 2 + .../org/opensearch/action/ActionModule.java | 6 +- .../common/settings/ClusterSettings.java | 3 + .../main/java/org/opensearch/node/Node.java | 39 +- .../SearchBackpressureService.java | 13 +- .../org/opensearch/wlm/QueryGroupService.java | 290 ++++++++++- .../org/opensearch/wlm/QueryGroupTask.java | 9 +- .../wlm/QueryGroupsStateAccessor.java | 64 +++ .../main/java/org/opensearch/wlm/WlmMode.java | 40 ++ .../wlm/WorkloadManagementSettings.java | 105 ++++ .../QueryGroupTaskCancellationService.java | 123 ++++- .../opensearch/wlm/stats/QueryGroupState.java | 29 +- .../opensearch/wlm/stats/QueryGroupStats.java | 11 +- ...QueryGroupResourceUsageTrackerService.java | 2 + .../SearchBackpressureServiceTests.java | 52 +- .../wlm/QueryGroupServiceTests.java | 489 ++++++++++++++++++ ...adManagementTransportInterceptorTests.java | 36 +- ...ueryGroupTaskCancellationServiceTests.java | 100 +++- ...eryGroupRequestOperationListenerTests.java | 97 +++- .../wlm/stats/QueryGroupStateTests.java | 10 +- .../wlm/stats/QueryGroupStatsTests.java | 6 +- 22 files changed, 1410 insertions(+), 117 deletions(-) create mode 100644 server/src/main/java/org/opensearch/wlm/QueryGroupsStateAccessor.java create mode 100644 server/src/main/java/org/opensearch/wlm/WlmMode.java create mode 100644 server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 59851363c66e9..1c5962feb51be 100644 --- 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/server/src/internalClusterTest/java/org/opensearch/search/backpressure/SearchBackpressureIT.java b/server/src/internalClusterTest/java/org/opensearch/search/backpressure/SearchBackpressureIT.java index fb84134120e00..40c9301ef4bce 100644 --- a/server/src/internalClusterTest/java/org/opensearch/search/backpressure/SearchBackpressureIT.java +++ b/server/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 { 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 index e9392467895fb..d7d54810e360c 100644 --- 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 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 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 index 3ef063b03d738..6eb5812fd2e98 100644 --- 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 index 40a834e6ea28f..e3f9f73bdbd31 100644 --- 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 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 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 index d0480f3dfad80..c4d78ff615610 100644 --- 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, 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,7 +147,8 @@ public SearchBackpressureService( NodeDuressTrackers nodeDuressTrackers, TaskResourceUsageTrackers searchTaskTrackers, TaskResourceUsageTrackers searchShardTaskTrackers, - TaskManager taskManager + TaskManager taskManager, + QueryGroupService queryGroupService ) { this.settings = settings; this.taskResourceTrackingService = taskResourceTrackingService; @@ -151,6 +156,7 @@ public SearchBackpressureService( this.threadPool = threadPool; this.nodeDuressTrackers = nodeDuressTrackers; this.taskManager = taskManager; + this.queryGroupService = queryGroupService; this.searchBackpressureStates = Map.of( SearchTask.class, @@ -346,6 +352,7 @@ List getTa .stream() .filter(type::isInstance) .map(type::cast) + .filter(queryGroupService::shouldSBPHandle) .collect(Collectors.toUnmodifiableList()); } diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java index 6545598dd9951..cda5916db26f3 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java @@ -8,36 +8,192 @@ package org.opensearch.wlm; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.action.search.SearchShardTask; +import org.opensearch.cluster.ClusterChangedEvent; +import org.opensearch.cluster.ClusterStateListener; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.lifecycle.AbstractLifecycleComponent; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.monitor.jvm.JvmStats; +import org.opensearch.monitor.process.ProcessProbe; +import org.opensearch.search.backpressure.trackers.NodeDuressTrackers; +import org.opensearch.search.backpressure.trackers.NodeDuressTrackers.NodeDuressTracker; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskResourceTrackingService; +import org.opensearch.threadpool.Scheduler; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.stats.QueryGroupStats; import org.opensearch.wlm.stats.QueryGroupStats.QueryGroupStatsHolder; +import java.io.IOException; import java.util.HashMap; +import java.util.HashSet; import java.util.Map; +import java.util.Optional; +import java.util.Set; + +import static org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService.TRACKED_RESOURCES; /** * 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 */ -public class QueryGroupService { - // This map does not need to be concurrent since we will process the cluster state change serially and update - // this map with new additions and deletions of entries. QueryGroupState is thread safe - private final Map queryGroupStateMap; +public class QueryGroupService extends AbstractLifecycleComponent + implements + ClusterStateListener, + TaskResourceTrackingService.TaskCompletionListener { + + private static final Logger logger = LogManager.getLogger(QueryGroupService.class); + + private final QueryGroupTaskCancellationService taskCancellationService; + private volatile Scheduler.Cancellable scheduledFuture; + private final ThreadPool threadPool; + private final ClusterService clusterService; + private final WorkloadManagementSettings workloadManagementSettings; + private Set activeQueryGroups; + private final Set deletedQueryGroups; + private final NodeDuressTrackers nodeDuressTrackers; + private final QueryGroupsStateAccessor queryGroupsStateAccessor; + + public QueryGroupService( + QueryGroupTaskCancellationService taskCancellationService, + ClusterService clusterService, + ThreadPool threadPool, + WorkloadManagementSettings workloadManagementSettings, + QueryGroupsStateAccessor queryGroupsStateAccessor + ) { + + this( + taskCancellationService, + clusterService, + threadPool, + workloadManagementSettings, + new NodeDuressTrackers( + Map.of( + ResourceType.CPU, + new NodeDuressTracker( + () -> workloadManagementSettings.getNodeLevelCpuCancellationThreshold() < ProcessProbe.getInstance() + .getProcessCpuPercent() / 100.0, + workloadManagementSettings::getDuressStreak + ), + ResourceType.MEMORY, + new NodeDuressTracker( + () -> workloadManagementSettings.getNodeLevelMemoryCancellationThreshold() <= JvmStats.jvmStats() + .getMem() + .getHeapUsedPercent() / 100.0, + workloadManagementSettings::getDuressStreak + ) + ) + ), + queryGroupsStateAccessor, + new HashSet<>(), + new HashSet<>() + ); + } + + public QueryGroupService( + QueryGroupTaskCancellationService taskCancellationService, + ClusterService clusterService, + ThreadPool threadPool, + WorkloadManagementSettings workloadManagementSettings, + NodeDuressTrackers nodeDuressTrackers, + QueryGroupsStateAccessor queryGroupsStateAccessor, + Set activeQueryGroups, + Set deletedQueryGroups + ) { + this.taskCancellationService = taskCancellationService; + this.clusterService = clusterService; + this.threadPool = threadPool; + this.workloadManagementSettings = workloadManagementSettings; + this.nodeDuressTrackers = nodeDuressTrackers; + this.activeQueryGroups = activeQueryGroups; + this.deletedQueryGroups = deletedQueryGroups; + this.queryGroupsStateAccessor = queryGroupsStateAccessor; + activeQueryGroups.forEach(queryGroup -> this.queryGroupsStateAccessor.addNewQueryGroup(queryGroup.get_id())); + this.queryGroupsStateAccessor.addNewQueryGroup(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get()); + this.clusterService.addListener(this); + } + + /** + * run at regular interval + */ + void doRun() { + if (workloadManagementSettings.getWlmMode() == WlmMode.DISABLED) { + return; + } + taskCancellationService.cancelTasks(nodeDuressTrackers::isNodeInDuress, activeQueryGroups, deletedQueryGroups); + taskCancellationService.pruneDeletedQueryGroups(deletedQueryGroups); + } + + /** + * {@link AbstractLifecycleComponent} lifecycle method + */ + @Override + protected void doStart() { + scheduledFuture = threadPool.scheduleWithFixedDelay(() -> { + try { + doRun(); + } catch (Exception e) { + logger.debug("Exception occurred in Query Sandbox service", e); + } + }, this.workloadManagementSettings.getQueryGroupServiceRunInterval(), ThreadPool.Names.GENERIC); + } - public QueryGroupService() { - this(new HashMap<>()); + @Override + protected void doStop() { + if (scheduledFuture != null) { + scheduledFuture.cancel(); + } } - public QueryGroupService(Map queryGroupStateMap) { - this.queryGroupStateMap = queryGroupStateMap; + @Override + protected void doClose() throws IOException {} + + @Override + public void clusterChanged(ClusterChangedEvent event) { + // Retrieve the current and previous cluster states + Metadata previousMetadata = event.previousState().metadata(); + Metadata currentMetadata = event.state().metadata(); + + // Extract the query groups from both the current and previous cluster states + Map previousQueryGroups = previousMetadata.queryGroups(); + Map currentQueryGroups = currentMetadata.queryGroups(); + + // Detect new query groups added in the current cluster state + for (String queryGroupName : currentQueryGroups.keySet()) { + if (!previousQueryGroups.containsKey(queryGroupName)) { + // New query group detected + QueryGroup newQueryGroup = currentQueryGroups.get(queryGroupName); + // Perform any necessary actions with the new query group + queryGroupsStateAccessor.addNewQueryGroup(newQueryGroup.get_id()); + } + } + + // Detect query groups deleted in the current cluster state + for (String queryGroupName : previousQueryGroups.keySet()) { + if (!currentQueryGroups.containsKey(queryGroupName)) { + // Query group deleted + QueryGroup deletedQueryGroup = previousQueryGroups.get(queryGroupName); + // Perform any necessary actions with the deleted query group + this.deletedQueryGroups.add(deletedQueryGroup); + queryGroupsStateAccessor.removeQueryGroup(deletedQueryGroup.get_id()); + } + } + this.activeQueryGroups = new HashSet<>(currentMetadata.queryGroups().values()); } /** * updates the failure stats for the query group + * * @param queryGroupId query group identifier */ public void incrementFailuresFor(final String queryGroupId) { - QueryGroupState queryGroupState = queryGroupStateMap.get(queryGroupId); + QueryGroupState queryGroupState = queryGroupsStateAccessor.getQueryGroupState(queryGroupId); // This can happen if the request failed for a deleted query group // or new queryGroup is being created and has not been acknowledged yet if (queryGroupState == null) { @@ -47,12 +203,11 @@ public void incrementFailuresFor(final String queryGroupId) { } /** - * * @return node level query group stats */ public QueryGroupStats nodeStats() { final Map statsHolderMap = new HashMap<>(); - for (Map.Entry queryGroupsState : queryGroupStateMap.entrySet()) { + for (Map.Entry queryGroupsState : queryGroupsStateAccessor.getQueryGroupStateMap().entrySet()) { final String queryGroupId = queryGroupsState.getKey(); final QueryGroupState currentState = queryGroupsState.getValue(); @@ -63,18 +218,113 @@ public QueryGroupStats nodeStats() { } /** - * * @param queryGroupId query group identifier */ public void rejectIfNeeded(String queryGroupId) { - if (queryGroupId == null) return; - boolean reject = false; - final StringBuilder reason = new StringBuilder(); - // TODO: At this point this is dummy and we need to decide whether to cancel the request based on last - // reported resource usage for the queryGroup. We also need to increment the rejection count here for the - // query group - if (reject) { - throw new OpenSearchRejectedExecutionException("QueryGroup " + queryGroupId + " is already contended." + reason.toString()); + if (workloadManagementSettings.getWlmMode() != WlmMode.ENABLED) { + return; + } + + if (queryGroupId == null || queryGroupId.equals(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get())) return; + QueryGroupState queryGroupState = queryGroupsStateAccessor.getQueryGroupState(queryGroupId); + + // This can happen if the request failed for a deleted query group + // or new queryGroup is being created and has not been acknowledged yet or invalid query group id + if (queryGroupState == null) { + return; + } + + // rejections will not happen for SOFT mode QueryGroups + Optional optionalQueryGroup = activeQueryGroups.stream().filter(x -> x.get_id().equals(queryGroupId)).findFirst(); + + if (optionalQueryGroup.isPresent() && optionalQueryGroup.get().getResiliencyMode() == MutableQueryGroupFragment.ResiliencyMode.SOFT) + return; + + optionalQueryGroup.ifPresent(queryGroup -> { + boolean reject = false; + final StringBuilder reason = new StringBuilder(); + for (ResourceType resourceType : TRACKED_RESOURCES) { + if (queryGroup.getResourceLimits().containsKey(resourceType)) { + final double threshold = getNormalisedRejectionThreshold( + queryGroup.getResourceLimits().get(resourceType), + resourceType + ); + final double lastRecordedUsage = queryGroupState.getResourceState().get(resourceType).getLastRecordedUsage(); + if (threshold < lastRecordedUsage) { + reject = true; + reason.append(resourceType) + .append(" limit is breaching for ENFORCED type QueryGroup: (") + .append(threshold) + .append(" < ") + .append(lastRecordedUsage) + .append("). "); + queryGroupState.getResourceState().get(resourceType).rejections.inc(); + // should not double count even if both the resource limits are breaching + break; + } + } + } + if (reject) { + queryGroupState.totalRejections.inc(); + throw new OpenSearchRejectedExecutionException( + "QueryGroup " + queryGroupId + " is already contended. " + reason.toString() + ); + } + }); + } + + private double getNormalisedRejectionThreshold(double limit, ResourceType resourceType) { + if (resourceType == ResourceType.CPU) { + return limit * workloadManagementSettings.getNodeLevelCpuRejectionThreshold(); + } else if (resourceType == ResourceType.MEMORY) { + return limit * workloadManagementSettings.getNodeLevelMemoryRejectionThreshold(); + } + throw new IllegalArgumentException(resourceType + " is not supported in WLM yet"); + } + + public Set getActiveQueryGroups() { + return activeQueryGroups; + } + + public Set getDeletedQueryGroups() { + return deletedQueryGroups; + } + + /** + * This method determines whether the task should be accounted by SBP if both features co-exist + * @param t QueryGroupTask + * @return whether or not SBP handle it + */ + public boolean shouldSBPHandle(Task t) { + QueryGroupTask task = (QueryGroupTask) t; + boolean isInvalidQueryGroupTask = true; + if (!task.getQueryGroupId().equals(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get())) { + isInvalidQueryGroupTask = activeQueryGroups.stream() + .noneMatch(queryGroup -> queryGroup.get_id().equals(task.getQueryGroupId())); + } + return workloadManagementSettings.getWlmMode() != WlmMode.ENABLED || isInvalidQueryGroupTask; + } + + @Override + public void onTaskCompleted(Task task) { + if (!(task instanceof QueryGroupTask)) { + return; + } + final QueryGroupTask queryGroupTask = (QueryGroupTask) task; + String queryGroupId = queryGroupTask.getQueryGroupId(); + + // set the default queryGroupId if not existing in the active query groups + String finalQueryGroupId = queryGroupId; + boolean exists = activeQueryGroups.stream().anyMatch(queryGroup -> queryGroup.get_id().equals(finalQueryGroupId)); + + if (!exists) { + queryGroupId = QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(); + } + + if (task instanceof SearchShardTask) { + queryGroupsStateAccessor.getQueryGroupState(queryGroupId).shardCompletions.inc(); + } else { + queryGroupsStateAccessor.getQueryGroupState(queryGroupId).completions.inc(); } } } diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java b/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java index a1cb766579d43..97c48bd828978 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java @@ -17,7 +17,6 @@ import org.opensearch.tasks.CancellableTask; import java.util.Map; -import java.util.Optional; import java.util.function.LongSupplier; import java.util.function.Supplier; @@ -82,9 +81,11 @@ public final String getQueryGroupId() { * @param threadContext current threadContext */ public final void setQueryGroupId(final ThreadContext threadContext) { - this.queryGroupId = Optional.ofNullable(threadContext) - .map(threadContext1 -> threadContext1.getHeader(QUERY_GROUP_ID_HEADER)) - .orElse(DEFAULT_QUERY_GROUP_ID_SUPPLIER.get()); + if (threadContext != null && threadContext.getHeader(QUERY_GROUP_ID_HEADER) != null) { + this.queryGroupId = threadContext.getHeader(QUERY_GROUP_ID_HEADER); + } else { + this.queryGroupId = DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(); + } } public long getElapsedTime() { diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupsStateAccessor.java b/server/src/main/java/org/opensearch/wlm/QueryGroupsStateAccessor.java new file mode 100644 index 0000000000000..7f93e41f12092 --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupsStateAccessor.java @@ -0,0 +1,64 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.wlm.stats.QueryGroupState; + +import java.util.HashMap; +import java.util.Map; + +/** + * This class is used to decouple {@link QueryGroupService} and {@link org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService} to share the + * {@link QueryGroupState}s + */ +public class QueryGroupsStateAccessor { + // This map does not need to be concurrent since we will process the cluster state change serially and update + // this map with new additions and deletions of entries. QueryGroupState is thread safe + private final Map queryGroupStateMap; + + public QueryGroupsStateAccessor() { + this(new HashMap<>()); + } + + public QueryGroupsStateAccessor(Map queryGroupStateMap) { + this.queryGroupStateMap = queryGroupStateMap; + } + + /** + * returns the query groups state + */ + public Map getQueryGroupStateMap() { + return queryGroupStateMap; + } + + /** + * return QueryGroupState for the given queryGroupId + * @param queryGroupId + * @return QueryGroupState for the given queryGroupId, if id is invalid return default query group state + */ + public QueryGroupState getQueryGroupState(String queryGroupId) { + return queryGroupStateMap.getOrDefault(queryGroupId, queryGroupStateMap.get(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get())); + } + + /** + * adds new QueryGroupState against given queryGroupId + * @param queryGroupId + */ + public void addNewQueryGroup(String queryGroupId) { + this.queryGroupStateMap.putIfAbsent(queryGroupId, new QueryGroupState()); + } + + /** + * removes QueryGroupState against given queryGroupId + * @param queryGroupId + */ + public void removeQueryGroup(String queryGroupId) { + this.queryGroupStateMap.remove(queryGroupId); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/WlmMode.java b/server/src/main/java/org/opensearch/wlm/WlmMode.java new file mode 100644 index 0000000000000..40407525cc24d --- /dev/null +++ b/server/src/main/java/org/opensearch/wlm/WlmMode.java @@ -0,0 +1,40 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.common.annotation.PublicApi; + +/** + * Enum to hold the values whether wlm is enabled or not + */ +@PublicApi(since = "2.18.0") +public enum WlmMode { + ENABLED("enabled"), + MONITOR_ONLY("monitor_only"), + DISABLED("disabled"); + + private final String name; + + WlmMode(String name) { + this.name = name; + } + + public String getName() { + return name; + } + + public static WlmMode fromName(String name) { + for (WlmMode wlmMode : values()) { + if (wlmMode.getName().equals(name)) { + return wlmMode; + } + } + throw new IllegalArgumentException(name + " is an invalid WlmMode"); + } +} diff --git a/server/src/main/java/org/opensearch/wlm/WorkloadManagementSettings.java b/server/src/main/java/org/opensearch/wlm/WorkloadManagementSettings.java index b3577c1b3219d..af25eedd7eed5 100644 --- a/server/src/main/java/org/opensearch/wlm/WorkloadManagementSettings.java +++ b/server/src/main/java/org/opensearch/wlm/WorkloadManagementSettings.java @@ -12,6 +12,7 @@ import org.opensearch.common.settings.ClusterSettings; import org.opensearch.common.settings.Setting; import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; /** * Main class to declare Workload Management related settings @@ -22,16 +23,66 @@ public class WorkloadManagementSettings { private static final Double DEFAULT_NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD = 0.9; private static final Double DEFAULT_NODE_LEVEL_CPU_REJECTION_THRESHOLD = 0.8; private static final Double DEFAULT_NODE_LEVEL_CPU_CANCELLATION_THRESHOLD = 0.9; + private static final Long DEFAULT_QUERYGROUP_SERVICE_RUN_INTERVAL_MILLIS = 1000L; public static final double NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD_MAX_VALUE = 0.95; public static final double NODE_LEVEL_MEMORY_REJECTION_THRESHOLD_MAX_VALUE = 0.9; public static final double NODE_LEVEL_CPU_CANCELLATION_THRESHOLD_MAX_VALUE = 0.95; public static final double NODE_LEVEL_CPU_REJECTION_THRESHOLD_MAX_VALUE = 0.9; + public static final String DEFAULT_WLM_MODE = "monitor_only"; private Double nodeLevelMemoryCancellationThreshold; private Double nodeLevelMemoryRejectionThreshold; private Double nodeLevelCpuCancellationThreshold; private Double nodeLevelCpuRejectionThreshold; + /** + * Setting name for QueryGroupService node duress streak + */ + public static final String QUERYGROUP_DURESS_STREAK_SETTING_NAME = "wlm.query_group.duress_streak"; + private int duressStreak; + public static final Setting QUERYGROUP_SERVICE_DURESS_STREAK_SETTING = Setting.intSetting( + QUERYGROUP_DURESS_STREAK_SETTING_NAME, + 3, + 3, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * Setting name for Query Group Service run interval + */ + public static final String QUERYGROUP_ENFORCEMENT_INTERVAL_SETTING_NAME = "wlm.query_group.enforcement_interval"; + + private TimeValue queryGroupServiceRunInterval; + /** + * Setting to control the run interval of Query Group Service + */ + public static final Setting QUERYGROUP_SERVICE_RUN_INTERVAL_SETTING = Setting.longSetting( + QUERYGROUP_ENFORCEMENT_INTERVAL_SETTING_NAME, + DEFAULT_QUERYGROUP_SERVICE_RUN_INTERVAL_MILLIS, + 1000, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * WLM mode setting name + */ + public static final String WLM_MODE_SETTING_NAME = "wlm.query_group.mode"; + + private volatile WlmMode wlmMode; + + /** + * WLM mode setting, which determines which mode WLM is operating in + */ + public static final Setting WLM_MODE_SETTING = new Setting( + WLM_MODE_SETTING_NAME, + DEFAULT_WLM_MODE, + WlmMode::fromName, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + /** * Setting name for node level memory based rejection threshold for QueryGroup service */ @@ -91,10 +142,13 @@ public class WorkloadManagementSettings { * @param clusterSettings - QueryGroup cluster settings */ public WorkloadManagementSettings(Settings settings, ClusterSettings clusterSettings) { + this.wlmMode = WLM_MODE_SETTING.get(settings); nodeLevelMemoryCancellationThreshold = NODE_LEVEL_MEMORY_CANCELLATION_THRESHOLD.get(settings); nodeLevelMemoryRejectionThreshold = NODE_LEVEL_MEMORY_REJECTION_THRESHOLD.get(settings); nodeLevelCpuCancellationThreshold = NODE_LEVEL_CPU_CANCELLATION_THRESHOLD.get(settings); nodeLevelCpuRejectionThreshold = NODE_LEVEL_CPU_REJECTION_THRESHOLD.get(settings); + this.queryGroupServiceRunInterval = TimeValue.timeValueMillis(QUERYGROUP_SERVICE_RUN_INTERVAL_SETTING.get(settings)); + duressStreak = QUERYGROUP_SERVICE_DURESS_STREAK_SETTING.get(settings); ensureRejectionThresholdIsLessThanCancellation( nodeLevelMemoryRejectionThreshold, @@ -113,6 +167,57 @@ public WorkloadManagementSettings(Settings settings, ClusterSettings clusterSett clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_MEMORY_REJECTION_THRESHOLD, this::setNodeLevelMemoryRejectionThreshold); clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CPU_CANCELLATION_THRESHOLD, this::setNodeLevelCpuCancellationThreshold); clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CPU_REJECTION_THRESHOLD, this::setNodeLevelCpuRejectionThreshold); + clusterSettings.addSettingsUpdateConsumer(WLM_MODE_SETTING, this::setWlmMode); + clusterSettings.addSettingsUpdateConsumer(QUERYGROUP_SERVICE_RUN_INTERVAL_SETTING, this::setQueryGroupServiceRunInterval); + clusterSettings.addSettingsUpdateConsumer(QUERYGROUP_SERVICE_DURESS_STREAK_SETTING, this::setDuressStreak); + } + + /** + * node duress streak getter + * @return current duressStreak value + */ + public int getDuressStreak() { + return duressStreak; + } + + /** + * node duress streak setter + * @param duressStreak new value + */ + private void setDuressStreak(int duressStreak) { + this.duressStreak = duressStreak; + } + + /** + * queryGroupServiceRunInterval setter + * @param newIntervalInMillis new value + */ + private void setQueryGroupServiceRunInterval(long newIntervalInMillis) { + this.queryGroupServiceRunInterval = TimeValue.timeValueMillis(newIntervalInMillis); + } + + /** + * queryGroupServiceRunInterval getter + * @return current queryGroupServiceRunInterval value + */ + public TimeValue getQueryGroupServiceRunInterval() { + return this.queryGroupServiceRunInterval; + } + + /** + * WlmMode setter + * @param mode new mode value + */ + private void setWlmMode(final WlmMode mode) { + this.wlmMode = mode; + } + + /** + * WlmMode getter + * @return the current wlmMode + */ + public WlmMode getWlmMode() { + return this.wlmMode; } /** diff --git a/server/src/main/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationService.java b/server/src/main/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationService.java index a2c97c8d8635b..e82a19c5f7af2 100644 --- a/server/src/main/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationService.java +++ b/server/src/main/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationService.java @@ -8,20 +8,26 @@ package org.opensearch.wlm.cancellation; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; import org.opensearch.cluster.metadata.QueryGroup; -import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.TaskCancellation; import org.opensearch.wlm.MutableQueryGroupFragment.ResiliencyMode; import org.opensearch.wlm.QueryGroupLevelResourceUsageView; import org.opensearch.wlm.QueryGroupTask; +import org.opensearch.wlm.QueryGroupsStateAccessor; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WlmMode; import org.opensearch.wlm.WorkloadManagementSettings; +import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService; import java.util.ArrayList; import java.util.Collection; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Set; import java.util.function.BooleanSupplier; import java.util.function.Consumer; import java.util.stream.Collectors; @@ -47,46 +53,78 @@ */ public class QueryGroupTaskCancellationService { public static final double MIN_VALUE = 1e-9; + private static final Logger log = LogManager.getLogger(QueryGroupTaskCancellationService.class); private final WorkloadManagementSettings workloadManagementSettings; private final TaskSelectionStrategy taskSelectionStrategy; private final QueryGroupResourceUsageTrackerService resourceUsageTrackerService; // a map of QueryGroupId to its corresponding QueryGroupLevelResourceUsageView object Map queryGroupLevelResourceUsageViews; - private final Collection activeQueryGroups; - private final Collection deletedQueryGroups; + private final QueryGroupsStateAccessor queryGroupStateAccessor; public QueryGroupTaskCancellationService( WorkloadManagementSettings workloadManagementSettings, TaskSelectionStrategy taskSelectionStrategy, QueryGroupResourceUsageTrackerService resourceUsageTrackerService, - Collection activeQueryGroups, - Collection deletedQueryGroups + QueryGroupsStateAccessor queryGroupStateAccessor ) { this.workloadManagementSettings = workloadManagementSettings; this.taskSelectionStrategy = taskSelectionStrategy; this.resourceUsageTrackerService = resourceUsageTrackerService; - this.activeQueryGroups = activeQueryGroups; - this.deletedQueryGroups = deletedQueryGroups; + this.queryGroupStateAccessor = queryGroupStateAccessor; } /** * Cancel tasks based on the implemented strategy. */ - public final void cancelTasks(BooleanSupplier isNodeInDuress) { + public void cancelTasks( + BooleanSupplier isNodeInDuress, + Collection activeQueryGroups, + Collection deletedQueryGroups + ) { queryGroupLevelResourceUsageViews = resourceUsageTrackerService.constructQueryGroupLevelUsageViews(); // cancel tasks from QueryGroups that are in Enforced mode that are breaching their resource limits - cancelTasks(ResiliencyMode.ENFORCED); + cancelTasks(ResiliencyMode.ENFORCED, activeQueryGroups); // if the node is in duress, cancel tasks accordingly. - handleNodeDuress(isNodeInDuress); + handleNodeDuress(isNodeInDuress, activeQueryGroups, deletedQueryGroups); + + updateResourceUsageInQueryGroupState(activeQueryGroups); + } + + private void updateResourceUsageInQueryGroupState(Collection activeQueryGroups) { + Set isSearchWorkloadRunning = new HashSet<>(); + for (Map.Entry queryGroupLevelResourceUsageViewEntry : queryGroupLevelResourceUsageViews + .entrySet()) { + isSearchWorkloadRunning.add(queryGroupLevelResourceUsageViewEntry.getKey()); + QueryGroupState queryGroupState = getQueryGroupState(queryGroupLevelResourceUsageViewEntry.getKey()); + TRACKED_RESOURCES.forEach(resourceType -> { + final double currentUsage = queryGroupLevelResourceUsageViewEntry.getValue().getResourceUsageData().get(resourceType); + queryGroupState.getResourceState().get(resourceType).setLastRecordedUsage(currentUsage); + }); + } + + activeQueryGroups.forEach(queryGroup -> { + if (!isSearchWorkloadRunning.contains(queryGroup.get_id())) { + TRACKED_RESOURCES.forEach( + resourceType -> getQueryGroupState(queryGroup.get_id()).getResourceState().get(resourceType).setLastRecordedUsage(0.0) + ); + } + }); } - private void handleNodeDuress(BooleanSupplier isNodeInDuress) { + private void handleNodeDuress( + BooleanSupplier isNodeInDuress, + Collection activeQueryGroups, + Collection deletedQueryGroups + ) { if (!isNodeInDuress.getAsBoolean()) { return; } // List of tasks to be executed in order if the node is in duress - List> duressActions = List.of(v -> cancelTasksFromDeletedQueryGroups(), v -> cancelTasks(ResiliencyMode.SOFT)); + List> duressActions = List.of( + v -> cancelTasksFromDeletedQueryGroups(deletedQueryGroups), + v -> cancelTasks(ResiliencyMode.SOFT, activeQueryGroups) + ); for (Consumer duressAction : duressActions) { if (!isNodeInDuress.getAsBoolean()) { @@ -96,8 +134,8 @@ private void handleNodeDuress(BooleanSupplier isNodeInDuress) { } } - private void cancelTasksFromDeletedQueryGroups() { - cancelTasks(getAllCancellableTasks(this.deletedQueryGroups)); + private void cancelTasksFromDeletedQueryGroups(Collection deletedQueryGroups) { + cancelTasks(getAllCancellableTasks(deletedQueryGroups)); } /** @@ -105,9 +143,9 @@ private void cancelTasksFromDeletedQueryGroups() { * * @return List of tasks that can be cancelled */ - List getAllCancellableTasks(ResiliencyMode resiliencyMode) { + List getAllCancellableTasks(ResiliencyMode resiliencyMode, Collection queryGroups) { return getAllCancellableTasks( - activeQueryGroups.stream().filter(queryGroup -> queryGroup.getResiliencyMode() == resiliencyMode).collect(Collectors.toList()) + queryGroups.stream().filter(queryGroup -> queryGroup.getResiliencyMode() == resiliencyMode).collect(Collectors.toList()) ); } @@ -118,6 +156,7 @@ List getAllCancellableTasks(ResiliencyMode resiliencyMode) { */ List getAllCancellableTasks(Collection queryGroups) { List taskCancellations = new ArrayList<>(); + final List onCancelCallbacks = new ArrayList<>(); for (QueryGroup queryGroup : queryGroups) { final List reasons = new ArrayList<>(); List selectedTasks = new ArrayList<>(); @@ -127,8 +166,7 @@ List getAllCancellableTasks(Collection queryGroups .calculateResourceUsage(selectedTasks); if (excessUsage > MIN_VALUE) { reasons.add(new TaskCancellation.Reason(generateReasonString(queryGroup, resourceType), 1)); - // TODO: We will need to add the cancellation callback for these resources for the queryGroup to reflect stats - + onCancelCallbacks.add(this.getResourceTypeOnCancelCallback(queryGroup.get_id(), resourceType)); // Only add tasks not already added to avoid double cancellations selectedTasks.addAll( taskSelectionStrategy.selectTasksForCancellation(getTasksFor(queryGroup), excessUsage, resourceType) @@ -140,8 +178,9 @@ List getAllCancellableTasks(Collection queryGroups } if (!reasons.isEmpty()) { + onCancelCallbacks.add(getQueryGroupState(queryGroup.get_id()).totalCancellations::inc); taskCancellations.addAll( - selectedTasks.stream().map(task -> createTaskCancellation(task, reasons)).collect(Collectors.toList()) + selectedTasks.stream().map(task -> new TaskCancellation(task, reasons, onCancelCallbacks)).collect(Collectors.toList()) ); } } @@ -164,16 +203,27 @@ private List getTasksFor(QueryGroup queryGroup) { return queryGroupLevelResourceUsageViews.get(queryGroup.get_id()).getActiveTasks(); } - private void cancelTasks(ResiliencyMode resiliencyMode) { - cancelTasks(getAllCancellableTasks(resiliencyMode)); + private void cancelTasks(ResiliencyMode resiliencyMode, Collection queryGroups) { + cancelTasks(getAllCancellableTasks(resiliencyMode, queryGroups)); } private void cancelTasks(List cancellableTasks) { - cancellableTasks.forEach(TaskCancellation::cancel); - } - private TaskCancellation createTaskCancellation(CancellableTask task, List reasons) { - return new TaskCancellation(task, reasons, List.of(this::callbackOnCancel)); + Consumer cancellationLoggingConsumer = (taskCancellation -> { + log.warn( + "Task {} is eligible for cancellation for reason {}", + taskCancellation.getTask().getId(), + taskCancellation.getReasonString() + ); + }); + Consumer cancellationConsumer = cancellationLoggingConsumer; + if (workloadManagementSettings.getWlmMode() == WlmMode.ENABLED) { + cancellationConsumer = (taskCancellation -> { + cancellationLoggingConsumer.accept(taskCancellation); + taskCancellation.cancel(); + }); + } + cancellableTasks.forEach(cancellationConsumer); } private double getExcessUsage(QueryGroup queryGroup, ResourceType resourceType) { @@ -199,7 +249,26 @@ private double getNormalisedThreshold(QueryGroup queryGroup, ResourceType resour return queryGroup.getResourceLimits().get(resourceType) * nodeLevelCancellationThreshold; } - private void callbackOnCancel() { - // TODO Implement callback logic here mostly used for Stats + private Runnable getResourceTypeOnCancelCallback(String queryGroupId, ResourceType resourceType) { + QueryGroupState queryGroupState = getQueryGroupState(queryGroupId); + return queryGroupState.getResourceState().get(resourceType).cancellations::inc; + } + + private QueryGroupState getQueryGroupState(String queryGroupId) { + assert queryGroupId != null : "queryGroupId should never be null at this point."; + + return queryGroupStateAccessor.getQueryGroupState(queryGroupId); + } + + /** + * Removes the queryGroups from deleted list if it doesn't have any tasks running + */ + public void pruneDeletedQueryGroups(Collection deletedQueryGroups) { + List currentDeletedQueryGroups = new ArrayList<>(deletedQueryGroups); + for (QueryGroup queryGroup : currentDeletedQueryGroups) { + if (queryGroupLevelResourceUsageViews.get(queryGroup.get_id()).getActiveTasks().isEmpty()) { + deletedQueryGroups.remove(queryGroup); + } + } } } diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java index 376d34dd7c8ca..cbc7046a79464 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupState.java @@ -19,14 +19,19 @@ */ public class QueryGroupState { /** - * completions at the query group level, this is a cumulative counter since the Opensearch start time + * co-ordinator level completions at the query group level, this is a cumulative counter since the Opensearch start time */ - final CounterMetric completions = new CounterMetric(); + public final CounterMetric completions = new CounterMetric(); + + /** + * shard level completions at the query group level, this is a cumulative counter since the Opensearch start time + */ + public final CounterMetric shardCompletions = new CounterMetric(); /** * rejections at the query group level, this is a cumulative counter since the OpenSearch start time */ - final CounterMetric totalRejections = new CounterMetric(); + public final CounterMetric totalRejections = new CounterMetric(); /** * this will track the cumulative failures in a query group @@ -36,7 +41,7 @@ public class QueryGroupState { /** * This will track total number of cancellations in the query group due to all resource type breaches */ - final CounterMetric totalCancellations = new CounterMetric(); + public final CounterMetric totalCancellations = new CounterMetric(); /** * This is used to store the resource type state both for CPU and MEMORY @@ -54,12 +59,20 @@ public QueryGroupState() { /** * - * @return completions in the query group + * @return co-ordinator completions in the query group */ public long getCompletions() { return completions.count(); } + /** + * + * @return shard completions in the query group + */ + public long getShardCompletions() { + return shardCompletions.count(); + } + /** * * @return rejections in the query group @@ -92,9 +105,9 @@ public Map getResourceState() { * This class holds the resource level stats for the query group */ public static class ResourceTypeState { - final ResourceType resourceType; - final CounterMetric cancellations = new CounterMetric(); - final CounterMetric rejections = new CounterMetric(); + public final ResourceType resourceType; + public final CounterMetric cancellations = new CounterMetric(); + public final CounterMetric rejections = new CounterMetric(); private double lastRecordedUsage = 0; public ResourceTypeState(ResourceType resourceType) { diff --git a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java index 2b389c2167778..9d74201de252b 100644 --- a/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java +++ b/server/src/main/java/org/opensearch/wlm/stats/QueryGroupStats.java @@ -91,7 +91,9 @@ public static class QueryGroupStatsHolder implements ToXContentObject, Writeable public static final String REJECTIONS = "rejections"; public static final String TOTAL_CANCELLATIONS = "total_cancellations"; public static final String FAILURES = "failures"; + public static final String SHARD_COMPLETIONS = "shard_completions"; private long completions; + private long shardCompletions; private long rejections; private long failures; private long totalCancellations; @@ -105,11 +107,13 @@ public QueryGroupStatsHolder( long rejections, long failures, long totalCancellations, + long shardCompletions, Map resourceStats ) { this.completions = completions; this.rejections = rejections; this.failures = failures; + this.shardCompletions = shardCompletions; this.totalCancellations = totalCancellations; this.resourceStats = resourceStats; } @@ -119,6 +123,7 @@ public QueryGroupStatsHolder(StreamInput in) throws IOException { this.rejections = in.readVLong(); this.failures = in.readVLong(); this.totalCancellations = in.readVLong(); + this.shardCompletions = in.readVLong(); this.resourceStats = in.readMap((i) -> ResourceType.fromName(i.readString()), ResourceStats::new); } @@ -140,6 +145,7 @@ public static QueryGroupStatsHolder from(QueryGroupState queryGroupState) { statsHolder.rejections = queryGroupState.getTotalRejections(); statsHolder.failures = queryGroupState.getFailures(); statsHolder.totalCancellations = queryGroupState.getTotalCancellations(); + statsHolder.shardCompletions = queryGroupState.getShardCompletions(); statsHolder.resourceStats = resourceStatsMap; return statsHolder; } @@ -155,6 +161,7 @@ public static void writeTo(StreamOutput out, QueryGroupStatsHolder statsHolder) out.writeVLong(statsHolder.rejections); out.writeVLong(statsHolder.failures); out.writeVLong(statsHolder.totalCancellations); + out.writeVLong(statsHolder.shardCompletions); out.writeMap(statsHolder.resourceStats, (o, val) -> o.writeString(val.getName()), ResourceStats::writeTo); } @@ -166,6 +173,7 @@ public void writeTo(StreamOutput out) throws IOException { @Override public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException { builder.field(COMPLETIONS, completions); + builder.field(SHARD_COMPLETIONS, shardCompletions); builder.field(REJECTIONS, rejections); builder.field(FAILURES, failures); builder.field(TOTAL_CANCELLATIONS, totalCancellations); @@ -187,6 +195,7 @@ public boolean equals(Object o) { QueryGroupStatsHolder that = (QueryGroupStatsHolder) o; return completions == that.completions && rejections == that.rejections + && shardCompletions == that.shardCompletions && Objects.equals(resourceStats, that.resourceStats) && failures == that.failures && totalCancellations == that.totalCancellations; @@ -194,7 +203,7 @@ public boolean equals(Object o) { @Override public int hashCode() { - return Objects.hash(completions, rejections, totalCancellations, failures, resourceStats); + return Objects.hash(completions, shardCompletions, rejections, totalCancellations, failures, resourceStats); } } diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index b23d9ff342139..19f7bf48d8421 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -47,6 +47,8 @@ public Map constructQueryGroupLevelUsa // Iterate over each QueryGroup entry for (Map.Entry> queryGroupEntry : tasksByQueryGroup.entrySet()) { + // refresh the resource stats + taskResourceTrackingService.refreshResourceStats(queryGroupEntry.getValue().toArray(new QueryGroupTask[0])); // Compute the QueryGroup resource usage final Map queryGroupUsage = new EnumMap<>(ResourceType.class); for (ResourceType resourceType : TRACKED_RESOURCES) { diff --git a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java index 62aec845a3333..6f132683171f3 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java +++ b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java @@ -39,6 +39,8 @@ import org.opensearch.test.transport.MockTransportService; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.QueryGroupService; +import org.opensearch.wlm.QueryGroupTask; import org.opensearch.wlm.ResourceType; import org.junit.After; import org.junit.Before; @@ -75,10 +77,12 @@ public class SearchBackpressureServiceTests extends OpenSearchTestCase { MockTransportService transportService; TaskManager taskManager; ThreadPool threadPool; + QueryGroupService queryGroupService; @Before public void setup() { threadPool = new TestThreadPool(getClass().getName()); + queryGroupService = mock(QueryGroupService.class); transportService = MockTransportService.createNewService(Settings.EMPTY, Version.CURRENT, threadPool, NoopTracer.INSTANCE); transportService.start(); transportService.acceptIncomingRequests(); @@ -120,9 +124,12 @@ public void testIsNodeInDuress() { new NodeDuressTrackers(duressTrackers), new TaskResourceUsageTrackers(), new TaskResourceUsageTrackers(), - taskManager + taskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + // Node not in duress. cpuUsage.set(0.0); heapUsage.set(0.0); @@ -163,9 +170,12 @@ public void testTrackerStateUpdateOnSearchTaskCompletion() { new NodeDuressTrackers(new EnumMap<>(ResourceType.class)), taskResourceUsageTrackers, new TaskResourceUsageTrackers(), - taskManager + taskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + for (int i = 0; i < 100; i++) { // service.onTaskCompleted(new SearchTask(1, "test", "test", () -> "Test", TaskId.EMPTY_TASK_ID, new HashMap<>())); service.onTaskCompleted(createMockTaskWithResourceStats(SearchTask.class, 100, 200, i)); @@ -194,9 +204,12 @@ public void testTrackerStateUpdateOnSearchShardTaskCompletion() { new NodeDuressTrackers(new EnumMap<>(ResourceType.class)), new TaskResourceUsageTrackers(), taskResourceUsageTrackers, - taskManager + taskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + // Record task completions to update the tracker state. Tasks other than SearchTask & SearchShardTask are ignored. service.onTaskCompleted(createMockTaskWithResourceStats(CancellableTask.class, 100, 200, 101)); for (int i = 0; i < 100; i++) { @@ -246,9 +259,12 @@ public void testSearchTaskInFlightCancellation() { new NodeDuressTrackers(duressTrackers), taskResourceUsageTrackers, new TaskResourceUsageTrackers(), - mockTaskManager + mockTaskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + // Run two iterations so that node is marked 'in duress' from the third iteration onwards. service.doRun(); service.doRun(); @@ -261,7 +277,7 @@ public void testSearchTaskInFlightCancellation() { when(settings.getSearchTaskSettings()).thenReturn(searchTaskSettings); // Create a mix of low and high resource usage SearchTasks (50 low + 25 high resource usage tasks). - Map activeSearchTasks = new HashMap<>(); + Map activeSearchTasks = new HashMap<>(); for (long i = 0; i < 75; i++) { if (i % 3 == 0) { activeSearchTasks.put(i, createMockTaskWithResourceStats(SearchTask.class, 500, taskHeapUsageBytes, i)); @@ -269,6 +285,7 @@ public void testSearchTaskInFlightCancellation() { activeSearchTasks.put(i, createMockTaskWithResourceStats(SearchTask.class, 100, taskHeapUsageBytes, i)); } } + activeSearchTasks.values().forEach(task -> task.setQueryGroupId(threadPool.getThreadContext())); doReturn(activeSearchTasks).when(mockTaskResourceTrackingService).getResourceAwareTasks(); // There are 25 SearchTasks eligible for cancellation but only 5 will be cancelled (burst limit). @@ -337,9 +354,12 @@ public void testSearchShardTaskInFlightCancellation() { nodeDuressTrackers, new TaskResourceUsageTrackers(), taskResourceUsageTrackers, - mockTaskManager + mockTaskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + // Run two iterations so that node is marked 'in duress' from the third iteration onwards. service.doRun(); service.doRun(); @@ -352,7 +372,7 @@ public void testSearchShardTaskInFlightCancellation() { when(settings.getSearchShardTaskSettings()).thenReturn(searchShardTaskSettings); // Create a mix of low and high resource usage tasks (60 low + 15 high resource usage tasks). - Map activeSearchShardTasks = new HashMap<>(); + Map activeSearchShardTasks = new HashMap<>(); for (long i = 0; i < 75; i++) { if (i % 5 == 0) { activeSearchShardTasks.put(i, createMockTaskWithResourceStats(SearchShardTask.class, 500, taskHeapUsageBytes, i)); @@ -360,6 +380,7 @@ public void testSearchShardTaskInFlightCancellation() { activeSearchShardTasks.put(i, createMockTaskWithResourceStats(SearchShardTask.class, 100, taskHeapUsageBytes, i)); } } + activeSearchShardTasks.values().forEach(task -> task.setQueryGroupId(threadPool.getThreadContext())); doReturn(activeSearchShardTasks).when(mockTaskResourceTrackingService).getResourceAwareTasks(); // There are 15 SearchShardTasks eligible for cancellation but only 10 will be cancelled (burst limit). @@ -437,9 +458,12 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { nodeDuressTrackers, taskResourceUsageTrackers, new TaskResourceUsageTrackers(), - mockTaskManager + mockTaskManager, + queryGroupService ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); + service.doRun(); service.doRun(); @@ -449,7 +473,7 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { when(settings.getSearchTaskSettings()).thenReturn(searchTaskSettings); // Create a mix of low and high resource usage tasks (60 low + 15 high resource usage tasks). - Map activeSearchTasks = new HashMap<>(); + Map activeSearchTasks = new HashMap<>(); for (long i = 0; i < 75; i++) { if (i % 5 == 0) { activeSearchTasks.put(i, createMockTaskWithResourceStats(SearchTask.class, 500, 800, i)); @@ -457,6 +481,7 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { activeSearchTasks.put(i, createMockTaskWithResourceStats(SearchTask.class, 100, 800, i)); } } + activeSearchTasks.values().forEach(task -> task.setQueryGroupId(threadPool.getThreadContext())); doReturn(activeSearchTasks).when(mockTaskResourceTrackingService).getResourceAwareTasks(); // this will trigger cancellation but these cancellation should only be cpu based @@ -531,10 +556,12 @@ public void testNonCancellationWhenSearchTrafficIsNotQualifyingForCancellation() nodeDuressTrackers, taskResourceUsageTrackers, new TaskResourceUsageTrackers(), - mockTaskManager + mockTaskManager, + queryGroupService ) ); + when(queryGroupService.shouldSBPHandle(any())).thenReturn(true); when(service.isHeapUsageDominatedBySearch(anyList(), anyDouble())).thenReturn(false); service.doRun(); @@ -546,15 +573,16 @@ public void testNonCancellationWhenSearchTrafficIsNotQualifyingForCancellation() when(settings.getSearchTaskSettings()).thenReturn(searchTaskSettings); // Create a mix of low and high resource usage tasks (60 low + 15 high resource usage tasks). - Map activeSearchTasks = new HashMap<>(); + Map activeSearchTasks = new HashMap<>(); for (long i = 0; i < 75; i++) { - Class taskType = randomBoolean() ? SearchTask.class : SearchShardTask.class; + Class taskType = randomBoolean() ? SearchTask.class : SearchShardTask.class; if (i % 5 == 0) { activeSearchTasks.put(i, createMockTaskWithResourceStats(taskType, 500, 800, i)); } else { activeSearchTasks.put(i, createMockTaskWithResourceStats(taskType, 100, 800, i)); } } + activeSearchTasks.values().forEach(task -> task.setQueryGroupId(threadPool.getThreadContext())); doReturn(activeSearchTasks).when(mockTaskResourceTrackingService).getResourceAwareTasks(); // this will trigger cancellation but the cancellation should not happen as the node is not is duress because of search traffic diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java new file mode 100644 index 0000000000000..c5cf0dac4f807 --- /dev/null +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java @@ -0,0 +1,489 @@ +/* + * SPDX-License-Identifier: Apache-2.0 + * + * The OpenSearch Contributors require contributions made to + * this file be licensed under the Apache-2.0 license or a + * compatible open source license. + */ + +package org.opensearch.wlm; + +import org.opensearch.action.search.SearchTask; +import org.opensearch.cluster.ClusterChangedEvent; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; +import org.opensearch.search.backpressure.trackers.NodeDuressTrackers; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.Scheduler; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; +import org.opensearch.wlm.cancellation.TaskSelectionStrategy; +import org.opensearch.wlm.stats.QueryGroupState; +import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService; + +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; +import java.util.function.BooleanSupplier; + +import org.mockito.ArgumentCaptor; +import org.mockito.Mockito; + +import static org.opensearch.wlm.tracker.ResourceUsageCalculatorTests.createMockTaskWithResourceStats; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +public class QueryGroupServiceTests extends OpenSearchTestCase { + private QueryGroupService queryGroupService; + private QueryGroupTaskCancellationService mockCancellationService; + private ClusterService mockClusterService; + private ThreadPool mockThreadPool; + private WorkloadManagementSettings mockWorkloadManagementSettings; + private Scheduler.Cancellable mockScheduledFuture; + private Map mockQueryGroupStateMap; + NodeDuressTrackers mockNodeDuressTrackers; + QueryGroupsStateAccessor mockQueryGroupsStateAccessor; + + public void setUp() throws Exception { + super.setUp(); + mockClusterService = Mockito.mock(ClusterService.class); + mockThreadPool = Mockito.mock(ThreadPool.class); + mockScheduledFuture = Mockito.mock(Scheduler.Cancellable.class); + mockWorkloadManagementSettings = Mockito.mock(WorkloadManagementSettings.class); + mockQueryGroupStateMap = new HashMap<>(); + mockNodeDuressTrackers = Mockito.mock(NodeDuressTrackers.class); + mockCancellationService = Mockito.mock(TestQueryGroupCancellationService.class); + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + new HashSet<>(), + new HashSet<>() + ); + } + + public void tearDown() throws Exception { + super.tearDown(); + mockThreadPool.shutdown(); + } + + public void testClusterChanged() { + ClusterChangedEvent mockClusterChangedEvent = Mockito.mock(ClusterChangedEvent.class); + ClusterState mockPreviousClusterState = Mockito.mock(ClusterState.class); + ClusterState mockClusterState = Mockito.mock(ClusterState.class); + Metadata mockPreviousMetadata = Mockito.mock(Metadata.class); + Metadata mockMetadata = Mockito.mock(Metadata.class); + QueryGroup addedQueryGroup = new QueryGroup( + "addedQueryGroup", + "4242", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5)), + 1L + ); + QueryGroup deletedQueryGroup = new QueryGroup( + "deletedQueryGroup", + "4241", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.MEMORY, 0.5)), + 1L + ); + Map previousQueryGroups = new HashMap<>(); + previousQueryGroups.put("4242", addedQueryGroup); + Map currentQueryGroups = new HashMap<>(); + currentQueryGroups.put("4241", deletedQueryGroup); + + when(mockClusterChangedEvent.previousState()).thenReturn(mockPreviousClusterState); + when(mockClusterChangedEvent.state()).thenReturn(mockClusterState); + when(mockPreviousClusterState.metadata()).thenReturn(mockPreviousMetadata); + when(mockClusterState.metadata()).thenReturn(mockMetadata); + when(mockPreviousMetadata.queryGroups()).thenReturn(previousQueryGroups); + when(mockMetadata.queryGroups()).thenReturn(currentQueryGroups); + queryGroupService.clusterChanged(mockClusterChangedEvent); + + Set currentQueryGroupsExpected = Set.of(currentQueryGroups.get("4241")); + Set previousQueryGroupsExpected = Set.of(previousQueryGroups.get("4242")); + + assertEquals(currentQueryGroupsExpected, queryGroupService.getActiveQueryGroups()); + assertEquals(previousQueryGroupsExpected, queryGroupService.getDeletedQueryGroups()); + } + + public void testDoStart_SchedulesTask() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + when(mockWorkloadManagementSettings.getQueryGroupServiceRunInterval()).thenReturn(TimeValue.timeValueSeconds(1)); + queryGroupService.doStart(); + Mockito.verify(mockThreadPool).scheduleWithFixedDelay(any(Runnable.class), any(TimeValue.class), eq(ThreadPool.Names.GENERIC)); + } + + public void testDoStop_CancelsScheduledTask() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + when(mockThreadPool.scheduleWithFixedDelay(any(), any(), any())).thenReturn(mockScheduledFuture); + queryGroupService.doStart(); + queryGroupService.doStop(); + Mockito.verify(mockScheduledFuture).cancel(); + } + + public void testDoRun_WhenModeEnabled() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + when(mockNodeDuressTrackers.isNodeInDuress()).thenReturn(true); + // Call the method + queryGroupService.doRun(); + + // Verify that refreshQueryGroups was called + + // Verify that cancelTasks was called with a BooleanSupplier + ArgumentCaptor booleanSupplierCaptor = ArgumentCaptor.forClass(BooleanSupplier.class); + Mockito.verify(mockCancellationService).cancelTasks(booleanSupplierCaptor.capture(), any(), any()); + + // Assert the behavior of the BooleanSupplier + BooleanSupplier capturedSupplier = booleanSupplierCaptor.getValue(); + assertTrue(capturedSupplier.getAsBoolean()); + + } + + public void testDoRun_WhenModeDisabled() { + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED); + when(mockNodeDuressTrackers.isNodeInDuress()).thenReturn(false); + queryGroupService.doRun(); + // Verify that refreshQueryGroups was called + + Mockito.verify(mockCancellationService, never()).cancelTasks(any(), any(), any()); + + } + + public void testRejectIfNeeded_whenQueryGroupIdIsNullOrDefaultOne() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.10)), + 1L + ); + Set activeQueryGroups = new HashSet<>() { + { + add(testQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + mockQueryGroupStateMap.put("queryGroupId1", new QueryGroupState()); + + Map spyMap = spy(mockQueryGroupStateMap); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + queryGroupService.rejectIfNeeded(null); + + verify(spyMap, never()).get(any()); + + queryGroupService.rejectIfNeeded(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get()); + verify(spyMap, never()).get(any()); + } + + public void testRejectIfNeeded_whenQueryGroupIsSoftMode() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.SOFT, Map.of(ResourceType.CPU, 0.10)), + 1L + ); + Set activeQueryGroups = new HashSet<>() { + { + add(testQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + QueryGroupState spyState = spy(new QueryGroupState()); + mockQueryGroupStateMap.put("queryGroupId1", spyState); + + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + + Map spyMap = spy(mockQueryGroupStateMap); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + queryGroupService.rejectIfNeeded("queryGroupId1"); + + verify(spyState, never()).getResourceState(); + } + + public void testRejectIfNeeded_whenQueryGroupIsEnforcedMode_andNotBreaching() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.10)), + 1L + ); + QueryGroup spuQueryGroup = spy(testQueryGroup); + Set activeQueryGroups = new HashSet<>() { + { + add(spuQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + QueryGroupState queryGroupState = new QueryGroupState(); + queryGroupState.getResourceState().get(ResourceType.CPU).setLastRecordedUsage(0.05); + + mockQueryGroupStateMap.put("queryGroupId1", queryGroupState); + + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + when(mockWorkloadManagementSettings.getNodeLevelCpuRejectionThreshold()).thenReturn(0.8); + queryGroupService.rejectIfNeeded("queryGroupId1"); + + // verify the check to compare the current usage and limit + // this should happen 3 times => 2 to check whether the resource limit has the TRACKED resource type and 1 to get the value + verify(spuQueryGroup, times(3)).getResourceLimits(); + assertEquals(0, queryGroupState.getResourceState().get(ResourceType.CPU).rejections.count()); + assertEquals(0, queryGroupState.totalRejections.count()); + } + + public void testRejectIfNeeded_whenQueryGroupIsEnforcedMode_andBreaching() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment( + MutableQueryGroupFragment.ResiliencyMode.ENFORCED, + Map.of(ResourceType.CPU, 0.10, ResourceType.MEMORY, 0.10) + ), + 1L + ); + QueryGroup spuQueryGroup = spy(testQueryGroup); + Set activeQueryGroups = new HashSet<>() { + { + add(spuQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + QueryGroupState queryGroupState = new QueryGroupState(); + queryGroupState.getResourceState().get(ResourceType.CPU).setLastRecordedUsage(0.18); + queryGroupState.getResourceState().get(ResourceType.MEMORY).setLastRecordedUsage(0.18); + QueryGroupState spyState = spy(queryGroupState); + + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + + mockQueryGroupStateMap.put("queryGroupId1", spyState); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + assertThrows(OpenSearchRejectedExecutionException.class, () -> queryGroupService.rejectIfNeeded("queryGroupId1")); + + // verify the check to compare the current usage and limit + // this should happen 3 times => 1 to check whether the resource limit has the TRACKED resource type and 1 to get the value + // because it will break out of the loop since the limits are breached + verify(spuQueryGroup, times(2)).getResourceLimits(); + assertEquals( + 1, + queryGroupState.getResourceState().get(ResourceType.CPU).rejections.count() + queryGroupState.getResourceState() + .get(ResourceType.MEMORY).rejections.count() + ); + assertEquals(1, queryGroupState.totalRejections.count()); + } + + public void testRejectIfNeeded_whenFeatureIsNotEnabled() { + QueryGroup testQueryGroup = new QueryGroup( + "testQueryGroup", + "queryGroupId1", + new MutableQueryGroupFragment(MutableQueryGroupFragment.ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.10)), + 1L + ); + Set activeQueryGroups = new HashSet<>() { + { + add(testQueryGroup); + } + }; + mockQueryGroupStateMap = new HashMap<>(); + mockQueryGroupStateMap.put("queryGroupId1", new QueryGroupState()); + + Map spyMap = spy(mockQueryGroupStateMap); + + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + new HashSet<>() + ); + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED); + + queryGroupService.rejectIfNeeded(testQueryGroup.get_id()); + verify(spyMap, never()).get(any()); + } + + public void testOnTaskCompleted() { + Task task = createMockTaskWithResourceStats(SearchTask.class, 100, 200, 0, 12); + mockThreadPool = new TestThreadPool("queryGroupServiceTests"); + mockThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, "testId"); + QueryGroupState queryGroupState = new QueryGroupState(); + mockQueryGroupStateMap.put("testId", queryGroupState); + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + new HashSet<>() { + { + add( + new QueryGroup( + "testQueryGroup", + "testId", + new MutableQueryGroupFragment( + MutableQueryGroupFragment.ResiliencyMode.ENFORCED, + Map.of(ResourceType.CPU, 0.10, ResourceType.MEMORY, 0.10) + ), + 1L + ) + ); + } + }, + new HashSet<>() + ); + + ((QueryGroupTask) task).setQueryGroupId(mockThreadPool.getThreadContext()); + queryGroupService.onTaskCompleted(task); + + assertEquals(1, queryGroupState.completions.count()); + + // test non QueryGroupTask + task = new Task(1, "simple", "test", "mock task", null, null); + queryGroupService.onTaskCompleted(task); + + // It should still be 1 + assertEquals(1, queryGroupState.completions.count()); + + mockThreadPool.shutdown(); + } + + public void testShouldSBPHandle() { + QueryGroupTask task = createMockTaskWithResourceStats(SearchTask.class, 100, 200, 0, 12); + QueryGroupState queryGroupState = new QueryGroupState(); + Set activeQueryGroups = new HashSet<>(); + mockQueryGroupStateMap.put("testId", queryGroupState); + mockQueryGroupsStateAccessor = new QueryGroupsStateAccessor(mockQueryGroupStateMap); + queryGroupService = new QueryGroupService( + mockCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + mockNodeDuressTrackers, + mockQueryGroupsStateAccessor, + activeQueryGroups, + Collections.emptySet() + ); + + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + + // Default queryGroupId + mockThreadPool = new TestThreadPool("queryGroupServiceTests"); + mockThreadPool.getThreadContext() + .putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get()); + task.setQueryGroupId(mockThreadPool.getThreadContext()); + assertTrue(queryGroupService.shouldSBPHandle(task)); + + mockThreadPool.shutdownNow(); + + // invalid queryGroup task + mockThreadPool = new TestThreadPool("queryGroupServiceTests"); + mockThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, "testId"); + task.setQueryGroupId(mockThreadPool.getThreadContext()); + assertTrue(queryGroupService.shouldSBPHandle(task)); + + // Valid query group task but wlm not enabled + when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.DISABLED); + activeQueryGroups.add( + new QueryGroup( + "testQueryGroup", + "testId", + new MutableQueryGroupFragment( + MutableQueryGroupFragment.ResiliencyMode.ENFORCED, + Map.of(ResourceType.CPU, 0.10, ResourceType.MEMORY, 0.10) + ), + 1L + ) + ); + assertTrue(queryGroupService.shouldSBPHandle(task)); + + } + + // This is needed to test the behavior of QueryGroupService#doRun method + static class TestQueryGroupCancellationService extends QueryGroupTaskCancellationService { + public TestQueryGroupCancellationService( + WorkloadManagementSettings workloadManagementSettings, + TaskSelectionStrategy taskSelectionStrategy, + QueryGroupResourceUsageTrackerService resourceUsageTrackerService, + QueryGroupsStateAccessor queryGroupsStateAccessor, + Collection activeQueryGroups, + Collection deletedQueryGroups + ) { + super(workloadManagementSettings, taskSelectionStrategy, resourceUsageTrackerService, queryGroupsStateAccessor); + } + + @Override + public void cancelTasks( + BooleanSupplier isNodeInDuress, + Collection activeQueryGroups, + Collection deletedQueryGroups + ) { + + } + } +} diff --git a/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportInterceptorTests.java b/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportInterceptorTests.java index 4668b845150a9..d4cd7b79455a3 100644 --- a/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportInterceptorTests.java +++ b/server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportInterceptorTests.java @@ -8,24 +8,56 @@ package org.opensearch.wlm; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; import org.opensearch.transport.TransportRequest; import org.opensearch.transport.TransportRequestHandler; import org.opensearch.wlm.WorkloadManagementTransportInterceptor.RequestHandler; +import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; + +import java.util.Collections; import static org.opensearch.threadpool.ThreadPool.Names.SAME; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class WorkloadManagementTransportInterceptorTests extends OpenSearchTestCase { - + private QueryGroupTaskCancellationService mockTaskCancellationService; + private ClusterService mockClusterService; + private ThreadPool mockThreadPool; + private WorkloadManagementSettings mockWorkloadManagementSettings; private ThreadPool threadPool; private WorkloadManagementTransportInterceptor sut; + private QueryGroupsStateAccessor stateAccessor; public void setUp() throws Exception { super.setUp(); + mockTaskCancellationService = mock(QueryGroupTaskCancellationService.class); + mockClusterService = mock(ClusterService.class); + mockThreadPool = mock(ThreadPool.class); + mockWorkloadManagementSettings = mock(WorkloadManagementSettings.class); threadPool = new TestThreadPool(getTestName()); - sut = new WorkloadManagementTransportInterceptor(threadPool, new QueryGroupService()); + stateAccessor = new QueryGroupsStateAccessor(); + + ClusterState state = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + when(mockClusterService.state()).thenReturn(state); + when(state.metadata()).thenReturn(metadata); + when(metadata.queryGroups()).thenReturn(Collections.emptyMap()); + sut = new WorkloadManagementTransportInterceptor( + threadPool, + new QueryGroupService( + mockTaskCancellationService, + mockClusterService, + mockThreadPool, + mockWorkloadManagementSettings, + stateAccessor + ) + ); } public void tearDown() throws Exception { diff --git a/server/src/test/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationServiceTests.java b/server/src/test/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationServiceTests.java index f7a49235efc69..13e8e2c527073 100644 --- a/server/src/test/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/cancellation/QueryGroupTaskCancellationServiceTests.java @@ -17,12 +17,16 @@ import org.opensearch.wlm.MutableQueryGroupFragment.ResiliencyMode; import org.opensearch.wlm.QueryGroupLevelResourceUsageView; import org.opensearch.wlm.QueryGroupTask; +import org.opensearch.wlm.QueryGroupsStateAccessor; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WlmMode; import org.opensearch.wlm.WorkloadManagementSettings; +import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService; import org.opensearch.wlm.tracker.ResourceUsageCalculatorTrackerServiceTests.TestClock; import org.junit.Before; +import java.util.ArrayList; import java.util.Collection; import java.util.Collections; import java.util.HashMap; @@ -31,7 +35,9 @@ import java.util.Map; import java.util.Set; import java.util.stream.Collectors; +import java.util.stream.IntStream; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -47,6 +53,7 @@ public class QueryGroupTaskCancellationServiceTests extends OpenSearchTestCase { private QueryGroupTaskCancellationService taskCancellation; private WorkloadManagementSettings workloadManagementSettings; private QueryGroupResourceUsageTrackerService resourceUsageTrackerService; + private QueryGroupsStateAccessor stateAccessor; @Before public void setup() { @@ -59,12 +66,13 @@ public void setup() { when(workloadManagementSettings.getNodeLevelCpuCancellationThreshold()).thenReturn(0.9); when(workloadManagementSettings.getNodeLevelMemoryCancellationThreshold()).thenReturn(0.9); resourceUsageTrackerService = mock(QueryGroupResourceUsageTrackerService.class); + stateAccessor = mock(QueryGroupsStateAccessor.class); + when(stateAccessor.getQueryGroupState(any())).thenReturn(new QueryGroupState()); taskCancellation = new QueryGroupTaskCancellationService( workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); } @@ -138,7 +146,7 @@ public void testGetCancellableTasksFrom_returnsTasksWhenBreachingThresholdForMem activeQueryGroups.add(queryGroup1); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); @@ -187,11 +195,10 @@ public void testGetCancellableTasksFrom_filtersQueryGroupCorrectly() { workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.SOFT); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.SOFT, activeQueryGroups); assertEquals(0, cancellableTasksFrom.size()); } @@ -219,19 +226,19 @@ public void testCancelTasks_cancelsGivenTasks() { workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); when(resourceUsageTrackerService.constructQueryGroupLevelUsageViews()).thenReturn(queryGroupLevelViews); - taskCancellation.cancelTasks(() -> false); + when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + taskCancellation.cancelTasks(() -> false, activeQueryGroups, deletedQueryGroups); assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); } @@ -281,13 +288,11 @@ public void testCancelTasks_cancelsTasksFromDeletedQueryGroups() { workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); - taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); @@ -298,7 +303,8 @@ public void testCancelTasks_cancelsTasksFromDeletedQueryGroups() { assertEquals(1001, cancellableTasksFromDeletedQueryGroups.get(1).getTask().getId()); when(resourceUsageTrackerService.constructQueryGroupLevelUsageViews()).thenReturn(queryGroupLevelViews); - taskCancellation.cancelTasks(() -> true); + when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + taskCancellation.cancelTasks(() -> true, activeQueryGroups, deletedQueryGroups); assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); @@ -352,12 +358,11 @@ public void testCancelTasks_does_not_cancelTasksFromDeletedQueryGroups_whenNodeN workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); @@ -368,7 +373,8 @@ public void testCancelTasks_does_not_cancelTasksFromDeletedQueryGroups_whenNodeN assertEquals(1001, cancellableTasksFromDeletedQueryGroups.get(1).getTask().getId()); when(resourceUsageTrackerService.constructQueryGroupLevelUsageViews()).thenReturn(queryGroupLevelViews); - taskCancellation.cancelTasks(() -> false); + when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + taskCancellation.cancelTasks(() -> false, activeQueryGroups, deletedQueryGroups); assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); @@ -411,24 +417,24 @@ public void testCancelTasks_cancelsGivenTasks_WhenNodeInDuress() { workloadManagementSettings, new MaximumResourceTaskSelectionStrategy(), resourceUsageTrackerService, - activeQueryGroups, - deletedQueryGroups + stateAccessor ); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, cancellableTasksFrom.size()); assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); - List cancellableTasksFrom1 = taskCancellation.getAllCancellableTasks(ResiliencyMode.SOFT); + List cancellableTasksFrom1 = taskCancellation.getAllCancellableTasks(ResiliencyMode.SOFT, activeQueryGroups); assertEquals(2, cancellableTasksFrom1.size()); assertEquals(5678, cancellableTasksFrom1.get(0).getTask().getId()); assertEquals(8765, cancellableTasksFrom1.get(1).getTask().getId()); when(resourceUsageTrackerService.constructQueryGroupLevelUsageViews()).thenReturn(queryGroupLevelViews); - taskCancellation.cancelTasks(() -> true); + when(workloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED); + taskCancellation.cancelTasks(() -> true, activeQueryGroups, deletedQueryGroups); assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); assertTrue(cancellableTasksFrom1.get(0).getTask().isCancelled()); @@ -456,7 +462,7 @@ public void testGetAllCancellableTasks_ReturnsNoTasksWhenNotBreachingThresholds( activeQueryGroups.add(queryGroup1); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List allCancellableTasks = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List allCancellableTasks = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertTrue(allCancellableTasks.isEmpty()); } @@ -479,7 +485,7 @@ public void testGetAllCancellableTasks_ReturnsTasksWhenBreachingThresholds() { activeQueryGroups.add(queryGroup1); taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; - List allCancellableTasks = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED); + List allCancellableTasks = taskCancellation.getAllCancellableTasks(ResiliencyMode.ENFORCED, activeQueryGroups); assertEquals(2, allCancellableTasks.size()); assertEquals(1234, allCancellableTasks.get(0).getTask().getId()); assertEquals(4321, allCancellableTasks.get(1).getTask().getId()); @@ -513,6 +519,48 @@ public void testGetCancellableTasksFrom_doesNotReturnTasksWhenQueryGroupIdNotFou assertEquals(0, cancellableTasksFrom.size()); } + public void testPruneDeletedQueryGroups() { + QueryGroup queryGroup1 = new QueryGroup( + "testQueryGroup1", + queryGroupId1, + new MutableQueryGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.2)), + 1L + ); + QueryGroup queryGroup2 = new QueryGroup( + "testQueryGroup2", + queryGroupId2, + new MutableQueryGroupFragment(ResiliencyMode.ENFORCED, Map.of(ResourceType.CPU, 0.1)), + 1L + ); + List deletedQueryGroups = new ArrayList<>(); + deletedQueryGroups.add(queryGroup1); + deletedQueryGroups.add(queryGroup2); + QueryGroupLevelResourceUsageView resourceUsageView1 = createResourceUsageViewMock(); + + List activeTasks = IntStream.range(0, 5).mapToObj(this::getRandomSearchTask).collect(Collectors.toList()); + when(resourceUsageView1.getActiveTasks()).thenReturn(activeTasks); + + QueryGroupLevelResourceUsageView resourceUsageView2 = createResourceUsageViewMock(); + when(resourceUsageView2.getActiveTasks()).thenReturn(new ArrayList<>()); + + queryGroupLevelViews.put(queryGroupId1, resourceUsageView1); + queryGroupLevelViews.put(queryGroupId2, resourceUsageView2); + + QueryGroupTaskCancellationService taskCancellation = new QueryGroupTaskCancellationService( + workloadManagementSettings, + new MaximumResourceTaskSelectionStrategy(), + resourceUsageTrackerService, + stateAccessor + ); + taskCancellation.queryGroupLevelResourceUsageViews = queryGroupLevelViews; + + taskCancellation.pruneDeletedQueryGroups(deletedQueryGroups); + + assertEquals(1, deletedQueryGroups.size()); + assertEquals(queryGroupId1, deletedQueryGroups.get(0).get_id()); + + } + private QueryGroupLevelResourceUsageView createResourceUsageViewMock() { QueryGroupLevelResourceUsageView mockView = mock(QueryGroupLevelResourceUsageView.class); when(mockView.getActiveTasks()).thenReturn(List.of(getRandomSearchTask(1234), getRandomSearchTask(4321))); diff --git a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java index 0307ff623c408..1127b50399d24 100644 --- a/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java +++ b/server/src/test/java/org/opensearch/wlm/listeners/QueryGroupRequestOperationListenerTests.java @@ -8,6 +8,9 @@ package org.opensearch.wlm.listeners; +import org.opensearch.cluster.ClusterState; +import org.opensearch.cluster.metadata.Metadata; +import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException; import org.opensearch.test.OpenSearchTestCase; @@ -15,12 +18,16 @@ import org.opensearch.threadpool.ThreadPool; import org.opensearch.wlm.QueryGroupService; import org.opensearch.wlm.QueryGroupTask; +import org.opensearch.wlm.QueryGroupsStateAccessor; import org.opensearch.wlm.ResourceType; +import org.opensearch.wlm.WorkloadManagementSettings; +import org.opensearch.wlm.cancellation.QueryGroupTaskCancellationService; import org.opensearch.wlm.stats.QueryGroupState; import org.opensearch.wlm.stats.QueryGroupStats; import java.io.IOException; import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -28,18 +35,24 @@ import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; public class QueryGroupRequestOperationListenerTests extends OpenSearchTestCase { public static final int ITERATIONS = 20; ThreadPool testThreadPool; QueryGroupService queryGroupService; - + private QueryGroupTaskCancellationService taskCancellationService; + private ClusterService mockClusterService; + private WorkloadManagementSettings mockWorkloadManagementSettings; Map queryGroupStateMap; String testQueryGroupId; QueryGroupRequestOperationListener sut; public void setUp() throws Exception { super.setUp(); + taskCancellationService = mock(QueryGroupTaskCancellationService.class); + mockClusterService = mock(ClusterService.class); + mockWorkloadManagementSettings = mock(WorkloadManagementSettings.class); queryGroupStateMap = new HashMap<>(); testQueryGroupId = "safjgagnakg-3r3fads"; testThreadPool = new TestThreadPool("RejectionTestThreadPool"); @@ -77,6 +90,21 @@ public void testValidQueryGroupRequestFailure() throws IOException { 0, 1, 0, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ), + QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(), + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + 0, + 0, + 0, Map.of( ResourceType.CPU, new QueryGroupStats.ResourceStats(0, 0, 0), @@ -93,8 +121,18 @@ public void testValidQueryGroupRequestFailure() throws IOException { public void testMultiThreadedValidQueryGroupRequestFailures() { queryGroupStateMap.put(testQueryGroupId, new QueryGroupState()); - - queryGroupService = new QueryGroupService(queryGroupStateMap); + QueryGroupsStateAccessor accessor = new QueryGroupsStateAccessor(queryGroupStateMap); + setupMockedQueryGroupsFromClusterState(); + queryGroupService = new QueryGroupService( + taskCancellationService, + mockClusterService, + testThreadPool, + mockWorkloadManagementSettings, + null, + accessor, + Collections.emptySet(), + Collections.emptySet() + ); sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); @@ -127,6 +165,21 @@ public void testMultiThreadedValidQueryGroupRequestFailures() { 0, ITERATIONS, 0, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ), + QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(), + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + 0, + 0, + 0, Map.of( ResourceType.CPU, new QueryGroupStats.ResourceStats(0, 0, 0), @@ -149,6 +202,21 @@ public void testInvalidQueryGroupFailure() throws IOException { 0, 0, 0, + 0, + Map.of( + ResourceType.CPU, + new QueryGroupStats.ResourceStats(0, 0, 0), + ResourceType.MEMORY, + new QueryGroupStats.ResourceStats(0, 0, 0) + ) + ), + QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(), + new QueryGroupStats.QueryGroupStatsHolder( + 0, + 0, + 1, + 0, + 0, Map.of( ResourceType.CPU, new QueryGroupStats.ResourceStats(0, 0, 0), @@ -169,12 +237,23 @@ private void assertSuccess( QueryGroupStats expectedStats, String threadContextQG_Id ) { - + QueryGroupsStateAccessor stateAccessor = new QueryGroupsStateAccessor(queryGroupStateMap); try (ThreadContext.StoredContext currentContext = testThreadPool.getThreadContext().stashContext()) { testThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, threadContextQG_Id); queryGroupStateMap.put(testQueryGroupId, new QueryGroupState()); - queryGroupService = new QueryGroupService(queryGroupStateMap); + setupMockedQueryGroupsFromClusterState(); + + queryGroupService = new QueryGroupService( + taskCancellationService, + mockClusterService, + testThreadPool, + mockWorkloadManagementSettings, + null, + stateAccessor, + Collections.emptySet(), + Collections.emptySet() + ); sut = new QueryGroupRequestOperationListener(queryGroupService, testThreadPool); sut.onRequestFailure(null, null); @@ -184,4 +263,12 @@ private void assertSuccess( } } + + private void setupMockedQueryGroupsFromClusterState() { + ClusterState state = mock(ClusterState.class); + Metadata metadata = mock(Metadata.class); + when(mockClusterService.state()).thenReturn(state); + when(state.metadata()).thenReturn(metadata); + when(metadata.queryGroups()).thenReturn(Collections.emptyMap()); + } } diff --git a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java index 576eec7be1888..566c4261d6878 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStateTests.java @@ -23,7 +23,13 @@ public void testRandomQueryGroupsStateUpdates() { for (int i = 0; i < 25; i++) { if (i % 5 == 0) { - updaterThreads.add(new Thread(() -> queryGroupState.completions.inc())); + updaterThreads.add(new Thread(() -> { + if (randomBoolean()) { + queryGroupState.completions.inc(); + } else { + queryGroupState.shardCompletions.inc(); + } + })); } else if (i % 5 == 1) { updaterThreads.add(new Thread(() -> { queryGroupState.totalRejections.inc(); @@ -57,7 +63,7 @@ public void testRandomQueryGroupsStateUpdates() { } }); - assertEquals(5, queryGroupState.getCompletions()); + assertEquals(5, queryGroupState.getCompletions() + queryGroupState.getShardCompletions()); assertEquals(5, queryGroupState.getTotalRejections()); final long sumOfRejectionsDueToResourceTypes = queryGroupState.getResourceState().get(ResourceType.CPU).rejections.count() diff --git a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java index 661c3a7beae40..ac6d19580dacb 100644 --- a/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java +++ b/server/src/test/java/org/opensearch/wlm/stats/QueryGroupStatsTests.java @@ -28,9 +28,10 @@ public void testToXContent() throws IOException { queryGroupId, new QueryGroupStats.QueryGroupStatsHolder( 123456789, + 13, 2, 0, - 13, + 1213718, Map.of(ResourceType.CPU, new QueryGroupStats.ResourceStats(0.3, 13, 2)) ) ); @@ -40,7 +41,7 @@ public void testToXContent() throws IOException { queryGroupStats.toXContent(builder, ToXContent.EMPTY_PARAMS); builder.endObject(); assertEquals( - "{\"query_groups\":{\"afakjklaj304041-afaka\":{\"completions\":123456789,\"rejections\":2,\"failures\":0,\"total_cancellations\":13,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", + "{\"query_groups\":{\"afakjklaj304041-afaka\":{\"completions\":123456789,\"shard_completions\":1213718,\"rejections\":13,\"failures\":2,\"total_cancellations\":0,\"cpu\":{\"current_usage\":0.3,\"cancellations\":13,\"rejections\":2}}}}", builder.toString() ); } @@ -60,6 +61,7 @@ protected QueryGroupStats createTestInstance() { randomNonNegativeLong(), randomNonNegativeLong(), randomNonNegativeLong(), + randomNonNegativeLong(), Map.of( ResourceType.CPU, new QueryGroupStats.ResourceStats(