Fix deadlock that can occur while merging group by results

benchmarks/src/test/java/org/apache/druid/benchmark/GroupByTypeInterfaceBenchmark.java

@@ -54,6 +54,7 @@
 import org.apache.druid.query.groupby.GroupByQueryConfig;
 import org.apache.druid.query.groupby.GroupByQueryQueryToolChest;
 import org.apache.druid.query.groupby.GroupByQueryRunnerFactory;
+import org.apache.druid.query.groupby.GroupByResourcesReservationPool;
 import org.apache.druid.query.groupby.GroupingEngine;
 import org.apache.druid.query.groupby.ResultRow;
 import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
@@ -372,19 +373,21 @@ public String getFormatString()
     };
 
     final Supplier<GroupByQueryConfig> configSupplier = Suppliers.ofInstance(config);
+    final GroupByResourcesReservationPool groupByResourcesReservationPool =
+        new GroupByResourcesReservationPool(mergePool, config);
     final GroupingEngine groupingEngine = new GroupingEngine(
         druidProcessingConfig,
         configSupplier,
         bufferPool,
-        mergePool,
+        groupByResourcesReservationPool,
         TestHelper.makeJsonMapper(),
         new ObjectMapper(new SmileFactory()),
         QueryBenchmarkUtil.NOOP_QUERYWATCHER
     );
 
     factory = new GroupByQueryRunnerFactory(
         groupingEngine,
-        new GroupByQueryQueryToolChest(groupingEngine)
+        new GroupByQueryQueryToolChest(groupingEngine, groupByResourcesReservationPool)
     );
   }
 

benchmarks/src/test/java/org/apache/druid/benchmark/query/CachingClusteredClientBenchmark.java

@@ -80,6 +80,7 @@
 import org.apache.druid.query.groupby.GroupByQueryQueryToolChest;
 import org.apache.druid.query.groupby.GroupByQueryRunnerFactory;
 import org.apache.druid.query.groupby.GroupByQueryRunnerTest;
+import org.apache.druid.query.groupby.GroupByResourcesReservationPool;
 import org.apache.druid.query.groupby.GroupingEngine;
 import org.apache.druid.query.groupby.ResultRow;
 import org.apache.druid.query.planning.DataSourceAnalysis;
@@ -356,16 +357,18 @@ private static GroupByQueryRunnerFactory makeGroupByQueryRunnerFactory(
         bufferSupplier,
         processingConfig.getNumMergeBuffers()
     );
+    final GroupByResourcesReservationPool groupByResourcesReservationPool =
+        new GroupByResourcesReservationPool(mergeBufferPool, config);
     final GroupingEngine groupingEngine = new GroupingEngine(
         processingConfig,
         configSupplier,
         bufferPool,
-        mergeBufferPool,
+        groupByResourcesReservationPool,
         mapper,
         mapper,
         QueryRunnerTestHelper.NOOP_QUERYWATCHER
     );
-    final GroupByQueryQueryToolChest toolChest = new GroupByQueryQueryToolChest(groupingEngine);
+    final GroupByQueryQueryToolChest toolChest = new GroupByQueryQueryToolChest(groupingEngine, groupByResourcesReservationPool);
     return new GroupByQueryRunnerFactory(groupingEngine, toolChest);
   }
 
@@ -469,7 +472,7 @@ private <T> List<T> runQuery()
             toolChestWarehouse.getToolChest(query)
         )
         .applyPreMergeDecoration()
-        .mergeResults()
+        .mergeResults(true)
         .applyPostMergeDecoration();
 
     //noinspection unchecked

benchmarks/src/test/java/org/apache/druid/benchmark/query/GroupByBenchmark.java

@@ -65,6 +65,7 @@
 import org.apache.druid.query.groupby.GroupByQueryConfig;
 import org.apache.druid.query.groupby.GroupByQueryQueryToolChest;
 import org.apache.druid.query.groupby.GroupByQueryRunnerFactory;
+import org.apache.druid.query.groupby.GroupByResourcesReservationPool;
 import org.apache.druid.query.groupby.GroupingEngine;
 import org.apache.druid.query.groupby.ResultRow;
 import org.apache.druid.query.groupby.orderby.DefaultLimitSpec;
