// 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 "agent/task_worker_pool.h"

#include <brpc/controller.h>

#include <fmt/format.h>

#include <gen_cpp/AgentService_types.h>

#include <gen_cpp/DataSinks_types.h>

#include <gen_cpp/HeartbeatService_types.h>

#include <gen_cpp/MasterService_types.h>

#include <gen_cpp/Status_types.h>

#include <gen_cpp/Types_types.h>

#include <unistd.h>

#include <algorithm>

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

#include <thrift/protocol/TDebugProtocol.h>

#include <atomic>

#include <chrono> // IWYU pragma: keep

#include <ctime>

#include <functional>

#include <memory>

#include <mutex>

#include <shared_mutex>

#include <sstream>

#include <string>

#include <thread>

#include <type_traits>

#include <utility>

#include <vector>

#include "agent/utils.h"

#include "cloud/cloud_delete_task.h"

#include "cloud/cloud_engine_calc_delete_bitmap_task.h"

#include "cloud/cloud_schema_change_job.h"

#include "cloud/cloud_snapshot_loader.h"

#include "cloud/cloud_snapshot_mgr.h"

#include "cloud/cloud_tablet.h"

#include "cloud/cloud_tablet_mgr.h"

#include "cloud/config.h"

#include "common/config.h"

#include "common/logging.h"

#include "common/metrics/doris_metrics.h"

#include "common/status.h"

#include "io/fs/file_system.h"

#include "io/fs/hdfs_file_system.h"

#include "io/fs/local_file_system.h"

#include "io/fs/obj_storage_client.h"

#include "io/fs/path.h"

#include "io/fs/remote_file_system.h"

#include "io/fs/s3_file_system.h"

#include "runtime/exec_env.h"

#include "runtime/fragment_mgr.h"

#include "runtime/index_policy/index_policy_mgr.h"

#include "runtime/memory/global_memory_arbitrator.h"

#include "runtime/snapshot_loader.h"

#include "runtime/user_function_cache.h"

#include "service/backend_options.h"

#include "storage/compaction/cumulative_compaction_time_series_policy.h"

#include "storage/data_dir.h"

#include "storage/olap_common.h"

#include "storage/rowset/rowset_meta.h"

#include "storage/snapshot/snapshot_manager.h"

#include "storage/storage_engine.h"

#include "storage/storage_policy.h"

#include "storage/tablet/tablet.h"

#include "storage/tablet/tablet_manager.h"

#include "storage/tablet/tablet_meta.h"

#include "storage/tablet/tablet_schema.h"

#include "storage/task/engine_batch_load_task.h"

#include "storage/task/engine_checksum_task.h"

#include "storage/task/engine_clone_task.h"

#include "storage/task/engine_cloud_index_change_task.h"

#include "storage/task/engine_index_change_task.h"

#include "storage/task/engine_publish_version_task.h"

#include "storage/task/engine_storage_migration_task.h"

#include "storage/txn/txn_manager.h"

#include "storage/utils.h"

#include "util/brpc_client_cache.h"

#include "util/debug_points.h"

#include "util/jni-util.h"

#include "util/mem_info.h"

#include "util/random.h"

#include "util/s3_util.h"

#include "util/stopwatch.hpp"

#include "util/threadpool.h"

#include "util/time.h"

#include "util/trace.h"

