// 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 "exec/scan/scanner_scheduler.h"

#include <algorithm>

#include <cstdint>

#include <functional>

#include <list>

#include <memory>

#include <ostream>

#include <string>

#include <utility>

#include "common/compiler_util.h" // IWYU pragma: keep

#include "common/config.h"

#include "common/exception.h"

#include "common/logging.h"

#include "common/status.h"

#include "core/block/block.h"

#include "exec/pipeline/pipeline_task.h"

#include "exec/scan/file_scanner.h"

#include "exec/scan/olap_scanner.h" // IWYU pragma: keep

#include "exec/scan/scan_node.h"

#include "exec/scan/scanner.h"

#include "exec/scan/scanner_context.h"

#include "runtime/exec_env.h"

#include "runtime/runtime_state.h"

#include "runtime/thread_context.h"

#include "runtime/workload_group/workload_group_manager.h"

#include "storage/tablet/tablet.h"

#include "util/async_io.h" // IWYU pragma: keep

#include "util/cpu_info.h"

#include "util/defer_op.h"

#include "util/thread.h"

#include "util/threadpool.h"

namespace doris {

Status ScannerScheduler::submit(std::shared_ptr<ScannerContext> ctx,

                                std::shared_ptr<ScanTask> scan_task) {

    if (ctx->done()) {

        return Status::OK();

    }

    auto task_lock = ctx->task_exec_ctx();

    if (task_lock == nullptr) {

        LOG(INFO) << "could not lock task execution context, query " << ctx->debug_string()

                  << " maybe finished";

        return Status::OK();

    }

    std::shared_ptr<ScannerDelegate> scanner_delegate = scan_task->scanner.lock();

    if (scanner_delegate == nullptr) {

        return Status::OK();

    }

    scanner_delegate->_scanner->start_wait_worker_timer();

    TabletStorageType type = scanner_delegate->_scanner->get_storage_type();

    auto sumbit_task = [&]() {

        auto work_func = [scanner_ref = scan_task, ctx]() {

            auto status = [&] {

                RETURN_IF_CATCH_EXCEPTION(_scanner_scan(ctx, scanner_ref));

                return Status::OK();

            }();

            if (!status.ok()) {

                scanner_ref->set_status(status);

                ctx->push_back_scan_task(scanner_ref);

                return true;

            }

            return scanner_ref->is_eos();

        };

        SimplifiedScanTask simple_scan_task = {work_func, ctx, scan_task};

        return this->submit_scan_task(simple_scan_task);

    };

    Status submit_status = sumbit_task();

    if (!submit_status.ok()) {

        // User will see TooManyTasks error. It looks like a more reasonable error.

        Status scan_task_status = Status::TooManyTasks(

                "Failed to submit scanner to scanner pool reason:" +

                std::string(submit_status.msg()) + "|type:" + std::to_string(type));

        scan_task->set_status(scan_task_status);

        return scan_task_status;

    }

    return Status::OK();

}

void handle_reserve_memory_failure(RuntimeState* state, std::shared_ptr<ScannerContext> ctx,

                                   const Status& st, size_t reserve_size) {

    ctx->clear_free_blocks();

    auto* local_state = ctx->local_state();

    auto debug_msg = fmt::format(

            "Query: {} , scanner try to reserve: {}, operator name {}, "

            "operator "

            "id: {}, "

            "task id: "

            "{}, failed: {}",

            print_id(state->query_id()), PrettyPrinter::print_bytes(reserve_size),

            local_state->get_name(), local_state->parent()->node_id(), state->task_id(),

            st.to_string());

    // PROCESS_MEMORY_EXCEEDED error msg alread contains process_mem_log_str

    if (!st.is<ErrorCode::PROCESS_MEMORY_EXCEEDED>()) {

        debug_msg += fmt::format(", debug info: {}", GlobalMemoryArbitrator::process_mem_log_str());

    }

    VLOG_DEBUG << debug_msg;

    state->get_query_ctx()->set_low_memory_mode();

}

void ScannerScheduler::_scanner_scan(std::shared_ptr<ScannerContext> ctx,

                                     std::shared_ptr<ScanTask> scan_task) {

    auto task_lock = ctx->task_exec_ctx();

    if (task_lock == nullptr) {

        return;

    }

    SCOPED_ATTACH_TASK(ctx->state());

    ctx->update_peak_running_scanner(1);

    Defer defer([&] { ctx->update_peak_running_scanner(-1); });

    std::shared_ptr<ScannerDelegate> scanner_delegate = scan_task->scanner.lock();

    if (scanner_delegate == nullptr) {

        return;

    }

    ScannerSPtr& scanner = scanner_delegate->_scanner;

    // for cpu hard limit, thread name should not be reset

    if (ctx->_should_reset_thread_name) {

        Thread::set_self_name("_scanner_scan");

    }

#ifndef __APPLE__

    // The configuration item is used to lower the priority of the scanner thread,

    // typically employed to ensure CPU scheduling for write operations.

    if (config::scan_thread_nice_value != 0 && scanner->get_name() != FileScanner::NAME) {

        Thread::set_thread_nice_value();

    }

#endif

    // we set and get counter according below order, to make sure the counter is updated before get_block, and the time of get_block is recorded in the counter.

    // 1. update_wait_worker_timer to make sure the time of waiting for worker thread is recorded in the timer

    // 2. start_scan_cpu_timer to make sure the cpu timer include the time of open and get_block, which is the real cpu time of scanner

    // 3. update_scan_cpu_timer when defer, to make sure the cpu timer include the time of open and get_block, which is the real cpu time of scanner

    // 4. start_wait_worker_timer when defer, to make sure the time of waiting for worker thread is recorded in the timer

    MonotonicStopWatch max_run_time_watch;

    max_run_time_watch.start();

    scanner->update_wait_worker_timer();

    scanner->start_scan_cpu_timer();

    bool need_update_profile = true;

    auto update_scanner_profile = [&]() {

        if (need_update_profile) {

            scanner->update_scan_cpu_timer();

            scanner->update_realtime_counters();

            need_update_profile = false;

        }

    };

    Status status = Status::OK();

    bool eos = false;

    Defer defer_scanner([&] {

        if (status.ok() && !eos) {

            // if status is not ok, it means the scanner is failed, and the counter may be not updated correctly, so no need to update counter again. if eos is true, it means the scanner is finished successfully, and the counter is updated correctly, so no need to update counter again.

            scanner->start_wait_worker_timer();

        }

    });

    ASSIGN_STATUS_IF_CATCH_EXCEPTION(

            RuntimeState* state = ctx->state(); DCHECK(nullptr != state);

            // scanner->open may alloc plenty amount of memory(read blocks of data),

            // so better to also check low memory and clear free blocks here.

            if (ctx->low_memory_mode()) { ctx->clear_free_blocks(); }

            if (!scanner->has_prepared()) {

                status = scanner->prepare();

                if (!status.ok()) {

                    eos = true;

                }

            }

            if (!eos && !scanner->is_open()) {

                status = scanner->open(state);

                if (!status.ok()) {

                    eos = true;

                }

                scanner->set_opened();

            }

            Status rf_status = scanner->try_append_late_arrival_runtime_filter();

            if (!rf_status.ok()) {

                LOG(WARNING) << "Failed to append late arrival runtime filter: "

                             << rf_status.to_string();

            }

            size_t raw_bytes_threshold = config::doris_scanner_row_bytes;

            if (ctx->low_memory_mode()) {

                ctx->clear_free_blocks();

                if (raw_bytes_threshold > ctx->low_memory_mode_scan_bytes_per_scanner()) {

                    raw_bytes_threshold = ctx->low_memory_mode_scan_bytes_per_scanner();

                }

            }

            size_t raw_bytes_read = 0;

            bool first_read = true; int64_t limit = scanner->limit();

            // If the first block is full, then it is true. Or the first block + second block > batch_size

            bool has_first_full_block = false;

            // During low memory mode, every scan task will return at most 2 block to reduce memory usage.

            while (!eos && raw_bytes_read < raw_bytes_threshold &&

                   (!ctx->low_memory_mode() || !has_first_full_block) &&

                   (!has_first_full_block || doris::thread_context()

                                                     ->thread_mem_tracker_mgr->limiter_mem_tracker()

                                                     ->check_limit(1))) {

                if (UNLIKELY(ctx->done())) {

                    eos = true;

                    break;

                }

                // If shared limit quota is exhausted, stop scanning.

                if (ctx->remaining_limit() == 0) {

                    eos = true;

                    break;

                }

                if (max_run_time_watch.elapsed_time() >

                    config::doris_scanner_max_run_time_ms * 1e6) {

                    break;

                }

                DEFER_RELEASE_RESERVED();

                BlockUPtr free_block;

                if (first_read) {

                    free_block = ctx->get_free_block(first_read);

                } else {

                    if (state->get_query_ctx()

                                ->resource_ctx()

                                ->task_controller()

                                ->is_enable_reserve_memory()) {

                        size_t block_avg_bytes = scanner->get_block_avg_bytes();

                        auto st = thread_context()->thread_mem_tracker_mgr->try_reserve(

                                block_avg_bytes);

                        if (!st.ok()) {

                            handle_reserve_memory_failure(state, ctx, st, block_avg_bytes);

                            break;

                        }

                    }

                    free_block = ctx->get_free_block(first_read);

                }

                if (free_block == nullptr) {

                    break;

                }

                // We got a new created block or a reused block.

                status = scanner->get_block_after_projects(state, free_block.get(), &eos);

                first_read = false;

                if (!status.ok()) {

                    LOG(WARNING) << "Scan thread read Scanner failed: " << status.to_string();

                    break;

                }

                // Check column type only after block is read successfully.

                // Or it may cause a crash when the block is not normal.

                _make_sure_virtual_col_is_materialized(scanner, free_block.get());

                // Shared limit quota: acquire rows from the context's shared pool.

                // Discard or truncate the block if quota is exhausted.

                if (free_block->rows() > 0) {

                    int64_t block_rows = free_block->rows();

                    int64_t granted = ctx->acquire_limit_quota(block_rows);

                    if (granted == 0) {

                        // No quota remaining, discard this block and mark eos.

                        ctx->return_free_block(std::move(free_block));

                        eos = true;

                        break;

                    } else if (granted < block_rows) {

                        // Partial quota: truncate block to granted rows and mark eos.

                        free_block->set_num_rows(granted);

                        eos = true;

                    }

                }

                // Projection will truncate useless columns, makes block size change.

                auto free_block_bytes = free_block->allocated_bytes();

                raw_bytes_read += free_block_bytes;

                if (!scan_task->cached_blocks.empty() &&

                    scan_task->cached_blocks.back().first->rows() + free_block->rows() <=

                            ctx->batch_size()) {

                    size_t block_size = scan_task->cached_blocks.back().first->allocated_bytes();

                    MutableBlock mutable_block(scan_task->cached_blocks.back().first.get());

                    status = mutable_block.merge(*free_block);

                    if (!status.ok()) {

                        LOG(WARNING) << "Block merge failed: " << status.to_string();

                        break;

                    }

                    scan_task->cached_blocks.back().second = mutable_block.allocated_bytes();

                    scan_task->cached_blocks.back().first.get()->set_columns(

                            std::move(mutable_block.mutable_columns()));

                    // Return block succeed or not, this free_block is not used by this scan task any more.

                    // If block can be reused, its memory usage will be added back.

                    ctx->return_free_block(std::move(free_block));

                    ctx->inc_block_usage(scan_task->cached_blocks.back().first->allocated_bytes() -

                                         block_size);

                } else {

                    if (!scan_task->cached_blocks.empty()) {

                        has_first_full_block = true;

                    }

                    ctx->inc_block_usage(free_block->allocated_bytes());

                    scan_task->cached_blocks.emplace_back(std::move(free_block), free_block_bytes);

                }

                // Per-scanner small-limit optimization: if limit is small (< batch_size),

                // return immediately instead of accumulating to raw_bytes_threshold.

                if (limit > 0 && limit < ctx->batch_size()) {

                    break;

                }

                if (scan_task->cached_blocks.back().first->rows() > 0) {

                    auto block_avg_bytes = (scan_task->cached_blocks.back().first->bytes() +

                                            scan_task->cached_blocks.back().first->rows() - 1) /

                                           scan_task->cached_blocks.back().first->rows() *

                                           ctx->batch_size();

                    scanner->update_block_avg_bytes(block_avg_bytes);

                }

                if (ctx->low_memory_mode()) {

                    ctx->clear_free_blocks();

                    if (raw_bytes_threshold > ctx->low_memory_mode_scan_bytes_per_scanner()) {

                        raw_bytes_threshold = ctx->low_memory_mode_scan_bytes_per_scanner();

                    }

                }

            } // end for while

            if (UNLIKELY(!status.ok())) {

                scan_task->set_status(status);

                eos = true;

            },

            status);

    if (UNLIKELY(!status.ok())) {

        scan_task->set_status(status);

        eos = true;

    }

    if (eos) {

        // If eos, scanner will call _collect_profile_before_close to update profile,

        // so we need update_scanner_profile here

        update_scanner_profile();

        scanner->mark_to_need_to_close();

    }

    scan_task->set_eos(eos);

    VLOG_DEBUG << fmt::format(

            "Scanner context {} has finished task, cached_block {} current scheduled task is "

            "{}, eos: {}, status: {}",

            ctx->ctx_id, scan_task->cached_blocks.size(), ctx->num_scheduled_scanners(), eos,

            status.to_string());

    ctx->push_back_scan_task(scan_task);

}

int ScannerScheduler::default_local_scan_thread_num() {

    return config::doris_scanner_thread_pool_thread_num > 0

                   ? config::doris_scanner_thread_pool_thread_num

                   : std::max(48, CpuInfo::num_cores() * 2);

}

int ScannerScheduler::default_remote_scan_thread_num() {

    int num = config::doris_max_remote_scanner_thread_pool_thread_num > 0

                      ? config::doris_max_remote_scanner_thread_pool_thread_num

                      : std::max(512, CpuInfo::num_cores() * 10);

    return std::max(num, default_local_scan_thread_num());

}

int ScannerScheduler::get_remote_scan_thread_queue_size() {

    return config::doris_remote_scanner_thread_pool_queue_size;

}

int ScannerScheduler::default_min_active_scan_threads() {

    return config::min_active_scan_threads > 0

                   ? config::min_active_scan_threads

                   : config::min_active_scan_threads = CpuInfo::num_cores() * 2;

}

int ScannerScheduler::default_min_active_file_scan_threads() {

    return config::min_active_file_scan_threads > 0

                   ? config::min_active_file_scan_threads

                   : config::min_active_file_scan_threads = CpuInfo::num_cores() * 8;

}

void ScannerScheduler::_make_sure_virtual_col_is_materialized(

        const std::shared_ptr<Scanner>& scanner, Block* free_block) {

#ifndef NDEBUG

    // Currently, virtual column can only be used on olap table.

    std::shared_ptr<OlapScanner> olap_scanner = std::dynamic_pointer_cast<OlapScanner>(scanner);

    if (olap_scanner == nullptr) {

        return;

    }

    if (free_block->rows() == 0) {

        return;

    }

    size_t idx = 0;

    for (const auto& entry : *free_block) {

        // Virtual column must be materialized on the end of SegmentIterator's next batch method.

        const ColumnNothing* column_nothing =

                check_and_get_column<ColumnNothing>(entry.column.get());

        if (column_nothing == nullptr) {

            idx++;

            continue;

        }

        std::vector<std::string> vcid_to_idx;

        for (const auto& pair : olap_scanner->_vir_cid_to_idx_in_block) {

            vcid_to_idx.push_back(fmt::format("{}-{}", pair.first, pair.second));

        }

        std::string error_msg = fmt::format(

                "Column in idx {} is nothing, block columns {}, normal_columns "

                "{}, "

                "vir_cid_to_idx_in_block_msg {}",

                idx, free_block->columns(), olap_scanner->_return_columns.size(),

                fmt::format("_vir_cid_to_idx_in_block:[{}]", fmt::join(vcid_to_idx, ",")));

        throw doris::Exception(ErrorCode::INTERNAL_ERROR, error_msg);

    }

#endif

}

Result<SharedListenableFuture<Void>> ScannerSplitRunner::process_for(std::chrono::nanoseconds) {

    _started = true;

    bool is_completed = _scan_func();

    if (is_completed) {

        _completion_future.set_value(Void {});

    }

    return SharedListenableFuture<Void>::create_ready(Void {});

}

bool ScannerSplitRunner::is_finished() {

    return _completion_future.is_done();

}

Status ScannerSplitRunner::finished_status() {

    return _completion_future.get_status();

}

bool ScannerSplitRunner::is_started() const {

    return _started.load();

}

bool ScannerSplitRunner::is_auto_reschedule() const {

    return false;

}

} // namespace doris