release ByteBuf off heap memory block

Author: Pengzna

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/allocator/MemoryAllocator.java

@@ -23,5 +23,7 @@ public interface MemoryAllocator {
 
     Object forceAllocate(long size);
 
-    void releaseMemory(long size);
+    void returnMemoryToManager(long size);
+
+    void releaseMemoryBlock(Object memoryBlock);
 }

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/allocator/NettyMemoryAllocator.java

@@ -21,6 +21,7 @@
 
 import io.netty.buffer.ByteBuf;
 import io.netty.buffer.PooledByteBufAllocator;
+import io.netty.util.ReferenceCountUtil;
 
 /**
  * This class makes fully use of Netty's efficient memory management strategy.
@@ -50,7 +51,24 @@ public ByteBuf tryToAllocate(long size) {
     }
 
     @Override
-    public void releaseMemory(long size) {
+    public void returnMemoryToManager(long size) {
         memoryManager.getCurrentOffHeapAllocatedMemory().addAndGet(-size);
     }
+
+    @Override
+    public void releaseMemoryBlock(Object memoryBlock) {
+        if (!(memoryBlock instanceof ByteBuf)) {
+            throw new IllegalArgumentException("memoryBlock must be ByteBuf");
+        }
+        ByteBuf buf = (ByteBuf) memoryBlock;
+        ReferenceCountUtil.safeRelease(buf, ReferenceCountUtil.refCnt(buf));
+    }
+
+    public static void main(String[] args) {
+        MemoryAllocator netty = new NettyMemoryAllocator(null);
+        ByteBuf buf = (ByteBuf) netty.forceAllocate(1024);
+        System.out.println(ReferenceCountUtil.refCnt(buf));
+        netty.releaseMemoryBlock(buf);
+        System.out.println(ReferenceCountUtil.refCnt(buf));
+    }
 }

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/allocator/OnHeapMemoryStrategy.java

@@ -48,7 +48,12 @@ public byte[] forceAllocate(long size) {
     }
 
     @Override
-    public void releaseMemory(long size) {
+    public void returnMemoryToManager(long size) {
         memoryManager.getCurrentOnHeapAllocatedMemory().addAndGet(-size);
     }
+
+    @Override
+    public void releaseMemoryBlock(Object memoryBlock) {
+        memoryBlock = null;
+    }
 }

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/MemoryConsumer.java

@@ -17,8 +17,12 @@
 
 package org.apache.hugegraph.memory.consumer;
 
+import java.util.List;
+
 import org.apache.hugegraph.memory.pool.MemoryPool;
 
+import io.netty.buffer.ByteBuf;
+
 /**
  * This interface is used by immutable, memory-heavy objects which will be stored in off heap.
  */
@@ -44,4 +48,6 @@ public interface MemoryConsumer {
     void releaseOriginalOnHeapVars();
 
     MemoryPool getOperatorMemoryPool();
+
+    List<ByteBuf> getAllOffHeapByteBuf();
 }

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/BinaryIdOffHeap.java

@@ -17,6 +17,9 @@
 
 package org.apache.hugegraph.memory.consumer.impl.id;
 
+import java.util.Collections;
+import java.util.List;
+
 import org.apache.hugegraph.backend.id.Id;
 import org.apache.hugegraph.backend.serializer.BinaryBackendEntry;
 import org.apache.hugegraph.memory.consumer.MemoryConsumer;
@@ -59,6 +62,11 @@ public MemoryPool getOperatorMemoryPool() {
         return this.memoryPool;
     }
 
