forked from opensearch-project/OpenSearch
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
add transport interceptor to populate queryGroupId in task headers
Signed-off-by: Kaushal Kumar <[email protected]>
- Loading branch information
1 parent
2e13b79
commit c926996
Showing
9 changed files
with
248 additions
and
17 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
19 changes: 19 additions & 0 deletions
19
server/src/main/java/org/opensearch/wlm/QueryGroupConstants.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,19 @@ | ||
| /* | ||
| * 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 java.util.function.Supplier; | ||
|
|
||
| /** | ||
| * This class will hold all the QueryGroup related constants | ||
| */ | ||
| public class QueryGroupConstants { | ||
| public static final String QUERY_GROUP_ID_HEADER = "queryGroupId"; | ||
| public static final Supplier<String> DEFAULT_QUERY_GROUP_ID_SUPPLIER = () -> "DEFAULT_QUERY_GROUP"; | ||
| } |
53 changes: 53 additions & 0 deletions
53
server/src/main/java/org/opensearch/wlm/SearchWorkloadTransportHandler.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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.search.fetch.ShardFetchRequest; | ||
| import org.opensearch.search.internal.InternalScrollSearchRequest; | ||
| import org.opensearch.search.internal.ShardSearchRequest; | ||
| import org.opensearch.search.query.QuerySearchRequest; | ||
| import org.opensearch.tasks.Task; | ||
| import org.opensearch.threadpool.ThreadPool; | ||
| import org.opensearch.transport.TransportChannel; | ||
| import org.opensearch.transport.TransportRequest; | ||
| import org.opensearch.transport.TransportRequestHandler; | ||
|
|
||
| /** | ||
| * This class is mainly used to populate the queryGroupId header | ||
| * @param <T> T is Search related request | ||
| */ | ||
| public class SearchWorkloadTransportHandler<T extends TransportRequest> implements TransportRequestHandler<T> { | ||
|
|
||
| private final ThreadPool threadPool; | ||
| TransportRequestHandler<T> actualHandler; | ||
|
|
||
| public SearchWorkloadTransportHandler(ThreadPool threadPool, TransportRequestHandler<T> actualHandler) { | ||
| this.threadPool = threadPool; | ||
| this.actualHandler = actualHandler; | ||
| } | ||
|
|
||
| @Override | ||
| public void messageReceived(T request, TransportChannel channel, Task task) throws Exception { | ||
| if (isSearchWorkloadRequest(request)) { | ||
| task.addHeader( | ||
| QueryGroupConstants.QUERY_GROUP_ID_HEADER, | ||
| threadPool.getThreadContext(), | ||
| QueryGroupConstants.DEFAULT_QUERY_GROUP_ID_SUPPLIER | ||
| ); | ||
| } | ||
| actualHandler.messageReceived(request, channel, task); | ||
| } | ||
|
|
||
| private boolean isSearchWorkloadRequest(TransportRequest request) { | ||
| return (request instanceof ShardSearchRequest) | ||
| || (request instanceof ShardFetchRequest) | ||
| || (request instanceof InternalScrollSearchRequest) | ||
| || (request instanceof QuerySearchRequest); | ||
| } | ||
| } |
37 changes: 37 additions & 0 deletions
37
server/src/main/java/org/opensearch/wlm/SearchWorkloadTransportInterceptor.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,37 @@ | ||
| /* | ||
| * 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.threadpool.ThreadPool; | ||
| import org.opensearch.transport.TransportInterceptor; | ||
| import org.opensearch.transport.TransportRequest; | ||
| import org.opensearch.transport.TransportRequestHandler; | ||
|
|
||
| /** | ||
| * This class is used to intercept search traffic requests and populate the queryGroupId header in task headers | ||
| * TODO: We still need to add this interceptor in {@link org.opensearch.node.Node} class to enable, | ||
| * leaving it until the feature is tested and done. | ||
| */ | ||
| public class SearchWorkloadTransportInterceptor implements TransportInterceptor { | ||
| private final ThreadPool threadPool; | ||
|
|
||
| public SearchWorkloadTransportInterceptor(ThreadPool threadPool) { | ||
| this.threadPool = threadPool; | ||
| } | ||
|
|
||
| @Override | ||
| public <T extends TransportRequest> TransportRequestHandler<T> interceptHandler( | ||
| String action, | ||
| String executor, | ||
| boolean forceExecution, | ||
| TransportRequestHandler<T> actualHandler | ||
| ) { | ||
| return new SearchWorkloadTransportHandler<T>(threadPool, actualHandler); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
84 changes: 84 additions & 0 deletions
84
server/src/test/java/org/opensearch/wlm/SearchWorkloadTransportHandlerTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,84 @@ | ||
| /* | ||
| * 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.action.index.IndexRequest; | ||
| import org.opensearch.search.internal.ShardSearchRequest; | ||
| import org.opensearch.tasks.Task; | ||
| import org.opensearch.test.OpenSearchTestCase; | ||
| import org.opensearch.threadpool.TestThreadPool; | ||
| import org.opensearch.threadpool.ThreadPool; | ||
| import org.opensearch.transport.TransportChannel; | ||
| import org.opensearch.transport.TransportRequest; | ||
| import org.opensearch.transport.TransportRequestHandler; | ||
|
|
||
| import java.util.Collections; | ||
|
|
||
| import static org.mockito.Mockito.any; | ||
| import static org.mockito.Mockito.mock; | ||
| import static org.mockito.Mockito.spy; | ||
| import static org.mockito.Mockito.times; | ||
| import static org.mockito.Mockito.verify; | ||
|
|
||
| public class SearchWorkloadTransportHandlerTests extends OpenSearchTestCase { | ||
| private SearchWorkloadTransportHandler<TransportRequest> sut; | ||
| private ThreadPool threadPool; | ||
|
|
||
| private TransportRequestHandler<TransportRequest> actualHandler; | ||
|
|
||
| public void setUp() throws Exception { | ||
| super.setUp(); | ||
| threadPool = new TestThreadPool(getTestName()); | ||
| actualHandler = new TestTransportRequestHandler<>(); | ||
|
|
||
| sut = new SearchWorkloadTransportHandler<>(threadPool, actualHandler); | ||
| } | ||
|
|
||
| public void tearDown() throws Exception { | ||
| super.tearDown(); | ||
| threadPool.shutdown(); | ||
| } | ||
|
|
||
| public void testMessageReceivedForSearchWorkload() throws Exception { | ||
| ShardSearchRequest request = mock(ShardSearchRequest.class); | ||
| Task spyTask = getSpyTask(); | ||
|
|
||
| sut.messageReceived(request, mock(TransportChannel.class), spyTask); | ||
|
|
||
| verify(spyTask, times(1)).addHeader( | ||
| QueryGroupConstants.QUERY_GROUP_ID_HEADER, | ||
| threadPool.getThreadContext(), | ||
| QueryGroupConstants.DEFAULT_QUERY_GROUP_ID_SUPPLIER | ||
| ); | ||
| } | ||
|
|
||
| public void testMessageReceivedForNonSearchWorkload() throws Exception { | ||
| IndexRequest indexRequest = mock(IndexRequest.class); | ||
| Task spyTask = getSpyTask(); | ||
| sut.messageReceived(indexRequest, mock(TransportChannel.class), spyTask); | ||
|
|
||
| verify(spyTask, times(0)).addHeader(any(), any(), any()); | ||
| } | ||
|
|
||
| private static Task getSpyTask() { | ||
| final Task task = new Task(123, "transport", "Search", "test task", null, Collections.emptyMap()); | ||
|
|
||
| return spy(task); | ||
| } | ||
|
|
||
| private static class TestTransportRequestHandler<T extends TransportRequest> implements TransportRequestHandler<T> { | ||
| int invokeCount = 0; | ||
|
|
||
| @Override | ||
| public void messageReceived(TransportRequest request, TransportChannel channel, Task task) throws Exception { | ||
| invokeCount += 1; | ||
| } | ||
|
|
||
| }; | ||
| } |
37 changes: 37 additions & 0 deletions
37
server/src/test/java/org/opensearch/wlm/SearchWorkloadTransportInterceptorTests.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,37 @@ | ||
| /* | ||
| * 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 org.opensearch.threadpool.TestThreadPool; | ||
| import org.opensearch.threadpool.ThreadPool; | ||
| import org.opensearch.transport.TransportRequest; | ||
| import org.opensearch.transport.TransportRequestHandler; | ||
|
|
||
| import static org.opensearch.threadpool.ThreadPool.Names.SAME; | ||
|
|
||
| public class SearchWorkloadTransportInterceptorTests extends OpenSearchTestCase { | ||
|
|
||
| private ThreadPool threadPool; | ||
| private SearchWorkloadTransportInterceptor sut; | ||
|
|
||
| public void setUp() throws Exception { | ||
| threadPool = new TestThreadPool(getTestName()); | ||
| sut = new SearchWorkloadTransportInterceptor(threadPool); | ||
| } | ||
|
|
||
| public void tearDown() throws Exception { | ||
| threadPool.shutdown(); | ||
| } | ||
|
|
||
| public void testInterceptHandler() { | ||
| TransportRequestHandler<TransportRequest> requestHandler = sut.interceptHandler("Search", SAME, false, null); | ||
| assertTrue(requestHandler instanceof SearchWorkloadTransportHandler); | ||
| } | ||
| } |