@@ -487,19 +488,21 @@ public String getFormatString()
     };
 
     final Supplier<GroupByQueryConfig> configSupplier = Suppliers.ofInstance(config);
+    final GroupByResourcesReservationPool groupByResourcesReservationPool =
+        new GroupByResourcesReservationPool(mergePool, config);
     final GroupingEngine groupingEngine = new GroupingEngine(
         druidProcessingConfig,
         configSupplier,
         bufferPool,
-        mergePool,
+        groupByResourcesReservationPool,
         TestHelper.makeJsonMapper(),
         new ObjectMapper(new SmileFactory()),
         QueryBenchmarkUtil.NOOP_QUERYWATCHER
     );
 
     factory = new GroupByQueryRunnerFactory(
         groupingEngine,
-        new GroupByQueryQueryToolChest(groupingEngine)
+        new GroupByQueryQueryToolChest(groupingEngine, groupByResourcesReservationPool)
     );
   }
 

extensions-contrib/distinctcount/src/test/java/org/apache/druid/query/aggregation/distinctcount/DistinctCountGroupByQueryTest.java

@@ -137,9 +137,9 @@ public void testGroupByWithDistinctCountAgg() throws Exception
     Iterable<ResultRow> results = FluentQueryRunner
         .create(factory.createRunner(incrementalIndexSegment), factory.getToolchest())
         .applyPreMergeDecoration()
-        .mergeResults()
+        .mergeResults(true)
         .applyPostMergeDecoration()