+    @Override
+    public List<ByteBuf> getAllOffHeapByteBuf() {
+        return Collections.singletonList(bytesOffHeap);
+    }
+
     @Override
     public void releaseOriginalOnHeapVars() {
         this.bytes = null;

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/EdgeIdOffHeap.java

@@ -18,6 +18,7 @@
 package org.apache.hugegraph.memory.consumer.impl.id;
 
 import java.nio.charset.StandardCharsets;
+import java.util.List;
 import java.util.Objects;
 
 import org.apache.hugegraph.backend.id.EdgeId;
@@ -29,6 +30,8 @@
 import org.apache.hugegraph.structure.HugeVertex;
 import org.apache.hugegraph.type.define.Directions;
 
+import com.google.common.collect.Lists;
+
 import io.netty.buffer.ByteBuf;
 
 // TODO: rewrite static method in EdgeId
@@ -140,6 +143,11 @@ public MemoryPool getOperatorMemoryPool() {
         return memoryPool;
     }
 
+    @Override
+    public List<ByteBuf> getAllOffHeapByteBuf() {
+        return Lists.newArrayList(this.sortValuesOffHeap, this.cacheOffHeap);
+    }
+
     @Override
     public EdgeId switchDirection() {
         Directions newDirection = this.direction.opposite();

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/LongIdOffHeap.java

@@ -19,6 +19,8 @@
 
 import static org.apache.hugegraph.backend.id.IdGenerator.compareType;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.Objects;
 
 import org.apache.hugegraph.backend.id.Id;
@@ -75,6 +77,11 @@ public MemoryPool getOperatorMemoryPool() {
         return memoryPool;
     }
 
+    @Override
+    public List<ByteBuf> getAllOffHeapByteBuf() {
+        return Collections.singletonList(idOffHeap);
+    }
+
     @Override
     public long asLong() {
         try {

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/ObjectIdOffHeap.java

@@ -17,6 +17,8 @@
 
 package org.apache.hugegraph.memory.consumer.impl.id;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.Objects;
 
 import org.apache.hugegraph.backend.id.IdGenerator;
@@ -63,6 +65,11 @@ public MemoryPool getOperatorMemoryPool() {
         return memoryPool;
     }
 
+    @Override
+    public List<ByteBuf> getAllOffHeapByteBuf() {
+        return Collections.singletonList(objectOffHeap);
+    }
+
     @Override
     public Object asObject() {
         return FurySerializationUtils.FURY.deserialize(ByteBufUtil.getBytes(objectOffHeap));

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/QueryIdOffHeap.java

@@ -18,6 +18,8 @@
 package org.apache.hugegraph.memory.consumer.impl.id;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.List;
 
 import org.apache.hugegraph.backend.cache.CachedBackendStore;
 import org.apache.hugegraph.backend.id.Id;
@@ -69,6 +71,11 @@ public MemoryPool getOperatorMemoryPool() {
         return memoryPool;
     }
 
+    @Override
+    public List<ByteBuf> getAllOffHeapByteBuf() {
+        return Collections.singletonList(queryOffHeap);
+    }
+
     @Override
     public boolean equals(Object other) {
         if (!(other instanceof QueryIdOffHeap)) {

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/StringIdOffHeap.java

@@ -20,6 +20,8 @@
 import static org.apache.hugegraph.backend.id.IdGenerator.compareType;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Collections;
+import java.util.List;
 
 import org.apache.hugegraph.backend.id.Id;
 import org.apache.hugegraph.backend.id.IdGenerator;
@@ -75,6 +77,11 @@ public MemoryPool getOperatorMemoryPool() {
         return memoryPool;
     }
 
+    @Override
+    public List<ByteBuf> getAllOffHeapByteBuf() {
+        return Collections.singletonList(idOffHeap);
+    }
+
     @Override
     public Object asObject() {
         return this.asString();

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/id/UuidIdOffHeap.java

@@ -20,6 +20,8 @@
 import static org.apache.hugegraph.backend.id.IdGenerator.compareType;
 
 import java.io.IOException;
+import java.util.Collections;
+import java.util.List;
 import java.util.UUID;
 
 import org.apache.hugegraph.backend.id.Id;
@@ -87,6 +89,11 @@ public MemoryPool getOperatorMemoryPool() {
         return memoryPool;
     }
 
+    @Override
+    public List<ByteBuf> getAllOffHeapByteBuf() {
+        return Collections.singletonList(idOffHeap);
+    }
+
     @Override
     public Object asObject() {
         return FurySerializationUtils.FURY.deserialize(

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/property/HugeEdgePropertyOffHeap.java

@@ -17,6 +17,8 @@
 
 package org.apache.hugegraph.memory.consumer.impl.property;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.NoSuchElementException;
 
 import org.apache.hugegraph.memory.consumer.MemoryConsumer;
@@ -68,6 +70,11 @@ public MemoryPool getOperatorMemoryPool() {
         return memoryPool;
     }
 
+    @Override
+    public List<ByteBuf> getAllOffHeapByteBuf() {
+        return Collections.singletonList(valueOffHeap);
+    }
+
     @Override
     public HugeEdgeProperty<V> switchEdgeOwner() {
         assert this.owner instanceof HugeEdge;

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/consumer/impl/property/HugeVertexPropertyOffHeap.java

@@ -17,6 +17,8 @@
 
 package org.apache.hugegraph.memory.consumer.impl.property;
 
+import java.util.Collections;
+import java.util.List;
 import java.util.NoSuchElementException;
 
 import org.apache.hugegraph.memory.consumer.MemoryConsumer;
@@ -67,6 +69,11 @@ public MemoryPool getOperatorMemoryPool() {
         return memoryPool;
     }
 
+    @Override
+    public List<ByteBuf> getAllOffHeapByteBuf() {
+        return Collections.singletonList(valueOffHeap);
+    }
+
     @Override
     public Object serialValue(boolean encodeNumber) {
         return this.pkey.serialValue(this.value(), encodeNumber);

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/pool/AbstractMemoryPool.java

@@ -25,6 +25,7 @@
 import java.util.concurrent.locks.ReentrantLock;
 
 import org.apache.hugegraph.memory.MemoryManager;
+import org.apache.hugegraph.memory.consumer.MemoryConsumer;
 import org.apache.hugegraph.memory.pool.impl.MemoryPoolStats;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -139,6 +140,11 @@ public Object tryToAcquireMemoryInternal(long bytes) {
         return null;
     }
 
+    @Override
+    public void bindMemoryConsumer(MemoryConsumer memoryConsumer) {
+        // default do nothing
+    }
+
     @Override
     public Object requireMemory(long bytes) {
         return null;

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/pool/MemoryPool.java

@@ -17,6 +17,7 @@
 
 package org.apache.hugegraph.memory.pool;
 
+import org.apache.hugegraph.memory.consumer.MemoryConsumer;
 import org.apache.hugegraph.memory.pool.impl.MemoryPoolStats;
 import org.apache.hugegraph.memory.util.QueryOutOfMemoryException;
 
@@ -49,4 +50,6 @@ public interface MemoryPool {
     MemoryPool getParentPool();
 
     MemoryPool findRootQueryPool();
+
+    void bindMemoryConsumer(MemoryConsumer memoryConsumer);
 }

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/pool/impl/OperatorMemoryPool.java

@@ -17,8 +17,12 @@
 
 package org.apache.hugegraph.memory.pool.impl;
 
+import java.util.HashSet;
+import java.util.Set;
+
 import org.apache.hugegraph.memory.MemoryManager;
 import org.apache.hugegraph.memory.allocator.MemoryAllocator;
+import org.apache.hugegraph.memory.consumer.MemoryConsumer;
 import org.apache.hugegraph.memory.pool.AbstractMemoryPool;
 import org.apache.hugegraph.memory.pool.MemoryPool;
 import org.apache.hugegraph.memory.util.MemoryManageUtils;
@@ -30,16 +34,27 @@ public class OperatorMemoryPool extends AbstractMemoryPool {
 
     private static final Logger LOG = LoggerFactory.getLogger(OperatorMemoryPool.class);
     private final MemoryAllocator memoryAllocator;
+    private final Set<MemoryConsumer> memoryConsumers;
 
     public OperatorMemoryPool(MemoryPool parent, String poolName,
                               MemoryAllocator memoryAllocator, MemoryManager memoryManager) {
         super(parent, poolName, memoryManager);
         this.memoryAllocator = memoryAllocator;
+        this.memoryConsumers = new HashSet<>();
+    }
+
+    @Override
+    public void bindMemoryConsumer(MemoryConsumer memoryConsumer) {
+        this.memoryConsumers.add(memoryConsumer);
     }
 
     @Override
     public synchronized void releaseSelf(String reason) {
-        this.memoryAllocator.releaseMemory(getAllocatedBytes());
+        this.memoryAllocator.returnMemoryToManager(getAllocatedBytes());
+        this.memoryConsumers.forEach(memoryConsumer -> {
+            memoryConsumer.getAllOffHeapByteBuf().forEach(memoryAllocator::releaseMemoryBlock);
+        });
+        this.memoryConsumers.clear();
         super.releaseSelf(reason);
         // TODO: release memory consumer, release byte buffer.
     }