Fix deadlock that can occur while merging group by results

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

@@ -376,15 +376,14 @@ public String getFormatString()
         druidProcessingConfig,
         configSupplier,
         bufferPool,
-        mergePool,
         TestHelper.makeJsonMapper(),
         new ObjectMapper(new SmileFactory()),
         QueryBenchmarkUtil.NOOP_QUERYWATCHER
     );
 
     factory = new GroupByQueryRunnerFactory(
         groupingEngine,
-        new GroupByQueryQueryToolChest(groupingEngine)
+        new GroupByQueryQueryToolChest(groupingEngine, mergePool)
     );
   }
 

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

@@ -360,12 +360,11 @@ private static GroupByQueryRunnerFactory makeGroupByQueryRunnerFactory(
         processingConfig,
         configSupplier,
         bufferPool,
-        mergeBufferPool,
         mapper,
         mapper,
         QueryRunnerTestHelper.NOOP_QUERYWATCHER
     );
-    final GroupByQueryQueryToolChest toolChest = new GroupByQueryQueryToolChest(groupingEngine);
+    final GroupByQueryQueryToolChest toolChest = new GroupByQueryQueryToolChest(groupingEngine, mergeBufferPool);
     return new GroupByQueryRunnerFactory(groupingEngine, toolChest);
   }
 

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

@@ -491,15 +491,14 @@ public String getFormatString()
         druidProcessingConfig,
         configSupplier,
         bufferPool,
-        mergePool,
         TestHelper.makeJsonMapper(),
         new ObjectMapper(new SmileFactory()),
         QueryBenchmarkUtil.NOOP_QUERYWATCHER
     );
 
     factory = new GroupByQueryRunnerFactory(
         groupingEngine,
-        new GroupByQueryQueryToolChest(groupingEngine)
+        new GroupByQueryQueryToolChest(groupingEngine, mergePool)
     );
   }
 

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

@@ -96,15 +96,14 @@ public int getNumThreads()
         },
         GroupByQueryConfig::new,
         new StupidPool<>("map-virtual-column-groupby-test", () -> ByteBuffer.allocate(1024)),