-        .run(QueryPlus.wrap(query))
+        .run(QueryPlus.wrap(GroupByQueryRunnerTestHelper.populateResourceId(query)))
         .toList();
 
     List<ResultRow> expectedResults = Arrays.asList(

extensions-contrib/materialized-view-selection/src/main/java/org/apache/druid/query/materializedview/MaterializedViewQueryQueryToolChest.java

@@ -48,16 +48,35 @@ public MaterializedViewQueryQueryToolChest(
   {
     this.warehouse = warehouse;
   }
-  
+
   @Override
   public QueryRunner mergeResults(QueryRunner runner)
   {
-    return new QueryRunner() {
+    return new QueryRunner()
+    {
+      @Override
+      public Sequence run(QueryPlus queryPlus, ResponseContext responseContext)
+      {
+        Query realQuery = getRealQuery(queryPlus.getQuery());
+        return warehouse.getToolChest(realQuery)
+                        .mergeResults(runner)
+                        .run(queryPlus.withQuery(realQuery), responseContext);
+      }
+    };
+  }
+
+  @Override
+  public QueryRunner mergeResults(QueryRunner runner, boolean willMergeRunner)
+  {
+    return new QueryRunner()
+    {
       @Override
       public Sequence run(QueryPlus queryPlus, ResponseContext responseContext)
       {
         Query realQuery = getRealQuery(queryPlus.getQuery());
-        return warehouse.getToolChest(realQuery).mergeResults(runner).run(queryPlus.withQuery(realQuery), responseContext);
+        return warehouse.getToolChest(realQuery)
+                        .mergeResults(runner, willMergeRunner)
+                        .run(queryPlus.withQuery(realQuery), responseContext);
       }
     };
   }

extensions-contrib/materialized-view-selection/src/test/java/org/apache/druid/query/materializedview/MaterializedViewQueryQueryToolChestTest.java

@@ -127,7 +127,7 @@ public void testDecorateObjectMapper() throws IOException
     QueryToolChest queryToolChest =
         new MaterializedViewQueryQueryToolChest(new MapQueryToolChestWarehouse(
             ImmutableMap.<Class<? extends Query>, QueryToolChest>builder()
-                .put(GroupByQuery.class, new GroupByQueryQueryToolChest(null))
+                .put(GroupByQuery.class, new GroupByQueryQueryToolChest(null, null))
                 .build()
         ));
 
@@ -186,7 +186,7 @@ public void testDecorateObjectMapperMaterializedViewQuery() throws IOException
     QueryToolChest materializedViewQueryQueryToolChest =
         new MaterializedViewQueryQueryToolChest(new MapQueryToolChestWarehouse(
             ImmutableMap.<Class<? extends Query>, QueryToolChest>builder()
-                .put(GroupByQuery.class, new GroupByQueryQueryToolChest(null))
+                .put(GroupByQuery.class, new GroupByQueryQueryToolChest(null, null))
                 .build()
         ));
 
@@ -245,7 +245,7 @@ public void testGetRealQuery()
     MaterializedViewQueryQueryToolChest materializedViewQueryQueryToolChest =
         new MaterializedViewQueryQueryToolChest(new MapQueryToolChestWarehouse(
             ImmutableMap.<Class<? extends Query>, QueryToolChest>builder()
-                .put(GroupByQuery.class, new GroupByQueryQueryToolChest(null))
+                .put(GroupByQuery.class, new GroupByQueryQueryToolChest(null, null))
                 .build()
         ));
 

extensions-contrib/virtual-columns/src/test/java/org/apache/druid/segment/MapVirtualColumnGroupByTest.java

@@ -42,6 +42,7 @@
 import org.apache.druid.query.groupby.GroupByQueryConfig;
 import org.apache.druid.query.groupby.GroupByQueryQueryToolChest;
 import org.apache.druid.query.groupby.GroupByQueryRunnerFactory;
+import org.apache.druid.query.groupby.GroupByResourcesReservationPool;
 import org.apache.druid.query.groupby.GroupingEngine;
 import org.apache.druid.query.groupby.ResultRow;
 import org.apache.druid.query.spec.MultipleIntervalSegmentSpec;
@@ -67,6 +68,9 @@ public class MapVirtualColumnGroupByTest extends InitializedNullHandlingTest
   public void setup() throws IOException
   {
     final IncrementalIndex incrementalIndex = MapVirtualColumnTestBase.generateIndex();
+    final GroupByQueryConfig config = new GroupByQueryConfig();
+    final GroupByResourcesReservationPool groupByResourcesReservationPool =
+        new GroupByResourcesReservationPool(new DefaultBlockingPool<>(() -> ByteBuffer.allocate(1024), 1), config);
     final GroupingEngine groupingEngine = new GroupingEngine(
         new DruidProcessingConfig()
         {
@@ -94,17 +98,17 @@ public int getNumThreads()
             return 1;
           }
         },
-        GroupByQueryConfig::new,
+        () -> config,
         new StupidPool<>("map-virtual-column-groupby-test", () -> ByteBuffer.allocate(1024)),
-        new DefaultBlockingPool<>(() -> ByteBuffer.allocate(1024), 1),
+        groupByResourcesReservationPool,
         TestHelper.makeJsonMapper(),
         new DefaultObjectMapper(),
         QueryRunnerTestHelper.NOOP_QUERYWATCHER
     );
 
     final GroupByQueryRunnerFactory factory = new GroupByQueryRunnerFactory(
         groupingEngine,
-        new GroupByQueryQueryToolChest(groupingEngine)
+        new GroupByQueryQueryToolChest(groupingEngine, groupByResourcesReservationPool)
     );
 
     runner = QueryRunnerTestHelper.makeQueryRunner(

extensions-core/histogram/src/test/java/org/apache/druid/query/aggregation/histogram/ApproximateHistogramGroupByQueryTest.java

@@ -184,7 +184,8 @@ public void testGroupByWithApproximateHistogramAgg()
         )
     );
 
-    Iterable<ResultRow> results = runner.run(QueryPlus.wrap(query)).toList();
+    Iterable<ResultRow> results = runner.run(QueryPlus.wrap(GroupByQueryRunnerTestHelper.populateResourceId(query)))
+                                        .toList();
     TestHelper.assertExpectedObjects(expectedResults, results, "approx-histo");
   }
 
@@ -231,7 +232,8 @@ public void testGroupByWithSameNameComplexPostAgg()
         )
     );
 
-    Iterable<ResultRow> results = runner.run(QueryPlus.wrap(query)).toList();
+    Iterable<ResultRow> results = runner.run(QueryPlus.wrap(GroupByQueryRunnerTestHelper.populateResourceId(query)))
+                                        .toList();
     TestHelper.assertExpectedObjects(expectedResults, results, "approx-histo");
   }
 }

