From a402ddb82cfb9be1643b9e35b095766c62605ada Mon Sep 17 00:00:00 2001
From: Peng Junzhi <201250214@smail.nju.edu.cn>
Date: Tue, 10 Dec 2024 11:26:56 +0800
Subject: [PATCH] wip

---
 .../hugegraph/api/traversers/KoutAPI.java      |  5 ++++-
 .../org/apache/hugegraph/util/Consumers.java   | 18 ++++++++++++++++--
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/api/traversers/KoutAPI.java b/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/api/traversers/KoutAPI.java
index cc62992245..7b09121789 100644
--- a/hugegraph-server/hugegraph-api/src/main/java/org/apache/hugegraph/api/traversers/KoutAPI.java
+++ b/hugegraph-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
index daa54ee958..6d63ac1cf1 100644
--- 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) {