From 4912d6b5da0f39f2877e5b23e22de586864b42f0 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 10:05:17 -0700 Subject: [PATCH 01/66] initial code for the sandbox resource tracking and cancellation framework Signed-off-by: Kiran Prakash --- .../opensearch/cluster/metadata/Metadata.java | 7 +- .../opensearch/cluster/metadata/Sandbox.java | 87 ++++++++++++ .../SandboxLevelResourceUsageView.java | 60 +++++++++ .../search/sandboxing/SandboxService.java | 103 ++++++++++++++ .../search/sandboxing/SandboxTask.java | 18 +++ .../AbstractTaskCancellation.java | 125 +++++++++++++++++ .../AbstractTaskSelectionStrategy.java | 49 +++++++ .../cancellation/DefaultTaskCancellation.java | 63 +++++++++ .../LongestRunningTaskFirstStrategy.java | 21 +++ .../ShortestRunningTaskFirstStrategy.java | 21 +++ .../cancellation/TaskSelectionStrategy.java | 31 +++++ .../sandboxing/cancellation/package-info.java | 12 ++ .../sandboxing/module/SandboxModule.java | 29 ++++ .../sandboxing/module/package-info.java | 13 ++ .../search/sandboxing/package-info.java | 12 ++ .../resourcetype/CpuTimeResourceType.java | 18 +++ .../resourcetype/JvmMemoryResourceType.java | 18 +++ .../resourcetype/SandboxResourceType.java | 27 ++++ .../sandboxing/resourcetype/package-info.java | 9 ++ .../SandboxResourceUsageTrackerService.java | 127 ++++++++++++++++++ .../tracker/SandboxUsageTracker.java | 24 ++++ .../sandboxing/tracker/package-info.java | 12 ++ ...RunningTaskFirstStrategyStrategyTests.java | 34 +++++ ...andboxCancellationStrategyTestHelpers.java | 73 ++++++++++ ...RunningTaskFirstStrategyStrategyTests.java | 34 +++++ .../TaskSelectionStrategyTests.java | 76 +++++++++++ ...ndboxResourceUsageTrackerServiceTests.java | 119 ++++++++++++++++ 27 files changed, 1221 insertions(+), 1 deletion(-) create mode 100644 server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/package-info.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/module/SandboxModule.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/module/package-info.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/package-info.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/resourcetype/package-info.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxUsageTracker.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/tracker/package-info.java create mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java create mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java create mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java create mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java create mode 100644 server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java 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 440b9e267cf0a..feabdd095da04 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -127,7 +127,12 @@ public boolean isSegmentReplicationEnabled(String indexName) { .orElse(false); } - /** + public Map sandboxes() { + // stub + return Collections.emptyMap(); + } + + /** * Context of the XContent. * * @opensearch.api diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java new file mode 100644 index 0000000000000..c868be280ea4b --- /dev/null +++ b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java @@ -0,0 +1,87 @@ +/* + * 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.SandboxResourceType; + +import java.util.Collections; +import java.util.List; + +@ExperimentalApi +public class Sandbox { + //TODO Kaushal should have implemented hashcode and equals + private SandboxMode mode; + + public SandboxMode getMode() { + return mode; + } + + public ResourceLimit getResourceLimitFor(SandboxResourceType resourceType) { + return null; + } + + public String getName() { + return ""; + } + + public String getId() { + return ""; + } + + public List getResourceLimits() { + return Collections.emptyList(); + } + + @ExperimentalApi + public class ResourceLimit { + public Long getThresholdInLong() { + return 0L; + } + + public SandboxResourceType getResourceType() { + return null; + } + + public Long getThreshold() { + return 0L; + } + } + + @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/sandboxing/SandboxLevelResourceUsageView.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java new file mode 100644 index 0000000000000..aa181bb85c3c0 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.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.sandboxing; + +import org.opensearch.common.annotation.ExperimentalApi; +import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.tasks.Task; + +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +@ExperimentalApi +public class SandboxLevelResourceUsageView { + + private final String sandboxId; + private final Map resourceUsage; + private final List activeTasks; + + public SandboxLevelResourceUsageView(String sandboxId) { + this.sandboxId = sandboxId; + this.resourceUsage = new HashMap<>(); + this.activeTasks = new ArrayList<>(); + } + + public SandboxLevelResourceUsageView(String sandboxId, Map resourceUsage, List activeTasks) { + this.sandboxId = sandboxId; + this.resourceUsage = resourceUsage; + this.activeTasks = activeTasks; + } + + public Map getResourceUsageData() { + return resourceUsage; + } + + public List getActiveTasks() { + return activeTasks; + } + + @Override + 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); + } + + @Override + public int hashCode() { + return Objects.hashCode(sandboxId); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java new file mode 100644 index 0000000000000..662258bb7fd03 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java @@ -0,0 +1,103 @@ +/* + * 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.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.cluster.metadata.Sandbox; +import org.opensearch.cluster.service.ClusterService; +import org.opensearch.common.inject.Inject; +import org.opensearch.common.lifecycle.AbstractLifecycleComponent; +import org.opensearch.search.sandboxing.cancellation.DefaultTaskCancellation; +import org.opensearch.search.sandboxing.cancellation.LongestRunningTaskFirstStrategy; +import org.opensearch.search.sandboxing.tracker.SandboxUsageTracker; +import org.opensearch.threadpool.Scheduler; +import org.opensearch.threadpool.ThreadPool; + +import java.io.IOException; +import java.util.HashSet; +import java.util.Map; +import java.util.Set; + +/** + * Main service which will run periodically to track and cancel resource constraint violating tasks in sandboxes + */ +public class SandboxService extends AbstractLifecycleComponent { + private static final Logger logger = LogManager.getLogger(SandboxService.class); + + private final SandboxUsageTracker sandboxUsageTracker; + private volatile Scheduler.Cancellable scheduledFuture; + private final ThreadPool threadPool; + private final ClusterService clusterService; + + /** + * Guice managed constructor + * + * @param sandboxUsageTracker + * @param sandboxPruner + * @param sandboxServiceSettings + * @param threadPool + */ + @Inject + public SandboxService( + SandboxUsageTracker sandboxUsageTracker, + ClusterService clusterService, + ThreadPool threadPool + ) { + this.sandboxUsageTracker = sandboxUsageTracker; + this.sandboxServiceSettings = sandboxServiceSettings; + this.sandboxPruner = sandboxPruner; + this.clusterService = clusterService; + this.threadPool = threadPool; + } + + /** + * run at regular interval + */ + private void doRun() { + Map sandboxLevelResourceUsageViews = sandboxUsageTracker.constructSandboxLevelUsageViews(); + Set activeSandboxes = getActiveSandboxes(); + DefaultTaskCancellation taskCancellation = new DefaultTaskCancellation( + new LongestRunningTaskFirstStrategy(), + sandboxLevelResourceUsageViews, + activeSandboxes + ); + taskCancellation.cancelTasks(); + // TODO Prune the sandboxes + } + + private Set getActiveSandboxes() { + return new HashSet<>(clusterService.state().metadata().sandboxes().values()); + } + + /** + * {@link AbstractLifecycleComponent} lifecycle method + */ + @Override + protected void doStart() { + scheduledFuture = threadPool.scheduleWithFixedDelay(() -> { + try { + doRun(); + } catch (Exception e) { + logger.debug("Exception occurred in Query Sandbox service", e); + } + }, sandboxServiceSettings.getRunIntervalMillis(), ThreadPool.Names.GENERIC); + } + + @Override + protected void doStop() { + if (scheduledFuture != null) { + scheduledFuture.cancel(); + } + } + + @Override + protected void doClose() throws IOException { + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java new file mode 100644 index 0000000000000..92aba3759f115 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java @@ -0,0 +1,18 @@ +/* + * 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; + +/** + * This interface can be implemented by tasks which will be tracked and monitored using {@link org.opensearch.cluster.metadata.ResourceLimitGroup} + */ +public interface SandboxTask { + void setSandboxId(String sandboxId); + + String getSandboxId(); +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java new file mode 100644 index 0000000000000..9bf96bf63578c --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java @@ -0,0 +1,125 @@ +/* + * 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.SandboxResourceType; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskCancellation; + +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; + +/** + * Abstract class that provides a structure for task cancellation. + * This class is extended by other classes to provide specific task cancellation strategies. + */ +public abstract class AbstractTaskCancellation { + // Strategy for selecting tasks to cancel + protected final TaskSelectionStrategy taskSelectionStrategy; + // Views of resource usage at the sandbox level + protected final Map sandboxLevelViews; + // Set of active sandboxes + protected final Set activeSandboxes; + + public AbstractTaskCancellation( + 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(); + } + } + + /** + * Abstract method to get the list of sandboxes from which tasks can be cancelled. + * This method needs to be implemented by subclasses. + * + * @return List of sandboxes + */ + abstract List getSandboxesToCancelFrom(); + + /** + * 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, SandboxResourceType resourceType) { + long reduceBy = getReduceBy(sandbox, resourceType); + return reduceBy > 0; + } + + private List getTaskCancellations(Sandbox sandbox, SandboxResourceType resourceType) { + return taskSelectionStrategy.selectTasksForCancellation( + getAllTasksInSandbox(sandbox.getId()), + getReduceBy(sandbox, resourceType), + resourceType) + .stream() + .map(task -> createTaskCancellation((CancellableTask) task)) + .collect(Collectors.toList()); + } + + private long getReduceBy(Sandbox sandbox, SandboxResourceType resourceType) { + return getUsage(sandbox, resourceType) - sandbox.getResourceLimitFor(resourceType).getThresholdInLong(); + } + + private Long getUsage(Sandbox sandbox, SandboxResourceType resourceType) { + return sandboxLevelViews.get(sandbox.getId()).getResourceUsageData().get(resourceType); + } + + private List getAllTasksInSandbox(String sandboxId) { + return sandboxLevelViews.get(sandboxId).getActiveTasks(); + } + + private TaskCancellation createTaskCancellation(CancellableTask task) { + // todo add reasons and callbacks + return new TaskCancellation(task, List.of(), List.of(this::callbackOnCancel)); + } + + private void callbackOnCancel() { + // todo Implement callback logic here + System.out.println("Task was cancelled."); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java new file mode 100644 index 0000000000000..de0cc19090a4d --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java @@ -0,0 +1,49 @@ +/* + * 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.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.tasks.Task; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.Comparator; +import java.util.List; +import java.util.stream.Collectors; + +public abstract class AbstractTaskSelectionStrategy implements TaskSelectionStrategy { + + public abstract Comparator sortingCondition(); + + @Override + public List selectTasksForCancellation(List tasks, long limit, SandboxResourceType resourceType) { + if (limit < 0) { + throw new IllegalArgumentException("reduceBy has to be greater than zero"); + } + if(limit == 0) { + return Collections.emptyList(); + } + + List sortedTasks = tasks.stream() + .sorted(sortingCondition()) + .collect(Collectors.toList()); + + List selectedTasks = new ArrayList<>(); + long accumulated = 0; + + for (Task task : sortedTasks) { + selectedTasks.add(task); + accumulated += resourceType.getResourceUsage(task); + if (accumulated >= limit) { + break; + } + } + return selectedTasks; + } +} 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 new file mode 100644 index 0000000000000..f5d21f8af5e4e --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java @@ -0,0 +1,63 @@ +/* + * 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.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; + +import java.util.ArrayList; +import java.util.List; +import java.util.Map; +import java.util.Set; + +public class DefaultTaskCancellation extends AbstractTaskCancellation { + public DefaultTaskCancellation( + TaskSelectionStrategy cancellationStrategy, + Map sandboxLevelViews, + Set activeSandboxes + ) { + super(cancellationStrategy, sandboxLevelViews, activeSandboxes); + } + + /** + * // TODO + * This should cover 3 scenarios + * - if node not in duress + * - pick sandboxes in enforced mode only + * - if node in duress + * - pick sandboxes in enforced mode + * - tasks running in deleted sandboxes with tasks running + * - pick sandboxes in enforced mode + */ + public List getSandboxesToCancelFrom() { + final List sandboxesToCancelFrom = new ArrayList<>(); + + for (Sandbox sandbox : this.activeSandboxes) { + Map currentResourceUsage = getResourceUsage(sandbox.getId()); + + for (Sandbox.ResourceLimit resourceLimit : sandbox.getResourceLimits()) { + if (isBreachingThreshold(currentResourceUsage, resourceLimit)) { + sandboxesToCancelFrom.add(sandbox); + break; + } + } + } + + return sandboxesToCancelFrom; + } + + private boolean isBreachingThreshold(Map currentResourceUsage, Sandbox.ResourceLimit resourceLimit) { + return currentResourceUsage.get(resourceLimit.getResourceType()) > resourceLimit.getThreshold(); + } + + private Map getResourceUsage(String sandboxId) { + return sandboxLevelViews.get(sandboxId).getResourceUsageData(); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java new file mode 100644 index 0000000000000..d97f1b26b5362 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java @@ -0,0 +1,21 @@ +/* + * 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.tasks.Task; + +import java.util.Comparator; + +public class LongestRunningTaskFirstStrategy extends AbstractTaskSelectionStrategy { + + @Override + public Comparator sortingCondition() { + return Comparator.comparingLong(Task::getStartTime).reversed(); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java new file mode 100644 index 0000000000000..f4e1c7f4fe9bc --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java @@ -0,0 +1,21 @@ +/* + * 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.tasks.Task; + +import java.util.Comparator; + +public class ShortestRunningTaskFirstStrategy extends AbstractTaskSelectionStrategy { + + @Override + public Comparator sortingCondition() { + return Comparator.comparingLong(Task::getStartTime); + } +} \ No newline at end of file diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java new file mode 100644 index 0000000000000..b5345be6dda8b --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.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.sandboxing.cancellation; + +import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.tasks.Task; + +import java.util.List; + +/** + * Interface for strategies to select tasks for cancellation. + * Implementations of this interface define how tasks are selected for cancellation based on resource usage. + */ +public interface TaskSelectionStrategy { + /** + * Determines which tasks should be cancelled based on the provided criteria. + * + * @param tasks List of tasks available for cancellation. + * @param limit The amount of tasks to select whose resources reach this limit + * @param resourceType The type of resource that needs to be reduced, guiding the selection process. + * + * @return List of tasks that should be cancelled. + */ + List selectTasksForCancellation(List tasks, long limit, SandboxResourceType resourceType); +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/package-info.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/package-info.java new file mode 100644 index 0000000000000..7bda195ba18ff --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/package-info.java @@ -0,0 +1,12 @@ +/* + * 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 cancellation related abstracts + */ +package org.opensearch.search.sandboxing.cancellation; 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 new file mode 100644 index 0000000000000..8da5f20ab022f --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/module/SandboxModule.java @@ -0,0 +1,29 @@ +/* + * 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.SandboxUsageTracker; +import org.opensearch.search.sandboxing.tracker.SandboxResourceUsageTrackerService; + +/** + * 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(); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/module/package-info.java b/server/src/main/java/org/opensearch/search/sandboxing/module/package-info.java new file mode 100644 index 0000000000000..5be0923cc3c62 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/module/package-info.java @@ -0,0 +1,13 @@ +/* + * 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. + */ + +/** + * Guice Module + */ + +package org.opensearch.search.sandboxing.module; diff --git a/server/src/main/java/org/opensearch/search/sandboxing/package-info.java b/server/src/main/java/org/opensearch/search/sandboxing/package-info.java new file mode 100644 index 0000000000000..3dfecf384df44 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/package-info.java @@ -0,0 +1,12 @@ +/* + * 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. + */ + +/** + * Query Sandboxing related artifacts + */ +package org.opensearch.search.sandboxing; diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java new file mode 100644 index 0000000000000..6c4dd7fbaa320 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java @@ -0,0 +1,18 @@ +/* + * 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; + +public class CpuTimeResourceType extends SandboxResourceType { + @Override + public long getResourceUsage(Task task) { + return task.getTotalResourceStats().getCpuTimeInNanos(); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java new file mode 100644 index 0000000000000..91f6e73364b5f --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java @@ -0,0 +1,18 @@ +/* + * 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; + +public class JvmMemoryResourceType extends SandboxResourceType { + @Override + public long getResourceUsage(Task task) { + return task.getTotalResourceStats().getMemoryInBytes(); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java new file mode 100644 index 0000000000000..9b849718e6a11 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java @@ -0,0 +1,27 @@ +/* + * 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; + +@ExperimentalApi +public abstract class SandboxResourceType { + public abstract long getResourceUsage(Task task); + + public static SandboxResourceType fromString(String type) { + if (type.equalsIgnoreCase("JVM")) { + return new JvmMemoryResourceType(); + } else if (type.equalsIgnoreCase("CPU")) { + return new CpuTimeResourceType(); + } 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 new file mode 100644 index 0000000000000..97d449aa49912 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/package-info.java @@ -0,0 +1,9 @@ +/* + * 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; \ No newline at end of file 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 new file mode 100644 index 0000000000000..84db827944c9f --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java @@ -0,0 +1,127 @@ +/* + * 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.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.SandboxTask; +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(SandboxResourceType.fromString("JVM")); + + private final TaskManager taskManager; + private final TaskResourceTrackingService taskResourceTrackingService; + + /** + * SandboxResourceTrackerService constructor + * + * @param taskManager + * @param taskResourceTrackingService + */ + @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 (SandboxResourceType 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/main/java/org/opensearch/search/sandboxing/tracker/SandboxUsageTracker.java b/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxUsageTracker.java new file mode 100644 index 0000000000000..b08e110e090f9 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxUsageTracker.java @@ -0,0 +1,24 @@ +/* + * 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.search.sandboxing.SandboxLevelResourceUsageView; + +import java.util.Map; + +/** + * This interface is mainly for tracking the resourceLimitGroup level resource usages + */ +public interface SandboxUsageTracker { + /** + * updates the current resource usage of resourceLimitGroups + */ + + Map constructSandboxLevelUsageViews(); +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/tracker/package-info.java b/server/src/main/java/org/opensearch/search/sandboxing/tracker/package-info.java new file mode 100644 index 0000000000000..fcf0c504d3752 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/tracker/package-info.java @@ -0,0 +1,12 @@ +/* + * 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. + */ + +/** + * ResourceLimitGroup resource tracking artifacts + */ +package org.opensearch.search.sandboxing.tracker; diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java new file mode 100644 index 0000000000000..8c19de92f1b9f --- /dev/null +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java @@ -0,0 +1,34 @@ +/* + * 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.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Arrays; +import java.util.List; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class LongestRunningTaskFirstStrategyStrategyTests extends OpenSearchTestCase { + public void testSortingCondition() { + Task task1 = mock(Task.class); + Task task2 = mock(Task.class); + Task task3 = mock(Task.class); + when(task1.getStartTime()).thenReturn(100L); + when(task2.getStartTime()).thenReturn(200L); + when(task3.getStartTime()).thenReturn(300L); + + List tasks = Arrays.asList(task1, task3, task2); + tasks.sort(new LongestRunningTaskFirstStrategy().sortingCondition()); + + assertEquals(Arrays.asList(task3, task2, task1), tasks); + } +} diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java new file mode 100644 index 0000000000000..d4a5fa0c6d3f1 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.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.sandboxing.cancellation; + +import org.opensearch.action.search.SearchAction; +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.core.tasks.resourcetracker.TaskResourceUsage; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.opensearch.test.OpenSearchTestCase.randomLong; +import static org.opensearch.test.OpenSearchTestCase.randomLongBetween; + +public class SandboxCancellationStrategyTestHelpers { + + public static List getListOfTasks(long totalMemory) { + List tasks = new ArrayList<>(); + + while (totalMemory > 0) { + long id = randomLong(); + final Task task = getRandomTask(id); + long initial_memory = randomLongBetween(1, 100); + + ResourceUsageMetric[] initialTaskResourceMetrics = new ResourceUsageMetric[]{ + new ResourceUsageMetric(ResourceStats.MEMORY, initial_memory) + }; + task.startThreadResourceTracking(id, ResourceStatsType.WORKER_STATS, initialTaskResourceMetrics); + + long memory = initial_memory + randomLongBetween(1, 10000); + + totalMemory -= memory - initial_memory; + + ResourceUsageMetric[] taskResourceMetrics = new ResourceUsageMetric[]{ + new ResourceUsageMetric(ResourceStats.MEMORY, memory), + }; + task.updateThreadResourceStats(id, ResourceStatsType.WORKER_STATS, taskResourceMetrics); + task.stopThreadResourceTracking(id, ResourceStatsType.WORKER_STATS); + tasks.add(task); + } + + return tasks; + } + + public static Task getRandomTask(long id) { + return new Task( + id, + "transport", + SearchAction.NAME, + "test description", + new TaskId(randomLong() + ":" + randomLong()), + Collections.emptyMap() + ); + } +} diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java new file mode 100644 index 0000000000000..2422d41ccd2b5 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java @@ -0,0 +1,34 @@ +/* + * 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.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Arrays; +import java.util.List; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class ShortestRunningTaskFirstStrategyStrategyTests extends OpenSearchTestCase { + public void testSortingCondition() { + Task task1 = mock(Task.class); + Task task2 = mock(Task.class); + Task task3 = mock(Task.class); + when(task1.getStartTime()).thenReturn(100L); + when(task2.getStartTime()).thenReturn(200L); + when(task3.getStartTime()).thenReturn(300L); + + List tasks = Arrays.asList(task1, task3, task2); + tasks.sort(new ShortestRunningTaskFirstStrategy().sortingCondition()); + + assertEquals(Arrays.asList(task1, task2, task3), tasks); + } +} diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java new file mode 100644 index 0000000000000..39634d9e026b0 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java @@ -0,0 +1,76 @@ +/* + * 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.core.tasks.resourcetracker.ResourceStats; +import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Comparator; +import java.util.List; + +public class TaskSelectionStrategyTests extends OpenSearchTestCase { + + private class TestTaskSelectionStrategy extends AbstractTaskSelectionStrategy { + @Override + public Comparator sortingCondition() { + return Comparator.comparingLong(Task::getStartTime).reversed(); + } + } + + public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsGreaterThanZero() { + TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); + long threshold = 100L; + long reduceBy = 50L; + SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); + List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); + + List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); + assertFalse(selectedTasks.isEmpty()); + assertTrue(tasksUsageMeetsThreshold(selectedTasks, reduceBy)); + } + + public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsLesserThanZero() { + TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); + long threshold = 100L; + long reduceBy = -50L; + SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); + List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); + + try { + testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); + } catch (Exception e) { + assertTrue(e instanceof IllegalArgumentException); + assertEquals("reduceBy has to be greater than zero", e.getMessage()); + } + } + + public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsEqualToZero() { + TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); + long threshold = 100L; + long reduceBy = 0; + SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); + List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); + + List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); + assertTrue(selectedTasks.isEmpty()); + } + + private boolean tasksUsageMeetsThreshold(List selectedTasks, long threshold) { + long memory = 0; + for(Task task : selectedTasks) { + memory += task.getTotalResourceUtilization(ResourceStats.MEMORY); + if(memory > threshold) { + return true; + } + } + return false; + } +} diff --git a/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java new file mode 100644 index 0000000000000..265461652133c --- /dev/null +++ b/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java @@ -0,0 +1,119 @@ +/* + * 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.tracking; + +import org.junit.After; +import org.junit.Before; +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.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.SandboxTask; +import org.opensearch.search.sandboxing.tracker.SandboxResourceUsageTrackerService; +import org.opensearch.tasks.CancellableTask; +import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskManager; +import org.opensearch.tasks.TaskResourceTrackingService; +import org.opensearch.test.OpenSearchTestCase; +import org.opensearch.threadpool.TestThreadPool; +import org.opensearch.threadpool.ThreadPool; + +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class SandboxResourceUsageTrackerServiceTests extends OpenSearchTestCase { + TestThreadPool threadPool; + TaskManager taskManager; + TaskResourceTrackingService mockTaskResourceTrackingService; + SandboxResourceUsageTrackerService sandboxResourceUsageTrackerService; + @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); + } + + @After + public void cleanup() { + ThreadPool.terminate(threadPool, 5, TimeUnit.SECONDS); + } + + public void testConstructSandboxLevelViews_CreatesSandboxLevelUsageView_WhenTasksArePresent() { + List sandboxIds = List.of("sandbox1", "sandbox2", "sandbox3"); + + Map activeSearchShardTasks = createActiveSearchShardTasks(sandboxIds); + when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); + Map stringSandboxLevelResourceUsageViewMap = sandboxResourceUsageTrackerService.constructSandboxLevelUsageViews(); + + for(String sandboxId : sandboxIds) { + assertEquals(400, (long) stringSandboxLevelResourceUsageViewMap.get(sandboxId).getResourceUsageData().get(SandboxResourceType.fromString("JVM"))); + assertEquals(2, stringSandboxLevelResourceUsageViewMap.get(sandboxId).getActiveTasks().size()); + } + } + + public void testConstructSandboxLevelViews_CreatesSandboxLevelUsageView_WhenTasksAreNotPresent() { + Map stringSandboxLevelResourceUsageViewMap = sandboxResourceUsageTrackerService.constructSandboxLevelUsageViews(); + assertTrue(stringSandboxLevelResourceUsageViewMap.isEmpty()); + } + + public void testConstructSandboxLevelUsageViews_WithTasksHavingDifferentResourceUsage() { + Map activeSearchShardTasks = new HashMap<>(); + activeSearchShardTasks.put(1L, createMockTask(SearchShardTask.class, 100, 200, "sandbox1")); + activeSearchShardTasks.put(2L, createMockTask(SearchShardTask.class, 200, 400, "sandbox1")); + when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); + + Map sandboxViews = sandboxResourceUsageTrackerService.constructSandboxLevelUsageViews(); + + assertEquals(600, (long) sandboxViews.get("sandbox1").getResourceUsageData().get(SandboxResourceType.fromString("JVM"))); + assertEquals(2, sandboxViews.get("sandbox1").getActiveTasks().size()); + } + + private Map createActiveSearchShardTasks(List sandboxIds) { + Map activeSearchShardTasks = new HashMap<>(); + long task_id = 0; + for (String sandboxId : sandboxIds) { + for (int i = 0; i < 2; i++) { + activeSearchShardTasks.put(++task_id, createMockTask(SearchShardTask.class, 100, 200, sandboxId)); + } + } + return activeSearchShardTasks; + } + + + private T createMockTask(Class type, long cpuUsage, long heapUsage, String sandboxId) { + T task = mock(type); + if (task instanceof SearchTask || task instanceof SearchShardTask) { + when(((SandboxTask) task).getSandboxId()).thenReturn(sandboxId); + } + when(task.getTotalResourceStats()).thenReturn(new TaskResourceUsage(cpuUsage, heapUsage)); + when(task.getStartTimeNanos()).thenReturn((long) 0); + + AtomicBoolean isCancelled = new AtomicBoolean(false); + doAnswer(invocation -> { + isCancelled.set(true); + return null; + }).when(task).cancel(anyString()); + doAnswer(invocation -> isCancelled.get()).when(task).isCancelled(); + + return task; + } +} From 546837d1d86cc3ee90b8e1bce944432bf58aadbe Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 10:16:05 -0700 Subject: [PATCH 02/66] Fix Failing Tests Signed-off-by: Kiran Prakash --- .../sandboxing/resourcetype/CpuTimeResourceType.java | 9 +++++++++ .../sandboxing/resourcetype/JvmMemoryResourceType.java | 9 +++++++++ .../SandboxCancellationStrategyTestHelpers.java | 8 -------- 3 files changed, 18 insertions(+), 8 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java index 6c4dd7fbaa320..e7aa8997b0450 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java @@ -15,4 +15,13 @@ public class CpuTimeResourceType extends SandboxResourceType { public long getResourceUsage(Task task) { return task.getTotalResourceStats().getCpuTimeInNanos(); } + @Override + public boolean equals(Object obj) { + return obj instanceof CpuTimeResourceType; + } + + @Override + public int hashCode() { + return "CPU".hashCode(); + } } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java index 91f6e73364b5f..be57f99559bfe 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java @@ -15,4 +15,13 @@ public class JvmMemoryResourceType extends SandboxResourceType { public long getResourceUsage(Task task) { return task.getTotalResourceStats().getMemoryInBytes(); } + @Override + public boolean equals(Object obj) { + return obj instanceof JvmMemoryResourceType; + } + + @Override + public int hashCode() { + return "JVM".hashCode(); + } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java index d4a5fa0c6d3f1..3d950bf853cad 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java @@ -13,20 +13,12 @@ import org.opensearch.core.tasks.resourcetracker.ResourceStats; import org.opensearch.core.tasks.resourcetracker.ResourceStatsType; import org.opensearch.core.tasks.resourcetracker.ResourceUsageMetric; -import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; -import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; -import org.opensearch.test.OpenSearchTestCase; import java.util.ArrayList; import java.util.Collections; import java.util.List; -import java.util.concurrent.atomic.AtomicBoolean; -import static org.mockito.ArgumentMatchers.anyString; -import static org.mockito.Mockito.doAnswer; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; import static org.opensearch.test.OpenSearchTestCase.randomLong; import static org.opensearch.test.OpenSearchTestCase.randomLongBetween; From 789cc38fdfc9d81af1e83c13092e630746de9eb6 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 10:18:23 -0700 Subject: [PATCH 03/66] spotless Apply Signed-off-by: Kiran Prakash --- .../opensearch/cluster/metadata/Metadata.java | 10 +- .../opensearch/cluster/metadata/Sandbox.java | 106 ++++++------ .../SandboxLevelResourceUsageView.java | 72 ++++---- .../AbstractTaskCancellation.java | 10 +- .../AbstractTaskSelectionStrategy.java | 50 +++--- .../cancellation/DefaultTaskCancellation.java | 78 ++++----- .../LongestRunningTaskFirstStrategy.java | 8 +- .../ShortestRunningTaskFirstStrategy.java | 10 +- .../resourcetype/CpuTimeResourceType.java | 25 +-- .../resourcetype/JvmMemoryResourceType.java | 25 +-- .../resourcetype/SandboxResourceType.java | 18 +- .../sandboxing/resourcetype/package-info.java | 2 +- .../SandboxResourceUsageTrackerService.java | 16 +- ...RunningTaskFirstStrategyStrategyTests.java | 22 +-- ...andboxCancellationStrategyTestHelpers.java | 60 ++++--- ...RunningTaskFirstStrategyStrategyTests.java | 22 +-- .../TaskSelectionStrategyTests.java | 90 +++++----- ...ndboxResourceUsageTrackerServiceTests.java | 156 +++++++++--------- 18 files changed, 389 insertions(+), 391 deletions(-) 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 feabdd095da04..a5f2c8d97a3be 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -127,12 +127,12 @@ public boolean isSegmentReplicationEnabled(String indexName) { .orElse(false); } - public Map sandboxes() { - // stub - return Collections.emptyMap(); - } + public Map sandboxes() { + // stub + return Collections.emptyMap(); + } - /** + /** * Context of the XContent. * * @opensearch.api diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java index c868be280ea4b..8db96dffe4ee0 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java @@ -16,72 +16,72 @@ @ExperimentalApi public class Sandbox { - //TODO Kaushal should have implemented hashcode and equals - private SandboxMode mode; + // TODO Kaushal should have implemented hashcode and equals + private SandboxMode mode; - public SandboxMode getMode() { - return mode; - } - - public ResourceLimit getResourceLimitFor(SandboxResourceType resourceType) { - return null; - } - - public String getName() { - return ""; - } - - public String getId() { - return ""; - } + public SandboxMode getMode() { + return mode; + } - public List getResourceLimits() { - return Collections.emptyList(); - } + public ResourceLimit getResourceLimitFor(SandboxResourceType resourceType) { + return null; + } - @ExperimentalApi - public class ResourceLimit { - public Long getThresholdInLong() { - return 0L; + public String getName() { + return ""; } - public SandboxResourceType getResourceType() { - return null; + public String getId() { + return ""; } - public Long getThreshold() { - return 0L; + public List getResourceLimits() { + return Collections.emptyList(); } - } - @ExperimentalApi - public enum SandboxMode { - SOFT("soft"), - ENFORCED("enforced"), - MONITOR("monitor"); + @ExperimentalApi + public class ResourceLimit { + public Long getThresholdInLong() { + return 0L; + } - private final String name; + public SandboxResourceType getResourceType() { + return null; + } - SandboxMode(String mode) { - this.name = mode; + public Long getThreshold() { + return 0L; + } } - public String getName() { - return name; - } + @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); + } + } - 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/sandboxing/SandboxLevelResourceUsageView.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java index aa181bb85c3c0..0a25cdf900d9a 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java @@ -21,40 +21,40 @@ @ExperimentalApi public class SandboxLevelResourceUsageView { - private final String sandboxId; - private final Map resourceUsage; - private final List activeTasks; - - public SandboxLevelResourceUsageView(String sandboxId) { - this.sandboxId = sandboxId; - this.resourceUsage = new HashMap<>(); - this.activeTasks = new ArrayList<>(); - } - - public SandboxLevelResourceUsageView(String sandboxId, Map resourceUsage, List activeTasks) { - this.sandboxId = sandboxId; - this.resourceUsage = resourceUsage; - this.activeTasks = activeTasks; - } - - public Map getResourceUsageData() { - return resourceUsage; - } - - public List getActiveTasks() { - return activeTasks; - } - - @Override - 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); - } - - @Override - public int hashCode() { - return Objects.hashCode(sandboxId); - } + private final String sandboxId; + private final Map resourceUsage; + private final List activeTasks; + + public SandboxLevelResourceUsageView(String sandboxId) { + this.sandboxId = sandboxId; + this.resourceUsage = new HashMap<>(); + this.activeTasks = new ArrayList<>(); + } + + public SandboxLevelResourceUsageView(String sandboxId, Map resourceUsage, List activeTasks) { + this.sandboxId = sandboxId; + this.resourceUsage = resourceUsage; + this.activeTasks = activeTasks; + } + + public Map getResourceUsageData() { + return resourceUsage; + } + + public List getActiveTasks() { + return activeTasks; + } + + @Override + 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); + } + + @Override + public int hashCode() { + return Objects.hashCode(sandboxId); + } } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java index 9bf96bf63578c..447ba66ebddde 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java @@ -93,12 +93,10 @@ private boolean shouldCancelTasks(Sandbox sandbox, SandboxResourceType resourceT private List getTaskCancellations(Sandbox sandbox, SandboxResourceType resourceType) { return taskSelectionStrategy.selectTasksForCancellation( - getAllTasksInSandbox(sandbox.getId()), - getReduceBy(sandbox, resourceType), - resourceType) - .stream() - .map(task -> createTaskCancellation((CancellableTask) task)) - .collect(Collectors.toList()); + getAllTasksInSandbox(sandbox.getId()), + getReduceBy(sandbox, resourceType), + resourceType + ).stream().map(task -> createTaskCancellation((CancellableTask) task)).collect(Collectors.toList()); } private long getReduceBy(Sandbox sandbox, SandboxResourceType resourceType) { diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java index de0cc19090a4d..a9943c30ea5b0 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java @@ -19,31 +19,29 @@ public abstract class AbstractTaskSelectionStrategy implements TaskSelectionStrategy { - public abstract Comparator sortingCondition(); - - @Override - public List selectTasksForCancellation(List tasks, long limit, SandboxResourceType resourceType) { - if (limit < 0) { - throw new IllegalArgumentException("reduceBy has to be greater than zero"); - } - if(limit == 0) { - return Collections.emptyList(); - } - - List sortedTasks = tasks.stream() - .sorted(sortingCondition()) - .collect(Collectors.toList()); - - List selectedTasks = new ArrayList<>(); - long accumulated = 0; - - for (Task task : sortedTasks) { - selectedTasks.add(task); - accumulated += resourceType.getResourceUsage(task); - if (accumulated >= limit) { - break; - } + public abstract Comparator sortingCondition(); + + @Override + public List selectTasksForCancellation(List tasks, long limit, SandboxResourceType resourceType) { + if (limit < 0) { + throw new IllegalArgumentException("reduceBy has to be greater than zero"); + } + if (limit == 0) { + return Collections.emptyList(); + } + + List sortedTasks = tasks.stream().sorted(sortingCondition()).collect(Collectors.toList()); + + List selectedTasks = new ArrayList<>(); + long accumulated = 0; + + for (Task task : sortedTasks) { + selectedTasks.add(task); + accumulated += resourceType.getResourceUsage(task); + if (accumulated >= limit) { + break; + } + } + return selectedTasks; } - return selectedTasks; - } } 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 index f5d21f8af5e4e..abf3933865c6c 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java @@ -9,8 +9,8 @@ package org.opensearch.search.sandboxing.cancellation; import org.opensearch.cluster.metadata.Sandbox; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; +import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; import java.util.ArrayList; import java.util.List; @@ -18,46 +18,46 @@ import java.util.Set; public class DefaultTaskCancellation extends AbstractTaskCancellation { - public DefaultTaskCancellation( - TaskSelectionStrategy cancellationStrategy, - Map sandboxLevelViews, - Set activeSandboxes - ) { - super(cancellationStrategy, sandboxLevelViews, activeSandboxes); - } - - /** - * // TODO - * This should cover 3 scenarios - * - if node not in duress - * - pick sandboxes in enforced mode only - * - if node in duress - * - pick sandboxes in enforced mode - * - tasks running in deleted sandboxes with tasks running - * - pick sandboxes in enforced mode - */ - public List getSandboxesToCancelFrom() { - final List sandboxesToCancelFrom = new ArrayList<>(); - - for (Sandbox sandbox : this.activeSandboxes) { - Map currentResourceUsage = getResourceUsage(sandbox.getId()); - - for (Sandbox.ResourceLimit resourceLimit : sandbox.getResourceLimits()) { - if (isBreachingThreshold(currentResourceUsage, resourceLimit)) { - sandboxesToCancelFrom.add(sandbox); - break; - } - } + public DefaultTaskCancellation( + TaskSelectionStrategy cancellationStrategy, + Map sandboxLevelViews, + Set activeSandboxes + ) { + super(cancellationStrategy, sandboxLevelViews, activeSandboxes); } - return sandboxesToCancelFrom; - } + /** + * // TODO + * This should cover 3 scenarios + * - if node not in duress + * - pick sandboxes in enforced mode only + * - if node in duress + * - pick sandboxes in enforced mode + * - tasks running in deleted sandboxes with tasks running + * - pick sandboxes in enforced mode + */ + public List getSandboxesToCancelFrom() { + final List sandboxesToCancelFrom = new ArrayList<>(); + + for (Sandbox sandbox : this.activeSandboxes) { + Map currentResourceUsage = getResourceUsage(sandbox.getId()); + + for (Sandbox.ResourceLimit resourceLimit : sandbox.getResourceLimits()) { + if (isBreachingThreshold(currentResourceUsage, resourceLimit)) { + sandboxesToCancelFrom.add(sandbox); + break; + } + } + } + + return sandboxesToCancelFrom; + } - private boolean isBreachingThreshold(Map currentResourceUsage, Sandbox.ResourceLimit resourceLimit) { - return currentResourceUsage.get(resourceLimit.getResourceType()) > resourceLimit.getThreshold(); - } + private boolean isBreachingThreshold(Map currentResourceUsage, Sandbox.ResourceLimit resourceLimit) { + return currentResourceUsage.get(resourceLimit.getResourceType()) > resourceLimit.getThreshold(); + } - private Map getResourceUsage(String sandboxId) { - return sandboxLevelViews.get(sandboxId).getResourceUsageData(); - } + private Map getResourceUsage(String sandboxId) { + return sandboxLevelViews.get(sandboxId).getResourceUsageData(); + } } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java index d97f1b26b5362..b4585933223bb 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java @@ -14,8 +14,8 @@ public class LongestRunningTaskFirstStrategy extends AbstractTaskSelectionStrategy { - @Override - public Comparator sortingCondition() { - return Comparator.comparingLong(Task::getStartTime).reversed(); - } + @Override + public Comparator sortingCondition() { + return Comparator.comparingLong(Task::getStartTime).reversed(); + } } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java index f4e1c7f4fe9bc..0036247b4c895 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java @@ -14,8 +14,8 @@ public class ShortestRunningTaskFirstStrategy extends AbstractTaskSelectionStrategy { - @Override - public Comparator sortingCondition() { - return Comparator.comparingLong(Task::getStartTime); - } -} \ No newline at end of file + @Override + public Comparator sortingCondition() { + return Comparator.comparingLong(Task::getStartTime); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java index e7aa8997b0450..9a854f6937a22 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java @@ -11,17 +11,18 @@ import org.opensearch.tasks.Task; public class CpuTimeResourceType extends SandboxResourceType { - @Override - public long getResourceUsage(Task task) { - return task.getTotalResourceStats().getCpuTimeInNanos(); - } - @Override - public boolean equals(Object obj) { - return obj instanceof CpuTimeResourceType; - } + @Override + public long getResourceUsage(Task task) { + return task.getTotalResourceStats().getCpuTimeInNanos(); + } - @Override - public int hashCode() { - return "CPU".hashCode(); - } + @Override + public boolean equals(Object obj) { + return obj instanceof CpuTimeResourceType; + } + + @Override + public int hashCode() { + return "CPU".hashCode(); + } } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java index be57f99559bfe..666dae6e49b73 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java @@ -11,17 +11,18 @@ import org.opensearch.tasks.Task; public class JvmMemoryResourceType extends SandboxResourceType { - @Override - public long getResourceUsage(Task task) { - return task.getTotalResourceStats().getMemoryInBytes(); - } - @Override - public boolean equals(Object obj) { - return obj instanceof JvmMemoryResourceType; - } + @Override + public long getResourceUsage(Task task) { + return task.getTotalResourceStats().getMemoryInBytes(); + } - @Override - public int hashCode() { - return "JVM".hashCode(); - } + @Override + public boolean equals(Object obj) { + return obj instanceof JvmMemoryResourceType; + } + + @Override + public int hashCode() { + return "JVM".hashCode(); + } } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java index 9b849718e6a11..f6414b6652121 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java @@ -13,15 +13,15 @@ @ExperimentalApi public abstract class SandboxResourceType { - public abstract long getResourceUsage(Task task); + public abstract long getResourceUsage(Task task); - public static SandboxResourceType fromString(String type) { - if (type.equalsIgnoreCase("JVM")) { - return new JvmMemoryResourceType(); - } else if (type.equalsIgnoreCase("CPU")) { - return new CpuTimeResourceType(); - } else { - throw new IllegalArgumentException("Unsupported resource type: " + type); + public static SandboxResourceType fromString(String type) { + if (type.equalsIgnoreCase("JVM")) { + return new JvmMemoryResourceType(); + } else if (type.equalsIgnoreCase("CPU")) { + return new CpuTimeResourceType(); + } 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 index 97d449aa49912..1f7f64b1ba534 100644 --- 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 @@ -6,4 +6,4 @@ * compatible open source license. */ -package org.opensearch.search.sandboxing.resourcetype; \ No newline at end of file +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 index 84db827944c9f..22905e2a3f6e4 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java @@ -10,8 +10,8 @@ import org.opensearch.common.inject.Inject; import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; import org.opensearch.search.sandboxing.SandboxTask; +import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; import org.opensearch.tasks.TaskResourceTrackingService; @@ -24,7 +24,7 @@ /** * This class tracks requests per Sandbox */ -//@ExperimentalApi +// @ExperimentalApi public class SandboxResourceUsageTrackerService implements SandboxUsageTracker, TaskManager.TaskEventListeners { public static final List TRACKED_RESOURCES = List.of(SandboxResourceType.fromString("JVM")); @@ -59,12 +59,12 @@ public Map constructSandboxLevelUsageView Map> tasksBySandbox = getTasksGroupedBySandbox(); Map> sandboxResourceUsage = getResourceUsageOfSandboxes(tasksBySandbox); - for(String sandboxId : tasksBySandbox.keySet()) { + for (String sandboxId : tasksBySandbox.keySet()) { SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView( sandboxId, sandboxResourceUsage.get(sandboxId), tasksBySandbox.get(sandboxId) - ); + ); sandboxViews.put(sandboxId, sandboxLevelResourceUsageView); } return sandboxViews; @@ -81,10 +81,7 @@ private Map> getTasksGroupedBySandbox() { .stream() .filter(SandboxTask.class::isInstance) .map(SandboxTask.class::cast) - .collect(Collectors.groupingBy( - SandboxTask::getSandboxId, - Collectors.mapping(task -> (Task) task, Collectors.toList()) - )); + .collect(Collectors.groupingBy(SandboxTask::getSandboxId, Collectors.mapping(task -> (Task) task, Collectors.toList()))); } /** @@ -122,6 +119,5 @@ private Map> getResourceUsageOfSandboxes( * @param task The completed task */ @Override - public void onTaskCompleted(Task task) { - } + public void onTaskCompleted(Task task) {} } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java index 8c19de92f1b9f..4c5c0aeb4fd88 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java @@ -18,17 +18,17 @@ import static org.mockito.Mockito.when; public class LongestRunningTaskFirstStrategyStrategyTests extends OpenSearchTestCase { - public void testSortingCondition() { - Task task1 = mock(Task.class); - Task task2 = mock(Task.class); - Task task3 = mock(Task.class); - when(task1.getStartTime()).thenReturn(100L); - when(task2.getStartTime()).thenReturn(200L); - when(task3.getStartTime()).thenReturn(300L); + public void testSortingCondition() { + Task task1 = mock(Task.class); + Task task2 = mock(Task.class); + Task task3 = mock(Task.class); + when(task1.getStartTime()).thenReturn(100L); + when(task2.getStartTime()).thenReturn(200L); + when(task3.getStartTime()).thenReturn(300L); - List tasks = Arrays.asList(task1, task3, task2); - tasks.sort(new LongestRunningTaskFirstStrategy().sortingCondition()); + List tasks = Arrays.asList(task1, task3, task2); + tasks.sort(new LongestRunningTaskFirstStrategy().sortingCondition()); - assertEquals(Arrays.asList(task3, task2, task1), tasks); - } + assertEquals(Arrays.asList(task3, task2, task1), tasks); + } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java index 3d950bf853cad..5308b526b820c 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java @@ -24,42 +24,40 @@ public class SandboxCancellationStrategyTestHelpers { - public static List getListOfTasks(long totalMemory) { - List tasks = new ArrayList<>(); + public static List getListOfTasks(long totalMemory) { + List tasks = new ArrayList<>(); - while (totalMemory > 0) { - long id = randomLong(); - final Task task = getRandomTask(id); - long initial_memory = randomLongBetween(1, 100); + while (totalMemory > 0) { + long id = randomLong(); + final Task task = getRandomTask(id); + long initial_memory = randomLongBetween(1, 100); - ResourceUsageMetric[] initialTaskResourceMetrics = new ResourceUsageMetric[]{ - new ResourceUsageMetric(ResourceStats.MEMORY, initial_memory) - }; - task.startThreadResourceTracking(id, ResourceStatsType.WORKER_STATS, initialTaskResourceMetrics); + ResourceUsageMetric[] initialTaskResourceMetrics = new ResourceUsageMetric[] { + new ResourceUsageMetric(ResourceStats.MEMORY, initial_memory) }; + task.startThreadResourceTracking(id, ResourceStatsType.WORKER_STATS, initialTaskResourceMetrics); - long memory = initial_memory + randomLongBetween(1, 10000); + long memory = initial_memory + randomLongBetween(1, 10000); - totalMemory -= memory - initial_memory; + totalMemory -= memory - initial_memory; - ResourceUsageMetric[] taskResourceMetrics = new ResourceUsageMetric[]{ - new ResourceUsageMetric(ResourceStats.MEMORY, memory), - }; - task.updateThreadResourceStats(id, ResourceStatsType.WORKER_STATS, taskResourceMetrics); - task.stopThreadResourceTracking(id, ResourceStatsType.WORKER_STATS); - tasks.add(task); - } + ResourceUsageMetric[] taskResourceMetrics = new ResourceUsageMetric[] { + new ResourceUsageMetric(ResourceStats.MEMORY, memory), }; + task.updateThreadResourceStats(id, ResourceStatsType.WORKER_STATS, taskResourceMetrics); + task.stopThreadResourceTracking(id, ResourceStatsType.WORKER_STATS); + tasks.add(task); + } - return tasks; - } + return tasks; + } - public static Task getRandomTask(long id) { - return new Task( - id, - "transport", - SearchAction.NAME, - "test description", - new TaskId(randomLong() + ":" + randomLong()), - Collections.emptyMap() - ); - } + public static Task getRandomTask(long id) { + return new Task( + id, + "transport", + SearchAction.NAME, + "test description", + new TaskId(randomLong() + ":" + randomLong()), + Collections.emptyMap() + ); + } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java index 2422d41ccd2b5..cdd3053e58373 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java @@ -18,17 +18,17 @@ import static org.mockito.Mockito.when; public class ShortestRunningTaskFirstStrategyStrategyTests extends OpenSearchTestCase { - public void testSortingCondition() { - Task task1 = mock(Task.class); - Task task2 = mock(Task.class); - Task task3 = mock(Task.class); - when(task1.getStartTime()).thenReturn(100L); - when(task2.getStartTime()).thenReturn(200L); - when(task3.getStartTime()).thenReturn(300L); + public void testSortingCondition() { + Task task1 = mock(Task.class); + Task task2 = mock(Task.class); + Task task3 = mock(Task.class); + when(task1.getStartTime()).thenReturn(100L); + when(task2.getStartTime()).thenReturn(200L); + when(task3.getStartTime()).thenReturn(300L); - List tasks = Arrays.asList(task1, task3, task2); - tasks.sort(new ShortestRunningTaskFirstStrategy().sortingCondition()); + List tasks = Arrays.asList(task1, task3, task2); + tasks.sort(new ShortestRunningTaskFirstStrategy().sortingCondition()); - assertEquals(Arrays.asList(task1, task2, task3), tasks); - } + assertEquals(Arrays.asList(task1, task2, task3), tasks); + } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java index 39634d9e026b0..c513cb74a54fd 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java @@ -18,59 +18,59 @@ public class TaskSelectionStrategyTests extends OpenSearchTestCase { - private class TestTaskSelectionStrategy extends AbstractTaskSelectionStrategy { - @Override - public Comparator sortingCondition() { - return Comparator.comparingLong(Task::getStartTime).reversed(); + private class TestTaskSelectionStrategy extends AbstractTaskSelectionStrategy { + @Override + public Comparator sortingCondition() { + return Comparator.comparingLong(Task::getStartTime).reversed(); + } } - } - public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsGreaterThanZero() { - TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); - long threshold = 100L; - long reduceBy = 50L; - SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); - List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); + public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsGreaterThanZero() { + TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); + long threshold = 100L; + long reduceBy = 50L; + SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); + List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); - List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); - assertFalse(selectedTasks.isEmpty()); - assertTrue(tasksUsageMeetsThreshold(selectedTasks, reduceBy)); - } + List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); + assertFalse(selectedTasks.isEmpty()); + assertTrue(tasksUsageMeetsThreshold(selectedTasks, reduceBy)); + } - public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsLesserThanZero() { - TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); - long threshold = 100L; - long reduceBy = -50L; - SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); - List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); + public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsLesserThanZero() { + TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); + long threshold = 100L; + long reduceBy = -50L; + SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); + List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); - try { - testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); - } catch (Exception e) { - assertTrue(e instanceof IllegalArgumentException); - assertEquals("reduceBy has to be greater than zero", e.getMessage()); + try { + testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); + } catch (Exception e) { + assertTrue(e instanceof IllegalArgumentException); + assertEquals("reduceBy has to be greater than zero", e.getMessage()); + } } - } - public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsEqualToZero() { - TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); - long threshold = 100L; - long reduceBy = 0; - SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); - List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); + public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsEqualToZero() { + TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); + long threshold = 100L; + long reduceBy = 0; + SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); + List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); - List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); - assertTrue(selectedTasks.isEmpty()); - } + List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); + assertTrue(selectedTasks.isEmpty()); + } - private boolean tasksUsageMeetsThreshold(List selectedTasks, long threshold) { - long memory = 0; - for(Task task : selectedTasks) { - memory += task.getTotalResourceUtilization(ResourceStats.MEMORY); - if(memory > threshold) { - return true; - } + private boolean tasksUsageMeetsThreshold(List selectedTasks, long threshold) { + long memory = 0; + for (Task task : selectedTasks) { + memory += task.getTotalResourceUtilization(ResourceStats.MEMORY); + if (memory > threshold) { + return true; + } + } + return false; } - return false; - } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java index 265461652133c..1120122f33f01 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java @@ -8,15 +8,13 @@ package org.opensearch.search.sandboxing.tracking; -import org.junit.After; -import org.junit.Before; 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.resourcetype.SandboxResourceType; import org.opensearch.search.sandboxing.SandboxTask; +import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; import org.opensearch.search.sandboxing.tracker.SandboxResourceUsageTrackerService; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; @@ -25,6 +23,8 @@ import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; +import org.junit.After; +import org.junit.Before; import java.util.Collections; import java.util.HashMap; @@ -39,81 +39,87 @@ import static org.mockito.Mockito.when; public class SandboxResourceUsageTrackerServiceTests extends OpenSearchTestCase { - TestThreadPool threadPool; - TaskManager taskManager; - TaskResourceTrackingService mockTaskResourceTrackingService; - SandboxResourceUsageTrackerService sandboxResourceUsageTrackerService; - @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); - } - - @After - public void cleanup() { - ThreadPool.terminate(threadPool, 5, TimeUnit.SECONDS); - } - - public void testConstructSandboxLevelViews_CreatesSandboxLevelUsageView_WhenTasksArePresent() { - List sandboxIds = List.of("sandbox1", "sandbox2", "sandbox3"); - - Map activeSearchShardTasks = createActiveSearchShardTasks(sandboxIds); - when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); - Map stringSandboxLevelResourceUsageViewMap = sandboxResourceUsageTrackerService.constructSandboxLevelUsageViews(); - - for(String sandboxId : sandboxIds) { - assertEquals(400, (long) stringSandboxLevelResourceUsageViewMap.get(sandboxId).getResourceUsageData().get(SandboxResourceType.fromString("JVM"))); - assertEquals(2, stringSandboxLevelResourceUsageViewMap.get(sandboxId).getActiveTasks().size()); + TestThreadPool threadPool; + TaskManager taskManager; + TaskResourceTrackingService mockTaskResourceTrackingService; + SandboxResourceUsageTrackerService sandboxResourceUsageTrackerService; + + @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); + } + + @After + public void cleanup() { + ThreadPool.terminate(threadPool, 5, TimeUnit.SECONDS); } - } - - public void testConstructSandboxLevelViews_CreatesSandboxLevelUsageView_WhenTasksAreNotPresent() { - Map stringSandboxLevelResourceUsageViewMap = sandboxResourceUsageTrackerService.constructSandboxLevelUsageViews(); - assertTrue(stringSandboxLevelResourceUsageViewMap.isEmpty()); - } - - public void testConstructSandboxLevelUsageViews_WithTasksHavingDifferentResourceUsage() { - Map activeSearchShardTasks = new HashMap<>(); - activeSearchShardTasks.put(1L, createMockTask(SearchShardTask.class, 100, 200, "sandbox1")); - activeSearchShardTasks.put(2L, createMockTask(SearchShardTask.class, 200, 400, "sandbox1")); - when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); - - Map sandboxViews = sandboxResourceUsageTrackerService.constructSandboxLevelUsageViews(); - - assertEquals(600, (long) sandboxViews.get("sandbox1").getResourceUsageData().get(SandboxResourceType.fromString("JVM"))); - assertEquals(2, sandboxViews.get("sandbox1").getActiveTasks().size()); - } - - private Map createActiveSearchShardTasks(List sandboxIds) { - Map activeSearchShardTasks = new HashMap<>(); - long task_id = 0; - for (String sandboxId : sandboxIds) { - for (int i = 0; i < 2; i++) { - activeSearchShardTasks.put(++task_id, createMockTask(SearchShardTask.class, 100, 200, sandboxId)); - } + + public void testConstructSandboxLevelViews_CreatesSandboxLevelUsageView_WhenTasksArePresent() { + List sandboxIds = List.of("sandbox1", "sandbox2", "sandbox3"); + + Map activeSearchShardTasks = createActiveSearchShardTasks(sandboxIds); + when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); + Map stringSandboxLevelResourceUsageViewMap = sandboxResourceUsageTrackerService + .constructSandboxLevelUsageViews(); + + for (String sandboxId : sandboxIds) { + assertEquals( + 400, + (long) stringSandboxLevelResourceUsageViewMap.get(sandboxId) + .getResourceUsageData() + .get(SandboxResourceType.fromString("JVM")) + ); + assertEquals(2, stringSandboxLevelResourceUsageViewMap.get(sandboxId).getActiveTasks().size()); + } + } + + public void testConstructSandboxLevelViews_CreatesSandboxLevelUsageView_WhenTasksAreNotPresent() { + Map stringSandboxLevelResourceUsageViewMap = sandboxResourceUsageTrackerService + .constructSandboxLevelUsageViews(); + assertTrue(stringSandboxLevelResourceUsageViewMap.isEmpty()); + } + + public void testConstructSandboxLevelUsageViews_WithTasksHavingDifferentResourceUsage() { + Map activeSearchShardTasks = new HashMap<>(); + activeSearchShardTasks.put(1L, createMockTask(SearchShardTask.class, 100, 200, "sandbox1")); + activeSearchShardTasks.put(2L, createMockTask(SearchShardTask.class, 200, 400, "sandbox1")); + when(mockTaskResourceTrackingService.getResourceAwareTasks()).thenReturn(activeSearchShardTasks); + + Map sandboxViews = sandboxResourceUsageTrackerService.constructSandboxLevelUsageViews(); + + assertEquals(600, (long) sandboxViews.get("sandbox1").getResourceUsageData().get(SandboxResourceType.fromString("JVM"))); + assertEquals(2, sandboxViews.get("sandbox1").getActiveTasks().size()); } - return activeSearchShardTasks; - } + private Map createActiveSearchShardTasks(List sandboxIds) { + Map activeSearchShardTasks = new HashMap<>(); + long task_id = 0; + for (String sandboxId : sandboxIds) { + for (int i = 0; i < 2; i++) { + activeSearchShardTasks.put(++task_id, createMockTask(SearchShardTask.class, 100, 200, sandboxId)); + } + } + return activeSearchShardTasks; + } - private T createMockTask(Class type, long cpuUsage, long heapUsage, String sandboxId) { - T task = mock(type); - if (task instanceof SearchTask || task instanceof SearchShardTask) { - when(((SandboxTask) task).getSandboxId()).thenReturn(sandboxId); + private T createMockTask(Class type, long cpuUsage, long heapUsage, String sandboxId) { + T task = mock(type); + if (task instanceof SearchTask || task instanceof SearchShardTask) { + when(((SandboxTask) task).getSandboxId()).thenReturn(sandboxId); + } + when(task.getTotalResourceStats()).thenReturn(new TaskResourceUsage(cpuUsage, heapUsage)); + when(task.getStartTimeNanos()).thenReturn((long) 0); + + AtomicBoolean isCancelled = new AtomicBoolean(false); + doAnswer(invocation -> { + isCancelled.set(true); + return null; + }).when(task).cancel(anyString()); + doAnswer(invocation -> isCancelled.get()).when(task).isCancelled(); + + return task; } - when(task.getTotalResourceStats()).thenReturn(new TaskResourceUsage(cpuUsage, heapUsage)); - when(task.getStartTimeNanos()).thenReturn((long) 0); - - AtomicBoolean isCancelled = new AtomicBoolean(false); - doAnswer(invocation -> { - isCancelled.set(true); - return null; - }).when(task).cancel(anyString()); - doAnswer(invocation -> isCancelled.get()).when(task).isCancelled(); - - return task; - } } From 60dfab3fc18a83ae352807068e2c0389c5efb3ce Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 10:30:25 -0700 Subject: [PATCH 04/66] Update SandboxService.java Signed-off-by: Kiran Prakash --- .../search/sandboxing/SandboxService.java | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java index 662258bb7fd03..78ce7510c0f5d 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java @@ -14,6 +14,7 @@ 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; @@ -24,6 +25,7 @@ import java.util.HashSet; import java.util.Map; import java.util.Set; +import java.util.concurrent.TimeUnit; /** * Main service which will run periodically to track and cancel resource constraint violating tasks in sandboxes @@ -40,8 +42,8 @@ public class SandboxService extends AbstractLifecycleComponent { * Guice managed constructor * * @param sandboxUsageTracker - * @param sandboxPruner - * @param sandboxServiceSettings +// * @param sandboxPruner +// * @param sandboxServiceSettings * @param threadPool */ @Inject @@ -52,7 +54,7 @@ public SandboxService( ) { this.sandboxUsageTracker = sandboxUsageTracker; this.sandboxServiceSettings = sandboxServiceSettings; - this.sandboxPruner = sandboxPruner; +// this.sandboxPruner = sandboxPruner; this.clusterService = clusterService; this.threadPool = threadPool; } @@ -82,12 +84,14 @@ private Set getActiveSandboxes() { @Override protected void doStart() { scheduledFuture = threadPool.scheduleWithFixedDelay(() -> { - try { - doRun(); - } catch (Exception e) { - logger.debug("Exception occurred in Query Sandbox service", e); - } - }, sandboxServiceSettings.getRunIntervalMillis(), ThreadPool.Names.GENERIC); + try { + doRun(); + } catch (Exception e) { + logger.debug("Exception occurred in Query Sandbox service", e); + } + }, + new TimeValue(1, TimeUnit.SECONDS), // TODO get this from SandboxServiceSettings + ThreadPool.Names.GENERIC); } @Override @@ -98,6 +102,5 @@ protected void doStop() { } @Override - protected void doClose() throws IOException { - } + protected void doClose() throws IOException {} } From 13882b86121766ba8f938b5a4a9f284695db4919 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 10:36:38 -0700 Subject: [PATCH 05/66] Update SandboxService.java Signed-off-by: Kiran Prakash --- .../org/opensearch/search/sandboxing/SandboxService.java | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java index 78ce7510c0f5d..3db9205e5b07e 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java @@ -34,7 +34,9 @@ public class SandboxService extends AbstractLifecycleComponent { private static final Logger logger = LogManager.getLogger(SandboxService.class); private final SandboxUsageTracker sandboxUsageTracker; +// private final SandboxPruner sandboxPruner; private volatile Scheduler.Cancellable scheduledFuture; +// private final SandboxServiceSettings sandboxServiceSettings; private final ThreadPool threadPool; private final ClusterService clusterService; @@ -49,11 +51,13 @@ public class SandboxService extends AbstractLifecycleComponent { @Inject public SandboxService( SandboxUsageTracker sandboxUsageTracker, +// SandboxServiceSettings sandboxServiceSettings, +// SandboxPruner sandboxPruner, ClusterService clusterService, ThreadPool threadPool ) { this.sandboxUsageTracker = sandboxUsageTracker; - this.sandboxServiceSettings = sandboxServiceSettings; +// this.sandboxServiceSettings = sandboxServiceSettings; // this.sandboxPruner = sandboxPruner; this.clusterService = clusterService; this.threadPool = threadPool; From 4fcfb79ebe0bab60be71ce747ff3d4674056e1a8 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 10:43:39 -0700 Subject: [PATCH 06/66] Update SandboxTask.java Signed-off-by: Kiran Prakash --- .../main/java/org/opensearch/search/sandboxing/SandboxTask.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java index 92aba3759f115..a674f27d71eff 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxTask.java @@ -9,7 +9,7 @@ package org.opensearch.search.sandboxing; /** - * This interface can be implemented by tasks which will be tracked and monitored using {@link org.opensearch.cluster.metadata.ResourceLimitGroup} + * This interface can be implemented by tasks which will be tracked and monitored using {@link org.opensearch.cluster.metadata.Sandbox} */ public interface SandboxTask { void setSandboxId(String sandboxId); From aec68f048bdef287df778a2421469cc20e432945 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 11:13:35 -0700 Subject: [PATCH 07/66] Add java docs Signed-off-by: Kiran Prakash --- .../SandboxLevelResourceUsageView.java | 14 +++++++++++++ .../AbstractTaskSelectionStrategy.java | 21 +++++++++++++++++++ .../cancellation/DefaultTaskCancellation.java | 20 ++++++++++++++++-- .../LongestRunningTaskFirstStrategy.java | 8 +++++++ .../ShortestRunningTaskFirstStrategy.java | 8 +++++++ .../resourcetype/CpuTimeResourceType.java | 9 ++++++++ .../resourcetype/JvmMemoryResourceType.java | 9 ++++++++ .../resourcetype/SandboxResourceType.java | 21 +++++++++++++++++++ 8 files changed, 108 insertions(+), 2 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java index 0a25cdf900d9a..c104170c2eee0 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java @@ -18,6 +18,10 @@ import java.util.Map; import java.util.Objects; +/** + * Represents the resource usage view at the sandbox level. + * This class holds the sandbox ID, the resource usage data, and the list of active tasks. + */ @ExperimentalApi public class SandboxLevelResourceUsageView { @@ -37,10 +41,20 @@ public SandboxLevelResourceUsageView(String sandboxId, Map getResourceUsageData() { return resourceUsage; } + /** + * Returns the list of active tasks. + * + * @return The list of active tasks + */ public List getActiveTasks() { return activeTasks; } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java index a9943c30ea5b0..5767975d3497e 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java @@ -17,10 +17,31 @@ import java.util.List; import java.util.stream.Collectors; +/** + * Represents an abstract task selection strategy. + * This class implements the TaskSelectionStrategy interface and provides a method to select tasks for cancellation based on a sorting condition. + * The specific sorting condition depends on the implementation. + */ public abstract class AbstractTaskSelectionStrategy implements TaskSelectionStrategy { + /** + * Returns a comparator that defines the sorting condition for tasks. + * The specific sorting condition depends on the implementation. + * + * @return The comparator + */ public abstract Comparator sortingCondition(); + /** + * Selects tasks for cancellation based on the provided limit and resource type. + * The tasks are sorted based on the sorting condition and then selected until the accumulated resource usage reaches the limit. + * + * @param tasks The list of tasks from which to select + * @param limit The limit on the accumulated resource usage + * @param resourceType The type of resource to consider + * @return The list of selected tasks + * @throws IllegalArgumentException If the limit is less than zero + */ @Override public List selectTasksForCancellation(List tasks, long limit, SandboxResourceType resourceType) { if (limit < 0) { 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 index abf3933865c6c..0e351ab2e90ea 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java @@ -17,6 +17,10 @@ import java.util.Map; import java.util.Set; +/** + * Represents the default task cancellation strategy. + * This class extends the AbstractTaskCancellation and provides a method to get the sandboxes from which tasks should be cancelled. + */ public class DefaultTaskCancellation extends AbstractTaskCancellation { public DefaultTaskCancellation( TaskSelectionStrategy cancellationStrategy, @@ -27,8 +31,7 @@ public DefaultTaskCancellation( } /** - * // TODO - * This should cover 3 scenarios + * // TODO create another Task cancellation implementation that does below * - if node not in duress * - pick sandboxes in enforced mode only * - if node in duress @@ -53,10 +56,23 @@ public List getSandboxesToCancelFrom() { return sandboxesToCancelFrom; } + /** + * 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) { return sandboxLevelViews.get(sandboxId).getResourceUsageData(); } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java index b4585933223bb..67f3ed18252d4 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java @@ -12,8 +12,16 @@ import java.util.Comparator; +/** + * Represents a task selection strategy that prioritizes the longest running tasks first. + */ public class LongestRunningTaskFirstStrategy extends AbstractTaskSelectionStrategy { + /** + * Returns a comparator that sorts tasks based on their start time in descending order. + * + * @return The comparator + */ @Override public Comparator sortingCondition() { return Comparator.comparingLong(Task::getStartTime).reversed(); diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java index 0036247b4c895..2f7acb43512d8 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java @@ -12,8 +12,16 @@ import java.util.Comparator; +/** + * Represents a task selection strategy that prioritizes the shortest running tasks first. + */ public class ShortestRunningTaskFirstStrategy extends AbstractTaskSelectionStrategy { + /** + * Returns a comparator that sorts tasks based on their start time in ascending order. + * + * @return The comparator + */ @Override public Comparator sortingCondition() { return Comparator.comparingLong(Task::getStartTime); diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java index 9a854f6937a22..675bfb0d38ff5 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java @@ -10,7 +10,16 @@ import org.opensearch.tasks.Task; +/** + * Represents the CPU time resource type. + */ public class CpuTimeResourceType extends SandboxResourceType { + /** + * 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(); diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java index 666dae6e49b73..b5408de9c61b8 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java @@ -10,7 +10,16 @@ import org.opensearch.tasks.Task; +/** + * Represents the JVM memory resource type. + */ public class JvmMemoryResourceType extends SandboxResourceType { + /** + * 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(); diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java index f6414b6652121..1b46a2b25197e 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java @@ -11,10 +11,31 @@ 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 SandboxResourceType { + /** + * 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 SandboxResourceType from a string. + * If the string is "JVM", a JvmMemoryResourceType is returned. + * If the string is "CPU", a CpuTimeResourceType is returned. + * If the string is not recognized, an IllegalArgumentException is thrown. + * + * @param type The string from which to create a SandboxResourceType + * @return The created SandboxResourceType + * @throws IllegalArgumentException If the string is not recognized + */ public static SandboxResourceType fromString(String type) { if (type.equalsIgnoreCase("JVM")) { return new JvmMemoryResourceType(); From d6918cbd97342cf6af65cd6cba94aa380fe580ce Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 11:15:23 -0700 Subject: [PATCH 08/66] spotless Signed-off-by: Kiran Prakash --- .../search/sandboxing/SandboxService.java | 31 ++++++++++--------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java index 3db9205e5b07e..1c18d0aa9b727 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java @@ -34,9 +34,9 @@ public class SandboxService extends AbstractLifecycleComponent { private static final Logger logger = LogManager.getLogger(SandboxService.class); private final SandboxUsageTracker sandboxUsageTracker; -// private final SandboxPruner sandboxPruner; + // private final SandboxPruner sandboxPruner; private volatile Scheduler.Cancellable scheduledFuture; -// private final SandboxServiceSettings sandboxServiceSettings; + // private final SandboxServiceSettings sandboxServiceSettings; private final ThreadPool threadPool; private final ClusterService clusterService; @@ -44,21 +44,21 @@ public class SandboxService extends AbstractLifecycleComponent { * Guice managed constructor * * @param sandboxUsageTracker -// * @param sandboxPruner -// * @param sandboxServiceSettings + // * @param sandboxPruner + // * @param sandboxServiceSettings * @param threadPool */ @Inject public SandboxService( SandboxUsageTracker sandboxUsageTracker, -// SandboxServiceSettings sandboxServiceSettings, -// SandboxPruner sandboxPruner, + // SandboxServiceSettings sandboxServiceSettings, + // SandboxPruner sandboxPruner, ClusterService clusterService, ThreadPool threadPool ) { this.sandboxUsageTracker = sandboxUsageTracker; -// this.sandboxServiceSettings = sandboxServiceSettings; -// this.sandboxPruner = sandboxPruner; + // this.sandboxServiceSettings = sandboxServiceSettings; + // this.sandboxPruner = sandboxPruner; this.clusterService = clusterService; this.threadPool = threadPool; } @@ -88,14 +88,15 @@ private Set getActiveSandboxes() { @Override protected void doStart() { scheduledFuture = threadPool.scheduleWithFixedDelay(() -> { - try { - doRun(); - } catch (Exception e) { - logger.debug("Exception occurred in Query Sandbox service", e); - } - }, + try { + doRun(); + } catch (Exception e) { + logger.debug("Exception occurred in Query Sandbox service", e); + } + }, new TimeValue(1, TimeUnit.SECONDS), // TODO get this from SandboxServiceSettings - ThreadPool.Names.GENERIC); + ThreadPool.Names.GENERIC + ); } @Override From d38ea3ab7276f3d5e7dd4b452a7faeb4b1ef0c7a Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 11:23:47 -0700 Subject: [PATCH 09/66] javadocs Signed-off-by: Kiran Prakash --- .../org/opensearch/cluster/metadata/Sandbox.java | 12 ++++++++++++ .../search/sandboxing/resourcetype/package-info.java | 3 +++ 2 files changed, 15 insertions(+) diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java index 8db96dffe4ee0..219f11e2014a2 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java @@ -14,6 +14,10 @@ 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 @@ -39,6 +43,10 @@ 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 { public Long getThresholdInLong() { @@ -54,6 +62,10 @@ public Long getThreshold() { } } + /* + Stub class only, this will be added in Kaushal's PR + This class wil be deleted after that. + */ @ExperimentalApi public enum SandboxMode { SOFT("soft"), 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 index 1f7f64b1ba534..f71eb9233abd4 100644 --- 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 @@ -6,4 +6,7 @@ * compatible open source license. */ +/** + * Package for Sandbox Resource Types + */ package org.opensearch.search.sandboxing.resourcetype; From b42f43987404fc751cd4f0ba1a2c8f638f312060 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 11:35:44 -0700 Subject: [PATCH 10/66] javadocs Signed-off-by: Kiran Prakash --- .../opensearch/cluster/metadata/Sandbox.java | 28 ++++++++++--------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java index 219f11e2014a2..ae5bf220df2f7 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java @@ -14,10 +14,10 @@ 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. -* */ +/** + * 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 @@ -43,12 +43,15 @@ public List getResourceLimits() { return Collections.emptyList(); } - /* - Stub class only, this will be added in Kaushal's PR - This class wil be deleted after that. - * */ + /** + * 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 + */ public Long getThresholdInLong() { return 0L; } @@ -62,10 +65,10 @@ public Long getThreshold() { } } - /* - Stub class only, this will be added in Kaushal's PR - This class wil be deleted after that. - */ + /** + * Stub class only, this will be added in Kaushal's PR + * This class wil be deleted after that. + */ @ExperimentalApi public enum SandboxMode { SOFT("soft"), @@ -94,6 +97,5 @@ public static SandboxMode fromName(String s) { throw new IllegalArgumentException("Invalid value for SandboxMode: " + s); } } - } } From ddc1576fdb453891d50cb64d5ac83599d212aef4 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 13:23:11 -0700 Subject: [PATCH 11/66] java docs Signed-off-by: Kiran Prakash --- .../java/org/opensearch/search/sandboxing/SandboxService.java | 4 ++-- .../tracker/SandboxResourceUsageTrackerService.java | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java index 1c18d0aa9b727..21fadbaddb3d1 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxService.java @@ -43,10 +43,10 @@ public class SandboxService extends AbstractLifecycleComponent { /** * Guice managed constructor * - * @param sandboxUsageTracker + * @param sandboxUsageTracker tracker service // * @param sandboxPruner // * @param sandboxServiceSettings - * @param threadPool + * @param threadPool threadpool this will be used to schedule the service */ @Inject public SandboxService( 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 index 22905e2a3f6e4..8bc69c70f1786 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java @@ -35,8 +35,8 @@ public class SandboxResourceUsageTrackerService implements SandboxUsageTracker, /** * SandboxResourceTrackerService constructor * - * @param taskManager - * @param taskResourceTrackingService + * @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( From ec4142a571ca8ee296c8881c2d2fe1b27a3f8768 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 15:02:56 -0700 Subject: [PATCH 12/66] Update AbstractTaskCancellation.java Signed-off-by: Kiran Prakash --- .../search/sandboxing/cancellation/AbstractTaskCancellation.java | 1 - 1 file changed, 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java index 447ba66ebddde..bf0ba64aaf862 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java @@ -118,6 +118,5 @@ private TaskCancellation createTaskCancellation(CancellableTask task) { private void callbackOnCancel() { // todo Implement callback logic here - System.out.println("Task was cancelled."); } } From f3764e5eb0a263293a1520679842f866a8c3f410 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 30 May 2024 15:55:31 -0700 Subject: [PATCH 13/66] Update SandboxModule.java Signed-off-by: Kiran Prakash --- .../opensearch/search/sandboxing/module/SandboxModule.java | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 index 8da5f20ab022f..3f7c8a11099c1 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/module/SandboxModule.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/module/SandboxModule.java @@ -9,8 +9,8 @@ package org.opensearch.search.sandboxing.module; import org.opensearch.common.inject.AbstractModule; -import org.opensearch.search.sandboxing.tracker.SandboxUsageTracker; import org.opensearch.search.sandboxing.tracker.SandboxResourceUsageTrackerService; +import org.opensearch.search.sandboxing.tracker.SandboxUsageTracker; /** * Module class for resource usage limiting related artifacts @@ -25,5 +25,7 @@ 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(); } } From 4b05aecfa8cc202420b7733728bd82937be218ad Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Tue, 4 Jun 2024 12:45:54 -0700 Subject: [PATCH 14/66] Some tests and stubs Signed-off-by: Kiran Prakash --- .../opensearch/cluster/metadata/Sandbox.java | 14 +- .../SandboxLevelResourceUsageViewTest.java | 70 +++++++++ .../DefaultTaskCancellationTest.java | 148 ++++++++++++++++++ 3 files changed, 229 insertions(+), 3 deletions(-) create mode 100644 server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java create mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java index ae5bf220df2f7..85c92897d11d0 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java @@ -52,16 +52,24 @@ public class ResourceLimit { /* Stub class only */ + SandboxResourceType resourceType; + Long threshold; + + ResourceLimit(SandboxResourceType resourceType, Long threshold) { + this.resourceType = resourceType; + this.threshold = threshold; + } + public Long getThresholdInLong() { - return 0L; + return threshold; } public SandboxResourceType getResourceType() { - return null; + return resourceType; } public Long getThreshold() { - return 0L; + return threshold; } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java new file mode 100644 index 0000000000000..cc95cf17ab0fa --- /dev/null +++ b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java @@ -0,0 +1,70 @@ +/* + * 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.SandboxResourceType; +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.SandboxCancellationStrategyTestHelpers.getRandomTask; + +public class SandboxLevelResourceUsageViewTest extends OpenSearchTestCase { + Map resourceUsage; + List activeTasks; + + public void setUp() throws Exception { + super.setUp(); + resourceUsage = Map.of( + SandboxResourceType.fromString("JVM"), 34L, + SandboxResourceType.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(SandboxResourceType.fromString("JVM")) == 34L && + resourceUsageData.get(SandboxResourceType.fromString("CPU")) == 12L; + } +} \ No newline at end of file diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java new file mode 100644 index 0000000000000..cb1ee67bbd8b7 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java @@ -0,0 +1,148 @@ +/* + * 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.junit.Before; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; +import org.opensearch.cluster.metadata.Sandbox; +import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; +import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.test.OpenSearchTestCase; + +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class DefaultTaskCancellationTest extends OpenSearchTestCase { + @Mock + private TaskSelectionStrategy mockStrategy; + + @Mock + private SandboxLevelResourceUsageView mockView; + + private Map sandboxLevelViews; + private Set activeSandboxes; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + sandboxLevelViews = new HashMap<>(); + activeSandboxes = new HashSet<>(); + } + + public void testConstructor() { + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + assertNotNull(cancellation); + } + + public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForDifferentResourceTypes() { + // setup mocks for sandbox1 + Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + // setup mocks for sandbox2 + Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 100L, 50L); + // add both sandboxes to active sandboxes + Collections.addAll(activeSandboxes, sandbox1, sandbox2); + // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + List result = cancellation.getSandboxesToCancelFrom(); + + // only sandbox1 should be returned as it is breaching the threshold + assertEquals(1, result.size()); + assertTrue(result.contains(sandbox1)); + assertFalse(result.contains(sandbox2)); + } + + public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForSameResourceType() { + // setup mocks for sandbox1 + Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + // setup mocks for sandbox2 + Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 100L, 50L); + // add both sandboxes to active sandboxes + Collections.addAll(activeSandboxes, sandbox1, sandbox2); + // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + List result = cancellation.getSandboxesToCancelFrom(); + + // only sandbox1 should be returned as it is breaching the threshold + assertEquals(1, result.size()); + assertTrue(result.contains(sandbox1)); + assertFalse(result.contains(sandbox2)); + } + + public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForDifferentResourceTypes() { + // setup mocks for sandbox1 + Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + // setup mocks for sandbox2 + Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 10L, 50L); + // add both sandboxes to active sandboxes + Collections.addAll(activeSandboxes, sandbox1, sandbox2); + // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + List result = cancellation.getSandboxesToCancelFrom(); + + // Both sandboxes should be returned as it is breaching the threshold + assertEquals(2, result.size()); + assertTrue(result.contains(sandbox1)); + assertTrue(result.contains(sandbox2)); + } + + public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForSameResourceType() { + // setup mocks for sandbox1 + Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + // setup mocks for sandbox2 + Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 10L, 50L); + // add both sandboxes to active sandboxes + Collections.addAll(activeSandboxes, sandbox1, sandbox2); + // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + List result = cancellation.getSandboxesToCancelFrom(); + + // Both sandboxes should be returned as it is breaching the threshold + assertEquals(2, result.size()); + assertTrue(result.contains(sandbox1)); + assertTrue(result.contains(sandbox2)); + } + + // Utility methods + private Sandbox createSandboxMock(String id, String resourceTypeStr, Long threshold, Long usage) { + Sandbox sandbox = Mockito.mock(Sandbox.class); + when(sandbox.getId()).thenReturn(id); + + Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); + when(sandbox.getResourceLimits()).thenReturn(Collections.singletonList(resourceLimitMock)); + + SandboxLevelResourceUsageView mockView = createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView); + + return sandbox; + } + + private Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeStr, Long threshold) { + Sandbox.ResourceLimit resourceLimitMock = mock(Sandbox.ResourceLimit.class); + SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); + when(resourceLimitMock.getResourceType()).thenReturn(resourceType); + when(resourceLimitMock.getThreshold()).thenReturn(threshold); + return resourceLimitMock; + } + + private SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) { + SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class); + SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); + when(mockView.getResourceUsageData()).thenReturn(Collections.singletonMap(resourceType, usage)); + return mockView; + } +} \ No newline at end of file From cdb43a10a34a13ab948d074fe9167987d9997295 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Tue, 4 Jun 2024 12:47:05 -0700 Subject: [PATCH 15/66] spotless Signed-off-by: Kiran Prakash --- .../SandboxLevelResourceUsageViewTest.java | 89 +++---- .../DefaultTaskCancellationTest.java | 245 +++++++++--------- 2 files changed, 162 insertions(+), 172 deletions(-) diff --git a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java index cc95cf17ab0fa..b04360129ef1d 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java @@ -18,53 +18,42 @@ import static org.opensearch.search.sandboxing.cancellation.SandboxCancellationStrategyTestHelpers.getRandomTask; public class SandboxLevelResourceUsageViewTest extends OpenSearchTestCase { - Map resourceUsage; - List activeTasks; - - public void setUp() throws Exception { - super.setUp(); - resourceUsage = Map.of( - SandboxResourceType.fromString("JVM"), 34L, - SandboxResourceType.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(SandboxResourceType.fromString("JVM")) == 34L && - resourceUsageData.get(SandboxResourceType.fromString("CPU")) == 12L; - } -} \ No newline at end of file + Map resourceUsage; + List activeTasks; + + public void setUp() throws Exception { + super.setUp(); + resourceUsage = Map.of(SandboxResourceType.fromString("JVM"), 34L, SandboxResourceType.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(SandboxResourceType.fromString("JVM")) == 34L + && resourceUsageData.get(SandboxResourceType.fromString("CPU")) == 12L; + } +} diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java index cb1ee67bbd8b7..c148e60238b3d 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java @@ -8,14 +8,11 @@ package org.opensearch.search.sandboxing.cancellation; -import org.junit.Before; -import org.mockito.Mock; -import org.mockito.Mockito; -import org.mockito.MockitoAnnotations; import org.opensearch.cluster.metadata.Sandbox; import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; import org.opensearch.test.OpenSearchTestCase; +import org.junit.Before; import java.util.Collections; import java.util.HashMap; @@ -24,125 +21,129 @@ import java.util.Map; import java.util.Set; +import org.mockito.Mock; +import org.mockito.Mockito; +import org.mockito.MockitoAnnotations; + import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; public class DefaultTaskCancellationTest extends OpenSearchTestCase { - @Mock - private TaskSelectionStrategy mockStrategy; - - @Mock - private SandboxLevelResourceUsageView mockView; - - private Map sandboxLevelViews; - private Set activeSandboxes; - - @Before - public void setup() { - MockitoAnnotations.openMocks(this); - sandboxLevelViews = new HashMap<>(); - activeSandboxes = new HashSet<>(); - } - - public void testConstructor() { - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - assertNotNull(cancellation); - } - - public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForDifferentResourceTypes() { - // setup mocks for sandbox1 - Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); - // setup mocks for sandbox2 - Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 100L, 50L); - // add both sandboxes to active sandboxes - Collections.addAll(activeSandboxes, sandbox1, sandbox2); - // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - List result = cancellation.getSandboxesToCancelFrom(); - - // only sandbox1 should be returned as it is breaching the threshold - assertEquals(1, result.size()); - assertTrue(result.contains(sandbox1)); - assertFalse(result.contains(sandbox2)); - } - - public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForSameResourceType() { - // setup mocks for sandbox1 - Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); - // setup mocks for sandbox2 - Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 100L, 50L); - // add both sandboxes to active sandboxes - Collections.addAll(activeSandboxes, sandbox1, sandbox2); - // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - List result = cancellation.getSandboxesToCancelFrom(); - - // only sandbox1 should be returned as it is breaching the threshold - assertEquals(1, result.size()); - assertTrue(result.contains(sandbox1)); - assertFalse(result.contains(sandbox2)); - } - - public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForDifferentResourceTypes() { - // setup mocks for sandbox1 - Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); - // setup mocks for sandbox2 - Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 10L, 50L); - // add both sandboxes to active sandboxes - Collections.addAll(activeSandboxes, sandbox1, sandbox2); - // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - List result = cancellation.getSandboxesToCancelFrom(); - - // Both sandboxes should be returned as it is breaching the threshold - assertEquals(2, result.size()); - assertTrue(result.contains(sandbox1)); - assertTrue(result.contains(sandbox2)); - } - - public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForSameResourceType() { - // setup mocks for sandbox1 - Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); - // setup mocks for sandbox2 - Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 10L, 50L); - // add both sandboxes to active sandboxes - Collections.addAll(activeSandboxes, sandbox1, sandbox2); - // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - List result = cancellation.getSandboxesToCancelFrom(); - - // Both sandboxes should be returned as it is breaching the threshold - assertEquals(2, result.size()); - assertTrue(result.contains(sandbox1)); - assertTrue(result.contains(sandbox2)); - } - - // Utility methods - private Sandbox createSandboxMock(String id, String resourceTypeStr, Long threshold, Long usage) { - Sandbox sandbox = Mockito.mock(Sandbox.class); - when(sandbox.getId()).thenReturn(id); - - Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); - when(sandbox.getResourceLimits()).thenReturn(Collections.singletonList(resourceLimitMock)); - - SandboxLevelResourceUsageView mockView = createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView); - - return sandbox; - } - - private Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeStr, Long threshold) { - Sandbox.ResourceLimit resourceLimitMock = mock(Sandbox.ResourceLimit.class); - SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); - when(resourceLimitMock.getResourceType()).thenReturn(resourceType); - when(resourceLimitMock.getThreshold()).thenReturn(threshold); - return resourceLimitMock; - } - - private SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) { - SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class); - SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); - when(mockView.getResourceUsageData()).thenReturn(Collections.singletonMap(resourceType, usage)); - return mockView; - } -} \ No newline at end of file + @Mock + private TaskSelectionStrategy mockStrategy; + + @Mock + private SandboxLevelResourceUsageView mockView; + + private Map sandboxLevelViews; + private Set activeSandboxes; + + @Before + public void setup() { + MockitoAnnotations.openMocks(this); + sandboxLevelViews = new HashMap<>(); + activeSandboxes = new HashSet<>(); + } + + public void testConstructor() { + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + assertNotNull(cancellation); + } + + public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForDifferentResourceTypes() { + // setup mocks for sandbox1 + Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + // setup mocks for sandbox2 + Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 100L, 50L); + // add both sandboxes to active sandboxes + Collections.addAll(activeSandboxes, sandbox1, sandbox2); + // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + List result = cancellation.getSandboxesToCancelFrom(); + + // only sandbox1 should be returned as it is breaching the threshold + assertEquals(1, result.size()); + assertTrue(result.contains(sandbox1)); + assertFalse(result.contains(sandbox2)); + } + + public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForSameResourceType() { + // setup mocks for sandbox1 + Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + // setup mocks for sandbox2 + Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 100L, 50L); + // add both sandboxes to active sandboxes + Collections.addAll(activeSandboxes, sandbox1, sandbox2); + // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + List result = cancellation.getSandboxesToCancelFrom(); + + // only sandbox1 should be returned as it is breaching the threshold + assertEquals(1, result.size()); + assertTrue(result.contains(sandbox1)); + assertFalse(result.contains(sandbox2)); + } + + public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForDifferentResourceTypes() { + // setup mocks for sandbox1 + Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + // setup mocks for sandbox2 + Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 10L, 50L); + // add both sandboxes to active sandboxes + Collections.addAll(activeSandboxes, sandbox1, sandbox2); + // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + List result = cancellation.getSandboxesToCancelFrom(); + + // Both sandboxes should be returned as it is breaching the threshold + assertEquals(2, result.size()); + assertTrue(result.contains(sandbox1)); + assertTrue(result.contains(sandbox2)); + } + + public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForSameResourceType() { + // setup mocks for sandbox1 + Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + // setup mocks for sandbox2 + Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 10L, 50L); + // add both sandboxes to active sandboxes + Collections.addAll(activeSandboxes, sandbox1, sandbox2); + // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom + DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); + List result = cancellation.getSandboxesToCancelFrom(); + + // Both sandboxes should be returned as it is breaching the threshold + assertEquals(2, result.size()); + assertTrue(result.contains(sandbox1)); + assertTrue(result.contains(sandbox2)); + } + + // Utility methods + private Sandbox createSandboxMock(String id, String resourceTypeStr, Long threshold, Long usage) { + Sandbox sandbox = Mockito.mock(Sandbox.class); + when(sandbox.getId()).thenReturn(id); + + Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); + when(sandbox.getResourceLimits()).thenReturn(Collections.singletonList(resourceLimitMock)); + + SandboxLevelResourceUsageView mockView = createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView); + + return sandbox; + } + + private Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeStr, Long threshold) { + Sandbox.ResourceLimit resourceLimitMock = mock(Sandbox.ResourceLimit.class); + SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); + when(resourceLimitMock.getResourceType()).thenReturn(resourceType); + when(resourceLimitMock.getThreshold()).thenReturn(threshold); + return resourceLimitMock; + } + + private SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) { + SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class); + SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); + when(mockView.getResourceUsageData()).thenReturn(Collections.singletonMap(resourceType, usage)); + return mockView; + } +} From 4d7436d66140500abd6132122371376f83e6b80b Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Tue, 4 Jun 2024 14:29:32 -0700 Subject: [PATCH 16/66] :server:testingConventions Signed-off-by: Kiran Prakash --- ...ageViewTest.java => SandboxLevelResourceUsageViewTests.java} | 2 +- ...kCancellationTest.java => DefaultTaskCancellationTests.java} | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename server/src/test/java/org/opensearch/search/sandboxing/{SandboxLevelResourceUsageViewTest.java => SandboxLevelResourceUsageViewTests.java} (97%) rename server/src/test/java/org/opensearch/search/sandboxing/cancellation/{DefaultTaskCancellationTest.java => DefaultTaskCancellationTests.java} (98%) diff --git a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java similarity index 97% rename from server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java rename to server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java index b04360129ef1d..a49a7d4cd70b2 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTest.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java @@ -17,7 +17,7 @@ import static org.opensearch.search.sandboxing.cancellation.SandboxCancellationStrategyTestHelpers.getRandomTask; -public class SandboxLevelResourceUsageViewTest extends OpenSearchTestCase { +public class SandboxLevelResourceUsageViewTests extends OpenSearchTestCase { Map resourceUsage; List activeTasks; diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java similarity index 98% rename from server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java rename to server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java index c148e60238b3d..6b85431b33ed1 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTest.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java @@ -28,7 +28,7 @@ import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -public class DefaultTaskCancellationTest extends OpenSearchTestCase { +public class DefaultTaskCancellationTests extends OpenSearchTestCase { @Mock private TaskSelectionStrategy mockStrategy; From 09589071e476a72fc628854328acacb889e9fa93 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 5 Jun 2024 01:13:19 -0700 Subject: [PATCH 17/66] Update AbstractTaskCancellation.java Signed-off-by: Kiran Prakash --- .../sandboxing/cancellation/AbstractTaskCancellation.java | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java index bf0ba64aaf862..eb1e92ddf0bc1 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java @@ -100,6 +100,10 @@ private List getTaskCancellations(Sandbox sandbox, SandboxReso } private long getReduceBy(Sandbox sandbox, SandboxResourceType resourceType) { + Long usage = getUsage(sandbox, resourceType); + if (usage == null) { + return 0; + } return getUsage(sandbox, resourceType) - sandbox.getResourceLimitFor(resourceType).getThresholdInLong(); } @@ -112,8 +116,8 @@ private List getAllTasksInSandbox(String sandboxId) { } private TaskCancellation createTaskCancellation(CancellableTask task) { - // todo add reasons and callbacks - return new TaskCancellation(task, List.of(), List.of(this::callbackOnCancel)); + // todo add correct reason and callbacks + return new TaskCancellation(task, List.of(new TaskCancellation.Reason("limits exceeded", 5)), List.of(this::callbackOnCancel)); } private void callbackOnCancel() { From 9870e04d6318f1acb296c20ae7199df1e6da95ca Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 5 Jun 2024 01:14:35 -0700 Subject: [PATCH 18/66] more tests Signed-off-by: Kiran Prakash --- .../SandboxResourceUsageTrackerService.java | 5 +- .../SandboxLevelResourceUsageViewTests.java | 2 +- .../AbstractTaskCancellationTests.java | 159 ++++++++++++++++++ .../DefaultTaskCancellationTests.java | 94 ++++++----- ...andboxCancellationStrategyTestHelpers.java | 63 ------- .../cancellation/SandboxTestHelpers.java | 109 ++++++++++++ .../TaskSelectionStrategyTests.java | 8 +- 7 files changed, 329 insertions(+), 111 deletions(-) create mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellationTests.java delete mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java create mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java 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 index 8bc69c70f1786..c51ef7f105efb 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java @@ -27,7 +27,10 @@ // @ExperimentalApi public class SandboxResourceUsageTrackerService implements SandboxUsageTracker, TaskManager.TaskEventListeners { - public static final List TRACKED_RESOURCES = List.of(SandboxResourceType.fromString("JVM")); + public static final List TRACKED_RESOURCES = List.of( + SandboxResourceType.fromString("JVM"), + SandboxResourceType.fromString("CPU") + ); private final TaskManager taskManager; private final TaskResourceTrackingService taskResourceTrackingService; diff --git a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java index a49a7d4cd70b2..2376db7eda1f9 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java @@ -15,7 +15,7 @@ import java.util.List; import java.util.Map; -import static org.opensearch.search.sandboxing.cancellation.SandboxCancellationStrategyTestHelpers.getRandomTask; +import static org.opensearch.search.sandboxing.cancellation.SandboxTestHelpers.getRandomTask; public class SandboxLevelResourceUsageViewTests extends OpenSearchTestCase { Map resourceUsage; diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellationTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellationTests.java new file mode 100644 index 0000000000000..f59f6ea0aa65f --- /dev/null +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellationTests.java @@ -0,0 +1,159 @@ +/* + * 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.SandboxResourceType; +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 AbstractTaskCancellationTests extends OpenSearchTestCase { + + private class TestTaskCancellationImpl extends AbstractTaskCancellation { + + public TestTaskCancellationImpl( + TaskSelectionStrategy taskSelectionStrategy, + Map sandboxLevelViews, + Set activeSandboxes + ) { + super(taskSelectionStrategy, sandboxLevelViews, activeSandboxes); + } + + @Override + List getSandboxesToCancelFrom() { + return new ArrayList<>(activeSandboxes); + } + } + + private TaskSelectionStrategy taskSelectionStrategy; + private Map sandboxLevelViews; + private Set activeSandboxes; + private AbstractTaskCancellation 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(SandboxResourceType.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); + SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView); + + List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(sandbox1); + assertEquals(2, cancellableTasksFrom.size()); + assertEquals(4321, cancellableTasksFrom.get(0).getTask().getId()); + assertEquals(1234, 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(SandboxResourceType.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(SandboxResourceType.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(4321, cancellableTasksFrom.get(0).getTask().getId()); + assertEquals(1234, 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(SandboxResourceType.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(SandboxResourceType.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(4321, allCancellableTasks.get(0).getTask().getId()); + assertEquals(1234, allCancellableTasks.get(1).getTask().getId()); + } +} 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 index 6b85431b33ed1..a7f144526a5ed 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java @@ -10,7 +10,6 @@ import org.opensearch.cluster.metadata.Sandbox; import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; import org.opensearch.test.OpenSearchTestCase; import org.junit.Before; @@ -22,12 +21,8 @@ import java.util.Set; import org.mockito.Mock; -import org.mockito.Mockito; import org.mockito.MockitoAnnotations; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - public class DefaultTaskCancellationTests extends OpenSearchTestCase { @Mock private TaskSelectionStrategy mockStrategy; @@ -52,9 +47,20 @@ public void testConstructor() { public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForDifferentResourceTypes() { // setup mocks for sandbox1 - Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + String id = "sandbox1"; + String resourceTypeStr = "CPU"; + long usage = 50L; + Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView); + // setup mocks for sandbox2 - Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 100L, 50L); + id = "sandbox2"; + resourceTypeStr = "JVM"; + usage = 50L; + Sandbox sandbox2 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 100L, usage); + SandboxLevelResourceUsageView mockView2 = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView2); // add both sandboxes to active sandboxes Collections.addAll(activeSandboxes, sandbox1, sandbox2); // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom @@ -69,9 +75,20 @@ public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForDiffe public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForSameResourceType() { // setup mocks for sandbox1 - Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + String id = "sandbox1"; + String resourceTypeStr = "CPU"; + long usage = 50L; + Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView); + // setup mocks for sandbox2 - Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 100L, 50L); + id = "sandbox2"; + resourceTypeStr = "CPU"; + usage = 50L; + Sandbox sandbox2 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 100L, usage); + SandboxLevelResourceUsageView mockView2 = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView2); // add both sandboxes to active sandboxes Collections.addAll(activeSandboxes, sandbox1, sandbox2); // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom @@ -86,9 +103,20 @@ public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForSameR public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForDifferentResourceTypes() { // setup mocks for sandbox1 - Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + String id = "sandbox1"; + String resourceTypeStr = "CPU"; + long usage = 50L; + Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView); + // setup mocks for sandbox2 - Sandbox sandbox2 = createSandboxMock("sandbox2", "JVM", 10L, 50L); + id = "sandbox2"; + resourceTypeStr = "JVM"; + usage = 50L; + Sandbox sandbox2 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + SandboxLevelResourceUsageView mockView2 = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView2); // add both sandboxes to active sandboxes Collections.addAll(activeSandboxes, sandbox1, sandbox2); // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom @@ -103,9 +131,20 @@ public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForDifferen public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForSameResourceType() { // setup mocks for sandbox1 - Sandbox sandbox1 = createSandboxMock("sandbox1", "CPU", 10L, 50L); + String id = "sandbox1"; + String resourceTypeStr = "CPU"; + long usage = 50L; + Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView); + // setup mocks for sandbox2 - Sandbox sandbox2 = createSandboxMock("sandbox2", "CPU", 10L, 50L); + id = "sandbox2"; + resourceTypeStr = "CPU"; + usage = 50L; + Sandbox sandbox2 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + SandboxLevelResourceUsageView mockView2 = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); + sandboxLevelViews.put(id, mockView2); // add both sandboxes to active sandboxes Collections.addAll(activeSandboxes, sandbox1, sandbox2); // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom @@ -117,33 +156,4 @@ public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForSameReso assertTrue(result.contains(sandbox1)); assertTrue(result.contains(sandbox2)); } - - // Utility methods - private Sandbox createSandboxMock(String id, String resourceTypeStr, Long threshold, Long usage) { - Sandbox sandbox = Mockito.mock(Sandbox.class); - when(sandbox.getId()).thenReturn(id); - - Sandbox.ResourceLimit resourceLimitMock = createResourceLimitMock(resourceTypeStr, threshold); - when(sandbox.getResourceLimits()).thenReturn(Collections.singletonList(resourceLimitMock)); - - SandboxLevelResourceUsageView mockView = createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView); - - return sandbox; - } - - private Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeStr, Long threshold) { - Sandbox.ResourceLimit resourceLimitMock = mock(Sandbox.ResourceLimit.class); - SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); - when(resourceLimitMock.getResourceType()).thenReturn(resourceType); - when(resourceLimitMock.getThreshold()).thenReturn(threshold); - return resourceLimitMock; - } - - private SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) { - SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class); - SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); - when(mockView.getResourceUsageData()).thenReturn(Collections.singletonMap(resourceType, usage)); - return mockView; - } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java deleted file mode 100644 index 5308b526b820c..0000000000000 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxCancellationStrategyTestHelpers.java +++ /dev/null @@ -1,63 +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.action.search.SearchAction; -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.tasks.Task; - -import java.util.ArrayList; -import java.util.Collections; -import java.util.List; - -import static org.opensearch.test.OpenSearchTestCase.randomLong; -import static org.opensearch.test.OpenSearchTestCase.randomLongBetween; - -public class SandboxCancellationStrategyTestHelpers { - - public static List getListOfTasks(long totalMemory) { - List tasks = new ArrayList<>(); - - while (totalMemory > 0) { - long id = randomLong(); - final Task task = getRandomTask(id); - long initial_memory = randomLongBetween(1, 100); - - ResourceUsageMetric[] initialTaskResourceMetrics = new ResourceUsageMetric[] { - new ResourceUsageMetric(ResourceStats.MEMORY, initial_memory) }; - task.startThreadResourceTracking(id, ResourceStatsType.WORKER_STATS, initialTaskResourceMetrics); - - long memory = initial_memory + randomLongBetween(1, 10000); - - totalMemory -= memory - initial_memory; - - ResourceUsageMetric[] taskResourceMetrics = new ResourceUsageMetric[] { - new ResourceUsageMetric(ResourceStats.MEMORY, memory), }; - task.updateThreadResourceStats(id, ResourceStatsType.WORKER_STATS, taskResourceMetrics); - task.stopThreadResourceTracking(id, ResourceStatsType.WORKER_STATS); - tasks.add(task); - } - - return tasks; - } - - public static Task getRandomTask(long id) { - return new Task( - id, - "transport", - SearchAction.NAME, - "test description", - new TaskId(randomLong() + ":" + randomLong()), - Collections.emptyMap() - ); - } -} diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java new file mode 100644 index 0000000000000..df29466cf1f5c --- /dev/null +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java @@ -0,0 +1,109 @@ +/* + * 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.action.search.SearchAction; +import org.opensearch.action.search.SearchTask; +import org.opensearch.cluster.metadata.Sandbox; +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.SandboxResourceType; +import org.opensearch.tasks.Task; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; + +import org.mockito.Mockito; + +import static org.opensearch.test.OpenSearchTestCase.randomLong; +import static org.opensearch.test.OpenSearchTestCase.randomLongBetween; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class SandboxTestHelpers { + + public static List getListOfTasks(long totalMemory) { + List tasks = new ArrayList<>(); + + while (totalMemory > 0) { + long id = randomLong(); + final Task task = getRandomTask(id); + long initial_memory = randomLongBetween(1, 100); + + ResourceUsageMetric[] initialTaskResourceMetrics = new ResourceUsageMetric[] { + new ResourceUsageMetric(ResourceStats.MEMORY, initial_memory) }; + task.startThreadResourceTracking(id, ResourceStatsType.WORKER_STATS, initialTaskResourceMetrics); + + long memory = initial_memory + randomLongBetween(1, 10000); + + totalMemory -= memory - initial_memory; + + ResourceUsageMetric[] taskResourceMetrics = new ResourceUsageMetric[] { + new ResourceUsageMetric(ResourceStats.MEMORY, memory), }; + task.updateThreadResourceStats(id, ResourceStatsType.WORKER_STATS, taskResourceMetrics); + task.stopThreadResourceTracking(id, ResourceStatsType.WORKER_STATS); + tasks.add(task); + } + + return tasks; + } + + public static Task getRandomTask(long id) { + return new Task( + id, + "transport", + SearchAction.NAME, + "test description", + new TaskId(randomLong() + ":" + randomLong()), + Collections.emptyMap() + ); + } + + public static Task getRandomSearchTask(long id) { + return new SearchTask( + id, + "transport", + SearchAction.NAME, + () -> "test description", + new TaskId(randomLong() + ":" + randomLong()), + Collections.emptyMap() + ); + } + + public static Sandbox createSandboxMock(String id, String resourceTypeStr, Long threshold, Long usage) { + Sandbox sandbox = Mockito.mock(Sandbox.class); + when(sandbox.getId()).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); + SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); + when(resourceLimitMock.getResourceType()).thenReturn(resourceType); + when(resourceLimitMock.getThreshold()).thenReturn(threshold); + when(resourceLimitMock.getThresholdInLong()).thenReturn(threshold); + return resourceLimitMock; + } + + public static SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) { + SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class); + SandboxResourceType resourceType = SandboxResourceType.fromString(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/TaskSelectionStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java index c513cb74a54fd..7936e940e9cc9 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java @@ -18,7 +18,7 @@ public class TaskSelectionStrategyTests extends OpenSearchTestCase { - private class TestTaskSelectionStrategy extends AbstractTaskSelectionStrategy { + public static class TestTaskSelectionStrategy extends AbstractTaskSelectionStrategy { @Override public Comparator sortingCondition() { return Comparator.comparingLong(Task::getStartTime).reversed(); @@ -30,7 +30,7 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsGrea long threshold = 100L; long reduceBy = 50L; SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); - List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); + List tasks = SandboxTestHelpers.getListOfTasks(threshold); List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); assertFalse(selectedTasks.isEmpty()); @@ -42,7 +42,7 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsLess long threshold = 100L; long reduceBy = -50L; SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); - List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); + List tasks = SandboxTestHelpers.getListOfTasks(threshold); try { testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); @@ -57,7 +57,7 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsEqua long threshold = 100L; long reduceBy = 0; SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); - List tasks = SandboxCancellationStrategyTestHelpers.getListOfTasks(threshold); + List tasks = SandboxTestHelpers.getListOfTasks(threshold); List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); assertTrue(selectedTasks.isEmpty()); From 837149b262d24019ddca88bf2f59b9920b692727 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 3 Jul 2024 16:21:51 -0700 Subject: [PATCH 19/66] addressing comments Signed-off-by: Kiran Prakash --- .../java/Sandbox/SandboxMode/ENFORCED.java | 11 + .../opensearch/cluster/metadata/Sandbox.java | 10 +- .../sandboxing/DefaultSandboxPruner.java | 66 ++++++ .../SandboxLevelResourceUsageView.java | 13 +- .../search/sandboxing/SandboxPruner.java | 21 ++ .../sandboxing/SandboxServiceSettings.java | 198 +++++++++++++++++ .../search/sandboxing/TaskData.java | 28 +++ .../AbstractTaskCancellation.java | 126 ----------- .../AbstractTaskSelectionStrategy.java | 27 ++- .../cancellation/DefaultTaskCancellation.java | 123 +++++++++-- .../LongestRunningTaskFirstStrategy.java | 2 +- ...ostResourceIntensiveTaskFirstStrategy.java | 27 +++ .../ShortestRunningTaskFirstStrategy.java | 2 +- .../cancellation/TaskCancellationContext.java | 35 +++ .../cancellation/TaskSelectionStrategy.java | 5 +- ...ResourceType.java => CpuTimeResource.java} | 4 +- ...sourceType.java => JvmMemoryResource.java} | 4 +- ...xResourceType.java => SystemResource.java} | 18 +- .../SandboxResourceUsageTrackerService.java | 18 +- .../SandboxLevelResourceUsageViewTests.java | 16 +- .../AbstractTaskCancellationTests.java | 159 -------------- .../DefaultTaskCancellationTests.java | 203 ++++++++++-------- ...RunningTaskFirstStrategyStrategyTests.java | 4 +- .../cancellation/SandboxTestHelpers.java | 13 +- .../TaskSelectionStrategyTests.java | 19 +- ...ndboxResourceUsageTrackerServiceTests.java | 8 +- 26 files changed, 697 insertions(+), 463 deletions(-) create mode 100644 server/src/main/java/Sandbox/SandboxMode/ENFORCED.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/DefaultSandboxPruner.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/SandboxPruner.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/SandboxServiceSettings.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/TaskData.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/MostResourceIntensiveTaskFirstStrategy.java create mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskCancellationContext.java rename server/src/main/java/org/opensearch/search/sandboxing/resourcetype/{CpuTimeResourceType.java => CpuTimeResource.java} (87%) rename server/src/main/java/org/opensearch/search/sandboxing/resourcetype/{JvmMemoryResourceType.java => JvmMemoryResource.java} (87%) rename server/src/main/java/org/opensearch/search/sandboxing/resourcetype/{SandboxResourceType.java => SystemResource.java} (71%) delete mode 100644 server/src/test/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellationTests.java diff --git a/server/src/main/java/Sandbox/SandboxMode/ENFORCED.java b/server/src/main/java/Sandbox/SandboxMode/ENFORCED.java new file mode 100644 index 0000000000000..ab8c2ee2cb60a --- /dev/null +++ b/server/src/main/java/Sandbox/SandboxMode/ENFORCED.java @@ -0,0 +1,11 @@ +/* + * 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 Sandbox.SandboxMode; + +public class ENFORCED {} diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java index 85c92897d11d0..a9014bf57bf64 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Sandbox.java @@ -9,7 +9,7 @@ package org.opensearch.cluster.metadata; import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.resourcetype.SystemResource; import java.util.Collections; import java.util.List; @@ -27,7 +27,7 @@ public SandboxMode getMode() { return mode; } - public ResourceLimit getResourceLimitFor(SandboxResourceType resourceType) { + public ResourceLimit getResourceLimitFor(SystemResource resourceType) { return null; } @@ -52,10 +52,10 @@ public class ResourceLimit { /* Stub class only */ - SandboxResourceType resourceType; + SystemResource resourceType; Long threshold; - ResourceLimit(SandboxResourceType resourceType, Long threshold) { + ResourceLimit(SystemResource resourceType, Long threshold) { this.resourceType = resourceType; this.threshold = threshold; } @@ -64,7 +64,7 @@ public Long getThresholdInLong() { return threshold; } - public SandboxResourceType getResourceType() { + public SystemResource getResourceType() { return resourceType; } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/DefaultSandboxPruner.java b/server/src/main/java/org/opensearch/search/sandboxing/DefaultSandboxPruner.java new file mode 100644 index 0000000000000..471c71b2cb205 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/DefaultSandboxPruner.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.sandboxing; + +import org.opensearch.cluster.metadata.Sandbox; + +import java.util.List; +import java.util.stream.Collectors; + +public class DefaultSandboxPruner implements SandboxPruner { + private List activeSandboxes; + private List toDeleteSandboxes; + private TaskData taskData; + + @Override + public void pruneSandboxes() { + toDeleteSandboxes = toDeleteSandboxes.stream().filter(this::hasUnfinishedTasks).collect(Collectors.toList()); + } + + @Override + public void deleteSandbox(String sandboxId) { + // should be called from the API and we need to delete from cluster metadata + // TODO may be we need to build a listener interface instead + if (hasUnfinishedTasks(sandboxId)) { + toDeleteSandboxes.add(sandboxId); + } + // remove this sandbox from the active sandboxes + activeSandboxes = activeSandboxes.stream() + .filter(resourceLimitGroup -> !resourceLimitGroup.getName().equals(sandboxId)) + .collect(Collectors.toList()); + } + + private boolean hasUnfinishedTasks(String sandboxName) { + return !taskData.getTasksBySandbox().get(sandboxName).isEmpty(); + } + + public List getActiveSandboxes() { + return activeSandboxes; + } + + public void setActiveSandboxes(List activeSandboxes) { + this.activeSandboxes = activeSandboxes; + } + + public List getToDeleteSandboxes() { + return toDeleteSandboxes; + } + + public void setToDeleteSandboxes(List toDeleteSandboxes) { + this.toDeleteSandboxes = toDeleteSandboxes; + } + + public TaskData getTaskData() { + return taskData; + } + + public void setTaskData(TaskData taskData) { + this.taskData = taskData; + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java index c104170c2eee0..a6da186487a45 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageView.java @@ -9,7 +9,7 @@ package org.opensearch.search.sandboxing; import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.resourcetype.SystemResource; import org.opensearch.tasks.Task; import java.util.ArrayList; @@ -19,14 +19,17 @@ import java.util.Objects; /** - * Represents the resource usage view at the sandbox level. + * 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. */ @ExperimentalApi public class SandboxLevelResourceUsageView { private final String sandboxId; - private final Map resourceUsage; + // 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 List activeTasks; public SandboxLevelResourceUsageView(String sandboxId) { @@ -35,7 +38,7 @@ public SandboxLevelResourceUsageView(String sandboxId) { this.activeTasks = new ArrayList<>(); } - public SandboxLevelResourceUsageView(String sandboxId, Map resourceUsage, List activeTasks) { + public SandboxLevelResourceUsageView(String sandboxId, Map resourceUsage, List activeTasks) { this.sandboxId = sandboxId; this.resourceUsage = resourceUsage; this.activeTasks = activeTasks; @@ -46,7 +49,7 @@ public SandboxLevelResourceUsageView(String sandboxId, Map getResourceUsageData() { + public Map getResourceUsageData() { return resourceUsage; } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxPruner.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxPruner.java new file mode 100644 index 0000000000000..40d7039333b85 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxPruner.java @@ -0,0 +1,21 @@ +/* + * 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; + +/** + * This interface is used to identify and completely remove deleted resourceLimitGroups which has been marked as deleted + * previously but had the tasks running at the time of deletion request + */ +public interface SandboxPruner { + // remove the deleted sandboxes from the system once all the tasks in that sandbox are completed/cancelled + void pruneSandboxes(); + + // accepts the request to delete the sandbox and marks the sandbox as deleted + void deleteSandbox(String sandboxId); +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxServiceSettings.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxServiceSettings.java new file mode 100644 index 0000000000000..d31d905c4503d --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/SandboxServiceSettings.java @@ -0,0 +1,198 @@ +/* + * 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.common.settings.ClusterSettings; +import org.opensearch.common.settings.Setting; +import org.opensearch.common.settings.Settings; +import org.opensearch.common.unit.TimeValue; + +/** + * Main class to declare the query ResourceLimitGrouping feature related settings + */ +public class SandboxServiceSettings { + private static final Long DEFAULT_RUN_INTERVAL_MILLIS = 1000L; + private static final Double DEFAULT_NODE_LEVEL_REJECTION_THRESHOLD = 0.8; + private static final Double DEFAULT_NODE_LEVEL_CANCELLATION_THRESHOLD = 0.9; + /** + * default max ResourceLimitGroup count on any node at any given point in time + */ + public static final int DEFAULT_MAX_RESOURCE_LIMIT_GROUP_COUNT_VALUE = 100; + + public static final String RESOURCE_LIMIT_GROUP_COUNT_SETTING_NAME = "node.resource_limit_group.max_count"; + public static final double NODE_LEVEL_CANCELLATION_THRESHOLD_MAX_VALUE = 0.95; + public static final double NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE = 0.90; + + private TimeValue runIntervalMillis; + private Double nodeLevelJvmCancellationThreshold; + private Double nodeLevelJvmRejectionThreshold; + private volatile int maxResourceLimitGroupCount; + /** + * max ResourceLimitGroup count setting + */ + public static final Setting MAX_RESOURCE_LIMIT_GROUP_COUNT = Setting.intSetting( + RESOURCE_LIMIT_GROUP_COUNT_SETTING_NAME, + DEFAULT_MAX_RESOURCE_LIMIT_GROUP_COUNT_VALUE, + 0, + (newVal) -> { + if (newVal > 100 || newVal < 1) throw new IllegalArgumentException( + RESOURCE_LIMIT_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]" + ); + }, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + /** + * Setting name for default ResourceLimitGroup count + */ + public static final String SERVICE_RUN_INTERVAL_MILLIS_SETTING_NAME = "resource_limit_group.service.run_interval_millis"; + /** + * Setting to control the run interval of QSB service + */ + private static final Setting RESOURCE_LIMIT_GROUP_RUN_INTERVAL_SETTING = Setting.longSetting( + SERVICE_RUN_INTERVAL_MILLIS_SETTING_NAME, + DEFAULT_RUN_INTERVAL_MILLIS, + 1, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * Setting name for node level rejection threshold for QSB + */ + public static final String NODE_REJECTION_THRESHOLD_SETTING_NAME = "resource_limit_group.node.rejection_threshold"; + /** + * Setting to control the rejection threshold + */ + public static final Setting NODE_LEVEL_REJECTION_THRESHOLD = Setting.doubleSetting( + NODE_REJECTION_THRESHOLD_SETTING_NAME, + DEFAULT_NODE_LEVEL_REJECTION_THRESHOLD, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + /** + * Setting name for node level cancellation threshold + */ + public static final String NODE_CANCELLATION_THRESHOLD_SETTING_NAME = "resource_limit_group.node.cancellation_threshold"; + /** + * Setting name for node level cancellation threshold + */ + public static final Setting NODE_LEVEL_CANCELLATION_THRESHOLD = Setting.doubleSetting( + NODE_CANCELLATION_THRESHOLD_SETTING_NAME, + DEFAULT_NODE_LEVEL_CANCELLATION_THRESHOLD, + Setting.Property.Dynamic, + Setting.Property.NodeScope + ); + + /** + * ResourceLimitGroup service settings constructor + * @param settings + * @param clusterSettings + */ + public SandboxServiceSettings(Settings settings, ClusterSettings clusterSettings) { + runIntervalMillis = new TimeValue(RESOURCE_LIMIT_GROUP_RUN_INTERVAL_SETTING.get(settings)); + nodeLevelJvmCancellationThreshold = NODE_LEVEL_CANCELLATION_THRESHOLD.get(settings); + nodeLevelJvmRejectionThreshold = NODE_LEVEL_REJECTION_THRESHOLD.get(settings); + maxResourceLimitGroupCount = MAX_RESOURCE_LIMIT_GROUP_COUNT.get(settings); + + ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); + + clusterSettings.addSettingsUpdateConsumer(MAX_RESOURCE_LIMIT_GROUP_COUNT, this::setMaxResourceLimitGroupCount); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CANCELLATION_THRESHOLD, this::setNodeLevelJvmCancellationThreshold); + clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_REJECTION_THRESHOLD, this::setNodeLevelJvmRejectionThreshold); + } + + /** + * Method to get runInterval for QSB + * @return runInterval in milliseconds for QSB Service + */ + public TimeValue getRunIntervalMillis() { + return runIntervalMillis; + } + + /** + * Method to set the new ResourceLimitGroup count + * @param newMaxResourceLimitGroupCount is the new maxResourceLimitGroupCount per node + */ + public void setMaxResourceLimitGroupCount(int newMaxResourceLimitGroupCount) { + if (newMaxResourceLimitGroupCount < 0) { + throw new IllegalArgumentException("node.ResourceLimitGroup.max_count can't be negative"); + } + this.maxResourceLimitGroupCount = newMaxResourceLimitGroupCount; + } + + /** + * Method to get the node level cancellation threshold + * @return current node level cancellation threshold + */ + public Double getNodeLevelJvmCancellationThreshold() { + return nodeLevelJvmCancellationThreshold; + } + + /** + * Method to set the node level cancellation threshold + * @param nodeLevelJvmCancellationThreshold sets the new node level cancellation threshold + * @throws IllegalArgumentException if the value is > 0.95 and cancellation < rejection threshold + */ + public void setNodeLevelJvmCancellationThreshold(Double nodeLevelJvmCancellationThreshold) { + if (Double.compare(nodeLevelJvmCancellationThreshold, NODE_LEVEL_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { + throw new IllegalArgumentException( + NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be greater than 0.95 as it pose a threat of node drop" + ); + } + + ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); + + this.nodeLevelJvmCancellationThreshold = nodeLevelJvmCancellationThreshold; + } + + /** + * Method to get the node level rejection threshold + * @return the current node level rejection threshold + */ + public Double getNodeLevelJvmRejectionThreshold() { + return nodeLevelJvmRejectionThreshold; + } + + /** + * Method to set the node level rejection threshold + * @param nodeLevelJvmRejectionThreshold sets the new rejection threshold + * @throws IllegalArgumentException if rejection > 0.90 and rejection < cancellation threshold + */ + public void setNodeLevelJvmRejectionThreshold(Double nodeLevelJvmRejectionThreshold) { + if (Double.compare(nodeLevelJvmRejectionThreshold, NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE) > 0) { + throw new IllegalArgumentException( + NODE_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" + ); + } + + ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); + + this.nodeLevelJvmRejectionThreshold = nodeLevelJvmRejectionThreshold; + } + + private void ensureRejectionThresholdIsLessThanCancellation( + Double nodeLevelJvmRejectionThreshold, + Double nodeLevelJvmCancellationThreshold + ) { + if (Double.compare(nodeLevelJvmCancellationThreshold, nodeLevelJvmRejectionThreshold) < 0) { + throw new IllegalArgumentException( + NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be less than " + NODE_REJECTION_THRESHOLD_SETTING_NAME + ); + } + } + + /** + * Method to get the current ResourceLimitGroup count + * @return the current max ResourceLimitGroup count + */ + public int getMaxResourceLimitGroupCount() { + return maxResourceLimitGroupCount; + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/TaskData.java b/server/src/main/java/org/opensearch/search/sandboxing/TaskData.java new file mode 100644 index 0000000000000..181867a8a7dbf --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/TaskData.java @@ -0,0 +1,28 @@ +/* + * 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.common.annotation.ExperimentalApi; +import org.opensearch.tasks.Task; + +import java.util.List; +import java.util.Map; + +@ExperimentalApi +public class TaskData { + private final Map> tasksBySandboxes; + + public TaskData(Map> tasksBySandboxes) { + this.tasksBySandboxes = tasksBySandboxes; + } + + public Map> getTasksBySandbox() { + return tasksBySandboxes; + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.java deleted file mode 100644 index eb1e92ddf0bc1..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellation.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.cancellation; - -import org.opensearch.cluster.metadata.Sandbox; -import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; -import org.opensearch.tasks.CancellableTask; -import org.opensearch.tasks.Task; -import org.opensearch.tasks.TaskCancellation; - -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; - -/** - * Abstract class that provides a structure for task cancellation. - * This class is extended by other classes to provide specific task cancellation strategies. - */ -public abstract class AbstractTaskCancellation { - // Strategy for selecting tasks to cancel - protected final TaskSelectionStrategy taskSelectionStrategy; - // Views of resource usage at the sandbox level - protected final Map sandboxLevelViews; - // Set of active sandboxes - protected final Set activeSandboxes; - - public AbstractTaskCancellation( - 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(); - } - } - - /** - * Abstract method to get the list of sandboxes from which tasks can be cancelled. - * This method needs to be implemented by subclasses. - * - * @return List of sandboxes - */ - abstract List getSandboxesToCancelFrom(); - - /** - * 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, SandboxResourceType resourceType) { - long reduceBy = getReduceBy(sandbox, resourceType); - return reduceBy > 0; - } - - private List getTaskCancellations(Sandbox sandbox, SandboxResourceType resourceType) { - return taskSelectionStrategy.selectTasksForCancellation( - getAllTasksInSandbox(sandbox.getId()), - getReduceBy(sandbox, resourceType), - resourceType - ).stream().map(task -> createTaskCancellation((CancellableTask) task)).collect(Collectors.toList()); - } - - private long getReduceBy(Sandbox sandbox, SandboxResourceType resourceType) { - Long usage = getUsage(sandbox, resourceType); - if (usage == null) { - return 0; - } - return getUsage(sandbox, resourceType) - sandbox.getResourceLimitFor(resourceType).getThresholdInLong(); - } - - private Long getUsage(Sandbox sandbox, SandboxResourceType resourceType) { - return sandboxLevelViews.get(sandbox.getId()).getResourceUsageData().get(resourceType); - } - - private List getAllTasksInSandbox(String sandboxId) { - return sandboxLevelViews.get(sandboxId).getActiveTasks(); - } - - private TaskCancellation createTaskCancellation(CancellableTask task) { - // todo add correct reason and callbacks - return new TaskCancellation(task, List.of(new TaskCancellation.Reason("limits exceeded", 5)), List.of(this::callbackOnCancel)); - } - - private void callbackOnCancel() { - // todo Implement callback logic here - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java index 5767975d3497e..242bfdbefa0b0 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskSelectionStrategy.java @@ -8,8 +8,10 @@ package org.opensearch.search.sandboxing.cancellation; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.resourcetype.SystemResource; +import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskCancellation; import java.util.ArrayList; import java.util.Collections; @@ -43,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, SandboxResourceType resourceType) { + public List selectTasksForCancellation(List tasks, long limit, SystemResource resourceType) { if (limit < 0) { throw new IllegalArgumentException("reduceBy has to be greater than zero"); } @@ -53,16 +55,27 @@ public List selectTasksForCancellation(List tasks, long limit, Sandb List sortedTasks = tasks.stream().sorted(sortingCondition()).collect(Collectors.toList()); - List selectedTasks = new ArrayList<>(); + List selectedTasks = new ArrayList<>(); long accumulated = 0; for (Task task : sortedTasks) { - selectedTasks.add(task); - accumulated += resourceType.getResourceUsage(task); - if (accumulated >= limit) { - break; + if (task instanceof CancellableTask) { + selectedTasks.add(createTaskCancellation((CancellableTask) task)); + accumulated += resourceType.getResourceUsage(task); + if (accumulated >= limit) { + break; + } } } return selectedTasks; } + + private TaskCancellation createTaskCancellation(CancellableTask task) { + // todo add correct reason and callbacks + return new TaskCancellation(task, List.of(new TaskCancellation.Reason("limits exceeded", 5)), List.of(this::callbackOnCancel)); + } + + private void callbackOnCancel() { + // todo Implement callback logic here + } } 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 index 0e351ab2e90ea..73901df716ada 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellation.java @@ -10,40 +10,75 @@ import org.opensearch.cluster.metadata.Sandbox; import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +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; /** - * Represents the default task cancellation strategy. - * This class extends the AbstractTaskCancellation and provides a method to get the sandboxes from which tasks should be cancelled. + * 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 extends AbstractTaskCancellation { +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 cancellationStrategy, + TaskSelectionStrategy taskSelectionStrategy, Map sandboxLevelViews, Set activeSandboxes ) { - super(cancellationStrategy, sandboxLevelViews, 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(); + } } /** - * // TODO create another Task cancellation implementation that does below - * - if node not in duress - * - pick sandboxes in enforced mode only - * - if node in duress - * - pick sandboxes in enforced mode - * - tasks running in deleted sandboxes with tasks running - * - pick sandboxes in enforced mode + * 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()); + 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)) { @@ -56,6 +91,59 @@ public List getSandboxesToCancelFrom() { 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. * @@ -63,7 +151,7 @@ public List getSandboxesToCancelFrom() { * @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) { + private boolean isBreachingThreshold(Map currentResourceUsage, Sandbox.ResourceLimit resourceLimit) { return currentResourceUsage.get(resourceLimit.getResourceType()) > resourceLimit.getThreshold(); } @@ -73,7 +161,10 @@ private boolean isBreachingThreshold(Map currentResou * @param sandboxId The ID of the sandbox * @return The resource usage of the sandbox */ - private Map getResourceUsage(String sandboxId) { + 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/cancellation/LongestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java index 67f3ed18252d4..eacdf7a4c15f4 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategy.java @@ -24,6 +24,6 @@ public class LongestRunningTaskFirstStrategy extends AbstractTaskSelectionStrate */ @Override public Comparator sortingCondition() { - return Comparator.comparingLong(Task::getStartTime).reversed(); + return Comparator.comparingLong(Task::getStartTime); } } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/MostResourceIntensiveTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/MostResourceIntensiveTaskFirstStrategy.java new file mode 100644 index 0000000000000..eeaa03d99cec8 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/MostResourceIntensiveTaskFirstStrategy.java @@ -0,0 +1,27 @@ +/* + * 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.search.sandboxing.resourcetype.SystemResource; +import org.opensearch.tasks.Task; + +import java.util.Comparator; + +public class MostResourceIntensiveTaskFirstStrategy extends AbstractTaskSelectionStrategy { + private final SystemResource resourceType; + + MostResourceIntensiveTaskFirstStrategy(SystemResource resourceType) { + this.resourceType = resourceType; + } + + @Override + public Comparator sortingCondition() { + return Comparator.comparingLong(resourceType::getResourceUsage).reversed(); + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java index 2f7acb43512d8..af186b01a8974 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategy.java @@ -24,6 +24,6 @@ public class ShortestRunningTaskFirstStrategy extends AbstractTaskSelectionStrat */ @Override public Comparator sortingCondition() { - return Comparator.comparingLong(Task::getStartTime); + return Comparator.comparingLong(Task::getStartTime).reversed(); } } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskCancellationContext.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskCancellationContext.java new file mode 100644 index 0000000000000..95129d71c153a --- /dev/null +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskCancellationContext.java @@ -0,0 +1,35 @@ +/* + * 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 java.util.Set; + +// TODO - can be deleted +public class TaskCancellationContext { + + // make this a sandbox level view where each view wil have both list of tasks & resource usage map + private final Set sandboxLevelViews; + private final Set activeSandboxes; + + public TaskCancellationContext(Set sandboxLevelViews, Set activeSandboxes) { + this.sandboxLevelViews = sandboxLevelViews; + this.activeSandboxes = activeSandboxes; + } + + public Set getSandboxLevelViews() { + return sandboxLevelViews; + } + + public Set getActiveSandboxes() { + return activeSandboxes; + } +} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java index b5345be6dda8b..61f1883aefaa5 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategy.java @@ -8,8 +8,9 @@ package org.opensearch.search.sandboxing.cancellation; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.resourcetype.SystemResource; import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskCancellation; import java.util.List; @@ -27,5 +28,5 @@ public interface TaskSelectionStrategy { * * @return List of tasks that should be cancelled. */ - List selectTasksForCancellation(List tasks, long limit, SandboxResourceType resourceType); + List selectTasksForCancellation(List tasks, long limit, SystemResource resourceType); } diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResource.java similarity index 87% rename from server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java rename to server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResource.java index 675bfb0d38ff5..808a76a9bc66a 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/CpuTimeResource.java @@ -13,7 +13,7 @@ /** * Represents the CPU time resource type. */ -public class CpuTimeResourceType extends SandboxResourceType { +public class CpuTimeResource extends SystemResource { /** * Returns the CPU time usage of the provided task. * @@ -27,7 +27,7 @@ public long getResourceUsage(Task task) { @Override public boolean equals(Object obj) { - return obj instanceof CpuTimeResourceType; + return obj instanceof CpuTimeResource; } @Override diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResource.java similarity index 87% rename from server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java rename to server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResource.java index b5408de9c61b8..2387954113a74 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/JvmMemoryResource.java @@ -13,7 +13,7 @@ /** * Represents the JVM memory resource type. */ -public class JvmMemoryResourceType extends SandboxResourceType { +public class JvmMemoryResource extends SystemResource { /** * Returns the memory usage of the provided task. * @@ -27,7 +27,7 @@ public long getResourceUsage(Task task) { @Override public boolean equals(Object obj) { - return obj instanceof JvmMemoryResourceType; + return obj instanceof JvmMemoryResource; } @Override diff --git a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SystemResource.java similarity index 71% rename from server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java rename to server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SystemResource.java index 1b46a2b25197e..525861a3356e9 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SandboxResourceType.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/resourcetype/SystemResource.java @@ -16,7 +16,7 @@ * This class is abstract and requires the implementation of the getResourceUsage method. */ @ExperimentalApi -public abstract class SandboxResourceType { +public abstract class SystemResource { /** * Returns the resource usage of the provided task. * The specific resource that this method returns depends on the implementation. @@ -27,20 +27,20 @@ public abstract class SandboxResourceType { public abstract long getResourceUsage(Task task); /** - * Creates a SandboxResourceType from a string. - * If the string is "JVM", a JvmMemoryResourceType is returned. - * If the string is "CPU", a CpuTimeResourceType is returned. + * 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 SandboxResourceType - * @return The created SandboxResourceType + * @param type The string from which to create a SystemResource + * @return The created SystemResource * @throws IllegalArgumentException If the string is not recognized */ - public static SandboxResourceType fromString(String type) { + public static SystemResource fromString(String type) { if (type.equalsIgnoreCase("JVM")) { - return new JvmMemoryResourceType(); + return new JvmMemoryResource(); } else if (type.equalsIgnoreCase("CPU")) { - return new CpuTimeResourceType(); + return new CpuTimeResource(); } else { throw new IllegalArgumentException("Unsupported resource type: " + type); } 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 index c51ef7f105efb..9dcb5e5006475 100644 --- a/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/sandboxing/tracker/SandboxResourceUsageTrackerService.java @@ -11,7 +11,7 @@ import org.opensearch.common.inject.Inject; import org.opensearch.search.sandboxing.SandboxLevelResourceUsageView; import org.opensearch.search.sandboxing.SandboxTask; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.resourcetype.SystemResource; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; import org.opensearch.tasks.TaskResourceTrackingService; @@ -27,9 +27,9 @@ // @ExperimentalApi public class SandboxResourceUsageTrackerService implements SandboxUsageTracker, TaskManager.TaskEventListeners { - public static final List TRACKED_RESOURCES = List.of( - SandboxResourceType.fromString("JVM"), - SandboxResourceType.fromString("CPU") + public static final List TRACKED_RESOURCES = List.of( + SystemResource.fromString("JVM"), + SystemResource.fromString("CPU") ); private final TaskManager taskManager; @@ -60,7 +60,7 @@ public Map constructSandboxLevelUsageView Map sandboxViews = new HashMap<>(); Map> tasksBySandbox = getTasksGroupedBySandbox(); - Map> sandboxResourceUsage = getResourceUsageOfSandboxes(tasksBySandbox); + Map> sandboxResourceUsage = getResourceUsageOfSandboxes(tasksBySandbox); for (String sandboxId : tasksBySandbox.keySet()) { SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView( @@ -93,8 +93,8 @@ private Map> getTasksGroupedBySandbox() { * @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<>(); + private Map> getResourceUsageOfSandboxes(Map> tasksBySandbox) { + Map> resourceUsageOfSandboxes = new HashMap<>(); // Iterate over each sandbox entry for (Map.Entry> sandboxEntry : tasksBySandbox.entrySet()) { @@ -102,11 +102,11 @@ private Map> getResourceUsageOfSandboxes( 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<>()); + Map sandboxUsage = resourceUsageOfSandboxes.computeIfAbsent(sandboxId, k -> new HashMap<>()); // Accumulate resource usage for each task in the sandbox for (Task task : tasks) { - for (SandboxResourceType resourceType : TRACKED_RESOURCES) { + for (SystemResource resourceType : TRACKED_RESOURCES) { long currentUsage = sandboxUsage.getOrDefault(resourceType, 0L); long taskUsage = resourceType.getResourceUsage(task); sandboxUsage.put(resourceType, currentUsage + taskUsage); diff --git a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java index 2376db7eda1f9..17e987d608d3e 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/SandboxLevelResourceUsageViewTests.java @@ -8,7 +8,7 @@ package org.opensearch.search.sandboxing; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.resourcetype.SystemResource; import org.opensearch.tasks.Task; import org.opensearch.test.OpenSearchTestCase; @@ -18,24 +18,24 @@ import static org.opensearch.search.sandboxing.cancellation.SandboxTestHelpers.getRandomTask; public class SandboxLevelResourceUsageViewTests extends OpenSearchTestCase { - Map resourceUsage; + Map resourceUsage; List activeTasks; public void setUp() throws Exception { super.setUp(); - resourceUsage = Map.of(SandboxResourceType.fromString("JVM"), 34L, SandboxResourceType.fromString("CPU"), 12L); + 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(); + Map resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData(); assertTrue(assertResourceUsageData(resourceUsageData)); } public void testGetResourceUsageDataDefault() { SandboxLevelResourceUsageView sandboxLevelResourceUsageView = new SandboxLevelResourceUsageView("1234"); - Map resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData(); + Map resourceUsageData = sandboxLevelResourceUsageView.getResourceUsageData(); assertTrue(resourceUsageData.isEmpty()); } @@ -52,8 +52,8 @@ public void testGetActiveTasksDefault() { assertTrue(activeTasks.isEmpty()); } - private boolean assertResourceUsageData(Map resourceUsageData) { - return resourceUsageData.get(SandboxResourceType.fromString("JVM")) == 34L - && resourceUsageData.get(SandboxResourceType.fromString("CPU")) == 12L; + 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/AbstractTaskCancellationTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellationTests.java deleted file mode 100644 index f59f6ea0aa65f..0000000000000 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/AbstractTaskCancellationTests.java +++ /dev/null @@ -1,159 +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.SandboxResourceType; -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 AbstractTaskCancellationTests extends OpenSearchTestCase { - - private class TestTaskCancellationImpl extends AbstractTaskCancellation { - - public TestTaskCancellationImpl( - TaskSelectionStrategy taskSelectionStrategy, - Map sandboxLevelViews, - Set activeSandboxes - ) { - super(taskSelectionStrategy, sandboxLevelViews, activeSandboxes); - } - - @Override - List getSandboxesToCancelFrom() { - return new ArrayList<>(activeSandboxes); - } - } - - private TaskSelectionStrategy taskSelectionStrategy; - private Map sandboxLevelViews; - private Set activeSandboxes; - private AbstractTaskCancellation 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(SandboxResourceType.fromString(resourceTypeStr))).thenReturn(resourceLimitMock); - SandboxLevelResourceUsageView mockView = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView); - - List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(sandbox1); - assertEquals(2, cancellableTasksFrom.size()); - assertEquals(4321, cancellableTasksFrom.get(0).getTask().getId()); - assertEquals(1234, 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(SandboxResourceType.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(SandboxResourceType.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(4321, cancellableTasksFrom.get(0).getTask().getId()); - assertEquals(1234, 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(SandboxResourceType.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(SandboxResourceType.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(4321, allCancellableTasks.get(0).getTask().getId()); - assertEquals(1234, allCancellableTasks.get(1).getTask().getId()); - } -} 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 index a7f144526a5ed..2d1233eb83286 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java @@ -10,150 +10,171 @@ 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.Collections; +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.Mock; import org.mockito.MockitoAnnotations; +import static org.opensearch.search.sandboxing.cancellation.SandboxTestHelpers.createResourceLimitMock; +import static org.mockito.Mockito.when; + public class DefaultTaskCancellationTests extends OpenSearchTestCase { - @Mock - private TaskSelectionStrategy mockStrategy; - @Mock - private SandboxLevelResourceUsageView mockView; + private 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 testConstructor() { - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - assertNotNull(cancellation); + 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(4321, cancellableTasksFrom.get(0).getTask().getId()); + assertEquals(1234, cancellableTasksFrom.get(1).getTask().getId()); } - public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForDifferentResourceTypes() { - // setup mocks for sandbox1 + public void testGetCancellableTasksFrom_returnsNoTasksWhenBreachingThreshold() { String id = "sandbox1"; String resourceTypeStr = "CPU"; long usage = 50L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + 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); - // setup mocks for sandbox2 - id = "sandbox2"; - resourceTypeStr = "JVM"; - usage = 50L; - Sandbox sandbox2 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 100L, usage); - SandboxLevelResourceUsageView mockView2 = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView2); - // add both sandboxes to active sandboxes - Collections.addAll(activeSandboxes, sandbox1, sandbox2); - // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - List result = cancellation.getSandboxesToCancelFrom(); - - // only sandbox1 should be returned as it is breaching the threshold - assertEquals(1, result.size()); - assertTrue(result.contains(sandbox1)); - assertFalse(result.contains(sandbox2)); + List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(sandbox1); + assertTrue(cancellableTasksFrom.isEmpty()); } - public void testGetSandboxesToCancelFrom_whenNotAllSandboxesAreBreachingForSameResourceType() { - // setup mocks for sandbox1 + public void testCancelTasks_cancelsGivenTasks() { String id = "sandbox1"; String resourceTypeStr = "CPU"; long usage = 50L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + 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); - - // setup mocks for sandbox2 - id = "sandbox2"; - resourceTypeStr = "CPU"; - usage = 50L; - Sandbox sandbox2 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 100L, usage); - SandboxLevelResourceUsageView mockView2 = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView2); - // add both sandboxes to active sandboxes - Collections.addAll(activeSandboxes, sandbox1, sandbox2); - // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - List result = cancellation.getSandboxesToCancelFrom(); - - // only sandbox1 should be returned as it is breaching the threshold - assertEquals(1, result.size()); - assertTrue(result.contains(sandbox1)); - assertFalse(result.contains(sandbox2)); + activeSandboxes.add(sandbox1); + + TestTaskCancellationImpl taskCancellation = new TestTaskCancellationImpl( + new TaskSelectionStrategyTests.TestTaskSelectionStrategy(), + sandboxLevelViews, + activeSandboxes + ); + + List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(); + assertEquals(2, cancellableTasksFrom.size()); + assertEquals(4321, cancellableTasksFrom.get(0).getTask().getId()); + assertEquals(1234, cancellableTasksFrom.get(1).getTask().getId()); + + taskCancellation.cancelTasks(); + assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); + assertTrue(cancellableTasksFrom.get(1).getTask().isCancelled()); } - public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForDifferentResourceTypes() { - // setup mocks for sandbox1 + public void testGetAllCancellableTasks_ReturnsNoTasksWhenNotBreachingThresholds() { String id = "sandbox1"; String resourceTypeStr = "CPU"; long usage = 50L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + 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); - // setup mocks for sandbox2 - id = "sandbox2"; - resourceTypeStr = "JVM"; - usage = 50L; - Sandbox sandbox2 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); - SandboxLevelResourceUsageView mockView2 = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView2); - // add both sandboxes to active sandboxes - Collections.addAll(activeSandboxes, sandbox1, sandbox2); - // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - List result = cancellation.getSandboxesToCancelFrom(); - - // Both sandboxes should be returned as it is breaching the threshold - assertEquals(2, result.size()); - assertTrue(result.contains(sandbox1)); - assertTrue(result.contains(sandbox2)); + List allCancellableTasks = taskCancellation.getAllCancellableTasks(); + assertTrue(allCancellableTasks.isEmpty()); } - public void testGetSandboxesToCancelFrom_whenAllSandboxesAreBreachingForSameResourceType() { - // setup mocks for sandbox1 + public void testGetAllCancellableTasks_ReturnsTasksWhenBreachingThresholds() { String id = "sandbox1"; String resourceTypeStr = "CPU"; - long usage = 50L; - Sandbox sandbox1 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); + 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(4321, allCancellableTasks.get(0).getTask().getId()); + assertEquals(1234, allCancellableTasks.get(1).getTask().getId()); + } - // setup mocks for sandbox2 - id = "sandbox2"; - resourceTypeStr = "CPU"; - usage = 50L; - Sandbox sandbox2 = SandboxTestHelpers.createSandboxMock(id, resourceTypeStr, 10L, usage); - SandboxLevelResourceUsageView mockView2 = SandboxTestHelpers.createResourceUsageViewMock(resourceTypeStr, usage); - sandboxLevelViews.put(id, mockView2); - // add both sandboxes to active sandboxes - Collections.addAll(activeSandboxes, sandbox1, sandbox2); - // create a new instance of DefaultTaskCancellation and call getSandboxesToCancelFrom - DefaultTaskCancellation cancellation = new DefaultTaskCancellation(mockStrategy, sandboxLevelViews, activeSandboxes); - List result = cancellation.getSandboxesToCancelFrom(); - - // Both sandboxes should be returned as it is breaching the threshold - assertEquals(2, result.size()); - assertTrue(result.contains(sandbox1)); - assertTrue(result.contains(sandbox2)); + 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(4321, cancellableTasksFrom.get(0).getTask().getId()); + assertEquals(1234, 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/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java index 4c5c0aeb4fd88..e06d0d6e20128 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java @@ -26,9 +26,9 @@ public void testSortingCondition() { when(task2.getStartTime()).thenReturn(200L); when(task3.getStartTime()).thenReturn(300L); - List tasks = Arrays.asList(task1, task3, task2); + List tasks = Arrays.asList(task2, task1, task3); tasks.sort(new LongestRunningTaskFirstStrategy().sortingCondition()); - assertEquals(Arrays.asList(task3, task2, task1), tasks); + assertEquals(Arrays.asList(task1, task2, task3), tasks); } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java index df29466cf1f5c..6c83548ae0b3a 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/SandboxTestHelpers.java @@ -16,12 +16,13 @@ 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.SandboxResourceType; +import org.opensearch.search.sandboxing.resourcetype.SystemResource; import org.opensearch.tasks.Task; import java.util.ArrayList; import java.util.Collections; import java.util.List; +import java.util.function.Supplier; import org.mockito.Mockito; @@ -37,7 +38,7 @@ public static List getListOfTasks(long totalMemory) { while (totalMemory > 0) { long id = randomLong(); - final Task task = getRandomTask(id); + final Task task = getRandomSearchTask(id); long initial_memory = randomLongBetween(1, 100); ResourceUsageMetric[] initialTaskResourceMetrics = new ResourceUsageMetric[] { @@ -92,7 +93,7 @@ public static Sandbox createSandboxMock(String id, String resourceTypeStr, Long public static Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeStr, Long threshold) { Sandbox.ResourceLimit resourceLimitMock = mock(Sandbox.ResourceLimit.class); - SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); + SystemResource resourceType = SystemResource.fromString(resourceTypeStr); when(resourceLimitMock.getResourceType()).thenReturn(resourceType); when(resourceLimitMock.getThreshold()).thenReturn(threshold); when(resourceLimitMock.getThresholdInLong()).thenReturn(threshold); @@ -101,9 +102,13 @@ public static Sandbox.ResourceLimit createResourceLimitMock(String resourceTypeS public static SandboxLevelResourceUsageView createResourceUsageViewMock(String resourceTypeStr, Long usage) { SandboxLevelResourceUsageView mockView = mock(SandboxLevelResourceUsageView.class); - SandboxResourceType resourceType = SandboxResourceType.fromString(resourceTypeStr); + SystemResource resourceType = SystemResource.fromString(resourceTypeStr); when(mockView.getResourceUsageData()).thenReturn(Collections.singletonMap(resourceType, usage)); when(mockView.getActiveTasks()).thenReturn(List.of(getRandomSearchTask(1234), getRandomSearchTask(4321))); return mockView; } + + private static Supplier descriptionSupplier() { + return () -> "test description"; + } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java index 7936e940e9cc9..fe89263a09c3c 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java @@ -9,8 +9,9 @@ package org.opensearch.search.sandboxing.cancellation; import org.opensearch.core.tasks.resourcetracker.ResourceStats; -import org.opensearch.search.sandboxing.resourcetype.SandboxResourceType; +import org.opensearch.search.sandboxing.resourcetype.SystemResource; import org.opensearch.tasks.Task; +import org.opensearch.tasks.TaskCancellation; import org.opensearch.test.OpenSearchTestCase; import java.util.Comparator; @@ -29,10 +30,10 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsGrea TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); long threshold = 100L; long reduceBy = 50L; - SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); + SystemResource resourceType = SystemResource.fromString("JVM"); List tasks = SandboxTestHelpers.getListOfTasks(threshold); - List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); + List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); assertFalse(selectedTasks.isEmpty()); assertTrue(tasksUsageMeetsThreshold(selectedTasks, reduceBy)); } @@ -41,7 +42,7 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsLess TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); long threshold = 100L; long reduceBy = -50L; - SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); + SystemResource resourceType = SystemResource.fromString("JVM"); List tasks = SandboxTestHelpers.getListOfTasks(threshold); try { @@ -56,17 +57,17 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsEqua TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); long threshold = 100L; long reduceBy = 0; - SandboxResourceType resourceType = SandboxResourceType.fromString("JVM"); + SystemResource resourceType = SystemResource.fromString("JVM"); List tasks = SandboxTestHelpers.getListOfTasks(threshold); - List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); + List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); assertTrue(selectedTasks.isEmpty()); } - private boolean tasksUsageMeetsThreshold(List selectedTasks, long threshold) { + private boolean tasksUsageMeetsThreshold(List selectedTasks, long threshold) { long memory = 0; - for (Task task : selectedTasks) { - memory += task.getTotalResourceUtilization(ResourceStats.MEMORY); + for (TaskCancellation task : selectedTasks) { + memory += task.getTask().getTotalResourceUtilization(ResourceStats.MEMORY); if (memory > threshold) { return true; } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java index 1120122f33f01..7a203d6afe9f2 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/tracking/SandboxResourceUsageTrackerServiceTests.java @@ -14,7 +14,7 @@ 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.SandboxResourceType; +import org.opensearch.search.sandboxing.resourcetype.SystemResource; import org.opensearch.search.sandboxing.tracker.SandboxResourceUsageTrackerService; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; @@ -68,9 +68,7 @@ public void testConstructSandboxLevelViews_CreatesSandboxLevelUsageView_WhenTask for (String sandboxId : sandboxIds) { assertEquals( 400, - (long) stringSandboxLevelResourceUsageViewMap.get(sandboxId) - .getResourceUsageData() - .get(SandboxResourceType.fromString("JVM")) + (long) stringSandboxLevelResourceUsageViewMap.get(sandboxId).getResourceUsageData().get(SystemResource.fromString("JVM")) ); assertEquals(2, stringSandboxLevelResourceUsageViewMap.get(sandboxId).getActiveTasks().size()); } @@ -90,7 +88,7 @@ public void testConstructSandboxLevelUsageViews_WithTasksHavingDifferentResource Map sandboxViews = sandboxResourceUsageTrackerService.constructSandboxLevelUsageViews(); - assertEquals(600, (long) sandboxViews.get("sandbox1").getResourceUsageData().get(SandboxResourceType.fromString("JVM"))); + assertEquals(600, (long) sandboxViews.get("sandbox1").getResourceUsageData().get(SystemResource.fromString("JVM"))); assertEquals(2, sandboxViews.get("sandbox1").getActiveTasks().size()); } From 2962437cbd7cdb46f05b9b4e1d5dae5e66e1c5ae Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 3 Jul 2024 16:53:46 -0700 Subject: [PATCH 20/66] revert some accidentally pushed files Signed-off-by: Kiran Prakash --- .../java/Sandbox/SandboxMode/ENFORCED.java | 11 - .../sandboxing/DefaultSandboxPruner.java | 66 ------ .../search/sandboxing/SandboxPruner.java | 21 -- .../sandboxing/SandboxServiceSettings.java | 198 ------------------ .../search/sandboxing/TaskData.java | 28 --- ...ostResourceIntensiveTaskFirstStrategy.java | 27 --- .../cancellation/TaskCancellationContext.java | 35 ---- 7 files changed, 386 deletions(-) delete mode 100644 server/src/main/java/Sandbox/SandboxMode/ENFORCED.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/DefaultSandboxPruner.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/SandboxPruner.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/SandboxServiceSettings.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/TaskData.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/MostResourceIntensiveTaskFirstStrategy.java delete mode 100644 server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskCancellationContext.java diff --git a/server/src/main/java/Sandbox/SandboxMode/ENFORCED.java b/server/src/main/java/Sandbox/SandboxMode/ENFORCED.java deleted file mode 100644 index ab8c2ee2cb60a..0000000000000 --- a/server/src/main/java/Sandbox/SandboxMode/ENFORCED.java +++ /dev/null @@ -1,11 +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 Sandbox.SandboxMode; - -public class ENFORCED {} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/DefaultSandboxPruner.java b/server/src/main/java/org/opensearch/search/sandboxing/DefaultSandboxPruner.java deleted file mode 100644 index 471c71b2cb205..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/DefaultSandboxPruner.java +++ /dev/null @@ -1,66 +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.cluster.metadata.Sandbox; - -import java.util.List; -import java.util.stream.Collectors; - -public class DefaultSandboxPruner implements SandboxPruner { - private List activeSandboxes; - private List toDeleteSandboxes; - private TaskData taskData; - - @Override - public void pruneSandboxes() { - toDeleteSandboxes = toDeleteSandboxes.stream().filter(this::hasUnfinishedTasks).collect(Collectors.toList()); - } - - @Override - public void deleteSandbox(String sandboxId) { - // should be called from the API and we need to delete from cluster metadata - // TODO may be we need to build a listener interface instead - if (hasUnfinishedTasks(sandboxId)) { - toDeleteSandboxes.add(sandboxId); - } - // remove this sandbox from the active sandboxes - activeSandboxes = activeSandboxes.stream() - .filter(resourceLimitGroup -> !resourceLimitGroup.getName().equals(sandboxId)) - .collect(Collectors.toList()); - } - - private boolean hasUnfinishedTasks(String sandboxName) { - return !taskData.getTasksBySandbox().get(sandboxName).isEmpty(); - } - - public List getActiveSandboxes() { - return activeSandboxes; - } - - public void setActiveSandboxes(List activeSandboxes) { - this.activeSandboxes = activeSandboxes; - } - - public List getToDeleteSandboxes() { - return toDeleteSandboxes; - } - - public void setToDeleteSandboxes(List toDeleteSandboxes) { - this.toDeleteSandboxes = toDeleteSandboxes; - } - - public TaskData getTaskData() { - return taskData; - } - - public void setTaskData(TaskData taskData) { - this.taskData = taskData; - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxPruner.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxPruner.java deleted file mode 100644 index 40d7039333b85..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxPruner.java +++ /dev/null @@ -1,21 +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; - -/** - * This interface is used to identify and completely remove deleted resourceLimitGroups which has been marked as deleted - * previously but had the tasks running at the time of deletion request - */ -public interface SandboxPruner { - // remove the deleted sandboxes from the system once all the tasks in that sandbox are completed/cancelled - void pruneSandboxes(); - - // accepts the request to delete the sandbox and marks the sandbox as deleted - void deleteSandbox(String sandboxId); -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/SandboxServiceSettings.java b/server/src/main/java/org/opensearch/search/sandboxing/SandboxServiceSettings.java deleted file mode 100644 index d31d905c4503d..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/SandboxServiceSettings.java +++ /dev/null @@ -1,198 +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.common.settings.ClusterSettings; -import org.opensearch.common.settings.Setting; -import org.opensearch.common.settings.Settings; -import org.opensearch.common.unit.TimeValue; - -/** - * Main class to declare the query ResourceLimitGrouping feature related settings - */ -public class SandboxServiceSettings { - private static final Long DEFAULT_RUN_INTERVAL_MILLIS = 1000L; - private static final Double DEFAULT_NODE_LEVEL_REJECTION_THRESHOLD = 0.8; - private static final Double DEFAULT_NODE_LEVEL_CANCELLATION_THRESHOLD = 0.9; - /** - * default max ResourceLimitGroup count on any node at any given point in time - */ - public static final int DEFAULT_MAX_RESOURCE_LIMIT_GROUP_COUNT_VALUE = 100; - - public static final String RESOURCE_LIMIT_GROUP_COUNT_SETTING_NAME = "node.resource_limit_group.max_count"; - public static final double NODE_LEVEL_CANCELLATION_THRESHOLD_MAX_VALUE = 0.95; - public static final double NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE = 0.90; - - private TimeValue runIntervalMillis; - private Double nodeLevelJvmCancellationThreshold; - private Double nodeLevelJvmRejectionThreshold; - private volatile int maxResourceLimitGroupCount; - /** - * max ResourceLimitGroup count setting - */ - public static final Setting MAX_RESOURCE_LIMIT_GROUP_COUNT = Setting.intSetting( - RESOURCE_LIMIT_GROUP_COUNT_SETTING_NAME, - DEFAULT_MAX_RESOURCE_LIMIT_GROUP_COUNT_VALUE, - 0, - (newVal) -> { - if (newVal > 100 || newVal < 1) throw new IllegalArgumentException( - RESOURCE_LIMIT_GROUP_COUNT_SETTING_NAME + " should be in range [1-100]" - ); - }, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - /** - * Setting name for default ResourceLimitGroup count - */ - public static final String SERVICE_RUN_INTERVAL_MILLIS_SETTING_NAME = "resource_limit_group.service.run_interval_millis"; - /** - * Setting to control the run interval of QSB service - */ - private static final Setting RESOURCE_LIMIT_GROUP_RUN_INTERVAL_SETTING = Setting.longSetting( - SERVICE_RUN_INTERVAL_MILLIS_SETTING_NAME, - DEFAULT_RUN_INTERVAL_MILLIS, - 1, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - - /** - * Setting name for node level rejection threshold for QSB - */ - public static final String NODE_REJECTION_THRESHOLD_SETTING_NAME = "resource_limit_group.node.rejection_threshold"; - /** - * Setting to control the rejection threshold - */ - public static final Setting NODE_LEVEL_REJECTION_THRESHOLD = Setting.doubleSetting( - NODE_REJECTION_THRESHOLD_SETTING_NAME, - DEFAULT_NODE_LEVEL_REJECTION_THRESHOLD, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - /** - * Setting name for node level cancellation threshold - */ - public static final String NODE_CANCELLATION_THRESHOLD_SETTING_NAME = "resource_limit_group.node.cancellation_threshold"; - /** - * Setting name for node level cancellation threshold - */ - public static final Setting NODE_LEVEL_CANCELLATION_THRESHOLD = Setting.doubleSetting( - NODE_CANCELLATION_THRESHOLD_SETTING_NAME, - DEFAULT_NODE_LEVEL_CANCELLATION_THRESHOLD, - Setting.Property.Dynamic, - Setting.Property.NodeScope - ); - - /** - * ResourceLimitGroup service settings constructor - * @param settings - * @param clusterSettings - */ - public SandboxServiceSettings(Settings settings, ClusterSettings clusterSettings) { - runIntervalMillis = new TimeValue(RESOURCE_LIMIT_GROUP_RUN_INTERVAL_SETTING.get(settings)); - nodeLevelJvmCancellationThreshold = NODE_LEVEL_CANCELLATION_THRESHOLD.get(settings); - nodeLevelJvmRejectionThreshold = NODE_LEVEL_REJECTION_THRESHOLD.get(settings); - maxResourceLimitGroupCount = MAX_RESOURCE_LIMIT_GROUP_COUNT.get(settings); - - ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); - - clusterSettings.addSettingsUpdateConsumer(MAX_RESOURCE_LIMIT_GROUP_COUNT, this::setMaxResourceLimitGroupCount); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_CANCELLATION_THRESHOLD, this::setNodeLevelJvmCancellationThreshold); - clusterSettings.addSettingsUpdateConsumer(NODE_LEVEL_REJECTION_THRESHOLD, this::setNodeLevelJvmRejectionThreshold); - } - - /** - * Method to get runInterval for QSB - * @return runInterval in milliseconds for QSB Service - */ - public TimeValue getRunIntervalMillis() { - return runIntervalMillis; - } - - /** - * Method to set the new ResourceLimitGroup count - * @param newMaxResourceLimitGroupCount is the new maxResourceLimitGroupCount per node - */ - public void setMaxResourceLimitGroupCount(int newMaxResourceLimitGroupCount) { - if (newMaxResourceLimitGroupCount < 0) { - throw new IllegalArgumentException("node.ResourceLimitGroup.max_count can't be negative"); - } - this.maxResourceLimitGroupCount = newMaxResourceLimitGroupCount; - } - - /** - * Method to get the node level cancellation threshold - * @return current node level cancellation threshold - */ - public Double getNodeLevelJvmCancellationThreshold() { - return nodeLevelJvmCancellationThreshold; - } - - /** - * Method to set the node level cancellation threshold - * @param nodeLevelJvmCancellationThreshold sets the new node level cancellation threshold - * @throws IllegalArgumentException if the value is > 0.95 and cancellation < rejection threshold - */ - public void setNodeLevelJvmCancellationThreshold(Double nodeLevelJvmCancellationThreshold) { - if (Double.compare(nodeLevelJvmCancellationThreshold, NODE_LEVEL_CANCELLATION_THRESHOLD_MAX_VALUE) > 0) { - throw new IllegalArgumentException( - NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be greater than 0.95 as it pose a threat of node drop" - ); - } - - ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); - - this.nodeLevelJvmCancellationThreshold = nodeLevelJvmCancellationThreshold; - } - - /** - * Method to get the node level rejection threshold - * @return the current node level rejection threshold - */ - public Double getNodeLevelJvmRejectionThreshold() { - return nodeLevelJvmRejectionThreshold; - } - - /** - * Method to set the node level rejection threshold - * @param nodeLevelJvmRejectionThreshold sets the new rejection threshold - * @throws IllegalArgumentException if rejection > 0.90 and rejection < cancellation threshold - */ - public void setNodeLevelJvmRejectionThreshold(Double nodeLevelJvmRejectionThreshold) { - if (Double.compare(nodeLevelJvmRejectionThreshold, NODE_LEVEL_REJECTION_THRESHOLD_MAX_VALUE) > 0) { - throw new IllegalArgumentException( - NODE_REJECTION_THRESHOLD_SETTING_NAME + " value not be greater than 0.90 as it pose a threat of node drop" - ); - } - - ensureRejectionThresholdIsLessThanCancellation(nodeLevelJvmRejectionThreshold, nodeLevelJvmCancellationThreshold); - - this.nodeLevelJvmRejectionThreshold = nodeLevelJvmRejectionThreshold; - } - - private void ensureRejectionThresholdIsLessThanCancellation( - Double nodeLevelJvmRejectionThreshold, - Double nodeLevelJvmCancellationThreshold - ) { - if (Double.compare(nodeLevelJvmCancellationThreshold, nodeLevelJvmRejectionThreshold) < 0) { - throw new IllegalArgumentException( - NODE_CANCELLATION_THRESHOLD_SETTING_NAME + " value should not be less than " + NODE_REJECTION_THRESHOLD_SETTING_NAME - ); - } - } - - /** - * Method to get the current ResourceLimitGroup count - * @return the current max ResourceLimitGroup count - */ - public int getMaxResourceLimitGroupCount() { - return maxResourceLimitGroupCount; - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/TaskData.java b/server/src/main/java/org/opensearch/search/sandboxing/TaskData.java deleted file mode 100644 index 181867a8a7dbf..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/TaskData.java +++ /dev/null @@ -1,28 +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.common.annotation.ExperimentalApi; -import org.opensearch.tasks.Task; - -import java.util.List; -import java.util.Map; - -@ExperimentalApi -public class TaskData { - private final Map> tasksBySandboxes; - - public TaskData(Map> tasksBySandboxes) { - this.tasksBySandboxes = tasksBySandboxes; - } - - public Map> getTasksBySandbox() { - return tasksBySandboxes; - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/MostResourceIntensiveTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/MostResourceIntensiveTaskFirstStrategy.java deleted file mode 100644 index eeaa03d99cec8..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/MostResourceIntensiveTaskFirstStrategy.java +++ /dev/null @@ -1,27 +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.search.sandboxing.resourcetype.SystemResource; -import org.opensearch.tasks.Task; - -import java.util.Comparator; - -public class MostResourceIntensiveTaskFirstStrategy extends AbstractTaskSelectionStrategy { - private final SystemResource resourceType; - - MostResourceIntensiveTaskFirstStrategy(SystemResource resourceType) { - this.resourceType = resourceType; - } - - @Override - public Comparator sortingCondition() { - return Comparator.comparingLong(resourceType::getResourceUsage).reversed(); - } -} diff --git a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskCancellationContext.java b/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskCancellationContext.java deleted file mode 100644 index 95129d71c153a..0000000000000 --- a/server/src/main/java/org/opensearch/search/sandboxing/cancellation/TaskCancellationContext.java +++ /dev/null @@ -1,35 +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 java.util.Set; - -// TODO - can be deleted -public class TaskCancellationContext { - - // make this a sandbox level view where each view wil have both list of tasks & resource usage map - private final Set sandboxLevelViews; - private final Set activeSandboxes; - - public TaskCancellationContext(Set sandboxLevelViews, Set activeSandboxes) { - this.sandboxLevelViews = sandboxLevelViews; - this.activeSandboxes = activeSandboxes; - } - - public Set getSandboxLevelViews() { - return sandboxLevelViews; - } - - public Set getActiveSandboxes() { - return activeSandboxes; - } -} From 3952631bd884feca70c95173baec76a7e985b2db Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 15 Jul 2024 09:55:25 -0700 Subject: [PATCH 21/66] resolve flakiness Signed-off-by: Kiran Prakash --- .../DefaultTaskCancellationTests.java | 18 +++++++++--------- ...tRunningTaskFirstStrategyStrategyTests.java | 2 +- .../TaskSelectionStrategyTests.java | 2 +- 3 files changed, 11 insertions(+), 11 deletions(-) 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 index 2d1233eb83286..d836089b74119 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/DefaultTaskCancellationTests.java @@ -29,7 +29,7 @@ public class DefaultTaskCancellationTests extends OpenSearchTestCase { - private class TestTaskCancellationImpl extends DefaultTaskCancellation { + private static class TestTaskCancellationImpl extends DefaultTaskCancellation { public TestTaskCancellationImpl( TaskSelectionStrategy taskSelectionStrategy, @@ -75,8 +75,8 @@ public void testGetCancellableTasksFrom_returnsTasksWhenBreachingThreshold() { List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(sandbox1); assertEquals(2, cancellableTasksFrom.size()); - assertEquals(4321, cancellableTasksFrom.get(0).getTask().getId()); - assertEquals(1234, cancellableTasksFrom.get(1).getTask().getId()); + assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); + assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); } public void testGetCancellableTasksFrom_returnsNoTasksWhenBreachingThreshold() { @@ -115,8 +115,8 @@ public void testCancelTasks_cancelsGivenTasks() { List cancellableTasksFrom = taskCancellation.getAllCancellableTasks(); assertEquals(2, cancellableTasksFrom.size()); - assertEquals(4321, cancellableTasksFrom.get(0).getTask().getId()); - assertEquals(1234, cancellableTasksFrom.get(1).getTask().getId()); + assertEquals(1234, cancellableTasksFrom.get(0).getTask().getId()); + assertEquals(4321, cancellableTasksFrom.get(1).getTask().getId()); taskCancellation.cancelTasks(); assertTrue(cancellableTasksFrom.get(0).getTask().isCancelled()); @@ -153,8 +153,8 @@ public void testGetAllCancellableTasks_ReturnsTasksWhenBreachingThresholds() { List allCancellableTasks = taskCancellation.getAllCancellableTasks(); assertEquals(2, allCancellableTasks.size()); - assertEquals(4321, allCancellableTasks.get(0).getTask().getId()); - assertEquals(1234, allCancellableTasks.get(1).getTask().getId()); + assertEquals(1234, allCancellableTasks.get(0).getTask().getId()); + assertEquals(4321, allCancellableTasks.get(1).getTask().getId()); } public void testGetCancellableTasksFrom_returnsTasksEvenWhenSandboxIdNotFound() { @@ -174,7 +174,7 @@ public void testGetCancellableTasksFrom_returnsTasksEvenWhenSandboxIdNotFound() List cancellableTasksFrom = taskCancellation.getCancellableTasksFrom(sandbox1); assertEquals(2, cancellableTasksFrom.size()); - assertEquals(4321, cancellableTasksFrom.get(0).getTask().getId()); - assertEquals(1234, cancellableTasksFrom.get(1).getTask().getId()); + 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/ShortestRunningTaskFirstStrategyStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java index cdd3053e58373..c88cb37781a10 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java @@ -29,6 +29,6 @@ public void testSortingCondition() { List tasks = Arrays.asList(task1, task3, task2); tasks.sort(new ShortestRunningTaskFirstStrategy().sortingCondition()); - assertEquals(Arrays.asList(task1, task2, task3), tasks); + assertEquals(Arrays.asList(task3, task2, task1), tasks); } } diff --git a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java index fe89263a09c3c..6fdae3a6fece9 100644 --- a/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/sandboxing/cancellation/TaskSelectionStrategyTests.java @@ -22,7 +22,7 @@ public class TaskSelectionStrategyTests extends OpenSearchTestCase { public static class TestTaskSelectionStrategy extends AbstractTaskSelectionStrategy { @Override public Comparator sortingCondition() { - return Comparator.comparingLong(Task::getStartTime).reversed(); + return Comparator.comparingLong(Task::getId); } } From 5ac3d83e79a619937b6307c85ccedeb27af52222 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 18 Jul 2024 15:02:01 -0700 Subject: [PATCH 22/66] renaming sandbox to querygroup and adjusting code based on merged PRs Signed-off-by: Kiran Prakash --- .../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 ------------------ 44 files changed, 953 insertions(+), 1047 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/cluster/metadata/Metadata.java b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java index a5f2c8d97a3be..ca7fa5a8ee5f9 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -127,7 +127,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()); - } -} From 9420e1b85ccf5d322955d1af1e06a1c101cf77d4 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 18 Jul 2024 17:45:20 -0700 Subject: [PATCH 23/66] jvm to memory Signed-off-by: Kiran Prakash --- ...QueryGroupResourceUsageTrackerService.java | 2 +- .../opensearch/search/resourcetypes/CPU.java | 2 +- .../resourcetypes/{JVM.java => Memory.java} | 20 +++++++++---------- .../search/resourcetypes/ResourceType.java | 8 ++++---- ...QueryGroupLevelResourceUsageViewTests.java | 4 ++-- .../TaskSelectionStrategyTests.java | 6 +++--- ...GroupResourceUsageTrackerServiceTests.java | 6 ++++-- 7 files changed, 25 insertions(+), 23 deletions(-) rename server/src/main/java/org/opensearch/search/resourcetypes/{JVM.java => Memory.java} (68%) 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 index 7e0fcd6bc7eea..5fcb0faff6d43 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java @@ -27,7 +27,7 @@ // @ExperimentalApi public class QueryGroupResourceUsageTrackerService implements QueryGroupUsageTracker, TaskManager.TaskEventListeners { - public static final List TRACKED_RESOURCES = List.of(ResourceType.fromName("JVM"), ResourceType.fromName("CPU")); + public static final List TRACKED_RESOURCES = List.of(ResourceType.fromName("Memory"), ResourceType.fromName("CPU")); private final TaskManager taskManager; private final TaskResourceTrackingService taskResourceTrackingService; diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java b/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java index bdb30aab970bd..725473600dbd4 100644 --- a/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java +++ b/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java @@ -33,7 +33,7 @@ public long getResourceUsage(Task task) { */ @Override public String getName() { - return "CPU"; + return "cpu"; } @Override diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/JVM.java b/server/src/main/java/org/opensearch/search/resourcetypes/Memory.java similarity index 68% rename from server/src/main/java/org/opensearch/search/resourcetypes/JVM.java rename to server/src/main/java/org/opensearch/search/resourcetypes/Memory.java index 4861e5a94b1df..515bc8e68edc9 100644 --- a/server/src/main/java/org/opensearch/search/resourcetypes/JVM.java +++ b/server/src/main/java/org/opensearch/search/resourcetypes/Memory.java @@ -12,15 +12,15 @@ 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. + * Represents the Memory resource type for tasks. + * This class provides methods to retrieve Memory usage and convert threshold percentages to long values. */ -public class JVM extends ResourceType { +public class Memory extends ResourceType { /** - * Returns the JVM usage of the provided task. + * Returns the Memory usage of the provided task. * - * @param task The task whose JVM usage is to be returned - * @return The JVM usage of the 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) { @@ -30,21 +30,21 @@ public long getResourceUsage(Task task) { /** * Returns the name of the resource type. * - * @return The name of the resource type, which is "JVM" + * @return The name of the resource type, which is "Memory" */ @Override public String getName() { - return "JVM"; + return "memory"; } @Override public boolean equals(Object obj) { - return obj instanceof JVM; + return obj instanceof Memory; } @Override public int hashCode() { - return "JVM".hashCode(); + return "Memory".hashCode(); } /** diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java index ee091dd70c1e9..bec0d71766be8 100644 --- a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java @@ -20,7 +20,7 @@ @PublicApi(since = "2.x") public abstract class ResourceType { public static ResourceType[] values() { - return new ResourceType[] { new CPU(), new JVM() }; + return new ResourceType[] { new CPU(), new Memory() }; } /** @@ -34,7 +34,7 @@ public static ResourceType[] values() { /** * Creates a SystemResource from a string. - * If the string is "JVM", a JVM is returned. + * If the string is "Memory", a Memory is returned. * If the string is "CPU", a CPU is returned. * If the string is not recognized, an IllegalArgumentException is thrown. * @@ -43,8 +43,8 @@ public static ResourceType[] values() { * @throws IllegalArgumentException If the string is not recognized */ public static ResourceType fromName(String type) { - if (type.equalsIgnoreCase("JVM")) { - return new JVM(); + if (type.equalsIgnoreCase("JVM") || type.equalsIgnoreCase("memory")) { + return new Memory(); } else if (type.equalsIgnoreCase("CPU")) { return new CPU(); } else { diff --git a/server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java index 804818667244d..6c9b2d861001c 100644 --- a/server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java +++ b/server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java @@ -23,7 +23,7 @@ public class QueryGroupLevelResourceUsageViewTests extends OpenSearchTestCase { public void setUp() throws Exception { super.setUp(); - resourceUsage = Map.of(ResourceType.fromName("JVM"), 34L, ResourceType.fromName("CPU"), 12L); + resourceUsage = Map.of(ResourceType.fromName("Memory"), 34L, ResourceType.fromName("CPU"), 12L); activeTasks = List.of(getRandomTask(4321)); } @@ -61,6 +61,6 @@ public void testGetActiveTasksDefault() { } private boolean assertResourceUsageData(Map resourceUsageData) { - return resourceUsageData.get(ResourceType.fromName("JVM")) == 34L && resourceUsageData.get(ResourceType.fromName("CPU")) == 12L; + return resourceUsageData.get(ResourceType.fromName("Memory")) == 34L && resourceUsageData.get(ResourceType.fromName("CPU")) == 12L; } } diff --git a/server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java index afcf83229c72b..5c13623d51253 100644 --- a/server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java +++ b/server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java @@ -30,7 +30,7 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsGrea TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); long threshold = 100L; long reduceBy = 50L; - ResourceType resourceType = ResourceType.fromName("JVM"); + ResourceType resourceType = ResourceType.fromName("Memory"); List tasks = QueryGroupTestHelpers.getListOfTasks(threshold); List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); @@ -42,7 +42,7 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsLess TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); long threshold = 100L; long reduceBy = -50L; - ResourceType resourceType = ResourceType.fromName("JVM"); + ResourceType resourceType = ResourceType.fromName("Memory"); List tasks = QueryGroupTestHelpers.getListOfTasks(threshold); try { @@ -57,7 +57,7 @@ public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsEqua TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); long threshold = 100L; long reduceBy = 0; - ResourceType resourceType = ResourceType.fromName("JVM"); + ResourceType resourceType = ResourceType.fromName("Memory"); List tasks = QueryGroupTestHelpers.getListOfTasks(threshold); List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); diff --git a/server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java index 27e923def2b95..46132d2dd5277 100644 --- a/server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java +++ b/server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java @@ -68,7 +68,9 @@ public void testConstructQueryGroupLevelViews_CreatesQueryGroupLevelUsageView_Wh for (String queryGroupId : queryGroupIds) { assertEquals( 400, - (long) stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId).getResourceUsageData().get(ResourceType.fromName("JVM")) + (long) stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId) + .getResourceUsageData() + .get(ResourceType.fromName("Memory")) ); assertEquals(2, stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId).getActiveTasks().size()); } @@ -89,7 +91,7 @@ public void testConstructQueryGroupLevelUsageViews_WithTasksHavingDifferentResou Map queryGroupViews = queryGroupResourceUsageTrackerService .constructQueryGroupLevelUsageViews(); - assertEquals(600, (long) queryGroupViews.get("queryGroup1").getResourceUsageData().get(ResourceType.fromName("JVM"))); + assertEquals(600, (long) queryGroupViews.get("queryGroup1").getResourceUsageData().get(ResourceType.fromName("Memory"))); assertEquals(2, queryGroupViews.get("queryGroup1").getActiveTasks().size()); } From 74d77363088820f9730a26512fa2337cffa47d7c Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 18 Jul 2024 17:55:24 -0700 Subject: [PATCH 24/66] missing java docs Signed-off-by: Kiran Prakash --- .../search/resourcetypes/ResourceType.java | 2 +- .../search/resourcetypes/package-info.java | 12 ++++++++++++ 2 files changed, 13 insertions(+), 1 deletion(-) create mode 100644 server/src/main/java/org/opensearch/search/resourcetypes/package-info.java diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java index bec0d71766be8..b9529765ce83d 100644 --- a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java @@ -15,7 +15,7 @@ import java.io.IOException; /** - * Enum to hold the resource type + * Abstract class representing a resource type. */ @PublicApi(since = "2.x") public abstract class ResourceType { diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/package-info.java b/server/src/main/java/org/opensearch/search/resourcetypes/package-info.java new file mode 100644 index 0000000000000..fae5b44bf82ca --- /dev/null +++ b/server/src/main/java/org/opensearch/search/resourcetypes/package-info.java @@ -0,0 +1,12 @@ +/* + * 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. + */ + +/** + * System Resource Type related artifacts + */ +package org.opensearch.search.resourcetypes; \ No newline at end of file From 9818cbb07073d0375954157007e18207c9a50790 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 18 Jul 2024 18:07:34 -0700 Subject: [PATCH 25/66] spotless Signed-off-by: Kiran Prakash --- .../java/org/opensearch/search/resourcetypes/package-info.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/package-info.java b/server/src/main/java/org/opensearch/search/resourcetypes/package-info.java index fae5b44bf82ca..da92fb5430313 100644 --- a/server/src/main/java/org/opensearch/search/resourcetypes/package-info.java +++ b/server/src/main/java/org/opensearch/search/resourcetypes/package-info.java @@ -9,4 +9,4 @@ /** * System Resource Type related artifacts */ -package org.opensearch.search.resourcetypes; \ No newline at end of file +package org.opensearch.search.resourcetypes; From b234c1898f5516e7fd6d85c01e0f64013ff78e51 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Fri, 19 Jul 2024 07:44:38 -0700 Subject: [PATCH 26/66] Update CHANGELOG.md Signed-off-by: Kiran Prakash --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 061e3280852e8..2d2f669f81f2f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,6 +12,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add basic aggregation support for derived fields ([#14618](https://github.com/opensearch-project/OpenSearch/pull/14618)) - Add ThreadContextPermission for markAsSystemContext and allow core to perform the method ([#15016](https://github.com/opensearch-project/OpenSearch/pull/15016)) - Add ThreadContextPermission for stashAndMergeHeaders and stashWithOrigin ([#15039](https://github.com/opensearch-project/OpenSearch/pull/15039)) +- QueryGroup resource tracking and cancellation framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) - [Concurrent Segment Search] Support composite aggregations with scripting ([#15072](https://github.com/opensearch-project/OpenSearch/pull/15072)) - Add `rangeQuery` and `regexpQuery` for `constant_keyword` field type ([#14711](https://github.com/opensearch-project/OpenSearch/pull/14711)) From 9a6f3aca341a4e0acfd6b36616479c8f417bb95e Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Fri, 19 Jul 2024 14:28:51 -0700 Subject: [PATCH 27/66] pluck cancellation changes out of this PR Signed-off-by: Kiran Prakash --- .../search/querygroup/QueryGroupService.java | 15 +- .../AbstractTaskSelectionStrategy.java | 81 -------- .../cancellation/DefaultTaskCancellation.java | 177 ------------------ .../LongestRunningTaskFirstStrategy.java | 29 --- .../ShortestRunningTaskFirstStrategy.java | 29 --- .../cancellation/TaskSelectionStrategy.java | 32 ---- .../querygroup/cancellation/package-info.java | 12 -- .../querygroup/module/QueryGroupModule.java | 2 - ...QueryGroupResourceUsageTrackerService.java | 1 - ...QueryGroupLevelResourceUsageViewTests.java | 2 +- .../querygroup/QueryGroupTestHelpers.java | 31 +++ .../DefaultTaskCancellationTests.java | 175 ----------------- ...RunningTaskFirstStrategyStrategyTests.java | 34 ---- .../cancellation/QueryGroupTestHelpers.java | 106 ----------- ...RunningTaskFirstStrategyStrategyTests.java | 34 ---- .../TaskSelectionStrategyTests.java | 77 -------- 16 files changed, 33 insertions(+), 804 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/search/querygroup/cancellation/AbstractTaskSelectionStrategy.java delete mode 100644 server/src/main/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellation.java delete mode 100644 server/src/main/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategy.java delete mode 100644 server/src/main/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategy.java delete mode 100644 server/src/main/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategy.java delete mode 100644 server/src/main/java/org/opensearch/search/querygroup/cancellation/package-info.java create mode 100644 server/src/test/java/org/opensearch/search/querygroup/QueryGroupTestHelpers.java delete mode 100644 server/src/test/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellationTests.java delete mode 100644 server/src/test/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java delete mode 100644 server/src/test/java/org/opensearch/search/querygroup/cancellation/QueryGroupTestHelpers.java delete mode 100644 server/src/test/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java delete mode 100644 server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java diff --git a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java index 5a4f465662d0e..f5f9feea038ca 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java +++ b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java @@ -15,8 +15,6 @@ import org.opensearch.common.inject.Inject; import org.opensearch.common.lifecycle.AbstractLifecycleComponent; import org.opensearch.common.unit.TimeValue; -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; @@ -34,9 +32,7 @@ public class QueryGroupService extends AbstractLifecycleComponent { private static final Logger logger = LogManager.getLogger(QueryGroupService.class); private final QueryGroupUsageTracker queryGroupUsageTracker; - // private final QueryGroupPruner queryGroupPruner; private volatile Scheduler.Cancellable scheduledFuture; - // private final QueryGroupServiceSettings queryGroupServiceSettings; private final ThreadPool threadPool; private final ClusterService clusterService; @@ -44,9 +40,7 @@ public class QueryGroupService extends AbstractLifecycleComponent { * Guice managed constructor * * @param queryGroupUsageTracker tracker service - // * @param queryGroupPruner - // * @param queryGroupServiceSettings - * @param threadPool threadpool this will be used to schedule the service + * @param threadPool threadPool this will be used to schedule the service */ @Inject public QueryGroupService( @@ -70,13 +64,6 @@ private void doRun() { Map queryGroupLevelResourceUsageViews = queryGroupUsageTracker .constructQueryGroupLevelUsageViews(); Set activeQueryGroups = getActiveQueryGroups(); - DefaultTaskCancellation taskCancellation = new DefaultTaskCancellation( - new LongestRunningTaskFirstStrategy(), - queryGroupLevelResourceUsageViews, - activeQueryGroups - ); - taskCancellation.cancelTasks(); - // TODO Prune the QueryGroups } private Set getActiveQueryGroups() { diff --git a/server/src/main/java/org/opensearch/search/querygroup/cancellation/AbstractTaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/AbstractTaskSelectionStrategy.java deleted file mode 100644 index 36c6ba60bd850..0000000000000 --- a/server/src/main/java/org/opensearch/search/querygroup/cancellation/AbstractTaskSelectionStrategy.java +++ /dev/null @@ -1,81 +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.querygroup.cancellation; - -import org.opensearch.search.resourcetypes.ResourceType; -import org.opensearch.tasks.CancellableTask; -import org.opensearch.tasks.Task; -import org.opensearch.tasks.TaskCancellation; - -import java.util.ArrayList; -import java.util.Collections; -import java.util.Comparator; -import java.util.List; -import java.util.stream.Collectors; - -/** - * Represents an abstract task selection strategy. - * This class implements the TaskSelectionStrategy interface and provides a method to select tasks for cancellation based on a sorting condition. - * The specific sorting condition depends on the implementation. - */ -public abstract class AbstractTaskSelectionStrategy implements TaskSelectionStrategy { - - /** - * Returns a comparator that defines the sorting condition for tasks. - * The specific sorting condition depends on the implementation. - * - * @return The comparator - */ - public abstract Comparator sortingCondition(); - - /** - * Selects tasks for cancellation based on the provided limit and resource type. - * The tasks are sorted based on the sorting condition and then selected until the accumulated resource usage reaches the limit. - * - * @param tasks The list of tasks from which to select - * @param limit The limit on the accumulated resource usage - * @param resourceType The type of resource to consider - * @return The list of selected tasks - * @throws IllegalArgumentException If the limit is less than zero - */ - @Override - public List selectTasksForCancellation(List tasks, long limit, ResourceType resourceType) { - if (limit < 0) { - throw new IllegalArgumentException("reduceBy has to be greater than zero"); - } - if (limit == 0) { - return Collections.emptyList(); - } - - List sortedTasks = tasks.stream().sorted(sortingCondition()).collect(Collectors.toList()); - - List selectedTasks = new ArrayList<>(); - long accumulated = 0; - - for (Task task : sortedTasks) { - if (task instanceof CancellableTask) { - selectedTasks.add(createTaskCancellation((CancellableTask) task)); - accumulated += resourceType.getResourceUsage(task); - if (accumulated >= limit) { - break; - } - } - } - return selectedTasks; - } - - private TaskCancellation createTaskCancellation(CancellableTask task) { - // todo add correct reason and callbacks - return new TaskCancellation(task, List.of(new TaskCancellation.Reason("limits exceeded", 5)), List.of(this::callbackOnCancel)); - } - - private void callbackOnCancel() { - // todo Implement callback logic here - } -} 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 deleted file mode 100644 index 5104fef64ce38..0000000000000 --- a/server/src/main/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellation.java +++ /dev/null @@ -1,177 +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.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/querygroup/cancellation/LongestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategy.java deleted file mode 100644 index 26eb280a7aac3..0000000000000 --- a/server/src/main/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategy.java +++ /dev/null @@ -1,29 +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.querygroup.cancellation; - -import org.opensearch.tasks.Task; - -import java.util.Comparator; - -/** - * Represents a task selection strategy that prioritizes the longest running tasks first. - */ -public class LongestRunningTaskFirstStrategy extends AbstractTaskSelectionStrategy { - - /** - * Returns a comparator that sorts tasks based on their start time in descending order. - * - * @return The comparator - */ - @Override - public Comparator sortingCondition() { - return Comparator.comparingLong(Task::getStartTime); - } -} diff --git a/server/src/main/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategy.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategy.java deleted file mode 100644 index d2335b90e8f31..0000000000000 --- a/server/src/main/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategy.java +++ /dev/null @@ -1,29 +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.querygroup.cancellation; - -import org.opensearch.tasks.Task; - -import java.util.Comparator; - -/** - * Represents a task selection strategy that prioritizes the shortest running tasks first. - */ -public class ShortestRunningTaskFirstStrategy extends AbstractTaskSelectionStrategy { - - /** - * Returns a comparator that sorts tasks based on their start time in ascending order. - * - * @return The comparator - */ - @Override - public Comparator sortingCondition() { - return Comparator.comparingLong(Task::getStartTime).reversed(); - } -} diff --git a/server/src/main/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategy.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategy.java deleted file mode 100644 index e467db15bedac..0000000000000 --- a/server/src/main/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategy.java +++ /dev/null @@ -1,32 +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.querygroup.cancellation; - -import org.opensearch.search.resourcetypes.ResourceType; -import org.opensearch.tasks.Task; -import org.opensearch.tasks.TaskCancellation; - -import java.util.List; - -/** - * Interface for strategies to select tasks for cancellation. - * Implementations of this interface define how tasks are selected for cancellation based on resource usage. - */ -public interface TaskSelectionStrategy { - /** - * Determines which tasks should be cancelled based on the provided criteria. - * - * @param tasks List of tasks available for cancellation. - * @param limit The amount of tasks to select whose resources reach this limit - * @param resourceType The type of resource that needs to be reduced, guiding the selection process. - * - * @return List of tasks that should be cancelled. - */ - List selectTasksForCancellation(List tasks, long limit, ResourceType resourceType); -} diff --git a/server/src/main/java/org/opensearch/search/querygroup/cancellation/package-info.java b/server/src/main/java/org/opensearch/search/querygroup/cancellation/package-info.java deleted file mode 100644 index fdcb1017f65aa..0000000000000 --- a/server/src/main/java/org/opensearch/search/querygroup/cancellation/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 cancellation related abstracts - */ -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 index f6c6c21a22cd2..aa96417ce328c 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/module/QueryGroupModule.java +++ b/server/src/main/java/org/opensearch/search/querygroup/module/QueryGroupModule.java @@ -25,7 +25,5 @@ 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/querygroup/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java index 5fcb0faff6d43..9153fcf0a7fc5 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java @@ -106,7 +106,6 @@ private Map> getResourceUsageOfQueryGroups(Map resourceUsage; diff --git a/server/src/test/java/org/opensearch/search/querygroup/QueryGroupTestHelpers.java b/server/src/test/java/org/opensearch/search/querygroup/QueryGroupTestHelpers.java new file mode 100644 index 0000000000000..67dcbdeae9e20 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/querygroup/QueryGroupTestHelpers.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; + +import org.opensearch.action.search.SearchAction; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.tasks.Task; + +import java.util.Collections; + +import static org.opensearch.test.OpenSearchTestCase.randomLong; + +public class QueryGroupTestHelpers { + + public static Task getRandomTask(long id) { + return new Task( + id, + "transport", + SearchAction.NAME, + "test description", + new TaskId(randomLong() + ":" + randomLong()), + Collections.emptyMap() + ); + } +} 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 deleted file mode 100644 index afdb3ac5c24b6..0000000000000 --- a/server/src/test/java/org/opensearch/search/querygroup/cancellation/DefaultTaskCancellationTests.java +++ /dev/null @@ -1,175 +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.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/querygroup/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java deleted file mode 100644 index f2918bff4cbdf..0000000000000 --- a/server/src/test/java/org/opensearch/search/querygroup/cancellation/LongestRunningTaskFirstStrategyStrategyTests.java +++ /dev/null @@ -1,34 +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.querygroup.cancellation; - -import org.opensearch.tasks.Task; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.Arrays; -import java.util.List; - -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class LongestRunningTaskFirstStrategyStrategyTests extends OpenSearchTestCase { - public void testSortingCondition() { - Task task1 = mock(Task.class); - Task task2 = mock(Task.class); - Task task3 = mock(Task.class); - when(task1.getStartTime()).thenReturn(100L); - when(task2.getStartTime()).thenReturn(200L); - when(task3.getStartTime()).thenReturn(300L); - - List tasks = Arrays.asList(task2, task1, task3); - tasks.sort(new LongestRunningTaskFirstStrategy().sortingCondition()); - - assertEquals(Arrays.asList(task1, task2, task3), tasks); - } -} diff --git a/server/src/test/java/org/opensearch/search/querygroup/cancellation/QueryGroupTestHelpers.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/QueryGroupTestHelpers.java deleted file mode 100644 index cdcea2375aa9e..0000000000000 --- a/server/src/test/java/org/opensearch/search/querygroup/cancellation/QueryGroupTestHelpers.java +++ /dev/null @@ -1,106 +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.querygroup.cancellation; - -import org.opensearch.action.search.SearchAction; -import org.opensearch.action.search.SearchTask; -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.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; - -import static org.opensearch.test.OpenSearchTestCase.randomLong; -import static org.opensearch.test.OpenSearchTestCase.randomLongBetween; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class QueryGroupTestHelpers { - - public static List getListOfTasks(long totalMemory) { - List tasks = new ArrayList<>(); - - while (totalMemory > 0) { - long id = randomLong(); - final Task task = getRandomSearchTask(id); - long initial_memory = randomLongBetween(1, 100); - - ResourceUsageMetric[] initialTaskResourceMetrics = new ResourceUsageMetric[] { - new ResourceUsageMetric(ResourceStats.MEMORY, initial_memory) }; - task.startThreadResourceTracking(id, ResourceStatsType.WORKER_STATS, initialTaskResourceMetrics); - - long memory = initial_memory + randomLongBetween(1, 10000); - - totalMemory -= memory - initial_memory; - - ResourceUsageMetric[] taskResourceMetrics = new ResourceUsageMetric[] { - new ResourceUsageMetric(ResourceStats.MEMORY, memory), }; - task.updateThreadResourceStats(id, ResourceStatsType.WORKER_STATS, taskResourceMetrics); - task.stopThreadResourceTracking(id, ResourceStatsType.WORKER_STATS); - tasks.add(task); - } - - return tasks; - } - - public static Task getRandomTask(long id) { - return new Task( - id, - "transport", - SearchAction.NAME, - "test description", - new TaskId(randomLong() + ":" + randomLong()), - Collections.emptyMap() - ); - } - - public static Task getRandomSearchTask(long id) { - return new SearchTask( - id, - "transport", - SearchAction.NAME, - () -> "test description", - new TaskId(randomLong() + ":" + randomLong()), - Collections.emptyMap() - ); - } - - public static QueryGroup createQueryGroupMock(String id, String resourceTypeStr, Long threshold, Long usage) { - QueryGroup queryGroupMock = Mockito.mock(QueryGroup.class); - when(queryGroupMock.get_id()).thenReturn(id); - - ResourceType resourceType = ResourceType.fromName(resourceTypeStr); - when(queryGroupMock.getResourceLimits()).thenReturn(Map.of(resourceType, threshold)); - when(queryGroupMock.getThresholdInLong(resourceType)).thenReturn(threshold); - return queryGroupMock; - } - - 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; - } - - private static Supplier descriptionSupplier() { - return () -> "test description"; - } -} diff --git a/server/src/test/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java deleted file mode 100644 index 147a6514a6a74..0000000000000 --- a/server/src/test/java/org/opensearch/search/querygroup/cancellation/ShortestRunningTaskFirstStrategyStrategyTests.java +++ /dev/null @@ -1,34 +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.querygroup.cancellation; - -import org.opensearch.tasks.Task; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.Arrays; -import java.util.List; - -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.when; - -public class ShortestRunningTaskFirstStrategyStrategyTests extends OpenSearchTestCase { - public void testSortingCondition() { - Task task1 = mock(Task.class); - Task task2 = mock(Task.class); - Task task3 = mock(Task.class); - when(task1.getStartTime()).thenReturn(100L); - when(task2.getStartTime()).thenReturn(200L); - when(task3.getStartTime()).thenReturn(300L); - - List tasks = Arrays.asList(task1, task3, task2); - tasks.sort(new ShortestRunningTaskFirstStrategy().sortingCondition()); - - assertEquals(Arrays.asList(task3, task2, task1), tasks); - } -} diff --git a/server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java b/server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java deleted file mode 100644 index 5c13623d51253..0000000000000 --- a/server/src/test/java/org/opensearch/search/querygroup/cancellation/TaskSelectionStrategyTests.java +++ /dev/null @@ -1,77 +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.querygroup.cancellation; - -import org.opensearch.core.tasks.resourcetracker.ResourceStats; -import org.opensearch.search.resourcetypes.ResourceType; -import org.opensearch.tasks.Task; -import org.opensearch.tasks.TaskCancellation; -import org.opensearch.test.OpenSearchTestCase; - -import java.util.Comparator; -import java.util.List; - -public class TaskSelectionStrategyTests extends OpenSearchTestCase { - - public static class TestTaskSelectionStrategy extends AbstractTaskSelectionStrategy { - @Override - public Comparator sortingCondition() { - return Comparator.comparingLong(Task::getId); - } - } - - public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsGreaterThanZero() { - TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); - long threshold = 100L; - long reduceBy = 50L; - ResourceType resourceType = ResourceType.fromName("Memory"); - List tasks = QueryGroupTestHelpers.getListOfTasks(threshold); - - List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); - assertFalse(selectedTasks.isEmpty()); - assertTrue(tasksUsageMeetsThreshold(selectedTasks, reduceBy)); - } - - public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsLesserThanZero() { - TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); - long threshold = 100L; - long reduceBy = -50L; - ResourceType resourceType = ResourceType.fromName("Memory"); - List tasks = QueryGroupTestHelpers.getListOfTasks(threshold); - - try { - testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); - } catch (Exception e) { - assertTrue(e instanceof IllegalArgumentException); - assertEquals("reduceBy has to be greater than zero", e.getMessage()); - } - } - - public void testSelectTasksToCancelSelectsTasksMeetingThreshold_ifReduceByIsEqualToZero() { - TaskSelectionStrategy testTaskSelectionStrategy = new TestTaskSelectionStrategy(); - long threshold = 100L; - long reduceBy = 0; - ResourceType resourceType = ResourceType.fromName("Memory"); - List tasks = QueryGroupTestHelpers.getListOfTasks(threshold); - - List selectedTasks = testTaskSelectionStrategy.selectTasksForCancellation(tasks, reduceBy, resourceType); - assertTrue(selectedTasks.isEmpty()); - } - - private boolean tasksUsageMeetsThreshold(List selectedTasks, long threshold) { - long memory = 0; - for (TaskCancellation task : selectedTasks) { - memory += task.getTask().getTotalResourceUtilization(ResourceStats.MEMORY); - if (memory > threshold) { - return true; - } - } - return false; - } -} From f4d6360db5d4c5432ec009604c4dec0f273390d3 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Fri, 19 Jul 2024 14:40:47 -0700 Subject: [PATCH 28/66] remove unused Signed-off-by: Kiran Prakash --- .../search/querygroup/QueryGroupService.java | 23 ++----------------- 1 file changed, 2 insertions(+), 21 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java index f5f9feea038ca..3bafbe85e9706 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java +++ b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java @@ -14,7 +14,6 @@ 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.querygroup.tracker.QueryGroupUsageTracker; import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.ThreadPool; @@ -23,7 +22,6 @@ import java.util.HashSet; import java.util.Map; import java.util.Set; -import java.util.concurrent.TimeUnit; /** * Main service which will run periodically to track and cancel resource constraint violating tasks in QueryGroups @@ -43,16 +41,8 @@ public class QueryGroupService extends AbstractLifecycleComponent { * @param threadPool threadPool this will be used to schedule the service */ @Inject - public QueryGroupService( - QueryGroupUsageTracker queryGroupUsageTracker, - // QueryGroupServiceSettings queryGroupServiceSettings, - // QueryGroupPruner queryGroupPruner, - ClusterService clusterService, - ThreadPool threadPool - ) { + public QueryGroupService(QueryGroupUsageTracker queryGroupUsageTracker, ClusterService clusterService, ThreadPool threadPool) { this.queryGroupUsageTracker = queryGroupUsageTracker; - // this.queryGroupServiceSettings = queryGroupServiceSettings; - // this.queryGroupPruner = queryGroupPruner; this.clusterService = clusterService; this.threadPool = threadPool; } @@ -75,16 +65,7 @@ private Set getActiveQueryGroups() { */ @Override protected void doStart() { - scheduledFuture = threadPool.scheduleWithFixedDelay(() -> { - try { - doRun(); - } catch (Exception e) { - logger.debug("Exception occurred in QueryGroup service", e); - } - }, - new TimeValue(1, TimeUnit.SECONDS), // TODO get this from QueryGroupServiceSettings - ThreadPool.Names.GENERIC - ); + // This method is intentionally left empty to ensure the service is never scheduled } @Override From cb039fc65e07090beb9ba153ade209fd2d0063d7 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 22 Jul 2024 12:06:38 -0700 Subject: [PATCH 29/66] remove cancellation related code and add more tests coverage Signed-off-by: Kiran Prakash --- .../cluster/metadata/QueryGroup.java | 11 ----- .../opensearch/search/resourcetypes/CPU.java | 11 ----- .../search/resourcetypes/Memory.java | 12 ----- .../search/resourcetypes/ResourceType.java | 8 ---- .../search/resourcetypes/CPUTests.java | 46 +++++++++++++++++++ .../search/resourcetypes/MemoryTests.java | 46 +++++++++++++++++++ .../resourcetypes/ResourceTypeTests.java | 30 ++++++++++++ 7 files changed, 122 insertions(+), 42 deletions(-) create mode 100644 server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java create mode 100644 server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java create mode 100644 server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java 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 0396ddf05b36a..3a2abc3ae7d48 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/QueryGroup.java @@ -227,17 +227,6 @@ 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/search/resourcetypes/CPU.java b/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java index 725473600dbd4..0bc2a309fa426 100644 --- a/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java +++ b/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java @@ -45,15 +45,4 @@ public boolean equals(Object obj) { 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/Memory.java b/server/src/main/java/org/opensearch/search/resourcetypes/Memory.java index 515bc8e68edc9..021381bd45a25 100644 --- a/server/src/main/java/org/opensearch/search/resourcetypes/Memory.java +++ b/server/src/main/java/org/opensearch/search/resourcetypes/Memory.java @@ -8,7 +8,6 @@ package org.opensearch.search.resourcetypes; -import org.opensearch.monitor.jvm.JvmStats; import org.opensearch.tasks.Task; /** @@ -46,15 +45,4 @@ public boolean equals(Object obj) { public int hashCode() { return "Memory".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 index b9529765ce83d..ea08fe013eae9 100644 --- a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java @@ -59,14 +59,6 @@ public static ResourceType fromName(String 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/test/java/org/opensearch/search/resourcetypes/CPUTests.java b/server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java new file mode 100644 index 0000000000000..9746ba3dc2dcb --- /dev/null +++ b/server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java @@ -0,0 +1,46 @@ +/* + * 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.core.tasks.resourcetracker.TaskResourceUsage; +import org.opensearch.tasks.Task; + +import org.mockito.Mockito; + +public class CPUTests extends ResourceTypeTests { + + public void testGetResourceUsage() { + Task mockTask = Mockito.mock(Task.class); + TaskResourceUsage mockStats = Mockito.mock(TaskResourceUsage.class); + Mockito.when(mockTask.getTotalResourceStats()).thenReturn(mockStats); + Mockito.when(mockStats.getCpuTimeInNanos()).thenReturn(1024L); + + CPU cpu = new CPU(); + assertEquals(1024L, cpu.getResourceUsage(mockTask)); + } + + public void testGetName() { + CPU cpu = new CPU(); + assertEquals("cpu", cpu.getName()); + } + + public void testEquals() { + CPU cpu1 = new CPU(); + CPU cpu2 = new CPU(); + Object notCPU = new Object(); + + assertTrue(cpu1.equals(cpu2)); + assertFalse(cpu1.equals(notCPU)); + } + + public void testHashCode() { + CPU cpu = new CPU(); + assertEquals("CPU".hashCode(), cpu.hashCode()); + } +} diff --git a/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java b/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java new file mode 100644 index 0000000000000..26c674fbfff28 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java @@ -0,0 +1,46 @@ +/* + * 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.core.tasks.resourcetracker.TaskResourceUsage; +import org.opensearch.tasks.Task; + +import org.mockito.Mockito; + +public class MemoryTests extends ResourceTypeTests { + + public void testGetResourceUsage() { + Task mockTask = Mockito.mock(Task.class); + TaskResourceUsage mockStats = Mockito.mock(TaskResourceUsage.class); + Mockito.when(mockTask.getTotalResourceStats()).thenReturn(mockStats); + Mockito.when(mockStats.getMemoryInBytes()).thenReturn(1024L); + + Memory memory = new Memory(); + assertEquals(1024L, memory.getResourceUsage(mockTask)); + } + + public void testGetName() { + Memory memory = new Memory(); + assertEquals("memory", memory.getName()); + } + + public void testEquals() { + Memory memory1 = new Memory(); + Memory memory2 = new Memory(); + Object notMemory = new Object(); + + assertTrue(memory1.equals(memory2)); + assertFalse(memory1.equals(notMemory)); + } + + public void testHashCode() { + Memory memory = new Memory(); + assertEquals("Memory".hashCode(), memory.hashCode()); + } +} diff --git a/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java b/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java new file mode 100644 index 0000000000000..9f4ae63369083 --- /dev/null +++ b/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java @@ -0,0 +1,30 @@ +/* + * 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 junit.framework.TestCase; + +import static org.junit.Assert.assertThrows; + +public class ResourceTypeTests extends TestCase { + + public void testFromNameMemory() { + ResourceType result = ResourceType.fromName("memory"); + assertTrue(result instanceof Memory); + } + + public void testFromNameCPU() { + ResourceType result = ResourceType.fromName("CPU"); + assertTrue(result instanceof CPU); + } + + public void testFromNameInvalid() { + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("invalid"); }); + } +} From d6d1a1ad5e47390da6636a2c87c176f72cb632fd Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 22 Jul 2024 12:19:36 -0700 Subject: [PATCH 30/66] us only memory and not jvm Signed-off-by: Kiran Prakash --- .../search/backpressure/SearchBackpressureService.java | 4 ++-- .../opensearch/search/resourcetypes/ResourceType.java | 2 +- .../cluster/metadata/QueryGroupMetadataTests.java | 2 +- .../opensearch/cluster/metadata/QueryGroupTests.java | 10 +++++----- .../backpressure/SearchBackpressureServiceTests.java | 10 +++++----- .../backpressure/trackers/NodeDuressTrackersTests.java | 8 ++++---- 6 files changed, 18 insertions(+), 18 deletions(-) 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 cb26590ba5778..27bca335dd0e3 100644 --- a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java +++ b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java @@ -70,7 +70,7 @@ public class SearchBackpressureService extends AbstractLifecycleComponent implem TaskResourceUsageTrackerType.CPU_USAGE_TRACKER, (nodeDuressTrackers) -> nodeDuressTrackers.isResourceInDuress(ResourceType.fromName("cpu")), TaskResourceUsageTrackerType.HEAP_USAGE_TRACKER, - (nodeDuressTrackers) -> isHeapTrackingSupported() && nodeDuressTrackers.isResourceInDuress(ResourceType.fromName("jvm")), + (nodeDuressTrackers) -> isHeapTrackingSupported() && nodeDuressTrackers.isResourceInDuress(ResourceType.fromName("memory")), TaskResourceUsageTrackerType.ELAPSED_TIME_TRACKER, (nodeDuressTrackers) -> true ); @@ -104,7 +104,7 @@ public SearchBackpressureService( ) ); put( - ResourceType.fromName("jvm"), + ResourceType.fromName("memory"), new NodeDuressTracker( () -> JvmStats.jvmStats().getMem().getHeapUsedPercent() / 100.0 >= settings.getNodeDuressSettings() .getHeapThreshold(), diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java index ea08fe013eae9..a770b403f652e 100644 --- a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java @@ -43,7 +43,7 @@ public static ResourceType[] values() { * @throws IllegalArgumentException If the string is not recognized */ public static ResourceType fromName(String type) { - if (type.equalsIgnoreCase("JVM") || type.equalsIgnoreCase("memory")) { + if (type.equalsIgnoreCase("memory")) { return new Memory(); } else if (type.equalsIgnoreCase("CPU")) { return new CPU(); 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 a5ad13f1a53e8..3cc034b6d3ba9 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -34,7 +34,7 @@ public void testToXContent() throws IOException { "test", "ajakgakg983r92_4242", QueryGroup.ResiliencyMode.ENFORCED, - Map.of(ResourceType.fromName("jvm"), 0.5), + Map.of(ResourceType.fromName("memory"), 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 3d3a3792cb82d..507e9e5f097d8 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupTests.java @@ -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.fromName("jvm"), randomDoubleBetween(0.0, 0.80, false)); + resourceLimit.put(ResourceType.fromName("memory"), 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.fromName("jvm"), (Object) 0.4), Instant.now().getMillis()) + () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.fromName("memory"), (Object) 0.4), Instant.now().getMillis()) ); } @@ -110,7 +110,7 @@ public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { "analytics", "_id", randomMode(), - Map.of(ResourceType.fromName("jvm"), (Object) randomDoubleBetween(1.1, 1.8, false)), + Map.of(ResourceType.fromName("memory"), (Object) randomDoubleBetween(1.1, 1.8, false)), Instant.now().getMillis() ) ); @@ -121,7 +121,7 @@ public void testValidQueryGroup() { "analytics", "_id", randomMode(), - Map.of(ResourceType.fromName("jvm"), randomDoubleBetween(0.01, 0.8, false)), + Map.of(ResourceType.fromName("memory"), 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.fromName("cpu"), 0.30, ResourceType.fromName("jvm"), 0.40), + Map.of(ResourceType.fromName("cpu"), 0.30, ResourceType.fromName("memory"), 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 0459d060555a7..b22a233303d9a 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java +++ b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java @@ -99,7 +99,7 @@ public void testIsNodeInDuress() { HashMap duressTrackers = new HashMap<>() { { - put(ResourceType.fromName("jvm"), heapUsageTracker); + put(ResourceType.fromName("memory"), heapUsageTracker); put(ResourceType.fromName("cpu"), cpuUsageTracker); } }; @@ -230,7 +230,7 @@ public void testSearchTaskInFlightCancellation() { HashMap duressTrackers = new HashMap<>() { { - put(ResourceType.fromName("jvm"), heapUsageTracker); + put(ResourceType.fromName("memory"), heapUsageTracker); put(ResourceType.fromName("cpu"), mockNodeDuressTracker); } }; @@ -305,7 +305,7 @@ public void testSearchShardTaskInFlightCancellation() { HashMap duressTrackers = new HashMap<>() { { - put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 3)); + put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> false, () -> 3)); put(ResourceType.fromName("cpu"), mockNodeDuressTracker); } }; @@ -398,7 +398,7 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { HashMap duressTrackers = new HashMap<>() { { - put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 3)); + put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> false, () -> 3)); put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> true, () -> 3)); } }; @@ -492,7 +492,7 @@ public void testNonCancellationWhenSearchTrafficIsNotQualifyingForCancellation() HashMap duressTrackers = new HashMap<>() { { - put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 3)); + put(ResourceType.fromName("memory"), 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 8d1b3b1fba268..930bd4ed0bdf8 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 @@ -19,7 +19,7 @@ public class NodeDuressTrackersTests extends OpenSearchTestCase { public void testNodeNotInDuress() { HashMap map = new HashMap<>() { { - put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 2)); + put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> false, () -> 2)); put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> false, () -> 2)); } }; @@ -34,7 +34,7 @@ public void testNodeNotInDuress() { public void testNodeInDuressWhenHeapInDuress() { HashMap map = new HashMap<>() { { - put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> true, () -> 3)); put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> false, () -> 1)); } }; @@ -51,7 +51,7 @@ public void testNodeInDuressWhenHeapInDuress() { public void testNodeInDuressWhenCPUInDuress() { HashMap map = new HashMap<>() { { - put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> false, () -> 1)); + put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> false, () -> 1)); put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> true, () -> 3)); } }; @@ -68,7 +68,7 @@ public void testNodeInDuressWhenCPUInDuress() { public void testNodeInDuressWhenCPUAndHeapInDuress() { HashMap map = new HashMap<>() { { - put(ResourceType.fromName("jvm"), new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> true, () -> 3)); put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> false, () -> 3)); } }; From 5da43db75e45443ce007ccb5a41dd35de5a260d1 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 22 Jul 2024 15:24:58 -0700 Subject: [PATCH 31/66] test conventions Signed-off-by: Kiran Prakash --- .../java/org/opensearch/search/resourcetypes/CPUTests.java | 3 ++- .../org/opensearch/search/resourcetypes/MemoryTests.java | 3 ++- .../opensearch/search/resourcetypes/ResourceTypeTests.java | 6 ++---- 3 files changed, 6 insertions(+), 6 deletions(-) diff --git a/server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java b/server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java index 9746ba3dc2dcb..26cfb3d433d7a 100644 --- a/server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java +++ b/server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java @@ -10,10 +10,11 @@ import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; import org.mockito.Mockito; -public class CPUTests extends ResourceTypeTests { +public class CPUTests extends OpenSearchTestCase { public void testGetResourceUsage() { Task mockTask = Mockito.mock(Task.class); diff --git a/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java b/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java index 26c674fbfff28..6d0b0b85ea121 100644 --- a/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java +++ b/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java @@ -10,10 +10,11 @@ import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; import org.opensearch.tasks.Task; +import org.opensearch.test.OpenSearchTestCase; import org.mockito.Mockito; -public class MemoryTests extends ResourceTypeTests { +public class MemoryTests extends OpenSearchTestCase { public void testGetResourceUsage() { Task mockTask = Mockito.mock(Task.class); diff --git a/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java b/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java index 9f4ae63369083..1958d26585cf8 100644 --- a/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java +++ b/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java @@ -8,11 +8,9 @@ package org.opensearch.search.resourcetypes; -import junit.framework.TestCase; +import org.opensearch.test.OpenSearchTestCase; -import static org.junit.Assert.assertThrows; - -public class ResourceTypeTests extends TestCase { +public class ResourceTypeTests extends OpenSearchTestCase { public void testFromNameMemory() { ResourceType result = ResourceType.fromName("memory"); From bdc34ab6c1ce864878c15ec5a6abe263a8c87655 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Sun, 28 Jul 2024 17:40:59 -0700 Subject: [PATCH 32/66] Bring back enum Signed-off-by: Kiran Prakash --- .../cluster/metadata/QueryGroup.java | 2 +- .../org/opensearch/search/ResourceType.java | 51 +++++++++++++++ .../SearchBackpressureService.java | 52 ++++++++------- .../trackers/NodeDuressTrackers.java | 4 +- .../opensearch/search/resourcetypes/CPU.java | 48 -------------- .../search/resourcetypes/Memory.java | 48 -------------- .../search/resourcetypes/ResourceType.java | 65 ------------------- .../search/resourcetypes/package-info.java | 12 ---- .../metadata/QueryGroupMetadataTests.java | 3 +- .../cluster/metadata/QueryGroupTests.java | 12 ++-- .../SearchBackpressureServiceTests.java | 41 ++++++------ .../trackers/NodeDuressTrackersTests.java | 30 ++++----- .../search/resourcetypes/CPUTests.java | 47 -------------- .../search/resourcetypes/MemoryTests.java | 47 -------------- .../resourcetypes/ResourceTypeTests.java | 28 -------- 15 files changed, 127 insertions(+), 363 deletions(-) create mode 100644 server/src/main/java/org/opensearch/search/ResourceType.java delete mode 100644 server/src/main/java/org/opensearch/search/resourcetypes/CPU.java delete mode 100644 server/src/main/java/org/opensearch/search/resourcetypes/Memory.java delete mode 100644 server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java delete mode 100644 server/src/main/java/org/opensearch/search/resourcetypes/package-info.java delete mode 100644 server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java delete mode 100644 server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java delete mode 100644 server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java 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 3a2abc3ae7d48..beaab198073df 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.resourcetypes.ResourceType; +import org.opensearch.search.ResourceType; import org.joda.time.Instant; import java.io.IOException; diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java new file mode 100644 index 0000000000000..b027e83ad59dd --- /dev/null +++ b/server/src/main/java/org/opensearch/search/ResourceType.java @@ -0,0 +1,51 @@ +/* + * 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; + } +} \ No newline at end of file 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 27bca335dd0e3..b9971d8498b91 100644 --- a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java +++ b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java @@ -18,6 +18,7 @@ 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; @@ -33,7 +34,6 @@ 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,6 +47,7 @@ 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; @@ -68,9 +69,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.fromName("cpu")), + (nodeDuressTrackers) -> nodeDuressTrackers.isResourceInDuress(ResourceType.CPU), TaskResourceUsageTrackerType.HEAP_USAGE_TRACKER, - (nodeDuressTrackers) -> isHeapTrackingSupported() && nodeDuressTrackers.isResourceInDuress(ResourceType.fromName("memory")), + (nodeDuressTrackers) -> isHeapTrackingSupported() && nodeDuressTrackers.isResourceInDuress(ResourceType.MEMORY), TaskResourceUsageTrackerType.ELAPSED_TIME_TRACKER, (nodeDuressTrackers) -> true ); @@ -93,26 +94,26 @@ public SearchBackpressureService( ThreadPool threadPool, TaskManager taskManager ) { - this(settings, taskResourceTrackingService, threadPool, System::nanoTime, new NodeDuressTrackers(new HashMap<>() { - { - put( - ResourceType.fromName("cpu"), - new NodeDuressTracker( - () -> ProcessProbe.getInstance().getProcessCpuPercent() / 100.0 >= settings.getNodeDuressSettings() - .getCpuThreshold(), - () -> settings.getNodeDuressSettings().getNumSuccessiveBreaches() - ) - ); - put( - ResourceType.fromName("memory"), - new NodeDuressTracker( - () -> JvmStats.jvmStats().getMem().getHeapUsedPercent() / 100.0 >= settings.getNodeDuressSettings() - .getHeapThreshold(), - () -> settings.getNodeDuressSettings().getNumSuccessiveBreaches() - ) - ); - } - }), + this(settings, taskResourceTrackingService, threadPool, System::nanoTime, new NodeDuressTrackers(new EnumMap<>(ResourceType.class) { + { + put( + ResourceType.CPU, + new NodeDuressTracker( + () -> ProcessProbe.getInstance().getProcessCpuPercent() / 100.0 >= settings.getNodeDuressSettings() + .getCpuThreshold(), + () -> settings.getNodeDuressSettings().getNumSuccessiveBreaches() + ) + ); + put( + ResourceType.MEMORY, + new NodeDuressTracker( + () -> JvmStats.jvmStats().getMem().getHeapUsedPercent() / 100.0 >= settings.getNodeDuressSettings() + .getHeapThreshold(), + () -> settings.getNodeDuressSettings().getNumSuccessiveBreaches() + ) + ); + } + }), getTrackers( settings.getSearchTaskSettings()::getCpuTimeNanosThreshold, settings.getSearchTaskSettings()::getHeapVarianceThreshold, @@ -255,6 +256,7 @@ void doRun() { /** * Had to define this method to help mock this static method to test the scenario where SearchTraffic should not be * penalised when not breaching the threshold + * * @param searchTasks inFlight co-ordinator requests * @param threshold miniumum jvm allocated bytes ratio w.r.t. available heap * @return a boolean value based on whether the threshold is breached @@ -298,6 +300,7 @@ private List addResourceTrackerBasedCancellations( /** * Method to reduce the taskCancellations into unique bunch + * * @param taskCancellations all task cancellations * @return unique task cancellations */ @@ -439,7 +442,8 @@ protected void doStop() { } @Override - protected void doClose() throws IOException {} + protected void doClose() throws IOException { + } public SearchBackpressureStats nodeStats() { List searchTasks = getTaskByType(SearchTask.class); 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 2f74889b034de..dc70f6de07e16 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.resourcetypes.ResourceType; +import org.opensearch.search.ResourceType; import java.util.Map; import java.util.function.BooleanSupplier; @@ -80,4 +80,4 @@ public boolean test() { return breaches.record(isNodeInDuress.getAsBoolean()) >= maxBreachAllowedSupplier.getAsInt(); } } -} +} \ No newline at end of file diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java b/server/src/main/java/org/opensearch/search/resourcetypes/CPU.java deleted file mode 100644 index 0bc2a309fa426..0000000000000 --- a/server/src/main/java/org/opensearch/search/resourcetypes/CPU.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.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(); - } -} diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/Memory.java b/server/src/main/java/org/opensearch/search/resourcetypes/Memory.java deleted file mode 100644 index 021381bd45a25..0000000000000 --- a/server/src/main/java/org/opensearch/search/resourcetypes/Memory.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.resourcetypes; - -import org.opensearch.tasks.Task; - -/** - * Represents the Memory resource type for tasks. - * This class provides methods to retrieve Memory usage and convert threshold percentages to long values. - */ -public class Memory extends ResourceType { - /** - * 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(); - } - - /** - * Returns the name of the resource type. - * - * @return The name of the resource type, which is "Memory" - */ - @Override - public String getName() { - return "memory"; - } - - @Override - public boolean equals(Object obj) { - return obj instanceof Memory; - } - - @Override - public int hashCode() { - return "Memory".hashCode(); - } -} diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java b/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java deleted file mode 100644 index a770b403f652e..0000000000000 --- a/server/src/main/java/org/opensearch/search/resourcetypes/ResourceType.java +++ /dev/null @@ -1,65 +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.resourcetypes; - -import org.opensearch.common.annotation.PublicApi; -import org.opensearch.core.common.io.stream.StreamOutput; -import org.opensearch.tasks.Task; - -import java.io.IOException; - -/** - * Abstract class representing a resource type. - */ -@PublicApi(since = "2.x") -public abstract class ResourceType { - public static ResourceType[] values() { - return new ResourceType[] { new CPU(), new Memory() }; - } - - /** - * 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 "Memory", a Memory 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("memory")) { - return new Memory(); - } 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(); - - public static void writeTo(StreamOutput out, ResourceType resourceType) throws IOException { - out.writeString(resourceType.getName()); - } -} diff --git a/server/src/main/java/org/opensearch/search/resourcetypes/package-info.java b/server/src/main/java/org/opensearch/search/resourcetypes/package-info.java deleted file mode 100644 index da92fb5430313..0000000000000 --- a/server/src/main/java/org/opensearch/search/resourcetypes/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. - */ - -/** - * System Resource Type related artifacts - */ -package org.opensearch.search.resourcetypes; 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 3cc034b6d3ba9..dbb22c71dfc5f 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -14,6 +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; @@ -34,7 +35,7 @@ public void testToXContent() throws IOException { "test", "ajakgakg983r92_4242", QueryGroup.ResiliencyMode.ENFORCED, - Map.of(ResourceType.fromName("memory"), 0.5), + Map.of(ResourceType.MEMORY, 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 507e9e5f097d8..c564f0778e6f0 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.resourcetypes.ResourceType; +import org.opensearch.search.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.fromName("memory"), randomDoubleBetween(0.0, 0.80, false)); + resourceLimit.put(ResourceType.MEMORY, 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.fromName("memory"), (Object) 0.4), Instant.now().getMillis()) + () -> new QueryGroup("analytics", "_id", null, Map.of(ResourceType.MEMORY, (Object) 0.4), Instant.now().getMillis()) ); } @@ -110,7 +110,7 @@ public void testInvalidResourceLimitWhenInvalidSystemResourceValueIsGiven() { "analytics", "_id", randomMode(), - Map.of(ResourceType.fromName("memory"), (Object) randomDoubleBetween(1.1, 1.8, false)), + Map.of(ResourceType.MEMORY, (Object) randomDoubleBetween(1.1, 1.8, false)), Instant.now().getMillis() ) ); @@ -121,7 +121,7 @@ public void testValidQueryGroup() { "analytics", "_id", randomMode(), - Map.of(ResourceType.fromName("memory"), randomDoubleBetween(0.01, 0.8, false)), + Map.of(ResourceType.MEMORY, 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.fromName("cpu"), 0.30, ResourceType.fromName("memory"), 0.40), + Map.of(ResourceType.CPU, 0.30, ResourceType.MEMORY, 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 b22a233303d9a..9f17a7d7f0a15 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java +++ b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java @@ -16,6 +16,7 @@ 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; @@ -28,7 +29,6 @@ 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,6 +45,7 @@ import java.io.IOException; import java.util.Collections; +import java.util.EnumMap; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -55,6 +56,8 @@ 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; @@ -97,10 +100,10 @@ public void testIsNodeInDuress() { NodeDuressTracker cpuUsageTracker = new NodeDuressTracker(() -> cpuUsage.get() >= 0.5, () -> 3); NodeDuressTracker heapUsageTracker = new NodeDuressTracker(() -> heapUsage.get() >= 0.5, () -> 3); - HashMap duressTrackers = new HashMap<>() { + EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(ResourceType.fromName("memory"), heapUsageTracker); - put(ResourceType.fromName("cpu"), cpuUsageTracker); + put(ResourceType.MEMORY, heapUsageTracker); + put(ResourceType.CPU, cpuUsageTracker); } }; @@ -157,7 +160,7 @@ public void testTrackerStateUpdateOnSearchTaskCompletion() { mockTaskResourceTrackingService, threadPool, mockTimeNanosSupplier, - new NodeDuressTrackers(new HashMap<>()), + new NodeDuressTrackers(new EnumMap<>(ResourceType.class)), taskResourceUsageTrackers, new TaskResourceUsageTrackers(), taskManager @@ -188,7 +191,7 @@ public void testTrackerStateUpdateOnSearchShardTaskCompletion() { mockTaskResourceTrackingService, threadPool, mockTimeNanosSupplier, - new NodeDuressTrackers(new HashMap<>()), + new NodeDuressTrackers(new EnumMap<>(ResourceType.class)), new TaskResourceUsageTrackers(), taskResourceUsageTrackers, taskManager @@ -228,10 +231,10 @@ public void testSearchTaskInFlightCancellation() { NodeDuressTracker heapUsageTracker = new NodeDuressTracker(() -> false, () -> 3); - HashMap duressTrackers = new HashMap<>() { + EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(ResourceType.fromName("memory"), heapUsageTracker); - put(ResourceType.fromName("cpu"), mockNodeDuressTracker); + put(MEMORY, heapUsageTracker); + put(CPU, mockNodeDuressTracker); } }; @@ -303,10 +306,10 @@ public void testSearchShardTaskInFlightCancellation() { LongSupplier mockTimeNanosSupplier = mockTime::get; NodeDuressTracker mockNodeDuressTracker = new NodeDuressTracker(() -> true, () -> 3); - HashMap duressTrackers = new HashMap<>() { + EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> false, () -> 3)); - put(ResourceType.fromName("cpu"), mockNodeDuressTracker); + put(MEMORY, new NodeDuressTracker(() -> false, () -> 3)); + put(CPU, mockNodeDuressTracker); } }; NodeDuressTrackers nodeDuressTrackers = new NodeDuressTrackers(duressTrackers); @@ -396,10 +399,10 @@ public void testNonCancellationOfHeapBasedTasksWhenHeapNotInDuress() { AtomicLong mockTime = new AtomicLong(0); LongSupplier mockTimeNanosSupplier = mockTime::get; - HashMap duressTrackers = new HashMap<>() { + EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> false, () -> 3)); - put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> true, () -> 3)); + put(MEMORY, new NodeDuressTracker(() -> false, () -> 3)); + put(CPU, new NodeDuressTracker(() -> true, () -> 3)); } }; @@ -490,10 +493,10 @@ public void testNonCancellationWhenSearchTrafficIsNotQualifyingForCancellation() AtomicLong mockTime = new AtomicLong(0); LongSupplier mockTimeNanosSupplier = mockTime::get; - HashMap duressTrackers = new HashMap<>() { + EnumMap duressTrackers = new EnumMap<>(ResourceType.class) { { - put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> false, () -> 3)); - put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> true, () -> 3)); + put(MEMORY, new NodeDuressTracker(() -> false, () -> 3)); + put(CPU, new NodeDuressTracker(() -> true, () -> 3)); } }; @@ -633,4 +636,4 @@ public int hashCode() { return Objects.hash(cancellationCount); } } -} +} \ No newline at end of file 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 930bd4ed0bdf8..bf99ee9a7cd8b 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.HashMap; +import java.util.EnumMap; public class NodeDuressTrackersTests extends OpenSearchTestCase { public void testNodeNotInDuress() { - HashMap map = new HashMap<>() { + EnumMap map = new EnumMap<>(ResourceType.class) { { - put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> false, () -> 2)); - put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> false, () -> 2)); + put(ResourceType.MEMORY, new NodeDuressTracker(() -> false, () -> 2)); + put(ResourceType.CPU, new NodeDuressTracker(() -> false, () -> 2)); } }; @@ -32,10 +32,10 @@ public void testNodeNotInDuress() { } public void testNodeInDuressWhenHeapInDuress() { - HashMap map = new HashMap<>() { + EnumMap map = new EnumMap<>(ResourceType.class) { { - put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> true, () -> 3)); - put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> false, () -> 1)); + put(ResourceType.MEMORY, new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.CPU, new NodeDuressTracker(() -> false, () -> 1)); } }; @@ -49,10 +49,10 @@ public void testNodeInDuressWhenHeapInDuress() { } public void testNodeInDuressWhenCPUInDuress() { - HashMap map = new HashMap<>() { + EnumMap map = new EnumMap<>(ResourceType.class) { { - put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> false, () -> 1)); - put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.MEMORY, new NodeDuressTracker(() -> false, () -> 1)); + put(ResourceType.CPU, new NodeDuressTracker(() -> true, () -> 3)); } }; @@ -66,10 +66,10 @@ public void testNodeInDuressWhenCPUInDuress() { } public void testNodeInDuressWhenCPUAndHeapInDuress() { - HashMap map = new HashMap<>() { + EnumMap map = new EnumMap<>(ResourceType.class) { { - put(ResourceType.fromName("memory"), new NodeDuressTracker(() -> true, () -> 3)); - put(ResourceType.fromName("cpu"), new NodeDuressTracker(() -> false, () -> 3)); + put(ResourceType.MEMORY, new NodeDuressTracker(() -> true, () -> 3)); + put(ResourceType.CPU, new NodeDuressTracker(() -> false, () -> 3)); } }; @@ -81,4 +81,4 @@ public void testNodeInDuressWhenCPUAndHeapInDuress() { // for the third time it should be in duress assertTrue(nodeDuressTrackers.isNodeInDuress()); } -} +} \ No newline at end of file diff --git a/server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java b/server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java deleted file mode 100644 index 26cfb3d433d7a..0000000000000 --- a/server/src/test/java/org/opensearch/search/resourcetypes/CPUTests.java +++ /dev/null @@ -1,47 +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.resourcetypes; - -import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; -import org.opensearch.tasks.Task; -import org.opensearch.test.OpenSearchTestCase; - -import org.mockito.Mockito; - -public class CPUTests extends OpenSearchTestCase { - - public void testGetResourceUsage() { - Task mockTask = Mockito.mock(Task.class); - TaskResourceUsage mockStats = Mockito.mock(TaskResourceUsage.class); - Mockito.when(mockTask.getTotalResourceStats()).thenReturn(mockStats); - Mockito.when(mockStats.getCpuTimeInNanos()).thenReturn(1024L); - - CPU cpu = new CPU(); - assertEquals(1024L, cpu.getResourceUsage(mockTask)); - } - - public void testGetName() { - CPU cpu = new CPU(); - assertEquals("cpu", cpu.getName()); - } - - public void testEquals() { - CPU cpu1 = new CPU(); - CPU cpu2 = new CPU(); - Object notCPU = new Object(); - - assertTrue(cpu1.equals(cpu2)); - assertFalse(cpu1.equals(notCPU)); - } - - public void testHashCode() { - CPU cpu = new CPU(); - assertEquals("CPU".hashCode(), cpu.hashCode()); - } -} diff --git a/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java b/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java deleted file mode 100644 index 6d0b0b85ea121..0000000000000 --- a/server/src/test/java/org/opensearch/search/resourcetypes/MemoryTests.java +++ /dev/null @@ -1,47 +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.resourcetypes; - -import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; -import org.opensearch.tasks.Task; -import org.opensearch.test.OpenSearchTestCase; - -import org.mockito.Mockito; - -public class MemoryTests extends OpenSearchTestCase { - - public void testGetResourceUsage() { - Task mockTask = Mockito.mock(Task.class); - TaskResourceUsage mockStats = Mockito.mock(TaskResourceUsage.class); - Mockito.when(mockTask.getTotalResourceStats()).thenReturn(mockStats); - Mockito.when(mockStats.getMemoryInBytes()).thenReturn(1024L); - - Memory memory = new Memory(); - assertEquals(1024L, memory.getResourceUsage(mockTask)); - } - - public void testGetName() { - Memory memory = new Memory(); - assertEquals("memory", memory.getName()); - } - - public void testEquals() { - Memory memory1 = new Memory(); - Memory memory2 = new Memory(); - Object notMemory = new Object(); - - assertTrue(memory1.equals(memory2)); - assertFalse(memory1.equals(notMemory)); - } - - public void testHashCode() { - Memory memory = new Memory(); - assertEquals("Memory".hashCode(), memory.hashCode()); - } -} diff --git a/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java b/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java deleted file mode 100644 index 1958d26585cf8..0000000000000 --- a/server/src/test/java/org/opensearch/search/resourcetypes/ResourceTypeTests.java +++ /dev/null @@ -1,28 +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.resourcetypes; - -import org.opensearch.test.OpenSearchTestCase; - -public class ResourceTypeTests extends OpenSearchTestCase { - - public void testFromNameMemory() { - ResourceType result = ResourceType.fromName("memory"); - assertTrue(result instanceof Memory); - } - - public void testFromNameCPU() { - ResourceType result = ResourceType.fromName("CPU"); - assertTrue(result instanceof CPU); - } - - public void testFromNameInvalid() { - assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("invalid"); }); - } -} From c6b2e69fcaf8f435aa61934df6b4ccb12e1b5f30 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Sun, 28 Jul 2024 17:43:05 -0700 Subject: [PATCH 33/66] Update SearchBackpressureService.java Signed-off-by: Kiran Prakash --- .../SearchBackpressureService.java | 42 +++++++++---------- 1 file changed, 19 insertions(+), 23 deletions(-) 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 b9971d8498b91..9f429118c11e8 100644 --- a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java +++ b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java @@ -94,25 +94,24 @@ public SearchBackpressureService( ThreadPool threadPool, TaskManager taskManager ) { - this(settings, taskResourceTrackingService, threadPool, System::nanoTime, new NodeDuressTrackers(new EnumMap<>(ResourceType.class) { - { - put( - ResourceType.CPU, - new NodeDuressTracker( - () -> ProcessProbe.getInstance().getProcessCpuPercent() / 100.0 >= settings.getNodeDuressSettings() - .getCpuThreshold(), - () -> settings.getNodeDuressSettings().getNumSuccessiveBreaches() - ) - ); - put( - ResourceType.MEMORY, - new NodeDuressTracker( - () -> JvmStats.jvmStats().getMem().getHeapUsedPercent() / 100.0 >= settings.getNodeDuressSettings() - .getHeapThreshold(), - () -> settings.getNodeDuressSettings().getNumSuccessiveBreaches() - ) - ); - } + this(settings, taskResourceTrackingService, threadPool, System::nanoTime, new NodeDuressTrackers(new EnumMap<>(ResourceType.class) { { + put( + ResourceType.CPU, + new NodeDuressTracker( + () -> ProcessProbe.getInstance().getProcessCpuPercent() / 100.0 >= settings.getNodeDuressSettings() + .getCpuThreshold(), + () -> settings.getNodeDuressSettings().getNumSuccessiveBreaches() + ) + ); + put( + ResourceType.MEMORY, + new NodeDuressTracker( + () -> JvmStats.jvmStats().getMem().getHeapUsedPercent() / 100.0 >= settings.getNodeDuressSettings() + .getHeapThreshold(), + () -> settings.getNodeDuressSettings().getNumSuccessiveBreaches() + ) + ); + } }), getTrackers( settings.getSearchTaskSettings()::getCpuTimeNanosThreshold, @@ -256,7 +255,6 @@ void doRun() { /** * Had to define this method to help mock this static method to test the scenario where SearchTraffic should not be * penalised when not breaching the threshold - * * @param searchTasks inFlight co-ordinator requests * @param threshold miniumum jvm allocated bytes ratio w.r.t. available heap * @return a boolean value based on whether the threshold is breached @@ -300,7 +298,6 @@ private List addResourceTrackerBasedCancellations( /** * Method to reduce the taskCancellations into unique bunch - * * @param taskCancellations all task cancellations * @return unique task cancellations */ @@ -442,8 +439,7 @@ protected void doStop() { } @Override - protected void doClose() throws IOException { - } + protected void doClose() throws IOException {} public SearchBackpressureStats nodeStats() { List searchTasks = getTaskByType(SearchTask.class); From 499e139d22c95e2df60fee18beec63b3c88c96fd Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Sun, 28 Jul 2024 17:46:53 -0700 Subject: [PATCH 34/66] revert changes Signed-off-by: Kiran Prakash --- .../org/opensearch/search/ResourceType.java | 56 +++++++++---------- .../SearchBackpressureService.java | 5 +- 2 files changed, 31 insertions(+), 30 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java index b027e83ad59dd..fe5ce4dd2bb50 100644 --- a/server/src/main/java/org/opensearch/search/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/ResourceType.java @@ -18,34 +18,34 @@ */ @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; - } + 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 + "]"); } - throw new IllegalArgumentException("Unknown resource type: [" + s + "]"); - } - public static void writeTo(StreamOutput out, ResourceType resourceType) throws IOException { - out.writeString(resourceType.getName()); - } + public static void writeTo(StreamOutput out, ResourceType resourceType) throws IOException { + out.writeString(resourceType.getName()); + } - public String getName() { - return name; - } -} \ No newline at end of file + 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 9f429118c11e8..c26c5d63a3573 100644 --- a/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java +++ b/server/src/main/java/org/opensearch/search/backpressure/SearchBackpressureService.java @@ -94,7 +94,8 @@ 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 EnumMap<>(ResourceType.class) { + { put( ResourceType.CPU, new NodeDuressTracker( @@ -112,7 +113,7 @@ public SearchBackpressureService( ) ); } - }), + }), getTrackers( settings.getSearchTaskSettings()::getCpuTimeNanosThreshold, settings.getSearchTaskSettings()::getHeapVarianceThreshold, From 538f8af581fc94157c76a5a54d1f5a51270ded57 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Sun, 28 Jul 2024 17:48:12 -0700 Subject: [PATCH 35/66] revert changes Signed-off-by: Kiran Prakash --- .../search/backpressure/trackers/NodeDuressTrackers.java | 2 +- .../search/backpressure/SearchBackpressureServiceTests.java | 2 +- .../search/backpressure/trackers/NodeDuressTrackersTests.java | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) 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 dc70f6de07e16..ae60a82fc2816 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 @@ -80,4 +80,4 @@ public boolean test() { return breaches.record(isNodeInDuress.getAsBoolean()) >= maxBreachAllowedSupplier.getAsInt(); } } -} \ No newline at end of file +} 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 9f17a7d7f0a15..15d0fcd10d701 100644 --- a/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java +++ b/server/src/test/java/org/opensearch/search/backpressure/SearchBackpressureServiceTests.java @@ -636,4 +636,4 @@ public int hashCode() { return Objects.hash(cancellationCount); } } -} \ No newline at end of file +} 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 bf99ee9a7cd8b..801576bdf89d4 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 @@ -81,4 +81,4 @@ public void testNodeInDuressWhenCPUAndHeapInDuress() { // for the third time it should be in duress assertTrue(nodeDuressTrackers.isNodeInDuress()); } -} \ No newline at end of file +} From 225dbdb9fc0fb08048ef20754c7554cae352db5d Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 29 Jul 2024 09:52:34 -0700 Subject: [PATCH 36/66] all required changes Signed-off-by: Kiran Prakash --- .../opensearch/cluster/metadata/Metadata.java | 7 +- .../search/querygroup/QueryGroupTask.java | 18 ----- .../search/wlm/QueryGroupConstants.java | 19 +++++ .../search/wlm/QueryGroupHelper.java | 29 ++++++++ .../QueryGroupLevelResourceUsageView.java | 4 +- .../QueryGroupService.java | 8 +- .../opensearch/search/wlm/QueryGroupTask.java | 74 +++++++++++++++++++ .../module/QueryGroupModule.java | 6 +- .../module/package-info.java | 2 +- .../{querygroup => wlm}/package-info.java | 2 +- ...QueryGroupResourceUsageTrackerService.java | 28 ++++--- .../tracker/QueryGroupUsageTracker.java | 8 +- .../tracker/package-info.java | 4 +- .../metadata/QueryGroupMetadataTests.java | 1 - ...QueryGroupLevelResourceUsageViewTests.java | 10 +-- .../QueryGroupTestHelpers.java | 2 +- ...GroupResourceUsageTrackerServiceTests.java | 18 ++--- 17 files changed, 171 insertions(+), 69 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/search/querygroup/QueryGroupTask.java create mode 100644 server/src/main/java/org/opensearch/search/wlm/QueryGroupConstants.java create mode 100644 server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java rename server/src/main/java/org/opensearch/search/{querygroup => wlm}/QueryGroupLevelResourceUsageView.java (95%) rename server/src/main/java/org/opensearch/search/{querygroup => wlm}/QueryGroupService.java (88%) create mode 100644 server/src/main/java/org/opensearch/search/wlm/QueryGroupTask.java rename server/src/main/java/org/opensearch/search/{querygroup => wlm}/module/QueryGroupModule.java (75%) rename server/src/main/java/org/opensearch/search/{querygroup => wlm}/module/package-info.java (82%) rename server/src/main/java/org/opensearch/search/{querygroup => wlm}/package-info.java (85%) rename server/src/main/java/org/opensearch/search/{querygroup => wlm}/tracker/QueryGroupResourceUsageTrackerService.java (79%) rename server/src/main/java/org/opensearch/search/{querygroup => wlm}/tracker/QueryGroupUsageTracker.java (58%) rename server/src/main/java/org/opensearch/search/{querygroup => wlm}/tracker/package-info.java (68%) rename server/src/test/java/org/opensearch/search/{querygroup => wlm}/QueryGroupLevelResourceUsageViewTests.java (84%) rename server/src/test/java/org/opensearch/search/{querygroup => wlm}/QueryGroupTestHelpers.java (94%) rename server/src/test/java/org/opensearch/search/{querygroup => wlm}/tracking/QueryGroupResourceUsageTrackerServiceTests.java (91%) 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 ca7fa5a8ee5f9..09bef2ddf9ee6 100644 --- a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java +++ b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java @@ -127,11 +127,6 @@ public boolean isSegmentReplicationEnabled(String indexName) { .orElse(false); } - public Map queryGroups() { - // stub - return Collections.emptyMap(); - } - /** * Context of the XContent. * @@ -1396,7 +1391,7 @@ public Builder put(final QueryGroup queryGroup) { return queryGroups(existing); } - private Map getQueryGroups() { + public Map getQueryGroups() { return Optional.ofNullable(this.customs.get(QueryGroupMetadata.TYPE)) .map(o -> (QueryGroupMetadata) o) .map(QueryGroupMetadata::queryGroups) diff --git a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupTask.java b/server/src/main/java/org/opensearch/search/querygroup/QueryGroupTask.java deleted file mode 100644 index d3ab8c2f0f8b6..0000000000000 --- a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupTask.java +++ /dev/null @@ -1,18 +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.querygroup; - -/** - * This interface can be implemented by tasks which will be tracked and monitored using {@link org.opensearch.cluster.metadata.QueryGroup} - */ -public interface QueryGroupTask { - void setQueryGroupId(String queryGroupId); - - String getQueryGroupId(); -} diff --git a/server/src/main/java/org/opensearch/search/wlm/QueryGroupConstants.java b/server/src/main/java/org/opensearch/search/wlm/QueryGroupConstants.java new file mode 100644 index 0000000000000..b64e65c6836bc --- /dev/null +++ b/server/src/main/java/org/opensearch/search/wlm/QueryGroupConstants.java @@ -0,0 +1,19 @@ +/* + * 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.wlm; + +import java.util.function.Supplier; + +/** + * This class will hold all the QueryGroup related constants + */ +public class QueryGroupConstants { + public static final String QUERY_GROUP_ID_HEADER = "queryGroupId"; + public static final Supplier DEFAULT_QUERY_GROUP_ID_SUPPLIER = () -> "DEFAULT_QUERY_GROUP"; +} diff --git a/server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java b/server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java new file mode 100644 index 0000000000000..4917ddd9ed29e --- /dev/null +++ b/server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java @@ -0,0 +1,29 @@ +/* + * 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.wlm; + +import org.opensearch.search.ResourceType; +import org.opensearch.tasks.Task; + +import java.util.Map; +import java.util.function.Function; + +public class QueryGroupHelper { + + private static final Map> resourceUsageCalculator = Map.of( + ResourceType.MEMORY, + (task) -> task.getTotalResourceStats().getMemoryInBytes(), + ResourceType.CPU, + (task) -> task.getTotalResourceStats().getCpuTimeInNanos() + ); + + public static long getResourceUsage(ResourceType resource, Task task) { + return resourceUsageCalculator.get(resource).apply(task); + } +} diff --git a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageView.java b/server/src/main/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageView.java similarity index 95% rename from server/src/main/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageView.java rename to server/src/main/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageView.java index 354d38862f746..faafc64ae2860 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageView.java +++ b/server/src/main/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageView.java @@ -6,10 +6,10 @@ * compatible open source license. */ -package org.opensearch.search.querygroup; +package org.opensearch.search.wlm; import org.opensearch.common.annotation.ExperimentalApi; -import org.opensearch.search.resourcetypes.ResourceType; +import org.opensearch.search.ResourceType; import org.opensearch.tasks.Task; import java.util.ArrayList; diff --git a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java b/server/src/main/java/org/opensearch/search/wlm/QueryGroupService.java similarity index 88% rename from server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java rename to server/src/main/java/org/opensearch/search/wlm/QueryGroupService.java index 3bafbe85e9706..32f5716caea50 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/QueryGroupService.java +++ b/server/src/main/java/org/opensearch/search/wlm/QueryGroupService.java @@ -6,15 +6,16 @@ * compatible open source license. */ -package org.opensearch.search.querygroup; +package org.opensearch.search.wlm; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.opensearch.cluster.metadata.Metadata; 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.search.querygroup.tracker.QueryGroupUsageTracker; +import org.opensearch.search.wlm.tracker.QueryGroupUsageTracker; import org.opensearch.threadpool.Scheduler; import org.opensearch.threadpool.ThreadPool; @@ -57,7 +58,8 @@ private void doRun() { } private Set getActiveQueryGroups() { - return new HashSet<>(clusterService.state().metadata().queryGroups().values()); + Map queryGroups = Metadata.builder(clusterService.state().metadata()).getQueryGroups(); + return new HashSet<>(queryGroups.values()); } /** diff --git a/server/src/main/java/org/opensearch/search/wlm/QueryGroupTask.java b/server/src/main/java/org/opensearch/search/wlm/QueryGroupTask.java new file mode 100644 index 0000000000000..e8b5ca4efd4e3 --- /dev/null +++ b/server/src/main/java/org/opensearch/search/wlm/QueryGroupTask.java @@ -0,0 +1,74 @@ +/* + * 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.wlm; + +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.opensearch.common.unit.TimeValue; +import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.core.tasks.TaskId; +import org.opensearch.tasks.CancellableTask; + +import java.util.Map; + +import static org.opensearch.search.SearchService.NO_TIMEOUT; + +/** + * Base class to define QueryGroup tasks + */ +public class QueryGroupTask extends CancellableTask { + + private static final Logger logger = LogManager.getLogger(QueryGroupTask.class); + private String queryGroupId; + + public QueryGroupTask(long id, String type, String action, String description, TaskId parentTaskId, Map headers) { + this(id, type, action, description, parentTaskId, headers, NO_TIMEOUT); + } + + public QueryGroupTask( + long id, + String type, + String action, + String description, + TaskId parentTaskId, + Map headers, + TimeValue cancelAfterTimeInterval + ) { + super(id, type, action, description, parentTaskId, headers, cancelAfterTimeInterval); + } + + /** + * This method should always be called after calling setQueryGroupId at least once on this object + * @return task queryGroupId + */ + public String getQueryGroupId() { + if (queryGroupId == null) { + logger.warn("QueryGroup _id can't be null, It should be set before accessing it. This is abnormal behaviour "); + } + return queryGroupId; + } + + /** + * sets the queryGroupId from threadContext into the task itself, + * This method was defined since the queryGroupId can only be evaluated after task creation + * @param threadContext current threadContext + */ + public void setQueryGroupId(final ThreadContext threadContext) { + this.queryGroupId = QueryGroupConstants.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(); + + if (threadContext != null && threadContext.getHeader(QueryGroupConstants.QUERY_GROUP_ID_HEADER) != null) { + this.queryGroupId = threadContext.getHeader(QueryGroupConstants.QUERY_GROUP_ID_HEADER); + } + } + + @Override + public boolean shouldCancelChildrenOnCancellation() { + return false; + } +} diff --git a/server/src/main/java/org/opensearch/search/querygroup/module/QueryGroupModule.java b/server/src/main/java/org/opensearch/search/wlm/module/QueryGroupModule.java similarity index 75% rename from server/src/main/java/org/opensearch/search/querygroup/module/QueryGroupModule.java rename to server/src/main/java/org/opensearch/search/wlm/module/QueryGroupModule.java index aa96417ce328c..ddc6954dd412d 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/module/QueryGroupModule.java +++ b/server/src/main/java/org/opensearch/search/wlm/module/QueryGroupModule.java @@ -6,11 +6,11 @@ * compatible open source license. */ -package org.opensearch.search.querygroup.module; +package org.opensearch.search.wlm.module; import org.opensearch.common.inject.AbstractModule; -import org.opensearch.search.querygroup.tracker.QueryGroupResourceUsageTrackerService; -import org.opensearch.search.querygroup.tracker.QueryGroupUsageTracker; +import org.opensearch.search.wlm.tracker.QueryGroupResourceUsageTrackerService; +import org.opensearch.search.wlm.tracker.QueryGroupUsageTracker; /** * Module class for resource usage limiting related artifacts diff --git a/server/src/main/java/org/opensearch/search/querygroup/module/package-info.java b/server/src/main/java/org/opensearch/search/wlm/module/package-info.java similarity index 82% rename from server/src/main/java/org/opensearch/search/querygroup/module/package-info.java rename to server/src/main/java/org/opensearch/search/wlm/module/package-info.java index 769192a288664..aadf19448caef 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/module/package-info.java +++ b/server/src/main/java/org/opensearch/search/wlm/module/package-info.java @@ -10,4 +10,4 @@ * Guice Module */ -package org.opensearch.search.querygroup.module; +package org.opensearch.search.wlm.module; diff --git a/server/src/main/java/org/opensearch/search/querygroup/package-info.java b/server/src/main/java/org/opensearch/search/wlm/package-info.java similarity index 85% rename from server/src/main/java/org/opensearch/search/querygroup/package-info.java rename to server/src/main/java/org/opensearch/search/wlm/package-info.java index 39dfc756a0a4e..a37e5ca764b54 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/package-info.java +++ b/server/src/main/java/org/opensearch/search/wlm/package-info.java @@ -9,4 +9,4 @@ /** * Query Sandboxing related artifacts */ -package org.opensearch.search.querygroup; +package org.opensearch.search.wlm; diff --git a/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java similarity index 79% rename from server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java rename to server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 9153fcf0a7fc5..01bea9077626a 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -6,29 +6,30 @@ * compatible open source license. */ -package org.opensearch.search.querygroup.tracker; +package org.opensearch.search.wlm.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.search.ResourceType; +import org.opensearch.search.wlm.QueryGroupHelper; +import org.opensearch.search.wlm.QueryGroupLevelResourceUsageView; +import org.opensearch.search.wlm.QueryGroupTask; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; import org.opensearch.tasks.TaskResourceTrackingService; +import java.util.EnumMap; import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.stream.Collectors; /** - * This class tracks requests per QueryGroup + * This class tracks resource usage per QueryGroup */ // @ExperimentalApi public class QueryGroupResourceUsageTrackerService implements QueryGroupUsageTracker, TaskManager.TaskEventListeners { - public static final List TRACKED_RESOURCES = List.of(ResourceType.fromName("Memory"), ResourceType.fromName("CPU")); - + public static final List TRACKED_RESOURCES = List.of(ResourceType.fromName("memory"), ResourceType.fromName("cpu")); private final TaskManager taskManager; private final TaskResourceTrackingService taskResourceTrackingService; @@ -57,7 +58,7 @@ public Map constructQueryGroupLevelUsa Map queryGroupViews = new HashMap<>(); Map> tasksByQueryGroup = getTasksGroupedByQueryGroup(); - Map> queryGroupResourceUsage = getResourceUsageOfQueryGroups(tasksByQueryGroup); + Map> queryGroupResourceUsage = getResourceUsageOfQueryGroups(tasksByQueryGroup); for (String queryGroupId : tasksByQueryGroup.keySet()) { QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( @@ -90,8 +91,8 @@ private Map> getTasksGroupedByQueryGroup() { * @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<>(); + private Map> getResourceUsageOfQueryGroups(Map> tasksByQueryGroup) { + Map> resourceUsageOfQueryGroups = new HashMap<>(); // Iterate over each QueryGroup entry for (Map.Entry> queryGroupEntry : tasksByQueryGroup.entrySet()) { @@ -99,13 +100,16 @@ private Map> getResourceUsageOfQueryGroups(Map tasks = queryGroupEntry.getValue(); // Prepare a usage map for the current QueryGroup, or retrieve the existing one - Map queryGroupUsage = resourceUsageOfQueryGroups.computeIfAbsent(queryGroupId, k -> new HashMap<>()); + EnumMap queryGroupUsage = resourceUsageOfQueryGroups.computeIfAbsent( + queryGroupId, + k -> new EnumMap<>(ResourceType.class) + ); // 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); + long taskUsage = QueryGroupHelper.getResourceUsage(resourceType, task); queryGroupUsage.put(resourceType, currentUsage + taskUsage); } } diff --git a/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupUsageTracker.java b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupUsageTracker.java similarity index 58% rename from server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupUsageTracker.java rename to server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupUsageTracker.java index 15e162a5e25a8..269971780daa9 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/tracker/QueryGroupUsageTracker.java +++ b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupUsageTracker.java @@ -6,18 +6,18 @@ * compatible open source license. */ -package org.opensearch.search.querygroup.tracker; +package org.opensearch.search.wlm.tracker; -import org.opensearch.search.querygroup.QueryGroupLevelResourceUsageView; +import org.opensearch.search.wlm.QueryGroupLevelResourceUsageView; import java.util.Map; /** - * This interface is mainly for tracking the resourceLimitGroup level resource usages + * This interface is mainly for tracking the queryGroup level resource usages */ public interface QueryGroupUsageTracker { /** - * updates the current resource usage of resourceLimitGroups + * updates the current resource usage of queryGroup */ Map constructQueryGroupLevelUsageViews(); diff --git a/server/src/main/java/org/opensearch/search/querygroup/tracker/package-info.java b/server/src/main/java/org/opensearch/search/wlm/tracker/package-info.java similarity index 68% rename from server/src/main/java/org/opensearch/search/querygroup/tracker/package-info.java rename to server/src/main/java/org/opensearch/search/wlm/tracker/package-info.java index 28c1bf3d03ded..1012749a845f8 100644 --- a/server/src/main/java/org/opensearch/search/querygroup/tracker/package-info.java +++ b/server/src/main/java/org/opensearch/search/wlm/tracker/package-info.java @@ -7,6 +7,6 @@ */ /** - * ResourceLimitGroup resource tracking artifacts + * QueryGroup resource tracking artifacts */ -package org.opensearch.search.querygroup.tracker; +package org.opensearch.search.wlm.tracker; 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 dbb22c71dfc5f..d70a9ce5e10cd 100644 --- a/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java +++ b/server/src/test/java/org/opensearch/cluster/metadata/QueryGroupMetadataTests.java @@ -15,7 +15,6 @@ 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; diff --git a/server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java similarity index 84% rename from server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java rename to server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java index f00d491a5ff94..98ff82248dd9a 100644 --- a/server/src/test/java/org/opensearch/search/querygroup/QueryGroupLevelResourceUsageViewTests.java +++ b/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java @@ -6,16 +6,16 @@ * compatible open source license. */ -package org.opensearch.search.querygroup; +package org.opensearch.search.wlm; -import org.opensearch.search.resourcetypes.ResourceType; +import org.opensearch.search.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.QueryGroupTestHelpers.getRandomTask; +import static org.opensearch.search.wlm.QueryGroupTestHelpers.getRandomTask; public class QueryGroupLevelResourceUsageViewTests extends OpenSearchTestCase { Map resourceUsage; @@ -23,7 +23,7 @@ public class QueryGroupLevelResourceUsageViewTests extends OpenSearchTestCase { public void setUp() throws Exception { super.setUp(); - resourceUsage = Map.of(ResourceType.fromName("Memory"), 34L, ResourceType.fromName("CPU"), 12L); + resourceUsage = Map.of(ResourceType.fromName("memory"), 34L, ResourceType.fromName("cpu"), 12L); activeTasks = List.of(getRandomTask(4321)); } @@ -61,6 +61,6 @@ public void testGetActiveTasksDefault() { } private boolean assertResourceUsageData(Map resourceUsageData) { - return resourceUsageData.get(ResourceType.fromName("Memory")) == 34L && resourceUsageData.get(ResourceType.fromName("CPU")) == 12L; + return resourceUsageData.get(ResourceType.fromName("memory")) == 34L && resourceUsageData.get(ResourceType.fromName("cpu")) == 12L; } } diff --git a/server/src/test/java/org/opensearch/search/querygroup/QueryGroupTestHelpers.java b/server/src/test/java/org/opensearch/search/wlm/QueryGroupTestHelpers.java similarity index 94% rename from server/src/test/java/org/opensearch/search/querygroup/QueryGroupTestHelpers.java rename to server/src/test/java/org/opensearch/search/wlm/QueryGroupTestHelpers.java index 67dcbdeae9e20..0403b7c399780 100644 --- a/server/src/test/java/org/opensearch/search/querygroup/QueryGroupTestHelpers.java +++ b/server/src/test/java/org/opensearch/search/wlm/QueryGroupTestHelpers.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.search.querygroup; +package org.opensearch.search.wlm; import org.opensearch.action.search.SearchAction; import org.opensearch.core.tasks.TaskId; diff --git a/server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/search/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java similarity index 91% rename from server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java rename to server/src/test/java/org/opensearch/search/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java index 46132d2dd5277..0b26777e73d2b 100644 --- a/server/src/test/java/org/opensearch/search/querygroup/tracking/QueryGroupResourceUsageTrackerServiceTests.java +++ b/server/src/test/java/org/opensearch/search/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java @@ -6,16 +6,16 @@ * compatible open source license. */ -package org.opensearch.search.querygroup.tracking; +package org.opensearch.search.wlm.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.querygroup.QueryGroupLevelResourceUsageView; -import org.opensearch.search.querygroup.QueryGroupTask; -import org.opensearch.search.querygroup.tracker.QueryGroupResourceUsageTrackerService; -import org.opensearch.search.resourcetypes.ResourceType; +import org.opensearch.search.ResourceType; +import org.opensearch.search.wlm.QueryGroupLevelResourceUsageView; +import org.opensearch.search.wlm.QueryGroupTask; +import org.opensearch.search.wlm.tracker.QueryGroupResourceUsageTrackerService; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; @@ -46,8 +46,8 @@ public class QueryGroupResourceUsageTrackerServiceTests extends OpenSearchTestCa @Before public void setup() { - taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet()); threadPool = new TestThreadPool(getTestName()); + taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet()); mockTaskResourceTrackingService = mock(TaskResourceTrackingService.class); queryGroupResourceUsageTrackerService = new QueryGroupResourceUsageTrackerService(taskManager, mockTaskResourceTrackingService); } @@ -68,9 +68,7 @@ public void testConstructQueryGroupLevelViews_CreatesQueryGroupLevelUsageView_Wh for (String queryGroupId : queryGroupIds) { assertEquals( 400, - (long) stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId) - .getResourceUsageData() - .get(ResourceType.fromName("Memory")) + (long) stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId).getResourceUsageData().get(ResourceType.MEMORY) ); assertEquals(2, stringQueryGroupLevelResourceUsageViewMap.get(queryGroupId).getActiveTasks().size()); } @@ -91,7 +89,7 @@ public void testConstructQueryGroupLevelUsageViews_WithTasksHavingDifferentResou Map queryGroupViews = queryGroupResourceUsageTrackerService .constructQueryGroupLevelUsageViews(); - assertEquals(600, (long) queryGroupViews.get("queryGroup1").getResourceUsageData().get(ResourceType.fromName("Memory"))); + assertEquals(600, (long) queryGroupViews.get("queryGroup1").getResourceUsageData().get(ResourceType.MEMORY)); assertEquals(2, queryGroupViews.get("queryGroup1").getActiveTasks().size()); } From 1907acc74245cb24ddb294d29557f311ad3a5312 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 29 Jul 2024 09:54:27 -0700 Subject: [PATCH 37/66] Update CHANGELOG.md Signed-off-by: Kiran Prakash --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2d2f669f81f2f..2f3d3aba65854 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ## [Unreleased 2.x] ### Added - Fix for hasInitiatedFetching to fix allocation explain and manual reroute APIs (([#14972](https://github.com/opensearch-project/OpenSearch/pull/14972)) +- [Workload Management] QueryGroup resource tracking and cancellation framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) - [Workload Management] Add queryGroupId to Task ([14708](https://github.com/opensearch-project/OpenSearch/pull/14708)) - Add setting to ignore throttling nodes for allocation of unassigned primaries in remote restore ([#14991](https://github.com/opensearch-project/OpenSearch/pull/14991)) - [Streaming Indexing] Enhance RestClient with a new streaming API support ([#14437](https://github.com/opensearch-project/OpenSearch/pull/14437)) From b59ae5b87d23a424944330b68ee0f5fb376af7f4 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 29 Jul 2024 10:04:06 -0700 Subject: [PATCH 38/66] cleanups Signed-off-by: Kiran Prakash --- .../org/opensearch/search/wlm/QueryGroupHelper.java | 13 +++++++++++++ .../org/opensearch/search/wlm/package-info.java | 2 +- 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java b/server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java index 4917ddd9ed29e..2a562474825c8 100644 --- a/server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java +++ b/server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java @@ -14,8 +14,14 @@ import java.util.Map; import java.util.function.Function; +/** + * Helper class for calculating resource usage for different resource types. + */ public class QueryGroupHelper { + /** + * A map that associates each {@link ResourceType} with a function that calculates the resource usage for a given {@link Task}. + */ private static final Map> resourceUsageCalculator = Map.of( ResourceType.MEMORY, (task) -> task.getTotalResourceStats().getMemoryInBytes(), @@ -23,6 +29,13 @@ public class QueryGroupHelper { (task) -> task.getTotalResourceStats().getCpuTimeInNanos() ); + /** + * Gets the resource usage for a given resource type and task. + * + * @param resource the resource type + * @param task the task for which to calculate resource usage + * @return the resource usage + */ public static long getResourceUsage(ResourceType resource, Task task) { return resourceUsageCalculator.get(resource).apply(task); } diff --git a/server/src/main/java/org/opensearch/search/wlm/package-info.java b/server/src/main/java/org/opensearch/search/wlm/package-info.java index a37e5ca764b54..1f1a3605a4a55 100644 --- a/server/src/main/java/org/opensearch/search/wlm/package-info.java +++ b/server/src/main/java/org/opensearch/search/wlm/package-info.java @@ -7,6 +7,6 @@ */ /** - * Query Sandboxing related artifacts + * This package contains workload management constructs */ package org.opensearch.search.wlm; From 46380d5507f5dd3d278148d768a93d94489f9af0 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 29 Jul 2024 10:16:47 -0700 Subject: [PATCH 39/66] Delete QueryGroupService.java Signed-off-by: Kiran Prakash --- .../search/wlm/QueryGroupService.java | 82 ------------------- 1 file changed, 82 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/search/wlm/QueryGroupService.java diff --git a/server/src/main/java/org/opensearch/search/wlm/QueryGroupService.java b/server/src/main/java/org/opensearch/search/wlm/QueryGroupService.java deleted file mode 100644 index 32f5716caea50..0000000000000 --- a/server/src/main/java/org/opensearch/search/wlm/QueryGroupService.java +++ /dev/null @@ -1,82 +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.wlm; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.opensearch.cluster.metadata.Metadata; -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.search.wlm.tracker.QueryGroupUsageTracker; -import org.opensearch.threadpool.Scheduler; -import org.opensearch.threadpool.ThreadPool; - -import java.io.IOException; -import java.util.HashSet; -import java.util.Map; -import java.util.Set; - -/** - * Main service which will run periodically to track and cancel resource constraint violating tasks in QueryGroups - */ -public class QueryGroupService extends AbstractLifecycleComponent { - private static final Logger logger = LogManager.getLogger(QueryGroupService.class); - - private final QueryGroupUsageTracker queryGroupUsageTracker; - private volatile Scheduler.Cancellable scheduledFuture; - private final ThreadPool threadPool; - private final ClusterService clusterService; - - /** - * Guice managed constructor - * - * @param queryGroupUsageTracker tracker service - * @param threadPool threadPool this will be used to schedule the service - */ - @Inject - public QueryGroupService(QueryGroupUsageTracker queryGroupUsageTracker, ClusterService clusterService, ThreadPool threadPool) { - this.queryGroupUsageTracker = queryGroupUsageTracker; - this.clusterService = clusterService; - this.threadPool = threadPool; - } - - /** - * run at regular interval - */ - private void doRun() { - Map queryGroupLevelResourceUsageViews = queryGroupUsageTracker - .constructQueryGroupLevelUsageViews(); - Set activeQueryGroups = getActiveQueryGroups(); - } - - private Set getActiveQueryGroups() { - Map queryGroups = Metadata.builder(clusterService.state().metadata()).getQueryGroups(); - return new HashSet<>(queryGroups.values()); - } - - /** - * {@link AbstractLifecycleComponent} lifecycle method - */ - @Override - protected void doStart() { - // This method is intentionally left empty to ensure the service is never scheduled - } - - @Override - protected void doStop() { - if (scheduledFuture != null) { - scheduledFuture.cancel(); - } - } - - @Override - protected void doClose() throws IOException {} -} From 4759c8d67012fe0d8893595d009178df054af5da Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 29 Jul 2024 10:19:58 -0700 Subject: [PATCH 40/66] cleanups Signed-off-by: Kiran Prakash --- CHANGELOG.md | 2 +- .../search/wlm/module/QueryGroupModule.java | 29 ------------------- .../search/wlm/module/package-info.java | 13 --------- 3 files changed, 1 insertion(+), 43 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/search/wlm/module/QueryGroupModule.java delete mode 100644 server/src/main/java/org/opensearch/search/wlm/module/package-info.java diff --git a/CHANGELOG.md b/CHANGELOG.md index 2f3d3aba65854..28c78bedcdb07 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ## [Unreleased 2.x] ### Added - Fix for hasInitiatedFetching to fix allocation explain and manual reroute APIs (([#14972](https://github.com/opensearch-project/OpenSearch/pull/14972)) -- [Workload Management] QueryGroup resource tracking and cancellation framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) +- [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) - [Workload Management] Add queryGroupId to Task ([14708](https://github.com/opensearch-project/OpenSearch/pull/14708)) - Add setting to ignore throttling nodes for allocation of unassigned primaries in remote restore ([#14991](https://github.com/opensearch-project/OpenSearch/pull/14991)) - [Streaming Indexing] Enhance RestClient with a new streaming API support ([#14437](https://github.com/opensearch-project/OpenSearch/pull/14437)) diff --git a/server/src/main/java/org/opensearch/search/wlm/module/QueryGroupModule.java b/server/src/main/java/org/opensearch/search/wlm/module/QueryGroupModule.java deleted file mode 100644 index ddc6954dd412d..0000000000000 --- a/server/src/main/java/org/opensearch/search/wlm/module/QueryGroupModule.java +++ /dev/null @@ -1,29 +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.wlm.module; - -import org.opensearch.common.inject.AbstractModule; -import org.opensearch.search.wlm.tracker.QueryGroupResourceUsageTrackerService; -import org.opensearch.search.wlm.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(); - } -} diff --git a/server/src/main/java/org/opensearch/search/wlm/module/package-info.java b/server/src/main/java/org/opensearch/search/wlm/module/package-info.java deleted file mode 100644 index aadf19448caef..0000000000000 --- a/server/src/main/java/org/opensearch/search/wlm/module/package-info.java +++ /dev/null @@ -1,13 +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. - */ - -/** - * Guice Module - */ - -package org.opensearch.search.wlm.module; From 464db0e3ac573c0352a5d499675f471920209af6 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 29 Jul 2024 11:24:43 -0700 Subject: [PATCH 41/66] Update QueryGroupLevelResourceUsageViewTests.java Signed-off-by: Kiran Prakash --- .../search/wlm/QueryGroupLevelResourceUsageViewTests.java | 1 + 1 file changed, 1 insertion(+) diff --git a/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java index 98ff82248dd9a..bfb35a8996746 100644 --- a/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java +++ b/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java @@ -27,6 +27,7 @@ public void setUp() throws Exception { activeTasks = List.of(getRandomTask(4321)); } + public void testGetResourceUsageData() { QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( "1234", From af8f84e667c3fb0ab18d48de47d957a9f49781d6 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 29 Jul 2024 11:24:53 -0700 Subject: [PATCH 42/66] Update QueryGroupLevelResourceUsageViewTests.java Signed-off-by: Kiran Prakash --- .../search/wlm/QueryGroupLevelResourceUsageViewTests.java | 1 - 1 file changed, 1 deletion(-) diff --git a/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java index bfb35a8996746..98ff82248dd9a 100644 --- a/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java +++ b/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java @@ -27,7 +27,6 @@ public void setUp() throws Exception { activeTasks = List.of(getRandomTask(4321)); } - public void testGetResourceUsageData() { QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( "1234", From 554db7f93be3e829b3cb7dd6c577bf9539b88bf5 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Mon, 29 Jul 2024 11:37:58 -0700 Subject: [PATCH 43/66] Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash --- .../wlm/tracker/QueryGroupResourceUsageTrackerService.java | 2 -- 1 file changed, 2 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 01bea9077626a..d3510919983b2 100644 --- a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -8,7 +8,6 @@ package org.opensearch.search.wlm.tracker; -import org.opensearch.common.inject.Inject; import org.opensearch.search.ResourceType; import org.opensearch.search.wlm.QueryGroupHelper; import org.opensearch.search.wlm.QueryGroupLevelResourceUsageView; @@ -39,7 +38,6 @@ public class QueryGroupResourceUsageTrackerService implements QueryGroupUsageTra * @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 From a4d39218ae8123f13b5d379696a4ac73b4f466ea Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Tue, 30 Jul 2024 14:04:49 -0700 Subject: [PATCH 44/66] Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash --- ...QueryGroupResourceUsageTrackerService.java | 28 ++++++++----------- 1 file changed, 12 insertions(+), 16 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java index d3510919983b2..6dec836aa02ca 100644 --- a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -20,6 +20,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.function.Function; import java.util.stream.Collectors; /** @@ -90,28 +91,23 @@ private Map> getTasksGroupedByQueryGroup() { * @return Map of resource usage for each QueryGroup */ private Map> getResourceUsageOfQueryGroups(Map> tasksByQueryGroup) { - Map> resourceUsageOfQueryGroups = new HashMap<>(); + // Prepare a usage map for the QueryGroups + Map> resourceUsageOfQueryGroups = tasksByQueryGroup.keySet() + .stream() + .collect(Collectors.toMap(Function.identity(), k -> new EnumMap<>(ResourceType.class))); // 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 - EnumMap queryGroupUsage = resourceUsageOfQueryGroups.computeIfAbsent( - queryGroupId, - k -> new EnumMap<>(ResourceType.class) - ); - - // 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 = QueryGroupHelper.getResourceUsage(resourceType, task); - queryGroupUsage.put(resourceType, currentUsage + taskUsage); + EnumMap queryGroupUsage = resourceUsageOfQueryGroups.get(queryGroupEntry.getKey()); + for (ResourceType resourceType : TRACKED_RESOURCES) { + long queryGroupResourceUsage = 0; + for (Task task : queryGroupEntry.getValue()) { + queryGroupResourceUsage += QueryGroupHelper.getResourceUsage(resourceType, task); } + queryGroupUsage.put(resourceType, queryGroupResourceUsage); } } + return resourceUsageOfQueryGroups; } From 59c51fd80bd805f9d6c1b21f864152792206a9c4 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 31 Jul 2024 06:13:13 -0700 Subject: [PATCH 45/66] Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash --- ...QueryGroupResourceUsageTrackerService.java | 56 ++++++------------- 1 file changed, 18 insertions(+), 38 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 6dec836aa02ca..749c5c4b8e2ca 100644 --- a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -20,7 +20,6 @@ import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.function.Function; import java.util.stream.Collectors; /** @@ -29,7 +28,7 @@ // @ExperimentalApi public class QueryGroupResourceUsageTrackerService implements QueryGroupUsageTracker, TaskManager.TaskEventListeners { - public static final List TRACKED_RESOURCES = List.of(ResourceType.fromName("memory"), ResourceType.fromName("cpu")); + public static final List TRACKED_RESOURCES = List.of(ResourceType.MEMORY, ResourceType.CPU); private final TaskManager taskManager; private final TaskResourceTrackingService taskResourceTrackingService; @@ -54,18 +53,26 @@ public QueryGroupResourceUsageTrackerService( */ @Override public Map constructQueryGroupLevelUsageViews() { - Map queryGroupViews = new HashMap<>(); + final Map> tasksByQueryGroup = getTasksGroupedByQueryGroup(); + final Map queryGroupViews = new HashMap<>(); - Map> tasksByQueryGroup = getTasksGroupedByQueryGroup(); - Map> queryGroupResourceUsage = getResourceUsageOfQueryGroups(tasksByQueryGroup); + // Iterate over each QueryGroup entry + for (Map.Entry> queryGroupEntry : tasksByQueryGroup.entrySet()) { + // Compute the QueryGroup usage + final EnumMap queryGroupUsage = new EnumMap<>(ResourceType.class); + for (ResourceType resourceType : TRACKED_RESOURCES) { + long queryGroupResourceUsage = 0; + for (Task task : queryGroupEntry.getValue()) { + queryGroupResourceUsage += QueryGroupHelper.getResourceUsage(resourceType, task); + } + queryGroupUsage.put(resourceType, queryGroupResourceUsage); + } - for (String queryGroupId : tasksByQueryGroup.keySet()) { - QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( - queryGroupId, - queryGroupResourceUsage.get(queryGroupId), - tasksByQueryGroup.get(queryGroupId) + // Add to the QueryGroup View + queryGroupViews.put( + queryGroupEntry.getKey(), + new QueryGroupLevelResourceUsageView(queryGroupEntry.getKey(), queryGroupUsage, queryGroupEntry.getValue()) ); - queryGroupViews.put(queryGroupId, queryGroupLevelResourceUsageView); } return queryGroupViews; } @@ -84,33 +91,6 @@ private Map> getTasksGroupedByQueryGroup() { .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) { - // Prepare a usage map for the QueryGroups - Map> resourceUsageOfQueryGroups = tasksByQueryGroup.keySet() - .stream() - .collect(Collectors.toMap(Function.identity(), k -> new EnumMap<>(ResourceType.class))); - - // Iterate over each QueryGroup entry - for (Map.Entry> queryGroupEntry : tasksByQueryGroup.entrySet()) { - EnumMap queryGroupUsage = resourceUsageOfQueryGroups.get(queryGroupEntry.getKey()); - for (ResourceType resourceType : TRACKED_RESOURCES) { - long queryGroupResourceUsage = 0; - for (Task task : queryGroupEntry.getValue()) { - queryGroupResourceUsage += QueryGroupHelper.getResourceUsage(resourceType, task); - } - queryGroupUsage.put(resourceType, queryGroupResourceUsage); - } - } - - return resourceUsageOfQueryGroups; - } - /** * Handles the completion of a task. * From 3ff68a279d312db98626b8cb1ebde8f03fc0879d Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 31 Jul 2024 06:20:10 -0700 Subject: [PATCH 46/66] Update CHANGELOG.md Signed-off-by: Kiran Prakash --- CHANGELOG.md | 3 --- 1 file changed, 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 28c78bedcdb07..2818134f2155b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -13,9 +13,6 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Add basic aggregation support for derived fields ([#14618](https://github.com/opensearch-project/OpenSearch/pull/14618)) - Add ThreadContextPermission for markAsSystemContext and allow core to perform the method ([#15016](https://github.com/opensearch-project/OpenSearch/pull/15016)) - Add ThreadContextPermission for stashAndMergeHeaders and stashWithOrigin ([#15039](https://github.com/opensearch-project/OpenSearch/pull/15039)) -- QueryGroup resource tracking and cancellation framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) -- [Concurrent Segment Search] Support composite aggregations with scripting ([#15072](https://github.com/opensearch-project/OpenSearch/pull/15072)) -- Add `rangeQuery` and `regexpQuery` for `constant_keyword` field type ([#14711](https://github.com/opensearch-project/OpenSearch/pull/14711)) ### Dependencies - Bump `netty` from 4.1.111.Final to 4.1.112.Final ([#15081](https://github.com/opensearch-project/OpenSearch/pull/15081)) From 0dd15f41b9abb2cb6ed2eec065e94cedf943a488 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 31 Jul 2024 11:55:16 -0700 Subject: [PATCH 47/66] rebasing with latest main Signed-off-by: Kiran Prakash --- .../search/wlm/QueryGroupConstants.java | 19 ----- .../opensearch/search/wlm/QueryGroupTask.java | 74 ------------------- .../opensearch/search/wlm/package-info.java | 12 --- .../{search => }/wlm/QueryGroupHelper.java | 2 +- .../wlm/QueryGroupLevelResourceUsageView.java | 2 +- ...QueryGroupResourceUsageTrackerService.java | 8 +- .../wlm/tracker/QueryGroupUsageTracker.java | 4 +- .../wlm/tracker/package-info.java | 2 +- ...QueryGroupLevelResourceUsageViewTests.java | 4 +- .../wlm/QueryGroupTestHelpers.java | 2 +- ...GroupResourceUsageTrackerServiceTests.java | 17 +++-- 11 files changed, 24 insertions(+), 122 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/search/wlm/QueryGroupConstants.java delete mode 100644 server/src/main/java/org/opensearch/search/wlm/QueryGroupTask.java delete mode 100644 server/src/main/java/org/opensearch/search/wlm/package-info.java rename server/src/main/java/org/opensearch/{search => }/wlm/QueryGroupHelper.java (97%) rename server/src/main/java/org/opensearch/{search => }/wlm/QueryGroupLevelResourceUsageView.java (98%) rename server/src/main/java/org/opensearch/{search => }/wlm/tracker/QueryGroupResourceUsageTrackerService.java (94%) rename server/src/main/java/org/opensearch/{search => }/wlm/tracker/QueryGroupUsageTracker.java (82%) rename server/src/main/java/org/opensearch/{search => }/wlm/tracker/package-info.java (85%) rename server/src/test/java/org/opensearch/{search => }/wlm/QueryGroupLevelResourceUsageViewTests.java (95%) rename server/src/test/java/org/opensearch/{search => }/wlm/QueryGroupTestHelpers.java (95%) rename server/src/test/java/org/opensearch/{search => }/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java (87%) diff --git a/server/src/main/java/org/opensearch/search/wlm/QueryGroupConstants.java b/server/src/main/java/org/opensearch/search/wlm/QueryGroupConstants.java deleted file mode 100644 index b64e65c6836bc..0000000000000 --- a/server/src/main/java/org/opensearch/search/wlm/QueryGroupConstants.java +++ /dev/null @@ -1,19 +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.wlm; - -import java.util.function.Supplier; - -/** - * This class will hold all the QueryGroup related constants - */ -public class QueryGroupConstants { - public static final String QUERY_GROUP_ID_HEADER = "queryGroupId"; - public static final Supplier DEFAULT_QUERY_GROUP_ID_SUPPLIER = () -> "DEFAULT_QUERY_GROUP"; -} diff --git a/server/src/main/java/org/opensearch/search/wlm/QueryGroupTask.java b/server/src/main/java/org/opensearch/search/wlm/QueryGroupTask.java deleted file mode 100644 index e8b5ca4efd4e3..0000000000000 --- a/server/src/main/java/org/opensearch/search/wlm/QueryGroupTask.java +++ /dev/null @@ -1,74 +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.wlm; - -import org.apache.logging.log4j.LogManager; -import org.apache.logging.log4j.Logger; -import org.opensearch.common.unit.TimeValue; -import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.core.tasks.TaskId; -import org.opensearch.tasks.CancellableTask; - -import java.util.Map; - -import static org.opensearch.search.SearchService.NO_TIMEOUT; - -/** - * Base class to define QueryGroup tasks - */ -public class QueryGroupTask extends CancellableTask { - - private static final Logger logger = LogManager.getLogger(QueryGroupTask.class); - private String queryGroupId; - - public QueryGroupTask(long id, String type, String action, String description, TaskId parentTaskId, Map headers) { - this(id, type, action, description, parentTaskId, headers, NO_TIMEOUT); - } - - public QueryGroupTask( - long id, - String type, - String action, - String description, - TaskId parentTaskId, - Map headers, - TimeValue cancelAfterTimeInterval - ) { - super(id, type, action, description, parentTaskId, headers, cancelAfterTimeInterval); - } - - /** - * This method should always be called after calling setQueryGroupId at least once on this object - * @return task queryGroupId - */ - public String getQueryGroupId() { - if (queryGroupId == null) { - logger.warn("QueryGroup _id can't be null, It should be set before accessing it. This is abnormal behaviour "); - } - return queryGroupId; - } - - /** - * sets the queryGroupId from threadContext into the task itself, - * This method was defined since the queryGroupId can only be evaluated after task creation - * @param threadContext current threadContext - */ - public void setQueryGroupId(final ThreadContext threadContext) { - this.queryGroupId = QueryGroupConstants.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(); - - if (threadContext != null && threadContext.getHeader(QueryGroupConstants.QUERY_GROUP_ID_HEADER) != null) { - this.queryGroupId = threadContext.getHeader(QueryGroupConstants.QUERY_GROUP_ID_HEADER); - } - } - - @Override - public boolean shouldCancelChildrenOnCancellation() { - return false; - } -} diff --git a/server/src/main/java/org/opensearch/search/wlm/package-info.java b/server/src/main/java/org/opensearch/search/wlm/package-info.java deleted file mode 100644 index 1f1a3605a4a55..0000000000000 --- a/server/src/main/java/org/opensearch/search/wlm/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. - */ - -/** - * This package contains workload management constructs - */ -package org.opensearch.search.wlm; diff --git a/server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java b/server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java similarity index 97% rename from server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java rename to server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java index 2a562474825c8..13b87b5010bf3 100644 --- a/server/src/main/java/org/opensearch/search/wlm/QueryGroupHelper.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.search.wlm; +package org.opensearch.wlm; import org.opensearch.search.ResourceType; import org.opensearch.tasks.Task; diff --git a/server/src/main/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageView.java b/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java similarity index 98% rename from server/src/main/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageView.java rename to server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java index faafc64ae2860..a01dad2c96697 100644 --- a/server/src/main/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageView.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.search.wlm; +package org.opensearch.wlm; import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.search.ResourceType; diff --git a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java similarity index 94% rename from server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java rename to server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 749c5c4b8e2ca..e702699398b22 100644 --- a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -6,15 +6,15 @@ * compatible open source license. */ -package org.opensearch.search.wlm.tracker; +package org.opensearch.wlm.tracker; import org.opensearch.search.ResourceType; -import org.opensearch.search.wlm.QueryGroupHelper; -import org.opensearch.search.wlm.QueryGroupLevelResourceUsageView; -import org.opensearch.search.wlm.QueryGroupTask; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; import org.opensearch.tasks.TaskResourceTrackingService; +import org.opensearch.wlm.QueryGroupHelper; +import org.opensearch.wlm.QueryGroupLevelResourceUsageView; +import org.opensearch.wlm.QueryGroupTask; import java.util.EnumMap; import java.util.HashMap; diff --git a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupUsageTracker.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupUsageTracker.java similarity index 82% rename from server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupUsageTracker.java rename to server/src/main/java/org/opensearch/wlm/tracker/QueryGroupUsageTracker.java index 269971780daa9..23fb8b1b45aac 100644 --- a/server/src/main/java/org/opensearch/search/wlm/tracker/QueryGroupUsageTracker.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupUsageTracker.java @@ -6,9 +6,9 @@ * compatible open source license. */ -package org.opensearch.search.wlm.tracker; +package org.opensearch.wlm.tracker; -import org.opensearch.search.wlm.QueryGroupLevelResourceUsageView; +import org.opensearch.wlm.QueryGroupLevelResourceUsageView; import java.util.Map; diff --git a/server/src/main/java/org/opensearch/search/wlm/tracker/package-info.java b/server/src/main/java/org/opensearch/wlm/tracker/package-info.java similarity index 85% rename from server/src/main/java/org/opensearch/search/wlm/tracker/package-info.java rename to server/src/main/java/org/opensearch/wlm/tracker/package-info.java index 1012749a845f8..86efc99355d3d 100644 --- a/server/src/main/java/org/opensearch/search/wlm/tracker/package-info.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/package-info.java @@ -9,4 +9,4 @@ /** * QueryGroup resource tracking artifacts */ -package org.opensearch.search.wlm.tracker; +package org.opensearch.wlm.tracker; diff --git a/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java similarity index 95% rename from server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java rename to server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java index 98ff82248dd9a..d9fd5b35c88b8 100644 --- a/server/src/test/java/org/opensearch/search/wlm/QueryGroupLevelResourceUsageViewTests.java +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.search.wlm; +package org.opensearch.wlm; import org.opensearch.search.ResourceType; import org.opensearch.tasks.Task; @@ -15,7 +15,7 @@ import java.util.List; import java.util.Map; -import static org.opensearch.search.wlm.QueryGroupTestHelpers.getRandomTask; +import static org.opensearch.wlm.QueryGroupTestHelpers.getRandomTask; public class QueryGroupLevelResourceUsageViewTests extends OpenSearchTestCase { Map resourceUsage; diff --git a/server/src/test/java/org/opensearch/search/wlm/QueryGroupTestHelpers.java b/server/src/test/java/org/opensearch/wlm/QueryGroupTestHelpers.java similarity index 95% rename from server/src/test/java/org/opensearch/search/wlm/QueryGroupTestHelpers.java rename to server/src/test/java/org/opensearch/wlm/QueryGroupTestHelpers.java index 0403b7c399780..b3fda77c662b2 100644 --- a/server/src/test/java/org/opensearch/search/wlm/QueryGroupTestHelpers.java +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupTestHelpers.java @@ -6,7 +6,7 @@ * compatible open source license. */ -package org.opensearch.search.wlm; +package org.opensearch.wlm; import org.opensearch.action.search.SearchAction; import org.opensearch.core.tasks.TaskId; diff --git a/server/src/test/java/org/opensearch/search/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java similarity index 87% rename from server/src/test/java/org/opensearch/search/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java rename to server/src/test/java/org/opensearch/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java index 0b26777e73d2b..1c1893c4cef37 100644 --- a/server/src/test/java/org/opensearch/search/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java @@ -6,16 +6,14 @@ * compatible open source license. */ -package org.opensearch.search.wlm.tracking; +package org.opensearch.wlm.tracking; import org.opensearch.action.search.SearchShardTask; import org.opensearch.action.search.SearchTask; import org.opensearch.common.settings.Settings; +import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; import org.opensearch.search.ResourceType; -import org.opensearch.search.wlm.QueryGroupLevelResourceUsageView; -import org.opensearch.search.wlm.QueryGroupTask; -import org.opensearch.search.wlm.tracker.QueryGroupResourceUsageTrackerService; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; import org.opensearch.tasks.TaskManager; @@ -23,6 +21,9 @@ import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; +import org.opensearch.wlm.QueryGroupLevelResourceUsageView; +import org.opensearch.wlm.QueryGroupTask; +import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService; import org.junit.After; import org.junit.Before; @@ -33,6 +34,7 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import static org.opensearch.wlm.QueryGroupTask.QUERY_GROUP_ID_HEADER; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.mock; @@ -107,7 +109,12 @@ private Map createActiveSearchShardTasks(List queryGroupIds) private T createMockTask(Class type, long cpuUsage, long heapUsage, String queryGroupId) { T task = mock(type); if (task instanceof SearchTask || task instanceof SearchShardTask) { - when(((QueryGroupTask) task).getQueryGroupId()).thenReturn(queryGroupId); + // Stash the current thread context to ensure that any existing context is preserved and restored after setting the query group + // ID. + try (ThreadContext.StoredContext ignore = threadPool.getThreadContext().stashContext()) { + threadPool.getThreadContext().putHeader(QUERY_GROUP_ID_HEADER, queryGroupId); + ((QueryGroupTask) task).setQueryGroupId(threadPool.getThreadContext()); + } } when(task.getTotalResourceStats()).thenReturn(new TaskResourceUsage(cpuUsage, heapUsage)); when(task.getStartTimeNanos()).thenReturn((long) 0); From f1383f1a37a5df95fc056e11b77a3714c8a45029 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 31 Jul 2024 14:01:48 -0700 Subject: [PATCH 48/66] remove experimental Signed-off-by: Kiran Prakash --- .../org/opensearch/wlm/QueryGroupLevelResourceUsageView.java | 1 - .../wlm/tracker/QueryGroupResourceUsageTrackerService.java | 1 - 2 files changed, 2 deletions(-) diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java b/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java index a01dad2c96697..3786688fe1ddf 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java @@ -23,7 +23,6 @@ * 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 QueryGroupLevelResourceUsageView { private final String queryGroupId; diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index e702699398b22..1c3a9405ccdb7 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -25,7 +25,6 @@ /** * This class tracks resource usage per QueryGroup */ -// @ExperimentalApi public class QueryGroupResourceUsageTrackerService implements QueryGroupUsageTracker, TaskManager.TaskEventListeners { public static final List TRACKED_RESOURCES = List.of(ResourceType.MEMORY, ResourceType.CPU); From ad3afde3d8cd76766cdd1729781c4aafc55740a8 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 31 Jul 2024 14:24:54 -0700 Subject: [PATCH 49/66] remove queryGroupId Signed-off-by: Kiran Prakash --- .../wlm/QueryGroupLevelResourceUsageView.java | 30 ++----------------- ...QueryGroupLevelResourceUsageViewTests.java | 14 --------- 2 files changed, 2 insertions(+), 42 deletions(-) diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java b/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java index 3786688fe1ddf..2fd743dc3f83f 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupLevelResourceUsageView.java @@ -8,37 +8,24 @@ package org.opensearch.wlm; -import org.opensearch.common.annotation.ExperimentalApi; import org.opensearch.search.ResourceType; import org.opensearch.tasks.Task; -import java.util.ArrayList; -import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Objects; /** * 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. + * This class holds the resource usage data and the list of active tasks. */ public class QueryGroupLevelResourceUsageView { - - 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 QueryGroupLevelResourceUsageView(String queryGroupId) { - this.queryGroupId = queryGroupId; - this.resourceUsage = new HashMap<>(); - this.activeTasks = new ArrayList<>(); - } - - public QueryGroupLevelResourceUsageView(String queryGroupId, Map resourceUsage, List activeTasks) { - this.queryGroupId = queryGroupId; + public QueryGroupLevelResourceUsageView(Map resourceUsage, List activeTasks) { this.resourceUsage = resourceUsage; this.activeTasks = activeTasks; } @@ -60,17 +47,4 @@ public Map getResourceUsageData() { public List getActiveTasks() { return activeTasks; } - - @Override - public boolean equals(Object o) { - if (this == o) return true; - if (o == null || getClass() != o.getClass()) return false; - QueryGroupLevelResourceUsageView that = (QueryGroupLevelResourceUsageView) o; - return Objects.equals(queryGroupId, that.queryGroupId); - } - - @Override - public int hashCode() { - return Objects.hashCode(queryGroupId); - } } diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java index d9fd5b35c88b8..c965957516bad 100644 --- a/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java @@ -29,7 +29,6 @@ public void setUp() throws Exception { public void testGetResourceUsageData() { QueryGroupLevelResourceUsageView queryGroupLevelResourceUsageView = new QueryGroupLevelResourceUsageView( - "1234", resourceUsage, activeTasks ); @@ -37,15 +36,8 @@ public void testGetResourceUsageData() { 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 ); @@ -54,12 +46,6 @@ public void testGetActiveTasks() { 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("memory")) == 34L && resourceUsageData.get(ResourceType.fromName("cpu")) == 12L; } From 4abe48eecc0baa3bf264161d4d3f10eeaa7fd5d5 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 31 Jul 2024 14:28:05 -0700 Subject: [PATCH 50/66] Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash --- .../wlm/tracker/QueryGroupResourceUsageTrackerService.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 1c3a9405ccdb7..4a7a46b75c3bd 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -70,7 +70,7 @@ public Map constructQueryGroupLevelUsa // Add to the QueryGroup View queryGroupViews.put( queryGroupEntry.getKey(), - new QueryGroupLevelResourceUsageView(queryGroupEntry.getKey(), queryGroupUsage, queryGroupEntry.getValue()) + new QueryGroupLevelResourceUsageView(queryGroupUsage, queryGroupEntry.getValue()) ); } return queryGroupViews; From 03fed4d9cac1aadd1cd285d6717b2a47c64b8bcb Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 31 Jul 2024 14:55:22 -0700 Subject: [PATCH 51/66] change code comments Signed-off-by: Kiran Prakash --- server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java b/server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java index 13b87b5010bf3..02158accea5ac 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java @@ -15,7 +15,7 @@ import java.util.function.Function; /** - * Helper class for calculating resource usage for different resource types. + * Helper class for containing static utility methods related to QueryGroups. */ public class QueryGroupHelper { From 03e0f4ac51e65227d5fe01ec489f23d7a6e54d45 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 1 Aug 2024 11:30:41 -0700 Subject: [PATCH 52/66] remmove QueryGroupUsageTracker Signed-off-by: Kiran Prakash --- ...QueryGroupResourceUsageTrackerService.java | 3 +-- .../wlm/tracker/QueryGroupUsageTracker.java | 24 ------------------- 2 files changed, 1 insertion(+), 26 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/wlm/tracker/QueryGroupUsageTracker.java diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 4a7a46b75c3bd..d25d6281a8670 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -25,7 +25,7 @@ /** * This class tracks resource usage per QueryGroup */ -public class QueryGroupResourceUsageTrackerService implements QueryGroupUsageTracker, TaskManager.TaskEventListeners { +public class QueryGroupResourceUsageTrackerService implements TaskManager.TaskEventListeners { public static final List TRACKED_RESOURCES = List.of(ResourceType.MEMORY, ResourceType.CPU); private final TaskManager taskManager; @@ -50,7 +50,6 @@ public QueryGroupResourceUsageTrackerService( * * @return Map of QueryGroup views */ - @Override public Map constructQueryGroupLevelUsageViews() { final Map> tasksByQueryGroup = getTasksGroupedByQueryGroup(); final Map queryGroupViews = new HashMap<>(); diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupUsageTracker.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupUsageTracker.java deleted file mode 100644 index 23fb8b1b45aac..0000000000000 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupUsageTracker.java +++ /dev/null @@ -1,24 +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.wlm.tracker; - -import org.opensearch.wlm.QueryGroupLevelResourceUsageView; - -import java.util.Map; - -/** - * This interface is mainly for tracking the queryGroup level resource usages - */ -public interface QueryGroupUsageTracker { - /** - * updates the current resource usage of queryGroup - */ - - Map constructQueryGroupLevelUsageViews(); -} From 0281cc5c48ee21c75758fbb6aa16d70ea3e885b5 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 1 Aug 2024 12:14:26 -0700 Subject: [PATCH 53/66] Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash --- .../wlm/tracker/QueryGroupResourceUsageTrackerService.java | 1 + 1 file changed, 1 insertion(+) diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index d25d6281a8670..6740579e70814 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -16,6 +16,7 @@ import org.opensearch.wlm.QueryGroupLevelResourceUsageView; import org.opensearch.wlm.QueryGroupTask; + import java.util.EnumMap; import java.util.HashMap; import java.util.List; From ffed5a9a7d759236ed6e6f5702ba3fe48e736f41 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 1 Aug 2024 12:14:36 -0700 Subject: [PATCH 54/66] Update QueryGroupResourceUsageTrackerService.java Signed-off-by: Kiran Prakash --- .../wlm/tracker/QueryGroupResourceUsageTrackerService.java | 1 - 1 file changed, 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 6740579e70814..d25d6281a8670 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -16,7 +16,6 @@ import org.opensearch.wlm.QueryGroupLevelResourceUsageView; import org.opensearch.wlm.QueryGroupTask; - import java.util.EnumMap; import java.util.HashMap; import java.util.List; From 21f37a6fc0e30861ffc177bf139369b25ba795c2 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 1 Aug 2024 15:21:12 -0700 Subject: [PATCH 55/66] remove QueryGroupTestHelpers Signed-off-by: Kiran Prakash --- ...QueryGroupLevelResourceUsageViewTests.java | 16 ++++++++-- .../opensearch/wlm/QueryGroupTestHelpers.java | 31 ------------------- 2 files changed, 14 insertions(+), 33 deletions(-) delete mode 100644 server/src/test/java/org/opensearch/wlm/QueryGroupTestHelpers.java diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java index c965957516bad..7f6419505fec2 100644 --- a/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupLevelResourceUsageViewTests.java @@ -8,15 +8,16 @@ package org.opensearch.wlm; +import org.opensearch.action.search.SearchAction; +import org.opensearch.core.tasks.TaskId; import org.opensearch.search.ResourceType; import org.opensearch.tasks.Task; import org.opensearch.test.OpenSearchTestCase; +import java.util.Collections; import java.util.List; import java.util.Map; -import static org.opensearch.wlm.QueryGroupTestHelpers.getRandomTask; - public class QueryGroupLevelResourceUsageViewTests extends OpenSearchTestCase { Map resourceUsage; List activeTasks; @@ -49,4 +50,15 @@ public void testGetActiveTasks() { private boolean assertResourceUsageData(Map resourceUsageData) { return resourceUsageData.get(ResourceType.fromName("memory")) == 34L && resourceUsageData.get(ResourceType.fromName("cpu")) == 12L; } + + private Task getRandomTask(long id) { + return new Task( + id, + "transport", + SearchAction.NAME, + "test description", + new TaskId(randomLong() + ":" + randomLong()), + Collections.emptyMap() + ); + } } diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupTestHelpers.java b/server/src/test/java/org/opensearch/wlm/QueryGroupTestHelpers.java deleted file mode 100644 index b3fda77c662b2..0000000000000 --- a/server/src/test/java/org/opensearch/wlm/QueryGroupTestHelpers.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.wlm; - -import org.opensearch.action.search.SearchAction; -import org.opensearch.core.tasks.TaskId; -import org.opensearch.tasks.Task; - -import java.util.Collections; - -import static org.opensearch.test.OpenSearchTestCase.randomLong; - -public class QueryGroupTestHelpers { - - public static Task getRandomTask(long id) { - return new Task( - id, - "transport", - SearchAction.NAME, - "test description", - new TaskId(randomLong() + ":" + randomLong()), - Collections.emptyMap() - ); - } -} From 01f74230608c192d5a5f478ceb8d89c45cb7a616 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Thu, 1 Aug 2024 15:42:42 -0700 Subject: [PATCH 56/66] cleanups Signed-off-by: Kiran Prakash --- .../tracker/QueryGroupResourceUsageTrackerService.java | 8 +------- .../QueryGroupResourceUsageTrackerServiceTests.java | 10 ++-------- 2 files changed, 3 insertions(+), 15 deletions(-) rename server/src/test/java/org/opensearch/wlm/{tracking => tracker}/QueryGroupResourceUsageTrackerServiceTests.java (93%) diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index d25d6281a8670..89d7e18f01c65 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -28,20 +28,14 @@ public class QueryGroupResourceUsageTrackerService implements TaskManager.TaskEventListeners { public static final List TRACKED_RESOURCES = List.of(ResourceType.MEMORY, ResourceType.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. */ - public QueryGroupResourceUsageTrackerService( - final TaskManager taskManager, - final TaskResourceTrackingService taskResourceTrackingService - ) { - this.taskManager = taskManager; + public QueryGroupResourceUsageTrackerService(TaskResourceTrackingService taskResourceTrackingService) { this.taskResourceTrackingService = taskResourceTrackingService; } diff --git a/server/src/test/java/org/opensearch/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java similarity index 93% rename from server/src/test/java/org/opensearch/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java rename to server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java index 1c1893c4cef37..6702a3948eb0e 100644 --- a/server/src/test/java/org/opensearch/wlm/tracking/QueryGroupResourceUsageTrackerServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java @@ -6,28 +6,24 @@ * compatible open source license. */ -package org.opensearch.wlm.tracking; +package org.opensearch.wlm.tracker; import org.opensearch.action.search.SearchShardTask; import org.opensearch.action.search.SearchTask; -import org.opensearch.common.settings.Settings; import org.opensearch.common.util.concurrent.ThreadContext; import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; import org.opensearch.search.ResourceType; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; -import org.opensearch.tasks.TaskManager; import org.opensearch.tasks.TaskResourceTrackingService; import org.opensearch.test.OpenSearchTestCase; import org.opensearch.threadpool.TestThreadPool; import org.opensearch.threadpool.ThreadPool; import org.opensearch.wlm.QueryGroupLevelResourceUsageView; import org.opensearch.wlm.QueryGroupTask; -import org.opensearch.wlm.tracker.QueryGroupResourceUsageTrackerService; import org.junit.After; import org.junit.Before; -import java.util.Collections; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -42,16 +38,14 @@ public class QueryGroupResourceUsageTrackerServiceTests extends OpenSearchTestCase { TestThreadPool threadPool; - TaskManager taskManager; TaskResourceTrackingService mockTaskResourceTrackingService; QueryGroupResourceUsageTrackerService queryGroupResourceUsageTrackerService; @Before public void setup() { threadPool = new TestThreadPool(getTestName()); - taskManager = new TaskManager(Settings.EMPTY, threadPool, Collections.emptySet()); mockTaskResourceTrackingService = mock(TaskResourceTrackingService.class); - queryGroupResourceUsageTrackerService = new QueryGroupResourceUsageTrackerService(taskManager, mockTaskResourceTrackingService); + queryGroupResourceUsageTrackerService = new QueryGroupResourceUsageTrackerService(mockTaskResourceTrackingService); } @After From 6eaa39ed0e6388a8eda40c42d99eba04b12ec14f Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Fri, 2 Aug 2024 16:45:19 -0700 Subject: [PATCH 57/66] remove queryGroupHelper Signed-off-by: Kiran Prakash --- .../org/opensearch/search/ResourceType.java | 20 +++++-- .../org/opensearch/wlm/QueryGroupHelper.java | 42 --------------- ...QueryGroupResourceUsageTrackerService.java | 14 +---- .../opensearch/search/ResourceTypeTests.java | 52 +++++++++++++++++++ 4 files changed, 71 insertions(+), 57 deletions(-) delete mode 100644 server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java create mode 100644 server/src/test/java/org/opensearch/search/ResourceTypeTests.java diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java index fe5ce4dd2bb50..8d47e1f72899e 100644 --- a/server/src/main/java/org/opensearch/search/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/ResourceType.java @@ -10,21 +10,25 @@ import org.opensearch.common.annotation.PublicApi; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.tasks.Task; import java.io.IOException; +import java.util.function.Function; /** * Enum to hold the resource type */ @PublicApi(since = "2.x") public enum ResourceType { - CPU("cpu"), - MEMORY("memory"); + CPU("cpu", task -> task.getTotalResourceStats().getCpuTimeInNanos()), + MEMORY("memory", task -> task.getTotalResourceStats().getMemoryInBytes()); private final String name; + private final Function getResourceUsage; - ResourceType(String name) { + ResourceType(String name, Function getResourceUsage) { this.name = name; + this.getResourceUsage = getResourceUsage; } /** @@ -48,4 +52,14 @@ public static void writeTo(StreamOutput out, ResourceType resourceType) throws I public String getName() { return name; } + + /** + * Gets the resource usage for a given resource type and task. + * + * @param task the task for which to calculate resource usage + * @return the resource usage + */ + public long getResourceUsage(Task task) { + return getResourceUsage.apply(task); + } } diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java b/server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java deleted file mode 100644 index 02158accea5ac..0000000000000 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupHelper.java +++ /dev/null @@ -1,42 +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.wlm; - -import org.opensearch.search.ResourceType; -import org.opensearch.tasks.Task; - -import java.util.Map; -import java.util.function.Function; - -/** - * Helper class for containing static utility methods related to QueryGroups. - */ -public class QueryGroupHelper { - - /** - * A map that associates each {@link ResourceType} with a function that calculates the resource usage for a given {@link Task}. - */ - private static final Map> resourceUsageCalculator = Map.of( - ResourceType.MEMORY, - (task) -> task.getTotalResourceStats().getMemoryInBytes(), - ResourceType.CPU, - (task) -> task.getTotalResourceStats().getCpuTimeInNanos() - ); - - /** - * Gets the resource usage for a given resource type and task. - * - * @param resource the resource type - * @param task the task for which to calculate resource usage - * @return the resource usage - */ - public static long getResourceUsage(ResourceType resource, Task task) { - return resourceUsageCalculator.get(resource).apply(task); - } -} diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 89d7e18f01c65..fb9932118eb02 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -10,9 +10,7 @@ import org.opensearch.search.ResourceType; import org.opensearch.tasks.Task; -import org.opensearch.tasks.TaskManager; import org.opensearch.tasks.TaskResourceTrackingService; -import org.opensearch.wlm.QueryGroupHelper; import org.opensearch.wlm.QueryGroupLevelResourceUsageView; import org.opensearch.wlm.QueryGroupTask; @@ -25,7 +23,7 @@ /** * This class tracks resource usage per QueryGroup */ -public class QueryGroupResourceUsageTrackerService implements TaskManager.TaskEventListeners { +public class QueryGroupResourceUsageTrackerService { public static final List TRACKED_RESOURCES = List.of(ResourceType.MEMORY, ResourceType.CPU); private final TaskResourceTrackingService taskResourceTrackingService; @@ -55,7 +53,7 @@ public Map constructQueryGroupLevelUsa for (ResourceType resourceType : TRACKED_RESOURCES) { long queryGroupResourceUsage = 0; for (Task task : queryGroupEntry.getValue()) { - queryGroupResourceUsage += QueryGroupHelper.getResourceUsage(resourceType, task); + queryGroupResourceUsage += resourceType.getResourceUsage(task); } queryGroupUsage.put(resourceType, queryGroupResourceUsage); } @@ -82,12 +80,4 @@ private Map> getTasksGroupedByQueryGroup() { .map(QueryGroupTask.class::cast) .collect(Collectors.groupingBy(QueryGroupTask::getQueryGroupId, Collectors.mapping(task -> (Task) task, Collectors.toList()))); } - - /** - * 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/search/ResourceTypeTests.java b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java new file mode 100644 index 0000000000000..9d2b183fcdb1b --- /dev/null +++ b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java @@ -0,0 +1,52 @@ +/* + * 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.action.search.SearchShardTask; +import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; +import org.opensearch.tasks.CancellableTask; + +import junit.framework.TestCase; + +import static org.junit.Assert.assertThrows; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +public class ResourceTypeTests extends TestCase { + + public void testFromName() { + assertSame(ResourceType.CPU, ResourceType.fromName("cpu")); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("CPU"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("CPU"); }); + + assertSame(ResourceType.MEMORY, ResourceType.fromName("memory")); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("Memory"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("MEMORY"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("JVM"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("Heap"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("Disk"); }); + } + + public void testGetName() { + assertEquals("cpu", ResourceType.CPU.getName()); + assertEquals("memory", ResourceType.MEMORY.getName()); + } + + public void testGetResourceUsage() { + SearchShardTask mockTask = createMockTask(SearchShardTask.class, 100, 200); + assertEquals(100, ResourceType.CPU.getResourceUsage(mockTask)); + assertEquals(200, ResourceType.MEMORY.getResourceUsage(mockTask)); + } + + private T createMockTask(Class type, long cpuUsage, long heapUsage) { + T task = mock(type); + when(task.getTotalResourceStats()).thenReturn(new TaskResourceUsage(cpuUsage, heapUsage)); + return task; + } +} From 96213e5e412b472fbfb37ea8126a08be27e45e1a Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Fri, 2 Aug 2024 16:56:43 -0700 Subject: [PATCH 58/66] Update ResourceTypeTests.java Signed-off-by: Kiran Prakash --- .../src/test/java/org/opensearch/search/ResourceTypeTests.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java index 9d2b183fcdb1b..61520516ee3ec 100644 --- a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java +++ b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java @@ -23,7 +23,7 @@ public class ResourceTypeTests extends TestCase { public void testFromName() { assertSame(ResourceType.CPU, ResourceType.fromName("cpu")); assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("CPU"); }); - assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("CPU"); }); + assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("Cpu"); }); assertSame(ResourceType.MEMORY, ResourceType.fromName("memory")); assertThrows(IllegalArgumentException.class, () -> { ResourceType.fromName("Memory"); }); From 44f4ff57a69c84ba75b02342a8e9e64dbd5bf4a0 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Fri, 2 Aug 2024 17:08:16 -0700 Subject: [PATCH 59/66] extend OpenSearchTestCase Signed-off-by: Kiran Prakash --- .../test/java/org/opensearch/search/ResourceTypeTests.java | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java index 61520516ee3ec..8568073e9a864 100644 --- a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java +++ b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java @@ -11,14 +11,13 @@ import org.opensearch.action.search.SearchShardTask; import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; import org.opensearch.tasks.CancellableTask; - -import junit.framework.TestCase; +import org.opensearch.test.OpenSearchTestCase; import static org.junit.Assert.assertThrows; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; -public class ResourceTypeTests extends TestCase { +public class ResourceTypeTests extends OpenSearchTestCase { public void testFromName() { assertSame(ResourceType.CPU, ResourceType.fromName("cpu")); From 2bf7d98410b8bead6fb7d28207577baf30f349c1 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Tue, 6 Aug 2024 11:39:07 -0700 Subject: [PATCH 60/66] pr comments Signed-off-by: Kiran Prakash --- server/src/main/java/org/opensearch/search/ResourceType.java | 5 +++-- .../wlm/tracker/QueryGroupResourceUsageTrackerService.java | 3 ++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java index 8d47e1f72899e..0cba2222a6e20 100644 --- a/server/src/main/java/org/opensearch/search/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/ResourceType.java @@ -10,6 +10,7 @@ import org.opensearch.common.annotation.PublicApi; import org.opensearch.core.common.io.stream.StreamOutput; +import org.opensearch.core.tasks.resourcetracker.ResourceStats; import org.opensearch.tasks.Task; import java.io.IOException; @@ -20,8 +21,8 @@ */ @PublicApi(since = "2.x") public enum ResourceType { - CPU("cpu", task -> task.getTotalResourceStats().getCpuTimeInNanos()), - MEMORY("memory", task -> task.getTotalResourceStats().getMemoryInBytes()); + CPU("cpu", task -> task.getTotalResourceUtilization(ResourceStats.CPU)), + MEMORY("memory", task -> task.getTotalResourceUtilization(ResourceStats.MEMORY)); private final String name; private final Function getResourceUsage; diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index fb9932118eb02..bfbf5d8a452d1 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -15,6 +15,7 @@ import org.opensearch.wlm.QueryGroupTask; import java.util.EnumMap; +import java.util.EnumSet; import java.util.HashMap; import java.util.List; import java.util.Map; @@ -25,7 +26,7 @@ */ public class QueryGroupResourceUsageTrackerService { - public static final List TRACKED_RESOURCES = List.of(ResourceType.MEMORY, ResourceType.CPU); + public static final EnumSet TRACKED_RESOURCES = EnumSet.allOf(ResourceType.class); private final TaskResourceTrackingService taskResourceTrackingService; /** From da459bbe8e26222588a1439eb6729d4cd780e4e9 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Tue, 6 Aug 2024 12:47:00 -0700 Subject: [PATCH 61/66] Update CHANGELOG.md Signed-off-by: Kiran Prakash --- CHANGELOG.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2818134f2155b..d7aa71b29a6c8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,13 +6,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), ## [Unreleased 2.x] ### Added - Fix for hasInitiatedFetching to fix allocation explain and manual reroute APIs (([#14972](https://github.com/opensearch-project/OpenSearch/pull/14972)) -- [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) - [Workload Management] Add queryGroupId to Task ([14708](https://github.com/opensearch-project/OpenSearch/pull/14708)) - Add setting to ignore throttling nodes for allocation of unassigned primaries in remote restore ([#14991](https://github.com/opensearch-project/OpenSearch/pull/14991)) - [Streaming Indexing] Enhance RestClient with a new streaming API support ([#14437](https://github.com/opensearch-project/OpenSearch/pull/14437)) - Add basic aggregation support for derived fields ([#14618](https://github.com/opensearch-project/OpenSearch/pull/14618)) - Add ThreadContextPermission for markAsSystemContext and allow core to perform the method ([#15016](https://github.com/opensearch-project/OpenSearch/pull/15016)) - Add ThreadContextPermission for stashAndMergeHeaders and stashWithOrigin ([#15039](https://github.com/opensearch-project/OpenSearch/pull/15039)) +- [Concurrent Segment Search] Support composite aggregations with scripting ([#15072](https://github.com/opensearch-project/OpenSearch/pull/15072)) +- Add `rangeQuery` and `regexpQuery` for `constant_keyword` field type ([#14711](https://github.com/opensearch-project/OpenSearch/pull/14711)) +- [Workload Management] QueryGroup resource tracking framework changes ([#13897](https://github.com/opensearch-project/OpenSearch/pull/13897)) ### Dependencies - Bump `netty` from 4.1.111.Final to 4.1.112.Final ([#15081](https://github.com/opensearch-project/OpenSearch/pull/15081)) From 2b3217b531bbf5d0cae4b483761b63b48079d6cf Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Tue, 6 Aug 2024 14:13:08 -0700 Subject: [PATCH 62/66] Update QueryGroupResourceUsageTrackerServiceTests.java Signed-off-by: Kiran Prakash --- .../tracker/QueryGroupResourceUsageTrackerServiceTests.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java b/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java index 6702a3948eb0e..967119583c25f 100644 --- a/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerServiceTests.java @@ -11,7 +11,7 @@ import org.opensearch.action.search.SearchShardTask; import org.opensearch.action.search.SearchTask; import org.opensearch.common.util.concurrent.ThreadContext; -import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; +import org.opensearch.core.tasks.resourcetracker.ResourceStats; import org.opensearch.search.ResourceType; import org.opensearch.tasks.CancellableTask; import org.opensearch.tasks.Task; @@ -110,7 +110,8 @@ private T createMockTask(Class type, long cpuUsag ((QueryGroupTask) task).setQueryGroupId(threadPool.getThreadContext()); } } - when(task.getTotalResourceStats()).thenReturn(new TaskResourceUsage(cpuUsage, heapUsage)); + when(task.getTotalResourceUtilization(ResourceStats.CPU)).thenReturn(cpuUsage); + when(task.getTotalResourceUtilization(ResourceStats.MEMORY)).thenReturn(heapUsage); when(task.getStartTimeNanos()).thenReturn((long) 0); AtomicBoolean isCancelled = new AtomicBoolean(false); From 6d94768f3339cd3b7113d3302791741044efab7f Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 7 Aug 2024 07:05:35 -0700 Subject: [PATCH 63/66] Update ResourceTypeTests.java Signed-off-by: Kiran Prakash --- .../test/java/org/opensearch/search/ResourceTypeTests.java | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java index 8568073e9a864..14f774db88eba 100644 --- a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java +++ b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java @@ -9,7 +9,7 @@ package org.opensearch.search; import org.opensearch.action.search.SearchShardTask; -import org.opensearch.core.tasks.resourcetracker.TaskResourceUsage; +import org.opensearch.core.tasks.resourcetracker.ResourceStats; import org.opensearch.tasks.CancellableTask; import org.opensearch.test.OpenSearchTestCase; @@ -45,7 +45,8 @@ public void testGetResourceUsage() { private T createMockTask(Class type, long cpuUsage, long heapUsage) { T task = mock(type); - when(task.getTotalResourceStats()).thenReturn(new TaskResourceUsage(cpuUsage, heapUsage)); + when(task.getTotalResourceUtilization(ResourceStats.CPU)).thenReturn(cpuUsage); + when(task.getTotalResourceUtilization(ResourceStats.MEMORY)).thenReturn(cpuUsage); return task; } } From 2f4380c1f006b477b2f85bc1b309b8e962ccbcad Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 7 Aug 2024 07:09:05 -0700 Subject: [PATCH 64/66] Update ResourceTypeTests.java Signed-off-by: Kiran Prakash --- .../src/test/java/org/opensearch/search/ResourceTypeTests.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java index 14f774db88eba..78827b8b1bdad 100644 --- a/server/src/test/java/org/opensearch/search/ResourceTypeTests.java +++ b/server/src/test/java/org/opensearch/search/ResourceTypeTests.java @@ -46,7 +46,7 @@ public void testGetResourceUsage() { private T createMockTask(Class type, long cpuUsage, long heapUsage) { T task = mock(type); when(task.getTotalResourceUtilization(ResourceStats.CPU)).thenReturn(cpuUsage); - when(task.getTotalResourceUtilization(ResourceStats.MEMORY)).thenReturn(cpuUsage); + when(task.getTotalResourceUtilization(ResourceStats.MEMORY)).thenReturn(heapUsage); return task; } } From fa6b2bab202cad265454912827473bae5c0f7847 Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 7 Aug 2024 07:34:33 -0700 Subject: [PATCH 65/66] Update ResourceType.java Signed-off-by: Kiran Prakash --- server/src/main/java/org/opensearch/search/ResourceType.java | 1 + 1 file changed, 1 insertion(+) diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java index 0cba2222a6e20..5226f4bd4b270 100644 --- a/server/src/main/java/org/opensearch/search/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/ResourceType.java @@ -51,6 +51,7 @@ public static void writeTo(StreamOutput out, ResourceType resourceType) throws I } public String getName() { + return name; } From 36b559c78b295ab8068e3cda6b6fb03c8a0860ee Mon Sep 17 00:00:00 2001 From: Kiran Prakash Date: Wed, 7 Aug 2024 07:34:41 -0700 Subject: [PATCH 66/66] Update ResourceType.java Signed-off-by: Kiran Prakash --- server/src/main/java/org/opensearch/search/ResourceType.java | 1 - 1 file changed, 1 deletion(-) diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java index 5226f4bd4b270..0cba2222a6e20 100644 --- a/server/src/main/java/org/opensearch/search/ResourceType.java +++ b/server/src/main/java/org/opensearch/search/ResourceType.java @@ -51,7 +51,6 @@ public static void writeTo(StreamOutput out, ResourceType resourceType) throws I } public String getName() { - return name; }