fix breaking testcases

Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com>
opensearch-project · Apr 19, 2024 · 5ba7535 · 5ba7535
1 parent 44e9d11
commit 5ba7535
Show file tree

Hide file tree

Showing 7 changed files with 35 additions and 11 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,6 +37,7 @@
 import org.opensearch.core.tasks.TaskId;
 import org.opensearch.search.fetch.ShardFetchSearchRequest;
 import org.opensearch.search.internal.ShardSearchRequest;
+import org.opensearch.search.resource_limit_group.ResourceLimitGroupTask;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.SearchBackpressureTask;
 
@@ -50,7 +51,7 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public class SearchShardTask extends CancellableTask implements SearchBackpressureTask {
+public class SearchShardTask extends CancellableTask implements SearchBackpressureTask, ResourceLimitGroupTask {
     // generating metadata in a lazy way since source can be quite big
     private final MemoizedSupplier<String> metadataSupplier;
     private String resourceLimitGroupId;

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,6 +35,7 @@
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.tasks.TaskId;
+import org.opensearch.search.resource_limit_group.ResourceLimitGroupTask;
 import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.SearchBackpressureTask;
 
@@ -49,7 +50,7 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public class SearchTask extends CancellableTask implements SearchBackpressureTask {
+public class SearchTask extends CancellableTask implements SearchBackpressureTask, ResourceLimitGroupTask {
     // generating description in a lazy way since source can be quite big
     private final Supplier<String> descriptionSupplier;
     private SearchProgressListener progressListener = SearchProgressListener.NOOP;

diff --git a/server/src/main/java/org/opensearch/cluster/metadata/ResourceLimitGroup.java b/server/src/main/java/org/opensearch/cluster/metadata/ResourceLimitGroup.java
@@ -22,6 +22,7 @@
 
 import java.io.IOException;
 import java.util.List;
+import java.util.Locale;
 import java.util.Objects;
 
 /**
@@ -118,7 +119,7 @@ public ResourceLimit(String resourceName, Double value) {
                 throw new IllegalArgumentException("resource value should be less than 1.0");
             }
 
-            if (!ALLOWED_RESOURCES.contains(resourceName.toLowerCase())) {
+            if (!ALLOWED_RESOURCES.contains(resourceName.toLowerCase(Locale.ROOT))) {
                 throw new IllegalArgumentException(
                     "resource has to be valid, valid resources " + ALLOWED_RESOURCES.stream().reduce((x, e) -> x + ", " + e).get()
                 );
@@ -175,6 +176,9 @@ public int hashCode() {
         }
     }
 
+    /**
+     * This enum models the different sandbox modes
+     */
     @ExperimentalApi
     public static enum ResourceLimitGroupMode {
         SOFT("soft"),

diff --git a/server/src/main/java/org/opensearch/search/resource_limit_group/ResourceLimitGroupTask.java b/server/src/main/java/org/opensearch/search/resource_limit_group/ResourceLimitGroupTask.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.resource_limit_group;
+
+/**
+ * This interface can be implemented by tasks which will be tracked and monitored using {@link org.opensearch.cluster.metadata.ResourceLimitGroup}
+ */
+public interface ResourceLimitGroupTask {
+    public void setResourceLimitGroupId(String sandboxResourceLimitGroup);
+
+    public String getResourceLimitGroupId();
+}
diff --git a/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java b/server/src/test/java/org/opensearch/action/search/AbstractSearchAsyncActionTests.java
@@ -179,7 +179,7 @@ private AbstractSearchAsyncAction<SearchPhaseResult> createAction(
             new GroupShardsIterator<>(Arrays.asList(shards)),
             timeProvider,
             ClusterState.EMPTY_STATE,
-            null,
+            new SearchTask(123, null, null, null, null, null),
             results,
             request.getMaxConcurrentShardRequests(),
             SearchResponse.Clusters.EMPTY,

diff --git a/server/src/test/java/org/opensearch/action/search/CanMatchPreFilterSearchPhaseTests.java b/server/src/test/java/org/opensearch/action/search/CanMatchPreFilterSearchPhaseTests.java
@@ -160,7 +160,7 @@ public void sendCanMatch(
             shardsIter,
             timeProvider,
             ClusterState.EMPTY_STATE,
-            null,
+            new SearchTask(1, null, null, null, null, null),
             (iter) -> new SearchPhase("test") {
                 @Override
                 public void run() throws IOException {
@@ -258,7 +258,7 @@ public void sendCanMatch(
             shardsIter,
             timeProvider,
             ClusterState.EMPTY_STATE,
-            null,
+            new SearchTask(1, null, null, null, null, null),
             (iter) -> new SearchPhase("test") {
                 @Override
                 public void run() throws IOException {
@@ -346,7 +346,7 @@ public void sendCanMatch(
             shardsIter,
             timeProvider,
             ClusterState.EMPTY_STATE,
-            null,
+            new SearchTask(1, null, null, null, null, null),
             (iter) -> {
                 return new WrappingSearchAsyncActionPhase(
                     new AbstractSearchAsyncAction<SearchPhaseResult>("test", logger, transportService, (cluster, node) -> {
@@ -478,7 +478,7 @@ public void sendCanMatch(
                 shardsIter,
                 timeProvider,
                 ClusterState.EMPTY_STATE,
-                null,
+                new SearchTask(1, null, null, null, null, null),
                 (iter) -> new SearchPhase("test") {
                     @Override
                     public void run() {
@@ -585,7 +585,7 @@ public void sendCanMatch(
                 shardsIter,
                 timeProvider,
                 ClusterState.EMPTY_STATE,
-                null,
+                new SearchTask(1, null, null, null, null, null),
                 (iter) -> new SearchPhase("test") {
                     @Override
                     public void run() {
@@ -690,7 +690,7 @@ public void sendCanMatch(
             shardsIter,
             timeProvider,
             ClusterState.EMPTY_STATE,
-            null,
+            new SearchTask(1, null, null, null, null, null),
             (iter) -> {
                 AbstractSearchAsyncAction<? extends SearchPhaseResult> action = new SearchDfsQueryAsyncAction(
                     logger,

diff --git a/server/src/test/java/org/opensearch/cluster/metadata/ResourceLimitGroupTests.java b/server/src/test/java/org/opensearch/cluster/metadata/ResourceLimitGroupTests.java
@@ -32,7 +32,7 @@ static ResourceLimitGroup createRandomResourceLimitGroup() {
     }
 
     private static String randomMode() {
-        return allowedModes.get((int) (Math.random() * allowedModes.size()));
+        return allowedModes.get(randomIntBetween(0, allowedModes.size() - 1));
     }
 
     /**
-Original file line number
+Diff line change
@@ Expand Up / @@ -32,7 +32,7 @@ static ResourceLimitGroup createRandomResourceLimitGroup() { @@
         }
         private static String randomMode() {
-            return allowedModes.get((int) (Math.random() * allowedModes.size()));
+            return allowedModes.get(randomIntBetween(0, allowedModes.size() - 1));
         }
         /**
@@ Expand Down @@