extensions-core/histogram/src/test/java/org/apache/druid/query/aggregation/histogram/FixedBucketsHistogramGroupByQueryTest.java

@@ -186,7 +186,8 @@ public void testGroupByWithFixedHistogramAgg()
         )
     );
 
-    Iterable<ResultRow> results = runner.run(QueryPlus.wrap(query)).toList();
+    Iterable<ResultRow> results = runner.run(QueryPlus.wrap(GroupByQueryRunnerTestHelper.populateResourceId(query)))
+                                        .toList();
     TestHelper.assertExpectedObjects(expectedResults, results, "fixed-histo");
   }
 
@@ -233,7 +234,8 @@ public void testGroupByWithSameNameComplexPostAgg()
         )
     );
 
-    Iterable<ResultRow> results = runner.run(QueryPlus.wrap(query)).toList();
+    Iterable<ResultRow> results = runner.run(QueryPlus.wrap(GroupByQueryRunnerTestHelper.populateResourceId(query)))
+                                        .toList();
     TestHelper.assertExpectedObjects(expectedResults, results, "fixed-histo");
   }
 }

integration-tests/src/main/java/org/apache/druid/testing/utils/ITRetryUtil.java

@@ -69,7 +69,10 @@ public static void retryUntil(
 
       LOG.info(
           "Attempt[%d/%d] did not pass: Task %s still not complete. Next retry in %d ms",
-          currentTry, retryCount, taskMessage, delayInMillis
+          currentTry,
+          retryCount,
+          taskMessage,
+          delayInMillis
       );
       try {
         Thread.sleep(delayInMillis);
@@ -83,10 +86,10 @@ public static void retryUntil(
     if (currentTry > retryCount) {
       if (lastException != null) {
         throw new ISE(
+            lastException,
             "Max number of retries[%d] exceeded for Task[%s]. Failing.",
             retryCount,
-            taskMessage,
-            lastException
+            taskMessage
         );
       } else {
         throw new ISE(

processing/src/main/java/org/apache/druid/query/FluentQueryRunner.java

@@ -90,9 +90,14 @@ public FluentQueryRunner<T> postProcess(PostProcessingOperator<T> postProcessing
     return from(postProcessing != null ? postProcessing.postProcess(baseRunner) : baseRunner);
   }
 
-  public FluentQueryRunner<T> mergeResults()
+  /**
+   * Delegates to {@link QueryToolChest#mergeResults(QueryRunner, boolean)}.
+   * 
+   * @see QueryToolChest#mergeResults(QueryRunner, boolean)
+   */
+  public FluentQueryRunner<T> mergeResults(boolean willMergeRunner)
   {
-    return from(toolChest.mergeResults(baseRunner));
+    return from(toolChest.mergeResults(baseRunner, willMergeRunner));
   }
 
   public FluentQueryRunner<T> map(final Function<QueryRunner<T>, QueryRunner<T>> mapFn)

processing/src/main/java/org/apache/druid/query/QueryContext.java

@@ -591,6 +591,10 @@ public boolean isWindowingStrictValidation()
     );
   }
 
+  public QueryResourceId getQueryResourceId()
+  {
+    return new QueryResourceId(getString(QueryContexts.QUERY_RESOURCE_ID));
+  }
 
   public String getBrokerServiceName()
   {

processing/src/main/java/org/apache/druid/query/QueryContexts.java

@@ -86,7 +86,9 @@ public class QueryContexts
   public static final String UNCOVERED_INTERVALS_LIMIT_KEY = "uncoveredIntervalsLimit";
   public static final String MIN_TOP_N_THRESHOLD = "minTopNThreshold";
   public static final String WINDOWING_STRICT_VALIDATION = "windowingStrictValidation";
-
+  // Unique identifier for the query, that is used to map the global shared resources (specifically merge buffers) to the
+  // query's runtime
+  public static final String QUERY_RESOURCE_ID = "queryResourceId";
 
   // SQL query context keys
   public static final String CTX_SQL_QUERY_ID = BaseQuery.SQL_QUERY_ID;