// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

#include "runtime/fragment_mgr.h"

#include <brpc/controller.h>

#include <bvar/latency_recorder.h>

#include <exprs/runtime_filter.h>

#include <fmt/format.h>

#include <gen_cpp/DorisExternalService_types.h>

#include <gen_cpp/FrontendService.h>

#include <gen_cpp/FrontendService_types.h>

#include <gen_cpp/HeartbeatService_types.h>

#include <gen_cpp/Metrics_types.h>

#include <gen_cpp/PaloInternalService_types.h>

#include <gen_cpp/PlanNodes_types.h>

#include <gen_cpp/Planner_types.h>

#include <gen_cpp/QueryPlanExtra_types.h>

#include <gen_cpp/RuntimeProfile_types.h>

#include <gen_cpp/Types_types.h>

#include <gen_cpp/internal_service.pb.h>

#include <pthread.h>

#include <stddef.h>

#include <sys/time.h>

#include <thrift/TApplicationException.h>

#include <thrift/Thrift.h>

#include <thrift/protocol/TDebugProtocol.h>

#include <thrift/transport/TTransportException.h>

#include <time.h>

#include <unistd.h>

#include <algorithm>

#include <atomic>

#include "common/status.h"

// IWYU pragma: no_include <bits/chrono.h>

#include <chrono> // IWYU pragma: keep

#include <cstdint>

#include <map>

#include <memory>

#include <mutex>

#include <sstream>

#include <unordered_map>

#include <unordered_set>

#include <utility>

#include "cloud/config.h"

#include "common/config.h"

#include "common/logging.h"

#include "common/object_pool.h"

#include "common/utils.h"

#include "gutil/strings/substitute.h"

#include "io/fs/stream_load_pipe.h"

#include "pipeline/pipeline_fragment_context.h"

#include "runtime/client_cache.h"

#include "runtime/descriptors.h"

#include "runtime/exec_env.h"

#include "runtime/frontend_info.h"

#include "runtime/memory/mem_tracker_limiter.h"

#include "runtime/primitive_type.h"

#include "runtime/query_context.h"

#include "runtime/runtime_filter_mgr.h"

#include "runtime/runtime_query_statistics_mgr.h"

#include "runtime/runtime_state.h"

#include "runtime/stream_load/new_load_stream_mgr.h"

#include "runtime/stream_load/stream_load_context.h"

#include "runtime/stream_load/stream_load_executor.h"

#include "runtime/thread_context.h"

#include "runtime/types.h"

#include "runtime/workload_group/workload_group.h"

#include "runtime/workload_group/workload_group_manager.h"

#include "runtime/workload_management/workload_query_info.h"

#include "service/backend_options.h"

#include "util/brpc_client_cache.h"

#include "util/debug_points.h"

#include "util/debug_util.h"

#include "util/doris_metrics.h"

#include "util/hash_util.hpp"

#include "util/mem_info.h"

#include "util/network_util.h"

#include "util/pretty_printer.h"

#include "util/runtime_profile.h"

#include "util/thread.h"

#include "util/threadpool.h"

#include "util/thrift_util.h"

#include "util/uid_util.h"

#include "util/url_coding.h"

#include "vec/runtime/shared_hash_table_controller.h"

#include "vec/runtime/vdatetime_value.h"