namespace doris {

using namespace ErrorCode;

namespace {

std::mutex s_task_signatures_mtx;

std::unordered_map<TTaskType::type, std::unordered_set<int64_t>> s_task_signatures;

std::atomic_ulong s_report_version(time(nullptr) * 100000);

void increase_report_version() {

    s_report_version.fetch_add(1, std::memory_order_relaxed);

}

// FIXME(plat1ko): Paired register and remove task info

bool register_task_info(const TTaskType::type task_type, int64_t signature) {

    if (task_type == TTaskType::type::PUSH_STORAGE_POLICY ||

        task_type == TTaskType::type::PUSH_COOLDOWN_CONF ||

        task_type == TTaskType::type::COMPACTION) {

        // no need to report task of these types

        return true;

    }

    if (signature == -1) { // No need to report task with unintialized signature

        return true;

    }

    std::lock_guard lock(s_task_signatures_mtx);

    auto& set = s_task_signatures[task_type];

    return set.insert(signature).second;

}

void remove_task_info(const TTaskType::type task_type, int64_t signature) {

    size_t queue_size;

    {

        std::lock_guard lock(s_task_signatures_mtx);

        auto& set = s_task_signatures[task_type];

        set.erase(signature);

        queue_size = set.size();

    }

    VLOG_NOTICE << "remove task info. type=" << task_type << ", signature=" << signature

                << ", queue_size=" << queue_size;

}

void finish_task(const TFinishTaskRequest& finish_task_request) {

    // Return result to FE

    TMasterResult result;

    uint32_t try_time = 0;

    constexpr int TASK_FINISH_MAX_RETRY = 3;

    while (try_time < TASK_FINISH_MAX_RETRY) {

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

        Status client_status =

                MasterServerClient::instance()->finish_task(finish_task_request, &result);

        if (client_status.ok()) {

            break;

        } else {

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

            LOG_WARNING("failed to finish task")

                    .tag("type", finish_task_request.task_type)

                    .tag("signature", finish_task_request.signature)

                    .error(result.status);

            try_time += 1;

        }

        sleep(1);

    }

}

Status get_tablet_info(StorageEngine& engine, const TTabletId tablet_id,

                       const TSchemaHash schema_hash, TTabletInfo* tablet_info) {

    tablet_info->__set_tablet_id(tablet_id);

    tablet_info->__set_schema_hash(schema_hash);

    return engine.tablet_manager()->report_tablet_info(tablet_info);

}

void random_sleep(int second) {

    Random rnd(static_cast<uint32_t>(UnixMillis()));

    sleep(rnd.Uniform(second) + 1);

}

void alter_tablet(StorageEngine& engine, const TAgentTaskRequest& agent_task_req, int64_t signature,

                  const TTaskType::type task_type, TFinishTaskRequest* finish_task_request) {

    Status status;

    std::string_view process_name = "alter tablet";

    // Check last schema change status, if failed delete tablet file

    // Do not need to adjust delete success or not

    // Because if delete failed create rollup will failed

    TTabletId new_tablet_id = 0;

    TSchemaHash new_schema_hash = 0;

    if (status.ok()) {

        new_tablet_id = agent_task_req.alter_tablet_req_v2.new_tablet_id;

        new_schema_hash = agent_task_req.alter_tablet_req_v2.new_schema_hash;

        auto mem_tracker = MemTrackerLimiter::create_shared(

                MemTrackerLimiter::Type::SCHEMA_CHANGE,

                fmt::format("EngineAlterTabletTask#baseTabletId={}:newTabletId={}",

                            std::to_string(agent_task_req.alter_tablet_req_v2.base_tablet_id),

                            std::to_string(agent_task_req.alter_tablet_req_v2.new_tablet_id),

                            engine.memory_limitation_bytes_per_thread_for_schema_change()));

        SCOPED_ATTACH_TASK(mem_tracker);

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

        Status res = Status::OK();

        try {

            LOG_INFO("start {}", process_name)

                    .tag("signature", agent_task_req.signature)

                    .tag("base_tablet_id", agent_task_req.alter_tablet_req_v2.base_tablet_id)

                    .tag("new_tablet_id", new_tablet_id)

                    .tag("mem_limit",

                         engine.memory_limitation_bytes_per_thread_for_schema_change());

            SchemaChangeJob job(engine, agent_task_req.alter_tablet_req_v2,

                                std::to_string(agent_task_req.alter_tablet_req_v2.__isset.job_id

                                                       ? agent_task_req.alter_tablet_req_v2.job_id

                                                       : 0));

            status = job.process_alter_tablet(agent_task_req.alter_tablet_req_v2);

        } catch (const Exception& e) {

            status = e.to_status();

        }

        if (!status.ok()) {

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

        }

    }

    if (status.ok()) {

        increase_report_version();

    }

    // Return result to fe

    finish_task_request->__set_backend(BackendOptions::get_local_backend());

    finish_task_request->__set_report_version(s_report_version);

    finish_task_request->__set_task_type(task_type);

    finish_task_request->__set_signature(signature);

    std::vector<TTabletInfo> finish_tablet_infos;

    if (status.ok()) {

        TTabletInfo tablet_info;

        status = get_tablet_info(engine, new_tablet_id, new_schema_hash, &tablet_info);

        if (status.ok()) {

            finish_tablet_infos.push_back(tablet_info);

        }

    }

    if (!status.ok() && !status.is<NOT_IMPLEMENTED_ERROR>()) {

        LOG_WARNING("failed to {}", process_name)

                .tag("signature", agent_task_req.signature)

                .tag("base_tablet_id", agent_task_req.alter_tablet_req_v2.base_tablet_id)

                .tag("new_tablet_id", new_tablet_id)

                .error(status);

    } else {

        finish_task_request->__set_finish_tablet_infos(finish_tablet_infos);

        LOG_INFO("successfully {}", process_name)

                .tag("signature", agent_task_req.signature)

                .tag("base_tablet_id", agent_task_req.alter_tablet_req_v2.base_tablet_id)

                .tag("new_tablet_id", new_tablet_id);

    }

    finish_task_request->__set_task_status(status.to_thrift());

}

void alter_cloud_tablet(CloudStorageEngine& engine, const TAgentTaskRequest& agent_task_req,

                        int64_t signature, const TTaskType::type task_type,

                        TFinishTaskRequest* finish_task_request) {

    Status status;

    std::string_view process_name = "alter tablet";

    // Check last schema change status, if failed delete tablet file

    // Do not need to adjust delete success or not

    // Because if delete failed create rollup will failed

    TTabletId new_tablet_id = 0;

    new_tablet_id = agent_task_req.alter_tablet_req_v2.new_tablet_id;

    auto mem_tracker = MemTrackerLimiter::create_shared(

            MemTrackerLimiter::Type::SCHEMA_CHANGE,

            fmt::format("EngineAlterTabletTask#baseTabletId={}:newTabletId={}",

                        std::to_string(agent_task_req.alter_tablet_req_v2.base_tablet_id),

                        std::to_string(agent_task_req.alter_tablet_req_v2.new_tablet_id),

                        engine.memory_limitation_bytes_per_thread_for_schema_change()));

    SCOPED_ATTACH_TASK(mem_tracker);

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

    LOG_INFO("start {}", process_name)

            .tag("signature", agent_task_req.signature)

            .tag("base_tablet_id", agent_task_req.alter_tablet_req_v2.base_tablet_id)

            .tag("new_tablet_id", new_tablet_id)

            .tag("mem_limit", engine.memory_limitation_bytes_per_thread_for_schema_change());

    DCHECK(agent_task_req.alter_tablet_req_v2.__isset.job_id);

    CloudSchemaChangeJob job(engine, std::to_string(agent_task_req.alter_tablet_req_v2.job_id),

                             agent_task_req.alter_tablet_req_v2.expiration);

    status = [&]() {

        HANDLE_EXCEPTION_IF_CATCH_EXCEPTION(

                job.process_alter_tablet(agent_task_req.alter_tablet_req_v2),

                [&](const doris::Exception& ex) {

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

                    job.clean_up_on_failure();

                });

        return Status::OK();

    }();

    if (status.ok()) {

        increase_report_version();

        LOG_INFO("successfully {}", process_name)

                .tag("signature", agent_task_req.signature)

                .tag("base_tablet_id", agent_task_req.alter_tablet_req_v2.base_tablet_id)

                .tag("new_tablet_id", new_tablet_id);

    } else {

        LOG_WARNING("failed to {}", process_name)

                .tag("signature", agent_task_req.signature)

                .tag("base_tablet_id", agent_task_req.alter_tablet_req_v2.base_tablet_id)

                .tag("new_tablet_id", new_tablet_id)

                .error(status);

    }

    // Return result to fe

    finish_task_request->__set_backend(BackendOptions::get_local_backend());

    finish_task_request->__set_report_version(s_report_version);

    finish_task_request->__set_task_type(task_type);

    finish_task_request->__set_signature(signature);

    finish_task_request->__set_task_status(status.to_thrift());

}

Status check_migrate_request(StorageEngine& engine, const TStorageMediumMigrateReq& req,

                             TabletSharedPtr& tablet, DataDir** dest_store) {

    int64_t tablet_id = req.tablet_id;

    tablet = engine.tablet_manager()->get_tablet(tablet_id);

    if (tablet == nullptr) {

        return Status::InternalError("could not find tablet {}", tablet_id);

    }

    if (req.__isset.data_dir) {

        // request specify the data dir

        *dest_store = engine.get_store(req.data_dir);

        if (*dest_store == nullptr) {

            return Status::InternalError("could not find data dir {}", req.data_dir);

        }

    } else {

        // this is a storage medium

        // get data dir by storage medium

        // judge case when no need to migrate

        uint32_t count = engine.available_storage_medium_type_count();

        if (count <= 1) {

            return Status::InternalError("available storage medium type count is less than 1");

        }

        // check current tablet storage medium

        TStorageMedium::type storage_medium = req.storage_medium;

        TStorageMedium::type src_storage_medium = tablet->data_dir()->storage_medium();

        if (src_storage_medium == storage_medium) {

            return Status::InternalError("tablet is already on specified storage medium {}",

                                         storage_medium);

        }

        // get a random store of specified storage medium

        auto stores = engine.get_stores_for_create_tablet(tablet->partition_id(), storage_medium);

        if (stores.empty()) {

            return Status::InternalError("failed to get root path for create tablet");

        }

        *dest_store = stores[0];

    }

    if (tablet->data_dir()->path() == (*dest_store)->path()) {

        LOG_WARNING("tablet is already on specified path").tag("path", tablet->data_dir()->path());

        return Status::Error<FILE_ALREADY_EXIST, false>("tablet is already on specified path: {}",

                                                        tablet->data_dir()->path());

    }

    // check local disk capacity

    int64_t tablet_size = tablet->tablet_local_size();

    if ((*dest_store)->reach_capacity_limit(tablet_size)) {

        return Status::Error<EXCEEDED_LIMIT>("reach the capacity limit of path {}, tablet_size={}",

                                             (*dest_store)->path(), tablet_size);

    }

    return Status::OK();

}

// Return `true` if report success

bool handle_report(const TReportRequest& request, const ClusterInfo* cluster_info,

                   std::string_view name) {

    TMasterResult result;

    Status status = MasterServerClient::instance()->report(request, &result);

    if (!status.ok()) [[unlikely]] {

        LOG_WARNING("failed to report {}", name)

                .tag("host", cluster_info->master_fe_addr.hostname)

                .tag("port", cluster_info->master_fe_addr.port)

                .error(status);

        return false;

    }

    else if (result.status.status_code != TStatusCode::OK) [[unlikely]] {

        LOG_WARNING("failed to report {}", name)

                .tag("host", cluster_info->master_fe_addr.hostname)

                .tag("port", cluster_info->master_fe_addr.port)

                .error(result.status);

        return false;

    }

    return true;

}

Status _submit_task(const TAgentTaskRequest& task,

                    std::function<Status(const TAgentTaskRequest&)> submit_op) {

    const TTaskType::type task_type = task.task_type;

    int64_t signature = task.signature;

    std::string type_str;

    EnumToString(TTaskType, task_type, type_str);

    VLOG_CRITICAL << "submitting task. type=" << type_str << ", signature=" << signature;

    if (!register_task_info(task_type, signature)) {

        LOG_WARNING("failed to register task").tag("type", type_str).tag("signature", signature);

        // Duplicated task request, just return OK

        return Status::OK();

    }

    // TODO(plat1ko): check task request member

    // Set the receiving time of task so that we can determine whether it is timed out later

    // exist a path task_worker_pool <- agent_server <- backend_service <- BackendService

    // use the arg BackendService_submit_tasks_args.tasks is not const, so modify is ok

    (const_cast<TAgentTaskRequest&>(task)).__set_recv_time(time(nullptr));

    auto st = submit_op(task);

    if (!st.ok()) [[unlikely]] {

        LOG_INFO("failed to submit task").tag("type", type_str).tag("signature", signature);

        return st;

    }

    LOG_INFO("successfully submit task").tag("type", type_str).tag("signature", signature);

    return Status::OK();

}

bvar::LatencyRecorder g_publish_version_latency("doris_pk", "publish_version");

bvar::Adder<uint64_t> ALTER_INVERTED_INDEX_count("task", "ALTER_INVERTED_INDEX");

bvar::Adder<uint64_t> CHECK_CONSISTENCY_count("task", "CHECK_CONSISTENCY");

bvar::Adder<uint64_t> UPLOAD_count("task", "UPLOAD");

bvar::Adder<uint64_t> DOWNLOAD_count("task", "DOWNLOAD");

bvar::Adder<uint64_t> MAKE_SNAPSHOT_count("task", "MAKE_SNAPSHOT");

bvar::Adder<uint64_t> RELEASE_SNAPSHOT_count("task", "RELEASE_SNAPSHOT");

bvar::Adder<uint64_t> MOVE_count("task", "MOVE");

bvar::Adder<uint64_t> COMPACTION_count("task", "COMPACTION");

bvar::Adder<uint64_t> PUSH_STORAGE_POLICY_count("task", "PUSH_STORAGE_POLICY");

bvar::Adder<uint64_t> PUSH_INDEX_POLICY_count("task", "PUSH_INDEX_POLICY");

bvar::Adder<uint64_t> PUSH_COOLDOWN_CONF_count("task", "PUSH_COOLDOWN_CONF");

bvar::Adder<uint64_t> CREATE_count("task", "CREATE_TABLE");

bvar::Adder<uint64_t> DROP_count("task", "DROP_TABLE");

bvar::Adder<uint64_t> PUBLISH_VERSION_count("task", "PUBLISH_VERSION");

bvar::Adder<uint64_t> CLEAR_TRANSACTION_TASK_count("task", "CLEAR_TRANSACTION_TASK");

bvar::Adder<uint64_t> DELETE_count("task", "DELETE");

bvar::Adder<uint64_t> PUSH_count("task", "PUSH");

bvar::Adder<uint64_t> UPDATE_TABLET_META_INFO_count("task", "UPDATE_TABLET_META_INFO");

bvar::Adder<uint64_t> ALTER_count("task", "ALTER_TABLE");

bvar::Adder<uint64_t> CLONE_count("task", "CLONE");

bvar::Adder<uint64_t> STORAGE_MEDIUM_MIGRATE_count("task", "STORAGE_MEDIUM_MIGRATE");

bvar::Adder<uint64_t> GC_BINLOG_count("task", "GC_BINLOG");

bvar::Adder<uint64_t> UPDATE_VISIBLE_VERSION_count("task", "UPDATE_VISIBLE_VERSION");

bvar::Adder<uint64_t> CALCULATE_DELETE_BITMAP_count("task", "CALCULATE_DELETE_BITMAP");

void add_task_count(const TAgentTaskRequest& task, int n) {

    // clang-format off

    switch (task.task_type) {

    #define ADD_TASK_COUNT(type) \

    case TTaskType::type:        \

        type##_count << n;       \

        return;

    ADD_TASK_COUNT(ALTER_INVERTED_INDEX)

    ADD_TASK_COUNT(CHECK_CONSISTENCY)

    ADD_TASK_COUNT(UPLOAD)

    ADD_TASK_COUNT(DOWNLOAD)

    ADD_TASK_COUNT(MAKE_SNAPSHOT)

    ADD_TASK_COUNT(RELEASE_SNAPSHOT)

    ADD_TASK_COUNT(MOVE)

    ADD_TASK_COUNT(COMPACTION)

    ADD_TASK_COUNT(PUSH_STORAGE_POLICY)

    ADD_TASK_COUNT(PUSH_INDEX_POLICY)

    ADD_TASK_COUNT(PUSH_COOLDOWN_CONF)

    ADD_TASK_COUNT(CREATE)

    ADD_TASK_COUNT(DROP)

    ADD_TASK_COUNT(PUBLISH_VERSION)

    ADD_TASK_COUNT(CLEAR_TRANSACTION_TASK)

    ADD_TASK_COUNT(UPDATE_TABLET_META_INFO)

    ADD_TASK_COUNT(CLONE)

    ADD_TASK_COUNT(STORAGE_MEDIUM_MIGRATE)

    ADD_TASK_COUNT(GC_BINLOG)

    ADD_TASK_COUNT(UPDATE_VISIBLE_VERSION)

    ADD_TASK_COUNT(CALCULATE_DELETE_BITMAP)

    #undef ADD_TASK_COUNT

    case TTaskType::REALTIME_PUSH:

    case TTaskType::PUSH:

        if (task.push_req.push_type == TPushType::LOAD_V2) {

            PUSH_count << n;

        } else if (task.push_req.push_type == TPushType::DELETE) {

            DELETE_count << n;

        }

        return;

    case TTaskType::ALTER:

    {

        ALTER_count << n;

        // cloud auto stop need sc jobs, a tablet's sc can also be considered a fragment

        if (n > 0) {

            // only count fragment when task is actually starting

            doris::g_fragment_executing_count << 1;

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

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

                                .count();

            g_fragment_last_active_time.set_value(now);

        }

        return;

    }

    default:

        return;

    }

    // clang-format on

}

bvar::Adder<uint64_t> report_task_total("report", "task_total");

bvar::Adder<uint64_t> report_task_failed("report", "task_failed");

bvar::Adder<uint64_t> report_disk_total("report", "disk_total");

bvar::Adder<uint64_t> report_disk_failed("report", "disk_failed");

bvar::Adder<uint64_t> report_tablet_total("report", "tablet_total");

bvar::Adder<uint64_t> report_tablet_failed("report", "tablet_failed");

bvar::Adder<uint64_t> report_index_policy_total("report", "index_policy_total");

bvar::Adder<uint64_t> report_index_policy_failed("report", "index_policy_failed");

} // namespace

TaskWorkerPool::TaskWorkerPool(

        std::string_view name, int worker_count,

        std::function<void(const TAgentTaskRequest& task)> callback,

        std::function<void(const TAgentTaskRequest& task)> pre_submit_callback)

