add memory check

vesoft-inc · Mar 23, 2023 · ff96a8c · ff96a8c
1 parent bd23161
commit ff96a8c
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git 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
@@ -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()));
  });
 }