From d3392a23ce15830276b066b5d500dd9cb6dd0736 Mon Sep 17 00:00:00 2001
From: Laksh Singla <lakshsingla@gmail.com>
Date: Thu, 12 Sep 2024 15:10:27 +0530
Subject: [PATCH] Cancel the group by processing tasks if the merging runner
 gets scheduled post the query timeout (#17037)

If the GroupByMergingQueryRunner gets scheduled after the query timeout, it fails to clean up the processing tasks that have been scheduled. This can lead to unnecessary processing being done for the tasks whos results won't get consumed.
---
 .../query/groupby/epinephelinae/GroupByMergingQueryRunner.java  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByMergingQueryRunner.java b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByMergingQueryRunner.java
index 6e2064d5f1f..75e8fa2c2ce 100644
--- a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByMergingQueryRunner.java
+++ b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByMergingQueryRunner.java
@@ -375,7 +375,7 @@ public class GroupByMergingQueryRunner implements QueryRunner<ResultRow>
       GuavaUtils.cancelAll(true, future, futures);
       throw new QueryInterruptedException(e);
     }
-    catch (TimeoutException e) {
+    catch (QueryTimeoutException | TimeoutException e) {
       log.info("Query timeout, cancelling pending results for query id [%s]", query.getId());
       GuavaUtils.cancelAll(true, future, futures);
       throw new QueryTimeoutException();