From b635ab2f7668372ddfe37f7d685f838feb6da675 Mon Sep 17 00:00:00 2001 From: Hu Shenggang Date: Mon, 17 Aug 2026 17:04:30 +0800 Subject: [PATCH] [fix](be) Account for TaskExecutor admission in scanner scheduling ### What problem does this PR solve? Issue Number: None Related PR: None Problem Summary: ScannerContext counted every submitted scanner task as active even when TaskExecutor kept excess leaf splits in its per-task admission queue. If admitted non-EOS scanners yielded while retaining their TaskExecutor slots, the hidden queued tasks could keep ScannerContext's scheduling margin closed and prevent those scanners from being re-enqueued. Expose the unadmitted leaf split count from TaskHandle and use the admission-aware margin only to re-enqueue splits already owned by TaskExecutor. First-time scanner submissions continue to use the original total in-flight margin and cannot consume admission-aware capacity. ### Release note Fix a possible file scan scheduling stall when ScannerContext concurrency exceeds TaskExecutor admission. ### Check List (For Author) - Test: - Unit Test: ScannerContextTest and TimeSharingTaskExecutorTest - Behavior changed: Yes (yielded admitted scanners can resume without increasing first-time TaskExecutor submissions) - Does this need documentation: No --- be/src/exec/scan/scanner_context.cpp | 186 +++++++++++------- be/src/exec/scan/scanner_context.h | 26 ++- be/src/exec/scan/task_executor/task_handle.h | 2 + .../time_sharing/time_sharing_task_handle.cpp | 5 + .../time_sharing/time_sharing_task_handle.h | 1 + .../time_sharing_task_executor_test.cpp | 2 + be/test/exec/scan/scanner_context_test.cpp | 142 ++++++++++++- 7 files changed, 276 insertions(+), 88 deletions(-) diff --git a/be/src/exec/scan/scanner_context.cpp b/be/src/exec/scan/scanner_context.cpp index 7fbf6ed951c997..d7e098bc2d9ae2 100644 --- a/be/src/exec/scan/scanner_context.cpp +++ b/be/src/exec/scan/scanner_context.cpp @@ -22,6 +22,8 @@ #include #include +#include +#include #include #include #include @@ -41,6 +43,7 @@ #include "exec/scan/scan_node.h" #include "exec/scan/scanner_scheduler.h" #include "exec/scan/task_executor/task_executor.h" +#include "exec/scan/task_executor/task_handle.h" #include "runtime/descriptors.h" #include "runtime/exec_env.h" #include "runtime/runtime_profile.h" @@ -178,6 +181,30 @@ int ScannerContext::_available_pickup_scanner_count() { return scanners; } +static Status init_task_executor(ScannerScheduler* scanner_scheduler, RuntimeState* state, + const std::string& ctx_id, + std::weak_ptr* task_executor, + std::shared_ptr* task_handle) { + if (auto* task_executor_scheduler = + dynamic_cast(scanner_scheduler)) { + std::shared_ptr executor = task_executor_scheduler->task_executor(); + *task_executor = executor; + TaskId task_id(fmt::format("{}-{}", print_id(state->query_id()), ctx_id)); + int initial_task_concurrency = + config::task_executor_initial_max_concurrency_per_task > 0 + ? config::task_executor_initial_max_concurrency_per_task + : std::max(48, CpuInfo::num_cores() * 2); + if (config::task_executor_max_concurrency_per_task > 0) { + initial_task_concurrency = std::min(initial_task_concurrency, + config::task_executor_max_concurrency_per_task); + } + *task_handle = DORIS_TRY(executor->create_task( + task_id, []() { return 0.0; }, initial_task_concurrency, + std::chrono::milliseconds(100), std::nullopt)); + } + return Status::OK(); +} + // After init function call, should not access _parent Status ScannerContext::init() { #ifndef BE_TEST @@ -197,20 +224,10 @@ Status ScannerContext::init() { auto scanner = _all_scanners.front().lock(); DCHECK(scanner != nullptr); - - if (auto* task_executor_scheduler = - dynamic_cast(_scanner_scheduler)) { - std::shared_ptr task_executor = task_executor_scheduler->task_executor(); - _task_executor = task_executor; - TaskId task_id(fmt::format("{}-{}", print_id(_state->query_id()), ctx_id)); - _task_handle = DORIS_TRY(task_executor->create_task( - task_id, []() { return 0.0; }, - config::task_executor_initial_max_concurrency_per_task > 0 - ? config::task_executor_initial_max_concurrency_per_task - : std::max(48, CpuInfo::num_cores() * 2), - std::chrono::milliseconds(100), std::nullopt)); - } #endif + + RETURN_IF_ERROR( + init_task_executor(_scanner_scheduler, _state, ctx_id, &_task_executor, &_task_handle)); // _max_bytes_in_queue controls the maximum memory that can be used by a single scan operator. // scan_queue_mem_limit on FE is 100MB by default, on backend we will make sure its actual value // is larger than 10MB. @@ -580,13 +597,19 @@ void ScannerContext::reestimated_block_mem_bytes(int64_t num) { int32_t ScannerContext::_get_margin(std::unique_lock& transfer_lock, std::unique_lock& scheduler_lock) { + return _get_margin(transfer_lock, scheduler_lock, _in_flight_tasks_num.load()); +} + +int32_t ScannerContext::_get_margin(std::unique_lock& transfer_lock, + std::unique_lock& scheduler_lock, + int32_t in_flight_tasks) { // Get effective max concurrency considering adaptive scheduling int32_t effective_max_concurrency = _available_pickup_scanner_count(); DCHECK_LE(effective_max_concurrency, _max_scan_concurrency); // margin_1 is used to ensure each scan operator could have at least _min_scan_concurrency scan tasks. - int32_t margin_1 = _min_scan_concurrency - - (cast_set(_completed_tasks.size()) + _in_flight_tasks_num); + int32_t margin_1 = + _min_scan_concurrency - (cast_set(_completed_tasks.size()) + in_flight_tasks); // margin_2 is used to ensure the scan scheduler could have at least _min_scan_concurrency_of_scan_scheduler scan tasks. int32_t margin_2 = @@ -596,7 +619,7 @@ int32_t ScannerContext::_get_margin(std::unique_lock& transfer_lock, // margin_3 is used to respect adaptive max concurrency limit int32_t margin_3 = std::max(effective_max_concurrency - - (cast_set(_completed_tasks.size()) + _in_flight_tasks_num), + (cast_set(_completed_tasks.size()) + in_flight_tasks), 1); if (margin_1 <= 0 && margin_2 <= 0) { @@ -611,24 +634,32 @@ int32_t ScannerContext::_get_margin(std::unique_lock& transfer_lock, if (low_memory_mode()) { // In low memory mode, we will limit the number of running scanners to `low_memory_mode_scanners()`. // So that we will not submit too many scan tasks to scheduler. - margin = std::min(low_memory_mode_scanners() - _in_flight_tasks_num, margin); + margin = std::min(low_memory_mode_scanners() - in_flight_tasks, margin); } VLOG_DEBUG << fmt::format( "[{}|{}] schedule scan task, margin_1: {} = {} - ({} + {}), margin_2: {} = {} - " "({} + {}), margin_3: {} = {} - ({} + {}), margin: {}, adaptive: {}", print_id(_query_id), ctx_id, margin_1, _min_scan_concurrency, _completed_tasks.size(), - _in_flight_tasks_num, margin_2, _min_scan_concurrency_of_scan_scheduler, + in_flight_tasks, margin_2, _min_scan_concurrency_of_scan_scheduler, _scanner_scheduler->get_active_threads(), _scanner_scheduler->get_queue_size(), - margin_3, effective_max_concurrency, _completed_tasks.size(), _in_flight_tasks_num, - margin, _enable_adaptive_scanners); + margin_3, effective_max_concurrency, _completed_tasks.size(), in_flight_tasks, margin, + _enable_adaptive_scanners); return margin; } +int32_t ScannerContext::_admitted_in_flight_tasks() const { + const int32_t queued_leaf_splits = _task_handle ? _task_handle->queued_leaf_splits() : 0; + DCHECK_LE(queued_leaf_splits, _in_flight_tasks_num); + return _in_flight_tasks_num - queued_leaf_splits; +} + // This function must be called with: // 1. _transfer_lock held. // 2. ScannerScheduler::_lock held. +// Keep the scheduling decision and submission together so both use the same locked state snapshot. +// NOLINTNEXTLINE(readability-function-cognitive-complexity,readability-function-size) Status ScannerContext::schedule_scan_task(std::shared_ptr current_scan_task, std::unique_lock& transfer_lock, std::unique_lock& scheduler_lock) { @@ -639,76 +670,74 @@ Status ScannerContext::schedule_scan_task(std::shared_ptr current_scan std::list> tasks_to_submit; - int32_t margin = _get_margin(transfer_lock, scheduler_lock); - - // margin is less than zero. Means this scan operator could not submit any scan task for now. - if (margin <= 0) { - // Be careful with current scan task. - // We need to add it back to task queue to make sure it could be resubmitted. - if (current_scan_task) { - // This usually happens when we should downgrade the concurrency. - current_scan_task->set_state(ScanTask::State::PENDING); - _pending_tasks.push(current_scan_task); - VLOG_DEBUG << fmt::format( - "{} push back scanner to task queue, because diff <= 0, _completed_tasks size " - "{}, _in_flight_tasks_num {}", - ctx_id, _completed_tasks.size(), _in_flight_tasks_num); - } + const int32_t submission_margin = std::max(_get_margin(transfer_lock, scheduler_lock), 0); + bool current_scan_task_scheduled = current_scan_task == nullptr; -#ifndef NDEBUG - // This DCHECK is necessary. - // We need to make sure each scan operator could have at least 1 scan tasks. - // Or this scan operator will not be re-scheduled. - if (!_pending_tasks.empty() && _in_flight_tasks_num == 0 && _completed_tasks.empty()) { - throw doris::Exception(ErrorCode::INTERNAL_ERROR, "Scanner scheduler logical error."); + // A yielded scanner keeps its TaskExecutor admission slot. Use the admission-aware margin + // only to re-enqueue splits that TaskExecutor already owns. In particular, a task with + // is_first_schedule=true must not consume this margin and create more TaskExecutor backlog. + int32_t admitted_in_flight_tasks = 0; + int32_t reschedule_margin = 0; + if (_task_handle) { + admitted_in_flight_tasks = _admitted_in_flight_tasks(); + reschedule_margin = + std::max(_get_margin(transfer_lock, scheduler_lock, admitted_in_flight_tasks), 0); + } + while (reschedule_margin-- > 0) { + const auto current_concurrency = cast_set( + _completed_tasks.size() + admitted_in_flight_tasks + tasks_to_submit.size()); + auto task_to_run = + _pull_next_scan_task(current_scan_task_scheduled ? nullptr : current_scan_task, + current_concurrency, true); + if (!task_to_run) { + break; } -#endif - - return Status::OK(); + current_scan_task_scheduled |= task_to_run == current_scan_task; + tasks_to_submit.push_back(std::move(task_to_run)); } - bool first_pull = true; - - while (margin-- > 0) { + // Rescheduled splits also consume ordinary ScannerContext margin. Only the remaining ordinary + // margin may pull a first-time scanner and submit a new split to TaskExecutor. + int32_t remaining_submission_margin = + std::max(submission_margin - cast_set(tasks_to_submit.size()), 0); + while (remaining_submission_margin-- > 0) { std::shared_ptr task_to_run; - const int32_t current_concurrency = cast_set( + const auto current_concurrency = cast_set( _completed_tasks.size() + _in_flight_tasks_num + tasks_to_submit.size()); - VLOG_DEBUG << fmt::format("{} currenct concurrency: {} = {} + {} + {}", ctx_id, + VLOG_DEBUG << fmt::format("{} current concurrency: {} = {} + {} + {}", ctx_id, current_concurrency, _completed_tasks.size(), _in_flight_tasks_num, tasks_to_submit.size()); - if (first_pull) { - task_to_run = _pull_next_scan_task(current_scan_task, current_concurrency); - if (task_to_run == nullptr) { - // In three situations we will get nullptr. - // 1. current_concurrency already reached _max_scan_concurrency. - // 2. all scanners are finished. - // 3. The shared LIMIT is exhausted while completed or in-flight tasks can still - // make progress. - if (current_scan_task) { - DCHECK(current_scan_task->cached_block == nullptr); - DCHECK(!current_scan_task->is_eos()); - if (current_scan_task->cached_block != nullptr || current_scan_task->is_eos()) { - // This should not happen. - throw doris::Exception(ErrorCode::INTERNAL_ERROR, - "Scanner scheduler logical error."); - } - // Current scan task is not scheduled, we need to add it back to task queue to make sure it could be resubmitted. - current_scan_task->set_state(ScanTask::State::PENDING); - _pending_tasks.push(current_scan_task); - } - } - first_pull = false; - } else { - task_to_run = _pull_next_scan_task(nullptr, current_concurrency); - } + task_to_run = _pull_next_scan_task( + current_scan_task_scheduled ? nullptr : current_scan_task, current_concurrency); if (task_to_run) { - tasks_to_submit.push_back(task_to_run); + current_scan_task_scheduled |= task_to_run == current_scan_task; + tasks_to_submit.push_back(std::move(task_to_run)); } else { break; } } + if (!current_scan_task_scheduled) { + DCHECK(current_scan_task->cached_block == nullptr); + DCHECK(!current_scan_task->is_eos()); + current_scan_task->set_state(ScanTask::State::PENDING); + _pending_tasks.push(current_scan_task); + VLOG_DEBUG << fmt::format( + "{} push back scanner to task queue, because no margin is available, " + "_completed_tasks size {}, _in_flight_tasks_num {}", + ctx_id, _completed_tasks.size(), _in_flight_tasks_num); + } + +#ifndef NDEBUG + // This DCHECK is necessary. We need to make sure each scan operator could have at least one + // task that can make progress, otherwise this scan operator will not be re-scheduled. + if (tasks_to_submit.empty() && !_pending_tasks.empty() && _in_flight_tasks_num == 0 && + _completed_tasks.empty()) { + throw doris::Exception(ErrorCode::INTERNAL_ERROR, "Scanner scheduler logical error."); + } +#endif + if (tasks_to_submit.empty()) { return Status::OK(); } @@ -730,7 +759,8 @@ Status ScannerContext::schedule_scan_task(std::shared_ptr current_scan } std::shared_ptr ScannerContext::_pull_next_scan_task( - std::shared_ptr current_scan_task, int32_t current_concurrency) { + std::shared_ptr current_scan_task, int32_t current_concurrency, + bool only_existing_split) { int32_t effective_max_concurrency = _max_scan_concurrency; if (_enable_adaptive_scanners) { effective_max_concurrency = _adaptive_processor->expected_scanners > 0 @@ -747,6 +777,7 @@ std::shared_ptr ScannerContext::_pull_next_scan_task( } if (current_scan_task != nullptr) { + DCHECK(!only_existing_split || !current_scan_task->is_first_schedule); if (current_scan_task->cached_block != nullptr || current_scan_task->is_eos()) { // This should not happen. throw doris::Exception(ErrorCode::INTERNAL_ERROR, "Scanner scheduler logical error."); @@ -755,6 +786,9 @@ std::shared_ptr ScannerContext::_pull_next_scan_task( } if (!_pending_tasks.empty()) { + if (only_existing_split && _pending_tasks.top()->is_first_schedule) { + return nullptr; + } // Do not submit more pending scanners after the shared LIMIT is exhausted while // completed or in-flight tasks can still make progress. If neither exists, allow pending // scanners to be submitted so they can report EOS and wake the pipeline task. diff --git a/be/src/exec/scan/scanner_context.h b/be/src/exec/scan/scanner_context.h index b5730897cdb286..4edcc07fcc7a06 100644 --- a/be/src/exec/scan/scanner_context.h +++ b/be/src/exec/scan/scanner_context.h @@ -93,7 +93,7 @@ class ScanTask { public: enum class State : int { PENDING, // not scheduled yet - IN_FLIGHT, // scheduled and running + IN_FLIGHT, // submitted to the scheduler; may still be waiting for TaskExecutor admission COMPLETED, // finished with result or error, waiting to be collected by scan node EOS, // finished and no more data, waiting to be collected by scan node }; @@ -289,10 +289,14 @@ class ScannerContext : public std::enable_shared_from_this, // accessed by both the scanner thread pool and the operator (get_block_from_queue). std::mutex _transfer_lock; - // Together, _completed_tasks and _in_flight_tasks_num represent all "occupied" concurrency - // slots. The scheduler uses their sum as the current concurrency: + // _in_flight_tasks_num includes every task submitted to the scheduler. New scanner submission + // always uses this total count. TaskExecutor can keep first-time submissions in its per-handle + // admission queue, so only the fallback that resumes yielded scanners uses: // - // current_concurrency = _completed_tasks.size() + _in_flight_tasks_num + // admitted_in_flight = _in_flight_tasks_num - TaskHandle::queued_leaf_splits() + // + // This admission-aware fallback can only re-enqueue an existing split. It must never grant + // margin for pulling another pending scanner into TaskExecutor. // // Lifecycle of a ScanTask: // _pending_tasks --(submit_scan_task)--> [thread pool] --(push_back_scan_task)--> @@ -312,9 +316,9 @@ class ScannerContext : public std::enable_shared_from_this, // _pull_next_scan_task() during scheduling. std::stack> _pending_tasks; - // Number of scan tasks currently submitted to the scanner scheduler thread pool - // (i.e. in-flight). Incremented by submit_scan_task() before submission and - // decremented by push_back_scan_task() when the thread pool returns the task. + // Number of scan tasks currently submitted to the scanner scheduler. This includes tasks + // waiting for TaskExecutor admission. Incremented by submit_scan_task() before submission and + // decremented by push_back_scan_task() when the scheduler returns the task. // Declared atomic so it can be read without _transfer_lock in non-critical paths, // but must be read under _transfer_lock whenever combined with _completed_tasks.size() // to form a consistent concurrency snapshot. @@ -347,10 +351,16 @@ class ScannerContext : public std::enable_shared_from_this, MOCK_REMOVE(const) int32_t _min_scan_concurrency = 1; std::shared_ptr _pull_next_scan_task(std::shared_ptr current_scan_task, - int32_t current_concurrency); + int32_t current_concurrency, + bool only_existing_split = false); int32_t _get_margin(std::unique_lock& transfer_lock, std::unique_lock& scheduler_lock); + int32_t _get_margin(std::unique_lock& transfer_lock, + std::unique_lock& scheduler_lock, + int32_t in_flight_tasks); + + int32_t _admitted_in_flight_tasks() const; // Memory-aware adaptive scheduling std::shared_ptr _scanner_mem_limiter = nullptr; diff --git a/be/src/exec/scan/task_executor/task_handle.h b/be/src/exec/scan/task_executor/task_handle.h index 69f79ad6bad43d..695c91b8ab3b81 100644 --- a/be/src/exec/scan/task_executor/task_handle.h +++ b/be/src/exec/scan/task_executor/task_handle.h @@ -28,6 +28,8 @@ class TaskHandle { virtual Status init() = 0; virtual bool is_closed() const = 0; virtual TaskId task_id() const = 0; + // Leaf splits accepted by TaskExecutor but not admitted to this task yet. + virtual int queued_leaf_splits() const = 0; }; } // namespace doris diff --git a/be/src/exec/scan/task_executor/time_sharing/time_sharing_task_handle.cpp b/be/src/exec/scan/task_executor/time_sharing/time_sharing_task_handle.cpp index ba55bb6dba38ea..61aae4f38a95aa 100644 --- a/be/src/exec/scan/task_executor/time_sharing/time_sharing_task_handle.cpp +++ b/be/src/exec/scan/task_executor/time_sharing/time_sharing_task_handle.cpp @@ -133,6 +133,11 @@ int TimeSharingTaskHandle::running_leaf_splits() const { return static_cast(_running_leaf_splits.size()); } +int TimeSharingTaskHandle::queued_leaf_splits() const { + std::lock_guard lock(_mutex); + return static_cast(_queued_leaf_splits.size()); +} + int64_t TimeSharingTaskHandle::scheduled_nanos() const { std::lock_guard lock(_mutex); return _scheduled_nanos; diff --git a/be/src/exec/scan/task_executor/time_sharing/time_sharing_task_handle.h b/be/src/exec/scan/task_executor/time_sharing/time_sharing_task_handle.h index 0b92830c4fff61..28afdbdf9bf623 100644 --- a/be/src/exec/scan/task_executor/time_sharing/time_sharing_task_handle.h +++ b/be/src/exec/scan/task_executor/time_sharing/time_sharing_task_handle.h @@ -57,6 +57,7 @@ class TimeSharingTaskHandle : public TaskHandle { bool enqueue_split(std::shared_ptr split); bool record_intermediate_split(std::shared_ptr split); int running_leaf_splits() const; + int queued_leaf_splits() const override; int64_t scheduled_nanos() const; std::shared_ptr poll_next_split(); void split_finished(std::shared_ptr split); diff --git a/be/test/exec/executor/time_sharing/time_sharing_task_executor_test.cpp b/be/test/exec/executor/time_sharing/time_sharing_task_executor_test.cpp index 792c01b3b23e08..990bb2f1fc3a8d 100644 --- a/be/test/exec/executor/time_sharing/time_sharing_task_executor_test.cpp +++ b/be/test/exec/executor/time_sharing/time_sharing_task_executor_test.cpp @@ -350,6 +350,8 @@ class TestingTaskHandle final : public TaskHandle { TaskId task_id() const override { return _task_id; } + int queued_leaf_splits() const override { return 0; } + private: TaskId _task_id; }; diff --git a/be/test/exec/scan/scanner_context_test.cpp b/be/test/exec/scan/scanner_context_test.cpp index 0e2aa0a9bde9d4..20ad815e70790e 100644 --- a/be/test/exec/scan/scanner_context_test.cpp +++ b/be/test/exec/scan/scanner_context_test.cpp @@ -23,12 +23,18 @@ #include #include +#include +#include +#include #include #include #include #include +#include #include +#include "common/cast_set.h" +#include "common/config.h" #include "common/object_pool.h" #include "core/block/block.h" #include "exec/operator/olap_scan_operator.h" @@ -37,6 +43,7 @@ #include "exec/scan/olap_scanner.h" #include "exec/scan/scan_node.h" #include "exec/scan/scanner_scheduler.h" +#include "exec/scan/task_executor/task_handle.h" #include "runtime/descriptors.h" #include "runtime/query_context.h" #include "storage/options.h" @@ -44,8 +51,33 @@ #include "storage/tablet/tablet.h" #include "storage/tablet/tablet_meta.h" #include "testutil/mock/mock_runtime_state.h" +#include "util/defer_op.h" namespace doris { +class YieldingScanner final : public Scanner { +public: + YieldingScanner(RuntimeState* state, ScanLocalStateBase* local_state, RuntimeProfile* profile, + std::atomic* produced_rows) + : Scanner(state, local_state, -1, profile), _produced_rows(produced_rows) {} + +protected: + Status _get_block_impl(RuntimeState* /*state*/, Block* block, bool* eof) override { + if (_has_returned_block) { + *eof = true; + return Status::OK(); + } + block->get_by_position(0).column->assert_mutable()->insert_data("x", 1); + _produced_rows->fetch_add(cast_set(block->rows()), std::memory_order_relaxed); + _has_returned_block = true; + *eof = false; + return Status::OK(); + } + +private: + std::atomic* _produced_rows; + bool _has_returned_block = false; +}; + class ScannerContextTest : public testing::Test { public: void SetUp() override { @@ -664,11 +696,18 @@ TEST_F(ScannerContextTest, pull_next_scan_task) { nullptr, scanner_context->_max_scan_concurrency - 1); EXPECT_EQ(pull_scan_task, nullptr); - scanner_context->_pending_tasks.push( - std::make_shared(std::make_shared(scanner))); + auto first_schedule_task = + std::make_shared(std::make_shared(scanner)); + scanner_context->_pending_tasks.push(first_schedule_task); pull_scan_task = scanner_context->_pull_next_scan_task( - nullptr, scanner_context->_max_scan_concurrency - 1); - EXPECT_NE(pull_scan_task, nullptr); + nullptr, scanner_context->_max_scan_concurrency - 1, true); + EXPECT_EQ(pull_scan_task, nullptr); + EXPECT_EQ(scanner_context->_pending_tasks.top(), first_schedule_task); + + first_schedule_task->is_first_schedule = false; + pull_scan_task = scanner_context->_pull_next_scan_task( + nullptr, scanner_context->_max_scan_concurrency - 1, true); + EXPECT_EQ(pull_scan_task, first_schedule_task); } TEST_F(ScannerContextTest, schedule_scan_task) { @@ -786,6 +825,101 @@ TEST_F(ScannerContextTest, schedule_scan_task) { scheduler_lock)); } +// Keep the end-to-end scheduler lifecycle in one test so the admission queue and scanner state +// transitions share the same fixture and executor instance. +// NOLINTNEXTLINE(readability-function-cognitive-complexity) +TEST_F(ScannerContextTest, task_executor_admission_queue_keeps_scanners_live) { + constexpr int scanner_count = 6; + constexpr int context_limit = 4; + constexpr int task_handle_limit = 2; + + const int old_initial_concurrency = config::task_executor_initial_max_concurrency_per_task; + const int old_max_concurrency = config::task_executor_max_concurrency_per_task; + config::task_executor_initial_max_concurrency_per_task = task_handle_limit; + config::task_executor_max_concurrency_per_task = task_handle_limit; + Defer restore_config {[&] { + config::task_executor_initial_max_concurrency_per_task = old_initial_concurrency; + config::task_executor_max_concurrency_per_task = old_max_concurrency; + }}; + + auto task_exec_ctx = std::make_shared(); + state->set_task_execution_context(task_exec_ctx); + auto scheduler = std::make_unique("scanner_liveness_test", + cgroup_cpu_ctl); + ASSERT_TRUE(scheduler->start(task_handle_limit, task_handle_limit, 100, 0).ok()); + Defer stop_scheduler {[&] { + state->get_query_ctx()->_scan_task_scheduler = nullptr; + scheduler->stop(); + }}; + state->get_query_ctx()->_scan_task_scheduler = scheduler.get(); + + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + scanner_count, TQueryCacheParam {}); + scan_operator->_shared_scan_limit.store(-1, std::memory_order_relaxed); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + olap_scan_local_state->_max_scan_concurrency = max_concurrency_counter.get(); + olap_scan_local_state->_min_scan_concurrency = min_concurrency_counter.get(); + olap_scan_local_state->_scan_cpu_timer = profile->add_counter("ScanCpuTime", TUnit::TIME_NS); + olap_scan_local_state->_rows_read_counter = profile->add_counter("RowsRead", TUnit::UNIT); + olap_scan_local_state->_parent = scan_operator.get(); + + std::atomic produced_rows = 0; + std::list> scanners; + for (int i = 0; i < scanner_count; ++i) { + ScannerSPtr scanner = std::make_shared( + state.get(), olap_scan_local_state.get(), profile.get(), &produced_rows); + scanner->_output_tuple_desc = output_tuple_desc; + scanner->_output_row_descriptor = nullptr; + ASSERT_TRUE(scanner->init(state.get(), {}).ok()); + scanners.push_back(std::make_shared(scanner)); + } + + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, output_row_descriptor, + scanners, -1, scan_dependency, &shared_limit, nullptr, nullptr, 0, false, + context_limit); + scanner_context->_newly_create_free_blocks_num = newly_create_free_blocks_num.get(); + scanner_context->_scanner_memory_used_counter = scanner_memory_used_counter.get(); + scanner_context->_min_scan_concurrency = 1; + // Force one initial submission batch above the TaskHandle limit, then remove the scheduler + // demand so a non-EOS scanner needs a free ScannerContext slot to be re-enqueued. + scanner_context->_min_scan_concurrency_of_scan_scheduler = scanner_count; + + ASSERT_TRUE(scanner_context->init().ok()); + EXPECT_EQ(scanner_context->_max_scan_concurrency, context_limit); + EXPECT_EQ(scanner_context->task_handle()->queued_leaf_splits(), + context_limit - task_handle_limit); + EXPECT_EQ(scanner_context->_pending_tasks.size(), scanner_count - context_limit); + scanner_context->_min_scan_concurrency_of_scan_scheduler = 0; + + bool eos = false; + bool checked_submission_backlog = false; + int peak_in_flight_tasks = scanner_context->num_scheduled_scanners(); + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (!eos && std::chrono::steady_clock::now() < deadline) { + if (!scan_dependency->ready()) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + continue; + } + Block block; + ASSERT_TRUE(scanner_context->get_block_from_queue(state.get(), &block, &eos, 0).ok()); + if (!checked_submission_backlog) { + // The yielded scanner is parked while another completed task occupies the last + // context slot; the two never-submitted scanners must remain pending as well. + EXPECT_EQ(scanner_context->_pending_tasks.size(), scanner_count - context_limit + 1); + checked_submission_backlog = true; + } + peak_in_flight_tasks = + std::max(peak_in_flight_tasks, scanner_context->num_scheduled_scanners()); + } + + EXPECT_TRUE(eos) << scanner_context->debug_string(); + EXPECT_EQ(produced_rows.load(std::memory_order_relaxed), scanner_count); + EXPECT_TRUE(checked_submission_backlog); + EXPECT_LE(peak_in_flight_tasks, context_limit); +} + TEST_F(ScannerContextTest, scan_queue_mem_limit) { state->_query_options.__set_scan_queue_mem_limit(100); ASSERT_EQ(state->scan_queue_mem_limit(), 100);