Implementation of recycling memory pool.

clients/src/main/java/org/apache/kafka/common/memory/RecyclingMemoryPool.java

@@ -0,0 +1,142 @@
+/*
+ * 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.
+ */
+
+package org.apache.kafka.common.memory;
+
+import java.nio.ByteBuffer;
+import java.util.Iterator;
+import java.util.LinkedList;
+import java.util.List;
+import org.apache.kafka.common.metrics.Sensor;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+
+/**
+ * An implementation of memory pool which recycles buffers of commonly used size.
+ * This memory pool is useful if most of the requested buffers' size are within close size range.
+ * In this case, instead of deallocate and reallocate the buffer, the memory pool will recycle the buffer for future use.
+ */
+public class RecyclingMemoryPool implements MemoryPool {
+    protected final Logger log = LoggerFactory.getLogger(getClass());
+    protected final int cacheableBufferSize;
+    protected final int bufferCacheCapacity;
+    protected int bufferCacheSlot;
+    protected final List<ByteBuffer> bufferCache;
+    protected volatile Sensor requestSensor;
+
+    public RecyclingMemoryPool(int cacheableBufferSize, int bufferCacheCapacity, Sensor requestSensor) {
+        if (bufferCacheCapacity <= 0 || cacheableBufferSize <= 0) {
+            throw new IllegalArgumentException(String.format("Must provide a positive cacheable buffer size and buffer cache " +
+                    "capacity, provided %d and %d respectively.", cacheableBufferSize, bufferCacheCapacity));
+        }
+        this.bufferCache = new LinkedList<>();
+        this.cacheableBufferSize = cacheableBufferSize;
+        this.bufferCacheCapacity = bufferCacheCapacity;
+        this.requestSensor = requestSensor;
+        this.bufferCacheSlot = 0;
+    }
+
+    @Override
+    public ByteBuffer tryAllocate(int sizeBytes) {
+        if (sizeBytes < 1) {
+            throw new IllegalArgumentException("requested size " + sizeBytes + "<=0");
+        }
+
+        ByteBuffer allocated = null;
+        if (sizeBytes > cacheableBufferSize / 2 && sizeBytes <= cacheableBufferSize) {
+            allocated = maybePopBufferFromCache();
+        }
+        if (allocated == null) {
+            allocated = ByteBuffer.allocate(sizeBytes);
+        }
+        bufferToBeAllocated(allocated);
+        return allocated;
+    }
+
+    /**
+     * If (1) there are recycled buffers in cache or (2) more buffer of size {@link #cacheableBufferSize} can be allocated,
+     * (allocate the buffer and) return the buffer to the client.
+     *
+     * @return The available buffer if any.
+     */
+    protected ByteBuffer maybePopBufferFromCache() {
+        boolean canAllocateMoreBuffer;
+        synchronized (this) {
+            if (!bufferCache.isEmpty()) {
+                Iterator<ByteBuffer> it = bufferCache.iterator();
+                ByteBuffer buffer = it.next();
+                it.remove();
+                buffer.clear();
+                return buffer;
+            }
+            canAllocateMoreBuffer = bufferCacheSlot < bufferCacheCapacity;
+            if (canAllocateMoreBuffer) {
+                bufferCacheSlot++;
+            }
+        }
+        if (canAllocateMoreBuffer) {
+            ByteBuffer byteBuffer = ByteBuffer.allocate(cacheableBufferSize);
+            bufferToBeAllocated(byteBuffer);
+            return byteBuffer;
+        }
+        return null;
+    }
+
+    @Override
+    public void release(ByteBuffer previouslyAllocated) {
+        if (previouslyAllocated == null) {
+            throw new IllegalArgumentException("provided null buffer");
+        }
+        if (previouslyAllocated.capacity() == cacheableBufferSize) {
+            maybeRecycleBufferToCache(previouslyAllocated);
+        }
+        bufferToBeReleased(previouslyAllocated);
+    }
+
+    protected synchronized void maybeRecycleBufferToCache(ByteBuffer previouslyAllocated) {
+        if (bufferCache.size() < bufferCacheCapacity) {
+            bufferCache.add(previouslyAllocated);
+        }
+    }
+
+    //allows subclasses to do their own bookkeeping (and validation) _before_ memory is returned to client code.
+    protected void bufferToBeAllocated(ByteBuffer justAllocated) {
+        this.requestSensor.record(justAllocated.capacity());
+        log.trace("allocated buffer of size {} ", justAllocated.capacity());
+    }
+
+    //allows subclasses to do their own bookkeeping (and validation) _before_ memory is marked as reclaimed.
+    protected void bufferToBeReleased(ByteBuffer justReleased) {
+        log.trace("released buffer of size {}", justReleased.capacity());
+    }
+
+    @Override
+    public long size() {
+        return Long.MAX_VALUE;
+    }
+
+    @Override
+    public long availableMemory() {
+        return Long.MAX_VALUE;
+    }
+
+    @Override
+    public boolean isOutOfMemory() {
+        return false;
+    }
+}

