Add changes to propagate queryGroupId across child requests and nodes (…

…opensearch-project#14614) (opensearch-project#14763) * add query group header propagator * apply spotless check * add new propagator in ThreadContext * spotlessApply * address comments * Bump com.microsoft.azure:msal4j from 1.15.1 to 1.16.0 in /plugins/repository-azure (opensearch-project#14610) * Bump com.microsoft.azure:msal4j in /plugins/repository-azure Bumps [com.microsoft.azure:msal4j](https://github.com/AzureAD/microsoft-authentication-library-for-java) from 1.15.1 to 1.16.0. - [Release notes](https://github.com/AzureAD/microsoft-authentication-library-for-java/releases) - [Changelog](https://github.com/AzureAD/microsoft-authentication-library-for-java/blob/dev/changelog.txt) - [Commits](AzureAD/microsoft-authentication-library-for-java@v1.15.1...v1.16.0) --- updated-dependencies: - dependency-name: com.microsoft.azure:msal4j dependency-type: direct:production update-type: version-update:semver-minor ... * Updating SHAs * Update changelog --------- * [Bugfix] Fix ICacheKeySerializerTests flakiness (opensearch-project#14564) * Fix testInvalidInput flakiness * Addressed andrross's comment * rerun security check --------- * Correct typo in method name (opensearch-project#14621) * Refactoring FilterPath.parse by using an iterative approach instead of recursion. (opensearch-project#14200) * Refactor FilterPath parse function (opensearch-project#12067) * Implement unit tests for FilterPathTests (opensearch-project#12067) * Write warn log if Filter is empty; Add comments (opensearch-project#12067) * Add changelog * Remove unnecessary log statement * Remove unused logger * Spotless apply * Remove incorrect changelog --------- * Removing String format in RemoteStoreMigrationAllocationDecider to optimise performance(opensearch-project#14612) * Clear templates before Adding; Use NamedWriteableAwareStreamInput for RemoteCustomMetadata; Correct the check for deciding upload of HashesOfConsistentSettings (opensearch-project#14513) * Clear templates before Adding; Use NamedWriteableAwareStreamInput for RemoteCustomMetadata * Correct the check for deciding upload of hashes of consistent settings * add changelog * add PR link changelog * Improve reroute performance by optimising List.removeAll in LocalShardsBalancer to filter remote search shard from relocation decision (opensearch-project#14613) * Fix assertion failure while deleting remote backed index (opensearch-project#14601) * Allow system index warning in OpenSearchRestTestCase.refreshAllIndices (opensearch-project#14635) * Allow system index warning * Add to CHANGELOG * Address code review comments --------- * Star tree codec changes (opensearch-project#14514) --------- * Bump com.github.spullara.mustache.java:compiler from 0.9.13 to 0.9.14 in /modules/lang-mustache (opensearch-project#14672) * Bump com.github.spullara.mustache.java:compiler Bumps [com.github.spullara.mustache.java:compiler](https://github.com/spullara/mustache.java) from 0.9.13 to 0.9.14. - [Commits](spullara/mustache.java@mustache.java-0.9.13...mustache.java-0.9.14) --- updated-dependencies: - dependency-name: com.github.spullara.mustache.java:compiler dependency-type: direct:production update-type: version-update:semver-patch ... * Updating SHAs * Update changelog --------- * Bump net.minidev:accessors-smart from 2.5.0 to 2.5.1 in /plugins/repository-azure (opensearch-project#14673) * Bump net.minidev:accessors-smart in /plugins/repository-azure Bumps [net.minidev:accessors-smart](https://github.com/netplex/json-smart-v2) from 2.5.0 to 2.5.1. - [Release notes](https://github.com/netplex/json-smart-v2/releases) - [Commits](netplex/json-smart-v2@2.5.0...2.5.1) --- updated-dependencies: - dependency-name: net.minidev:accessors-smart dependency-type: direct:production update-type: version-update:semver-patch ... * Updating SHAs * Update changelog --------- * move query group thread context propagator out of ThreadContext --------- (cherry picked from commit d33d24e) Signed-off-by: Kaushal Kumar <ravi.kaushal97@gmail.com> Signed-off-by: dependabot[bot] <support@github.com> Signed-off-by: Peter Alfonsi <petealft@amazon.com> Signed-off-by: vatsal <vatsal.v.anand@gmail.com> Signed-off-by: Siddhant Deshmukh <deshsid@amazon.com> Signed-off-by: RS146BIJAY <rishavsagar4b1@gmail.com> Signed-off-by: Sooraj Sinha <soosinha@amazon.com> Signed-off-by: Sachin Kale <kalsac@amazon.com> Signed-off-by: Craig Perkins <cwperx@amazon.com> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: dependabot[bot] <49699333+dependabot[bot]@users.noreply.github.com> Co-authored-by: dependabot[bot] <dependabot[bot]@users.noreply.github.com> Co-authored-by: Peter Alfonsi <peter.alfonsi@gmail.com> Co-authored-by: Peter Alfonsi <petealft@amazon.com> Co-authored-by: Vatsal <36672090+imvtsl@users.noreply.github.com> Co-authored-by: Siddhant Deshmukh <deshsid@amazon.com> Co-authored-by: Robin Friedmann <robinfriedmann.rf@gmail.com> Co-authored-by: rishavz_sagar <rishavsagar4b1@gmail.com> Co-authored-by: Sooraj Sinha <81695996+soosinha@users.noreply.github.com> Co-authored-by: Sachin Kale <sachinpkale@gmail.com> Co-authored-by: Craig Perkins <cwperx@amazon.com> Co-authored-by: Bharathwaj G <bharath78910@gmail.com> Signed-off-by: kkewwei <kkewwei@163.com>
kkewwei · Jul 24, 2024 · 48b29b1 · 48b29b1
1 parent 01d0cb1
commit 48b29b1
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 0 deletions.
diff --git a/CHANGELOG.md b/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))
 - Add matchesPluginSystemIndexPattern to SystemIndexRegistry ([#14750](https://github.com/opensearch-project/OpenSearch/pull/14750))

diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagator.java b/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;
+    }
+}
diff --git a/server/src/main/java/org/opensearch/wlm/package-info.java b/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;
diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagatorTests.java b/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"));
+    }
+}