ENABLE_FACTORY_CREATOR(DelayReleaseToken);

DelayReleaseToken(std::unique_ptr<ThreadPoolToken>&& token) { token_ = std::move(token); }

~DelayReleaseToken() override = default;

void run() override {}

std::unique_ptr<ThreadPoolToken> token_;

switch (queryType) {

case QuerySource::INTERNAL_FRONTEND:

return "INTERNAL_FRONTEND";

case QuerySource::STREAM_LOAD:

return "STREAM_LOAD";

case QuerySource::GROUP_COMMIT_LOAD:

return "EXTERNAL_QUERY";

case QuerySource::ROUTINE_LOAD:

return "ROUTINE_LOAD";

case QuerySource::EXTERNAL_CONNECTOR:

return "EXTERNAL_CONNECTOR";

default:

return "UNKNOWN";

}

const TQueryOptions& query_options,

TNetworkAddress coord_addr, bool is_nereids,

TNetworkAddress current_connect_fe,

QuerySource query_type) {

auto ctx = QueryContext::create_shared(query_id, exec_env, query_options, coord_addr,

is_nereids, current_connect_fe, query_type);

ctx->init_query_task_controller();

return ctx;

const TQueryOptions& query_options, TNetworkAddress coord_addr,

bool is_nereids, TNetworkAddress current_connect_fe,

QuerySource query_source)

: _timeout_second(-1),

_query_id(std::move(query_id)),

_exec_env(exec_env),

_is_nereids(is_nereids),

_query_options(query_options),