        : _callback(std::move(callback)), _pre_submit_callback(std::move(pre_submit_callback)) {

    auto st = ThreadPoolBuilder(fmt::format("TaskWP_{}", name))

                      .set_min_threads(worker_count)

                      .set_max_threads(worker_count)

                      .build(&_thread_pool);

    CHECK(st.ok()) << name << ": " << st;

}

TaskWorkerPool::~TaskWorkerPool() {

    stop();

}

void TaskWorkerPool::stop() {

    if (_stopped.exchange(true)) {

        return;

    }

    if (_thread_pool) {

        _thread_pool->shutdown();

    }

}

Status TaskWorkerPool::submit_task(const TAgentTaskRequest& task) {

    return _submit_task(task, [this](auto&& task) {

        if (_pre_submit_callback) {

            _pre_submit_callback(task);

        }

        add_task_count(task, 1);

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

            _callback(task);

            add_task_count(task, -1);

        });

    });

}

PriorTaskWorkerPool::PriorTaskWorkerPool(

        const std::string& name, int normal_worker_count, int high_prior_worker_count,

        std::function<void(const TAgentTaskRequest& task)> callback)

        : _callback(std::move(callback)) {

    for (int i = 0; i < normal_worker_count; ++i) {

        auto st = Thread::create(

                "Normal", name, [this] { normal_loop(); }, &_workers.emplace_back());

        CHECK(st.ok()) << name << ": " << st;

    }

    for (int i = 0; i < high_prior_worker_count; ++i) {

        auto st = Thread::create(

                "HighPrior", name, [this] { high_prior_loop(); }, &_workers.emplace_back());

        CHECK(st.ok()) << name << ": " << st;

    }

}

