From 32c5494e974877806b34b8e44a876610a70fe756 Mon Sep 17 00:00:00 2001
From: Himanshu <g.himanshu@gmail.com>
Date: Mon, 31 Oct 2016 17:24:07 -0500
Subject: [PATCH] eagerly allocate the intermediate computation buffers (#3628)

---
 .../druid/benchmark/query/GroupByBenchmark.java  |  1 +
 .../io/druid/benchmark/query/TopNBenchmark.java  |  2 +-
 .../java/io/druid/collections/StupidPool.java    | 16 +++++++++++++---
 .../io/druid/guice/DruidProcessingModule.java    |  1 +
 4 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/benchmarks/src/main/java/io/druid/benchmark/query/GroupByBenchmark.java b/benchmarks/src/main/java/io/druid/benchmark/query/GroupByBenchmark.java
index 13221006416d..a520ff338d17 100644
--- a/benchmarks/src/main/java/io/druid/benchmark/query/GroupByBenchmark.java
+++ b/benchmarks/src/main/java/io/druid/benchmark/query/GroupByBenchmark.java
@@ -309,6 +309,7 @@ public void setup() throws IOException
 
     StupidPool<ByteBuffer> bufferPool = new StupidPool<>(
         new OffheapBufferGenerator("compute", 250_000_000),
+        0,
         Integer.MAX_VALUE
     );
 
diff --git a/benchmarks/src/main/java/io/druid/benchmark/query/TopNBenchmark.java b/benchmarks/src/main/java/io/druid/benchmark/query/TopNBenchmark.java
index 12d1ac9cf8ca..a09b899eece9 100644
--- a/benchmarks/src/main/java/io/druid/benchmark/query/TopNBenchmark.java
+++ b/benchmarks/src/main/java/io/druid/benchmark/query/TopNBenchmark.java
@@ -272,7 +272,7 @@ public void setup() throws IOException
     }
 
     factory = new TopNQueryRunnerFactory(
-        new StupidPool<>(new OffheapBufferGenerator("compute", 250000000), Integer.MAX_VALUE),
+        new StupidPool<>(new OffheapBufferGenerator("compute", 250000000), 0, Integer.MAX_VALUE),
         new TopNQueryQueryToolChest(new TopNQueryConfig(), QueryBenchmarkUtil.NoopIntervalChunkingQueryRunnerDecorator()),
         QueryBenchmarkUtil.NOOP_QUERYWATCHER
     );
diff --git a/common/src/main/java/io/druid/collections/StupidPool.java b/common/src/main/java/io/druid/collections/StupidPool.java
index 406de3143337..f6cb2cfd8781 100644
--- a/common/src/main/java/io/druid/collections/StupidPool.java
+++ b/common/src/main/java/io/druid/collections/StupidPool.java
@@ -19,8 +19,8 @@
 
 package io.druid.collections;
 
+import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
-
 import io.druid.java.util.common.ISE;
 import io.druid.java.util.common.logger.Logger;
 
@@ -45,17 +45,27 @@ public StupidPool(
       Supplier<T> generator
   )
   {
-    this.generator = generator;
-    this.objectsCacheMaxCount = Integer.MAX_VALUE;
+    this(generator, 0, Integer.MAX_VALUE);
   }
 
   public StupidPool(
       Supplier<T> generator,
+      int initCount,
       int objectsCacheMaxCount
   )
   {
+    Preconditions.checkArgument(
+        initCount <= objectsCacheMaxCount,
+        "initCount[%s] must be less/equal to objectsCacheMaxCount[%s]",
+        initCount,
+        objectsCacheMaxCount
+    );
     this.generator = generator;
     this.objectsCacheMaxCount = objectsCacheMaxCount;
+
+    for (int i = 0; i < initCount; i++) {
+      objects.add(generator.get());
+    }
   }
 
   public ResourceHolder<T> take()
diff --git a/server/src/main/java/io/druid/guice/DruidProcessingModule.java b/server/src/main/java/io/druid/guice/DruidProcessingModule.java
index 6c48bc4fea21..dca3916d967b 100644
--- a/server/src/main/java/io/druid/guice/DruidProcessingModule.java
+++ b/server/src/main/java/io/druid/guice/DruidProcessingModule.java
@@ -110,6 +110,7 @@ public StupidPool<ByteBuffer> getIntermediateResultsPool(DruidProcessingConfig c
     verifyDirectMemory(config);
     return new StupidPool<>(
         new OffheapBufferGenerator("intermediate processing", config.intermediateComputeSizeBytes()),
+        config.getNumThreads(),
         config.poolCacheMaxCount()
     );
   }