_query_source(query_source) {

_init_resource_context();

SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(query_mem_tracker());

_query_watcher.start();

_shared_hash_table_controller.reset(new vectorized::SharedHashTableController());

_execution_dependency =

pipeline::Dependency::create_unique(-1, -1, "ExecutionDependency", false);

_memory_sufficient_dependency =

pipeline::Dependency::create_unique(-1, -1, "MemorySufficientDependency", true);

_runtime_filter_mgr = std::make_unique<RuntimeFilterMgr>(

TUniqueId(), RuntimeFilterParamsContext::create(this), query_mem_tracker(), true);

_timeout_second = query_options.execution_timeout;

bool is_query_type_valid = query_options.query_type == TQueryType::SELECT ||

query_options.query_type == TQueryType::LOAD ||

query_options.query_type == TQueryType::EXTERNAL;

DCHECK_EQ(is_query_type_valid, true);

this->coord_addr = coord_addr;

// current_connect_fe is used for report query statistics

this->current_connect_fe = current_connect_fe;

// external query has no current_connect_fe

if (query_options.query_type != TQueryType::EXTERNAL) {

bool is_report_fe_addr_valid =

!this->current_connect_fe.hostname.empty() && this->current_connect_fe.port != 0;

DCHECK_EQ(is_report_fe_addr_valid, true);

}

clock_gettime(CLOCK_MONOTONIC, &this->_query_arrival_timestamp);

DorisMetrics::instance()->query_ctx_cnt->increment(1);

bool has_query_mem_limit = _query_options.__isset.mem_limit && (_query_options.mem_limit > 0);

int64_t bytes_limit = has_query_mem_limit ? _query_options.mem_limit : -1;

if (bytes_limit > MemInfo::mem_limit() || bytes_limit == -1) {

VLOG_NOTICE << "Query memory limit " << PrettyPrinter::print(bytes_limit, TUnit::BYTES)

<< " exceeds process memory limit of "

<< PrettyPrinter::print(MemInfo::mem_limit(), TUnit::BYTES)

<< " OR is -1. Using process memory limit instead.";

bytes_limit = MemInfo::mem_limit();

}

// If the query is a pure load task(streamload, routine load, group commit), then it should not use

// memlimit per query to limit their memory usage.

if (is_pure_load_task()) {

bytes_limit = MemInfo::mem_limit();

}

std::shared_ptr<MemTrackerLimiter> query_mem_tracker;

if (_query_options.query_type == TQueryType::SELECT) {

query_mem_tracker = MemTrackerLimiter::create_shared(

MemTrackerLimiter::Type::QUERY, fmt::format("Query#Id={}", print_id(_query_id)),

bytes_limit);

} else if (_query_options.query_type == TQueryType::LOAD) {

query_mem_tracker = MemTrackerLimiter::create_shared(

MemTrackerLimiter::Type::LOAD, fmt::format("Load#Id={}", print_id(_query_id)),

bytes_limit);

} else if (_query_options.query_type == TQueryType::EXTERNAL) { // spark/flink/etc..

query_mem_tracker = MemTrackerLimiter::create_shared(

MemTrackerLimiter::Type::QUERY, fmt::format("External#Id={}", print_id(_query_id)),

bytes_limit);

} else {

LOG(FATAL) << "__builtin_unreachable";

__builtin_unreachable();

}

if (_query_options.__isset.is_report_success && _query_options.is_report_success) {

query_mem_tracker->enable_print_log_usage();

}

query_mem_tracker->set_enable_reserve_memory(_query_options.__isset.enable_reserve_memory &&

_query_options.enable_reserve_memory);

_user_set_mem_limit = bytes_limit;

_adjusted_mem_limit = bytes_limit;

_resource_ctx->memory_context()->set_mem_tracker(query_mem_tracker);

_resource_ctx = ResourceContext::create_shared();

_resource_ctx->set_memory_context(QueryContext::QueryMemoryContext::create());

_init_query_mem_tracker();

_resource_ctx->set_task_controller(QueryContext::QueryTaskController::create(this));

_resource_ctx->task_controller()->set_task_id(_query_id);

_resource_ctx->task_controller()->set_fe_addr(current_connect_fe);

_resource_ctx->task_controller()->set_query_type(_query_options.query_type);

#ifndef BE_TEST

_exec_env->runtime_query_statistics_mgr()->register_resource_context(print_id(_query_id),

_resource_ctx);

#endif

set_execution_dependency_ready();

_exec_status.update(reason);

if (query_mem_tracker() && !reason.ok()) {

query_mem_tracker()->set_is_query_cancelled(!reason.ok());

}

if (sufficient) {

{

_memory_sufficient_dependency->set_ready();

std::lock_guard l(_paused_mutex);

_paused_reason = Status::OK();

}

} else {

_memory_sufficient_dependency->block();

++_paused_count;

}

if (!_exec_status.update(new_status)) {

return;

}

// Tasks should be always runnable.

_execution_dependency->set_always_ready();

_memory_sufficient_dependency->set_always_ready();

if ((new_status.is<ErrorCode::MEM_LIMIT_EXCEEDED>() ||

new_status.is<ErrorCode::MEM_ALLOC_FAILED>()) &&

_query_options.__isset.dump_heap_profile_when_mem_limit_exceeded &&

_query_options.dump_heap_profile_when_mem_limit_exceeded) {

// if query is cancelled because of query mem limit exceeded, dump heap profile

// at the time of cancellation can get the most accurate memory usage for problem analysis

auto wg = workload_group();

auto log_str = fmt::format(

"Query {} canceled because of memory limit exceeded, dumping memory "

"detail profiles. wg: {}. {}",

print_id(_query_id), wg ? wg->debug_string() : "null",

doris::ProcessProfile::instance()->memory_profile()->process_memory_detail_str());

LOG_LONG_STRING(INFO, log_str);

std::string dot = HeapProfiler::instance()->dump_heap_profile_to_dot();

if (!dot.empty()) {

dot += "\n-------------------------------------------------------\n";

dot += "Copy the text after `digraph` in the above output to "

"http://www.webgraphviz.com to generate a dot graph.\n"

"after start heap profiler, if there is no operation, will print `No nodes "

"to "

"print`."

"If there are many errors: `addr2line: Dwarf Error`,"

"or other FAQ, reference doc: "

"https://doris.apache.org/community/developer-guide/debug-tool/#4-qa\n";

auto log_str =

fmt::format("Query {}, dump heap profile to dot: {}", print_id(_query_id), dot);

LOG_LONG_STRING(INFO, log_str);

}

}

set_ready_to_execute(new_status);

cancel_all_pipeline_context(new_status, fragment_id);

std::vector<std::weak_ptr<pipeline::PipelineFragmentContext>> ctx_to_cancel;

{

std::lock_guard<std::mutex> lock(_pipeline_map_write_lock);

for (auto& [f_id, f_context] : _fragment_id_to_pipeline_ctx) {

if (fragment_id == f_id) {

continue;

}

ctx_to_cancel.push_back(f_context);

}

}

for (auto& f_context : ctx_to_cancel) {

if (auto pipeline_ctx = f_context.lock()) {

pipeline_ctx->cancel(reason);

}

}

std::vector<std::weak_ptr<pipeline::PipelineFragmentContext>> ctx_to_print;

fmt::memory_buffer debug_string_buffer;

size_t i = 0;

{

fmt::format_to(debug_string_buffer, "{} pipeline fragment contexts in query {}. \n",

_fragment_id_to_pipeline_ctx.size(), print_id(_query_id));

{

std::lock_guard<std::mutex> lock(_pipeline_map_write_lock);

for (auto& [f_id, f_context] : _fragment_id_to_pipeline_ctx) {

ctx_to_print.push_back(f_context);

}

}

for (auto& f_context : ctx_to_print) {

if (auto pipeline_ctx = f_context.lock()) {

auto elapsed = pipeline_ctx->elapsed_time() / 1000000000.0;

fmt::format_to(debug_string_buffer,

"No.{} (elapse_second={}s, fragment_id={}) : {}\n", i, elapsed,

pipeline_ctx->get_fragment_id(), pipeline_ctx->debug_string());

i++;

}

}

}

return fmt::to_string(debug_string_buffer);

const int fragment_id, std::shared_ptr<pipeline::PipelineFragmentContext> pip_ctx) {

std::lock_guard<std::mutex> lock(_pipeline_map_write_lock);

_fragment_id_to_pipeline_ctx.insert({fragment_id, pip_ctx});

_resource_ctx->set_workload_group(wg);

// Should add query first, then the workload group will not be deleted.

// see task_group_manager::delete_workload_group_by_ids

workload_group()->add_mem_tracker_limiter(query_mem_tracker());

workload_group()->get_query_scheduler(&_task_scheduler, &_scan_task_scheduler,

&_remote_scan_task_scheduler);

int fragment_id, const std::vector<std::shared_ptr<TRuntimeProfileTree>>& pipeline_profiles,

std::shared_ptr<TRuntimeProfileTree> load_channel_profile) {

if (pipeline_profiles.empty()) {

std::string msg = fmt::format("Add pipeline profile failed, query {}, fragment {}",

print_id(this->_query_id), fragment_id);

LOG_ERROR(msg);

DCHECK(false) << msg;

return;

}

#ifndef NDEBUG

for (const auto& p : pipeline_profiles) {

DCHECK(p != nullptr) << fmt::format("Add pipeline profile failed, query {}, fragment {}",

print_id(this->_query_id), fragment_id);

}

#endif

std::lock_guard<std::mutex> l(_profile_mutex);

VLOG_ROW << fmt::format(

"Query add fragment profile, query {}, fragment {}, pipeline profile count {} ",

print_id(this->_query_id), fragment_id, pipeline_profiles.size());

_profile_map.insert(std::make_pair(fragment_id, pipeline_profiles));

if (load_channel_profile != nullptr) {

_load_channel_profile_map.insert(std::make_pair(fragment_id, load_channel_profile));

}

std::lock_guard<std::mutex> lg(_profile_mutex);

for (auto& [fragment_id, fragment_profile] : _profile_map) {

std::shared_ptr<TRuntimeProfileTree> load_channel_profile = nullptr;

if (_load_channel_profile_map.contains(fragment_id)) {

load_channel_profile = _load_channel_profile_map[fragment_id];

}

ExecEnv::GetInstance()->runtime_query_statistics_mgr()->register_fragment_profile(

_query_id, this->coord_addr, fragment_id, fragment_profile, load_channel_profile);

}

ExecEnv::GetInstance()->runtime_query_statistics_mgr()->trigger_report_profile();

bool* has_running_task) const {

*revocable_size = 0;

for (auto&& [fragment_id, fragment_wptr] : _fragment_id_to_pipeline_ctx) {

auto fragment_ctx = fragment_wptr.lock();

if (!fragment_ctx) {

continue;

}

*revocable_size += fragment_ctx->get_revocable_size(has_running_task);

// Should wait for all tasks are not running before revoking memory.

if (*has_running_task) {

break;

}

}

*memory_usage = query_mem_tracker()->consumption();

size_t revocable_size = 0;

for (auto&& [fragment_id, fragment_wptr] : _fragment_id_to_pipeline_ctx) {

auto fragment_ctx = fragment_wptr.lock();

if (!fragment_ctx) {

continue;

}

bool has_running_task = false;

revocable_size += fragment_ctx->get_revocable_size(&has_running_task);

// Should wait for all tasks are not running before revoking memory.

if (has_running_task) {

return 0;

}

}

return revocable_size;

std::vector<std::pair<size_t, pipeline::PipelineTask*>> tasks;

std::vector<std::shared_ptr<pipeline::PipelineFragmentContext>> fragments;

for (auto&& [fragment_id, fragment_wptr] : _fragment_id_to_pipeline_ctx) {

auto fragment_ctx = fragment_wptr.lock();

if (!fragment_ctx) {

continue;

}

auto tasks_of_fragment = fragment_ctx->get_revocable_tasks();

for (auto* task : tasks_of_fragment) {

tasks.emplace_back(task->get_revocable_size(), task);

}

fragments.emplace_back(std::move(fragment_ctx));

}

std::sort(tasks.begin(), tasks.end(), [](auto&& l, auto&& r) { return l.first > r.first; });

// Do not use memlimit, use current memory usage.

// For example, if current limit is 1.6G, but current used is 1G, if reserve failed

// should free 200MB memory, not 300MB

const auto target_revoking_size = (int64_t)(query_mem_tracker()->consumption() * 0.2);

size_t revoked_size = 0;

size_t total_revokable_size = 0;

std::vector<pipeline::PipelineTask*> chosen_tasks;

for (auto&& [revocable_size, task] : tasks) {

// Only revoke the largest task to ensure memory is used as much as possible

// break;

if (revoked_size < target_revoking_size) {

chosen_tasks.emplace_back(task);

revoked_size += revocable_size;

}

total_revokable_size += revocable_size;

}

std::weak_ptr<QueryContext> this_ctx = shared_from_this();

auto spill_context = std::make_shared<pipeline::SpillContext>(

chosen_tasks.size(), _query_id, [this_ctx](pipeline::SpillContext* context) {

auto query_context = this_ctx.lock();

if (!query_context) {

return;

}

LOG(INFO) << query_context->debug_string() << ", context: " << ((void*)context)

<< " all spill tasks done, resume it.";

query_context->set_memory_sufficient(true);

});

LOG(INFO) << fmt::format(

"{}, spill context: {}, revokable mem: {}/{}, tasks count: {}/{}", this->debug_string(),

((void*)spill_context.get()), PrettyPrinter::print_bytes(revoked_size),

PrettyPrinter::print_bytes(total_revokable_size), chosen_tasks.size(), tasks.size());

for (auto* task : chosen_tasks) {

RETURN_IF_ERROR(task->revoke_memory(spill_context));

}

return Status::OK();

std::vector<pipeline::PipelineTask*> tasks;

for (auto&& [fragment_id, fragment_wptr] : _fragment_id_to_pipeline_ctx) {

auto fragment_ctx = fragment_wptr.lock();

if (!fragment_ctx) {

continue;

}

auto tasks_of_fragment = fragment_ctx->get_revocable_tasks();

tasks.insert(tasks.end(), tasks_of_fragment.cbegin(), tasks_of_fragment.cend());

}

return tasks;

std::lock_guard l(_paused_mutex);

return fmt::format(

"QueryId={}, Memory [Used={}, Limit={}, Peak={}], Spill[RunningSpillTaskCnt={}, "

"TotalPausedPeriodSecs={}, LatestPausedReason={}]",

print_id(_query_id),

PrettyPrinter::print(query_mem_tracker()->consumption(), TUnit::BYTES),

PrettyPrinter::print(query_mem_tracker()->limit(), TUnit::BYTES),

PrettyPrinter::print(query_mem_tracker()->peak_consumption(), TUnit::BYTES),

_revoking_tasks_count,

_memory_sufficient_dependency->watcher_elapse_time() / NANOS_PER_SEC,

_paused_reason.to_string());

std::unordered_map<int, std::vector<std::shared_ptr<TRuntimeProfileTree>>> res;

for (const auto& [fragment_id, fragment_ctx_wptr] : _fragment_id_to_pipeline_ctx) {

if (auto fragment_ctx = fragment_ctx_wptr.lock()) {

if (fragment_ctx == nullptr) {

std::string msg =

fmt::format("PipelineFragmentContext is nullptr, query {} fragment_id: {}",

print_id(_query_id), fragment_id);

LOG_ERROR(msg);

DCHECK(false) << msg;

continue;

}

auto profile = fragment_ctx->collect_realtime_profile();

if (profile.empty()) {

std::string err_msg = fmt::format(

"Get nothing when collecting profile, query {}, fragment_id: {}",

print_id(_query_id), fragment_id);

LOG_ERROR(err_msg);

DCHECK(false) << err_msg;

continue;

}

res.insert(std::make_pair(fragment_id, profile));

}

}

return res;

TReportExecStatusParams exec_status;

auto realtime_query_profile = _collect_realtime_query_profile();

std::vector<std::shared_ptr<TRuntimeProfileTree>> load_channel_profiles;

for (auto load_channel_profile : _load_channel_profile_map) {

if (load_channel_profile.second != nullptr) {

load_channel_profiles.push_back(load_channel_profile.second);

}

}

exec_status = RuntimeQueryStatisticsMgr::create_report_exec_status_params(

this->_query_id, std::move(realtime_query_profile), std::move(load_channel_profiles),

/*is_done=*/false);

return exec_status;