renaming sandbox to querygroup and adjusting code based on merged PRs

Signed-off-by: Kiran Prakash <[email protected]>
kiranprakash154 · Jul 18, 2024 · 3e953a1 · 3e953a1
1 parent 3766b32
commit 3e953a1
Show file tree

Hide file tree

Showing 46 changed files with 967 additions and 1,061 deletions.
diff --git a/server/src/main/java/org/opensearch/action/search/SearchShardTask.java b/server/src/main/java/org/opensearch/action/search/SearchShardTask.java
@@ -37,7 +37,7 @@
 import org.opensearch.core.tasks.TaskId;
 import org.opensearch.search.fetch.ShardFetchSearchRequest;
 import org.opensearch.search.internal.ShardSearchRequest;
-import org.opensearch.search.sandboxing.SandboxTask;
+import org.opensearch.search.querygroup.QueryGroupTask;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.SearchBackpressureTask;
 
@@ -51,10 +51,10 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public class SearchShardTask extends CancellableTask implements SearchBackpressureTask, SandboxTask {
+public class SearchShardTask extends CancellableTask implements SearchBackpressureTask, QueryGroupTask {
     // generating metadata in a lazy way since source can be quite big
     private final MemoizedSupplier<String> metadataSupplier;
-    private String sandboxId;
+    private String queryGroupId;
 
     public SearchShardTask(long id, String type, String action, String description, TaskId parentTaskId, Map<String, String> headers) {
         this(id, type, action, description, parentTaskId, headers, () -> "");
@@ -88,12 +88,12 @@ public boolean shouldCancelChildrenOnCancellation() {
     }
 
     @Override
-    public void setSandboxId(String sandboxId) {
-        this.sandboxId = sandboxId;
+    public void setQueryGroupId(String queryGroupId) {
+        this.queryGroupId = queryGroupId;
     }
 
     @Override
-    public String getSandboxId() {
-        return "sandboxId";
+    public String getQueryGroupId() {
+        return "queryGroupId";
     }
 }
diff --git a/server/src/main/java/org/opensearch/action/search/SearchTask.java b/server/src/main/java/org/opensearch/action/search/SearchTask.java
@@ -35,7 +35,7 @@
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.tasks.TaskId;
-import org.opensearch.search.sandboxing.SandboxTask;
+import org.opensearch.search.querygroup.QueryGroupTask;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.SearchBackpressureTask;
 
@@ -50,11 +50,11 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public class SearchTask extends CancellableTask implements SearchBackpressureTask, SandboxTask {
+public class SearchTask extends CancellableTask implements SearchBackpressureTask, QueryGroupTask {
     // generating description in a lazy way since source can be quite big
     private final Supplier<String> descriptionSupplier;
     private SearchProgressListener progressListener = SearchProgressListener.NOOP;
-    private String sandboxId;
+    private String queryGroupId;
 
     public SearchTask(
         long id,
@@ -109,12 +109,12 @@ public boolean shouldCancelChildrenOnCancellation() {
         return true;
     }
 
-    public void setSandboxId(String sandboxId) {
-        this.sandboxId = sandboxId;
+    public void setQueryGroupId(String queryGroupId) {
+        this.queryGroupId = queryGroupId;
     }
 
     @Override
-    public String getSandboxId() {
-        return "sandboxId";
+    public String getQueryGroupId() {
+        return "queryGroupId";
     }
 }
diff --git a/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java b/server/src/main/java/org/opensearch/cluster/metadata/Metadata.java
@@ -126,7 +126,7 @@ public boolean isSegmentReplicationEnabled(String indexName) {
             .orElse(false);
     }
 
-    public Map<String, Sandbox> sandboxes() {
+    public Map<String, QueryGroup> 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
@@ -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
diff --git a/server/src/main/java/org/opensearch/search/ResourceType.java b/server/src/main/java/org/opensearch/search/ResourceType.java
diff --git 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<TaskResourceUsageTrackerType, Function<NodeDuressTrackers, Boolean>> 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,18 +93,18 @@ 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(),
                         () -> settings.getNodeDuressSettings().getNumSuccessiveBreaches()
                     )
                 );
                 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
@@ -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;