Skip to content
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

Add search backpressure service check for query group tasks #17576

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
+16 −4
Open

Add search backpressure service check for query group tasks #17576

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
466bc42
add validation check for SBP service
kaushalmahi12 Mar 11, 2025
12b8635
add SBP should not handle the task tracking UT
kaushalmahi12 Mar 11, 2025
728d2bd
add CHANGELOG entry
kaushalmahi12 Mar 11, 2025
d0f9e14
fix broken UT
kaushalmahi12 Mar 12, 2025
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
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
- Fix exists queries on nested flat_object fields throws exception ([#16803](https://github.com/opensearch-project/OpenSearch/pull/16803))
- Add highlighting for wildcard search on `match_only_text` field ([#17101](https://github.com/opensearch-project/OpenSearch/pull/17101))
- Fix illegal argument exception when creating a PIT ([#16781](https://github.com/opensearch-project/OpenSearch/pull/16781))
- Fix NPE in node stats due to QueryGroupTasks ([#17576](https://github.com/opensearch-project/OpenSearch/pull/17576))

### Security

Expand Down
4 changes: 2 additions & 2 deletions server/src/main/java/org/opensearch/wlm/QueryGroupService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -331,7 +331,7 @@ public Set<QueryGroup> getDeletedQueryGroups() {
public boolean shouldSBPHandle(Task t) {
QueryGroupTask task = (QueryGroupTask) t;
boolean isInvalidQueryGroupTask = true;
if (!task.getQueryGroupId().equals(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get())) {
if (task.isQueryGroupSet() && !task.getQueryGroupId().equals(QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get())) {
isInvalidQueryGroupTask = activeQueryGroups.stream()
.noneMatch(queryGroup -> queryGroup.get_id().equals(task.getQueryGroupId()));
}
Expand All @@ -340,7 +340,7 @@ public boolean shouldSBPHandle(Task t) {

@Override
public void onTaskCompleted(Task task) {
if (!(task instanceof QueryGroupTask)) {
if (!(task instanceof QueryGroupTask) || !((QueryGroupTask) task).isQueryGroupSet()) {
return;
}
final QueryGroupTask queryGroupTask = (QueryGroupTask) task;
Expand Down
15 changes: 13 additions & 2 deletions server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -395,7 +395,7 @@ public void testRejectIfNeeded_whenFeatureIsNotEnabled() {
}

public void testOnTaskCompleted() {
Task task = createMockTaskWithResourceStats(SearchTask.class, 100, 200, 0, 12);
Task task = new SearchTask(12, "", "", () -> "", null, null);
mockThreadPool = new TestThreadPool("queryGroupServiceTests");
mockThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, "testId");
QueryGroupState queryGroupState = new QueryGroupState();
Expand Down Expand Up @@ -442,7 +442,7 @@ public void testOnTaskCompleted() {
}

public void testShouldSBPHandle() {
QueryGroupTask task = createMockTaskWithResourceStats(SearchTask.class, 100, 200, 0, 12);
SearchTask task = createMockTaskWithResourceStats(SearchTask.class, 100, 200, 0, 12);
QueryGroupState queryGroupState = new QueryGroupState();
Set<QueryGroup> activeQueryGroups = new HashSet<>();
mockQueryGroupStateMap.put("testId", queryGroupState);
Expand All @@ -464,6 +464,8 @@ public void testShouldSBPHandle() {
mockThreadPool = new TestThreadPool("queryGroupServiceTests");
mockThreadPool.getThreadContext()
.putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get());
// we haven't set the queryGroupId yet SBP should still track the task for cancellation
assertTrue(queryGroupService.shouldSBPHandle(task));
task.setQueryGroupId(mockThreadPool.getThreadContext());
assertTrue(queryGroupService.shouldSBPHandle(task));

Expand All @@ -490,6 +492,15 @@ public void testShouldSBPHandle() {
);
assertTrue(queryGroupService.shouldSBPHandle(task));

mockThreadPool.shutdownNow();

// test the case when SBP should not track the task
when(mockWorkloadManagementSettings.getWlmMode()).thenReturn(WlmMode.ENABLED);
task = new SearchTask(1, "", "test", () -> "", null, null);
mockThreadPool = new TestThreadPool("queryGroupServiceTests");
mockThreadPool.getThreadContext().putHeader(QueryGroupTask.QUERY_GROUP_ID_HEADER, "testId");
task.setQueryGroupId(mockThreadPool.getThreadContext());
assertFalse(queryGroupService.shouldSBPHandle(task));
}

private static Set<QueryGroup> getActiveQueryGroups(
Expand Down
Loading