[Backport 2.x] Add changes to propagate queryGroupId across child requests and nodes

CHANGELOG.md

@@ -14,6 +14,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix race condition while parsing derived fields from search definition ([14445](https://github.com/opensearch-project/OpenSearch/pull/14445))
 - Add `strict_allow_templates` dynamic mapping option ([#14555](https://github.com/opensearch-project/OpenSearch/pull/14555))
 - Add allowlist setting for ingest-common and search-pipeline-common processors ([#14439](https://github.com/opensearch-project/OpenSearch/issues/14439))
+- [Workload Management] add queryGroupId header propagator across requests and nodes ([#14614](https://github.com/opensearch-project/OpenSearch/pull/14614))
 - Create SystemIndexRegistry with helper method matchesSystemIndex ([#14415](https://github.com/opensearch-project/OpenSearch/pull/14415))
 - Print reason why parent task was cancelled ([#14604](https://github.com/opensearch-project/OpenSearch/issues/14604))
 

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

@@ -0,0 +1,53 @@
+/*
+ * 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.common.util.concurrent.ThreadContextStatePropagator;
+
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * This class is used to propagate QueryGroup related headers to request and nodes
+ */
+public class QueryGroupThreadContextStatePropagator implements ThreadContextStatePropagator {
+    // TODO: move this constant to QueryGroupService class once the QueryGroup monitoring framework PR is ready
+    public static List<String> PROPAGATED_HEADERS = List.of("queryGroupId");
+
+    /**
+     * @param source current context transient headers
+     * @return the map of header and their values to be propagated across request threadContexts
+     */
+    @Override
+    @SuppressWarnings("removal")
+    public Map<String, Object> transients(Map<String, Object> source) {
+        final Map<String, Object> transientHeaders = new HashMap<>();
+
+        for (String headerName : PROPAGATED_HEADERS) {
+            transientHeaders.compute(headerName, (k, v) -> source.get(headerName));
+        }
+        return transientHeaders;
+    }
+
+    /**
+     * @param source current context headers
+     * @return map of header and their values to be propagated across nodes
+     */
+    @Override
+    @SuppressWarnings("removal")
+    public Map<String, String> headers(Map<String, Object> source) {
+        final Map<String, String> propagatedHeaders = new HashMap<>();
+
+        for (String headerName : PROPAGATED_HEADERS) {
+            propagatedHeaders.compute(headerName, (k, v) -> (String) source.get(headerName));
+        }
+        return propagatedHeaders;
+    }
+}

server/src/main/java/org/opensearch/wlm/package-info.java

@@ -0,0 +1,13 @@
+/*
+ * 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.
+ */
+
+/**
+ * This package contains workload management constructs
+ */
+
+package org.opensearch.wlm;

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

@@ -0,0 +1,30 @@
+/*
+ * 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 java.util.Map;
+
+public class QueryGroupThreadContextStatePropagatorTests extends OpenSearchTestCase {
+
+    public void testTransients() {
+        QueryGroupThreadContextStatePropagator sut = new QueryGroupThreadContextStatePropagator();
+        Map<String, Object> source = Map.of("queryGroupId", "adgarja0r235te");
+        Map<String, Object> transients = sut.transients(source);
+        assertEquals("adgarja0r235te", transients.get("queryGroupId"));
+    }
+
+    public void testHeaders() {
+        QueryGroupThreadContextStatePropagator sut = new QueryGroupThreadContextStatePropagator();
+        Map<String, Object> source = Map.of("queryGroupId", "adgarja0r235te");
+        Map<String, String> headers = sut.headers(source);
+        assertEquals("adgarja0r235te", headers.get("queryGroupId"));
+    }
+}