namespace doris {

DEFINE_GAUGE_METRIC_PROTOTYPE_2ARG(fragment_instance_count, MetricUnit::NOUNIT);

DEFINE_GAUGE_METRIC_PROTOTYPE_2ARG(timeout_canceled_fragment_count, MetricUnit::NOUNIT);

DEFINE_GAUGE_METRIC_PROTOTYPE_2ARG(fragment_thread_pool_queue_size, MetricUnit::NOUNIT);

bvar::LatencyRecorder g_fragmentmgr_prepare_latency("doris_FragmentMgr", "prepare");

bvar::Adder<uint64_t> g_fragment_executing_count("fragment_executing_count");

bvar::Status<uint64_t> g_fragment_last_active_time(

        "fragment_last_active_time", duration_cast<std::chrono::milliseconds>(

                                             std::chrono::system_clock::now().time_since_epoch())

                                             .count());

uint64_t get_fragment_executing_count() {

    return g_fragment_executing_count.get_value();

}

uint64_t get_fragment_last_active_time() {

    return g_fragment_last_active_time.get_value();

}

std::string to_load_error_http_path(const std::string& file_name) {

    if (file_name.empty()) {

        return "";

    }

    if (file_name.compare(0, 4, "http") == 0) {

        return file_name;

    }

    std::stringstream url;

    url << "http://" << get_host_port(BackendOptions::get_localhost(), config::webserver_port)

        << "/api/_load_error_log?"

        << "file=" << file_name;

    return url.str();

}

using apache::thrift::TException;

using apache::thrift::transport::TTransportException;

static Status _do_fetch_running_queries_rpc(const FrontendInfo& fe_info,

                                            std::unordered_set<TUniqueId>& query_set) {

    TFetchRunningQueriesResult rpc_result;

    TFetchRunningQueriesRequest rpc_request;

    Status client_status;

    const int32 timeout_ms = 3 * 1000;

    FrontendServiceConnection rpc_client(ExecEnv::GetInstance()->frontend_client_cache(),

                                         fe_info.info.coordinator_address, timeout_ms,

                                         &client_status);

    // Abort this fe.

    if (!client_status.ok()) {

        LOG_WARNING("Failed to get client for {}, reason is {}",

                    PrintThriftNetworkAddress(fe_info.info.coordinator_address),

                    client_status.to_string());

        return Status::InternalError("Failed to get client for {}, reason is {}",

                                     PrintThriftNetworkAddress(fe_info.info.coordinator_address),

                                     client_status.to_string());

    }

    // do rpc

    try {

        try {

            rpc_client->fetchRunningQueries(rpc_result, rpc_request);

        } catch (const apache::thrift::transport::TTransportException& e) {

            LOG_WARNING("Transport exception reason: {}, reopening", e.what());

            client_status = rpc_client.reopen(config::thrift_rpc_timeout_ms);

            if (!client_status.ok()) {

                LOG_WARNING("Reopen failed, reason: {}", client_status.to_string_no_stack());

                return Status::InternalError("Reopen failed, reason: {}",

                                             client_status.to_string_no_stack());

            }

            rpc_client->fetchRunningQueries(rpc_result, rpc_request);

        }

    } catch (apache::thrift::TException& e) {

        // During upgrading cluster or meet any other network error.

        LOG_WARNING("Failed to fetch running queries from {}, reason: {}",

                    PrintThriftNetworkAddress(fe_info.info.coordinator_address), e.what());

        return Status::InternalError("Failed to fetch running queries from {}, reason: {}",

                                     PrintThriftNetworkAddress(fe_info.info.coordinator_address),

                                     e.what());

    }

    // Avoid logic error in frontend.

    if (rpc_result.__isset.status == false || rpc_result.status.status_code != TStatusCode::OK) {

        LOG_WARNING("Failed to fetch running queries from {}, reason: {}",

                    PrintThriftNetworkAddress(fe_info.info.coordinator_address),

                    doris::to_string(rpc_result.status.status_code));

        return Status::InternalError("Failed to fetch running queries from {}, reason: {}",

                                     PrintThriftNetworkAddress(fe_info.info.coordinator_address),

                                     doris::to_string(rpc_result.status.status_code));

    }

    if (rpc_result.__isset.running_queries == false) {

        return Status::InternalError("Failed to fetch running queries from {}, reason: {}",

                                     PrintThriftNetworkAddress(fe_info.info.coordinator_address),

                                     "running_queries is not set");

    }

    query_set = std::unordered_set<TUniqueId>(rpc_result.running_queries.begin(),

                                              rpc_result.running_queries.end());

    return Status::OK();

};

static std::map<int64_t, std::unordered_set<TUniqueId>> _get_all_running_queries_from_fe() {

    const std::map<TNetworkAddress, FrontendInfo>& running_fes =

            ExecEnv::GetInstance()->get_running_frontends();

    std::map<int64_t, std::unordered_set<TUniqueId>> result;

    std::vector<FrontendInfo> qualified_fes;

    for (const auto& fe : running_fes) {

        // Only consider normal frontend.

        if (fe.first.port != 0 && fe.second.info.process_uuid != 0) {

            qualified_fes.push_back(fe.second);

        } else {

            return {};

        }

    }

    for (const auto& fe_addr : qualified_fes) {

        const int64_t process_uuid = fe_addr.info.process_uuid;

        std::unordered_set<TUniqueId> query_set;

        Status st = _do_fetch_running_queries_rpc(fe_addr, query_set);

        if (!st.ok()) {

            // Empty result, cancel worker will not do anything

            return {};

        }

        // frontend_info and process_uuid has been checked in rpc threads.

        result[process_uuid] = query_set;

    }

    return result;

}

FragmentMgr::FragmentMgr(ExecEnv* exec_env)

