From 3e953a19b309c6c383e0ae08b9aa1f5345dc93ad Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 18 Jul 2024 15:02:01 -0700 Subject: [PATCH] renaming sandbox to querygroup and adjusting code based on merged PRs Signed-off-by: Kiran Prakash --- .../action/search/SearchShardTask.java | 14 +- .../opensearch/action/search/SearchTask.java | 14 +- .../opensearch/cluster/metadata/Metadata.java | 2 +- .../cluster/metadata/QueryGroup.java | 13 +- .../opensearch/cluster/metadata/Sandbox.java | 109 ----------- .../org/opensearch/search/ResourceType.java | 51 ----- .../SearchBackpressureService.java | 13 +- .../trackers/NodeDuressTrackers.java | 2 +- .../QueryGroupLevelResourceUsageView.java} | 36 ++-- .../QueryGroupService.java} | 61 +++--- .../QueryGroupTask.java} | 10 +- .../AbstractTaskSelectionStrategy.java | 6 +- .../cancellation/DefaultTaskCancellation.java | 177 +++++++++++++++++ .../LongestRunningTaskFirstStrategy.java | 2 +- .../ShortestRunningTaskFirstStrategy.java | 2 +- .../cancellation/TaskSelectionStrategy.java | 6 +- .../cancellation/package-info.java | 2 +- .../querygroup/module/QueryGroupModule.java | 31 +++ .../module/package-info.java | 2 +- .../package-info.java | 2 +- ...QueryGroupResourceUsageTrackerService.java | 124 ++++++++++++ .../tracker/QueryGroupUsageTracker.java} | 8 +- .../tracker/package-info.java | 2 +- .../opensearch/search/resourcetypes/CPU.java | 59 ++++++ .../opensearch/search/resourcetypes/JVM.java | 60 ++++++ .../search/resourcetypes/ResourceType.java | 73 +++++++ .../cancellation/DefaultTaskCancellation.java | 170 ----------------- .../sandboxing/module/SandboxModule.java | 31 --- .../resourcetype/CpuTimeResource.java | 37 ---- .../resourcetype/JvmMemoryResource.java | 37 ---- .../resourcetype/SystemResource.java | 48 ----- .../sandboxing/resourcetype/package-info.java | 12 -- .../SandboxResourceUsageTrackerService.java | 126 ------------ .../metadata/QueryGroupMetadataTests.java | 4 +- .../cluster/metadata/QueryGroupTests.java | 12 +- .../SearchBackpressureServiceTests.java | 39 ++-- .../trackers/NodeDuressTrackersTests.java | 28 +-- ...QueryGroupLevelResourceUsageViewTests.java | 66 +++++++ .../DefaultTaskCancellationTests.java | 175 +++++++++++++++++ ...RunningTaskFirstStrategyStrategyTests.java | 2 +- .../cancellation/QueryGroupTestHelpers.java} | 40 ++-- ...RunningTaskFirstStrategyStrategyTests.java | 2 +- .../TaskSelectionStrategyTests.java | 16 +- ...roupResourceUsageTrackerServiceTests.java} | 63 +++--- .../SandboxLevelResourceUsageViewTests.java | 59 ------ .../DefaultTaskCancellationTests.java | 180 ------------------ 46 files changed, 967 insertions(+), 1061 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java delete mode 100644 server/src/main/java/org/opensearch/search/ResourceType.java rename server/src/main/java/org/opensearch/search/{sandboxing/SandboxLevelResourceUsageView.java => querygroup/QueryGroupLevelResourceUsageView.java} (51%) rename server/src/main/java/org/opensearch/search/{sandboxing/SandboxService.java => querygroup/QueryGroupService.java} (54%) rename server/src/main/java/org/opensearch/search/{sandboxing/SandboxTask.java => querygroup/QueryGroupTask.java} (55%) rename server/src/main/java/org/opensearch/search/{sandboxing => querygroup}/cancellation/AbstractTaskSelectionStrategy.java (94%) create mode 100644 server/src/main/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellation.java rename server/src/main/java/org/opensearch/search/{sandboxing => querygroup}/cancellation/LongestRunningTaskFirstStrategy.java (93%) rename server/src/main/java/org/opensearch/search/{sandboxing => querygroup}/cancellation/ShortestRunningTaskFirstStrategy.java (93%) rename server/src/main/java/org/opensearch/search/{sandboxing => querygroup}/cancellation/TaskSelectionStrategy.java (85%) rename server/src/main/java/org/opensearch/search/{sandboxing => querygroup}/cancellation/package-info.java (82%) create mode 100644 server/src/main/java/org/opensearch/search/querygroup/module/QueryGroupModule.java rename server/src/main/java/org/opensearch/search/{sandboxing => querygroup}/module/package-info.java (82%) rename server/src/main/java/org/opensearch/search/{sandboxing => querygroup}/package-info.java (85%) create mode 100644 server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java rename server/src/main/java/org/opensearch/search/{sandboxing/tracker/SandboxUsageTracker.java => querygroup/tracker/QueryGroupUsageTracker.java} (61%) rename server/src/main/java/org/opensearch/search/{sandboxing => querygroup}/tracker/package-info.java (84%) create mode 100644 server/src/main/java/org/opensearch/search/resourcetypes/CPU.java create mode 100644 server/src/main/java/org/opensearch/search/resourcetypes/JVM.java create mode 100644 server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/module/SandboxModule.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResource.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResource.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SystemResource.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/resourcetype/package-info.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java create mode 100644 server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java create mode 100644 server/src/test/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellationTests.java rename server/src/test/java/org/opensearch/search/{sandboxing => querygroup}/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java (94%) rename server/src/test/java/org/opensearch/search/{sandboxing/cancellation/SandboxTestHelpers.java => querygroup/cancellation/QueryGroupTestHelpers.java} (66%) rename server/src/test/java/org/opensearch/search/{sandboxing => querygroup}/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java (94%) rename server/src/test/java/org/opensearch/search/{sandboxing => querygroup}/cancellation/TaskSelectionStrategyTests.java (82%) rename server/src/test/java/org/opensearch/search/{sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java => querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java} (55%) delete mode 100644 server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java delete mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java diff --git a/server/src/main/java/org/opensearch/action/search/SearchShardTask.java b/server/src/main/java/org/opensearch/action/search/SearchShardTask.java index fa935c8792581..1e3adb387a788 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchShardTask.java +++ b/server/src/main/java/org/opensearch/action/search/SearchShardTask.java @@ -37,7 +37,7 @@ import org.opensearch.core.tasks.TaskId; import org.opensearch.search.fetch.ShardFetchSearchRequest; import org.opensearch.search.internal.ShardSearchRequest; -import org.opensearch.search.sandboxing.SandboxTask; +import org.opensearch.search.querygroup.QueryGroupTask; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.SearchBackpressureTask; @@ -51,10 +51,10 @@ * @opensearch.api */ @PublicApi(since = "1.0.0") -public class SearchShardTask extends CancellableTask implements SearchBackpressureTask, SandboxTask { +public class SearchShardTask extends CancellableTask implements SearchBackpressureTask, QueryGroupTask { // generating metadata in a lazy way since source can be quite big private final MemoizedSupplier metadataSupplier; - private String sandboxId; + private String queryGroupId; public SearchShardTask(long id, String type, String action, String description, TaskId parentTaskId, Map headers) { this(id, type, action, description, parentTaskId, headers, () -> ""); @@ -88,12 +88,12 @@ public boolean shouldCancelChildrenOnCancellation() { } @Override - public void setSandboxId(String sandboxId) { - this.sandboxId = sandboxId; + public void setQueryGroupId(String queryGroupId) { + this.queryGroupId = queryGroupId; } @Override - public String getSandboxId() { - return "sandboxId"; + public String getQueryGroupId() { + return "queryGroupId"; } } diff --git a/server/src/main/java/org/opensearch/action/search/SearchTask.java b/server/src/main/java/org/opensearch/action/search/SearchTask.java index d775c39ab477d..70fe32e89c11f 100644 --- a/server/src/main/java/org/opensearch/action/search/SearchTask.java +++ b/server/src/main/java/org/opensearch/action/search/SearchTask.java @@ -35,7 +35,7 @@ import org.opensearch.common.annotation.PublicApi; import org.opensearch.common.unit.TimeValue; import org.opensearch.core.tasks.TaskId; -import org.opensearch.search.sandboxing.SandboxTask; +import org.opensearch.search.querygroup.QueryGroupTask; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.SearchBackpressureTask; @@ -50,11 +50,11 @@ * @opensearch.api */ @PublicApi(since = "1.0.0") -public class SearchTask extends CancellableTask implements SearchBackpressureTask, SandboxTask { +public class SearchTask extends CancellableTask implements SearchBackpressureTask, QueryGroupTask { // generating description in a lazy way since source can be quite big private final Supplier descriptionSupplier; private SearchProgressListener progressListener = SearchProgressListener.NOOP; - private String sandboxId; + private String queryGroupId; public SearchTask( long id, @@ -109,12 +109,12 @@ public boolean shouldCancelChildrenOnCancellation() { return true; } - public void setSandboxId(String sandboxId) { - this.sandboxId = sandboxId; + public void setQueryGroupId(String queryGroupId) { + this.queryGroupId = queryGroupId; } @Override - public String getSandboxId() { - return "sandboxId"; + public String getQueryGroupId() { + return "queryGroupId"; } } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java index df1d75309bdb3..5cdd05ac2a8ff 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -126,7 +126,7 @@ public boolean isSegmentReplicationEnabled(String indexName) { .orElse(false); } - public Map sandboxes() { + public Map queryGroups() { // stub return Collections.emptyMap(); } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java index beaab198073df..0396ddf05b36a 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -17,7 +17,7 @@ import org.opensearch.core.xcontent.ToXContentObject; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.search.ResourceType; +import org.opensearch.search.resourcetypes.ResourceType; import org.joda.time.Instant; import java.io.IOException; @@ -227,6 +227,17 @@ public String get_id() { return _id; } + /** + * Converts the threshold percentage for the given resource type to a long value. + * + * @param resourceType The resource type for which the threshold is to be converted + * @return The threshold value in long format + */ + public long getThresholdInLong(ResourceType resourceType) { + Double thresholdInPercentage = (Double) resourceLimits.get(resourceType); + return resourceType.convertThresholdPercentageToLong(thresholdInPercentage); + } + public long getUpdatedAtInMillis() { return updatedAtInMillis; } diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java deleted file mode 100644 index a9014bf57bf64..0000000000000 --- a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java +++ /dev/null @@ -1,109 +0,0 @@ -/* - * 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.cluster.metadata; - -import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; - -import java.util.Collections; -import java.util.List; - -/** - * Stub class only, this will be added in Kaushal's PR - * This class wil be deleted after that. - */ -@ExperimentalApi -public class Sandbox { - // TODO Kaushal should have implemented hashcode and equals - private SandboxMode mode; - - public SandboxMode getMode() { - return mode; - } - - public ResourceLimit getResourceLimitFor(SystemResource resourceType) { - return null; - } - - public String getName() { - return ""; - } - - public String getId() { - return ""; - } - - public List getResourceLimits() { - return Collections.emptyList(); - } - - /** - * Stub class only, this will be added in Kaushal's PR - * This class wil be deleted after that. - */ - @ExperimentalApi - public class ResourceLimit { - /* - Stub class only - */ - SystemResource resourceType; - Long threshold; - - ResourceLimit(SystemResource resourceType, Long threshold) { - this.resourceType = resourceType; - this.threshold = threshold; - } - - public Long getThresholdInLong() { - return threshold; - } - - public SystemResource getResourceType() { - return resourceType; - } - - public Long getThreshold() { - return threshold; - } - } - - /** - * Stub class only, this will be added in Kaushal's PR - * This class wil be deleted after that. - */ - @ExperimentalApi - public enum SandboxMode { - SOFT("soft"), - ENFORCED("enforced"), - MONITOR("monitor"); - - private final String name; - - SandboxMode(String mode) { - this.name = mode; - } - - public String getName() { - return name; - } - - public static SandboxMode fromName(String s) { - switch (s) { - case "soft": - return SOFT; - case "enforced": - return ENFORCED; - case "monitor": - return MONITOR; - default: - throw new IllegalArgumentException("Invalid value for SandboxMode: " + s); - } - } - } -} diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java deleted file mode 100644 index fe5ce4dd2bb50..0000000000000 --- a/server/src/main/java/org/opensearch/search/ResourceType.java +++ /dev/null @@ -1,51 +0,0 @@ -/* - * 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.search; - -import org.opensearch.common.annotation.PublicApi; -import org.opensearch.core.common.io.stream.StreamOutput; - -import java.io.IOException; - -/** - * Enum to hold the resource type - */ -@PublicApi(since = "2.x") -public enum ResourceType { - CPU("cpu"), - MEMORY("memory"); - - private final String name; - - ResourceType(String name) { - this.name = name; - } - - /** - * The string match here is case-sensitive - * @param s name matching the resource type name - * @return a {@link ResourceType} - */ - public static ResourceType fromName(String s) { - for (ResourceType resourceType : values()) { - if (resourceType.getName().equals(s)) { - return resourceType; - } - } - throw new IllegalArgumentException("Unknown resource type: [" + s + "]"); - } - - public static void writeTo(StreamOutput out, ResourceType resourceType) throws IOException { - out.writeString(resourceType.getName()); - } - - public String getName() { - return name; - } -} 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 c26c5d63a3573..cb26590ba5778 100644 --- a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java +++ b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java @@ -18,7 +18,6 @@ import org.opensearch.common.settings.Setting; import org.opensearch.monitor.jvm.JvmStats; import org.opensearch.monitor.process.ProcessProbe; -import org.opensearch.search.ResourceType; import org.opensearch.search.backpressure.settings.SearchBackpressureMode; import org.opensearch.search.backpressure.settings.SearchBackpressureSettings; import org.opensearch.search.backpressure.settings.SearchShardTaskSettings; @@ -34,6 +33,7 @@ import org.opensearch.search.backpressure.trackers.TaskResourceUsageTrackerType; import org.opensearch.search.backpressure.trackers.TaskResourceUsageTrackers; import org.opensearch.search.backpressure.trackers.TaskResourceUsageTrackers.TaskResourceUsageTracker; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.SearchBackpressureTask; import org.opensearch.tasks.Task; @@ -47,7 +47,6 @@ import java.io.IOException; import java.util.ArrayList; import java.util.Collections; -import java.util.EnumMap; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -69,9 +68,9 @@ public class SearchBackpressureService extends AbstractLifecycleComponent implem private static final Logger logger = LogManager.getLogger(SearchBackpressureService.class); private static final Map> trackerApplyConditions = Map.of( TaskResourceUsageTrackerType.CPU_USAGE_TRACKER, - (nodeDuressTrackers) -> nodeDuressTrackers.isResourceInDuress(ResourceType.CPU), + (nodeDuressTrackers) -> nodeDuressTrackers.isResourceInDuress(ResourceType.fromName("cpu")), TaskResourceUsageTrackerType.HEAP_USAGE_TRACKER, - (nodeDuressTrackers) -> isHeapTrackingSupported() && nodeDuressTrackers.isResourceInDuress(ResourceType.MEMORY), + (nodeDuressTrackers) -> isHeapTrackingSupported() && nodeDuressTrackers.isResourceInDuress(ResourceType.fromName("jvm")), TaskResourceUsageTrackerType.ELAPSED_TIME_TRACKER, (nodeDuressTrackers) -> true ); @@ -94,10 +93,10 @@ public SearchBackpressureService( ThreadPool threadPool, TaskManager taskManager ) { - this(settings, taskResourceTrackingService, threadPool, System::nanoTime, new NodeDuressTrackers(new EnumMap<>(ResourceType.class) { + this(settings, taskResourceTrackingService, threadPool, System::nanoTime, new NodeDuressTrackers(new HashMap<>() { { put( - ResourceType.CPU, + ResourceType.fromName("cpu"), new NodeDuressTracker( () -> ProcessProbe.getInstance().getProcessCpuPercent() / 100.0 >= settings.getNodeDuressSettings() .getCpuThreshold(), @@ -105,7 +104,7 @@ public SearchBackpressureService( ) ); put( - ResourceType.MEMORY, + ResourceType.fromName("jvm"), new NodeDuressTracker( () -> JvmStats.jvmStats().getMem().getHeapUsedPercent() / 100.0 >= settings.getNodeDuressSettings() .getHeapThreshold(), diff --git a/server/src/main/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackers.java b/server/src/main/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackers.java index ae60a82fc2816..2f74889b034de 100644 --- a/server/src/main/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackers.java +++ b/server/src/main/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackers.java @@ -9,7 +9,7 @@ package org.opensearch.search.backpressure.trackers; import org.opensearch.common.util.Streak; -import org.opensearch.search.ResourceType; +import org.opensearch.search.resourcetypes.ResourceType; import java.util.Map; import java.util.function.BooleanSupplier; diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageView.java similarity index 51% rename from server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java rename to server/src/main/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageView.java index a6da186487a45..354d38862f746 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java +++ b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageView.java @@ -6,10 +6,10 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing; +package org.opensearch.search.querygroup; import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.tasks.Task; import java.util.ArrayList; @@ -19,27 +19,27 @@ import java.util.Objects; /** - * Represents the point in time view of resource usage of a sandbox and - * has a 1:1 relation with a sandbox. - * This class holds the sandbox ID, the resource usage data, and the list of active tasks. + * 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 QueryGroup ID, the resource usage data, and the list of active tasks. */ @ExperimentalApi -public class SandboxLevelResourceUsageView { +public class QueryGroupLevelResourceUsageView { - private final String sandboxId; - // resourceUsage holds the resource usage data for a sandbox at a point in time - private final Map resourceUsage; - // activeTasks holds the list of active tasks for a sandbox at a point in time + private final String queryGroupId; + // resourceUsage holds the resource usage data for a QueryGroup at a point in time + private final Map resourceUsage; + // activeTasks holds the list of active tasks for a QueryGroup at a point in time private final List activeTasks; - public SandboxLevelResourceUsageView(String sandboxId) { - this.sandboxId = sandboxId; + public QueryGroupLevelResourceUsageView(String queryGroupId) { + this.queryGroupId = queryGroupId; this.resourceUsage = new HashMap<>(); this.activeTasks = new ArrayList<>(); } - public SandboxLevelResourceUsageView(String sandboxId, Map resourceUsage, List activeTasks) { - this.sandboxId = sandboxId; + public QueryGroupLevelResourceUsageView(String queryGroupId, Map resourceUsage, List activeTasks) { + this.queryGroupId = queryGroupId; this.resourceUsage = resourceUsage; this.activeTasks = activeTasks; } @@ -49,7 +49,7 @@ public SandboxLevelResourceUsageView(String sandboxId, Map * * @return The map of resource usage data */ - public Map getResourceUsageData() { + public Map getResourceUsageData() { return resourceUsage; } @@ -66,12 +66,12 @@ public List getActiveTasks() { public boolean equals(Object o) { if (this == o) return true; if (o == null || getClass() != o.getClass()) return false; - SandboxLevelResourceUsageView that = (SandboxLevelResourceUsageView) o; - return Objects.equals(sandboxId, that.sandboxId); + QueryGroupLevelResourceUsageView that = (QueryGroupLevelResourceUsageView) o; + return Objects.equals(queryGroupId, that.queryGroupId); } @Override public int hashCode() { - return Objects.hashCode(sandboxId); + return Objects.hashCode(queryGroupId); } } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java similarity index 54% rename from server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java rename to server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java index 21fadbaddb3d1..5a4f465662d0e 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java +++ b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java @@ -6,18 +6,18 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing; +package org.opensearch.search.querygroup; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.opensearch.cluster.metadata.Sandbox; +import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.cluster.service.ClusterService; import org.opensearch.common.inject.Inject; import org.opensearch.common.lifecycle.AbstractLifecycleComponent; import org.opensearch.common.unit.TimeValue; -import org.opensearch.search.sandboxing.cancellation.DefaultTaskCancellation; -import org.opensearch.search.sandboxing.cancellation.LongestRunningTaskFirstStrategy; -import org.opensearch.search.sandboxing.tracker.SandboxUsageTracker; +import org.opensearch.search.querygroup.cancellation.DefaultTaskCancellation; +import org.opensearch.search.querygroup.cancellation.LongestRunningTaskFirstStrategy; +import org.opensearch.search.querygroup.tracker.QueryGroupUsageTracker; import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.ThreadPool; @@ -28,37 +28,37 @@ import java.util.concurrent.TimeUnit; /** - * Main service which will run periodically to track and cancel resource constraint violating tasks in sandboxes + * Main service which will run periodically to track and cancel resource constraint violating tasks in QueryGroups */ -public class SandboxService extends AbstractLifecycleComponent { - private static final Logger logger = LogManager.getLogger(SandboxService.class); +public class QueryGroupService extends AbstractLifecycleComponent { + private static final Logger logger = LogManager.getLogger(QueryGroupService.class); - private final SandboxUsageTracker sandboxUsageTracker; - // private final SandboxPruner sandboxPruner; + private final QueryGroupUsageTracker queryGroupUsageTracker; + // private final QueryGroupPruner queryGroupPruner; private volatile Scheduler.Cancellable scheduledFuture; - // private final SandboxServiceSettings sandboxServiceSettings; + // private final QueryGroupServiceSettings queryGroupServiceSettings; private final ThreadPool threadPool; private final ClusterService clusterService; /** * Guice managed constructor * - * @param sandboxUsageTracker tracker service - // * @param sandboxPruner - // * @param sandboxServiceSettings + * @param queryGroupUsageTracker tracker service + // * @param queryGroupPruner + // * @param queryGroupServiceSettings * @param threadPool threadpool this will be used to schedule the service */ @Inject - public SandboxService( - SandboxUsageTracker sandboxUsageTracker, - // SandboxServiceSettings sandboxServiceSettings, - // SandboxPruner sandboxPruner, + public QueryGroupService( + QueryGroupUsageTracker queryGroupUsageTracker, + // QueryGroupServiceSettings queryGroupServiceSettings, + // QueryGroupPruner queryGroupPruner, ClusterService clusterService, ThreadPool threadPool ) { - this.sandboxUsageTracker = sandboxUsageTracker; - // this.sandboxServiceSettings = sandboxServiceSettings; - // this.sandboxPruner = sandboxPruner; + this.queryGroupUsageTracker = queryGroupUsageTracker; + // this.queryGroupServiceSettings = queryGroupServiceSettings; + // this.queryGroupPruner = queryGroupPruner; this.clusterService = clusterService; this.threadPool = threadPool; } @@ -67,19 +67,20 @@ public SandboxService( * run at regular interval */ private void doRun() { - Map sandboxLevelResourceUsageViews = sandboxUsageTracker.constructSandboxLevelUsageViews(); - Set activeSandboxes = getActiveSandboxes(); + Map queryGroupLevelResourceUsageViews = queryGroupUsageTracker + .constructQueryGroupLevelUsageViews(); + Set activeQueryGroups = getActiveQueryGroups(); DefaultTaskCancellation taskCancellation = new DefaultTaskCancellation( new LongestRunningTaskFirstStrategy(), - sandboxLevelResourceUsageViews, - activeSandboxes + queryGroupLevelResourceUsageViews, + activeQueryGroups ); taskCancellation.cancelTasks(); - // TODO Prune the sandboxes + // TODO Prune the QueryGroups } - private Set getActiveSandboxes() { - return new HashSet<>(clusterService.state().metadata().sandboxes().values()); + private Set getActiveQueryGroups() { + return new HashSet<>(clusterService.state().metadata().queryGroups().values()); } /** @@ -91,10 +92,10 @@ protected void doStart() { try { doRun(); } catch (Exception e) { - logger.debug("Exception occurred in Query Sandbox service", e); + logger.debug("Exception occurred in QueryGroup service", e); } }, - new TimeValue(1, TimeUnit.SECONDS), // TODO get this from SandboxServiceSettings + new TimeValue(1, TimeUnit.SECONDS), // TODO get this from QueryGroupServiceSettings ThreadPool.Names.GENERIC ); } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupTask.java similarity index 55% rename from server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java rename to server/src/main/java/org/opensearch/search/querygroup/QueryGroupTask.java index a674f27d71eff..d3ab8c2f0f8b6 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java +++ b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupTask.java @@ -6,13 +6,13 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing; +package org.opensearch.search.querygroup; /** - * This interface can be implemented by tasks which will be tracked and monitored using {@link org.opensearch.cluster.metadata.Sandbox} + * This interface can be implemented by tasks which will be tracked and monitored using {@link org.opensearch.cluster.metadata.QueryGroup} */ -public interface SandboxTask { - void setSandboxId(String sandboxId); +public interface QueryGroupTask { + void setQueryGroupId(String queryGroupId); - String getSandboxId(); + String getQueryGroupId(); } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/AbstractTaskSelectionStrategy.java similarity index 94% rename from server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java rename to server/src/main/java/org/opensearch/search/querygroup/cancellation/AbstractTaskSelectionStrategy.java index 242bfdbefa0b0..36c6ba60bd850 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java +++ b/server/src/main/java/org/opensearch/search/querygroup/cancellation/AbstractTaskSelectionStrategy.java @@ -6,9 +6,9 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.cancellation; +package org.opensearch.search.querygroup.cancellation; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskCancellation; @@ -45,7 +45,7 @@ public abstract class AbstractTaskSelectionStrategy implements TaskSelectionStra * @throws IllegalArgumentException If the limit is less than zero */ @Override - public List selectTasksForCancellation(List tasks, long limit, SystemResource resourceType) { + public List selectTasksForCancellation(List tasks, long limit, ResourceType resourceType) { if (limit < 0) { throw new IllegalArgumentException("reduceBy has to be greater than zero"); } diff --git a/server/src/main/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellation.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellation.java new file mode 100644 index 0000000000000..5104fef64ce38 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellation.java @@ -0,0 +1,177 @@ +/* + * 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.search.querygroup.cancellation; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.search.querygroup.QueryGroupLevelResourceUsageView; +import org.opensearch.search.resourcetypes.ResourceType; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskCancellation; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; + +import static org.opensearch.search.querygroup.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 TaskSelectionStrategy} to identify tasks that exceed + * predefined resource usage limits and are therefore eligible for cancellation. + * + *

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.

+ * + *

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.

+ * + * @see TaskSelectionStrategy + * @see QueryGroup + * @see ResourceType + */ +public class DefaultTaskCancellation { + protected final TaskSelectionStrategy taskSelectionStrategy; + // a map of QueryGroupId to its corresponding QueryGroupLevelResourceUsageView object + protected final Map queryGroupLevelViews; + protected final Set activeQueryGroups; + + public DefaultTaskCancellation( + TaskSelectionStrategy taskSelectionStrategy, + Map queryGroupLevelViews, + Set activeQueryGroups + ) { + this.taskSelectionStrategy = taskSelectionStrategy; + this.queryGroupLevelViews = queryGroupLevelViews; + this.activeQueryGroups = activeQueryGroups; + } + + /** + * Cancel tasks based on the implemented strategy. + */ + public final void cancelTasks() { + List cancellableTasks = getAllCancellableTasks(); + for (TaskCancellation taskCancellation : cancellableTasks) { + taskCancellation.cancel(); + } + } + + /** + * returns the list of QueryGroups breaching their resource limits. + * + * @return List of QueryGroups + */ + public List getQueryGroupsToCancelFrom() { + final List queryGroupsToCancelFrom = new ArrayList<>(); + + for (QueryGroup queryGroup : this.activeQueryGroups) { + Map currentResourceUsage = getResourceUsage(queryGroup.get_id()); + // if(currentResourceUsage == null) { + // // skip if the QueryGroup is not found + // continue; + // } + + Map resourceLimits = queryGroup.getResourceLimits(); + for (ResourceType resourceType : TRACKED_RESOURCES) { + if (resourceLimits.containsKey(resourceType)) { + long threshold = queryGroup.getThresholdInLong(resourceType); + + if (isBreachingThreshold(currentResourceUsage, resourceType, threshold)) { + queryGroupsToCancelFrom.add(queryGroup); + break; + } + } + } + + } + + return queryGroupsToCancelFrom; + } + + /** + * Get all cancellable tasks from the QueryGroups. + * + * @return List of tasks that can be cancelled + */ + protected List getAllCancellableTasks() { + return getQueryGroupsToCancelFrom().stream() + .flatMap(queryGroup -> getCancellableTasksFrom(queryGroup).stream()) + .collect(Collectors.toList()); + } + + /** + * 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 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) { + long reduceBy = getReduceBy(queryGroup, resourceType); + return reduceBy > 0; + } + + private List getTaskCancellations(QueryGroup queryGroup, ResourceType resourceType) { + return taskSelectionStrategy.selectTasksForCancellation( + getAllTasksInQueryGroup(queryGroup.get_id()), + getReduceBy(queryGroup, resourceType), + resourceType + ); + } + + private long getReduceBy(QueryGroup queryGroup, ResourceType resourceType) { + Long usage = getUsage(queryGroup, resourceType); + if (usage == null) { + return 0; + } + return getUsage(queryGroup, resourceType) - queryGroup.getThresholdInLong(resourceType); + } + + private Long getUsage(QueryGroup queryGroup, ResourceType resourceType) { + return queryGroupLevelViews.get(queryGroup.get_id()).getResourceUsageData().get(resourceType); + } + + private List getAllTasksInQueryGroup(String queryGroupId) { + return queryGroupLevelViews.get(queryGroupId).getActiveTasks(); + } + + /** + * Checks if the current resource usage is breaching the threshold of the provided resource limit. + * + * @param currentResourceUsage The current resource usage + * @param resourceType The resource type to check against + * @param threshold Threshold of the query group + * @return true if the current resource usage is breaching the threshold, false otherwise + */ + private boolean isBreachingThreshold(Map currentResourceUsage, ResourceType resourceType, long threshold) { + return currentResourceUsage.get(resourceType) > threshold; + } + + /** + * Returns the resource usage of the QueryGroup with the provided ID. + * + * @param queryGroupId The ID of the QueryGroup + * @return The resource usage of the QueryGroup + */ + private Map getResourceUsage(String queryGroupId) { + // if(QueryGroupLevelViews.get(queryGroupId) == null) { + // return null; + // } + return queryGroupLevelViews.get(queryGroupId).getResourceUsageData(); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategy.java similarity index 93% rename from server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java rename to server/src/main/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategy.java index eacdf7a4c15f4..26eb280a7aac3 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java +++ b/server/src/main/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategy.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.cancellation; +package org.opensearch.search.querygroup.cancellation; import org.opensearch.tasks.Task; diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategy.java similarity index 93% rename from server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java rename to server/src/main/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategy.java index af186b01a8974..d2335b90e8f31 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java +++ b/server/src/main/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategy.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.cancellation; +package org.opensearch.search.querygroup.cancellation; import org.opensearch.tasks.Task; diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategy.java similarity index 85% rename from server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java rename to server/src/main/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategy.java index 61f1883aefaa5..e467db15bedac 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java +++ b/server/src/main/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategy.java @@ -6,9 +6,9 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.cancellation; +package org.opensearch.search.querygroup.cancellation; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskCancellation; @@ -28,5 +28,5 @@ public interface TaskSelectionStrategy { * * @return List of tasks that should be cancelled. */ - List selectTasksForCancellation(List tasks, long limit, SystemResource resourceType); + List selectTasksForCancellation(List tasks, long limit, ResourceType resourceType); } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/package-info.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/package-info.java similarity index 82% rename from server/src/main/java/org/opensearch/search/sandboxing/cancellation/package-info.java rename to server/src/main/java/org/opensearch/search/querygroup/cancellation/package-info.java index 7bda195ba18ff..fdcb1017f65aa 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/package-info.java +++ b/server/src/main/java/org/opensearch/search/querygroup/cancellation/package-info.java @@ -9,4 +9,4 @@ /** * Package for cancellation related abstracts */ -package org.opensearch.search.sandboxing.cancellation; +package org.opensearch.search.querygroup.cancellation; diff --git a/server/src/main/java/org/opensearch/search/querygroup/module/QueryGroupModule.java b/server/src/main/java/org/opensearch/search/querygroup/module/QueryGroupModule.java new file mode 100644 index 0000000000000..f6c6c21a22cd2 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/querygroup/module/QueryGroupModule.java @@ -0,0 +1,31 @@ +/* + * 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.search.querygroup.module; + +import org.opensearch.common.inject.AbstractModule; +import org.opensearch.search.querygroup.tracker.QueryGroupResourceUsageTrackerService; +import org.opensearch.search.querygroup.tracker.QueryGroupUsageTracker; + +/** + * Module class for resource usage limiting related artifacts + */ +public class QueryGroupModule extends AbstractModule { + + /** + * Default constructor + */ + public QueryGroupModule() {} + + @Override + protected void configure() { + bind(QueryGroupUsageTracker.class).to(QueryGroupResourceUsageTrackerService.class).asEagerSingleton(); + // bind(AbstractTaskCancellation.class).to(QueryGroupResourceUsageTrackerService.class).asEagerSingleton(); + // bind(QueryGroupPruner.class).to(QueryGroupResourceUsageTrackerService.class).asEagerSingleton(); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/module/package-info.java b/server/src/main/java/org/opensearch/search/querygroup/module/package-info.java similarity index 82% rename from server/src/main/java/org/opensearch/search/sandboxing/module/package-info.java rename to server/src/main/java/org/opensearch/search/querygroup/module/package-info.java index 5be0923cc3c62..769192a288664 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/module/package-info.java +++ b/server/src/main/java/org/opensearch/search/querygroup/module/package-info.java @@ -10,4 +10,4 @@ * Guice Module */ -package org.opensearch.search.sandboxing.module; +package org.opensearch.search.querygroup.module; diff --git a/server/src/main/java/org/opensearch/search/sandboxing/package-info.java b/server/src/main/java/org/opensearch/search/querygroup/package-info.java similarity index 85% rename from server/src/main/java/org/opensearch/search/sandboxing/package-info.java rename to server/src/main/java/org/opensearch/search/querygroup/package-info.java index 3dfecf384df44..39dfc756a0a4e 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/package-info.java +++ b/server/src/main/java/org/opensearch/search/querygroup/package-info.java @@ -9,4 +9,4 @@ /** * Query Sandboxing related artifacts */ -package org.opensearch.search.sandboxing; +package org.opensearch.search.querygroup; diff --git a/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java new file mode 100644 index 0000000000000..7e0fcd6bc7eea --- /dev/null +++ b/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java @@ -0,0 +1,124 @@ +/* + * 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.search.querygroup.tracker; + +import org.opensearch.common.inject.Inject; +import org.opensearch.search.querygroup.QueryGroupLevelResourceUsageView; +import org.opensearch.search.querygroup.QueryGroupTask; +import org.opensearch.search.resourcetypes.ResourceType; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskManager; +import org.opensearch.tasks.TaskResourceTrackingService; + +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; + +/** + * This class tracks requests per QueryGroup + */ +// @ExperimentalApi +public class QueryGroupResourceUsageTrackerService implements QueryGroupUsageTracker, TaskManager.TaskEventListeners { + + public static final List TRACKED_RESOURCES = List.of(ResourceType.fromName("JVM"), ResourceType.fromName("CPU")); + + private final TaskManager taskManager; + private final TaskResourceTrackingService taskResourceTrackingService; + + /** + * QueryGroupResourceTrackerService constructor + * + * @param taskManager Task Manager service for keeping track of currently running tasks on the nodes + * @param taskResourceTrackingService Service that helps track resource usage of tasks running on a node. + */ + @Inject + public QueryGroupResourceUsageTrackerService( + final TaskManager taskManager, + final TaskResourceTrackingService taskResourceTrackingService + ) { + this.taskManager = taskManager; + this.taskResourceTrackingService = taskResourceTrackingService; + } + + /** + * Constructs a map of QueryGroupLevelResourceUsageView instances for each QueryGroup. + * + * @return Map of QueryGroup views + */ + @Override + public Map constructQueryGroupLevelUsageViews() { + Map queryGroupViews = new HashMap<>(); + + Map> tasksByQueryGroup = getTasksGroupedByQueryGroup(); + Map> queryGroupResourceUsage = getResourceUsageOfQueryGroups(tasksByQueryGroup); + + for (String queryGroupId : tasksByQueryGroup.keySet()) { + QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( + queryGroupId, + queryGroupResourceUsage.get(queryGroupId), + tasksByQueryGroup.get(queryGroupId) + ); + queryGroupViews.put(queryGroupId, queryGroupLevelResourceUsageView); + } + return queryGroupViews; + } + + /** + * Groups tasks by their associated QueryGroup. + * + * @return Map of tasks grouped by QueryGroup + */ + private Map> getTasksGroupedByQueryGroup() { + return taskResourceTrackingService.getResourceAwareTasks() + .values() + .stream() + .filter(QueryGroupTask.class::isInstance) + .map(QueryGroupTask.class::cast) + .collect(Collectors.groupingBy(QueryGroupTask::getQueryGroupId, Collectors.mapping(task -> (Task) task, Collectors.toList()))); + } + + /** + * Calculates the resource usage of each QueryGroup. + * + * @param tasksByQueryGroup Map of tasks grouped by QueryGroup + * @return Map of resource usage for each QueryGroup + */ + private Map> getResourceUsageOfQueryGroups(Map> tasksByQueryGroup) { + Map> resourceUsageOfQueryGroups = new HashMap<>(); + + // Iterate over each QueryGroup entry + for (Map.Entry> queryGroupEntry : tasksByQueryGroup.entrySet()) { + String queryGroupId = queryGroupEntry.getKey(); + List tasks = queryGroupEntry.getValue(); + + // Prepare a usage map for the current QueryGroup, or retrieve the existing one + Map queryGroupUsage = resourceUsageOfQueryGroups.computeIfAbsent(queryGroupId, k -> new HashMap<>()); + + // Accumulate resource usage for each task in the QueryGroup + for (Task task : tasks) { + for (ResourceType resourceType : TRACKED_RESOURCES) { + long currentUsage = queryGroupUsage.getOrDefault(resourceType, 0L); + long taskUsage = resourceType.getResourceUsage(task); + // task.getTotalResourceStats().getCpuTimeInNanos(); + queryGroupUsage.put(resourceType, currentUsage + taskUsage); + } + } + } + return resourceUsageOfQueryGroups; + } + + /** + * Handles the completion of a task. + * + * @param task The completed task + */ + @Override + public void onTaskCompleted(Task task) {} +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxUsageTracker.java b/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupUsageTracker.java similarity index 61% rename from server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxUsageTracker.java rename to server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupUsageTracker.java index b08e110e090f9..15e162a5e25a8 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxUsageTracker.java +++ b/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupUsageTracker.java @@ -6,19 +6,19 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.tracker; +package org.opensearch.search.querygroup.tracker; -import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; +import org.opensearch.search.querygroup.QueryGroupLevelResourceUsageView; import java.util.Map; /** * This interface is mainly for tracking the resourceLimitGroup level resource usages */ -public interface SandboxUsageTracker { +public interface QueryGroupUsageTracker { /** * updates the current resource usage of resourceLimitGroups */ - Map constructSandboxLevelUsageViews(); + Map constructQueryGroupLevelUsageViews(); } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/tracker/package-info.java b/server/src/main/java/org/opensearch/search/querygroup/tracker/package-info.java similarity index 84% rename from server/src/main/java/org/opensearch/search/sandboxing/tracker/package-info.java rename to server/src/main/java/org/opensearch/search/querygroup/tracker/package-info.java index fcf0c504d3752..28c1bf3d03ded 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/tracker/package-info.java +++ b/server/src/main/java/org/opensearch/search/querygroup/tracker/package-info.java @@ -9,4 +9,4 @@ /** * ResourceLimitGroup resource tracking artifacts */ -package org.opensearch.search.sandboxing.tracker; +package org.opensearch.search.querygroup.tracker; diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java b/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java new file mode 100644 index 0000000000000..bdb30aab970bd --- /dev/null +++ b/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java @@ -0,0 +1,59 @@ +/* + * 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.search.resourcetypes; + +import org.opensearch.tasks.Task; + +/** + * Represents the CPU resource type for tasks. + * This class provides methods to retrieve CPU usage and convert threshold percentages to long values. + */ +public class CPU extends ResourceType { + /** + * Returns the CPU time usage of the provided task. + * + * @param task The task whose CPU time usage is to be returned + * @return The CPU time usage of the task + */ + @Override + public long getResourceUsage(Task task) { + return task.getTotalResourceStats().getCpuTimeInNanos(); + } + + /** + * Returns the name of the resource type. + * + * @return The name of the resource type, which is "CPU" + */ + @Override + public String getName() { + return "CPU"; + } + + @Override + public boolean equals(Object obj) { + return obj instanceof CPU; + } + + @Override + public int hashCode() { + return "CPU".hashCode(); + } + + /** + * Converts the given threshold percentage to a long value that can be compared. + * + * @param threshold The threshold percentage to be converted + * @return The threshold value in nanoseconds + */ + // TODO: Implement this method + public long convertThresholdPercentageToLong(Double threshold) { + return (long) (threshold * Runtime.getRuntime().availableProcessors()); + } +} diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/JVM.java b/server/src/main/java/org/opensearch/search/resourcetypes/JVM.java new file mode 100644 index 0000000000000..4861e5a94b1df --- /dev/null +++ b/server/src/main/java/org/opensearch/search/resourcetypes/JVM.java @@ -0,0 +1,60 @@ +/* + * 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.search.resourcetypes; + +import org.opensearch.monitor.jvm.JvmStats; +import org.opensearch.tasks.Task; + +/** + * Represents the JVM resource type for tasks. + * This class provides methods to retrieve JVM usage and convert threshold percentages to long values. + */ +public class JVM extends ResourceType { + /** + * Returns the JVM usage of the provided task. + * + * @param task The task whose JVM usage is to be returned + * @return The JVM usage of the task + */ + @Override + public long getResourceUsage(Task task) { + return task.getTotalResourceStats().getMemoryInBytes(); + } + + /** + * Returns the name of the resource type. + * + * @return The name of the resource type, which is "JVM" + */ + @Override + public String getName() { + return "JVM"; + } + + @Override + public boolean equals(Object obj) { + return obj instanceof JVM; + } + + @Override + public int hashCode() { + return "JVM".hashCode(); + } + + /** + * Converts the given threshold percentage to a long value that can be compared. + * + * @param threshold The threshold percentage to be converted + * @return The threshold value in bytes + */ + @Override + public long convertThresholdPercentageToLong(Double threshold) { + return (long) (threshold * JvmStats.jvmStats().getMem().getHeapMax().getBytes()); + } +} diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java new file mode 100644 index 0000000000000..ee091dd70c1e9 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java @@ -0,0 +1,73 @@ +/* + * 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.search.resourcetypes; + +import org.opensearch.common.annotation.PublicApi; +import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.tasks.Task; + +import java.io.IOException; + +/** + * Enum to hold the resource type + */ +@PublicApi(since = "2.x") +public abstract class ResourceType { + public static ResourceType[] values() { + return new ResourceType[] { new CPU(), new JVM() }; + } + + /** + * Returns the resource usage of the provided task. + * The specific resource that this method returns depends on the implementation. + * + * @param task The task whose resource usage is to be returned + * @return The resource usage of the task + */ + public abstract long getResourceUsage(Task task); + + /** + * Creates a SystemResource from a string. + * If the string is "JVM", a JVM is returned. + * If the string is "CPU", a CPU is returned. + * If the string is not recognized, an IllegalArgumentException is thrown. + * + * @param type The string from which to create a SystemResource + * @return The created SystemResource + * @throws IllegalArgumentException If the string is not recognized + */ + public static ResourceType fromName(String type) { + if (type.equalsIgnoreCase("JVM")) { + return new JVM(); + } else if (type.equalsIgnoreCase("CPU")) { + return new CPU(); + } else { + throw new IllegalArgumentException("Unsupported resource type: " + type); + } + } + + /** + * Returns the name of the resource type. + * + * @return The name of the resource type + */ + public abstract String getName(); + + /** + * Converts the given threshold percentage to a long value that can be compared. + * + * @param threshold The threshold percentage to be converted + * @return The threshold value in long format + */ + public abstract long convertThresholdPercentageToLong(Double threshold); + + public static void writeTo(StreamOutput out, ResourceType resourceType) throws IOException { + out.writeString(resourceType.getName()); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java deleted file mode 100644 index 73901df716ada..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java +++ /dev/null @@ -1,170 +0,0 @@ -/* - * 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.search.sandboxing.cancellation; - -import org.opensearch.cluster.metadata.Sandbox; -import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; -import org.opensearch.tasks.Task; -import org.opensearch.tasks.TaskCancellation; - -import java.util.ArrayList; -import java.util.List; -import java.util.Map; -import java.util.Set; -import java.util.stream.Collectors; - -import static org.opensearch.search.sandboxing.tracker.SandboxResourceUsageTrackerService.TRACKED_RESOURCES; - -/** - * Manages the cancellation of tasks enforced by sandbox thresholds on resource usage criteria. - * This class utilizes a strategy pattern through {@link TaskSelectionStrategy} to identify tasks that exceed - * predefined resource usage limits and are therefore eligible for cancellation. - * - *

The cancellation process is initiated by evaluating the resource usage of each sandbox against its - * resource limits. Tasks that contribute to exceeding these limits are selected for cancellation based on the - * implemented task selection strategy.

- * - *

Instances of this class are configured with a map linking sandbox IDs to their corresponding resource usage - * views, a set of active sandboxes, and a task selection strategy. These components collectively facilitate the - * identification and cancellation of tasks that threaten to breach sandbox resource limits.

- * - * @see TaskSelectionStrategy - * @see Sandbox - * @see SystemResource - */ -public class DefaultTaskCancellation { - protected final TaskSelectionStrategy taskSelectionStrategy; - // a map of sandboxId to its corresponding SandboxLevelResourceUsageView object - protected final Map sandboxLevelViews; - protected final Set activeSandboxes; - - public DefaultTaskCancellation( - TaskSelectionStrategy taskSelectionStrategy, - Map sandboxLevelViews, - Set activeSandboxes - ) { - this.taskSelectionStrategy = taskSelectionStrategy; - this.sandboxLevelViews = sandboxLevelViews; - this.activeSandboxes = activeSandboxes; - } - - /** - * Cancel tasks based on the implemented strategy. - */ - public final void cancelTasks() { - List cancellableTasks = getAllCancellableTasks(); - for (TaskCancellation taskCancellation : cancellableTasks) { - taskCancellation.cancel(); - } - } - - /** - * returns the list of sandboxes breaching their resource limits. - * - * @return List of sandboxes - */ - public List getSandboxesToCancelFrom() { - final List sandboxesToCancelFrom = new ArrayList<>(); - - for (Sandbox sandbox : this.activeSandboxes) { - Map currentResourceUsage = getResourceUsage(sandbox.getId()); - // if(currentResourceUsage == null) { - // // skip if the sandbox is not found - // continue; - // } - - for (Sandbox.ResourceLimit resourceLimit : sandbox.getResourceLimits()) { - if (isBreachingThreshold(currentResourceUsage, resourceLimit)) { - sandboxesToCancelFrom.add(sandbox); - break; - } - } - } - - return sandboxesToCancelFrom; - } - - /** - * Get all cancellable tasks from the sandboxes. - * - * @return List of tasks that can be cancelled - */ - protected List getAllCancellableTasks() { - return getSandboxesToCancelFrom().stream() - .flatMap(sandbox -> getCancellableTasksFrom(sandbox).stream()) - .collect(Collectors.toList()); - } - - /** - * Get cancellable tasks from a specific sandbox. - * - * @param sandbox The sandbox from which to get cancellable tasks - * @return List of tasks that can be cancelled - */ - protected List getCancellableTasksFrom(Sandbox sandbox) { - return TRACKED_RESOURCES.stream() - .filter(resourceType -> shouldCancelTasks(sandbox, resourceType)) - .flatMap(resourceType -> getTaskCancellations(sandbox, resourceType).stream()) - .collect(Collectors.toList()); - } - - private boolean shouldCancelTasks(Sandbox sandbox, SystemResource resourceType) { - long reduceBy = getReduceBy(sandbox, resourceType); - return reduceBy > 0; - } - - private List getTaskCancellations(Sandbox sandbox, SystemResource resourceType) { - return taskSelectionStrategy.selectTasksForCancellation( - getAllTasksInSandbox(sandbox.getId()), - getReduceBy(sandbox, resourceType), - resourceType - ); - } - - private long getReduceBy(Sandbox sandbox, SystemResource resourceType) { - Long usage = getUsage(sandbox, resourceType); - if (usage == null) { - return 0; - } - return getUsage(sandbox, resourceType) - sandbox.getResourceLimitFor(resourceType).getThresholdInLong(); - } - - private Long getUsage(Sandbox sandbox, SystemResource resourceType) { - return sandboxLevelViews.get(sandbox.getId()).getResourceUsageData().get(resourceType); - } - - private List getAllTasksInSandbox(String sandboxId) { - return sandboxLevelViews.get(sandboxId).getActiveTasks(); - } - - /** - * Checks if the current resource usage is breaching the threshold of the provided resource limit. - * - * @param currentResourceUsage The current resource usage - * @param resourceLimit The resource limit to check against - * @return true if the current resource usage is breaching the threshold, false otherwise - */ - private boolean isBreachingThreshold(Map currentResourceUsage, Sandbox.ResourceLimit resourceLimit) { - return currentResourceUsage.get(resourceLimit.getResourceType()) > resourceLimit.getThreshold(); - } - - /** - * Returns the resource usage of the sandbox with the provided ID. - * - * @param sandboxId The ID of the sandbox - * @return The resource usage of the sandbox - */ - private Map getResourceUsage(String sandboxId) { - // if(sandboxLevelViews.get(sandboxId) == null) { - // return null; - // } - return sandboxLevelViews.get(sandboxId).getResourceUsageData(); - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/module/SandboxModule.java b/server/src/main/java/org/opensearch/search/sandboxing/module/SandboxModule.java deleted file mode 100644 index 3f7c8a11099c1..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/module/SandboxModule.java +++ /dev/null @@ -1,31 +0,0 @@ -/* - * 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.search.sandboxing.module; - -import org.opensearch.common.inject.AbstractModule; -import org.opensearch.search.sandboxing.tracker.SandboxResourceUsageTrackerService; -import org.opensearch.search.sandboxing.tracker.SandboxUsageTracker; - -/** - * Module class for resource usage limiting related artifacts - */ -public class SandboxModule extends AbstractModule { - - /** - * Default constructor - */ - public SandboxModule() {} - - @Override - protected void configure() { - bind(SandboxUsageTracker.class).to(SandboxResourceUsageTrackerService.class).asEagerSingleton(); - // bind(AbstractTaskCancellation.class).to(SandboxResourceUsageTrackerService.class).asEagerSingleton(); - // bind(SandboxPruner.class).to(SandboxResourceUsageTrackerService.class).asEagerSingleton(); - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResource.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResource.java deleted file mode 100644 index 808a76a9bc66a..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResource.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * 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.search.sandboxing.resourcetype; - -import org.opensearch.tasks.Task; - -/** - * Represents the CPU time resource type. - */ -public class CpuTimeResource extends SystemResource { - /** - * Returns the CPU time usage of the provided task. - * - * @param task The task whose CPU time usage is to be returned - * @return The CPU time usage of the task - */ - @Override - public long getResourceUsage(Task task) { - return task.getTotalResourceStats().getCpuTimeInNanos(); - } - - @Override - public boolean equals(Object obj) { - return obj instanceof CpuTimeResource; - } - - @Override - public int hashCode() { - return "CPU".hashCode(); - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResource.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResource.java deleted file mode 100644 index 2387954113a74..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResource.java +++ /dev/null @@ -1,37 +0,0 @@ -/* - * 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.search.sandboxing.resourcetype; - -import org.opensearch.tasks.Task; - -/** - * Represents the JVM memory resource type. - */ -public class JvmMemoryResource extends SystemResource { - /** - * Returns the memory usage of the provided task. - * - * @param task The task whose memory usage is to be returned - * @return The memory usage of the task - */ - @Override - public long getResourceUsage(Task task) { - return task.getTotalResourceStats().getMemoryInBytes(); - } - - @Override - public boolean equals(Object obj) { - return obj instanceof JvmMemoryResource; - } - - @Override - public int hashCode() { - return "JVM".hashCode(); - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SystemResource.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SystemResource.java deleted file mode 100644 index 525861a3356e9..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SystemResource.java +++ /dev/null @@ -1,48 +0,0 @@ -/* - * 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.search.sandboxing.resourcetype; - -import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.tasks.Task; - -/** - * Represents a type of resource that can be used in a sandbox. - * This class is abstract and requires the implementation of the getResourceUsage method. - */ -@ExperimentalApi -public abstract class SystemResource { - /** - * Returns the resource usage of the provided task. - * The specific resource that this method returns depends on the implementation. - * - * @param task The task whose resource usage is to be returned - * @return The resource usage of the task - */ - public abstract long getResourceUsage(Task task); - - /** - * Creates a SystemResource from a string. - * If the string is "JVM", a JvmMemoryResource is returned. - * If the string is "CPU", a CpuTimeResource is returned. - * If the string is not recognized, an IllegalArgumentException is thrown. - * - * @param type The string from which to create a SystemResource - * @return The created SystemResource - * @throws IllegalArgumentException If the string is not recognized - */ - public static SystemResource fromString(String type) { - if (type.equalsIgnoreCase("JVM")) { - return new JvmMemoryResource(); - } else if (type.equalsIgnoreCase("CPU")) { - return new CpuTimeResource(); - } else { - throw new IllegalArgumentException("Unsupported resource type: " + type); - } - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/package-info.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/package-info.java deleted file mode 100644 index f71eb9233abd4..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/package-info.java +++ /dev/null @@ -1,12 +0,0 @@ -/* - * 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 for Sandbox Resource Types - */ -package org.opensearch.search.sandboxing.resourcetype; diff --git a/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java deleted file mode 100644 index 9dcb5e5006475..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java +++ /dev/null @@ -1,126 +0,0 @@ -/* - * 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.search.sandboxing.tracker; - -import org.opensearch.common.inject.Inject; -import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; -import org.opensearch.search.sandboxing.SandboxTask; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; -import org.opensearch.tasks.Task; -import org.opensearch.tasks.TaskManager; -import org.opensearch.tasks.TaskResourceTrackingService; - -import java.util.HashMap; -import java.util.List; -import java.util.Map; -import java.util.stream.Collectors; - -/** - * This class tracks requests per Sandbox - */ -// @ExperimentalApi -public class SandboxResourceUsageTrackerService implements SandboxUsageTracker, TaskManager.TaskEventListeners { - - public static final List TRACKED_RESOURCES = List.of( - SystemResource.fromString("JVM"), - SystemResource.fromString("CPU") - ); - - private final TaskManager taskManager; - private final TaskResourceTrackingService taskResourceTrackingService; - - /** - * SandboxResourceTrackerService constructor - * - * @param taskManager Task Manager service for keeping track of currently running tasks on the nodes - * @param taskResourceTrackingService Service that helps track resource usage of tasks running on a node. - */ - @Inject - public SandboxResourceUsageTrackerService( - final TaskManager taskManager, - final TaskResourceTrackingService taskResourceTrackingService - ) { - this.taskManager = taskManager; - this.taskResourceTrackingService = taskResourceTrackingService; - } - - /** - * Constructs a map of SandboxLevelResourceUsageView instances for each sandbox. - * - * @return Map of sandbox views - */ - @Override - public Map constructSandboxLevelUsageViews() { - Map sandboxViews = new HashMap<>(); - - Map> tasksBySandbox = getTasksGroupedBySandbox(); - Map> sandboxResourceUsage = getResourceUsageOfSandboxes(tasksBySandbox); - - for (String sandboxId : tasksBySandbox.keySet()) { - SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView( - sandboxId, - sandboxResourceUsage.get(sandboxId), - tasksBySandbox.get(sandboxId) - ); - sandboxViews.put(sandboxId, sandboxLevelResourceUsageView); - } - return sandboxViews; - } - - /** - * Groups tasks by their associated sandbox. - * - * @return Map of tasks grouped by sandbox - */ - private Map> getTasksGroupedBySandbox() { - return taskResourceTrackingService.getResourceAwareTasks() - .values() - .stream() - .filter(SandboxTask.class::isInstance) - .map(SandboxTask.class::cast) - .collect(Collectors.groupingBy(SandboxTask::getSandboxId, Collectors.mapping(task -> (Task) task, Collectors.toList()))); - } - - /** - * Calculates the resource usage of each sandbox. - * - * @param tasksBySandbox Map of tasks grouped by sandbox - * @return Map of resource usage for each sandbox - */ - private Map> getResourceUsageOfSandboxes(Map> tasksBySandbox) { - Map> resourceUsageOfSandboxes = new HashMap<>(); - - // Iterate over each sandbox entry - for (Map.Entry> sandboxEntry : tasksBySandbox.entrySet()) { - String sandboxId = sandboxEntry.getKey(); - List tasks = sandboxEntry.getValue(); - - // Prepare a usage map for the current sandbox, or retrieve the existing one - Map sandboxUsage = resourceUsageOfSandboxes.computeIfAbsent(sandboxId, k -> new HashMap<>()); - - // Accumulate resource usage for each task in the sandbox - for (Task task : tasks) { - for (SystemResource resourceType : TRACKED_RESOURCES) { - long currentUsage = sandboxUsage.getOrDefault(resourceType, 0L); - long taskUsage = resourceType.getResourceUsage(task); - sandboxUsage.put(resourceType, currentUsage + taskUsage); - } - } - } - return resourceUsageOfSandboxes; - } - - /** - * Handles the completion of a task. - * - * @param task The completed task - */ - @Override - public void onTaskCompleted(Task task) {} -} diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java index d70a9ce5e10cd..a5ad13f1a53e8 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -14,7 +14,7 @@ import org.opensearch.core.common.io.stream.Writeable; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.search.ResourceType; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.test.AbstractDiffableSerializationTestCase; import java.io.IOException; @@ -34,7 +34,7 @@ public void testToXContent() throws IOException { "test", "ajakgakg983r92_4242", QueryGroup.ResiliencyMode.ENFORCED, - Map.of(ResourceType.MEMORY, 0.5), + Map.of(ResourceType.fromName("jvm"), 0.5), updatedAt ) ) diff --git a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java index c564f0778e6f0..3d3a3792cb82d 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java @@ -14,7 +14,7 @@ import org.opensearch.core.xcontent.ToXContent; import org.opensearch.core.xcontent.XContentBuilder; import org.opensearch.core.xcontent.XContentParser; -import org.opensearch.search.ResourceType; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.test.AbstractSerializingTestCase; import org.joda.time.Instant; @@ -35,7 +35,7 @@ public class QueryGroupTests extends AbstractSerializingTestCase { static QueryGroup createRandomQueryGroup(String _id) { String name = randomAlphaOfLength(10); Map resourceLimit = new HashMap<>(); - resourceLimit.put(ResourceType.MEMORY, randomDoubleBetween(0.0, 0.80, false)); + resourceLimit.put(ResourceType.fromName("jvm"), randomDoubleBetween(0.0, 0.80, false)); return new QueryGroup(name, _id, randomMode(), resourceLimit, Instant.now().getMillis()); } @@ -99,7 +99,7 @@ public void testEmptyResourceLimits() { public void testIllegalQueryGroupMode() { assertThrows( NullPointerException.class, - () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.MEMORY, (Object) 0.4), Instant.now().getMillis()) + () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.fromName("jvm"), (Object) 0.4), Instant.now().getMillis()) ); } @@ -110,7 +110,7 @@ public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { "analytics", "_id", randomMode(), - Map.of(ResourceType.MEMORY, (Object) randomDoubleBetween(1.1, 1.8, false)), + Map.of(ResourceType.fromName("jvm"), (Object) randomDoubleBetween(1.1, 1.8, false)), Instant.now().getMillis() ) ); @@ -121,7 +121,7 @@ public void testValidQueryGroup() { "analytics", "_id", randomMode(), - Map.of(ResourceType.MEMORY, randomDoubleBetween(0.01, 0.8, false)), + Map.of(ResourceType.fromName("jvm"), randomDoubleBetween(0.01, 0.8, false)), Instant.ofEpochMilli(1717187289).getMillis() ); @@ -141,7 +141,7 @@ public void testToXContent() throws IOException { "TestQueryGroup", queryGroupId, QueryGroup.ResiliencyMode.ENFORCED, - Map.of(ResourceType.CPU, 0.30, ResourceType.MEMORY, 0.40), + Map.of(ResourceType.fromName("cpu"), 0.30, ResourceType.fromName("jvm"), 0.40), currentTimeInMillis ); XContentBuilder builder = JsonXContent.contentBuilder(); 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 15d0fcd10d701..0459d060555a7 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java +++ b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java @@ -16,7 +16,6 @@ import org.opensearch.core.common.io.stream.StreamInput; import org.opensearch.core.common.io.stream.StreamOutput; import org.opensearch.core.xcontent.XContentBuilder; -import org.opensearch.search.ResourceType; import org.opensearch.search.backpressure.settings.SearchBackpressureMode; import org.opensearch.search.backpressure.settings.SearchBackpressureSettings; import org.opensearch.search.backpressure.settings.SearchShardTaskSettings; @@ -29,6 +28,7 @@ import org.opensearch.search.backpressure.trackers.TaskResourceUsageTrackerType; import org.opensearch.search.backpressure.trackers.TaskResourceUsageTrackers; import org.opensearch.search.backpressure.trackers.TaskResourceUsageTrackers.TaskResourceUsageTracker; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskCancellation; @@ -45,7 +45,6 @@ import java.io.IOException; import java.util.Collections; -import java.util.EnumMap; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -56,8 +55,6 @@ import java.util.concurrent.atomic.AtomicReference; import java.util.function.LongSupplier; -import static org.opensearch.search.ResourceType.CPU; -import static org.opensearch.search.ResourceType.MEMORY; import static org.opensearch.search.backpressure.SearchBackpressureTestHelpers.createMockTaskWithResourceStats; import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; @@ -100,10 +97,10 @@ public void testIsNodeInDuress() { NodeDuressTracker cpuUsageTracker = new NodeDuressTracker(() -> cpuUsage.get() >= 0.5, () -> 3); NodeDuressTracker heapUsageTracker = new NodeDuressTracker(() -> heapUsage.get() >= 0.5, () -> 3); - EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { + HashMap duressTrackers = new HashMap<>() { { - put(ResourceType.MEMORY, heapUsageTracker); - put(ResourceType.CPU, cpuUsageTracker); + put(ResourceType.fromName("jvm"), heapUsageTracker); + put(ResourceType.fromName("cpu"), cpuUsageTracker); } }; @@ -160,7 +157,7 @@ public void testTrackerStateUpdateOnSearchTaskCompletion() { mockTaskResourceTrackingService, threadPool, mockTimeNanosSupplier, - new NodeDuressTrackers(new EnumMap<>(ResourceType.class)), + new NodeDuressTrackers(new HashMap<>()), taskResourceUsageTrackers, new TaskResourceUsageTrackers(), taskManager @@ -191,7 +188,7 @@ public void testTrackerStateUpdateOnSearchShardTaskCompletion() { mockTaskResourceTrackingService, threadPool, mockTimeNanosSupplier, - new NodeDuressTrackers(new EnumMap<>(ResourceType.class)), + new NodeDuressTrackers(new HashMap<>()), new TaskResourceUsageTrackers(), taskResourceUsageTrackers, taskManager @@ -231,10 +228,10 @@ public void testSearchTaskInFlightCancellation() { NodeDuressTracker heapUsageTracker = new NodeDuressTracker(() -> false, () -> 3); - EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { + HashMap duressTrackers = new HashMap<>() { { - put(MEMORY, heapUsageTracker); - put(CPU, mockNodeDuressTracker); + put(ResourceType.fromName("jvm"), heapUsageTracker); + put(ResourceType.fromName("cpu"), mockNodeDuressTracker); } }; @@ -306,10 +303,10 @@ public void testSearchShardTaskInFlightCancellation() { LongSupplier mockTimeNanosSupplier = mockTime::get; NodeDuressTracker mockNodeDuressTracker = new NodeDuressTracker(() -> true, () -> 3); - EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { + HashMap duressTrackers = new HashMap<>() { { - put(MEMORY, new NodeDuressTracker(() -> false, () -> 3)); - put(CPU, mockNodeDuressTracker); + put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 3)); + put(ResourceType.fromName("cpu"), mockNodeDuressTracker); } }; NodeDuressTrackers nodeDuressTrackers = new NodeDuressTrackers(duressTrackers); @@ -399,10 +396,10 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { AtomicLong mockTime = new AtomicLong(0); LongSupplier mockTimeNanosSupplier = mockTime::get; - EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { + HashMap duressTrackers = new HashMap<>() { { - put(MEMORY, new NodeDuressTracker(() -> false, () -> 3)); - put(CPU, new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 3)); + put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> true, () -> 3)); } }; @@ -493,10 +490,10 @@ public void testNonCancellationWhenSearchTrafficIsNotQualifyingForCancellation() AtomicLong mockTime = new AtomicLong(0); LongSupplier mockTimeNanosSupplier = mockTime::get; - EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { + HashMap duressTrackers = new HashMap<>() { { - put(MEMORY, new NodeDuressTracker(() -> false, () -> 3)); - put(CPU, new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 3)); + put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> true, () -> 3)); } }; diff --git a/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java b/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java index 801576bdf89d4..8d1b3b1fba268 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java +++ b/server/src/test/java/org/opensearch/search/backpressure/trackers/NodeDuressTrackersTests.java @@ -8,19 +8,19 @@ package org.opensearch.search.backpressure.trackers; -import org.opensearch.search.ResourceType; import org.opensearch.search.backpressure.trackers.NodeDuressTrackers.NodeDuressTracker; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.test.OpenSearchTestCase; -import java.util.EnumMap; +import java.util.HashMap; public class NodeDuressTrackersTests extends OpenSearchTestCase { public void testNodeNotInDuress() { - EnumMap map = new EnumMap<>(ResourceType.class) { + HashMap map = new HashMap<>() { { - put(ResourceType.MEMORY, new NodeDuressTracker(() -> false, () -> 2)); - put(ResourceType.CPU, new NodeDuressTracker(() -> false, () -> 2)); + put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 2)); + put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> false, () -> 2)); } }; @@ -32,10 +32,10 @@ public void testNodeNotInDuress() { } public void testNodeInDuressWhenHeapInDuress() { - EnumMap map = new EnumMap<>(ResourceType.class) { + HashMap map = new HashMap<>() { { - put(ResourceType.MEMORY, new NodeDuressTracker(() -> true, () -> 3)); - put(ResourceType.CPU, new NodeDuressTracker(() -> false, () -> 1)); + put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> false, () -> 1)); } }; @@ -49,10 +49,10 @@ public void testNodeInDuressWhenHeapInDuress() { } public void testNodeInDuressWhenCPUInDuress() { - EnumMap map = new EnumMap<>(ResourceType.class) { + HashMap map = new HashMap<>() { { - put(ResourceType.MEMORY, new NodeDuressTracker(() -> false, () -> 1)); - put(ResourceType.CPU, new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 1)); + put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> true, () -> 3)); } }; @@ -66,10 +66,10 @@ public void testNodeInDuressWhenCPUInDuress() { } public void testNodeInDuressWhenCPUAndHeapInDuress() { - EnumMap map = new EnumMap<>(ResourceType.class) { + HashMap map = new HashMap<>() { { - put(ResourceType.MEMORY, new NodeDuressTracker(() -> true, () -> 3)); - put(ResourceType.CPU, new NodeDuressTracker(() -> false, () -> 3)); + put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> false, () -> 3)); } }; diff --git a/server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java new file mode 100644 index 0000000000000..804818667244d --- /dev/null +++ b/server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java @@ -0,0 +1,66 @@ +/* + * 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.search.querygroup; + +import org.opensearch.search.resourcetypes.ResourceType; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.List; +import java.util.Map; + +import static org.opensearch.search.querygroup.cancellation.QueryGroupTestHelpers.getRandomTask; + +public class QueryGroupLevelResourceUsageViewTests extends OpenSearchTestCase { + Map resourceUsage; + List activeTasks; + + public void setUp() throws Exception { + super.setUp(); + resourceUsage = Map.of(ResourceType.fromName("JVM"), 34L, ResourceType.fromName("CPU"), 12L); + activeTasks = List.of(getRandomTask(4321)); + } + + public void testGetResourceUsageData() { + QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( + "1234", + resourceUsage, + activeTasks + ); + Map resourceUsageData = queryGroupLevelResourceUsageView.getResourceUsageData(); + assertTrue(assertResourceUsageData(resourceUsageData)); + } + + public void testGetResourceUsageDataDefault() { + QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView("1234"); + Map resourceUsageData = queryGroupLevelResourceUsageView.getResourceUsageData(); + assertTrue(resourceUsageData.isEmpty()); + } + + public void testGetActiveTasks() { + QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( + "1234", + resourceUsage, + activeTasks + ); + List activeTasks = queryGroupLevelResourceUsageView.getActiveTasks(); + assertEquals(1, activeTasks.size()); + assertEquals(4321, activeTasks.get(0).getId()); + } + + public void testGetActiveTasksDefault() { + QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView("1234"); + List activeTasks = queryGroupLevelResourceUsageView.getActiveTasks(); + assertTrue(activeTasks.isEmpty()); + } + + private boolean assertResourceUsageData(Map resourceUsageData) { + return resourceUsageData.get(ResourceType.fromName("JVM")) == 34L && resourceUsageData.get(ResourceType.fromName("CPU")) == 12L; + } +} diff --git a/server/src/test/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellationTests.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellationTests.java new file mode 100644 index 0000000000000..afdb3ac5c24b6 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellationTests.java @@ -0,0 +1,175 @@ +/* + * 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.search.querygroup.cancellation; + +import org.opensearch.cluster.metadata.QueryGroup; +import org.opensearch.search.querygroup.QueryGroupLevelResourceUsageView; +import org.opensearch.tasks.TaskCancellation; +import org.opensearch.test.OpenSearchTestCase; +import org.junit.Before; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import org.mockito.MockitoAnnotations; + +public class DefaultTaskCancellationTests extends OpenSearchTestCase { + + private static class TestTaskCancellationImpl extends DefaultTaskCancellation { + + public TestTaskCancellationImpl( + TaskSelectionStrategy taskSelectionStrategy, + Map queryGroupLevelViews, + Set activeQueryGroups + ) { + super(taskSelectionStrategy, queryGroupLevelViews, activeQueryGroups); + } + + public List getQueryGroupsToCancelFrom() { + return new ArrayList<>(activeQueryGroups); + } + } + + private TaskSelectionStrategy taskSelectionStrategy; + private Map queryGroupLevelViews; + private Set activeQueryGroups; + private DefaultTaskCancellation taskCancellation; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + queryGroupLevelViews = new HashMap<>(); + activeQueryGroups = new HashSet<>(); + taskCancellation = new TestTaskCancellationImpl( + new TaskSelectionStrategyTests.TestTaskSelectionStrategy(), + queryGroupLevelViews, + activeQueryGroups + ); + } + + public void testGetCancellableTasksFrom_returnsTasksWhenBreachingThreshold() { + String id = "queryGroup1"; + String resourceTypeStr = "CPU"; + long usage = 50L; + long threshold = 10L; + QueryGroup queryGroup1 = QueryGroupTestHelpers.createQueryGroupMock(id, resourceTypeStr, threshold, usage); + // QueryGroup.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); + // when(queryGroup1.getThresholdInLong(ResourceType.fromName(resourceTypeStr))).thenReturn(resourceLimitMock); + QueryGroupLevelResourceUsageView mockView = QueryGroupTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + queryGroupLevelViews.put(id, mockView); + + List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(queryGroup1); + assertEquals(2, cancellableTasksFrom.size()); + assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); + assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); + } + + // public void testGetCancellableTasksFrom_returnsNoTasksWhenBreachingThreshold() { + // String id = "querygroup1"; + // String resourceTypeStr = "CPU"; + // long usage = 50L; + // long threshold = 100L; + // QueryGroup querygroup1 = QueryGroupTestHelpers.createQueryGroupMock(id, resourceTypeStr, threshold, usage); + // QueryGroup.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); + // when(querygroup1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); + // QueryGroupLevelResourceUsageView mockView = QueryGroupTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + // queryGroupLevelViews.put(id, mockView); + // activeQueryGroups.add(querygroup1); + // + // List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(querygroup1); + // assertTrue(cancellableTasksFrom.isEmpty()); + // } + // + // public void testCancelTasks_cancelsGivenTasks() { + // String id = "querygroup1"; + // String resourceTypeStr = "CPU"; + // long usage = 50L; + // long threshold = 10L; + // QueryGroup querygroup1 = QueryGroupTestHelpers.createQueryGroupMock(id, resourceTypeStr, threshold, usage); + // QueryGroup.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); + // when(querygroup1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); + // QueryGroupLevelResourceUsageView mockView = QueryGroupTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + // queryGroupLevelViews.put(id, mockView); + // activeQueryGroups.add(querygroup1); + // + // TestTaskCancellationImpl taskCancellation = new TestTaskCancellationImpl( + // new TaskSelectionStrategyTests.TestTaskSelectionStrategy(), + // queryGroupLevelViews, + // activeQueryGroups + // ); + // + // List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(); + // assertEquals(2, cancellableTasksFrom.size()); + // assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); + // assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); + // + // taskCancellation.cancelTasks(); + // assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); + // assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); + // } + // + // public void testGetAllCancellableTasks_ReturnsNoTasksWhenNotBreachingThresholds() { + // String id = "querygroup1"; + // String resourceTypeStr = "CPU"; + // long usage = 50L; + // long threshold = 100L; + // QueryGroup querygroup1 = QueryGroupTestHelpers.createQueryGroupMock(id, resourceTypeStr, threshold, usage); + // QueryGroup.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); + // when(querygroup1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); + // QueryGroupLevelResourceUsageView mockView = QueryGroupTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + // queryGroupLevelViews.put(id, mockView); + // activeQueryGroups.add(querygroup1); + // + // List allCancellableTasks = taskCancellation.getAllCancellableTasks(); + // assertTrue(allCancellableTasks.isEmpty()); + // } + // + // public void testGetAllCancellableTasks_ReturnsTasksWhenBreachingThresholds() { + // String id = "querygroup1"; + // String resourceTypeStr = "CPU"; + // long usage = 100L; + // long threshold = 50L; + // QueryGroup querygroup1 = QueryGroupTestHelpers.createQueryGroupMock(id, resourceTypeStr, threshold, usage); + // QueryGroup.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); + // when(querygroup1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); + // QueryGroupLevelResourceUsageView mockView = QueryGroupTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + // queryGroupLevelViews.put(id, mockView); + // activeQueryGroups.add(querygroup1); + // + // List allCancellableTasks = taskCancellation.getAllCancellableTasks(); + // assertEquals(2, allCancellableTasks.size()); + // assertEquals(1234, allCancellableTasks.get(0).getTask().getId()); + // assertEquals(4321, allCancellableTasks.get(1).getTask().getId()); + // } + // + // public void testGetCancellableTasksFrom_returnsTasksEvenWhenquerygroupIdNotFound() { + // String querygroup_id1 = "querygroup1"; + // String querygroup_id2 = "querygroup2"; + // String resourceTypeStr = "CPU"; + // long usage = 50L; + // long threshold = 10L; + // QueryGroup querygroup1 = QueryGroupTestHelpers.createQueryGroupMock(querygroup_id1, resourceTypeStr, threshold, usage); + // QueryGroup querygroup2 = QueryGroupTestHelpers.createQueryGroupMock(querygroup_id2, resourceTypeStr, threshold, usage); + // QueryGroup.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); + // when(querygroup1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); + // QueryGroupLevelResourceUsageView mockView = QueryGroupTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + // queryGroupLevelViews.put(querygroup_id1, mockView); + // activeQueryGroups.add(querygroup1); + // activeQueryGroups.add(querygroup2); + // + // List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(querygroup1); + // assertEquals(2, cancellableTasksFrom.size()); + // assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); + // assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); + // } +} diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java similarity index 94% rename from server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java rename to server/src/test/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java index e06d0d6e20128..f2918bff4cbdf 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.cancellation; +package org.opensearch.search.querygroup.cancellation; import org.opensearch.tasks.Task; import org.opensearch.test.OpenSearchTestCase; diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/QueryGroupTestHelpers.java similarity index 66% rename from server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java rename to server/src/test/java/org/opensearch/search/querygroup/cancellation/QueryGroupTestHelpers.java index 6c83548ae0b3a..cdcea2375aa9e 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java +++ b/server/src/test/java/org/opensearch/search/querygroup/cancellation/QueryGroupTestHelpers.java @@ -6,22 +6,23 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.cancellation; +package org.opensearch.search.querygroup.cancellation; import org.opensearch.action.search.SearchAction; import org.opensearch.action.search.SearchTask; -import org.opensearch.cluster.metadata.Sandbox; +import org.opensearch.cluster.metadata.QueryGroup; import org.opensearch.core.tasks.TaskId; import org.opensearch.core.tasks.resourcetracker.ResourceStats; import org.opensearch.core.tasks.resourcetracker.ResourceStatsType; import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric; -import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; +import org.opensearch.search.querygroup.QueryGroupLevelResourceUsageView; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.tasks.Task; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.function.Supplier; import org.mockito.Mockito; @@ -31,7 +32,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -public class SandboxTestHelpers { +public class QueryGroupTestHelpers { public static List getListOfTasks(long totalMemory) { List tasks = new ArrayList<>(); @@ -81,28 +82,19 @@ public static Task getRandomSearchTask(long id) { ); } - public static Sandbox createSandboxMock(String id, String resourceTypeStr, Long threshold, Long usage) { - Sandbox sandbox = Mockito.mock(Sandbox.class); - when(sandbox.getId()).thenReturn(id); + public static QueryGroup createQueryGroupMock(String id, String resourceTypeStr, Long threshold, Long usage) { + QueryGroup queryGroupMock = Mockito.mock(QueryGroup.class); + when(queryGroupMock.get_id()).thenReturn(id); - Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); - when(sandbox.getResourceLimits()).thenReturn(Collections.singletonList(resourceLimitMock)); - - return sandbox; - } - - public static Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeStr, Long threshold) { - Sandbox.ResourceLimit resourceLimitMock = mock(Sandbox.ResourceLimit.class); - SystemResource resourceType = SystemResource.fromString(resourceTypeStr); - when(resourceLimitMock.getResourceType()).thenReturn(resourceType); - when(resourceLimitMock.getThreshold()).thenReturn(threshold); - when(resourceLimitMock.getThresholdInLong()).thenReturn(threshold); - return resourceLimitMock; + ResourceType resourceType = ResourceType.fromName(resourceTypeStr); + when(queryGroupMock.getResourceLimits()).thenReturn(Map.of(resourceType, threshold)); + when(queryGroupMock.getThresholdInLong(resourceType)).thenReturn(threshold); + return queryGroupMock; } - public static SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) { - SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class); - SystemResource resourceType = SystemResource.fromString(resourceTypeStr); + public static QueryGroupLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) { + QueryGroupLevelResourceUsageView mockView = mock(QueryGroupLevelResourceUsageView.class); + ResourceType resourceType = ResourceType.fromName(resourceTypeStr); when(mockView.getResourceUsageData()).thenReturn(Collections.singletonMap(resourceType, usage)); when(mockView.getActiveTasks()).thenReturn(List.of(getRandomSearchTask(1234), getRandomSearchTask(4321))); return mockView; diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java similarity index 94% rename from server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java rename to server/src/test/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java index c88cb37781a10..147a6514a6a74 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.cancellation; +package org.opensearch.search.querygroup.cancellation; import org.opensearch.tasks.Task; import org.opensearch.test.OpenSearchTestCase; diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java similarity index 82% rename from server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java rename to server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java index 6fdae3a6fece9..afcf83229c72b 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java @@ -6,10 +6,10 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.cancellation; +package org.opensearch.search.querygroup.cancellation; import org.opensearch.core.tasks.resourcetracker.ResourceStats; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskCancellation; import org.opensearch.test.OpenSearchTestCase; @@ -30,8 +30,8 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsGrea TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); long threshold = 100L; long reduceBy = 50L; - SystemResource resourceType = SystemResource.fromString("JVM"); - List tasks = SandboxTestHelpers.getListOfTasks(threshold); + ResourceType resourceType = ResourceType.fromName("JVM"); + List tasks = QueryGroupTestHelpers.getListOfTasks(threshold); List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); assertFalse(selectedTasks.isEmpty()); @@ -42,8 +42,8 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsLess TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); long threshold = 100L; long reduceBy = -50L; - SystemResource resourceType = SystemResource.fromString("JVM"); - List tasks = SandboxTestHelpers.getListOfTasks(threshold); + ResourceType resourceType = ResourceType.fromName("JVM"); + List tasks = QueryGroupTestHelpers.getListOfTasks(threshold); try { testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); @@ -57,8 +57,8 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsEqua TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); long threshold = 100L; long reduceBy = 0; - SystemResource resourceType = SystemResource.fromString("JVM"); - List tasks = SandboxTestHelpers.getListOfTasks(threshold); + ResourceType resourceType = ResourceType.fromName("JVM"); + List tasks = QueryGroupTestHelpers.getListOfTasks(threshold); List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); assertTrue(selectedTasks.isEmpty()); diff --git a/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java similarity index 55% rename from server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java rename to server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java index 7a203d6afe9f2..27e923def2b95 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java +++ b/server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java @@ -6,16 +6,16 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.tracking; +package org.opensearch.search.querygroup.tracking; import org.opensearch.action.search.SearchShardTask; import org.opensearch.action.search.SearchTask; import org.opensearch.common.settings.Settings; import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; -import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; -import org.opensearch.search.sandboxing.SandboxTask; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; -import org.opensearch.search.sandboxing.tracker.SandboxResourceUsageTrackerService; +import org.opensearch.search.querygroup.QueryGroupLevelResourceUsageView; +import org.opensearch.search.querygroup.QueryGroupTask; +import org.opensearch.search.querygroup.tracker.QueryGroupResourceUsageTrackerService; +import org.opensearch.search.resourcetypes.ResourceType; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; @@ -38,18 +38,18 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -public class SandboxResourceUsageTrackerServiceTests extends OpenSearchTestCase { +public class QueryGroupResourceUsageTrackerServiceTests extends OpenSearchTestCase { TestThreadPool threadPool; TaskManager taskManager; TaskResourceTrackingService mockTaskResourceTrackingService; - SandboxResourceUsageTrackerService sandboxResourceUsageTrackerService; + QueryGroupResourceUsageTrackerService queryGroupResourceUsageTrackerService; @Before public void setup() { taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet()); threadPool = new TestThreadPool(getTestName()); mockTaskResourceTrackingService = mock(TaskResourceTrackingService.class); - sandboxResourceUsageTrackerService = new SandboxResourceUsageTrackerService(taskManager, mockTaskResourceTrackingService); + queryGroupResourceUsageTrackerService = new QueryGroupResourceUsageTrackerService(taskManager, mockTaskResourceTrackingService); } @After @@ -57,56 +57,57 @@ public void cleanup() { ThreadPool.terminate(threadPool, 5, TimeUnit.SECONDS); } - public void testConstructSandboxLevelViews_CreatesSandboxLevelUsageView_WhenTasksArePresent() { - List sandboxIds = List.of("sandbox1", "sandbox2", "sandbox3"); + public void testConstructQueryGroupLevelViews_CreatesQueryGroupLevelUsageView_WhenTasksArePresent() { + List queryGroupIds = List.of("queryGroup1", "queryGroup2", "queryGroup3"); - Map activeSearchShardTasks = createActiveSearchShardTasks(sandboxIds); + Map activeSearchShardTasks = createActiveSearchShardTasks(queryGroupIds); when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); - Map stringSandboxLevelResourceUsageViewMap = sandboxResourceUsageTrackerService - .constructSandboxLevelUsageViews(); + Map stringQueryGroupLevelResourceUsageViewMap = queryGroupResourceUsageTrackerService + .constructQueryGroupLevelUsageViews(); - for (String sandboxId : sandboxIds) { + for (String queryGroupId : queryGroupIds) { assertEquals( 400, - (long) stringSandboxLevelResourceUsageViewMap.get(sandboxId).getResourceUsageData().get(SystemResource.fromString("JVM")) + (long) stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId).getResourceUsageData().get(ResourceType.fromName("JVM")) ); - assertEquals(2, stringSandboxLevelResourceUsageViewMap.get(sandboxId).getActiveTasks().size()); + assertEquals(2, stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId).getActiveTasks().size()); } } - public void testConstructSandboxLevelViews_CreatesSandboxLevelUsageView_WhenTasksAreNotPresent() { - Map stringSandboxLevelResourceUsageViewMap = sandboxResourceUsageTrackerService - .constructSandboxLevelUsageViews(); - assertTrue(stringSandboxLevelResourceUsageViewMap.isEmpty()); + public void testConstructQueryGroupLevelViews_CreatesQueryGroupLevelUsageView_WhenTasksAreNotPresent() { + Map stringQueryGroupLevelResourceUsageViewMap = queryGroupResourceUsageTrackerService + .constructQueryGroupLevelUsageViews(); + assertTrue(stringQueryGroupLevelResourceUsageViewMap.isEmpty()); } - public void testConstructSandboxLevelUsageViews_WithTasksHavingDifferentResourceUsage() { + public void testConstructQueryGroupLevelUsageViews_WithTasksHavingDifferentResourceUsage() { Map activeSearchShardTasks = new HashMap<>(); - activeSearchShardTasks.put(1L, createMockTask(SearchShardTask.class, 100, 200, "sandbox1")); - activeSearchShardTasks.put(2L, createMockTask(SearchShardTask.class, 200, 400, "sandbox1")); + activeSearchShardTasks.put(1L, createMockTask(SearchShardTask.class, 100, 200, "queryGroup1")); + activeSearchShardTasks.put(2L, createMockTask(SearchShardTask.class, 200, 400, "queryGroup1")); when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); - Map sandboxViews = sandboxResourceUsageTrackerService.constructSandboxLevelUsageViews(); + Map queryGroupViews = queryGroupResourceUsageTrackerService + .constructQueryGroupLevelUsageViews(); - assertEquals(600, (long) sandboxViews.get("sandbox1").getResourceUsageData().get(SystemResource.fromString("JVM"))); - assertEquals(2, sandboxViews.get("sandbox1").getActiveTasks().size()); + assertEquals(600, (long) queryGroupViews.get("queryGroup1").getResourceUsageData().get(ResourceType.fromName("JVM"))); + assertEquals(2, queryGroupViews.get("queryGroup1").getActiveTasks().size()); } - private Map createActiveSearchShardTasks(List sandboxIds) { + private Map createActiveSearchShardTasks(List queryGroupIds) { Map activeSearchShardTasks = new HashMap<>(); long task_id = 0; - for (String sandboxId : sandboxIds) { + for (String queryGroupId : queryGroupIds) { for (int i = 0; i < 2; i++) { - activeSearchShardTasks.put(++task_id, createMockTask(SearchShardTask.class, 100, 200, sandboxId)); + activeSearchShardTasks.put(++task_id, createMockTask(SearchShardTask.class, 100, 200, queryGroupId)); } } return activeSearchShardTasks; } - private T createMockTask(Class type, long cpuUsage, long heapUsage, String sandboxId) { + private T createMockTask(Class type, long cpuUsage, long heapUsage, String queryGroupId) { T task = mock(type); if (task instanceof SearchTask || task instanceof SearchShardTask) { - when(((SandboxTask) task).getSandboxId()).thenReturn(sandboxId); + when(((QueryGroupTask) task).getQueryGroupId()).thenReturn(queryGroupId); } when(task.getTotalResourceStats()).thenReturn(new TaskResourceUsage(cpuUsage, heapUsage)); when(task.getStartTimeNanos()).thenReturn((long) 0); diff --git a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java deleted file mode 100644 index 17e987d608d3e..0000000000000 --- a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java +++ /dev/null @@ -1,59 +0,0 @@ -/* - * 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.search.sandboxing; - -import org.opensearch.search.sandboxing.resourcetype.SystemResource; -import org.opensearch.tasks.Task; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.List; -import java.util.Map; - -import static org.opensearch.search.sandboxing.cancellation.SandboxTestHelpers.getRandomTask; - -public class SandboxLevelResourceUsageViewTests extends OpenSearchTestCase { - Map resourceUsage; - List activeTasks; - - public void setUp() throws Exception { - super.setUp(); - resourceUsage = Map.of(SystemResource.fromString("JVM"), 34L, SystemResource.fromString("CPU"), 12L); - activeTasks = List.of(getRandomTask(4321)); - } - - public void testGetResourceUsageData() { - SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234", resourceUsage, activeTasks); - Map resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData(); - assertTrue(assertResourceUsageData(resourceUsageData)); - } - - public void testGetResourceUsageDataDefault() { - SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234"); - Map resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData(); - assertTrue(resourceUsageData.isEmpty()); - } - - public void testGetActiveTasks() { - SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234", resourceUsage, activeTasks); - List activeTasks = sandboxLevelResourceUsageView.getActiveTasks(); - assertEquals(1, activeTasks.size()); - assertEquals(4321, activeTasks.get(0).getId()); - } - - public void testGetActiveTasksDefault() { - SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234"); - List activeTasks = sandboxLevelResourceUsageView.getActiveTasks(); - assertTrue(activeTasks.isEmpty()); - } - - private boolean assertResourceUsageData(Map resourceUsageData) { - return resourceUsageData.get(SystemResource.fromString("JVM")) == 34L - && resourceUsageData.get(SystemResource.fromString("CPU")) == 12L; - } -} diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java deleted file mode 100644 index d836089b74119..0000000000000 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java +++ /dev/null @@ -1,180 +0,0 @@ -/* - * 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.search.sandboxing.cancellation; - -import org.opensearch.cluster.metadata.Sandbox; -import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; -import org.opensearch.search.sandboxing.resourcetype.SystemResource; -import org.opensearch.tasks.TaskCancellation; -import org.opensearch.test.OpenSearchTestCase; -import org.junit.Before; - -import java.util.ArrayList; -import java.util.HashMap; -import java.util.HashSet; -import java.util.List; -import java.util.Map; -import java.util.Set; - -import org.mockito.MockitoAnnotations; - -import static org.opensearch.search.sandboxing.cancellation.SandboxTestHelpers.createResourceLimitMock; -import static org.mockito.Mockito.when; - -public class DefaultTaskCancellationTests extends OpenSearchTestCase { - - private static class TestTaskCancellationImpl extends DefaultTaskCancellation { - - public TestTaskCancellationImpl( - TaskSelectionStrategy taskSelectionStrategy, - Map sandboxLevelViews, - Set activeSandboxes - ) { - super(taskSelectionStrategy, sandboxLevelViews, activeSandboxes); - } - - @Override - public List getSandboxesToCancelFrom() { - return new ArrayList<>(activeSandboxes); - } - } - - private TaskSelectionStrategy taskSelectionStrategy; - private Map sandboxLevelViews; - private Set activeSandboxes; - private DefaultTaskCancellation taskCancellation; - - @Before - public void setup() { - MockitoAnnotations.openMocks(this); - sandboxLevelViews = new HashMap<>(); - activeSandboxes = new HashSet<>(); - taskCancellation = new TestTaskCancellationImpl( - new TaskSelectionStrategyTests.TestTaskSelectionStrategy(), - sandboxLevelViews, - activeSandboxes - ); - } - - public void testGetCancellableTasksFrom_returnsTasksWhenBreachingThreshold() { - String id = "sandbox1"; - String resourceTypeStr = "CPU"; - long usage = 50L; - long threshold = 10L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, threshold, usage); - Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); - when(sandbox1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); - SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView); - - List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(sandbox1); - assertEquals(2, cancellableTasksFrom.size()); - assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); - assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); - } - - public void testGetCancellableTasksFrom_returnsNoTasksWhenBreachingThreshold() { - String id = "sandbox1"; - String resourceTypeStr = "CPU"; - long usage = 50L; - long threshold = 100L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, threshold, usage); - Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); - when(sandbox1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); - SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView); - activeSandboxes.add(sandbox1); - - List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(sandbox1); - assertTrue(cancellableTasksFrom.isEmpty()); - } - - public void testCancelTasks_cancelsGivenTasks() { - String id = "sandbox1"; - String resourceTypeStr = "CPU"; - long usage = 50L; - long threshold = 10L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, threshold, usage); - Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); - when(sandbox1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); - SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView); - activeSandboxes.add(sandbox1); - - TestTaskCancellationImpl taskCancellation = new TestTaskCancellationImpl( - new TaskSelectionStrategyTests.TestTaskSelectionStrategy(), - sandboxLevelViews, - activeSandboxes - ); - - List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(); - assertEquals(2, cancellableTasksFrom.size()); - assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); - assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); - - taskCancellation.cancelTasks(); - assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); - assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); - } - - public void testGetAllCancellableTasks_ReturnsNoTasksWhenNotBreachingThresholds() { - String id = "sandbox1"; - String resourceTypeStr = "CPU"; - long usage = 50L; - long threshold = 100L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, threshold, usage); - Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); - when(sandbox1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); - SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView); - activeSandboxes.add(sandbox1); - - List allCancellableTasks = taskCancellation.getAllCancellableTasks(); - assertTrue(allCancellableTasks.isEmpty()); - } - - public void testGetAllCancellableTasks_ReturnsTasksWhenBreachingThresholds() { - String id = "sandbox1"; - String resourceTypeStr = "CPU"; - long usage = 100L; - long threshold = 50L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, threshold, usage); - Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); - when(sandbox1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); - SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView); - activeSandboxes.add(sandbox1); - - List allCancellableTasks = taskCancellation.getAllCancellableTasks(); - assertEquals(2, allCancellableTasks.size()); - assertEquals(1234, allCancellableTasks.get(0).getTask().getId()); - assertEquals(4321, allCancellableTasks.get(1).getTask().getId()); - } - - public void testGetCancellableTasksFrom_returnsTasksEvenWhenSandboxIdNotFound() { - String sandbox_id1 = "sandbox1"; - String sandbox_id2 = "sandbox2"; - String resourceTypeStr = "CPU"; - long usage = 50L; - long threshold = 10L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(sandbox_id1, resourceTypeStr, threshold, usage); - Sandbox sandbox2 = SandboxTestHelpers.createSandboxMock(sandbox_id2, resourceTypeStr, threshold, usage); - Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); - when(sandbox1.getResourceLimitFor(SystemResource.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); - SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(sandbox_id1, mockView); - activeSandboxes.add(sandbox1); - activeSandboxes.add(sandbox2); - - List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(sandbox1); - assertEquals(2, cancellableTasksFrom.size()); - assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); - assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); - } -}