Add support parallel combine in brokers

benchmarks/src/main/resources/log4j2.xml

@@ -0,0 +1,32 @@
+<?xml version="1.0" encoding="UTF-8" ?>
+<!--
+  ~ Licensed to the Apache Software Foundation (ASF) under one
+  ~ or more contributor license agreements.  See the NOTICE file
+  ~ distributed with this work for additional information
+  ~ regarding copyright ownership.  The ASF licenses this file
+  ~ to you under the Apache License, Version 2.0 (the
+  ~ "License"); you may not use this file except in compliance
+  ~ with the License.  You may obtain a copy of the License at
+  ~
+  ~   http://www.apache.org/licenses/LICENSE-2.0
+  ~
+  ~ Unless required by applicable law or agreed to in writing,
+  ~ software distributed under the License is distributed on an
+  ~ "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+  ~ KIND, either express or implied.  See the License for the
+  ~ specific language governing permissions and limitations
+  ~ under the License.
+  -->
+
+<Configuration status="WARN">
+  <Appenders>
+    <Console name="Console" target="SYSTEM_OUT">
+      <PatternLayout pattern="%d{ISO8601} %p [%t] %c - %m%n"/>
+    </Console>
+  </Appenders>
+  <Loggers>
+    <Root level="info">
+      <AppenderRef ref="Console"/>
+    </Root>
+  </Loggers>
+</Configuration>

core/src/main/java/org/apache/druid/collections/BlockingPool.java

@@ -24,8 +24,16 @@
 
 public interface BlockingPool<T>
 {
+  /**
+   * Returns the total pool size.
+   */
   int maxSize();
 
+  /**
+   * Poll all available resources from the pool. If there's no available resource, it returns an empty list.
+   */
+  List<ReferenceCountingResourceHolder<T>> pollAll();
+
   /**
    * Take a resource from the pool, waiting up to the
    * specified wait time if necessary for an element to become available.
@@ -51,16 +59,45 @@ public interface BlockingPool<T>
    * @param elementNum number of resources to take
    * @param timeoutMs  maximum time to wait for resources, in milliseconds.
    *
-   * @return a list of resource holders. An empty list is returned if {@code elementNum} resources aren't available.
+   * @return {@link TakeBatchResult} containing a list of resource holders if it succeeds. An empty list should be
+   *         returned if there aren't enough available resources. The result also contains the number of remaining
+   *         resources after this call no matter it succeeded or not.
    */
-  List<ReferenceCountingResourceHolder<T>> takeBatch(int elementNum, long timeoutMs);
+  TakeBatchResult<T> takeBatch(int elementNum, long timeoutMs);
 
   /**
    * Take resources from the pool, waiting if necessary until the elements of the given number become available.
    *
    * @param elementNum number of resources to take
    *
-   * @return a list of resource holders. An empty list is returned if {@code elementNum} resources aren't available.
+   * @return {@link TakeBatchResult} containing a list of resource holders and the number of remaining resources.
    */
-  List<ReferenceCountingResourceHolder<T>> takeBatch(int elementNum);
+  TakeBatchResult<T> takeBatch(int elementNum);
+
+  class TakeBatchResult<T>
+  {
+    private final List<ReferenceCountingResourceHolder<T>> elements;
+    private final int numAvailableElements;
+
+    public TakeBatchResult(List<ReferenceCountingResourceHolder<T>> elements, int numAvailableElements)
+    {
+      this.elements = elements;
+      this.numAvailableElements = numAvailableElements;
+    }
+
+    public boolean isOk()
+    {
+      return elements.size() > 0;
+    }
+
+    public List<ReferenceCountingResourceHolder<T>> getElements()
+    {
+      return elements;
+    }
+
+    public int getNumAvailableElements()
+    {
+      return numAvailableElements;
+    }
+  }
 }

core/src/main/java/org/apache/druid/collections/DefaultBlockingPool.java

@@ -71,11 +71,28 @@ public int maxSize()
   }
 
   @VisibleForTesting
-  public int getPoolSize()
+  public int available()
   {
     return objects.size();
   }
 
+  @Override
+  public List<ReferenceCountingResourceHolder<T>> pollAll()
+  {
+    checkInitialized();
+    final ReentrantLock lock = this.lock;
+    lock.lock();
+    try {
+      final List<T> list = new ArrayList<>(objects.size());
+      list.addAll(objects);
+      objects.clear();
+      return list.stream().map(this::wrapObject).collect(Collectors.toList());
+    }
+    finally {
+      lock.unlock();
+    }
+  }
+
   @Override
   @Nullable
   public ReferenceCountingResourceHolder<T> take(final long timeoutMs)
@@ -160,25 +177,30 @@ private T takeObject() throws InterruptedException
   }
 
   @Override
-  public List<ReferenceCountingResourceHolder<T>> takeBatch(final int elementNum, final long timeoutMs)
+  public TakeBatchResult<T> takeBatch(final int elementNum, final long timeoutMs)
   {
+    Preconditions.checkArgument(elementNum > 0, "elementNum should be positive");
     Preconditions.checkArgument(timeoutMs >= 0, "timeoutMs must be a non-negative value, but was [%s]", timeoutMs);
     checkInitialized();
     try {
       final List<T> objects = timeoutMs > 0 ? pollObjects(elementNum, timeoutMs) : pollObjects(elementNum);
-      return objects.stream().map(this::wrapObject).collect(Collectors.toList());
+      return new TakeBatchResult<>(objects.stream().map(this::wrapObject).collect(Collectors.toList()), objects.size());
     }
     catch (InterruptedException e) {
       throw new RuntimeException(e);
     }
   }
 
   @Override
-  public List<ReferenceCountingResourceHolder<T>> takeBatch(final int elementNum)
+  public TakeBatchResult<T> takeBatch(final int elementNum)
   {
+    Preconditions.checkArgument(elementNum > 0, "elementNum should be positive");
     checkInitialized();
     try {
-      return takeObjects(elementNum).stream().map(this::wrapObject).collect(Collectors.toList());
+      return new TakeBatchResult<>(
+          takeObjects(elementNum).stream().map(this::wrapObject).collect(Collectors.toList()),
+          objects.size()
+      );
     }
     catch (InterruptedException e) {
       throw new RuntimeException(e);

core/src/main/java/org/apache/druid/collections/DummyBlockingPool.java

@@ -44,6 +44,12 @@ public int maxSize()
     return 0;
   }
 
+  @Override
+  public List<ReferenceCountingResourceHolder<T>> pollAll()
+  {
+    throw new UnsupportedOperationException();
+  }
+
   @Override
   public ReferenceCountingResourceHolder<T> take(long timeoutMs)
   {
@@ -57,13 +63,13 @@ public ReferenceCountingResourceHolder<T> take()
   }
 
   @Override
-  public List<ReferenceCountingResourceHolder<T>> takeBatch(int elementNum, long timeoutMs)
+  public TakeBatchResult<T> takeBatch(int elementNum, long timeoutMs)
   {
     throw new UnsupportedOperationException();
   }
 
   @Override
-  public List<ReferenceCountingResourceHolder<T>> takeBatch(int elementNum)
+  public TakeBatchResult<T> takeBatch(int elementNum)
   {
     throw new UnsupportedOperationException();
   }