PriorTaskWorkerPool::~PriorTaskWorkerPool() {

    stop();

}

void PriorTaskWorkerPool::stop() {

    {

        std::lock_guard lock(_mtx);

        if (_stopped) {

            return;

        }

        _stopped = true;

    }

    _normal_condv.notify_all();

    _high_prior_condv.notify_all();

    for (auto&& w : _workers) {

        if (w) {

            w->join();

        }

    }

}

Status PriorTaskWorkerPool::submit_task(const TAgentTaskRequest& task) {

    return _submit_task(task, [this](auto&& task) {

        auto req = std::make_unique<TAgentTaskRequest>(task);

        add_task_count(*req, 1);

        if (req->__isset.priority && req->priority == TPriority::HIGH) {

            std::lock_guard lock(_mtx);

            _high_prior_queue.push_back(std::move(req));

            _high_prior_condv.notify_one();

            _normal_condv.notify_one();

        } else {

            std::lock_guard lock(_mtx);

            _normal_queue.push_back(std::move(req));

            _normal_condv.notify_one();

        }

        return Status::OK();

    });

}

Status PriorTaskWorkerPool::submit_high_prior_and_cancel_low(TAgentTaskRequest& task) {

    const TTaskType::type task_type = task.task_type;

    int64_t signature = task.signature;

    std::string type_str;

    EnumToString(TTaskType, task_type, type_str);

    auto req = std::make_unique<TAgentTaskRequest>(task);

    DCHECK(req->__isset.priority && req->priority == TPriority::HIGH);

    do {

        std::lock_guard lock(s_task_signatures_mtx);

        auto& set = s_task_signatures[task_type];

        if (!set.contains(signature)) {

            // If it doesn't exist, put it directly into the priority queue

            add_task_count(*req, 1);

            set.insert(signature);

            std::lock_guard temp_lock(_mtx);

            _high_prior_queue.push_back(std::move(req));

            _high_prior_condv.notify_one();

            _normal_condv.notify_one();

            break;

        } else {

            std::lock_guard temp_lock(_mtx);

            for (auto it = _normal_queue.begin(); it != _normal_queue.end();) {

                // If it exists in the normal queue, cancel the task in the normal queue

                if ((*it)->signature == signature) {

                    _normal_queue.erase(it);                     // cancel the original task

                    _high_prior_queue.push_back(std::move(req)); // add the new task to the queue

                    _high_prior_condv.notify_one();

                    _normal_condv.notify_one();

                    break;

                } else {

                    ++it; // doesn't meet the condition, continue to the next one

                }

            }

            // If it exists in the high priority queue, no operation is needed

            LOG_INFO("task has already existed in high prior queue.").tag("signature", signature);

        }

    } while (false);

    // Set the receiving time of task so that we can determine whether it is timed out later

    task.__set_recv_time(time(nullptr));

    LOG_INFO("successfully submit task").tag("type", type_str).tag("signature", signature);

    return Status::OK();

}

