remove untagged tasks from WLM tracking

Signed-off-by: Kaushal Kumar <[email protected]>
kaushalmahi12 · Jan 17, 2025 · 0c37980 · 0c37980
1 parent 36d43ef
commit 0c37980
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git 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<String> 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<String, String> 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
@@ -76,6 +76,7 @@ private Map<String, List<QueryGroupTask>> getTasksGroupedByQueryGroup() {
             .stream()
             .filter(QueryGroupTask.class::isInstance)
             .map(QueryGroupTask.class::cast)
+            .filter(QueryGroupTask::isEligibleForTracking)
             .collect(Collectors.groupingBy(QueryGroupTask::getQueryGroupId, Collectors.mapping(task -> task, Collectors.toList())));
     }
 }