clients/src/test/java/org/apache/kafka/common/memory/RecyclingMemoryPoolTest.java

@@ -0,0 +1,75 @@
+/*
+ * 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.
+ */
+
+package org.apache.kafka.common.memory;
+
+import java.nio.ByteBuffer;
+import org.apache.kafka.common.metrics.Metrics;
+import org.apache.kafka.common.metrics.Sensor;
+import org.junit.Assert;
+import org.junit.Test;
+
+
+public class RecyclingMemoryPoolTest {
+    private static final int TWO_KILOBYTES = 2048;
+    private static final int CACHEABLE_BUFFER_SIZE = 1024;
+    private static final int BUFFER_CACHE_CAPACITY = 2;
+    private static final Sensor ALLOCATE_SENSOR = new Metrics().sensor("allocate_sensor");
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testNegativeAllocation() {
+        RecyclingMemoryPool memoryPool = new RecyclingMemoryPool(CACHEABLE_BUFFER_SIZE, BUFFER_CACHE_CAPACITY, ALLOCATE_SENSOR);
+        memoryPool.tryAllocate(-1);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testZeroAllocation() {
+        RecyclingMemoryPool memoryPool = new RecyclingMemoryPool(CACHEABLE_BUFFER_SIZE, BUFFER_CACHE_CAPACITY, ALLOCATE_SENSOR);
+        memoryPool.tryAllocate(0);
+    }
+
+    @Test(expected = IllegalArgumentException.class)
+    public void testNullRelease() {
+        RecyclingMemoryPool memoryPool = new RecyclingMemoryPool(CACHEABLE_BUFFER_SIZE, BUFFER_CACHE_CAPACITY, ALLOCATE_SENSOR);
+        memoryPool.release(null);
+    }
+
+    @Test
+    public void testAllocation() {
+        RecyclingMemoryPool memoryPool = new RecyclingMemoryPool(CACHEABLE_BUFFER_SIZE, BUFFER_CACHE_CAPACITY, ALLOCATE_SENSOR);
+        ByteBuffer buffer1 = memoryPool.tryAllocate(TWO_KILOBYTES);
+        ByteBuffer buffer2 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE);
+        ByteBuffer buffer3 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE * 2 / 3);
+        ByteBuffer buffer4 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE);
+
+        memoryPool.release(buffer1);
+        ByteBuffer reuse1 = memoryPool.tryAllocate(TWO_KILOBYTES);
+        // Compare the references
+        Assert.assertNotEquals(System.identityHashCode(reuse1), System.identityHashCode(buffer1));
+
+        memoryPool.release(buffer2);
+        memoryPool.release(buffer3);
+        memoryPool.release(buffer4);
+        ByteBuffer reuse2 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE);
+        ByteBuffer reuse3 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE * 2 / 3);
+        ByteBuffer reuse4 = memoryPool.tryAllocate(CACHEABLE_BUFFER_SIZE);
+
+        Assert.assertEquals(System.identityHashCode(reuse2), System.identityHashCode(buffer2));
+        Assert.assertEquals(System.identityHashCode(reuse3), System.identityHashCode(buffer3));
+        Assert.assertNotEquals(System.identityHashCode(reuse4), System.identityHashCode(buffer4));
+    }
+}

core/src/main/scala/kafka/network/SocketServer.scala

@@ -37,7 +37,7 @@ import kafka.server.{BrokerReconfigurable, KafkaConfig}
 import kafka.utils._
 import org.apache.kafka.common.config.ConfigException
 import org.apache.kafka.common.{KafkaException, Reconfigurable}
-import org.apache.kafka.common.memory.{MemoryPool, SimpleMemoryPool}
+import org.apache.kafka.common.memory.{MemoryPool, RecyclingMemoryPool, SimpleMemoryPool}
 import org.apache.kafka.common.metrics._
 import org.apache.kafka.common.metrics.stats.Total
 import org.apache.kafka.common.metrics.stats.Percentiles.BucketSizing
