Skip to content

Commit

Permalink
Make ThreadContext.stashAndMergeHeaders package-private
Browse files Browse the repository at this point in the history
Signed-off-by: Craig Perkins <cwperx@amazon.com>
  • Loading branch information
cwperks committed Jul 26, 2024
1 parent 157d277 commit 56bebf6
Show file tree
Hide file tree
Showing 3 changed files with 48 additions and 16 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -2140,18 +2139,6 @@ public ActionFuture<GetViewAction.Response> updateView(CreateViewAction.Request

@Override
public Client filterWithHeader(Map<String, String> headers) {
return new FilterClient(this) {
@Override
protected <Request extends ActionRequest, Response extends ActionResponse> void doExecute(
ActionType<Response> action,
Request request,
ActionListener<Response> listener
) {
ThreadContext threadContext = threadPool().getThreadContext();
try (ThreadContext.StoredContext ctx = threadContext.stashAndMergeHeaders(headers)) {
super.doExecute(action, request, listener);
}
}
};
return new StashAndMergeHeadersFilterClient(this, headers);
}
}
Original file line number Diff line number Diff line change
@@ -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<String, String> headers;

public StashAndMergeHeadersFilterClient(Client in, Map<String, String> headers) {
super(in);
this.headers = headers;
}

@Override
protected <Request extends ActionRequest, Response extends ActionResponse> void doExecute(
ActionType<Response> action,
Request request,
ActionListener<Response> listener
) {
ThreadContext threadContext = threadPool().getThreadContext();
try (ThreadContext.StoredContext ctx = threadContext.stashAndMergeHeaders(headers)) {
super.doExecute(action, request, listener);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -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<String, String> headers) {
StoredContext stashAndMergeHeaders(Map<String, String> headers) {
final ThreadContextStruct context = threadLocal.get();
Map<String, String> newHeader = new HashMap<>(headers);
newHeader.putAll(context.requestHeaders);
Expand Down

0 comments on commit 56bebf6

Please sign in to comment.