Implememt auto release technique for Blob Allocator

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/binaryallocator/BinaryAllocator.java

@@ -21,6 +21,7 @@
 
 import org.apache.iotdb.commons.binaryallocator.arena.Arena;
 import org.apache.iotdb.commons.binaryallocator.arena.ArenaStrategy;
+import org.apache.iotdb.commons.binaryallocator.autoreleaser.Releaser;
 import org.apache.iotdb.commons.binaryallocator.config.AllocatorConfig;
 import org.apache.iotdb.commons.binaryallocator.evictor.Evictor;
 import org.apache.iotdb.commons.binaryallocator.metric.BinaryAllocatorMetrics;
@@ -33,7 +34,11 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.lang.ref.ReferenceQueue;
 import java.time.Duration;
+import java.util.Collections;
+import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
 
 public class BinaryAllocator {
@@ -49,14 +54,22 @@ public class BinaryAllocator {
 
   private final BinaryAllocatorMetrics metrics;
   private Evictor sampleEvictor;
+  private Releaser autoReleaser;
   private static final ThreadLocal<ThreadArenaRegistry> arenaRegistry =
       ThreadLocal.withInitial(ThreadArenaRegistry::new);
 
-  private static final int WARNING_GC_TIME_PERCENTAGE = 10;
-  private static final int HALF_GC_TIME_PERCENTAGE = 20;
+  private static final int WARNING_GC_TIME_PERCENTAGE = 20;
+  private static final int HALF_GC_TIME_PERCENTAGE = 25;
   private static final int SHUTDOWN_GC_TIME_PERCENTAGE = 30;
   private static final int RESTART_GC_TIME_PERCENTAGE = 5;
 
+  public final ReferenceQueue<PooledBinary> referenceQueue = new ReferenceQueue<>();
+
+  // JDK 9+ Cleaner uses double-linked list and synchronized to manage references, which has worse
+  // performance than lock-free hash set
+  public final Set<PooledBinaryPhantomReference> phantomRefs =
+      Collections.newSetFromMap(new ConcurrentHashMap<>());
+
   public BinaryAllocator(AllocatorConfig allocatorConfig) {
     this.allocatorConfig = allocatorConfig;
 
@@ -87,8 +100,14 @@ public synchronized void start() {
     sampleEvictor =
         new SampleEvictor(
             ThreadName.BINARY_ALLOCATOR_SAMPLE_EVICTOR.getName(),
-            allocatorConfig.durationEvictorShutdownTimeout);
-    sampleEvictor.startEvictor(allocatorConfig.durationBetweenEvictorRuns);
+            allocatorConfig.durationShutdownTimeout,
+            allocatorConfig.durationBetweenEvictorRuns);
+    sampleEvictor.start();
+    autoReleaser =
+        new AutoReleaser(
+            ThreadName.BINARY_ALLOCATOR_AUTO_RELEASER.getName(),
+            allocatorConfig.durationShutdownTimeout);
+    autoReleaser.start();
   }
 
   public synchronized void close(boolean forceClose) {
@@ -99,13 +118,14 @@ public synchronized void close(boolean forceClose) {
       state.set(BinaryAllocatorState.PENDING);
     }
 
-    sampleEvictor.stopEvictor();
+    sampleEvictor.stop();
+    autoReleaser.stop();
     for (Arena arena : heapArenas) {
       arena.close();
     }
   }
 
-  public PooledBinary allocateBinary(int reqCapacity) {
+  public PooledBinary allocateBinary(int reqCapacity, boolean autoRelease) {
     if (reqCapacity < allocatorConfig.minAllocateSize
         || reqCapacity > allocatorConfig.maxAllocateSize
         || state.get() != BinaryAllocatorState.OPEN) {
@@ -114,7 +134,7 @@ public PooledBinary allocateBinary(int reqCapacity) {
 
     Arena arena = arenaStrategy.choose(heapArenas);
 
-    return new PooledBinary(arena.allocate(reqCapacity), reqCapacity, arena.getArenaID());
+    return arena.allocate(reqCapacity, autoRelease);
   }
 
   public void deallocateBinary(PooledBinary binary) {
@@ -125,7 +145,7 @@ public void deallocateBinary(PooledBinary binary) {
       int arenaIndex = binary.getArenaIndex();
       if (arenaIndex != -1) {
         Arena arena = heapArenas[arenaIndex];
-        arena.deallocate(binary.getValues());
+        arena.deallocate(binary);
       }
     }
   }
@@ -168,11 +188,13 @@ public static BinaryAllocator getInstance() {
   }
 
   private static class BinaryAllocatorHolder {
+
     private static final BinaryAllocator INSTANCE =
         new BinaryAllocator(AllocatorConfig.DEFAULT_CONFIG);
   }
 
   private static class ThreadArenaRegistry {
+
     private Arena threadArenaBinding = null;
 
     public Arena getArena() {
@@ -199,6 +221,7 @@ protected void finalize() {
   }
 
   private static class LeastUsedArenaStrategy implements ArenaStrategy {
+
     @Override
     public Arena choose(Arena[] arenas) {
       Arena boundArena = arenaRegistry.get().getArena();
@@ -250,8 +273,9 @@ public void runGcEviction(long curGcTimePercent) {
 
   public class SampleEvictor extends Evictor {
 
-    public SampleEvictor(String name, Duration evictorShutdownTimeoutDuration) {
-      super(name, evictorShutdownTimeoutDuration);
+    public SampleEvictor(
+        String name, Duration evictorShutdownTimeoutDuration, Duration durationBetweenEvictorRuns) {
+      super(name, evictorShutdownTimeoutDuration, durationBetweenEvictorRuns);
     }
 
     @Override
@@ -263,4 +287,26 @@ public void run() {
       metrics.updateSampleEvictionCounter(evictedSize);
     }
   }
+
+  /** Process phantomly reachable objects and return their byte arrays to pool. */
+  public class AutoReleaser extends Releaser {
+
+    public AutoReleaser(String name, Duration shutdownTimeoutDuration) {
+      super(name, shutdownTimeoutDuration);
+    }
+
+    @Override
+    public void run() {
+      PooledBinaryPhantomReference ref;
+      try {
+        while ((ref = (PooledBinaryPhantomReference) referenceQueue.remove()) != null) {
+          phantomRefs.remove(ref);
+          ref.slabRegion.deallocate(ref.byteArray);
+        }
+      } catch (InterruptedException e) {
+        LOGGER.info("{} exits due to interruptedException.", name);
+        Thread.currentThread().interrupt();
+      }
+    }
+  }
 }

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/binaryallocator/PooledBinaryPhantomReference.java

@@ -0,0 +1,42 @@
+/*
+ * 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.iotdb.commons.binaryallocator;
+
+import org.apache.iotdb.commons.binaryallocator.arena.Arena;
+
+import org.apache.tsfile.utils.PooledBinary;
+
+import java.lang.ref.PhantomReference;
+import java.lang.ref.ReferenceQueue;
+
+public class PooledBinaryPhantomReference extends PhantomReference<PooledBinary> {
+  public final byte[] byteArray;
+  public Arena.SlabRegion slabRegion;
+
+  public PooledBinaryPhantomReference(
+      PooledBinary referent,
+      ReferenceQueue<? super PooledBinary> q,
+      byte[] byteArray,
+      Arena.SlabRegion region) {
+    super(referent, q);
+    this.byteArray = byteArray;
+    this.slabRegion = region;
+  }
+}

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/binaryallocator/arena/Arena.java

@@ -20,10 +20,15 @@
 package org.apache.iotdb.commons.binaryallocator.arena;
 
 import org.apache.iotdb.commons.binaryallocator.BinaryAllocator;
+import org.apache.iotdb.commons.binaryallocator.PooledBinaryPhantomReference;
 import org.apache.iotdb.commons.binaryallocator.config.AllocatorConfig;
 import org.apache.iotdb.commons.binaryallocator.ema.AdaptiveWeightedAverage;
 import org.apache.iotdb.commons.binaryallocator.utils.SizeClasses;
 
+import org.apache.tsfile.utils.PooledBinary;
+
+import java.lang.ref.ReferenceQueue;
+import java.util.Set;
 import java.util.concurrent.ConcurrentLinkedQueue;
 import java.util.concurrent.atomic.AtomicInteger;
 
@@ -39,6 +44,9 @@ public class Arena {
 
   private int sampleCount;
 
+  private final ReferenceQueue<PooledBinary> referenceQueue;
+  private final Set<PooledBinaryPhantomReference> phantomRefs;
+
   public Arena(
       BinaryAllocator allocator, SizeClasses sizeClasses, int id, AllocatorConfig allocatorConfig) {
     this.binaryAllocator = allocator;
@@ -52,20 +60,31 @@ public Arena(
     }
 
     sampleCount = 0;
+    referenceQueue = binaryAllocator.referenceQueue;
+    phantomRefs = binaryAllocator.phantomRefs;
   }
 
   public int getArenaID() {
     return arenaID;
   }
 
-  public byte[] allocate(int reqCapacity) {
+  public PooledBinary allocate(int reqCapacity, boolean autoRelease) {
     final int sizeIdx = sizeClasses.size2SizeIdx(reqCapacity);
-    return regions[sizeIdx].allocate();
+    byte[] data = regions[sizeIdx].allocate();
+    if (autoRelease) {
+      PooledBinary binary = new PooledBinary(data, reqCapacity, -1);
+      PooledBinaryPhantomReference ref =
+          new PooledBinaryPhantomReference(binary, referenceQueue, data, regions[sizeIdx]);
+      phantomRefs.add(ref);
+      return binary;
+    } else {
+      return new PooledBinary(data, reqCapacity, arenaID);
+    }
   }
 
-  public void deallocate(byte[] bytes) {
-    final int sizeIdx = sizeClasses.size2SizeIdx(bytes.length);
-    regions[sizeIdx].deallocate(bytes);
+  public void deallocate(PooledBinary binary) {
+    final int sizeIdx = sizeClasses.size2SizeIdx(binary.getLength());
+    regions[sizeIdx].deallocate(binary.getValues());
   }
 
   public long evict(double ratio) {
@@ -146,8 +165,13 @@ public long runSampleEviction() {
     return evictedSize;
   }
 
-  private static class SlabRegion {
+  public static class SlabRegion {
     private final int byteArraySize;
+
+    // Current implementation uses ConcurrentLinkedQueue for simplicity
+    // TODO: Can be optimized with more efficient lock-free approaches:
+    // 1. No need for strict FIFO, it's just an object pool
+    // 2. Use segmented arrays/queues with per-segment counters to reduce contention
     private final ConcurrentLinkedQueue<byte[]> queue;
 
     private final AtomicInteger allocationsFromAllocator;

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/binaryallocator/autoreleaser/Releaser.java

@@ -0,0 +1,91 @@
+/*
+ * 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.iotdb.commons.binaryallocator.autoreleaser;
+
+import org.apache.iotdb.commons.concurrent.IoTDBThreadPoolFactory;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.time.Duration;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+
+public abstract class Releaser implements Runnable {
+  private static final Logger LOGGER = LoggerFactory.getLogger(Releaser.class);
+
+  private Future<?> future;
+  protected final String name;
+  private final Duration shutdownTimeoutDuration;
+
+  private ExecutorService executor;
+
+  public Releaser(String name, Duration shutdownTimeoutDuration) {
+    this.name = name;
+    this.shutdownTimeoutDuration = shutdownTimeoutDuration;
+  }
+
+  /** Cancels the future. */
+  void cancel() {
+    future.cancel(false);
+  }
+
+  @Override
+  public abstract void run();
+
+  void setFuture(final Future<?> future) {
+    this.future = future;
+  }
+
+  @Override
+  public String toString() {
+    return getClass().getName() + " [future=" + future + "]";
+  }
+
+  public void start() {
+    if (null == executor) {
+      executor = IoTDBThreadPoolFactory.newSingleThreadExecutor(name);
+    }
+    final Future<?> future = executor.submit(this);
+    this.setFuture(future);
+  }
+
+  public void stop() {
+    if (executor == null) {
+      return;
+    }
+
+    LOGGER.info("Stopping {}", name);
+
+    cancel();
+    executor.shutdown();
+    try {
+      boolean result =
+          executor.awaitTermination(shutdownTimeoutDuration.toMillis(), TimeUnit.MILLISECONDS);
+      if (!result) {
+        LOGGER.info("unable to stop auto releaser after {} ms", shutdownTimeoutDuration.toMillis());
+      }
+    } catch (final InterruptedException ignored) {
+      Thread.currentThread().interrupt();
+    }
+    executor = null;
+  }
+}

iotdb-core/node-commons/src/main/java/org/apache/iotdb/commons/binaryallocator/config/AllocatorConfig.java

@@ -37,8 +37,8 @@ public class AllocatorConfig {
   public boolean enableBinaryAllocator =
       CommonDescriptor.getInstance().getConfig().isEnableBinaryAllocator();
 
-  /** Maximum wait time in milliseconds when shutting down the evictor */
-  public Duration durationEvictorShutdownTimeout = Duration.ofMillis(1000L);
+  /** Maximum wait time in milliseconds when shutting down the evictor and autoReleaser */
+  public Duration durationShutdownTimeout = Duration.ofMillis(1000L);
 
   /** Time interval in milliseconds between two consecutive evictor runs */
   public Duration durationBetweenEvictorRuns = Duration.ofMillis(1000L);