From ff96a8c534a5b999294aec71979256380edad479 Mon Sep 17 00:00:00 2001
From: jimingquan <mingquan.ji@vesoft.com>
Date: Thu, 23 Mar 2023 10:24:18 +0800
Subject: [PATCH] add memory check

---
 src/graph/executor/StorageAccessExecutor.cpp |  1 +
 src/graph/executor/algo/AllPathsExecutor.cpp | 12 ++++++++++--
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/src/graph/executor/StorageAccessExecutor.cpp b/src/graph/executor/StorageAccessExecutor.cpp
index 31ed39150d8..a1d691dc3ea 100644
--- a/src/graph/executor/StorageAccessExecutor.cpp
+++ b/src/graph/executor/StorageAccessExecutor.cpp
@@ -178,6 +178,7 @@ folly::Future<std::vector<Value>> StorageAccessExecutor::getProps(
           param, std::move(vertices), vertexPropPtr, nullptr, nullptr, false, {}, -1, nullptr)
       .via(runner())
       .thenValue([this](PropRpcResponse &&resp) {
+        memory::MemoryCheckGuard guard;
         addStats(resp);
         return handlePropResp(std::move(resp));
       });
diff --git a/src/graph/executor/algo/AllPathsExecutor.cpp b/src/graph/executor/algo/AllPathsExecutor.cpp
index 8b994096822..373d7f8b00a 100644
--- a/src/graph/executor/algo/AllPathsExecutor.cpp
+++ b/src/graph/executor/algo/AllPathsExecutor.cpp
@@ -405,8 +405,9 @@ folly::Future<std::vector<Row>> AllPathsExecutor::doBuildPath(
       }));
     }
   }
-  return folly::collect(futures).via(runner()).thenValue(
-      [pathPtr = std::move(currentPathPtr)](std::vector<std::vector<Row>>&& paths) {
+  return folly::collect(futures)
+      .via(runner())
+      .thenValue([pathPtr = std::move(currentPathPtr)](std::vector<std::vector<Row>>&& paths) {
         memory::MemoryCheckGuard guard;
         std::vector<Row> result = std::move(*pathPtr);
         for (auto& path : paths) {
@@ -418,6 +419,13 @@ folly::Future<std::vector<Row>> AllPathsExecutor::doBuildPath(
                         std::make_move_iterator(path.end()));
         }
         return result;
+      })
+      .thenError(folly::tag_t<std::bad_alloc>{},
+                 [](const std::bad_alloc&) {
+                   return folly::makeFuture<Status>(Executor::memoryExceededStatus());
+                 })
+      .thenError(folly::tag_t<std::exception>{}, [](const std::exception& e) {
+        return folly::makeFuture<Status>(std::runtime_error(e.what()));
       });
 }