content: /*
   * 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 io.netty.buffer;
  
  import io.netty.util.internal.OutOfDirectMemoryError;
  import io.netty.util.internal.StringUtil;
  import java.lang.reflect.Field;
  import java.nio.ByteBuffer;
  import java.util.concurrent.atomic.AtomicLong;
  import org.apache.arrow.memory.OutOfMemoryException;
  import org.apache.arrow.memory.util.AssertionUtil;
  import org.apache.arrow.memory.util.LargeMemoryUtil;
  
  /**
   * The base allocator that we use for all of Arrow's memory management. Returns
   * UnsafeDirectLittleEndian buffers.
   */
  public class PooledByteBufAllocatorL {
  
    private static final org.slf4j.Logger memoryLogger =
        org.slf4j.LoggerFactory.getLogger("arrow.allocator");
  
    private static final int MEMORY_LOGGER_FREQUENCY_SECONDS = 60;
    public final UnsafeDirectLittleEndian empty;
    private final AtomicLong hugeBufferSize = new AtomicLong(0);
    private final AtomicLong hugeBufferCount = new AtomicLong(0);
    private final AtomicLong normalBufferSize = new AtomicLong(0);
    private final AtomicLong normalBufferCount = new AtomicLong(0);
    private final InnerAllocator allocator;
  
    public PooledByteBufAllocatorL() {
      allocator = new InnerAllocator();
      empty = new UnsafeDirectLittleEndian(new DuplicatedByteBuf(Unpooled.EMPTY_BUFFER));
    }
  
    /** Returns a {@linkplain UnsafeDirectLittleEndian} of the given size. */
    public UnsafeDirectLittleEndian allocate(long size) {
      try {
        return allocator.directBuffer(LargeMemoryUtil.checkedCastToInt(size), Integer.MAX_VALUE);
      } catch (OutOfMemoryError e) {
        /*
         * OutOfDirectMemoryError is thrown by Netty when we exceed the direct memory limit defined by
         * -XX:MaxDirectMemorySize. OutOfMemoryError with "Direct buffer memory" message is thrown by
         * java.nio.Bits when we exceed the direct memory limit. This should never be hit in practice
         * as Netty is expected to throw an OutOfDirectMemoryError first.
         */
        if (e instanceof OutOfDirectMemoryError || "Direct buffer memory".equals(e.getMessage())) {
          throw new OutOfMemoryException("Failure allocating buffer.", e);
        }
        throw e;
      }
    }
  
    public int getChunkSize() {
      return allocator.chunkSize();
    }
  
    public long getHugeBufferSize() {
      return hugeBufferSize.get();
    }
  
    public long getHugeBufferCount() {
      return hugeBufferCount.get();
    }
  
    public long getNormalBufferSize() {
      return normalBufferSize.get();
    }
  
    public long getNormalBufferCount() {
      return normalBufferSize.get();
    }
  
    private static class AccountedUnsafeDirectLittleEndian extends UnsafeDirectLittleEndian {
  
      private final long initialCapacity;
      private final AtomicLong count;
      private final AtomicLong size;
  
      private AccountedUnsafeDirectLittleEndian(LargeBuffer buf, AtomicLong count, AtomicLong size) {
        super(buf);
        this.initialCapacity = buf.capacity();
        this.count = count;
        this.size = size;
      }
  
      private AccountedUnsafeDirectLittleEndian(
          PooledUnsafeDirectByteBuf buf, AtomicLong count, AtomicLong size) {
        super(buf);
        this.initialCapacity = buf.capacity();
        this.count = count;
        this.size = size;
      }
  
      @Override
      public ByteBuf copy() {
        throw new UnsupportedOperationException("copy method is not supported");
      }
  
      @Override
      public ByteBuf copy(int index, int length) {
        throw new UnsupportedOperationException("copy method is not supported");
      }
  
      @Override
      public boolean release(int decrement) {
        boolean released = super.release(decrement);
        if (released) {
          count.decrementAndGet();
          size.addAndGet(-initialCapacity);
        }
        return released;
      }
    }
  
    private class InnerAllocator extends PooledByteBufAllocator {
  
      private final PoolArena<ByteBuffer>[] directArenas;
      private final MemoryStatusThread statusThread;
  
      public InnerAllocator() {
        super(true);
  
        try {
          Field f = PooledByteBufAllocator.class.getDeclaredField("directArenas");
          f.setAccessible(true);
          this.directArenas = (PoolArena<ByteBuffer>[]) f.get(this);
        } catch (Exception e) {
          throw new RuntimeException(
              "Failure while initializing allocator.  Unable to retrieve direct arenas field.", e);
        }
  
        if (memoryLogger.isTraceEnabled()) {
          statusThread = new MemoryStatusThread(this);
          statusThread.start();
        } else {
          statusThread = null;
        }
      }
  
      private UnsafeDirectLittleEndian newDirectBufferL(int initialCapacity, int maxCapacity) {
        PoolThreadCache cache = threadCache();
        PoolArena<ByteBuffer> directArena = cache.directArena;
  
        if (directArena != null) {
  
          if (initialCapacity > chunkSize()) {
            // This is beyond chunk size so we'll allocate separately.
            ByteBuf buf = UnpooledByteBufAllocator.DEFAULT.directBuffer(initialCapacity, maxCapacity);
  
            hugeBufferSize.addAndGet(buf.capacity());
            hugeBufferCount.incrementAndGet();
  
            // logger.debug("Allocating huge buffer of size {}", initialCapacity, new Exception());
            return new AccountedUnsafeDirectLittleEndian(
                new LargeBuffer(buf), hugeBufferCount, hugeBufferSize);
          } else {
            // within chunk, use arena.
            ByteBuf buf = directArena.allocate(cache, initialCapacity, maxCapacity);
            if (!(buf instanceof PooledUnsafeDirectByteBuf)) {
              fail();
            }
  
            if (!AssertionUtil.ASSERT_ENABLED) {
              return new UnsafeDirectLittleEndian((PooledUnsafeDirectByteBuf) buf);
            }
  
            normalBufferSize.addAndGet(buf.capacity());
            normalBufferCount.incrementAndGet();
  
            return new AccountedUnsafeDirectLittleEndian(
                (PooledUnsafeDirectByteBuf) buf, normalBufferCount, normalBufferSize);
          }
  
        } else {
          throw fail();
        }
      }
  
      private UnsupportedOperationException fail() {
        return new UnsupportedOperationException(
            "Arrow requires that the JVM used supports access sun.misc.Unsafe.  This platform "
                + "didn't provide that functionality.");
      }
  
      @Override
      public UnsafeDirectLittleEndian directBuffer(int initialCapacity, int maxCapacity) {
        if (initialCapacity == 0 && maxCapacity == 0) {
          newDirectBuffer(initialCapacity, maxCapacity);
        }
        validate(initialCapacity, maxCapacity);
        return newDirectBufferL(initialCapacity, maxCapacity);
      }
  
      @Override
      public ByteBuf heapBuffer(int initialCapacity, int maxCapacity) {
        throw new UnsupportedOperationException("Arrow doesn't support using heap buffers.");
      }
  
      private void validate(int initialCapacity, int maxCapacity) {
        if (initialCapacity < 0) {
          throw new IllegalArgumentException(
              "initialCapacity: " + initialCapacity + " (expected: 0+)");
        }
        if (initialCapacity > maxCapacity) {
          throw new IllegalArgumentException(
              String.format(
                  "initialCapacity: %d (expected: not greater than maxCapacity(%d)",
                  initialCapacity, maxCapacity));
        }
      }
  
      @Override
      public String toString() {
        StringBuilder buf = new StringBuilder();
        buf.append(directArenas.length);
        buf.append(" direct arena(s):");
        buf.append(StringUtil.NEWLINE);
        for (PoolArena<ByteBuffer> a : directArenas) {
          buf.append(a);
        }
  
        buf.append("Large buffers outstanding: ");
        buf.append(hugeBufferCount.get());
        buf.append(" totaling ");
        buf.append(hugeBufferSize.get());
        buf.append(" bytes.");
        buf.append('\n');
        buf.append("Normal buffers outstanding: ");
        buf.append(normalBufferCount.get());
        buf.append(" totaling ");
        buf.append(normalBufferSize.get());
        buf.append(" bytes.");
        return buf.toString();
      }
  
      private class MemoryStatusThread extends Thread {
        private final InnerAllocator allocator;
  
        public MemoryStatusThread(InnerAllocator allocator) {
          super("allocation.logger");
          this.setDaemon(true);
          this.allocator = allocator;
        }
  
        @Override
        public void run() {
          while (true) {
            memoryLogger.trace("Memory Usage: \n{}", allocator);
            try {
              Thread.sleep(MEMORY_LOGGER_FREQUENCY_SECONDS * 1000);
            } catch (InterruptedException e) {
              return;
            }
          }
        }
      }
    }
  }
  