Skip to content

Fix phantom WLM rejections and stats for coordinator and shard search tasks #22839

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LilyCaroline17 wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Fix phantom WLM rejections and stats for coordinator and shard search tasks #22839

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
f66d002
Fix phantom WLM rejections for coordinator search tasks
LilyCaroline17 Jul 23, 2026
09b367a
Exclude rejected search tasks from WLM completion counts
LilyCaroline17 Aug 26, 2026
d82976c
Merge branch 'main' into fix-coordinator-search-rejection
peterzhuamazon Aug 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions server/src/main/java/org/opensearch/action/search/StreamTransportSearchAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
import org.opensearch.transport.StreamTransportService;
import org.opensearch.transport.Transport;
import org.opensearch.transport.client.node.NodeClient;
import org.opensearch.wlm.WorkloadGroupService;

import java.util.Map;
import java.util.Set;
Expand Down Expand Up @@ -59,7 +60,8 @@ public StreamTransportSearchAction(
SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory,
Tracer tracer,
TaskResourceTrackingService taskResourceTrackingService,
IndicesService indicesService
IndicesService indicesService,
WorkloadGroupService workloadGroupService
) {
super(
client,
Expand All @@ -78,7 +80,8 @@ public StreamTransportSearchAction(
searchRequestOperationsCompositeListenerFactory,
tracer,
taskResourceTrackingService,
indicesService
indicesService,
workloadGroupService
);
}

Expand Down
26 changes: 21 additions & 5 deletions server/src/main/java/org/opensearch/action/search/TransportSearchAction.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,7 @@
import org.opensearch.core.common.breaker.CircuitBreaker;
import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
import org.opensearch.core.common.io.stream.Writeable;
import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
import org.opensearch.core.index.Index;
import org.opensearch.core.index.shard.ShardId;
import org.opensearch.core.indices.breaker.CircuitBreakerService;
Expand Down Expand Up @@ -110,6 +111,7 @@
import org.opensearch.transport.client.Client;
import org.opensearch.transport.client.OriginSettingClient;
import org.opensearch.transport.client.node.NodeClient;
import org.opensearch.wlm.WorkloadGroupService;
import org.opensearch.wlm.WorkloadGroupTask;

import java.util.ArrayList;
Expand Down Expand Up @@ -194,6 +196,8 @@ public class TransportSearchAction extends HandledTransportAction<SearchRequest,

private final SearchIndexPruningService searchIndexPruningService;

private final WorkloadGroupService workloadGroupService;

@Inject
public TransportSearchAction(
NodeClient client,
Expand All @@ -212,7 +216,8 @@ public TransportSearchAction(
SearchRequestOperationsCompositeListenerFactory searchRequestOperationsCompositeListenerFactory,
Tracer tracer,
TaskResourceTrackingService taskResourceTrackingService,
IndicesService indicesService
IndicesService indicesService,
WorkloadGroupService workloadGroupService
) {
super(SearchAction.NAME, transportService, actionFilters, (Writeable.Reader<SearchRequest>) SearchRequest::new);
this.client = client;
Expand Down Expand Up @@ -240,6 +245,7 @@ public TransportSearchAction(
clusterService.getClusterSettings(),
new ClusterStateFieldDomainProvider()
);
this.workloadGroupService = workloadGroupService;
}

private Map<String, AliasFilter> buildPerIndexAliasFilter(
Expand Down Expand Up @@ -483,14 +489,24 @@ void executeRequest(
originalSearchRequest,
taskResourceTrackingService::getTaskResourceUsageFromThreadContext
);
searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);

// At this point either the QUERY_GROUP_ID header will be present in ThreadContext either via ActionFilter
// or HTTP header (HTTP header will be deprecated once ActionFilter is implemented)
if (task instanceof WorkloadGroupTask) {
((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext());
if (task instanceof WorkloadGroupTask workloadGroupTask) {
// Coordinator-task admission point. Runs before onRequestStart (keeps the in-flight gauge balanced) and
// before setWorkloadGroupId (so a rejected task is not counted in total_completions on task completion).
try {
workloadGroupService.rejectIfNeeded(
threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER)
);
} catch (OpenSearchRejectedExecutionException e) {
updatedListener.onFailure(e);

@dzane17 dzane17 Aug 25, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Right now a rejected request will count toward both total_rejections and total_completions.

You might be able to skip total_completions increment by callling setWorkloadGroupId() after the rejection check.

if (task instanceof WorkloadGroupTask workloadGroupTask) {
    // Admission must precede onRequestStart (in-flight gauge) and setWorkloadGroupId
    // (so a rejected task is not counted in total_completions).
    try {
        workloadGroupService.rejectIfNeeded(threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER));
    } catch (OpenSearchRejectedExecutionException e) {
        updatedListener.onFailure(e);
        return;
    }
    workloadGroupTask.setWorkloadGroupId(threadPool.getThreadContext());
}

@LilyCaroline17 LilyCaroline17 Aug 26, 2026 •
edited
Loading

Copy link
Copy Markdown
Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Oh, I see what you mean. onTaskCompleted in WorkloadGroupService would include it in the total_completions counter since setWorkloadGroupId sets isWorkloadGroupSet to true. Good catch, thanks! Updating now.

return;
}
workloadGroupTask.setWorkloadGroupId(threadPool.getThreadContext());
}

searchRequestContext.getSearchRequestOperationsListener().onRequestStart(searchRequestContext);

PipelinedRequest searchRequest;
ActionListener<SearchResponse> listener;
try {
Expand Down
4 changes: 2 additions & 2 deletions server/src/main/java/org/opensearch/wlm/WorkloadManagementTransportInterceptor.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,9 +56,9 @@ public RequestHandler(ThreadPool threadPool, TransportRequestHandler<T> actualHa
@Override
public void messageReceived(T request, TransportChannel channel, Task task) throws Exception {
if (isSearchWorkloadRequest(task)) {
// Reject before setWorkloadGroupId so a rejected task is not tagged and thus not counted as a phantom completion.
workloadGroupService.rejectIfNeeded(threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER));
((WorkloadGroupTask) task).setWorkloadGroupId(threadPool.getThreadContext());
final String workloadGroupId = ((WorkloadGroupTask) (task)).getWorkloadGroupId();
workloadGroupService.rejectIfNeeded(workloadGroupId);
}
actualHandler.messageReceived(request, channel, task);
}
Expand Down
3 changes: 1 addition & 2 deletions server/src/main/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListener.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,8 +42,7 @@ public WorkloadGroupRequestOperationListener(WorkloadGroupService workloadGroupS
*/
@Override
protected void onRequestStart(SearchRequestContext searchRequestContext) {
final String workloadGroupId = threadPool.getThreadContext().getHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER);
workloadGroupService.rejectIfNeeded(workloadGroupId);
// Do not reject here: CompositeListener swallows exceptions thrown in onRequestStart.
WorkloadGroup workloadGroup = workloadGroupService.getCurrentWorkloadGroup();
applyWorkloadGroupSearchSettings(workloadGroup, searchRequestContext.getRequest());
}
Expand Down
89 changes: 88 additions & 1 deletion server/src/test/java/org/opensearch/action/search/TransportSearchActionTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -66,6 +66,7 @@
import org.opensearch.core.common.Strings;
import org.opensearch.core.common.io.stream.NamedWriteableRegistry;
import org.opensearch.core.common.transport.TransportAddress;
import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
import org.opensearch.core.index.Index;
import org.opensearch.core.index.shard.ShardId;
import org.opensearch.core.indices.breaker.CircuitBreakerService;
Expand Down Expand Up @@ -113,6 +114,8 @@
import org.opensearch.transport.TransportRequestOptions;
import org.opensearch.transport.TransportService;
import org.opensearch.transport.client.node.NodeClient;
import org.opensearch.wlm.WorkloadGroupService;
import org.opensearch.wlm.WorkloadGroupTask;

import java.util.ArrayList;
import java.util.Arrays;
Expand All @@ -125,6 +128,7 @@
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.BiFunction;
Expand All @@ -135,7 +139,11 @@
import static org.hamcrest.CoreMatchers.containsString;
import static org.hamcrest.CoreMatchers.instanceOf;
import static org.hamcrest.CoreMatchers.startsWith;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;

public class TransportSearchActionTests extends OpenSearchTestCase {
Expand Down Expand Up @@ -1251,7 +1259,8 @@ public void testResolveIndices() {
new SearchRequestOperationsCompositeListenerFactory(),
mock(Tracer.class),
mock(TaskResourceTrackingService.class),
mock(IndicesService.class)
mock(IndicesService.class),
mock(WorkloadGroupService.class)
);

// Actual test cases start here:
Expand Down Expand Up @@ -1292,4 +1301,82 @@ public void testResolveIndices() {
}
}
}

public void testCoordinatorSearchTaskRejectedBeforeRequestStart() {
ClusterService clusterService = mock(ClusterService.class);
when(clusterService.getClusterSettings()).thenReturn(
new ClusterSettings(Settings.EMPTY, ClusterSettings.BUILT_IN_CLUSTER_SETTINGS)
);

WorkloadGroupService workloadGroupService = mock(WorkloadGroupService.class);
doThrow(new OpenSearchRejectedExecutionException("WorkloadGroup is already contended.")).when(workloadGroupService)
.rejectIfNeeded(anyString());

// onRequestStart must not run on rejection, else the in-flight gauge leaks.
AtomicBoolean requestStarted = new AtomicBoolean(false);
SearchRequestOperationsListener requestStartTracker = new SearchRequestOperationsListener() {
@Override
protected void onRequestStart(SearchRequestContext searchRequestContext) {
requestStarted.set(true);
}
};

TransportSearchAction action = new TransportSearchAction(
mock(NodeClient.class),
threadPool,
mock(CircuitBreakerService.class),
mock(TransportService.class),
mock(SearchService.class),
mock(SearchTransportService.class),
new SearchPhaseController(new NamedWriteableRegistry(Collections.emptyList()), (searchSourceBuilder) -> null),
clusterService,
mock(ActionFilters.class),
new IndexNameExpressionResolver(new ThreadContext(Settings.EMPTY)),
new NamedWriteableRegistry(Collections.emptyList()),
mock(SearchPipelineService.class),
mock(MetricsRegistry.class),
new SearchRequestOperationsCompositeListenerFactory(requestStartTracker),
NoopTracer.INSTANCE,
mock(TaskResourceTrackingService.class),
mock(IndicesService.class),
workloadGroupService
);

SearchTask task = new SearchTask(0, "transport", SearchAction.NAME, () -> "test", null, Collections.emptyMap());

AtomicReference<Exception> failure = new AtomicReference<>();
try (ThreadContext.StoredContext ignored = threadPool.getThreadContext().stashContext()) {
threadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, "test-workload-group");
action.executeRequest(
task,
new SearchRequest(),
(TransportSearchAction.SearchAsyncActionProvider) (
searchTask,
searchRequest,
executor,
shardIterators,
timeProvider,
connectionLookup,
clusterState,
aliasFilter,
concreteIndexBoosts,
indexRoutings,
listener,
preFilter,
tp,
clusters,
searchRequestContext) -> {
return null;
},
ActionListener.wrap(r -> fail("expected rejection"), failure::set)
);
}

assertThat(failure.get(), instanceOf(OpenSearchRejectedExecutionException.class));
assertFalse("onRequestStart must not run for a rejected request", requestStarted.get());
// A rejected task must not be tagged, otherwise onTaskCompleted would count it as a phantom completion.
assertFalse("rejected task must not have its workload group id set", task.isWorkloadGroupSet());
// Admission reads the header directly (before setWorkloadGroupId); pin that value since the stub matches anyString().
verify(workloadGroupService).rejectIfNeeded(eq("test-workload-group"));
}
}
4 changes: 3 additions & 1 deletion server/src/test/java/org/opensearch/snapshots/SnapshotResiliencyTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -251,6 +251,7 @@
import org.opensearch.transport.TransportService;
import org.opensearch.transport.client.AdminClient;
import org.opensearch.transport.client.node.NodeClient;
import org.opensearch.wlm.WorkloadGroupService;
import org.junit.After;
import org.junit.Before;

