code: /*
   * Copyright 2018 Mitsunori Komatsu (komamitsu)
   *
   * Licensed 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.komamitsu.fluency.buffer;
  
  import org.msgpack.core.annotations.VisibleForTesting;
  
  import java.nio.ByteBuffer;
  import java.util.HashMap;
  import java.util.Map;
  import java.util.concurrent.LinkedBlockingQueue;
  import java.util.concurrent.atomic.AtomicLong;
  
  class BufferPool
  {
      @VisibleForTesting final Map<Integer, LinkedBlockingQueue<ByteBuffer>> bufferPool = new HashMap<>();
      private final AtomicLong allocatedSize = new AtomicLong();
      private final int initialBufferSize;
      private final long maxBufferSize;
      private final boolean jvmHeapBufferMode;
  
      public BufferPool(int initialBufferSize, long maxBufferSize)
      {
          this(initialBufferSize, maxBufferSize, false);
      }
  
      public BufferPool(int initialBufferSize, long maxBufferSize, boolean jvmHeapBufferMode)
      {
          this.initialBufferSize = initialBufferSize;
          this.maxBufferSize = maxBufferSize;
          this.jvmHeapBufferMode = jvmHeapBufferMode;
      }
  
      public ByteBuffer acquireBuffer(int bufferSize)
      {
          int normalizedBufferSize = initialBufferSize;
          while (normalizedBufferSize < bufferSize) {
              normalizedBufferSize *= 2;
          }
  
          LinkedBlockingQueue<ByteBuffer> buffers;
          synchronized (bufferPool) {
              buffers = bufferPool.get(normalizedBufferSize);
              if (buffers == null) {
                  buffers = new LinkedBlockingQueue<>();
                  bufferPool.put(normalizedBufferSize, buffers);
              }
          }
  
          ByteBuffer buffer = buffers.poll();
          if (buffer != null) {
              return buffer;
          }
  
          /*
          synchronized (allocatedSize) {
              if (allocatedSize.get() + normalizedBufferSize > maxBufferSize) {
                  return null;    // `null` means the buffer is full.
              }
              allocatedSize.addAndGet(normalizedBufferSize);
              return ByteBuffer.allocateDirect(normalizedBufferSize);
          }
          */
  
          while (true) {
              long currentAllocatedSize = allocatedSize.get();
              if (currentAllocatedSize + normalizedBufferSize > maxBufferSize) {
                  releaseBuffers();
                  return null;    // `null` means the buffer is full.
              }
              if (currentAllocatedSize == allocatedSize.getAndAdd(normalizedBufferSize)) {
                  ByteBuffer buf;
                  if (jvmHeapBufferMode) {
                      buf = ByteBuffer.allocate(normalizedBufferSize);
                  }
                  else {
                      buf = ByteBuffer.allocateDirect(normalizedBufferSize);
                  }
                  return buf;
              }
              allocatedSize.getAndAdd(-normalizedBufferSize);
          }
      }
  
      public void returnBuffer(ByteBuffer byteBuffer)
      {
          LinkedBlockingQueue<ByteBuffer> buffers = bufferPool.get(byteBuffer.capacity());
          if (buffers == null) {
              throw new IllegalStateException("`buffers` shouldn't be null");
          }
  
          byteBuffer.position(0);
          byteBuffer.limit(byteBuffer.capacity());
          buffers.offer(byteBuffer);
      }
  
      public long getAllocatedSize()
      {
          return allocatedSize.get();
      }
  
      public void releaseBuffers()
      {
          // TODO: Stop releasing all the buffers when having released buffers to some extent
          synchronized (bufferPool) {
              for (Map.Entry<Integer, LinkedBlockingQueue<ByteBuffer>> entry : bufferPool.entrySet()) {
                  ByteBuffer buffer;
                  while ((buffer = entry.getValue().poll()) != null) {
                      allocatedSize.addAndGet(-buffer.capacity());
                  }
              }
          }
      }
  
      public boolean getJvmHeapBufferMode()
      {
          return jvmHeapBufferMode;
      }
  
      @Override
      public String toString()
      {
          return "BufferPool{" +
                  "bufferPool=" + bufferPool +
                  ", allocatedSize=" + allocatedSize +
                  ", initialBufferSize=" + initialBufferSize +
                  ", maxBufferSize=" + maxBufferSize +
                  ", jvmHeapBufferMode=" + jvmHeapBufferMode +
                  '}';
      }
  }
  