remove resource usages object from headers (#16532)

Signed-off-by: Chenyang Ji <[email protected]> (cherry picked from commit 80ca32f) Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Nov 1, 2024 · d8ff817 · d8ff817
1 parent bdb725d
commit d8ff817
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -23,6 +23,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix get index settings API doesn't show `number_of_routing_shards` setting when it was explicitly set ([#16294](https://github.com/opensearch-project/OpenSearch/pull/16294))
 - Revert changes to upload remote state manifest using minimum codec version([#16403](https://github.com/opensearch-project/OpenSearch/pull/16403))
 - Ensure index templates are not applied to system indices ([#16418](https://github.com/opensearch-project/OpenSearch/pull/16418))
+- Remove resource usages object from search response headers ([#16532](https://github.com/opensearch-project/OpenSearch/pull/16532))
 
 ### Security
 

diff --git a/server/src/main/java/org/opensearch/action/search/SearchTaskRequestOperationsListener.java b/server/src/main/java/org/opensearch/action/search/SearchTaskRequestOperationsListener.java
@@ -25,6 +25,14 @@ public SearchTaskRequestOperationsListener(TaskResourceTrackingService taskResou
 
     @Override
     public void onRequestEnd(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        // Refresh the coordinator node level resource usages
         taskResourceTrackingService.refreshResourceStats(context.getTask());
+        // Remove the shard level resource usages from thread context
+        taskResourceTrackingService.removeTaskResourceUsage();
+    }
+
+    @Override
+    public void onRequestFailure(SearchPhaseContext context, SearchRequestContext searchRequestContext) {
+        taskResourceTrackingService.removeTaskResourceUsage();
     }
 }
diff --git a/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java b/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java
@@ -563,6 +563,15 @@ public void updateResponseHeader(final String key, final String value, final Fun
         threadLocal.set(threadLocal.get().putResponse(key, value, uniqueValue, maxWarningHeaderCount, maxWarningHeaderSize, true));
     }
 
+    /**
+     * Remove the {@code value} for the specified {@code key}.
+     *
+     * @param key         the header name
+     */
+    public void removeResponseHeader(final String key) {
+        threadLocal.get().responseHeaders.remove(key);
+    }
+
     /**
      * Saves the current thread context and wraps command in a Runnable that restores that context before running command. If
      * <code>command</code> has already been passed through this method then it is returned unaltered rather than wrapped twice.

diff --git a/server/src/main/java/org/opensearch/tasks/TaskResourceTrackingService.java b/server/src/main/java/org/opensearch/tasks/TaskResourceTrackingService.java
@@ -328,6 +328,13 @@ public void writeTaskResourceUsage(SearchShardTask task, String nodeId) {
         }
     }
 
+    /**
+     * Remove the current task level resource usages.
+     */
+    public void removeTaskResourceUsage() {
+        threadPool.getThreadContext().removeResponseHeader(TASK_RESOURCE_USAGE);
+    }
+
     /**
      * Get the task resource usages from {@link ThreadContext}
      *