Expand Down Expand Up @@ -2410,7 +2411,8 @@ public void onFailure(final Exception e) {
searchRequestOperationsCompositeListenerFactory,
NoopTracer.INSTANCE,
new TaskResourceTrackingService(settings, clusterSettings, threadPool),
mockIndicesService
mockIndicesService,
mock(WorkloadGroupService.class)
)
);
actions.put(
Expand Down
12 changes: 9 additions & 3 deletions server/src/test/java/org/opensearch/wlm/WorkloadManagementTransportRequestHandlerTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@

import java.util.Collections;

import static org.mockito.Mockito.anyString;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
Expand Down Expand Up @@ -51,17 +51,23 @@ public void tearDown() throws Exception {
public void testMessageReceivedForSearchWorkload_nonRejectionCase() throws Exception {
ShardSearchRequest request = mock(ShardSearchRequest.class);
WorkloadGroupTask spyTask = getSpyTask();
doNothing().when(workloadGroupService).rejectIfNeeded(anyString());
doNothing().when(workloadGroupService).rejectIfNeeded(any());
sut.messageReceived(request, mock(TransportChannel.class), spyTask);
assertTrue(sut.isSearchWorkloadRequest(spyTask));
// Admitted task is tagged and forwarded to the wrapped handler.
assertTrue(spyTask.isWorkloadGroupSet());
assertEquals(1, actualHandler.invokeCount);
}

public void testMessageReceivedForSearchWorkload_RejectionCase() throws Exception {
ShardSearchRequest request = mock(ShardSearchRequest.class);
WorkloadGroupTask spyTask = getSpyTask();
doThrow(OpenSearchRejectedExecutionException.class).when(workloadGroupService).rejectIfNeeded(anyString());
doThrow(OpenSearchRejectedExecutionException.class).when(workloadGroupService).rejectIfNeeded(any());

assertThrows(OpenSearchRejectedExecutionException.class, () -> sut.messageReceived(request, mock(TransportChannel.class), spyTask));
// A rejected task must not be tagged (else onTaskCompleted would count it as a phantom completion) and must not be forwarded.
assertFalse(spyTask.isWorkloadGroupSet());
assertEquals(0, actualHandler.invokeCount);
}

public void testMessageReceivedForNonSearchWorkload() throws Exception {
Expand Down
18 changes: 0 additions & 18 deletions ...rc/test/java/org/opensearch/wlm/listeners/WorkloadGroupRequestOperationListenerTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
import org.opensearch.common.settings.Settings;
import org.opensearch.common.unit.TimeValue;
import org.opensearch.common.util.concurrent.ThreadContext;
import org.opensearch.core.concurrency.OpenSearchRejectedExecutionException;
import org.opensearch.search.builder.SearchSourceBuilder;
import org.opensearch.test.OpenSearchTestCase;
import org.opensearch.threadpool.TestThreadPool;
Expand All @@ -40,8 +39,6 @@
import java.util.List;
import java.util.Map;

import static org.mockito.Mockito.doNothing;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;

Expand Down Expand Up @@ -82,21 +79,6 @@ public void tearDown() throws Exception {
testThreadPool.shutdown();
}

public void testRejectionCase() {
final String testWorkloadGroupId = "asdgasgkajgkw3141_3rt4t";
testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, testWorkloadGroupId);
doThrow(OpenSearchRejectedExecutionException.class).when(workloadGroupService).rejectIfNeeded(testWorkloadGroupId);
assertThrows(OpenSearchRejectedExecutionException.class, () -> sut.onRequestStart(mockSearchRequestContext));
}

public void testNonRejectionCase() {
final String testWorkloadGroupId = "asdgasgkajgkw3141_3rt4t";
testThreadPool.getThreadContext().putHeader(WorkloadGroupTask.WORKLOAD_GROUP_ID_HEADER, testWorkloadGroupId);
doNothing().when(workloadGroupService).rejectIfNeeded(testWorkloadGroupId);

sut.onRequestStart(mockSearchRequestContext);
}

public void testValidWorkloadGroupRequestFailure() throws IOException {

WorkloadGroupStats expectedStats = new WorkloadGroupStats(
Expand Down
Loading