Add queryGroupId to search workload tasks at co-ordinator and data node level

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ## [Unreleased 2.x]
 ### Added
 - Fix for hasInitiatedFetching to fix allocation explain and manual reroute APIs (([#14972](https://github.com/opensearch-project/OpenSearch/pull/14972))
+- [Workload Management] Add queryGroupId to Task ([14708](https://github.com/opensearch-project/OpenSearch/pull/14708))
 - Add basic aggregation support for derived fields ([#14618](https://github.com/opensearch-project/OpenSearch/pull/14618))
 
 ### Dependencies

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

@@ -37,8 +37,8 @@
 import org.opensearch.core.tasks.TaskId;
 import org.opensearch.search.fetch.ShardFetchSearchRequest;
 import org.opensearch.search.internal.ShardSearchRequest;
-import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.SearchBackpressureTask;
+import org.opensearch.wlm.QueryGroupTask;
 
 import java.util.Map;
 import java.util.function.Supplier;
@@ -50,7 +50,7 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public class SearchShardTask extends CancellableTask implements SearchBackpressureTask {
+public class SearchShardTask extends QueryGroupTask implements SearchBackpressureTask {
     // generating metadata in a lazy way since source can be quite big
     private final MemoizedSupplier<String> metadataSupplier;
 

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

@@ -35,8 +35,8 @@
 import org.opensearch.common.annotation.PublicApi;
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.core.tasks.TaskId;
-import org.opensearch.tasks.CancellableTask;
 import org.opensearch.tasks.SearchBackpressureTask;
+import org.opensearch.wlm.QueryGroupTask;
 
 import java.util.Map;
 import java.util.function.Supplier;
@@ -49,7 +49,7 @@
  * @opensearch.api
  */
 @PublicApi(since = "1.0.0")
-public class SearchTask extends CancellableTask implements SearchBackpressureTask {
+public class SearchTask extends QueryGroupTask implements SearchBackpressureTask {
     // generating description in a lazy way since source can be quite big
     private final Supplier<String> descriptionSupplier;
     private SearchProgressListener progressListener = SearchProgressListener.NOOP;

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

@@ -101,6 +101,7 @@
 import org.opensearch.transport.RemoteTransportException;
 import org.opensearch.transport.Transport;
 import org.opensearch.transport.TransportService;
+import org.opensearch.wlm.QueryGroupTask;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -442,6 +443,12 @@ private void executeRequest(
             );
             searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);
 
+            // At this point either the QUERY_GROUP_ID header will be present in ThreadContext either via ActionFilter
+            // or HTTP header (HTTP header will be deprecated once ActionFilter is implemented)
+            if (task instanceof QueryGroupTask) {
+                ((QueryGroupTask) task).setQueryGroupId(threadPool.getThreadContext());
+            }
+
             PipelinedRequest searchRequest;
             ActionListener<SearchResponse> listener;
             try {

server/src/main/java/org/opensearch/node/Node.java

@@ -263,6 +263,7 @@
 import org.opensearch.transport.TransportService;
 import org.opensearch.usage.UsageService;
 import org.opensearch.watcher.ResourceWatcherService;
+import org.opensearch.wlm.WorkloadManagementTransportInterceptor;
 
 import javax.net.ssl.SNIHostName;
 
@@ -1047,14 +1048,21 @@ protected Node(
                 admissionControlService
             );
 
+            WorkloadManagementTransportInterceptor workloadManagementTransportInterceptor = new WorkloadManagementTransportInterceptor(
+                threadPool
+            );
+
             final Collection<SecureSettingsFactory> secureSettingsFactories = pluginsService.filterPlugins(Plugin.class)
                 .stream()
                 .map(p -> p.getSecureSettingFactory(settings))
                 .filter(Optional::isPresent)
                 .map(Optional::get)
                 .collect(Collectors.toList());
 
-            List<TransportInterceptor> transportInterceptors = List.of(admissionControlTransportInterceptor);
+            List<TransportInterceptor> transportInterceptors = List.of(
+                admissionControlTransportInterceptor,
+                workloadManagementTransportInterceptor
+            );
             final NetworkModule networkModule = new NetworkModule(
                 settings,
                 pluginsService.filterPlugins(NetworkPlugin.class),

server/src/main/java/org/opensearch/wlm/QueryGroupTask.java

@@ -0,0 +1,76 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.common.util.concurrent.ThreadContext;
+import org.opensearch.core.tasks.TaskId;
+import org.opensearch.tasks.CancellableTask;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.function.Supplier;
+
+import static org.opensearch.search.SearchService.NO_TIMEOUT;
+
+/**
+ * Base class to define QueryGroup tasks
+ */
+public class QueryGroupTask extends CancellableTask {
+
+    private static final Logger logger = LogManager.getLogger(QueryGroupTask.class);
+    public static final String QUERY_GROUP_ID_HEADER = "queryGroupId";
+    public static final Supplier<String> DEFAULT_QUERY_GROUP_ID_SUPPLIER = () -> "DEFAULT_QUERY_GROUP";
+    private String queryGroupId;
+
+    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);
+    }
+
+    public QueryGroupTask(
+        long id,
+        String type,
+        String action,
+        String description,
+        TaskId parentTaskId,
+        Map<String, String> headers,
+        TimeValue cancelAfterTimeInterval
+    ) {
+        super(id, type, action, description, parentTaskId, headers, cancelAfterTimeInterval);
+    }
+
+    /**
+     * This method should always be called after calling setQueryGroupId at least once on this object
+     * @return task queryGroupId
+     */
+    public final String getQueryGroupId() {
+        if (queryGroupId == null) {
+            logger.warn("QueryGroup _id can't be null, It should be set before accessing it. This is abnormal behaviour ");
+        }
+        return queryGroupId;
+    }
+
+    /**
+     * sets the queryGroupId from threadContext into the task itself,
+     * This method was defined since the queryGroupId can only be evaluated after task creation
+     * @param threadContext current threadContext
+     */
+    public final void setQueryGroupId(final ThreadContext threadContext) {
+        this.queryGroupId = Optional.ofNullable(threadContext)
+            .map(threadContext1 -> threadContext1.getHeader(QUERY_GROUP_ID_HEADER))
+            .orElse(DEFAULT_QUERY_GROUP_ID_SUPPLIER.get());
+    }
+
+    @Override
+    public boolean shouldCancelChildrenOnCancellation() {
+        return false;
+    }
+}

server/src/main/java/org/opensearch/wlm/WorkloadManagementTransportInterceptor.java

@@ -0,0 +1,64 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.tasks.Task;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportChannel;
+import org.opensearch.transport.TransportInterceptor;
+import org.opensearch.transport.TransportRequest;
+import org.opensearch.transport.TransportRequestHandler;
+
+/**
+ * This class is used to intercept search traffic requests and populate the queryGroupId header in task headers
+ */
+public class WorkloadManagementTransportInterceptor implements TransportInterceptor {
+    private final ThreadPool threadPool;
+
+    public WorkloadManagementTransportInterceptor(ThreadPool threadPool) {
+        this.threadPool = threadPool;
+    }
+
+    @Override
+    public <T extends TransportRequest> TransportRequestHandler<T> interceptHandler(
+        String action,
+        String executor,
+        boolean forceExecution,
+        TransportRequestHandler<T> actualHandler
+    ) {
+        return new RequestHandler<T>(threadPool, actualHandler);
+    }
+
+    /**
+     * This class is mainly used to populate the queryGroupId header
+     * @param <T> T is Search related request
+     */
+    public static class RequestHandler<T extends TransportRequest> implements TransportRequestHandler<T> {
+
+        private final ThreadPool threadPool;
+        TransportRequestHandler<T> actualHandler;
+
+        public RequestHandler(ThreadPool threadPool, TransportRequestHandler<T> actualHandler) {
+            this.threadPool = threadPool;
+            this.actualHandler = actualHandler;
+        }
+
+        @Override
+        public void messageReceived(T request, TransportChannel channel, Task task) throws Exception {
+            if (isSearchWorkloadRequest(task)) {
+                ((QueryGroupTask) task).setQueryGroupId(threadPool.getThreadContext());
+            }
+            actualHandler.messageReceived(request, channel, task);
+        }
+
+        boolean isSearchWorkloadRequest(Task task) {
+            return task instanceof QueryGroupTask;
+        }
+    }
+}

server/src/test/java/org/opensearch/wlm/QueryGroupTaskTests.java

@@ -0,0 +1,44 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.threadpool.TestThreadPool;
+import org.opensearch.threadpool.ThreadPool;
+
+import java.util.Collections;
+
+import static org.opensearch.wlm.QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER;
+import static org.opensearch.wlm.QueryGroupTask.QUERY_GROUP_ID_HEADER;
+
+public class QueryGroupTaskTests extends OpenSearchTestCase {
+    private ThreadPool threadPool;
+    private QueryGroupTask sut;
+
+    public void setUp() throws Exception {
+        super.setUp();
+        threadPool = new TestThreadPool(getTestName());
+        sut = new QueryGroupTask(123, "transport", "Search", "test task", null, Collections.emptyMap());
+    }
+
+    public void tearDown() throws Exception {
+        super.tearDown();
+        threadPool.shutdown();
+    }
+
+    public void testSuccessfulSetQueryGroupId() {
+        sut.setQueryGroupId(threadPool.getThreadContext());
+        assertEquals(DEFAULT_QUERY_GROUP_ID_SUPPLIER.get(), sut.getQueryGroupId());
+
+        threadPool.getThreadContext().putHeader(QUERY_GROUP_ID_HEADER, "akfanglkaglknag2332");
+
+        sut.setQueryGroupId(threadPool.getThreadContext());
+        assertEquals("akfanglkaglknag2332", sut.getQueryGroupId());
+    }
+}

server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportInterceptorTests.java

@@ -0,0 +1,40 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ */
+
+package org.opensearch.wlm;
+
+import org.opensearch.test.OpenSearchTestCase;
+import org.opensearch.threadpool.TestThreadPool;
+import org.opensearch.threadpool.ThreadPool;
+import org.opensearch.transport.TransportRequest;
+import org.opensearch.transport.TransportRequestHandler;
+import org.opensearch.wlm.WorkloadManagementTransportInterceptor.RequestHandler;
+
+import static org.opensearch.threadpool.ThreadPool.Names.SAME;
+
+public class WorkloadManagementTransportInterceptorTests extends OpenSearchTestCase {
+
+    private ThreadPool threadPool;
+    private WorkloadManagementTransportInterceptor sut;
+
+    public void setUp() throws Exception {
+        super.setUp();
+        threadPool = new TestThreadPool(getTestName());
+        sut = new WorkloadManagementTransportInterceptor(threadPool);
+    }
+
+    public void tearDown() throws Exception {
+        super.tearDown();
+        threadPool.shutdown();
+    }
+
+    public void testInterceptHandler() {
+        TransportRequestHandler<TransportRequest> requestHandler = sut.interceptHandler("Search", SAME, false, null);
+        assertTrue(requestHandler instanceof RequestHandler);
+    }
+}