fix some bugs: arbitration & suspend

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

@@ -17,8 +17,10 @@
 
 package org.apache.hugegraph.memory;
 
+import java.util.Map;
 import java.util.PriorityQueue;
 import java.util.Queue;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Future;
@@ -38,18 +40,22 @@
 public class MemoryManager {
 
     private static final Logger LOG = LoggerFactory.getLogger(MemoryManager.class);
+    private static final int ARBITRATE_MEMORY_THREAD_NUM = 12;
     private static final String QUERY_MEMORY_POOL_NAME_PREFIX = "QueryMemoryPool";
     private static final String ARBITRATE_MEMORY_POOL_NAME = "ArbitrateMemoryPool";
-    private static final String DELIMINATOR = "_";
-    private static final int ARBITRATE_MEMORY_THREAD_NUM = 12;
+    public static final String DELIMINATOR = "_";
+
     // TODO: read it from conf, current 1G
     public static final long MAX_MEMORY_CAPACITY_IN_BYTES = Bytes.GB;
     private final AtomicLong currentAvailableMemoryInBytes =
             new AtomicLong(MAX_MEMORY_CAPACITY_IN_BYTES);
     private final AtomicLong currentOffHeapAllocatedMemoryInBytes = new AtomicLong(0);
     private final AtomicLong currentOnHeapAllocatedMemoryInBytes = new AtomicLong(0);
+
     private final Queue<MemoryPool> queryMemoryPools =
             new PriorityQueue<>((o1, o2) -> (int) (o2.getFreeBytes() - o1.getFreeBytes()));
+    private final Map<String, MemoryPool> threadName2TaskMemoryPoolMap = new ConcurrentHashMap<>();
+
     private final MemoryArbitrator memoryArbitrator;
     private final ExecutorService arbitrateExecutor;
 
@@ -128,6 +134,21 @@ public synchronized long handleRequestFromQueryPool(long size) {
         return size;
     }
 
+    /**
+     * Used by task thread to find its memory pool to release self's memory resource when exiting.
+     */
+    public MemoryPool getCorrespondingTaskMemoryPool(String threadName) {
+        return threadName2TaskMemoryPoolMap.getOrDefault(threadName, null);
+    }
+
+    public void bindCorrespondingTaskMemoryPool(String threadName, MemoryPool memoryPool) {
+        threadName2TaskMemoryPoolMap.computeIfAbsent(threadName, key -> memoryPool);
+    }
+
+    public void removeCorrespondingTaskMemoryPool(String threadName) {
+        threadName2TaskMemoryPoolMap.remove(threadName);
+    }
+
     public Queue<MemoryPool> getCurrentQueryMemoryPools() {
         return new PriorityQueue<>(queryMemoryPools);
     }

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/memory/arbitrator/MemoryArbitrator.java

@@ -21,7 +21,7 @@
 
 public interface MemoryArbitrator {
 
-    long MAX_WAIT_TIME_FOR_LOCAL_RECLAIM = 1000;
+    long MAX_WAIT_TIME_FOR_LOCAL_RECLAIM = 2000;
 
     long MAX_WAIT_TIME_FOR_GLOBAL_RECLAIM = 5000;
 

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

@@ -27,17 +27,19 @@
 import org.apache.hugegraph.memory.MemoryManager;
 import org.apache.hugegraph.memory.consumer.MemoryConsumer;
 import org.apache.hugegraph.memory.pool.impl.MemoryPoolStats;
+import org.jetbrains.annotations.TestOnly;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 public abstract class AbstractMemoryPool implements MemoryPool {
 
     private static final Logger LOG = LoggerFactory.getLogger(AbstractMemoryPool.class);
-    private final Queue<MemoryPool> children =
+    protected final Queue<MemoryPool> children =
             new PriorityQueue<>((o1, o2) -> (int) (o2.getFreeBytes() - o1.getFreeBytes()));
     protected final MemoryManager memoryManager;
-    protected final ReentrantLock arbitrationLock = new ReentrantLock();
-    protected final Condition condition = arbitrationLock.newCondition();
+    // Allocation, deAllocation, arbitration must be serial which is controlled by this lock.
+    protected final ReentrantLock memoryActionLock = new ReentrantLock();
+    protected final Condition condition = memoryActionLock.newCondition();
     protected final AtomicBoolean isBeingArbitrated = new AtomicBoolean(false);
     protected final MemoryPoolStats stats;
     protected boolean isClosed = false;
@@ -60,7 +62,7 @@ public long tryToReclaimLocalMemory(long neededBytes) {
         long totalReclaimedBytes = 0;
         long currentNeededBytes = neededBytes;
         try {
-            this.arbitrationLock.lock();
+            this.memoryActionLock.lock();
             this.isBeingArbitrated.set(true);
             for (MemoryPool child : this.children) {
                 long reclaimedMemory = child.tryToReclaimLocalMemory(currentNeededBytes);
@@ -79,20 +81,23 @@ public long tryToReclaimLocalMemory(long neededBytes) {
                      totalReclaimedBytes, neededBytes);
             return totalReclaimedBytes;
         } finally {
-            this.stats.setNumShrinks(this.stats.getNumShrinks() + 1);
+            if (totalReclaimedBytes > 0) {
+                this.stats.setNumShrinks(this.stats.getNumShrinks() + 1);
+            }
             this.stats.setAllocatedBytes(
                     this.stats.getAllocatedBytes() - totalReclaimedBytes);
             this.isBeingArbitrated.set(false);
-            this.arbitrationLock.unlock();
             this.condition.signalAll();
+            this.memoryActionLock.unlock();
         }
     }
 
     /**
      * called when one layer pool is successfully executed and exited.
      */
     @Override
-    public synchronized void releaseSelf(String reason) {
+    public void releaseSelf(String reason) {
+        this.memoryActionLock.lock();
         try {
             if (this.isBeingArbitrated.get()) {
                 this.condition.await();
@@ -109,6 +114,7 @@ public synchronized void releaseSelf(String reason) {
             LOG.error("Failed to release self because ", e);
             Thread.currentThread().interrupt();
         } finally {
+            this.memoryActionLock.unlock();
             // Make these objs be GCed by JVM quickly.
             this.parent = null;
             this.children.clear();
@@ -198,4 +204,9 @@ public MemoryPool findRootQueryPool() {
         }
         return getParentPool().findRootQueryPool();
     }
+
+    @TestOnly
+    public int getChildrenCount() {
+        return this.children.size();
+    }
 }

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

@@ -20,6 +20,7 @@
 import org.apache.hugegraph.memory.consumer.MemoryConsumer;
 import org.apache.hugegraph.memory.pool.impl.MemoryPoolStats;
 import org.apache.hugegraph.memory.util.QueryOutOfMemoryException;
+import org.jetbrains.annotations.TestOnly;
 
 public interface MemoryPool {
 
@@ -51,5 +52,10 @@ public interface MemoryPool {
 
     MemoryPool findRootQueryPool();
 
+    MemoryPool addChildPool();
+
     void bindMemoryConsumer(MemoryConsumer memoryConsumer);
+
+    @TestOnly
+    int getChildrenCount();
 }

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

@@ -22,10 +22,14 @@ public class MemoryPoolStats {
     private final String memoryPoolName;
     private long maxCapacity;
     private long usedBytes;
+    // it represents the cumulative used bytes.
     private long cumulativeBytes;
     private long allocatedBytes;
 
+    // it represents the shrinking num of allocatedBytes
     private long numShrinks;
+    // for query pool, it represents the enlarging num of maxCapacity; for other pools, it
+    // represents the enlarging num of allocatedBytes
     private long numExpands;
     private long numAborts;
 

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

@@ -43,20 +43,27 @@ public OperatorMemoryPool(MemoryPool parent, String poolName,
         this.memoryConsumers = new HashSet<>();
     }
 
+    @Override
+    public MemoryPool addChildPool() {
+        throw new UnsupportedOperationException();
+    }
+
     @Override
     public void bindMemoryConsumer(MemoryConsumer memoryConsumer) {
         this.memoryConsumers.add(memoryConsumer);
     }
 
     @Override
-    public synchronized void releaseSelf(String reason) {
+    public void releaseSelf(String reason) {
+        super.releaseSelf(reason);
+        // since it is already closed, its stats will not be updated. so here we can use its
+        // stats out of memoryActionLock.
         this.memoryAllocator.returnMemoryToManager(getAllocatedBytes());
+        // release memory consumer, release byte buffer.
         this.memoryConsumers.forEach(memoryConsumer -> {
             memoryConsumer.getAllMemoryBlock().forEach(memoryAllocator::releaseMemoryBlock);
         });
         this.memoryConsumers.clear();
-        super.releaseSelf(reason);
-        // TODO: release memory consumer, release byte buffer.
     }
 
     @Override
@@ -66,11 +73,12 @@ public long tryToReclaimLocalMemory(long neededBytes) {
             return 0;
         }
         LOG.info("[{}] tryToReclaimLocalMemory: neededBytes={}", this, neededBytes);
+        long reclaimableBytes = 0;
         try {
-            this.arbitrationLock.lock();
+            this.memoryActionLock.lock();
             this.isBeingArbitrated.set(true);
             // 1. try to reclaim self free memory
-            long reclaimableBytes = getFreeBytes();
+            reclaimableBytes = getFreeBytes();
             // try its best to reclaim memory
             if (reclaimableBytes <= neededBytes) {
                 // 2. update stats
@@ -93,9 +101,12 @@ public long tryToReclaimLocalMemory(long neededBytes) {
 
             return neededBytes;
         } finally {
+            if (reclaimableBytes > 0) {
+                this.stats.setNumShrinks(this.stats.getNumShrinks() + 1);
+            }
             this.isBeingArbitrated.set(false);
-            this.arbitrationLock.unlock();
             this.condition.signalAll();
+            this.memoryActionLock.unlock();
         }
     }
 
@@ -106,9 +117,9 @@ public long tryToReclaimLocalMemory(long neededBytes) {
     public Object requireMemory(long bytes) {
         try {
             // use lock to ensure the atomicity of the two-step operation
-            this.arbitrationLock.lock();
-            long realBytes = requestMemoryInternal(bytes);
-            return tryToAcquireMemoryInternal(realBytes);
+            this.memoryActionLock.lock();
+            long ignoredRealAllocatedBytes = requestMemoryInternal(bytes);
+            return tryToAcquireMemoryInternal(bytes);
         } catch (QueryOutOfMemoryException e) {
             // Abort this query
             LOG.warn("[{}] detected an OOM exception when request memory, will ABORT this " +
@@ -117,7 +128,7 @@ public Object requireMemory(long bytes) {
             findRootQueryPool().releaseSelf(String.format(e.getMessage()));
             return null;
         } finally {
-            this.arbitrationLock.unlock();
+            this.memoryActionLock.unlock();
         }
     }
 

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

@@ -17,8 +17,11 @@
 
 package org.apache.hugegraph.memory.pool.impl;
 
+import static org.apache.hugegraph.memory.MemoryManager.DELIMINATOR;
+
 import org.apache.hugegraph.memory.MemoryManager;
 import org.apache.hugegraph.memory.pool.AbstractMemoryPool;
+import org.apache.hugegraph.memory.pool.MemoryPool;
 import org.apache.hugegraph.memory.util.MemoryManageUtils;
 import org.apache.hugegraph.util.Bytes;
 import org.slf4j.Logger;
@@ -27,6 +30,7 @@
 public class QueryMemoryPool extends AbstractMemoryPool {
 
     private static final Logger LOG = LoggerFactory.getLogger(QueryMemoryPool.class);
+    private static final String TASK_MEMORY_POOL_NAME_PREFIX = "TaskMemoryPool";
     // TODO: read from conf
     private static final long QUERY_POOL_MAX_CAPACITY = Bytes.MB * 100;
 
@@ -35,6 +39,18 @@ public QueryMemoryPool(String poolName, MemoryManager memoryManager) {
         this.stats.setMaxCapacity(QUERY_POOL_MAX_CAPACITY);
     }
 
+    @Override
+    public MemoryPool addChildPool() {
+        int count = this.children.size();
+        String poolName =
+                TASK_MEMORY_POOL_NAME_PREFIX + DELIMINATOR + count + DELIMINATOR +
+                System.currentTimeMillis();
+        MemoryPool taskMemoryPool = new TaskMemoryPool(this, poolName, this.memoryManager);
+        this.children.add(taskMemoryPool);
+        LOG.info("QueryPool-{} added task memory pool {}", this, taskMemoryPool);
+        return taskMemoryPool;
+    }
+
     @Override
     public long requestMemoryInternal(long bytes) {
         if (this.isClosed) {
@@ -83,8 +99,10 @@ private long tryToExpandSelfCapacity(long size) {
     private long requestMemoryThroughArbitration(long bytes) {
         LOG.info("[{}] try to request memory from manager through arbitration: size={}", this,
                  bytes);
-        this.stats.setNumExpands(this.stats.getNumExpands() + 1);
         long reclaimedBytes = this.memoryManager.triggerLocalArbitration(this, bytes);
+        if (reclaimedBytes > 0) {
+            this.stats.setNumExpands(this.stats.getNumExpands() + 1);
+        }
         // 1. if arbitrate successes, update stats and return success
         if (reclaimedBytes - bytes >= 0) {
             // here we don't update capacity & reserved & allocated, because memory is

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

@@ -17,7 +17,10 @@
 
 package org.apache.hugegraph.memory.pool.impl;
 
+import static org.apache.hugegraph.memory.MemoryManager.DELIMINATOR;
+
 import org.apache.hugegraph.memory.MemoryManager;
+import org.apache.hugegraph.memory.allocator.NettyMemoryAllocator;
 import org.apache.hugegraph.memory.pool.AbstractMemoryPool;
 import org.apache.hugegraph.memory.pool.MemoryPool;
 import org.apache.hugegraph.memory.util.QueryOutOfMemoryException;
@@ -27,11 +30,33 @@
 public class TaskMemoryPool extends AbstractMemoryPool {
 
     private static final Logger LOG = LoggerFactory.getLogger(TaskMemoryPool.class);
+    private static final String OPERATOR_MEMORY_POOL_NAME_PREFIX = "OperatorMemoryPool";
 
     public TaskMemoryPool(MemoryPool parent, String poolName, MemoryManager memoryManager) {
         super(parent, poolName, memoryManager);
     }
 
+    @Override
+    public synchronized void releaseSelf(String reason) {
+        this.memoryManager.removeCorrespondingTaskMemoryPool(Thread.currentThread().getName());
+        super.releaseSelf(reason);
+    }
+
+    @Override
+    public MemoryPool addChildPool() {
+        int count = this.children.size();
+        String poolName =
+                OPERATOR_MEMORY_POOL_NAME_PREFIX + DELIMINATOR + count + DELIMINATOR +
+                System.currentTimeMillis();
+        MemoryPool operatorPool =
+                new OperatorMemoryPool(this, poolName,
+                                       new NettyMemoryAllocator(this.memoryManager),
+                                       this.memoryManager);
+        this.children.add(operatorPool);
+        LOG.info("TaskPool-{} added operator memory pool {}", this, operatorPool);
+        return operatorPool;
+    }
+
     @Override
     public long requestMemoryInternal(long bytes) throws QueryOutOfMemoryException {
         if (this.isClosed) {
@@ -45,6 +70,7 @@ public long requestMemoryInternal(long bytes) throws QueryOutOfMemoryException {
             long parentRes = getParentPool().requestMemoryInternal(bytes);
             if (parentRes > 0) {
                 this.stats.setAllocatedBytes(this.stats.getAllocatedBytes() + parentRes);
+                this.stats.setNumExpands(this.stats.getNumExpands() + 1);
             }
             return parentRes;
         } catch (InterruptedException e) {

hugegraph-server/hugegraph-test/src/main/java/org/apache/hugegraph/core/memory/MemoryConsumerTest.java

@@ -1,2 +1,22 @@
-package org.apache.hugegraph.core.memory;public class MemoryConsumerTest {
+/*
+ * 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.hugegraph.core.memory;
+
+public class MemoryConsumerTest extends MemoryManageTest {
+
 }