feat(server): support MemoryManagement for graph query framework

hugegraph-server/hugegraph-core/pom.xml

@@ -96,6 +96,11 @@
             <groupId>org.apache.tinkerpop</groupId>
             <artifactId>gremlin-driver</artifactId>
         </dependency>
+        <dependency>
+            <groupId>org.apache.fury</groupId>
+            <artifactId>fury-core</artifactId>
+            <version>0.9.0-SNAPSHOT</version>
+        </dependency>
 
         <!-- jraft -->
         <dependency>

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/cache/CachedBackendStore.java

@@ -215,16 +215,21 @@ public Number queryNumber(Query query) {
     /**
      * Query as an Id for cache
      */
-    static class QueryId implements Id {
+    public static class QueryId implements Id {
 
-        private String query;
-        private int hashCode;
+        protected String query;
+        protected int hashCode;
 
         public QueryId(Query q) {
             this.query = q.toString();
             this.hashCode = q.hashCode();
         }
 
+        public QueryId(String query, int hashCode) {
+            this.query = query;
+            this.hashCode = hashCode;
+        }
+
         @Override
         public IdType type() {
             return IdType.UNKNOWN;

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/id/EdgeId.java

@@ -32,8 +32,8 @@
  * > sortKeys > target-vertex-id }
  * NOTE:
  * <p>1. for edges with edgeLabelType = NORMAL: edgeLabelId = parentEdgeLabelId = subEdgeLabelId;
- *       for edges with edgeLabelType = PARENT: edgeLabelId = subEdgeLabelId, parentEdgeLabelId =
- *       edgeLabelId.fatherId
+ * for edges with edgeLabelType = PARENT: edgeLabelId = subEdgeLabelId, parentEdgeLabelId =
+ * edgeLabelId.fatherId
  * <p>2.if we use `entry.type()` which is IN or OUT as a part of id,
  * an edge's id will be different due to different directions (belongs
  * to 2 owner vertex)
@@ -49,15 +49,14 @@ public class EdgeId implements Id {
             HugeKeys.OTHER_VERTEX
     };
 
-    private final Id ownerVertexId;
-    private final Directions direction;
-    private final Id edgeLabelId;
-    private final Id subLabelId;
-    private final String sortValues;
-    private final Id otherVertexId;
-
-    private final boolean directed;
-    private String cache;
+    protected final Id ownerVertexId;
+    protected final Id edgeLabelId;
+    protected final Id subLabelId;
+    protected final Id otherVertexId;
+    protected final Directions direction;
+    protected final boolean directed;
+    protected String sortValues;
+    protected String cache;
 
     public EdgeId(HugeVertex ownerVertex, Directions direction,
                   Id edgeLabelId, Id subLabelId, String sortValues, HugeVertex otherVertex) {

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/id/IdGenerator.java

@@ -119,15 +119,15 @@ public static IdType idType(Id id) {
         return IdType.UNKNOWN;
     }
 
-    private static int compareType(Id id1, Id id2) {
+    public static int compareType(Id id1, Id id2) {
         return idType(id1).ordinal() - idType(id2).ordinal();
     }
 
     /****************************** id defines ******************************/
 
-    public static final class StringId implements Id {
+    public static class StringId implements Id {
 
-        private final String id;
+        protected String id;
 
         public StringId(String id) {
             E.checkArgument(!id.isEmpty(), "The id can't be empty");
@@ -196,11 +196,11 @@ public String toString() {
         }
     }
 
-    public static final class LongId extends Number implements Id {
+    public static class LongId extends Number implements Id {
 
         private static final long serialVersionUID = -7732461469037400190L;
 
-        private final long id;
+        protected Long id;
 
         public LongId(long id) {
             this.id = id;
@@ -270,7 +270,7 @@ public String toString() {
 
         @Override
         public int intValue() {
-            return (int) this.id;
+            return this.id.intValue();
         }
 
         @Override
@@ -289,9 +289,9 @@ public double doubleValue() {
         }
     }
 
-    public static final class UuidId implements Id {
+    public static class UuidId implements Id {
 
-        private final UUID uuid;
+        protected UUID uuid;
 
         public UuidId(String string) {
             this(StringEncoding.uuid(string));
@@ -379,9 +379,9 @@ public String toString() {
     /**
      * This class is just used by backend store for wrapper object as Id
      */
-    public static final class ObjectId implements Id {
+    public static class ObjectId implements Id {
 
-        private final Object object;
+        protected Object object;
 
         public ObjectId(Object object) {
             E.checkNotNull(object, "object");

hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/backend/serializer/BinaryBackendEntry.java

@@ -49,7 +49,8 @@ public BinaryBackendEntry(HugeType type, byte[] bytes, boolean enablePartition)
     }
 
     // FIXME: `enablePartition` is unused here
-    public BinaryBackendEntry(HugeType type, byte[] bytes, boolean enablePartition, boolean isOlap) {
+    public BinaryBackendEntry(HugeType type, byte[] bytes, boolean enablePartition,
+                              boolean isOlap) {
         this(type, BytesBuffer.wrap(bytes).parseOlapId(type, isOlap));
     }
 
@@ -207,10 +208,10 @@ public int hashCode() {
         return this.id().hashCode() ^ this.columns.size();
     }
 
-    public static final class BinaryId implements Id {
+    public static class BinaryId implements Id {
 
-        private final byte[] bytes;
-        private final Id id;
+        protected byte[] bytes;
+        protected Id id;
 
         public BinaryId(byte[] bytes, Id id) {
             this.bytes = bytes;

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

@@ -0,0 +1,180 @@
+/*
+ * 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.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;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicLong;
+
+import org.apache.hugegraph.memory.arbitrator.MemoryArbitrator;
+import org.apache.hugegraph.memory.arbitrator.MemoryArbitratorImpl;
+import org.apache.hugegraph.memory.pool.MemoryPool;
+import org.apache.hugegraph.memory.pool.impl.QueryMemoryPool;
+import org.apache.hugegraph.util.Bytes;
+import org.apache.hugegraph.util.ExecutorUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+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";
+    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;
+
+    private MemoryManager() {
+        this.memoryArbitrator = new MemoryArbitratorImpl(this);
+        this.arbitrateExecutor = ExecutorUtil.newFixedThreadPool(ARBITRATE_MEMORY_THREAD_NUM,
+                                                                 ARBITRATE_MEMORY_POOL_NAME);
+    }
+
+    public MemoryPool addQueryMemoryPool() {
+        int count = queryMemoryPools.size();
+        String poolName =
+                QUERY_MEMORY_POOL_NAME_PREFIX + DELIMINATOR + count + DELIMINATOR +
+                System.currentTimeMillis();
+        MemoryPool queryPool = new QueryMemoryPool(poolName, this);
+        queryMemoryPools.add(queryPool);
+        LOG.info("Manager added query memory pool {}", queryPool);
+        return queryPool;
+    }
+
+    public void gcQueryMemoryPool(MemoryPool pool) {
+        LOG.info("Manager gc query memory pool {}", pool);
+        queryMemoryPools.remove(pool);
+        long reclaimedMemory = pool.getAllocatedBytes();
+        pool.releaseSelf(String.format("GC query memory pool %s", pool), false);
+        currentAvailableMemoryInBytes.addAndGet(reclaimedMemory);
+    }
+
+    public long triggerLocalArbitration(MemoryPool targetPool, long neededBytes) {
+        LOG.info("LocalArbitration triggered by {}: needed bytes={}", targetPool, neededBytes);
+        Future<Long> future =
+                arbitrateExecutor.submit(
+                        () -> memoryArbitrator.reclaimLocally(targetPool, neededBytes));
+        try {
+            return future.get(MemoryArbitrator.MAX_WAIT_TIME_FOR_LOCAL_RECLAIM,
+                              TimeUnit.MILLISECONDS);
+        } catch (TimeoutException e) {
+            LOG.warn("MemoryManager: arbitration locally for {} timed out", targetPool, e);
+        } catch (InterruptedException | ExecutionException e) {
+            LOG.error("MemoryManager: arbitration locally for {} interrupted or failed",
+                      targetPool,
+                      e);
+        }
+        return 0;
+    }
+
+    public long triggerGlobalArbitration(MemoryPool requestPool, long neededBytes) {
+        LOG.info("GlobalArbitration triggered by {}: needed bytes={}", requestPool, neededBytes);
+        Future<Long> future =
+                arbitrateExecutor.submit(
+                        () -> memoryArbitrator.reclaimGlobally(requestPool, neededBytes));
+        try {
+            return future.get(MemoryArbitrator.MAX_WAIT_TIME_FOR_GLOBAL_RECLAIM,
+                              TimeUnit.MILLISECONDS);
+        } catch (TimeoutException e) {
+            LOG.warn("MemoryManager: arbitration globally for {} timed out", requestPool, e);
+        } catch (InterruptedException | ExecutionException e) {
+            LOG.error("MemoryManager: arbitration globally for {} interrupted or failed",
+                      requestPool, e);
+        }
+        return 0;
+    }
+
+    public synchronized long handleRequestFromQueryPool(long size) {
+        // 1. check whole memory capacity.
+        if (currentAvailableMemoryInBytes.get() < size) {
+            LOG.info("There isn't enough memory for query pool to expand itself: " +
+                     "requestSize={}, remainingCapacity={}", size,
+                     currentAvailableMemoryInBytes.get());
+            return -1;
+        }
+        currentAvailableMemoryInBytes.addAndGet(-size);
+        LOG.info("Expand query pool successfully: " +
+                 "requestSize={}, afterThisExpandingRemainingCapacity={}", size,
+                 currentAvailableMemoryInBytes.get());
+        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);
+    }
+
+    public void consumeAvailableMemory(long size) {
+        currentAvailableMemoryInBytes.addAndGet(-size);
+    }
+
+    public AtomicLong getCurrentOnHeapAllocatedMemoryInBytes() {
+        return currentOnHeapAllocatedMemoryInBytes;
+    }
+
+    public AtomicLong getCurrentOffHeapAllocatedMemoryInBytes() {
+        return currentOffHeapAllocatedMemoryInBytes;
+    }
+
+    private static class MemoryManagerHolder {
+
+        private static final MemoryManager INSTANCE = new MemoryManager();
+
+        private MemoryManagerHolder() {
+            // empty constructor
+        }
+    }
+
+    public static MemoryManager getInstance() {
+        return MemoryManagerHolder.INSTANCE;
+    }
+}

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

@@ -0,0 +1,29 @@
+/*
+ * 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.memory.allocator;
+
+public interface MemoryAllocator {
+
+    Object tryToAllocate(long size);
+
+    Object forceAllocate(long size);
+
+    void returnMemoryToManager(long size);
+
+    void releaseMemoryBlock(Object memoryBlock);
+}