From 56bebf63bbefb20eedab206c11b15790977cf6c1 Mon Sep 17 00:00:00 2001 From: Craig Perkins Date: Thu, 25 Jul 2024 23:15:45 -0400 Subject: [PATCH 1/2] Make ThreadContext.stashAndMergeHeaders package-private Signed-off-by: Craig Perkins --- .../client/support/AbstractClient.java | 17 +------ .../StashAndMergeHeadersFilterClient.java | 45 +++++++++++++++++++ .../common/util/concurrent/ThreadContext.java | 2 +- 3 files changed, 48 insertions(+), 16 deletions(-) create mode 100644 server/src/main/java/org/opensearch/common/util/concurrent/StashAndMergeHeadersFilterClient.java diff --git a/server/src/main/java/org/opensearch/client/support/AbstractClient.java b/server/src/main/java/org/opensearch/client/support/AbstractClient.java index 6c6049f04231b..6af63f425d6a2 100644 --- a/server/src/main/java/org/opensearch/client/support/AbstractClient.java +++ b/server/src/main/java/org/opensearch/client/support/AbstractClient.java @@ -408,14 +408,13 @@ import org.opensearch.client.AdminClient; import org.opensearch.client.Client; import org.opensearch.client.ClusterAdminClient; -import org.opensearch.client.FilterClient; import org.opensearch.client.IndicesAdminClient; import org.opensearch.client.OpenSearchClient; import org.opensearch.cluster.metadata.IndexMetadata.APIBlock; import org.opensearch.common.Nullable; import org.opensearch.common.action.ActionFuture; import org.opensearch.common.settings.Settings; -import org.opensearch.common.util.concurrent.ThreadContext; +import org.opensearch.common.util.concurrent.StashAndMergeHeadersFilterClient; import org.opensearch.core.action.ActionListener; import org.opensearch.core.action.ActionResponse; import org.opensearch.core.common.bytes.BytesReference; @@ -2140,18 +2139,6 @@ public ActionFuture updateView(CreateViewAction.Request @Override public Client filterWithHeader(Map headers) { - return new FilterClient(this) { - @Override - protected void doExecute( - ActionType action, - Request request, - ActionListener listener - ) { - ThreadContext threadContext = threadPool().getThreadContext(); - try (ThreadContext.StoredContext ctx = threadContext.stashAndMergeHeaders(headers)) { - super.doExecute(action, request, listener); - } - } - }; + return new StashAndMergeHeadersFilterClient(this, headers); } } diff --git a/server/src/main/java/org/opensearch/common/util/concurrent/StashAndMergeHeadersFilterClient.java b/server/src/main/java/org/opensearch/common/util/concurrent/StashAndMergeHeadersFilterClient.java new file mode 100644 index 0000000000000..77655f838f826 --- /dev/null +++ b/server/src/main/java/org/opensearch/common/util/concurrent/StashAndMergeHeadersFilterClient.java @@ -0,0 +1,45 @@ +/* + * 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.common.util.concurrent; + +import org.opensearch.action.ActionRequest; +import org.opensearch.action.ActionType; +import org.opensearch.client.Client; +import org.opensearch.client.FilterClient; +import org.opensearch.core.action.ActionListener; +import org.opensearch.core.action.ActionResponse; + +import java.util.Map; + +/** + * This class is used for simulating request headers for testing. Used in TasksIT to + * assert that request headers are carried to task info + * + * @opensearch.internal + */ +public class StashAndMergeHeadersFilterClient extends FilterClient { + private Map headers; + + public StashAndMergeHeadersFilterClient(Client in, Map headers) { + super(in); + this.headers = headers; + } + + @Override + protected void doExecute( + ActionType action, + Request request, + ActionListener listener + ) { + ThreadContext threadContext = threadPool().getThreadContext(); + try (ThreadContext.StoredContext ctx = threadContext.stashAndMergeHeaders(headers)) { + super.doExecute(action, request, listener); + } + } +} diff --git a/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java b/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java index 906a27e9f398c..3ed9149a87743 100644 --- a/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java +++ b/server/src/main/java/org/opensearch/common/util/concurrent/ThreadContext.java @@ -217,7 +217,7 @@ public StoredContext stashWithOrigin(String origin) { * The removed context can be restored when closing the returned {@link StoredContext}. The merge strategy is that headers * that are already existing are preserved unless they are defaults. */ - public StoredContext stashAndMergeHeaders(Map headers) { + StoredContext stashAndMergeHeaders(Map headers) { final ThreadContextStruct context = threadLocal.get(); Map newHeader = new HashMap<>(headers); newHeader.putAll(context.requestHeaders); From fadc5feef7d3df079830c503041b25748c7df3fc Mon Sep 17 00:00:00 2001 From: Craig Perkins Date: Thu, 25 Jul 2024 23:18:43 -0400 Subject: [PATCH 2/2] Add to CHANGELOG Signed-off-by: Craig Perkins --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e88a084f7d7f6..7d6308211d52b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), - Bump `org.apache.commons:commons-lang3` from 3.14.0 to 3.15.0 ([#14861](https://github.com/opensearch-project/OpenSearch/pull/14861)) ### Changed +- Make ThreadContext.stashAndMergeHeaders package-private ([#14983](https://github.com/opensearch-project/OpenSearch/pull/14983)) ### Deprecated