Add AutoBufferLedger (apache#31)

zhouyuan · Apr 26, 2022 · edc3973 · edc3973
1 parent ffde238
commit edc3973
Show file tree

Hide file tree

Showing 16 changed files with 1,176 additions and 408 deletions.
diff --git a/java/dataset/src/main/java/org/apache/arrow/dataset/jni/DirectReservationListener.java b/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();
   }
 }
diff --git a/java/dataset/src/main/java/org/apache/arrow/dataset/jni/UnsafeRecordBatchSerializer.java b/java/dataset/src/main/java/org/apache/arrow/dataset/jni/UnsafeRecordBatchSerializer.java
@@ -35,10 +35,7 @@
 import org.apache.arrow.flatbuf.Message;
 import org.apache.arrow.flatbuf.MessageHeader;
 import org.apache.arrow.flatbuf.RecordBatch;
-import org.apache.arrow.memory.ArrowBuf;
-import org.apache.arrow.memory.BufferAllocator;
-import org.apache.arrow.memory.BufferLedger;
-import org.apache.arrow.memory.NativeUnderlyingMemory;
+import org.apache.arrow.memory.*;
 import org.apache.arrow.memory.util.LargeMemoryUtil;
 import org.apache.arrow.util.Preconditions;
 import org.apache.arrow.vector.compression.NoCompressionCodec;
@@ -122,19 +119,19 @@ public static ArrowRecordBatch deserializeUnsafe(
             throw new IllegalArgumentException("Buffer count mismatch between metadata and native managed refs");
         }
 
-        final ArrayList<ArrowBuf> buffers = new ArrayList<>();
-        for (int i = 0; i < batchMeta.buffersLength(); i++) {
-            final Buffer bufferMeta = batchMeta.buffers(i);
-            final KeyValue keyValue = metaMessage.customMetadata(i); // custom metadata containing native buffer refs
-            final byte[] refDecoded = Base64.getDecoder().decode(keyValue.value());
-            final long nativeBufferRef = ByteBuffer.wrap(refDecoded).order(ByteOrder.LITTLE_ENDIAN).getLong();
-            final int size = LargeMemoryUtil.checkedCastToInt(bufferMeta.length());
-            final NativeUnderlyingMemory am = NativeUnderlyingMemory.create(allocator,
-                    size, nativeBufferRef, bufferMeta.offset());
-            BufferLedger ledger = am.associate(allocator);
-            ArrowBuf buf = new ArrowBuf(ledger, null, size, bufferMeta.offset());
-            buffers.add(buf);
-        }
+    final ArrayList<ArrowBuf> buffers = new ArrayList<>();
+    for (int i = 0; i < batchMeta.buffersLength(); i++) {
+      final Buffer bufferMeta = batchMeta.buffers(i);
+      final KeyValue keyValue = metaMessage.customMetadata(i); // custom metadata containing native buffer refs
+      final byte[] refDecoded = Base64.getDecoder().decode(keyValue.value());
+      final long nativeBufferRef = ByteBuffer.wrap(refDecoded).order(ByteOrder.LITTLE_ENDIAN).getLong();
+      final int size = LargeMemoryUtil.checkedCastToInt(bufferMeta.length());
+      final NativeUnderlyingMemory am = NativeUnderlyingMemory.create(allocator,
+          size, nativeBufferRef, bufferMeta.offset());
+      ReferenceManager rm = am.createReferenceManager(allocator);
+      ArrowBuf buf = new ArrowBuf(rm, null, size, bufferMeta.offset());
+      buffers.add(buf);
+    }
 
         try {
             final int numRows = LargeMemoryUtil.checkedCastToInt(batchMeta.length());
@@ -263,4 +260,4 @@ public void close() throws Exception {
             delegated.close();
         }
     }
-}
+}
diff --git a/java/dataset/src/main/java/org/apache/arrow/memory/NativeUnderlyingMemory.java b/java/dataset/src/main/java/org/apache/arrow/memory/NativeUnderlyingMemory.java
@@ -60,8 +60,8 @@ public static NativeUnderlyingMemory create(BufferAllocator bufferAllocator, int
     return new NativeUnderlyingMemory(bufferAllocator, size, nativeInstanceId, address);
   }
 
-  public BufferLedger associate(BufferAllocator allocator) {
-    return super.associate(allocator);
+  public ReferenceManager createReferenceManager(BufferAllocator allocator) {
+    return super.associate(allocator).newReferenceManager();
   }
 
   @Override

diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/AllocationManager.java b/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

diff --git a/java/memory/memory-core/src/main/java/org/apache/arrow/memory/ArrowBuf.java b/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);
+    }
+  }
 }