remove resource usages object from headers

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

Author: ansjcy

server/src/main/java/org/opensearch/action/search/SearchTaskRequestOperationsListener.java

@@ -25,6 +25,9 @@ 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();
     }
 }

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