Kp/qsb cancellation

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 2.x]
 ### Added
 - Fix for hasInitiatedFetching to fix allocation explain and manual reroute APIs (([#14972](https://github.com/opensearch-project/OpenSearch/pull/14972))
+- [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897))
 - [Workload Management] Add queryGroupId to Task ([14708](https://github.com/opensearch-project/OpenSearch/pull/14708))
 - Add setting to ignore throttling nodes for allocation of unassigned primaries in remote restore ([#14991](https://github.com/opensearch-project/OpenSearch/pull/14991))
 - Add basic aggregation support for derived fields ([#14618](https://github.com/opensearch-project/OpenSearch/pull/14618))

server/src/main/java/org/opensearch/cluster/metadata/Metadata.java

@@ -1391,7 +1391,7 @@ public Builder put(final QueryGroup queryGroup) {
             return queryGroups(existing);
         }
 
-        private Map<String, QueryGroup> getQueryGroups() {
+        public Map<String, QueryGroup> getQueryGroups() {
             return Optional.ofNullable(this.customs.get(QueryGroupMetadata.TYPE))
                 .map(o -> (QueryGroupMetadata) o)
                 .map(QueryGroupMetadata::queryGroups)

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

@@ -0,0 +1,42 @@
+/*
+ * 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.search.ResourceType;
+import org.opensearch.tasks.Task;
+
+import java.util.Map;
+import java.util.function.Function;
+
+/**
+ * Helper class for calculating resource usage for different resource types.
+ */
+public class QueryGroupHelper {
+
+    /**
+     * A map that associates each {@link ResourceType} with a function that calculates the resource usage for a given {@link Task}.
+     */
+    private static final Map<ResourceType, Function<Task, Long>> resourceUsageCalculator = Map.of(
+        ResourceType.MEMORY,
+        (task) -> task.getTotalResourceStats().getMemoryInBytes(),
+        ResourceType.CPU,
+        (task) -> task.getTotalResourceStats().getCpuTimeInNanos()
+    );
+
+    /**
+     * Gets the resource usage for a given resource type and task.
+     *
+     * @param resource the resource type
+     * @param task the task for which to calculate resource usage
+     * @return the resource usage
+     */
+    public static long getResourceUsage(ResourceType resource, Task task) {
+        return resourceUsageCalculator.get(resource).apply(task);
+    }
+}

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

@@ -0,0 +1,50 @@
+/*
+ * 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.search.ResourceType;
+import org.opensearch.tasks.Task;
+
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Represents the point in time view of resource usage of a QueryGroup and
+ * has a 1:1 relation with a QueryGroup.
+ * This class holds the resource usage data and the list of active tasks.
+ */
+public class QueryGroupLevelResourceUsageView {
+    // resourceUsage holds the resource usage data for a QueryGroup at a point in time
+    private final Map<ResourceType, Long> resourceUsage;
+    // activeTasks holds the list of active tasks for a QueryGroup at a point in time
+    private final List<Task> activeTasks;
+
+    public QueryGroupLevelResourceUsageView(Map<ResourceType, Long> resourceUsage, List<Task> activeTasks) {
+        this.resourceUsage = resourceUsage;
+        this.activeTasks = activeTasks;
+    }
+
+    /**
+     * Returns the resource usage data.
+     *
+     * @return The map of resource usage data
+     */
+    public Map<ResourceType, Long> getResourceUsageData() {
+        return resourceUsage;
+    }
+
+    /**
+     * Returns the list of active tasks.
+     *
+     * @return The list of active tasks
+     */
+    public List<Task> getActiveTasks() {
+        return activeTasks;
+    }
+}

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

@@ -0,0 +1,82 @@
+/*
+ * 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.search.ResourceType;
+import org.opensearch.tasks.CancellableTask;
+import org.opensearch.tasks.Task;
+import org.opensearch.tasks.TaskCancellation;
+import org.opensearch.wlm.QueryGroupHelper;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.Comparator;
+import java.util.List;
+import java.util.stream.Collectors;
+
+/**
+ * Represents an abstract task selection strategy.
+ * This class implements the TaskSelectionStrategy interface and provides a method to select tasks for cancellation based on a sorting condition.
+ * The specific sorting condition depends on the implementation.
+ */
+public abstract class AbstractTaskSelectionStrategy implements TaskSelectionStrategy {
+
+    /**
+     * Returns a comparator that defines the sorting condition for tasks.
+     * The specific sorting condition depends on the implementation.
+     *
+     * @return The comparator
+     */
+    public abstract Comparator<Task> sortingCondition();
+
+    /**
+     * Selects tasks for cancellation based on the provided limit and resource type.
+     * The tasks are sorted based on the sorting condition and then selected until the accumulated resource usage reaches the limit.
+     *
+     * @param tasks The list of tasks from which to select
+     * @param limit The limit on the accumulated resource usage
+     * @param resourceType The type of resource to consider
+     * @return The list of selected tasks
+     * @throws IllegalArgumentException If the limit is less than zero
+     */
+    @Override
+    public List<TaskCancellation> selectTasksForCancellation(List<Task> tasks, long limit, ResourceType resourceType) {
+        if (limit < 0) {
+            throw new IllegalArgumentException("reduceBy has to be greater than zero");
+        }
+        if (limit == 0) {
+            return Collections.emptyList();
+        }
+
+        List<Task> sortedTasks = tasks.stream().sorted(sortingCondition()).collect(Collectors.toList());
+
+        List<TaskCancellation> selectedTasks = new ArrayList<>();
+        long accumulated = 0;
+
+        for (Task task : sortedTasks) {
+            if (task instanceof CancellableTask) {
+                selectedTasks.add(createTaskCancellation((CancellableTask) task));
+                accumulated += QueryGroupHelper.getResourceUsage(resourceType, task);
+                if (accumulated >= limit) {
+                    break;
+                }
+            }
+        }
+        return selectedTasks;
+    }
+
+    private TaskCancellation createTaskCancellation(CancellableTask task) {
+        // TODO add correct reason and callbacks
+        return new TaskCancellation(task, List.of(new TaskCancellation.Reason("limits exceeded", 5)), List.of(this::callbackOnCancel));
+    }
+
+    private void callbackOnCancel() {
+        // todo Implement callback logic here mostly used for Stats
+    }
+}