wip

apache · Dec 10, 2024 · a402ddb · a402ddb
1 parent 29b7845
commit a402ddb
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/...graph-server/hugegraph-api/src/main/java/org/apache/hugegraph/api/traversers/KoutAPI.java b/...graph-server/hugegraph-api/src/main/java/org/apache/hugegraph/api/traversers/KoutAPI.java
@@ -98,7 +98,10 @@ public String get(@Context GraphManager manager,
                   nearest, maxDegree, capacity, limit);
         MemoryPool queryPool = MemoryManager.getInstance().addQueryMemoryPool();
         MemoryPool currentTaskPool = queryPool.addChildPool("kout-main-task");
-        MemoryManager.getInstance().bindCorrespondingTaskMemoryPool(Thread.currentThread().getName(), (TaskMemoryPool) currentTaskPool);
+        MemoryManager.getInstance()
+                     .bindCorrespondingTaskMemoryPool(Thread.currentThread().getName(),
+                                                      (TaskMemoryPool) currentTaskPool);
+        MemoryPool currentOperationPool = currentTaskPool.addChildPool("kout-main-operation");
 
         ApiMeasurer measure = new ApiMeasurer();
 

diff --git a/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/util/Consumers.java b/hugegraph-server/hugegraph-core/src/main/java/org/apache/hugegraph/util/Consumers.java
@@ -35,6 +35,9 @@
 
 import org.apache.hugegraph.HugeException;
 import org.apache.hugegraph.config.CoreOptions;
+import org.apache.hugegraph.memory.MemoryManager;
+import org.apache.hugegraph.memory.pool.MemoryPool;
+import org.apache.hugegraph.memory.pool.impl.TaskMemoryPool;
 import org.apache.hugegraph.task.TaskManager.ContextCallable;
 import org.slf4j.Logger;
 
@@ -105,12 +108,23 @@ public void start(String name) {
                  this.workers, name, this.queueSize);
         for (int i = 0; i < this.workers; i++) {
             this.runningFutures.add(
-                    this.executor.submit(new ContextCallable<>(this::runAndDone)));
+                    this.executor.submit(
+                            new ContextCallable<>(() -> this.runAndDone(MemoryManager.getInstance()
+                                                                                     .getCorrespondingTaskMemoryPool(
+                                                                                             Thread.currentThread()
+                                                                                                   .getName())
+                                                                                     .findRootQueryPool()))));
         }
     }
 
-    private Void runAndDone() {
+    private Void runAndDone(MemoryPool queryPool) {
         try {
+            MemoryPool currentTaskPool = queryPool.addChildPool("kout-consume-task");
+            MemoryManager.getInstance()
+                         .bindCorrespondingTaskMemoryPool(Thread.currentThread().getName(),
+                                                          (TaskMemoryPool) currentTaskPool);
+            MemoryPool currentOperationPool =
+                    currentTaskPool.addChildPool("kout-consume-operation");
             this.run();
         } catch (Throwable e) {
             if (e instanceof StopExecution) {