From 466bc4247d91a59059ed1c906233f634f9698446 Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Tue, 11 Mar 2025 14:38:21 -0700 Subject: [PATCH 1/4] add validation check for SBP service Signed-off-by: Kaushal Kumar --- .../src/main/java/org/opensearch/wlm/QueryGroupService.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java index 14002a2b38134..83fb2fa879084 100644 --- a/server/src/main/java/org/opensearch/wlm/QueryGroupService.java +++ b/server/src/main/java/org/opensearch/wlm/QueryGroupService.java @@ -331,7 +331,7 @@ public Set 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())); } @@ -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; From 12b8635133e882e0691521341c6e8fd3bbf2d676 Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Tue, 11 Mar 2025 15:12:51 -0700 Subject: [PATCH 2/4] add SBP should not handle the task tracking UT Signed-off-by: Kaushal Kumar --- .../org/opensearch/wlm/QueryGroupServiceTests.java | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java index f22759ce968aa..ccc199413156e 100644 --- a/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java @@ -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 activeQueryGroups = new HashSet<>(); mockQueryGroupStateMap.put("testId", queryGroupState); @@ -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)); @@ -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 getActiveQueryGroups( From 728d2bd38cce8e516a854dd07e2680ea9dcdc4ce Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Tue, 11 Mar 2025 15:36:01 -0700 Subject: [PATCH 3/4] add CHANGELOG entry Signed-off-by: Kaushal Kumar --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index db39772ceef87..49e27315be3f6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 From d0f9e148ccbcea54655e9d8d2fbf003a0f94d639 Mon Sep 17 00:00:00 2001 From: Kaushal Kumar Date: Wed, 12 Mar 2025 10:00:57 -0700 Subject: [PATCH 4/4] fix broken UT Signed-off-by: Kaushal Kumar --- .../test/java/org/opensearch/wlm/QueryGroupServiceTests.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java index ccc199413156e..579d65846f69b 100644 --- a/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java +++ b/server/src/test/java/org/opensearch/wlm/QueryGroupServiceTests.java @@ -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();