From 80ca32fe9f99165a322e115605027a4df7695f1e Mon Sep 17 00:00:00 2001
From: Chenyang Ji <cyji@amazon.com>
Date: Thu, 31 Oct 2024 18:32:03 -0700
Subject: [PATCH] remove resource usages object from headers (#16532)

Signed-off-by: Chenyang Ji <cyji@amazon.com>
---
 CHANGELOG.md                                             | 1 +
 .../search/SearchTaskRequestOperationsListener.java      | 8 ++++++++
 .../opensearch/common/util/concurrent/ThreadContext.java | 9 +++++++++
 .../opensearch/tasks/TaskResourceTrackingService.java    | 7 +++++++
 4 files changed, 25 insertions(+)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 303e2708b6677..5209f46229edf 100644
--- a/CHANGELOG.md
+++ b/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
 
diff --git a/server/src/main/java/org/opensearch/action/search/SearchTaskRequestOperationsListener.java b/server/src/main/java/org/opensearch/action/search/SearchTaskRequestOperationsListener.java
index 4434d71793b23..ee111b563b747 100644
--- 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
index 070e18481f2a3..75a7ef94978d4 100644
--- a/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java
+++ b/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.
diff --git a/server/src/main/java/org/opensearch/tasks/TaskResourceTrackingService.java b/server/src/main/java/org/opensearch/tasks/TaskResourceTrackingService.java
index ca1957cdb1633..a184673a8fa2f 100644
--- 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}
      *