void PriorTaskWorkerPool::normal_loop() {

    while (true) {

        std::unique_ptr<TAgentTaskRequest> req;

        {

            std::unique_lock lock(_mtx);

            _normal_condv.wait(lock, [&] {

                return !_normal_queue.empty() || !_high_prior_queue.empty() || _stopped;

            });

            if (_stopped) {

                return;

            }

            if (!_high_prior_queue.empty()) {

                req = std::move(_high_prior_queue.front());

                _high_prior_queue.pop_front();

            } else if (!_normal_queue.empty()) {

                req = std::move(_normal_queue.front());

                _normal_queue.pop_front();

            } else {

                continue;

            }

        }

        _callback(*req);

        add_task_count(*req, -1);

    }

}

void PriorTaskWorkerPool::high_prior_loop() {

    while (true) {

        std::unique_ptr<TAgentTaskRequest> req;

        {

            std::unique_lock lock(_mtx);

            _high_prior_condv.wait(lock, [&] { return !_high_prior_queue.empty() || _stopped; });

            if (_stopped) {

                return;

            }

            if (_high_prior_queue.empty()) {

                continue;

            }

            req = std::move(_high_prior_queue.front());

            _high_prior_queue.pop_front();

        }

        _callback(*req);

        add_task_count(*req, -1);

    }

}

