fix memory conservation bug

Author: Pengzna

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

@@ -38,7 +38,7 @@ public long reclaimLocally(MemoryPool queryPool, long neededBytes, MemoryPool re
         long startTime = System.currentTimeMillis();
         long res = queryPool.tryToReclaimLocalMemory(neededBytes, requestingPool);
         LOG.info("[{}] reclaim local memory: {} bytes, took {} ms",
-                 Thread.currentThread().getName(),
+                 queryPool,
                  res,
                  System.currentTimeMillis() - startTime);
         return res;
@@ -65,7 +65,7 @@ public long reclaimGlobally(MemoryPool queryPool, long neededBytes) {
             }
         }
         LOG.info("[{}] reclaim global memory: {} bytes, took {} ms",
-                 Thread.currentThread().getName(),
+                 queryPool,
                  totalReclaimedBytes,
                  System.currentTimeMillis() - startTime);
         return totalReclaimedBytes;

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

@@ -94,6 +94,9 @@ private long reclaimChildren(long neededBytes, MemoryPool requestingPool) {
         long totalReclaimedBytes = 0;
         long currentNeededBytes = neededBytes;
         for (MemoryPool child : this.children) {
+            if (child.equals(requestingPool)) {
+                continue;
+            }
             long reclaimedMemory =
                     child.tryToReclaimLocalMemory(currentNeededBytes, requestingPool);
             if (reclaimedMemory > 0) {
@@ -217,16 +220,15 @@ public MemoryPool getParentPool() {
         return this.parent;
     }
 
-    @Override
-    public String getName() {
-        return this.stats.getMemoryPoolName();
-    }
-
     @Override
     public String toString() {
         return this.getName();
     }
 
+    public String getName() {
+        return this.stats.getMemoryPoolName();
+    }
+
     @Override
     public MemoryPool findRootQueryPool() {
         if (this.parent == null) {

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

@@ -92,8 +92,6 @@ public interface MemoryPool {
 
     long getMaxCapacityBytes();
 
-    String getName();
-
     MemoryPool getParentPool();
 
     MemoryPool findRootQueryPool();

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

@@ -189,16 +189,23 @@ public long requestMemoryInternal(long size, MemoryPool requestingPool) throws
             // 3. call father
             long fatherRes =
                     getParentPool().requestMemoryInternal(neededMemorySize, requestingPool);
-            if (fatherRes < 0) {
+            if (fatherRes <= 0) {
                 LOG.error("[{}] requestMemory failed because of OOM, request size={}", this,
                           size);
+                // if parentRes <= 0, indicating we don't get enough memory bytes.
+                // But we still need to allocate these memory bytes to operatorPool to ensure
+                // memory is conserved when parentRes < 0
+                if (fatherRes < 0) {
+                    this.stats.setAllocatedBytes(this.stats.getAllocatedBytes() - fatherRes);
+                    this.stats.setNumExpands(this.stats.getNumExpands() + 1);
+                }
                 this.stats.setNumAborts(this.stats.getNumAborts() + 1);
                 throw new OutOfMemoryException(String.format("%s requestMemory failed " +
                                                              "because of OOM, request " +
                                                              "size=%s", this, size));
             }
             // 4. update stats
-            this.stats.setAllocatedBytes(this.stats.getAllocatedBytes() + neededMemorySize);
+            this.stats.setAllocatedBytes(this.stats.getAllocatedBytes() + fatherRes);
             this.stats.setNumExpands(this.stats.getNumExpands() + 1);
             LOG.debug("[{}] requestMemory success: requestedMemorySize={}", this, fatherRes);
             return fatherRes;

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

@@ -105,7 +105,7 @@ public long requestMemoryInternal(long bytes, MemoryPool requestingPool) {
                 }
             }
             if (requestedMemoryFromManager > 0) {
-                this.memoryManager.consumeAvailableMemory(bytes);
+                this.memoryManager.consumeAvailableMemory(requestedMemoryFromManager);
             }
             return requestedMemoryFromManager == 0 ? requestedMemoryFromArbitration :
                    requestedMemoryFromManager;
@@ -134,30 +134,30 @@ private long requestMemoryThroughArbitration(long bytes, MemoryPool requestingPo
                  this.getSnapShot());
         long reclaimedBytes =
                 this.memoryManager.triggerLocalArbitration(this, bytes, requestingPool);
-        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
             // reclaimed from queryPool itself.
-            return bytes;
+            return reclaimedBytes;
         } else {
             // 2. if still not enough, try to reclaim globally
             long globalArbitrationNeededBytes = bytes - reclaimedBytes;
             long globalReclaimedBytes = this.memoryManager.triggerGlobalArbitration(this,
                                                                                     globalArbitrationNeededBytes);
             reclaimedBytes += globalReclaimedBytes;
-            // 3. if memory is enough, update stats and return success
-            if (reclaimedBytes - bytes >= 0) {
-                // add capacity
+            // add capacity whether arbitration failed or not.
+            // NOTE: only bytes gained from global arbitration can be added to stats.
+            if (globalReclaimedBytes != 0) {
                 this.stats.setMaxCapacity(this.stats.getMaxCapacity() + globalReclaimedBytes);
-                this.stats.setAllocatedBytes(this.stats.getAllocatedBytes() + bytes);
+                this.stats.setAllocatedBytes(this.stats.getAllocatedBytes() + globalReclaimedBytes);
                 this.stats.setNumExpands(this.stats.getNumExpands() + 1);
-                return bytes;
+            }
+            // 3. if memory is enough, update stats and return success
+            if (reclaimedBytes - bytes >= 0) {
+                return reclaimedBytes;
             }
         }
-        // 4. if arbitrate fails, return -1, indicating that request failed.
-        return -1;
+        // 4. if arbitrate fails, return -reclaimedBytes, indicating that request failed.
+        return -reclaimedBytes;
     }
 }

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

@@ -100,6 +100,12 @@ public long requestMemoryInternal(long bytes, MemoryPool requestingPool) throws
             if (parentRes > 0) {
                 this.stats.setAllocatedBytes(this.stats.getAllocatedBytes() + parentRes);
                 this.stats.setNumExpands(this.stats.getNumExpands() + 1);
+            } else if (parentRes < 0){
+                // if parentRes < 0, indicating we don't get enough memory bytes. But we still
+                // need to allocate these memory bytes to operatorPool to ensure memory is
+                // conserved.
+                this.stats.setAllocatedBytes(this.stats.getAllocatedBytes() - parentRes);
+                this.stats.setNumExpands(this.stats.getNumExpands() + 1);
             }
             return parentRes;
         } catch (InterruptedException e) {