        : _exec_env(exec_env), _stop_background_threads_latch(1) {

    _entity = DorisMetrics::instance()->metric_registry()->register_entity("FragmentMgr");

    INT_UGAUGE_METRIC_REGISTER(_entity, timeout_canceled_fragment_count);

    auto s = Thread::create(

            "FragmentMgr", "cancel_timeout_plan_fragment", [this]() { this->cancel_worker(); },

            &_cancel_thread);

    CHECK(s.ok()) << s.to_string();

    s = ThreadPoolBuilder("FragmentMgrAsyncWorkThreadPool")

                .set_min_threads(config::fragment_mgr_asynic_work_pool_thread_num_min)

                .set_max_threads(config::fragment_mgr_asynic_work_pool_thread_num_max)

                .set_max_queue_size(config::fragment_mgr_asynic_work_pool_queue_size)

                .build(&_thread_pool);

    REGISTER_HOOK_METRIC(fragment_thread_pool_queue_size,

                         [this]() { return _thread_pool->get_queue_size(); });

    CHECK(s.ok()) << s.to_string();

}

FragmentMgr::~FragmentMgr() = default;

void FragmentMgr::stop() {

    DEREGISTER_HOOK_METRIC(fragment_instance_count);

    DEREGISTER_HOOK_METRIC(fragment_thread_pool_queue_size);

    _stop_background_threads_latch.count_down();

    if (_cancel_thread) {

        _cancel_thread->join();

    }

    // Only me can delete

    {

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

        _query_ctx_map.clear();

        _pipeline_map.clear();

    }

    _thread_pool->shutdown();

}

std::string FragmentMgr::to_http_path(const std::string& file_name) {

    std::stringstream url;

    url << "http://" << BackendOptions::get_localhost() << ":" << config::webserver_port

        << "/api/_download_load?"

        << "token=" << _exec_env->token() << "&file=" << file_name;

    return url.str();

}

Status FragmentMgr::trigger_pipeline_context_report(

        const ReportStatusRequest req, std::shared_ptr<pipeline::PipelineFragmentContext>&& ctx) {

    return _thread_pool->submit_func([this, req, ctx]() {

        SCOPED_ATTACH_TASK(ctx->get_query_ctx()->query_mem_tracker);

        coordinator_callback(req);

        if (!req.done) {

            ctx->refresh_next_report_time();

        }

    });

}

// There can only be one of these callbacks in-flight at any moment, because

// it is only invoked from the executor's reporting thread.

// Also, the reported status will always reflect the most recent execution status,

// including the final status when execution finishes.

void FragmentMgr::coordinator_callback(const ReportStatusRequest& req) {

    DCHECK(req.status.ok() || req.done); // if !status.ok() => done

    if (req.coord_addr.hostname == "external") {

        // External query (flink/spark read tablets) not need to report to FE.

        return;

    }

    Status exec_status = req.status;

    Status coord_status;

    FrontendServiceConnection coord(_exec_env->frontend_client_cache(), req.coord_addr,

                                    &coord_status);

    if (!coord_status.ok()) {

        std::stringstream ss;

        UniqueId uid(req.query_id.hi, req.query_id.lo);

        static_cast<void>(req.cancel_fn(Status::InternalError(

                "query_id: {}, couldn't get a client for {}, reason is {}", uid.to_string(),

                PrintThriftNetworkAddress(req.coord_addr), coord_status.to_string())));

        return;

    }

    TReportExecStatusParams params;

    params.protocol_version = FrontendServiceVersion::V1;

    params.__set_query_id(req.query_id);

    params.__set_backend_num(req.backend_num);

    params.__set_fragment_instance_id(req.fragment_instance_id);

    params.__set_fragment_id(req.fragment_id);

    params.__set_status(exec_status.to_thrift());

    params.__set_done(req.done);

    params.__set_query_type(req.runtime_state->query_type());

    params.__isset.profile = false;

    DCHECK(req.runtime_state != nullptr);

    if (req.runtime_state->query_type() == TQueryType::LOAD) {

        params.__set_loaded_rows(req.runtime_state->num_rows_load_total());

        params.__set_loaded_bytes(req.runtime_state->num_bytes_load_total());

    } else {

        DCHECK(!req.runtime_states.empty());

        if (!req.runtime_state->output_files().empty()) {

            params.__isset.delta_urls = true;

            for (auto& it : req.runtime_state->output_files()) {

                params.delta_urls.push_back(to_http_path(it));

            }

        }

        if (!params.delta_urls.empty()) {

            params.__isset.delta_urls = true;

        }

    }

    // load rows

    static std::string s_dpp_normal_all = "dpp.norm.ALL";

    static std::string s_dpp_abnormal_all = "dpp.abnorm.ALL";

    static std::string s_unselected_rows = "unselected.rows";

    int64_t num_rows_load_success = 0;

    int64_t num_rows_load_filtered = 0;

    int64_t num_rows_load_unselected = 0;

    if (req.runtime_state->num_rows_load_total() > 0 ||

        req.runtime_state->num_rows_load_filtered() > 0 ||

        req.runtime_state->num_finished_range() > 0) {

        params.__isset.load_counters = true;

        num_rows_load_success = req.runtime_state->num_rows_load_success();

        num_rows_load_filtered = req.runtime_state->num_rows_load_filtered();

        num_rows_load_unselected = req.runtime_state->num_rows_load_unselected();

        params.__isset.fragment_instance_reports = true;

        TFragmentInstanceReport t;

        t.__set_fragment_instance_id(req.runtime_state->fragment_instance_id());

        t.__set_num_finished_range(req.runtime_state->num_finished_range());

        t.__set_loaded_rows(req.runtime_state->num_rows_load_total());

        t.__set_loaded_bytes(req.runtime_state->num_bytes_load_total());

        params.fragment_instance_reports.push_back(t);

    } else if (!req.runtime_states.empty()) {

        for (auto* rs : req.runtime_states) {

            if (rs->num_rows_load_total() > 0 || rs->num_rows_load_filtered() > 0 ||

                req.runtime_state->num_finished_range() > 0) {

                params.__isset.load_counters = true;

                num_rows_load_success += rs->num_rows_load_success();

                num_rows_load_filtered += rs->num_rows_load_filtered();

                num_rows_load_unselected += rs->num_rows_load_unselected();

                params.__isset.fragment_instance_reports = true;

                TFragmentInstanceReport t;

                t.__set_fragment_instance_id(rs->fragment_instance_id());

                t.__set_num_finished_range(rs->num_finished_range());

                t.__set_loaded_rows(rs->num_rows_load_total());

                t.__set_loaded_bytes(rs->num_bytes_load_total());

                params.fragment_instance_reports.push_back(t);

            }

        }

    }

    params.load_counters.emplace(s_dpp_normal_all, std::to_string(num_rows_load_success));

    params.load_counters.emplace(s_dpp_abnormal_all, std::to_string(num_rows_load_filtered));

    params.load_counters.emplace(s_unselected_rows, std::to_string(num_rows_load_unselected));

    if (!req.runtime_state->get_error_log_file_path().empty()) {

        params.__set_tracking_url(

                to_load_error_http_path(req.runtime_state->get_error_log_file_path()));

    } else if (!req.runtime_states.empty()) {

        for (auto* rs : req.runtime_states) {

            if (!rs->get_error_log_file_path().empty()) {

                params.__set_tracking_url(to_load_error_http_path(rs->get_error_log_file_path()));

            }

            if (rs->wal_id() > 0) {

                params.__set_txn_id(rs->wal_id());

                params.__set_label(rs->import_label());

            }

        }

    }

    if (!req.runtime_state->export_output_files().empty()) {

        params.__isset.export_files = true;

        params.export_files = req.runtime_state->export_output_files();

    } else if (!req.runtime_states.empty()) {

        for (auto* rs : req.runtime_states) {

            if (!rs->export_output_files().empty()) {

                params.__isset.export_files = true;

                params.export_files.insert(params.export_files.end(),

                                           rs->export_output_files().begin(),

                                           rs->export_output_files().end());

            }

        }

    }

    if (!req.runtime_state->tablet_commit_infos().empty()) {

        params.__isset.commitInfos = true;

        params.commitInfos.reserve(req.runtime_state->tablet_commit_infos().size());

        for (auto& info : req.runtime_state->tablet_commit_infos()) {

            params.commitInfos.push_back(info);

        }

    } else if (!req.runtime_states.empty()) {

        for (auto* rs : req.runtime_states) {

            if (!rs->tablet_commit_infos().empty()) {

                params.__isset.commitInfos = true;

                params.commitInfos.insert(params.commitInfos.end(),

                                          rs->tablet_commit_infos().begin(),

                                          rs->tablet_commit_infos().end());

            }

        }

    }

    if (!req.runtime_state->error_tablet_infos().empty()) {

        params.__isset.errorTabletInfos = true;

        params.errorTabletInfos.reserve(req.runtime_state->error_tablet_infos().size());

        for (auto& info : req.runtime_state->error_tablet_infos()) {

            params.errorTabletInfos.push_back(info);

        }

    } else if (!req.runtime_states.empty()) {

        for (auto* rs : req.runtime_states) {

            if (!rs->error_tablet_infos().empty()) {

                params.__isset.errorTabletInfos = true;

                params.errorTabletInfos.insert(params.errorTabletInfos.end(),

                                               rs->error_tablet_infos().begin(),

                                               rs->error_tablet_infos().end());

            }

        }

    }

    if (!req.runtime_state->hive_partition_updates().empty()) {

        params.__isset.hive_partition_updates = true;

        params.hive_partition_updates.reserve(req.runtime_state->hive_partition_updates().size());

        for (auto& hive_partition_update : req.runtime_state->hive_partition_updates()) {

            params.hive_partition_updates.push_back(hive_partition_update);

        }

    } else if (!req.runtime_states.empty()) {

        for (auto* rs : req.runtime_states) {

            if (!rs->hive_partition_updates().empty()) {

                params.__isset.hive_partition_updates = true;

                params.hive_partition_updates.insert(params.hive_partition_updates.end(),

                                                     rs->hive_partition_updates().begin(),

                                                     rs->hive_partition_updates().end());

            }

        }

    }

    if (!req.runtime_state->iceberg_commit_datas().empty()) {

        params.__isset.iceberg_commit_datas = true;

        params.iceberg_commit_datas.reserve(req.runtime_state->iceberg_commit_datas().size());

        for (auto& iceberg_commit_data : req.runtime_state->iceberg_commit_datas()) {

            params.iceberg_commit_datas.push_back(iceberg_commit_data);

        }

    } else if (!req.runtime_states.empty()) {

        for (auto* rs : req.runtime_states) {

            if (!rs->iceberg_commit_datas().empty()) {

                params.__isset.iceberg_commit_datas = true;

                params.iceberg_commit_datas.insert(params.iceberg_commit_datas.end(),

                                                   rs->iceberg_commit_datas().begin(),

                                                   rs->iceberg_commit_datas().end());

            }

        }

    }

    // Send new errors to coordinator

    req.runtime_state->get_unreported_errors(&(params.error_log));

    params.__isset.error_log = (!params.error_log.empty());

    if (_exec_env->cluster_info()->backend_id != 0) {

        params.__set_backend_id(_exec_env->cluster_info()->backend_id);

    }

    TReportExecStatusResult res;

    Status rpc_status;

    VLOG_DEBUG << "reportExecStatus params is "

               << apache::thrift::ThriftDebugString(params).c_str();

    if (!exec_status.ok()) {

        LOG(WARNING) << "report error status: " << exec_status.msg()

                     << " to coordinator: " << req.coord_addr

                     << ", query id: " << print_id(req.query_id);

    }

    try {

        try {

            coord->reportExecStatus(res, params);

        } catch (TTransportException& e) {

            LOG(WARNING) << "Retrying ReportExecStatus. query id: " << print_id(req.query_id)

                         << ", instance id: " << print_id(req.fragment_instance_id) << " to "

                         << req.coord_addr << ", err: " << e.what();

            rpc_status = coord.reopen();

            if (!rpc_status.ok()) {

                // we need to cancel the execution of this fragment

                req.cancel_fn(rpc_status);

                return;

            }

            coord->reportExecStatus(res, params);

        }

        rpc_status = Status::create<false>(res.status);

    } catch (TException& e) {

        rpc_status = Status::InternalError("ReportExecStatus() to {} failed: {}",

                                           PrintThriftNetworkAddress(req.coord_addr), e.what());

    }

    if (!rpc_status.ok()) {

        LOG_INFO("Going to cancel query {} since report exec status got rpc failed: {}",

                 print_id(req.query_id), rpc_status.to_string());

        // we need to cancel the execution of this fragment

        req.cancel_fn(rpc_status);

    }

}

static void empty_function(RuntimeState*, Status*) {}

Status FragmentMgr::exec_plan_fragment(const TExecPlanFragmentParams& params,

                                       const QuerySource query_source) {

    return Status::InternalError("Non-pipeline is disabled!");

}

Status FragmentMgr::exec_plan_fragment(const TPipelineFragmentParams& params,

                                       const QuerySource query_source) {

    if (params.txn_conf.need_txn) {

        std::shared_ptr<StreamLoadContext> stream_load_ctx =

                std::make_shared<StreamLoadContext>(_exec_env);

        stream_load_ctx->db = params.txn_conf.db;

        stream_load_ctx->db_id = params.txn_conf.db_id;

        stream_load_ctx->table = params.txn_conf.tbl;

        stream_load_ctx->txn_id = params.txn_conf.txn_id;

        stream_load_ctx->id = UniqueId(params.query_id);

        stream_load_ctx->put_result.__set_pipeline_params(params);

        stream_load_ctx->use_streaming = true;

        stream_load_ctx->load_type = TLoadType::MANUL_LOAD;

        stream_load_ctx->load_src_type = TLoadSourceType::RAW;

        stream_load_ctx->label = params.import_label;

        stream_load_ctx->format = TFileFormatType::FORMAT_CSV_PLAIN;

        stream_load_ctx->timeout_second = 3600;

        stream_load_ctx->auth.token = params.txn_conf.token;

        stream_load_ctx->need_commit_self = true;

        stream_load_ctx->need_rollback = true;

        auto pipe = std::make_shared<io::StreamLoadPipe>(

                io::kMaxPipeBufferedBytes /* max_buffered_bytes */, 64 * 1024 /* min_chunk_size */,

                -1 /* total_length */, true /* use_proto */);

        stream_load_ctx->body_sink = pipe;

        stream_load_ctx->pipe = pipe;

        stream_load_ctx->max_filter_ratio = params.txn_conf.max_filter_ratio;

        RETURN_IF_ERROR(

                _exec_env->new_load_stream_mgr()->put(stream_load_ctx->id, stream_load_ctx));

        RETURN_IF_ERROR(_exec_env->stream_load_executor()->execute_plan_fragment(stream_load_ctx));

        return Status::OK();

    } else {

        return exec_plan_fragment(params, query_source, empty_function);

    }

}

// Stage 2. prepare finished. then get FE instruction to execute

Status FragmentMgr::start_query_execution(const PExecPlanFragmentStartRequest* request) {

    TUniqueId query_id;

    query_id.__set_hi(request->query_id().hi());

    query_id.__set_lo(request->query_id().lo());

    std::shared_ptr<QueryContext> q_ctx = nullptr;

    {

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

        q_ctx = _get_or_erase_query_ctx(query_id);

    }

    if (q_ctx) {

        q_ctx->set_ready_to_execute(Status::OK());

        LOG_INFO("Query {} start execution", print_id(query_id));

    } else {

        return Status::InternalError(

                "Failed to get query fragments context. Query may be "

                "timeout or be cancelled. host: {}",

                BackendOptions::get_localhost());

    }

    return Status::OK();

}

void FragmentMgr::remove_pipeline_context(

        std::shared_ptr<pipeline::PipelineFragmentContext> f_context) {

    {

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

        auto query_id = f_context->get_query_id();

        int64 now = duration_cast<std::chrono::milliseconds>(

                            std::chrono::system_clock::now().time_since_epoch())

                            .count();

        g_fragment_executing_count << -1;

        g_fragment_last_active_time.set_value(now);

        _pipeline_map.erase({query_id, f_context->get_fragment_id()});

    }

}

std::shared_ptr<QueryContext> FragmentMgr::_get_or_erase_query_ctx(const TUniqueId& query_id) {

    auto search = _query_ctx_map.find(query_id);

    if (search != _query_ctx_map.end()) {

        if (auto q_ctx = search->second.lock()) {

            return q_ctx;

        } else {

            LOG(WARNING) << "Query context (query id = " << print_id(query_id)

                         << ") has been released.";

            _query_ctx_map.erase(search);

            return nullptr;

        }

    }

    return nullptr;

}

std::shared_ptr<QueryContext> FragmentMgr::get_or_erase_query_ctx_with_lock(

        const TUniqueId& query_id) {

    std::unique_lock<std::mutex> lock(_lock);

    return _get_or_erase_query_ctx(query_id);

}

template <typename Params>

Status FragmentMgr::_get_query_ctx(const Params& params, TUniqueId query_id, bool pipeline,

                                   QuerySource query_source,

                                   std::shared_ptr<QueryContext>& query_ctx) {

    DBUG_EXECUTE_IF("FragmentMgr._get_query_ctx.failed", {

        return Status::InternalError("FragmentMgr._get_query_ctx.failed, query id {}",

                                     print_id(query_id));

    });

    if (params.is_simplified_param) {

        // Get common components from _query_ctx_map

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

        if (auto q_ctx = _get_or_erase_query_ctx(query_id)) {

            query_ctx = q_ctx;

        } else {

            return Status::InternalError(

                    "Failed to get query fragments context. Query {} may be timeout or be "

                    "cancelled. host: {}",

                    print_id(query_id), BackendOptions::get_localhost());

        }

    } else {

        // Find _query_ctx_map, in case some other request has already

        // create the query fragments context.

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

        if (auto q_ctx = _get_or_erase_query_ctx(query_id)) {

            query_ctx = q_ctx;

            return Status::OK();

        }

        // First time a fragment of a query arrived. print logs.

        LOG(INFO) << "query_id: " << print_id(query_id) << ", coord_addr: " << params.coord

                  << ", total fragment num on current host: " << params.fragment_num_on_host

                  << ", fe process uuid: " << params.query_options.fe_process_uuid

                  << ", query type: " << params.query_options.query_type

                  << ", report audit fe:" << params.current_connect_fe;

        // This may be a first fragment request of the query.

        // Create the query fragments context.

        query_ctx = QueryContext::create_shared(query_id, _exec_env, params.query_options,

                                                params.coord, params.is_nereids,

                                                params.current_connect_fe, query_source);

        SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(query_ctx->query_mem_tracker);

        RETURN_IF_ERROR(DescriptorTbl::create(&(query_ctx->obj_pool), params.desc_tbl,

                                              &(query_ctx->desc_tbl)));

        // set file scan range params

        if (params.__isset.file_scan_params) {

            query_ctx->file_scan_range_params_map = params.file_scan_params;

        }

        query_ctx->query_globals = params.query_globals;

        if (params.__isset.resource_info) {

            query_ctx->user = params.resource_info.user;

            query_ctx->group = params.resource_info.group;

            query_ctx->set_rsc_info = true;

        }

        _set_scan_concurrency(params, query_ctx.get());

        if (params.__isset.workload_groups && !params.workload_groups.empty()) {

            uint64_t tg_id = params.workload_groups[0].id;

            WorkloadGroupPtr workload_group_ptr =

                    _exec_env->workload_group_mgr()->get_task_group_by_id(tg_id);

            if (workload_group_ptr != nullptr) {

                RETURN_IF_ERROR(workload_group_ptr->add_query(query_id, query_ctx));

                RETURN_IF_ERROR(query_ctx->set_workload_group(workload_group_ptr));

                _exec_env->runtime_query_statistics_mgr()->set_workload_group_id(print_id(query_id),

                                                                                 tg_id);

            } else {

                LOG(WARNING) << "Query/load id: " << print_id(query_ctx->query_id())

                             << "can't find its workload group " << tg_id;

            }

        }

        // There is some logic in query ctx's dctor, we could not check if exists and delete the

        // temp query ctx now. For example, the query id maybe removed from workload group's queryset.

        _query_ctx_map.insert(std::make_pair(query_ctx->query_id(), query_ctx));

    }

    return Status::OK();

}

Status FragmentMgr::exec_plan_fragment(const TExecPlanFragmentParams& params,

                                       QuerySource query_source, const FinishCallback& cb) {

    return Status::InternalError("Non-pipeline is disabled!");

}

std::string FragmentMgr::dump_pipeline_tasks(int64_t duration) {

    fmt::memory_buffer debug_string_buffer;

    size_t i = 0;

    {

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

        fmt::format_to(debug_string_buffer,

                       "{} pipeline fragment contexts are still running! duration_limit={}\n",

                       _pipeline_map.size(), duration);

        timespec now;

        clock_gettime(CLOCK_MONOTONIC, &now);

        for (auto& it : _pipeline_map) {

            auto elapsed = it.second->elapsed_time() / 1000000000.0;

            if (elapsed < duration) {

                // Only display tasks which has been running for more than {duration} seconds.

                continue;

            }

            auto timeout_second = it.second->timeout_second();

            fmt::format_to(

                    debug_string_buffer,

                    "No.{} (elapse_second={}s, query_timeout_second={}s, is_timeout={}) : {}\n", i,

                    elapsed, timeout_second, it.second->is_timeout(now), it.second->debug_string());

            i++;

        }

    }

    return fmt::to_string(debug_string_buffer);

}

std::string FragmentMgr::dump_pipeline_tasks(TUniqueId& query_id) {

    if (auto q_ctx = _get_or_erase_query_ctx(query_id)) {

        return q_ctx->print_all_pipeline_context();

    } else {

        return fmt::format("Query context (query id = {}) not found. \n", print_id(query_id));

    }

}

Status FragmentMgr::exec_plan_fragment(const TPipelineFragmentParams& params,

                                       QuerySource query_source, const FinishCallback& cb) {

    VLOG_ROW << "query: " << print_id(params.query_id) << " exec_plan_fragment params is "

             << apache::thrift::ThriftDebugString(params).c_str();

    // sometimes TExecPlanFragmentParams debug string is too long and glog

    // will truncate the log line, so print query options seperately for debuggin purpose

    VLOG_ROW << "query: " << print_id(params.query_id) << "query options is "

             << apache::thrift::ThriftDebugString(params.query_options).c_str();

    std::shared_ptr<QueryContext> query_ctx;

    RETURN_IF_ERROR(_get_query_ctx(params, params.query_id, true, query_source, query_ctx));

    SCOPED_ATTACH_TASK(query_ctx.get());

    int64_t duration_ns = 0;

    std::shared_ptr<pipeline::PipelineFragmentContext> context =

            std::make_shared<pipeline::PipelineFragmentContext>(

                    query_ctx->query_id(), params.fragment_id, query_ctx, _exec_env, cb,

                    std::bind<Status>(std::mem_fn(&FragmentMgr::trigger_pipeline_context_report),

                                      this, std::placeholders::_1, std::placeholders::_2));

    {

        SCOPED_RAW_TIMER(&duration_ns);

        Status prepare_st = Status::OK();

        ASSIGN_STATUS_IF_CATCH_EXCEPTION(prepare_st = context->prepare(params, _thread_pool.get()),

                                         prepare_st);

        if (!prepare_st.ok()) {

            query_ctx->cancel(prepare_st, params.fragment_id);

            query_ctx->set_execution_dependency_ready();

            return prepare_st;

        }

    }

    g_fragmentmgr_prepare_latency << (duration_ns / 1000);

    DBUG_EXECUTE_IF("FragmentMgr.exec_plan_fragment.failed",

                    { return Status::Aborted("FragmentMgr.exec_plan_fragment.failed"); });

    std::shared_ptr<RuntimeFilterMergeControllerEntity> handler;

    RETURN_IF_ERROR(_runtimefilter_controller.add_entity(

            params.local_params[0], params.query_id, params.query_options, &handler,

            RuntimeFilterParamsContext::create(context->get_runtime_state())));

    if (handler) {

        query_ctx->set_merge_controller_handler(handler);

    }

    {

        // (query_id, fragment_id) is executed only on one BE, locks _pipeline_map.

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

        for (const auto& local_param : params.local_params) {

            const TUniqueId& fragment_instance_id = local_param.fragment_instance_id;

            auto iter = _pipeline_map.find({params.query_id, params.fragment_id});

            if (iter != _pipeline_map.end()) {

                return Status::InternalError(

                        "exec_plan_fragment query_id({}) input duplicated fragment_id({})",

                        print_id(params.query_id), params.fragment_id);

            }

            query_ctx->fragment_instance_ids.push_back(fragment_instance_id);

        }

        int64 now = duration_cast<std::chrono::milliseconds>(

                            std::chrono::system_clock::now().time_since_epoch())

                            .count();

        g_fragment_executing_count << 1;

        g_fragment_last_active_time.set_value(now);

        // TODO: simplify this mapping

        _pipeline_map.insert({{params.query_id, params.fragment_id}, context});

    }

    if (!params.__isset.need_wait_execution_trigger || !params.need_wait_execution_trigger) {

        query_ctx->set_ready_to_execute_only();

    }

    query_ctx->set_pipeline_context(params.fragment_id, context);

    RETURN_IF_ERROR(context->submit());

    return Status::OK();

}

template <typename Param>

void FragmentMgr::_set_scan_concurrency(const Param& params, QueryContext* query_ctx) {

#ifndef BE_TEST

    // If the token is set, the scan task will use limited_scan_pool in scanner scheduler.

    // Otherwise, the scan task will use local/remote scan pool in scanner scheduler

    if (params.query_options.__isset.resource_limit &&

        params.query_options.resource_limit.__isset.cpu_limit) {

        query_ctx->set_thread_token(params.query_options.resource_limit.cpu_limit, false);

    }

#endif

}

void FragmentMgr::cancel_query(const TUniqueId query_id, const Status reason) {

    std::shared_ptr<QueryContext> query_ctx = nullptr;

    std::vector<TUniqueId> all_instance_ids;

    {

        std::lock_guard<std::mutex> state_lock(_lock);

        if (auto q_ctx = _get_or_erase_query_ctx(query_id)) {

            query_ctx = q_ctx;

            // Copy instanceids to avoid concurrent modification.

            // And to reduce the scope of lock.

            all_instance_ids = query_ctx->fragment_instance_ids;

        } else {

            LOG(WARNING) << "Query " << print_id(query_id)

                         << " does not exists, failed to cancel it";

            return;

        }

    }

    query_ctx->cancel(reason);

    {

        std::lock_guard<std::mutex> state_lock(_lock);

        _query_ctx_map.erase(query_id);

    }

    LOG(INFO) << "Query " << print_id(query_id)

              << " is cancelled and removed. Reason: " << reason.to_string();

}

void FragmentMgr::cancel_worker() {

    LOG(INFO) << "FragmentMgr cancel worker start working.";

    timespec check_invalid_query_last_timestamp;

    clock_gettime(CLOCK_MONOTONIC, &check_invalid_query_last_timestamp);

    do {

        std::vector<TUniqueId> queries_lost_coordinator;

        std::vector<TUniqueId> queries_timeout;

        std::vector<TUniqueId> queries_pipeline_task_leak;

        // Fe process uuid -> set<QueryId>

        std::map<int64_t, std::unordered_set<TUniqueId>> running_queries_on_all_fes;

        const std::map<TNetworkAddress, FrontendInfo>& running_fes =

                ExecEnv::GetInstance()->get_running_frontends();

        timespec now;

        clock_gettime(CLOCK_MONOTONIC, &now);

        if (config::enable_pipeline_task_leakage_detect &&

            now.tv_sec - check_invalid_query_last_timestamp.tv_sec >

                    config::pipeline_task_leakage_detect_period_secs) {

            check_invalid_query_last_timestamp = now;

            running_queries_on_all_fes = _get_all_running_queries_from_fe();

        } else {

            running_queries_on_all_fes.clear();

        }

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

        {

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

            ctx.reserve(_pipeline_map.size());

            for (auto& pipeline_itr : _pipeline_map) {

                ctx.push_back(pipeline_itr.second);

            }

        }

        for (auto& c : ctx) {

            c->clear_finished_tasks();

        }