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..6a864ee6b6195 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,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(); } } 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 * command 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} *