-        new DefaultBlockingPool<>(() -> ByteBuffer.allocate(1024), 1),
         TestHelper.makeJsonMapper(),
         new DefaultObjectMapper(),
         QueryRunnerTestHelper.NOOP_QUERYWATCHER
     );
 
     final GroupByQueryRunnerFactory factory = new GroupByQueryRunnerFactory(
         groupingEngine,
-        new GroupByQueryQueryToolChest(groupingEngine)
+        new GroupByQueryQueryToolChest(groupingEngine, new DefaultBlockingPool<>(() -> ByteBuffer.allocate(1024), 1))
     );
 
     runner = QueryRunnerTestHelper.makeQueryRunner(

extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/exec/LoadedSegmentDataProvider.java

@@ -144,7 +144,7 @@ public <RowType, QueryType> Pair<DataServerQueryStatus, Yielder<RowType>> fetchR
     final int numRetriesOnMissingSegments = preparedQuery.context().getNumRetriesOnMissingSegments(DEFAULT_NUM_TRIES);
 
     log.debug("Querying severs[%s] for segment[%s], retries:[%d]", servers, segmentDescriptor, numRetriesOnMissingSegments);
-    final ResponseContext responseContext = new DefaultResponseContext();
+    final ResponseContext responseContext = DefaultResponseContext.createEmpty();
 
     Pair<DataServerQueryStatus, Yielder<RowType>> statusSequencePair;
     try {

processing/src/main/java/org/apache/druid/query/context/ConcurrentResponseContext.java

@@ -21,6 +21,7 @@
 
 import org.apache.druid.guice.annotations.PublicApi;
 
+import java.util.Collections;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
@@ -30,16 +31,34 @@
 @PublicApi
 public class ConcurrentResponseContext extends ResponseContext
 {
+
+  private final ConcurrentHashMap<Key, Object> delegate;
+
+  private ConcurrentResponseContext()
+  {
+    this(Collections.emptyMap());
+  }
+
+  private ConcurrentResponseContext(final Map<Key, Object> delegate)
+  {
+    this.delegate = new ConcurrentHashMap<>(delegate);
+  }
+
   public static ConcurrentResponseContext createEmpty()
   {
     return new ConcurrentResponseContext();
   }
 
-  private final ConcurrentHashMap<Key, Object> delegate = new ConcurrentHashMap<>();
 
   @Override
   protected Map<Key, Object> getDelegate()
   {
     return delegate;
   }
+
+  @Override
+  public ResponseContext clone()
+  {
+    return new ConcurrentResponseContext(delegate);
+  }
 }

processing/src/main/java/org/apache/druid/query/context/DefaultResponseContext.java

@@ -21,6 +21,7 @@
 
 import org.apache.druid.guice.annotations.PublicApi;
 
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -30,16 +31,33 @@
 @PublicApi
 public class DefaultResponseContext extends ResponseContext
 {
+
+  private final HashMap<Key, Object> delegate;
+
+  private DefaultResponseContext()
+  {
+    this(Collections.emptyMap());
+  }
+
+  private DefaultResponseContext(final Map<Key, Object> delegate)
+  {
+    this.delegate = new HashMap<>(delegate);
+  }
+
   public static DefaultResponseContext createEmpty()
   {
     return new DefaultResponseContext();
   }
 
-  private final HashMap<Key, Object> delegate = new HashMap<>();
-
   @Override
   protected Map<Key, Object> getDelegate()
   {
     return delegate;
   }
+
+  @Override
+  public ResponseContext clone()
+  {
+    return new DefaultResponseContext(delegate);
+  }
 }

processing/src/main/java/org/apache/druid/query/context/ResponseContext.java

@@ -143,7 +143,7 @@ public abstract static class AbstractKey implements Key
     private final boolean canDrop;
     private final Function<JsonParser, Object> parseFunction;
 
-    AbstractKey(String name, boolean inHeader, boolean canDrop, Class<?> serializedClass)
+    public AbstractKey(String name, boolean inHeader, boolean canDrop, Class<?> serializedClass)
     {
       this.name = name;
       this.inHeader = inHeader;
@@ -158,7 +158,7 @@ public abstract static class AbstractKey implements Key
       };
     }
 
-    AbstractKey(String name, boolean inHeader, boolean canDrop, TypeReference<?> serializedTypeReference)
+    public AbstractKey(String name, boolean inHeader, boolean canDrop, TypeReference<?> serializedTypeReference)
     {
       this.name = name;
       this.inHeader = inHeader;
@@ -557,6 +557,9 @@ public Key find(String name)
 
   protected abstract Map<Key, Object> getDelegate();
 
+  @Override
+  public abstract ResponseContext clone();
+
   public Map<String, Object> toMap()
   {
     return CollectionUtils.mapKeys(getDelegate(), k -> k.getName());

processing/src/main/java/org/apache/druid/query/groupby/GroupByQueryQueryToolChest.java

@@ -34,6 +34,7 @@
 import com.google.common.base.Supplier;
 import com.google.common.collect.Lists;
 import com.google.inject.Inject;
+import org.apache.druid.collections.BlockingPool;
 import org.apache.druid.data.input.Row;
 import org.apache.druid.frame.Frame;
 import org.apache.druid.frame.FrameType;
@@ -42,12 +43,14 @@
 import org.apache.druid.frame.write.FrameWriterFactory;
 import org.apache.druid.frame.write.FrameWriterUtils;
 import org.apache.druid.frame.write.FrameWriters;
+import org.apache.druid.guice.annotations.Merging;
 import org.apache.druid.java.util.common.ISE;
 import org.apache.druid.java.util.common.Pair;
 import org.apache.druid.java.util.common.granularity.Granularity;
 import org.apache.druid.java.util.common.guava.MappedSequence;
 import org.apache.druid.java.util.common.guava.Sequence;
 import org.apache.druid.java.util.common.guava.Sequences;
+import org.apache.druid.java.util.common.io.Closer;
 import org.apache.druid.java.util.common.jackson.JacksonUtils;
 import org.apache.druid.query.CacheStrategy;
 import org.apache.druid.query.DataSource;
@@ -75,6 +78,7 @@
 
 import java.io.Closeable;
 import java.io.IOException;
+import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.BitSet;
 import java.util.Comparator;
@@ -102,23 +106,26 @@ public class GroupByQueryQueryToolChest extends QueryToolChest<ResultRow, GroupB
   private final GroupingEngine groupingEngine;
   private final GroupByQueryConfig queryConfig;
   private final GroupByQueryMetricsFactory queryMetricsFactory;
+  private final BlockingPool<ByteBuffer> mergeBufferPool;
 
   @VisibleForTesting
-  public GroupByQueryQueryToolChest(GroupingEngine groupingEngine)
+  public GroupByQueryQueryToolChest(GroupingEngine groupingEngine, BlockingPool<ByteBuffer> mergeBufferPool)
   {
-    this(groupingEngine, GroupByQueryConfig::new, DefaultGroupByQueryMetricsFactory.instance());
+    this(groupingEngine, GroupByQueryConfig::new, DefaultGroupByQueryMetricsFactory.instance(), mergeBufferPool);
   }
 
   @Inject
   public GroupByQueryQueryToolChest(
       GroupingEngine groupingEngine,
       Supplier<GroupByQueryConfig> queryConfigSupplier,
-      GroupByQueryMetricsFactory queryMetricsFactory
+      GroupByQueryMetricsFactory queryMetricsFactory,
+      @Merging BlockingPool<ByteBuffer> mergeBufferPool
   )
   {
     this.groupingEngine = groupingEngine;
     this.queryConfig = queryConfigSupplier.get();
     this.queryMetricsFactory = queryMetricsFactory;
+    this.mergeBufferPool = mergeBufferPool;
   }
 
   @Override
@@ -152,16 +159,34 @@ private Sequence<ResultRow> initAndMergeGroupByResults(
       ResponseContext context
   )
   {
-    final GroupByQueryResources resource = groupingEngine.prepareResource(query);
+    // .. 1. Historicals, Broker -> Which is using localWalker
+    // MerginV2 ->
+    ResponseContext clonedContext = context.clone();
+    Boolean usesGroupByMergingQueryRunner = (Boolean) query
+        .getContext()
+        .getOrDefault(GroupByUtils.CTX_KEY_RUNNER_MERGES_USING_GROUP_BY_MERGING_QUERY_RUNNER_V2, true);
+    final GroupByQueryResources resource = GroupingEngine.prepareResource(
+        query,
+        mergeBufferPool,
+        usesGroupByMergingQueryRunner,
+        queryConfig
+    );
+    if (usesGroupByMergingQueryRunner) {
+      clonedContext.add(GroupByUtils.RESPONSE_KEY_GROUP_BY_MERGING_QUERY_RUNNER_BUFFERS, resource);
+    }
     try {
       final Sequence<ResultRow> mergedSequence = mergeGroupByResults(
           query,
           resource,
           runner,
-          context
+          clonedContext
       );
-
-      return Sequences.withBaggage(mergedSequence, resource);
+      Closer closer = Closer.create();
+      closer.register(resource);
+      closer.register(() ->
+                          clonedContext.remove(GroupByUtils.RESPONSE_KEY_GROUP_BY_MERGING_QUERY_RUNNER_BUFFERS)
+      );
+      return Sequences.withBaggage(mergedSequence, closer);
     }
     catch (Exception e) {
       // Error creating the Sequence; release resources.