Add AutoBufferLedger (apache#31)

java/dataset/src/main/java/org/apache/arrow/dataset/jni/DirectReservationListener.java

@@ -21,6 +21,7 @@
 import java.lang.reflect.Method;
 import java.util.concurrent.atomic.AtomicLong;
 
+import org.apache.arrow.memory.util.MemoryUtil;
 import org.apache.arrow.util.VisibleForTesting;
 
 /**
@@ -29,20 +30,6 @@
  * "-XX:MaxDirectMemorySize".
  */
 public class DirectReservationListener implements ReservationListener {
-  private final Method methodReserve;
-  private final Method methodUnreserve;
-
-  private DirectReservationListener() {
-    try {
-      final Class<?> classBits = Class.forName("java.nio.Bits");
-      methodReserve = classBits.getDeclaredMethod("reserveMemory", long.class, int.class);
-      methodReserve.setAccessible(true);
-      methodUnreserve = classBits.getDeclaredMethod("unreserveMemory", long.class, int.class);
-      methodUnreserve.setAccessible(true);
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
-  }
 
   private static final DirectReservationListener INSTANCE = new DirectReservationListener();
 
@@ -55,43 +42,22 @@ public static DirectReservationListener instance() {
    */
   @Override
   public void reserve(long size) {
-    try {
-      if (size > Integer.MAX_VALUE) {
-        throw new IllegalArgumentException("reserve size should not be larger than Integer.MAX_VALUE (0x7fffffff)");
-      }
-      methodReserve.invoke(null, (int) size, (int) size);
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
+    MemoryUtil.reserveDirectMemory(size);
   }
 
   /**
    * Unreserve bytes by invoking java.nio.java.Bitjava.nio.Bitss#unreserveMemory.
    */
   @Override
   public void unreserve(long size) {
-    try {
-      if (size > Integer.MAX_VALUE) {
-        throw new IllegalArgumentException("unreserve size should not be larger than Integer.MAX_VALUE (0x7fffffff)");
-      }
-      methodUnreserve.invoke(null, (int) size, (int) size);
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
+    MemoryUtil.unreserveDirectMemory(size);
   }
 
   /**
    * Get current reservation of jVM direct memory. Visible for testing.
    */
   @VisibleForTesting
   public long getCurrentDirectMemReservation() {
-    try {
-      final Class<?> classBits = Class.forName("java.nio.Bits");
-      final Field f = classBits.getDeclaredField("reservedMemory");
-      f.setAccessible(true);
-      return ((AtomicLong) f.get(null)).get();
-    } catch (Exception e) {
-      throw new RuntimeException(e);
-    }
+    return MemoryUtil.getCurrentDirectMemReservation();
   }
 }

java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java

@@ -106,7 +106,7 @@ private BufferLedger associate(final BufferAllocator allocator, final boolean re
         return ledger;
       }
 
-      ledger = new BufferLedger(allocator, this);
+      ledger = allocator.getBufferLedgerFactory().create(allocator, this);
 
       if (retain) {
         // the new reference manager will have a ref count of 1

java/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java

@@ -1094,6 +1094,15 @@ public long getId() {
     return id;
   }
 
+  /**
+   * Create a logger of this {@link ArrowBuf}.
+   *
+   * @return the newly created logger
+   */
+  Logger createLogger() {
+    return new Logger(id, memoryAddress(), length, historicalLog);
+  }
+
   /**
    * Prints information of this buffer into <code>sb</code> at the given
    * indentation and verbosity level.
@@ -1103,12 +1112,7 @@ public long getId() {
    *
    */
   public void print(StringBuilder sb, int indent, Verbosity verbosity) {
-    CommonUtil.indent(sb, indent).append(toString());
-
-    if (BaseAllocator.DEBUG && verbosity.includeHistoricalLog) {
-      sb.append("\n");
-      historicalLog.buildHistory(sb, indent + 1, verbosity.includeStackTraces);
-    }
+    new Logger(id, addr, length, historicalLog).print(sb, indent, verbosity);
   }
 
   /**
@@ -1199,4 +1203,55 @@ public ArrowBuf clear() {
     this.readerIndex = this.writerIndex = 0;
     return this;
   }
+
+  /**
+   * Initialize the reader and writer index.
+   * @param readerIndex index to read from
+   * @param writerIndex index to write to
+   * @return this
+   */
+  @Deprecated
+  public ArrowBuf setIndex(int readerIndex, int writerIndex) {
+    if (readerIndex >= 0 && readerIndex <= writerIndex && writerIndex <= this.capacity()) {
+      this.readerIndex = readerIndex;
+      this.writerIndex = writerIndex;
+      return this;
+    } else {
+      throw new IndexOutOfBoundsException(String.format("readerIndex: %d, writerIndex: %d " +
+       "(expected:0 <= readerIndex <= writerIndex <= capacity(%d))", readerIndex, writerIndex, this.capacity()));
+    }
+  }
+
+  /**
+   * Create a logger for an {@link ArrowBuf}. This is currently used in debugging or historical logging
+   * in code of {@link BufferLedger} to avoid directly holding a strong reference to {@link ArrowBuf}.
+   * So that GC could be able to involved in auto cleaning logic in {@link AutoBufferLedger}.
+   */
+  static class Logger {
+    private final long id;
+    private final long addr;
+    private final long length;
+    private final HistoricalLog historicalLog;
+
+    public Logger(long id, long addr, long length, HistoricalLog historicalLog) {
+      this.id = id;
+      this.addr = addr;
+      this.length = length;
+      this.historicalLog = historicalLog;
+    }
+
+    public void print(StringBuilder sb, int indent, Verbosity verbosity) {
+      CommonUtil.indent(sb, indent).append(toString());
+
+      if (BaseAllocator.DEBUG && verbosity.includeHistoricalLog) {
+        sb.append("\n");
+        historicalLog.buildHistory(sb, indent + 1, verbosity.includeStackTraces);
+      }
+    }
+
+    @Override
+    public String toString() {
+      return String.format("ArrowBuf.Logger[%d], address:%d, length:%d", id, addr, length);
+    }
+  }
 }

java/memory/memory-core/src/main/java/org/apache/arrow/memory/AutoBufferLedger.java

@@ -0,0 +1,206 @@
+/*
+ * 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.arrow.memory;
+
+import sun.misc.Cleaner;
+
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+
+/**
+ * An alternative implementation of {@link BufferLedger}. The reference is auto managed by JVM garbage collector
+ * comparing to {@link LegacyBufferLedger}. Explicit calls to reference management methods such as
+ * {@link #retain()} and {@link #release()} will be ignored.
+ *
+ * <p>
+ * Note when this implementation, the accurate release time of the underlying {@link AllocationManager} may become
+ * unpredictable because we are relying GC to do clean-up. As a result, it's recommended to specify very large
+ * allocation limit (e.g. {@link Integer#MAX_VALUE}) to the corresponding {@link BufferAllocator} to avoid
+ * unexpected allocation failures.
+ * </p>
+ *
+ * <p>
+ * Also, to let the GC be aware of these allocations when off-heap based
+ * {@link AllocationManager}s are used, it's required to also add the allocated sizes to JVM direct
+ * memory counter (which can be limited by specifying JVM option "-XX:MaxDirectMemorySize"). To
+ * achieve this one can simply set allocator's {@link AllocationListener} to
+ * {@link DirectAllocationListener}.
+ * JVM should ensure that garbage collection will be performed once total reservation reached the limit.
+ * </p>
+ */
+public class AutoBufferLedger extends BufferLedger {
+
+  public static class Factory implements BufferLedger.Factory, AutoCloseable {
+    private AutoBufferLedger tail = null;
+
+    @Override
+    public BufferLedger create(BufferAllocator allocator, AllocationManager allocationManager) {
+      return new AutoBufferLedger(allocator, allocationManager, this);
+    }
+
+    private void link(AutoBufferLedger ledger) {
+      synchronized (this) {
+        if (ledger.next != null || ledger.prev != null) {
+          throw new IllegalStateException("already linked");
+        }
+        if (tail == null) {
+          tail = ledger;
+          return;
+        }
+        tail.next = ledger;
+        ledger.prev = tail;
+        tail = ledger;
+      }
+    }
+
+    private void unlink(AutoBufferLedger ledger) {
+      synchronized (this) {
+        if (ledger.next == ledger) {
+          return;
+        }
+        if (ledger.prev == ledger) {
+          throw new IllegalStateException();
+        }
+        if (ledger == tail) {
+          tail = ledger.prev;
+        }
+        if (ledger.prev != null) {
+          ledger.prev.next = ledger.next;
+        }
+        if (ledger.next != null) {
+          ledger.next.prev = ledger.prev;
+        }
+        ledger.prev = ledger;
+        ledger.next = ledger;
+      }
+    }
+
+    @Override
+    public void close() {
+      synchronized (this) {
+        while (tail != null) {
+          final AutoBufferLedger tmp = tail.prev;
+          tail.destruct();
+          tail = tmp;
+        }
+      }
+    }
+  }
+
+  public static Factory newFactory() {
+    return new Factory();
+  }
+
+  private volatile long lDestructionTime = 0;
+  private final AtomicInteger refCount = new AtomicInteger(0);
+  private final AtomicBoolean destructed = new AtomicBoolean(false);
+  private final Factory factory;
+
+  private AutoBufferLedger prev = null;
+  private AutoBufferLedger next = null;
+
+  AutoBufferLedger(BufferAllocator allocator, AllocationManager allocationManager,
+      Factory factory) {
+    super(allocator, allocationManager);
+    this.factory = factory;
+    factory.link(this);
+  }
+
+  @Override
+  protected long getDestructionTime() {
+    return lDestructionTime;
+  }
+
+  @Override
+  protected ReferenceManager newReferenceManager() {
+    reserve0();
+    final ReferenceManager rm = new BaseReferenceManager(this);
+    Cleaner.create(rm, new LedgerDeallocator());
+    return rm;
+  }
+
+  @Override
+  public int getRefCount() {
+    return refCount.get();
+  }
+
+  @Override
+  protected void increment() {
+    // no-op
+  }
+
+  @Override
+  public boolean release() {
+    return false;
+  }
+
+  @Override
+  public boolean release(int decrement) {
+    return false;
+  }
+
+  @Override
+  public void retain() {
+
+  }
+
+  @Override
+  public void retain(int increment) {
+
+  }
+
+  private void reserve0() {
+    if (refCount.getAndAdd(1) == 0) {
+      // no-op
+    }
+  }
+
+  private void release0() {
+    if (refCount.addAndGet(-1) == 0) {
+      destruct();
+    }
+  }
+
+  private void destruct() {
+    if (!destructed.compareAndSet(false, true)) {
+      return;
+    }
+    synchronized (getAllocationManager()) {
+      final AllocationManager am = getAllocationManager();
+      lDestructionTime = System.nanoTime();
+      am.release(this);
+    }
+    factory.unlink(this);
+  }
+
+  /**
+   * Release hook will be invoked by JVM cleaner.
+   *
+   * @see #newReferenceManager()
+   */
+  private class LedgerDeallocator implements Runnable {
+
+    private LedgerDeallocator() {
+    }
+
+    @Override
+    public void run() {
+      release0();
+    }
+  }
+}