ReportWorker::ReportWorker(std::string name, const ClusterInfo* cluster_info, int report_interval_s,

                           std::function<void()> callback)

        : _name(std::move(name)) {

    auto report_loop = [this, cluster_info, report_interval_s, callback = std::move(callback)] {

        auto& engine = ExecEnv::GetInstance()->storage_engine();

        engine.register_report_listener(this);

        while (true) {

            {

                std::unique_lock lock(_mtx);

                _condv.wait_for(lock, std::chrono::seconds(report_interval_s),

                                [&] { return _stopped || _signal; });

                if (_stopped) {

                    break;

                }

                if (_signal) {

                    // Consume received signal

                    _signal = false;

                }

            }

            if (cluster_info->master_fe_addr.port == 0) {

                // port == 0 means not received heartbeat yet

                LOG(INFO) << "waiting to receive first heartbeat from frontend before doing report";

                continue;

            }

            callback();

        }

        engine.deregister_report_listener(this);

    };

    auto st = Thread::create("ReportWorker", _name, report_loop, &_thread);

    CHECK(st.ok()) << _name << ": " << st;

}

ReportWorker::~ReportWorker() {

    stop();

}

void ReportWorker::notify() {

    {

        std::lock_guard lock(_mtx);

        _signal = true;

    }

    _condv.notify_all();

}

