Add search backpressure service check for query group tasks (#17576) (#17591)

* add validation check for SBP service



* add SBP should not handle the task tracking UT



* add CHANGELOG entry



* fix broken UT



* address comments



---------


(cherry picked from commit 8ee5eeb)

Signed-off-by: Kaushal Kumar <[email protected]>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: opensearch-trigger-bot[bot]

CHANGELOG.md

@@ -39,6 +39,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 - Fix visit of inner query for FunctionScoreQueryBuilder ([#16776](https://github.com/opensearch-project/OpenSearch/pull/16776))
 - Fix case insensitive and escaped query on wildcard ([#16827](https://github.com/opensearch-project/OpenSearch/pull/16827))
 - 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
 

server/src/main/java/org/opensearch/wlm/QueryGroupService.java

@@ -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() && !QueryGroupTask.DEFAULT_QUERY_GROUP_ID_SUPPLIER.get().equals(task.getQueryGroupId())) {
             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;

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();
@@ -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);
@@ -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<QueryGroup> getActiveQueryGroups(