diff --git a/CHANGELOG.md b/CHANGELOG.md index 6bfd98ceaea80..6f666dbf3b8d5 100644 --- 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)) diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagator.java b/server/src/main/java/org/opensearch/wlm/QueryGroupThreadContextStatePropagator.java new file mode 100644 index 0000000000000..06d223907082e --- /dev/null +++ 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 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 transients(Map source) { + final Map 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 headers(Map source) { + final Map 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 new file mode 100644 index 0000000000000..fa4731d95cc34 --- /dev/null +++ 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 new file mode 100644 index 0000000000000..ad5d7f569a56e --- /dev/null +++ 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 source = Map.of("queryGroupId", "adgarja0r235te"); + Map transients = sut.transients(source); + assertEquals("adgarja0r235te", transients.get("queryGroupId")); + } + + public void testHeaders() { + QueryGroupThreadContextStatePropagator sut = new QueryGroupThreadContextStatePropagator(); + Map source = Map.of("queryGroupId", "adgarja0r235te"); + Map headers = sut.headers(source); + assertEquals("adgarja0r235te", headers.get("queryGroupId")); + } +}