remove resource usages object from headers (#16532)

Signed-off-by: Chenyang Ji <[email protected]>

CHANGELOG.md

@@ -26,6 +26,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
 

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();
     }
 }

server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java

@@ -547,6 +547,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.

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}
      *