Add cancellation framework changes in wlm

CHANGELOG.md

@@ -21,6 +21,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - [Workload Management] Add Get QueryGroup API Logic ([14709](https://github.com/opensearch-project/OpenSearch/pull/14709))
 - [Workload Management] Add Settings for Workload Management feature ([#15028](https://github.com/opensearch-project/OpenSearch/pull/15028))
 - [Workload Management] Add Update QueryGroup API Logic ([#14775](https://github.com/opensearch-project/OpenSearch/pull/14775))
+- [Workload Management] QueryGroup resource cancellation framework changes ([#15651](https://github.com/opensearch-project/OpenSearch/pull/15651))
 - [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897))
 - Support filtering on a large list encoded by bitmap ([#14774](https://github.com/opensearch-project/OpenSearch/pull/14774))
 - Add slice execution listeners to SearchOperationListener interface ([#15153](https://github.com/opensearch-project/OpenSearch/pull/15153))

server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java

@@ -8,7 +8,7 @@
 
 package org.opensearch.wlm;
 
-import org.opensearch.tasks.Task;
+import org.opensearch.wlm.tracker.QueryGroupResourceUsage;
 
 import java.util.List;
 import java.util.Map;
@@ -20,11 +20,11 @@
  */
 public class QueryGroupLevelResourceUsageView {
     // resourceUsage holds the resource usage data for a QueryGroup at a point in time
-    private final Map<ResourceType, Long> resourceUsage;
+    private final Map<ResourceType, QueryGroupResourceUsage> resourceUsage;
     // activeTasks holds the list of active tasks for a QueryGroup at a point in time
-    private final List<Task> activeTasks;
+    private final List<QueryGroupTask> activeTasks;
 
-    public QueryGroupLevelResourceUsageView(Map<ResourceType, Long> resourceUsage, List<Task> activeTasks) {
+    public QueryGroupLevelResourceUsageView(Map<ResourceType, QueryGroupResourceUsage> resourceUsage, List<QueryGroupTask> activeTasks) {
         this.resourceUsage = resourceUsage;
         this.activeTasks = activeTasks;
     }
@@ -34,7 +34,7 @@ public QueryGroupLevelResourceUsageView(Map<ResourceType, Long> resourceUsage, L
      *
      * @return The map of resource usage data
      */
-    public Map<ResourceType, Long> getResourceUsageData() {
+    public Map<ResourceType, QueryGroupResourceUsage> getResourceUsageData() {
         return resourceUsage;
     }
 
@@ -43,7 +43,7 @@ public Map<ResourceType, Long> getResourceUsageData() {
      *
      * @return The list of active tasks
      */
-    public List<Task> getActiveTasks() {
+    public List<QueryGroupTask> getActiveTasks() {
         return activeTasks;
     }
 }

server/src/main/java/org/opensearch/wlm/ResourceType.java

@@ -10,12 +10,9 @@
 
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.core.common.io.stream.StreamOutput;
-import org.opensearch.core.tasks.resourcetracker.ResourceStats;
-import org.opensearch.tasks.Task;
 
 import java.io.IOException;
 import java.util.List;
-import java.util.function.Function;
 
 /**
  * Enum to hold the resource type
@@ -24,18 +21,16 @@
  */
 @PublicApi(since = "2.17.0")
 public enum ResourceType {
-    CPU("cpu", task -> task.getTotalResourceUtilization(ResourceStats.CPU), true),
-    MEMORY("memory", task -> task.getTotalResourceUtilization(ResourceStats.MEMORY), true);
+    CPU("cpu", true),
+    MEMORY("memory", true);
 
     private final String name;
-    private final Function<Task, Long> getResourceUsage;
     private final boolean statsEnabled;
 
     private static List<ResourceType> sortedValues = List.of(CPU, MEMORY);
 
-    ResourceType(String name, Function<Task, Long> getResourceUsage, boolean statsEnabled) {
+    ResourceType(String name, boolean statsEnabled) {
         this.name = name;
-        this.getResourceUsage = getResourceUsage;
         this.statsEnabled = statsEnabled;
     }
 
@@ -61,16 +56,6 @@ public String getName() {
         return name;
     }
 
-    /**
-     * Gets the resource usage for a given resource type and task.
-     *
-     * @param task the task for which to calculate resource usage
-     * @return the resource usage
-     */
-    public long getResourceUsage(Task task) {
-        return getResourceUsage.apply(task);
-    }
-
     public boolean hasStatsEnabled() {
         return statsEnabled;
     }

server/src/main/java/org/opensearch/wlm/cancellation/DefaultTaskCancellation.java

@@ -0,0 +1,244 @@
+/*
+ * 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.cancellation;
+
+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.ResourceType;
+import org.opensearch.wlm.WorkloadManagementSettings;
+import org.opensearch.wlm.tracker.QueryGroupResourceUsage;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.List;
+import java.util.Map;
+import java.util.function.BooleanSupplier;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+
+import static org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService.TRACKED_RESOURCES;
+
+/**
+ * Manages the cancellation of tasks enforced by QueryGroup thresholds on resource usage criteria.
+ * This class utilizes a strategy pattern through {@link DefaultTaskSelectionStrategy} to identify tasks that exceed
+ * predefined resource usage limits and are therefore eligible for cancellation.
+ *
+ * <p>The cancellation process is initiated by evaluating the resource usage of each QueryGroup against its
+ * resource limits. Tasks that contribute to exceeding these limits are selected for cancellation based on the
+ * implemented task selection strategy.</p>
+ *
+ * <p>Instances of this class are configured with a map linking QueryGroup IDs to their corresponding resource usage
+ * views, a set of active QueryGroups, and a task selection strategy. These components collectively facilitate the
+ * identification and cancellation of tasks that threaten to breach QueryGroup resource limits.</p>
+ *
+ * @see DefaultTaskSelectionStrategy
+ * @see QueryGroup
+ * @see ResourceType
+ */
+public class DefaultTaskCancellation {
+    public static final double MIN_VALUE = 1e-9;
+
+    protected final WorkloadManagementSettings workloadManagementSettings;
+    protected final DefaultTaskSelectionStrategy defaultTaskSelectionStrategy;
+    // a map of QueryGroupId to its corresponding QueryGroupLevelResourceUsageView object
+    protected final Map<String, QueryGroupLevelResourceUsageView> queryGroupLevelResourceUsageViews;
+    protected final Collection<QueryGroup> activeQueryGroups;
+    protected final Collection<QueryGroup> deletedQueryGroups;
+    protected BooleanSupplier isNodeInDuress;
+
+    public DefaultTaskCancellation(
+        WorkloadManagementSettings workloadManagementSettings,
+        DefaultTaskSelectionStrategy defaultTaskSelectionStrategy,
+        Map<String, QueryGroupLevelResourceUsageView> queryGroupLevelResourceUsageViews,
+        Collection<QueryGroup> activeQueryGroups,
+        Collection<QueryGroup> deletedQueryGroups,
+        BooleanSupplier isNodeInDuress
+    ) {
+        this.workloadManagementSettings = workloadManagementSettings;
+        this.defaultTaskSelectionStrategy = defaultTaskSelectionStrategy;
+        this.queryGroupLevelResourceUsageViews = queryGroupLevelResourceUsageViews;
+        this.activeQueryGroups = activeQueryGroups;
+        this.deletedQueryGroups = deletedQueryGroups;
+        this.isNodeInDuress = isNodeInDuress;
+    }
+
+    /**
+     * Cancel tasks based on the implemented strategy.
+     */
+    public final void cancelTasks() {
+        // cancel tasks from QueryGroups that are in Enforced mode that are breaching their resource limits
+        cancelTasks(ResiliencyMode.ENFORCED);
+        // if the node is in duress, cancel tasks accordingly.
+        handleNodeDuress();
+    }
+
+    private void handleNodeDuress() {
+        if (!isNodeInDuress.getAsBoolean()) {
+            return;
+        }
+        // List of tasks to be executed in order if the node is in duress
+        List<Consumer<Void>> duressActions = List.of(v -> cancelTasksFromDeletedQueryGroups(), v -> cancelTasks(ResiliencyMode.SOFT));
+
+        for (Consumer<Void> duressAction : duressActions) {
+            if (!isNodeInDuress.getAsBoolean()) {
+                break;
+            }
+            duressAction.accept(null);
+        }
+    }
+
+    private void cancelTasksFromDeletedQueryGroups() {
+        cancelTasks(getAllCancellableTasks(this.deletedQueryGroups));
+    }
+
+    /**
+     * Get all cancellable tasks from the QueryGroups.
+     *
+     * @return List of tasks that can be cancelled
+     */
+    protected List<TaskCancellation> getAllCancellableTasks(ResiliencyMode resiliencyMode) {
+        return getAllCancellableTasks(getQueryGroupsToCancelFrom(resiliencyMode));
+    }
+
+    /**
+     * Get all cancellable tasks from the given QueryGroups.
+     *
+     * @return List of tasks that can be cancelled
+     */
+    protected List<TaskCancellation> getAllCancellableTasks(Collection<QueryGroup> queryGroups) {
+        return queryGroups.stream().flatMap(queryGroup -> getCancellableTasksFrom(queryGroup).stream()).collect(Collectors.toList());
+    }
+
+    /**
+     * returns the list of QueryGroups breaching their resource limits.
+     *
+     * @return List of QueryGroups
+     */
+    private List<QueryGroup> getQueryGroupsToCancelFrom(ResiliencyMode resiliencyMode) {
+        final List<QueryGroup> queryGroupsToCancelFrom = new ArrayList<>();
+
+        for (QueryGroup queryGroup : this.activeQueryGroups) {
+            if (queryGroup.getResiliencyMode() != resiliencyMode) {
+                continue;
+            }
+            Map<ResourceType, QueryGroupResourceUsage> queryGroupResourcesUsage = queryGroupLevelResourceUsageViews.get(queryGroup.get_id())
+                .getResourceUsageData();
+
+            for (ResourceType resourceType : TRACKED_RESOURCES) {
+                if (queryGroup.getResourceLimits().containsKey(resourceType)) {
+                    final QueryGroupResourceUsage queryGroupResourceUsage = queryGroupResourcesUsage.get(resourceType);
+                    if (queryGroupResourceUsage.isBreachingThresholdFor(queryGroup, workloadManagementSettings)) {
+                        queryGroupsToCancelFrom.add(queryGroup);
+                        break;
+                    }
+
+                }
+            }
+        }
+
+        return queryGroupsToCancelFrom;
+    }
+
+    private void cancelTasks(ResiliencyMode resiliencyMode) {
+        cancelTasks(getAllCancellableTasks(resiliencyMode));
+    }
+
+    private void cancelTasks(List<TaskCancellation> cancellableTasks) {
+        cancellableTasks.forEach(TaskCancellation::cancel);
+    }
+
+    /**
+     * Get cancellable tasks from a specific queryGroup.
+     *
+     * @param queryGroup The QueryGroup from which to get cancellable tasks
+     * @return List of tasks that can be cancelled
+     */
+    protected List<TaskCancellation> getCancellableTasksFrom(QueryGroup queryGroup) {
+        return TRACKED_RESOURCES.stream()
+            .filter(resourceType -> shouldCancelTasks(queryGroup, resourceType))
+            .flatMap(resourceType -> getTaskCancellations(queryGroup, resourceType).stream())
+            .collect(Collectors.toList());
+    }
+
+    private boolean shouldCancelTasks(QueryGroup queryGroup, ResourceType resourceType) {
+        if (queryGroup == null || !queryGroupLevelResourceUsageViews.containsKey(queryGroup.get_id())) {
+            return false;
+        }
+        QueryGroupLevelResourceUsageView queryGroupResourceUsageView = queryGroupLevelResourceUsageViews.get(queryGroup.get_id());
+        return queryGroupResourceUsageView.getResourceUsageData()
+            .get(resourceType)
+            .isBreachingThresholdFor(queryGroup, workloadManagementSettings);
+    }
+
+    private List<TaskCancellation> getTaskCancellations(QueryGroup queryGroup, ResourceType resourceType) {
+        List<QueryGroupTask> selectedTasksToCancel = defaultTaskSelectionStrategy.selectTasksForCancellation(
+            queryGroupLevelResourceUsageViews.get(queryGroup.get_id()).getActiveTasks(),
+            getReduceBy(queryGroup, resourceType),
+            resourceType
+        );
+        List<TaskCancellation> taskCancellations = new ArrayList<>();
+        for (QueryGroupTask task : selectedTasksToCancel) {
+            String cancellationReason = createCancellationReason(queryGroup, task, resourceType);
+            taskCancellations.add(createTaskCancellation(task, cancellationReason));
+        }
+        return taskCancellations;
+    }
+
+    private String createCancellationReason(QueryGroup querygroup, QueryGroupTask task, ResourceType resourceType) {
+        Double thresholdInPercent = getThresholdInPercent(querygroup, resourceType);
+        return "[Workload Management] Cancelling Task ID : "
+            + task.getId()
+            + " from QueryGroup ID : "
+            + querygroup.get_id()
+            + " breached the resource limit of : "
+            + thresholdInPercent
+            + " for resource type : "
+            + resourceType.getName();
+    }
+
+    private Double getThresholdInPercent(QueryGroup querygroup, ResourceType resourceType) {
+        return querygroup.getResourceLimits().get(resourceType) * 100;
+    }
+
+    private TaskCancellation createTaskCancellation(CancellableTask task, String cancellationReason) {
+        return new TaskCancellation(task, List.of(new TaskCancellation.Reason(cancellationReason, 5)), List.of(this::callbackOnCancel));
+    }
+
+    protected List<TaskCancellation> getTaskCancellationsForDeletedQueryGroup(QueryGroup queryGroup) {
+        List<QueryGroupTask> tasks = queryGroupLevelResourceUsageViews.get(queryGroup.get_id()).getActiveTasks();
+
+        List<TaskCancellation> taskCancellations = new ArrayList<>();
+        for (QueryGroupTask task : tasks) {
+            String cancellationReason = "[Workload Management] Cancelling Task ID : "
+                + task.getId()
+                + " from QueryGroup ID : "
+                + queryGroup.get_id();
+            taskCancellations.add(createTaskCancellation(task, cancellationReason));
+        }
+        return taskCancellations;
+    }
+
+    private double getReduceBy(QueryGroup queryGroup, ResourceType resourceType) {
+        if (queryGroup.getResourceLimits().get(resourceType) == null
+            || !queryGroupLevelResourceUsageViews.containsKey(queryGroup.get_id())) {
+            return 0;
+        }
+        final QueryGroupLevelResourceUsageView queryGroupLevelResourceUsage = queryGroupLevelResourceUsageViews.get(queryGroup.get_id());
+        final QueryGroupResourceUsage queryGroupResourceUsage = queryGroupLevelResourceUsage.getResourceUsageData().get(resourceType);
+        return queryGroupResourceUsage.getReduceByFor(queryGroup, workloadManagementSettings);
+    }
+
+    private void callbackOnCancel() {
+        // TODO Implement callback logic here mostly used for Stats
+    }
+}