1

xinyiZzz · May 23, 2024 · 54818c9 · 54818c9
1 parent e4f41db
commit 54818c9
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 0 deletions.
diff --git a/be/src/common/config.cpp b/be/src/common/config.cpp
@@ -131,6 +131,8 @@ DEFINE_mString(process_full_gc_size, "20%");
 // If false, cancel query when the memory used exceeds exec_mem_limit, same as before.
 DEFINE_mBool(enable_query_memory_overcommit, "true");
 
+DEFINE_mBool(enable_query_enforce_exec, "false");
+
 DEFINE_mBool(disable_memory_gc, "false");
 
 DEFINE_mInt64(large_memory_check_bytes, "2147483648");

diff --git a/be/src/common/config.h b/be/src/common/config.h
@@ -171,6 +171,8 @@ DECLARE_mString(process_full_gc_size);
 // used memory and the exec_mem_limit will be canceled.
 // If false, cancel query when the memory used exceeds exec_mem_limit, same as before.
 DECLARE_mBool(enable_query_memory_overcommit);
+// Default is false, BE process exceed memory limit will reject new Query execution.
+DECLARE_mBool(enable_query_enforce_exec);
 //waibibabu
 // gc will release cache, cancel task, and task will wait for gc to release memory,
 // default gc strategy is conservative, if you want to exclude the interference of gc, let it be true

diff --git a/be/src/runtime/fragment_mgr.cpp b/be/src/runtime/fragment_mgr.cpp
@@ -720,6 +720,16 @@ Status FragmentMgr::exec_plan_fragment(const TExecPlanFragmentParams& params,
     VLOG_ROW << "query options is "
              << apache::thrift::ThriftDebugString(params.query_options).c_str();
     const TUniqueId& fragment_instance_id = params.params.fragment_instance_id;
+    int64_t query_reserved_memory = 0; // TODO
+    if (!config::enable_query_enforce_exec &&
+        doris::GlobalMemoryArbitrator::is_exceed_hard_mem_limit(query_reserved_memory)) {
+        std::string errmsg =
+                fmt::format(" try exec_plan_fragment failed, fragmentInstanceId {}, but {}",
+                            print_id(fragment_instance_id),
+                            doris::GlobalMemoryArbitrator::process_limit_exceeded_errmsg_str());
+        LOG(INFO) << errmsg;
+        return Status::MemoryLimitExceeded(errmsg);
+    }
     {
         std::lock_guard<std::mutex> lock(_lock);
         auto iter = _fragment_instance_map.find(fragment_instance_id);
@@ -839,6 +849,15 @@ Status FragmentMgr::exec_plan_fragment(const TPipelineFragmentParams& params,
     VLOG_ROW << "query: " << print_id(params.query_id) << "query options is "
              << apache::thrift::ThriftDebugString(params.query_options).c_str();
 
+    int64_t query_reserved_memory = 0; // TODO
+    if (!config::enable_query_enforce_exec &&
+        doris::GlobalMemoryArbitrator::is_exceed_hard_mem_limit(query_reserved_memory)) {
+        std::string errmsg = fmt::format(
+                " try exec_plan_fragment failed, queryId {}, {}", print_id(params.query_id),
+                doris::GlobalMemoryArbitrator::process_limit_exceeded_errmsg_str());
+        LOG(INFO) << errmsg;
+        return Status::MemoryLimitExceeded(errmsg);
+    }
     std::shared_ptr<QueryContext> query_ctx;
     RETURN_IF_ERROR(_get_query_ctx(params, params.query_id, true, query_ctx));
     SCOPED_ATTACH_TASK_WITH_ID(query_ctx->query_mem_tracker, params.query_id);