Restrict the max depth of plan tree in optimizer (#5050)

vesoft-inc · Dec 14, 2022 · 0937dbd · 0937dbd
1 parent 0dce639
commit 0937dbd
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 11 deletions.
diff --git a/src/graph/optimizer/Optimizer.cpp b/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

diff --git a/src/graph/optimizer/Optimizer.h b/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_;