@@ -96,6 +96,7 @@ class SocketServer(val config: KafkaConfig,
   // At current stage, we do not know the max decrypted request size, temporarily set it to 10MB.
   memoryPoolAllocationSensor.add(memoryPoolAllocateSizePercentilesMetricName, new Percentiles(400, 0.0, 10485760, BucketSizing.CONSTANT, percentiles:_*))
   private val memoryPool = if (config.queuedMaxBytes > 0) new SimpleMemoryPool(config.queuedMaxBytes, config.socketRequestMaxBytes, false, memoryPoolUsageSensor, memoryPoolAllocationSensor)
+                           else if (config.socketRequestCommonBytes > 0) new RecyclingMemoryPool(config.socketRequestCommonBytes, config.socketRequestBufferCacheSize, memoryPoolAllocationSensor)
                            else MemoryPool.NONE
   // data-plane
   private val dataPlaneProcessors = new ConcurrentHashMap[Int, Processor]()

core/src/main/scala/kafka/server/KafkaConfig.scala

@@ -81,6 +81,8 @@ object Defaults {
   val SocketSendBufferBytes: Int = 100 * 1024
   val SocketReceiveBufferBytes: Int = 100 * 1024
   val SocketRequestMaxBytes: Int = 100 * 1024 * 1024
+  val SocketRequestCommonBytes: Int = -1
+  val SocketRequestBufferCacheSize: Int = 0
   val RequestMaxLocalTimeMs = Long.MaxValue
   val MaxConnectionsPerIp: Int = Int.MaxValue
   val MaxConnectionsPerIpOverrides: String = ""
@@ -320,6 +322,8 @@ object KafkaConfig {
   val SocketReceiveBufferBytesProp = "socket.receive.buffer.bytes"
   val RequestMaxLocalTimeMsProp = "request.max.local.time.ms"
   val SocketRequestMaxBytesProp = "socket.request.max.bytes"
+  val SocketRequestCommonBytesProp = "socket.request.common.bytes"
+  val SocketRequestBufferCacheSizeProp = "socket.request.buffer.cache.size"
   val MaxConnectionsPerIpProp = "max.connections.per.ip"
   val MaxConnectionsPerIpOverridesProp = "max.connections.per.ip.overrides"
   val MaxConnectionsProp = "max.connections"
@@ -617,6 +621,8 @@ object KafkaConfig {
     "takes longer than this time the broker will kill itself as violating this timeout is a symptom of a more serious broker zombie state." +
     " It is useful to observe the RequestDequeuePollIntervalMs metric to find a suitable setting for this configuration."
   val SocketRequestMaxBytesDoc = "The maximum number of bytes in a socket request"
+  val SocketRequestCommonBytesDoc = "The common size in bytes of a socket request"
+  val SocketRequestBufferCacheSizeDoc = "The maximal number of cache slot recycling memory pool will keep"
   val MaxConnectionsPerIpDoc = "The maximum number of connections we allow from each ip address. This can be set to 0 if there are overrides " +
     s"configured using $MaxConnectionsPerIpOverridesProp property. New connections from the ip address are dropped if the limit is reached."
   val MaxConnectionsPerIpOverridesDoc = "A comma-separated list of per-ip or hostname overrides to the default maximum number of connections. " +
@@ -952,6 +958,8 @@ object KafkaConfig {
       .define(RequestMaxLocalTimeMsProp, LONG, Defaults.RequestMaxLocalTimeMs, atLeast(1), MEDIUM, RequestMaxLocalTimeMsDoc)
       .define(SocketReceiveBufferBytesProp, INT, Defaults.SocketReceiveBufferBytes, HIGH, SocketReceiveBufferBytesDoc)
       .define(SocketRequestMaxBytesProp, INT, Defaults.SocketRequestMaxBytes, atLeast(1), HIGH, SocketRequestMaxBytesDoc)
+      .define(SocketRequestCommonBytesProp, INT, Defaults.SocketRequestCommonBytes, MEDIUM, SocketRequestCommonBytesDoc)
+      .define(SocketRequestBufferCacheSizeProp, INT, Defaults.SocketRequestBufferCacheSize, atLeast(0), MEDIUM, SocketRequestBufferCacheSizeDoc)
       .define(MaxConnectionsPerIpProp, INT, Defaults.MaxConnectionsPerIp, atLeast(0), MEDIUM, MaxConnectionsPerIpDoc)
       .define(MaxConnectionsPerIpOverridesProp, STRING, Defaults.MaxConnectionsPerIpOverrides, MEDIUM, MaxConnectionsPerIpOverridesDoc)
       .define(MaxConnectionsProp, INT, Defaults.MaxConnections, atLeast(0), MEDIUM, MaxConnectionsDoc)
@@ -1259,6 +1267,8 @@ class KafkaConfig(val props: java.util.Map[_, _], doLog: Boolean, dynamicConfigO
   val socketSendBufferBytes = getInt(KafkaConfig.SocketSendBufferBytesProp)
   val socketReceiveBufferBytes = getInt(KafkaConfig.SocketReceiveBufferBytesProp)
   val socketRequestMaxBytes = getInt(KafkaConfig.SocketRequestMaxBytesProp)
+  val socketRequestCommonBytes = getInt(KafkaConfig.SocketRequestCommonBytesProp)
+  val socketRequestBufferCacheSize = getInt(KafkaConfig.SocketRequestBufferCacheSizeProp)
   val requestMaxLocalTimeMs = getLong(KafkaConfig.RequestMaxLocalTimeMsProp)
   val maxConnectionsPerIp = getInt(KafkaConfig.MaxConnectionsPerIpProp)
   val maxConnectionsPerIpOverrides: Map[String, Int] =