void ReportWorker::stop() {

    {

        std::lock_guard lock(_mtx);

        if (_stopped) {

            return;

        }

        _stopped = true;

    }

    _condv.notify_all();

    if (_thread) {

        _thread->join();

    }

}

void alter_cloud_index_callback(CloudStorageEngine& engine, const TAgentTaskRequest& req) {

    const auto& alter_inverted_index_rq = req.alter_inverted_index_req;

    LOG(INFO) << "[index_change]get alter index task. signature=" << req.signature

              << ", tablet_id=" << alter_inverted_index_rq.tablet_id

              << ", job_id=" << alter_inverted_index_rq.job_id;

    Status status = Status::OK();

    auto tablet_ptr = engine.tablet_mgr().get_tablet(alter_inverted_index_rq.tablet_id);

    if (tablet_ptr != nullptr) {

        EngineCloudIndexChangeTask engine_task(engine, req.alter_inverted_index_req);

        status = engine_task.execute();

    } else {

        status = Status::NotFound("could not find tablet {}", alter_inverted_index_rq.tablet_id);

    }

    // Return result to fe

    TFinishTaskRequest finish_task_request;

    finish_task_request.__set_backend(BackendOptions::get_local_backend());

    finish_task_request.__set_task_type(req.task_type);

    finish_task_request.__set_signature(req.signature);

    if (!status.ok()) {

        LOG(WARNING) << "[index_change]failed to alter inverted index task, signature="

                     << req.signature << ", tablet_id=" << alter_inverted_index_rq.tablet_id

                     << ", job_id=" << alter_inverted_index_rq.job_id << ", error=" << status;

    } else {

        LOG(INFO) << "[index_change]successfully alter inverted index task, signature="

                  << req.signature << ", tablet_id=" << alter_inverted_index_rq.tablet_id

                  << ", job_id=" << alter_inverted_index_rq.job_id;

    }

    finish_task_request.__set_task_status(status.to_thrift());

    finish_task(finish_task_request);

    remove_task_info(req.task_type, req.signature);

}

void alter_inverted_index_callback(StorageEngine& engine, const TAgentTaskRequest& req) {

    const auto& alter_inverted_index_rq = req.alter_inverted_index_req;

    LOG(INFO) << "get alter inverted index task. signature=" << req.signature

              << ", tablet_id=" << alter_inverted_index_rq.tablet_id

              << ", job_id=" << alter_inverted_index_rq.job_id;

    Status status = Status::OK();

    auto tablet_ptr = engine.tablet_manager()->get_tablet(alter_inverted_index_rq.tablet_id);

    if (tablet_ptr != nullptr) {

        EngineIndexChangeTask engine_task(engine, alter_inverted_index_rq);

        SCOPED_ATTACH_TASK(engine_task.mem_tracker());

        status = engine_task.execute();

    } else {

        status = Status::NotFound("could not find tablet {}", alter_inverted_index_rq.tablet_id);

    }

    // Return result to fe

    TFinishTaskRequest finish_task_request;

    finish_task_request.__set_backend(BackendOptions::get_local_backend());

    finish_task_request.__set_task_type(req.task_type);

    finish_task_request.__set_signature(req.signature);

    std::vector<TTabletInfo> finish_tablet_infos;

    if (!status.ok()) {

        LOG(WARNING) << "failed to alter inverted index task, signature=" << req.signature

                     << ", tablet_id=" << alter_inverted_index_rq.tablet_id

                     << ", job_id=" << alter_inverted_index_rq.job_id << ", error=" << status;

    } else {

        LOG(INFO) << "successfully alter inverted index task, signature=" << req.signature

                  << ", tablet_id=" << alter_inverted_index_rq.tablet_id

                  << ", job_id=" << alter_inverted_index_rq.job_id;

        TTabletInfo tablet_info;

        status = get_tablet_info(engine, alter_inverted_index_rq.tablet_id,

                                 alter_inverted_index_rq.schema_hash, &tablet_info);

        if (status.ok()) {

            finish_tablet_infos.push_back(tablet_info);

        }

        finish_task_request.__set_finish_tablet_infos(finish_tablet_infos);

    }

    finish_task_request.__set_task_status(status.to_thrift());

    finish_task(finish_task_request);

    remove_task_info(req.task_type, req.signature);

}

