Addressing UT failure and comments

Signed-off-by: PritLadani <[email protected]>
opensearch-project · Mar 11, 2023 · dfaae2d · dfaae2d
1 parent 08eb7b7
commit dfaae2d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/server/src/main/java/org/opensearch/OpenSearchException.java b/server/src/main/java/org/opensearch/OpenSearchException.java
@@ -259,12 +259,12 @@ protected Map<String, List<String>> getHeaders() {
      */
     public RestStatus status() {
         Throwable cause = unwrapCause();
-        if (cause.getCause() instanceof TaskCancelledException
-            && ((TaskCancelledException) cause.getCause()).status() == RestStatus.TOO_MANY_REQUESTS) {
-            return ((TaskCancelledException) cause.getCause()).status();
-        } else if (cause == this) {
+        if (cause == this) {
             return RestStatus.INTERNAL_SERVER_ERROR;
         } else {
+            if (cause.getCause() != cause) {
+                return ExceptionsHelper.status(cause.getCause());
+            }
             return ExceptionsHelper.status(cause);
         }
     }

diff --git a/server/src/main/java/org/opensearch/tasks/TaskCancellationService.java b/server/src/main/java/org/opensearch/tasks/TaskCancellationService.java
@@ -66,6 +66,7 @@ public class TaskCancellationService {
     public static final String BAN_PARENT_ACTION_NAME = "internal:admin/tasks/ban";
     public static final String REASON_PARENT_CANCELLED_HIGH_RESOURCE_CONSUMPTION =
         "The parent task was cancelled due to high resource consumption";
+    public static final String USAGE_EXCEEDED = "usage exceeded";
     private static final Logger logger = LogManager.getLogger(TaskCancellationService.class);
     private final TransportService transportService;
     private final TaskManager taskManager;
@@ -275,6 +276,6 @@ public static void throwTaskCancelledException(String reason) {
     }
 
     private static boolean isRejection(String reason) {
-        return (reason.contains("usage exceeded") || REASON_PARENT_CANCELLED_HIGH_RESOURCE_CONSUMPTION.equals(reason));
+        return (reason != null && (reason.contains(USAGE_EXCEEDED) || REASON_PARENT_CANCELLED_HIGH_RESOURCE_CONSUMPTION.equals(reason)));
     }
 }