Skip to content

Commit

Permalink
remove resource usages object from headers
Browse files Browse the repository at this point in the history
Signed-off-by: Chenyang Ji <cyji@amazon.com>
  • Loading branch information
ansjcy committed Oct 31, 2024
1 parent 4b284c5 commit 93f2012
Show file tree
Hide file tree
Showing 3 changed files with 19 additions and 0 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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();
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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}
*
Expand Down

0 comments on commit 93f2012

Please sign in to comment.