Restrict the max depth of plan tree in optimizer

src/graph/optimizer/Optimizer.cpp

@@ -22,6 +22,7 @@ using nebula::graph::Select;
 using nebula::graph::SingleDependencyNode;
 
 DEFINE_bool(enable_optimizer_property_pruner_rule, true, "");
+DEFINE_uint64(max_plan_depth, 512, "The max depth of plan tree");
 
 namespace nebula {
 namespace opt {
@@ -36,6 +37,7 @@ StatusOr<const PlanNode *> Optimizer::findBestPlan(QueryContext *qctx) {
  auto root = qctx->plan()->root();
  auto spaceID = qctx->rctx()->session()->space().id;
 
+ NG_RETURN_IF_ERROR(checkPlanDepth(root));
  auto ret = prepare(optCtx.get(), root);
  NG_RETURN_IF_ERROR(ret);
  auto rootGroup = std::move(ret).value();
@@ -158,15 +160,9 @@ bool findArgumentRefPlanNodeInPath(const std::vector<const PlanNode *> &path, Pl
 }
 
 Status rewriteArgumentInputVarInternal(PlanNode *root,
- uint16_t stackDepth,
+
  bool &hasArgument,
  std::vector<const PlanNode *> &path) {
- const uint16_t kMaxStackDepth = 512u;
- if (stackDepth > kMaxStackDepth) {
- return Status::Error("The depth of plan tree has exceeded the max %u level", kMaxStackDepth);
- }
- stackDepth++;
-
  if (!root) return Status::OK();
 
  path.push_back(root);
@@ -185,15 +181,15 @@ Status rewriteArgumentInputVarInternal(PlanNode *root,
  }
  case 1: {
  auto *dep = const_cast<PlanNode *>(root->dep());
- NG_RETURN_IF_ERROR(rewriteArgumentInputVarInternal(dep, stackDepth, hasArgument, path));
+ NG_RETURN_IF_ERROR(rewriteArgumentInputVarInternal(dep, hasArgument, path));
  break;
  }
  case 2: {
  auto *bpn = static_cast<BinaryInputNode *>(root);
  auto *left = const_cast<PlanNode *>(bpn->left());
- NG_RETURN_IF_ERROR(rewriteArgumentInputVarInternal(left, stackDepth, hasArgument, path));
+ NG_RETURN_IF_ERROR(rewriteArgumentInputVarInternal(left, hasArgument, path));
  auto *right = const_cast<PlanNode *>(bpn->right());
- NG_RETURN_IF_ERROR(rewriteArgumentInputVarInternal(right, stackDepth, hasArgument, path));
+ NG_RETURN_IF_ERROR(rewriteArgumentInputVarInternal(right, hasArgument, path));
  break;
  }
  default: {
@@ -220,7 +216,30 @@ Status rewriteArgumentInputVarInternal(PlanNode *root,
 Status Optimizer::rewriteArgumentInputVar(PlanNode *root) {
  bool hasArgument = false;
  std::vector<const PlanNode *> path;
- return rewriteArgumentInputVarInternal(root, 0, hasArgument, path);
+ return rewriteArgumentInputVarInternal(root, hasArgument, path);
+}
+
+Status Optimizer::checkPlanDepth(const PlanNode *root) const {
+ std::queue<const PlanNode *> queue;
+ queue.push(root);
+ size_t depth = 0;
+ while (!queue.empty()) {
+ size_t size = queue.size();
+ for (size_t i = 0; i < size; ++i) {
+ const PlanNode *node = queue.front();
+ queue.pop();
+ for (size_t j = 0; j < node->numDeps(); j++) {
+ queue.push(node->dep(j));
+ }
+ }
+ ++depth;
+ if (depth > FLAGS_max_plan_depth) {
+ return Status::Error("The depth of plan tree has exceeded the max %lu level",
+ FLAGS_max_plan_depth);
+ }
+ }
+
+ return Status::OK();
 }
 
 } // namespace opt

src/graph/optimizer/Optimizer.h

@@ -49,6 +49,8 @@ class Optimizer final {
 
  static Status rewriteArgumentInputVar(graph::PlanNode *root);
 
+ Status checkPlanDepth(const graph::PlanNode *root) const;
+
  static constexpr int8_t kMaxIterationRound = 5;
 
  std::vector<const RuleSet *> ruleSets_;