From 0c37980420b0e21618ce974246e290f8ddcaf3c8 Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Fri, 17 Jan 2025 11:53:25 -0800 Subject: [PATCH] remove untagged tasks from WLM tracking Signed-off-by: Kaushal Kumar --- server/src/main/java/org/opensearch/wlm/QueryGroupTask.java | 6 ++++++ .../wlm/tracker/QueryGroupResourceUsageTrackerService.java | 1 + 2 files changed, 7 insertions(+) diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java b/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java index 97c48bd828978..b815bf47ef934 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupTask.java @@ -33,6 +33,7 @@ public class QueryGroupTask extends CancellableTask { public static final Supplier DEFAULT_QUERY_GROUP_ID_SUPPLIER = () -> "DEFAULT_QUERY_GROUP"; private final LongSupplier nanoTimeSupplier; private String queryGroupId; + private boolean isEligibleForTracking = false; public QueryGroupTask(long id, String type, String action, String description, TaskId parentTaskId, Map headers) { this(id, type, action, description, parentTaskId, headers, NO_TIMEOUT, System::nanoTime); @@ -81,6 +82,7 @@ public final String getQueryGroupId() { * @param threadContext current threadContext */ public final void setQueryGroupId(final ThreadContext threadContext) { + isEligibleForTracking = true; if (threadContext != null && threadContext.getHeader(QUERY_GROUP_ID_HEADER) != null) { this.queryGroupId = threadContext.getHeader(QUERY_GROUP_ID_HEADER); } else { @@ -92,6 +94,10 @@ public long getElapsedTime() { return nanoTimeSupplier.getAsLong() - getStartTimeNanos(); } + public boolean isEligibleForTracking() { + return isEligibleForTracking; + } + @Override public boolean shouldCancelChildrenOnCancellation() { return false; diff --git a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java index 19f7bf48d8421..b2e27da19a210 100644 --- a/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java +++ b/server/src/main/java/org/opensearch/wlm/tracker/QueryGroupResourceUsageTrackerService.java @@ -76,6 +76,7 @@ private Map> getTasksGroupedByQueryGroup() { .stream() .filter(QueryGroupTask.class::isInstance) .map(QueryGroupTask.class::cast) + .filter(QueryGroupTask::isEligibleForTracking) .collect(Collectors.groupingBy(QueryGroupTask::getQueryGroupId, Collectors.mapping(task -> task, Collectors.toList()))); } }