void update_tablet_meta_callback(StorageEngine& engine, const TAgentTaskRequest& req) {

    LOG(INFO) << "get update tablet meta task. signature=" << req.signature;

    Status status;

    const auto& update_tablet_meta_req = req.update_tablet_meta_info_req;

    for (const auto& tablet_meta_info : update_tablet_meta_req.tabletMetaInfos) {

        auto tablet = engine.tablet_manager()->get_tablet(tablet_meta_info.tablet_id);

        if (tablet == nullptr) {

            status = Status::NotFound("tablet not found");

            LOG(WARNING) << "could not find tablet when update tablet meta. tablet_id="

                         << tablet_meta_info.tablet_id;

            continue;

        }

        bool need_to_save = false;

        if (tablet_meta_info.__isset.partition_id) {

            // for fix partition_id = 0

            LOG(WARNING) << "change be tablet id: " << tablet->tablet_meta()->tablet_id()

                         << "partition id from : " << tablet->tablet_meta()->partition_id()

                         << " to : " << tablet_meta_info.partition_id;

            auto succ = engine.tablet_manager()->update_tablet_partition_id(

                    tablet_meta_info.partition_id, tablet->tablet_meta()->tablet_id());

            if (!succ) {

                std::string err_msg = fmt::format(

                        "change be tablet id : {} partition_id : {} failed",

                        tablet->tablet_meta()->tablet_id(), tablet_meta_info.partition_id);

                LOG(WARNING) << err_msg;

                status = Status::InvalidArgument(err_msg);

                continue;

            }

            need_to_save = true;

        }

        if (tablet_meta_info.__isset.storage_policy_id) {

            tablet->tablet_meta()->set_storage_policy_id(tablet_meta_info.storage_policy_id);

            need_to_save = true;

        }

        if (tablet_meta_info.__isset.is_in_memory) {

            tablet->tablet_meta()->mutable_tablet_schema()->set_is_in_memory(

                    tablet_meta_info.is_in_memory);

            std::shared_lock rlock(tablet->get_header_lock());

            for (auto& [_, rowset_meta] : tablet->tablet_meta()->all_mutable_rs_metas()) {

                rowset_meta->tablet_schema()->set_is_in_memory(tablet_meta_info.is_in_memory);

            }

            tablet->tablet_schema_unlocked()->set_is_in_memory(tablet_meta_info.is_in_memory);

            need_to_save = true;

        }

        if (tablet_meta_info.__isset.compaction_policy) {

            if (tablet_meta_info.compaction_policy != CUMULATIVE_SIZE_BASED_POLICY &&

                tablet_meta_info.compaction_policy != CUMULATIVE_TIME_SERIES_POLICY) {

                status = Status::InvalidArgument(

                        "invalid compaction policy, only support for size_based or "

                        "time_series");

                continue;

            }

            tablet->tablet_meta()->set_compaction_policy(tablet_meta_info.compaction_policy);

            need_to_save = true;

        }

        if (tablet_meta_info.__isset.time_series_compaction_goal_size_mbytes) {

            if (tablet->tablet_meta()->compaction_policy() != CUMULATIVE_TIME_SERIES_POLICY) {

                status = Status::InvalidArgument(

                        "only time series compaction policy support time series config");

                continue;

            }

            tablet->tablet_meta()->set_time_series_compaction_goal_size_mbytes(

                    tablet_meta_info.time_series_compaction_goal_size_mbytes);

            need_to_save = true;

        }

        if (tablet_meta_info.__isset.time_series_compaction_file_count_threshold) {

            if (tablet->tablet_meta()->compaction_policy() != CUMULATIVE_TIME_SERIES_POLICY) {

                status = Status::InvalidArgument(

                        "only time series compaction policy support time series config");

                continue;

            }

            tablet->tablet_meta()->set_time_series_compaction_file_count_threshold(

                    tablet_meta_info.time_series_compaction_file_count_threshold);