diff --git a/be/src/exec/scan/scanner_context.cpp b/be/src/exec/scan/scanner_context.cpp index c3f1cd55503413..c4180cfcc2a1ad 100644 --- a/be/src/exec/scan/scanner_context.cpp +++ b/be/src/exec/scan/scanner_context.cpp @@ -327,9 +327,9 @@ void ScannerContext::return_free_block(BlockUPtr block) { Status ScannerContext::submit_scan_task(std::shared_ptr scan_task, std::unique_lock& /*transfer_lock*/) { // increase _num_finished_scanners no matter the scan_task is submitted successfully or not. - // since if submit failed, it will be added back by ScannerContext::push_back_scan_task + // since if submit failed, it will be added back by ScannerContext::push_completed_scan_task // and _num_finished_scanners will be reduced. - // if submit succeed, it will be also added back by ScannerContext::push_back_scan_task + // if submit succeed, it will be also added back by ScannerContext::push_completed_scan_task // see ScannerScheduler::_scanner_scan. _in_flight_tasks_num++; return _scanner_scheduler->submit(shared_from_this(), scan_task); @@ -339,7 +339,7 @@ void ScannerContext::clear_free_blocks() { clear_blocks(_free_blocks); } -void ScannerContext::push_back_scan_task(std::shared_ptr scan_task) { +void ScannerContext::push_completed_scan_task(std::shared_ptr scan_task) { if (scan_task->status_ok()) { if (scan_task->cached_block && scan_task->cached_block->rows() > 0) { Status st = validate_block_schema(scan_task->cached_block.get()); @@ -349,6 +349,8 @@ void ScannerContext::push_back_scan_task(std::shared_ptr scan_task) { } } + // Publishing the result and releasing its in-flight slot must be atomic. Otherwise a worker + // could observe an available slot before the operator can observe this completed task. std::lock_guard l(_transfer_lock); if (!scan_task->status_ok()) { _process_status = scan_task->get_status(); @@ -407,11 +409,6 @@ Status ScannerContext::get_block_from_queue(RuntimeState* state, Block* block, b if (scan_task->is_eos()) { // 1. if eos, record a finished scanner. _num_finished_scanners++; - RETURN_IF_ERROR(_scanner_scheduler->schedule_scan_task(shared_from_this(), nullptr, l)); - } else { - scan_task->set_state(ScanTask::State::IN_FLIGHT); - RETURN_IF_ERROR( - _scanner_scheduler->schedule_scan_task(shared_from_this(), scan_task, l)); } } @@ -420,6 +417,12 @@ Status ScannerContext::get_block_from_queue(RuntimeState* state, Block* block, b (_is_shared_scan_limit_exhausted() && _in_flight_tasks_num == 0))) { _set_scanner_done(); _is_finished = true; + } else if (scan_task != nullptr) { + // Check the terminal state before scheduling more work. In particular, after the shared + // LIMIT is exhausted, submitting another Context runnable can turn a completed query into + // a queue-capacity error even though there is no result left to drain. + RETURN_IF_ERROR(_scanner_scheduler->schedule_scan_task( + shared_from_this(), scan_task->is_eos() ? nullptr : scan_task, l)); } *eos = done(); @@ -541,15 +544,16 @@ std::string ScannerContext::debug_string() { return fmt::format( "_query_id: {}, id: {}, total scanners: {}, pending tasks: {}, completed tasks: {}," " _should_stop: {}, _is_finished: {}, free blocks: {}," - " limit: {}, _in_flight_tasks_num: {}, remaining_limit: {}, _num_running_scanners: {}, " - "_max_thread_num: {}," + " limit: {}, remaining_limit: {}, _in_flight_tasks_num: {}, _is_context_queued: {}, " + "_num_finished_scanners: {}, _max_scan_concurrency: {}, expected_scanners: {}," " _max_bytes_in_queue: {}, _ins_idx: {}, _enable_adaptive_scanners: {}, " "_mem_share_arb: {}, _scanner_mem_limiter: {}", print_id(_query_id), ctx_id, _all_scanners.size(), _pending_tasks.size(), _completed_tasks.size(), _should_stop, _is_finished, _free_blocks.size_approx(), limit, _shared_scan_limit->load(std::memory_order_relaxed), _in_flight_tasks_num, - _num_finished_scanners, _max_scan_concurrency, _max_bytes_in_queue, _ins_idx, - _enable_adaptive_scanners, + _is_context_queued, _num_finished_scanners, _max_scan_concurrency, + _enable_adaptive_scanners ? _adaptive_processor->expected_scanners : -1, + _max_bytes_in_queue, _ins_idx, _enable_adaptive_scanners, _enable_adaptive_scanners ? _mem_share_arb->debug_string() : "NULL", _enable_adaptive_scanners ? _scanner_mem_limiter->debug_string() : "NULL"); } @@ -562,6 +566,108 @@ bool ScannerContext::_is_shared_scan_limit_exhausted() const { return limit >= 0 && _shared_scan_limit->load(std::memory_order_acquire) <= 0; } +bool ScannerContext::is_context_queued(const std::unique_lock& transfer_lock) const { + DORIS_CHECK(transfer_lock.owns_lock()); + return _is_context_queued; +} + +void ScannerContext::set_context_queued(bool queued, + const std::unique_lock& transfer_lock) { + DORIS_CHECK(transfer_lock.owns_lock()); + DORIS_CHECK(_is_context_queued != queued); + _is_context_queued = queued; +} + +void ScannerContext::set_context_failure(const Status& failure, + const std::unique_lock& transfer_lock) { + DORIS_CHECK(transfer_lock.owns_lock()); + DORIS_CHECK(!failure.ok()); + _process_status = failure; + _is_finished = true; + _set_scanner_done(); +} + +void ScannerContext::push_pending_scan_task(std::shared_ptr scan_task, + const std::unique_lock& transfer_lock) { + DORIS_CHECK(transfer_lock.owns_lock()); + DORIS_CHECK(scan_task != nullptr); + DORIS_CHECK(scan_task->cached_block == nullptr); + DORIS_CHECK(!scan_task->is_eos()); + // The state transition documents that this is an admission queue, not a completed-result queue. + scan_task->set_state(ScanTask::State::PENDING); + _pending_tasks.push(std::move(scan_task)); +} + +bool ScannerContext::can_admit_scan_task(const std::unique_lock& transfer_lock) const { + DORIS_CHECK(transfer_lock.owns_lock()); + if (done() || _pending_tasks.empty()) { + return false; + } + + int32_t effective_max_concurrency = _max_scan_concurrency; + if (_enable_adaptive_scanners) { + // expected_scanners is the adaptive ceiling as refreshed by _available_pickup_scanner_count(). + // Zero is a real allocation: MemLimiter::available_scanner_count() hands nothing to a later + // instance when the node-wide scanner budget is smaller than the operator parallelism. It + // is not a "not initialized" marker, so it must not fall back to _max_scan_concurrency; + // the escape below still lets one task make progress. + effective_max_concurrency = _adaptive_processor->expected_scanners; + } + if (low_memory_mode()) { + effective_max_concurrency = std::min(effective_max_concurrency, low_memory_mode_scanners()); + } + // Mirror the scheduler-wide budget of _get_margin(): once the pool has no slack, a Context is + // held at its minimum outstanding scanners instead of ramping to its maximum. Both counters are + // read here under _transfer_lock exactly as the TaskExecutor path reads them. + if (_scanner_scheduler->get_active_threads() + _scanner_scheduler->get_queue_size() >= + _min_scan_concurrency_of_scan_scheduler) { + effective_max_concurrency = + std::min(effective_max_concurrency, std::max(1, _min_scan_concurrency)); + } + + // Completed blocks still occupy a concurrency slot until the operator consumes them. Counting + // both collections prevents a fast producer from exceeding the per-Context scanner limit. + const int32_t current_concurrency = + cast_set(_completed_tasks.size()) + _in_flight_tasks_num; + // Keep one task progressing even if an adaptive limit reaches zero. Otherwise no worker can + // publish a result and wake the operator to make another scheduling decision. + return current_concurrency == 0 || current_concurrency < effective_max_concurrency; +} + +std::shared_ptr ScannerContext::try_get_next_scan_task( + const std::unique_lock& transfer_lock) { + if (_enable_adaptive_scanners) { + // Refresh expected_scanners and feed current block estimates back to the memory limiter, + // exactly as _get_margin() does on the TaskExecutor path. can_admit_scan_task() only reads + // the cached value, so ThreadPool admission would otherwise never apply adaptive limits. + static_cast(_available_pickup_scanner_count()); + } + if (!can_admit_scan_task(transfer_lock)) { + VLOG_DEBUG << fmt::format( + "[{}|{}] refuse admission, pending: {}, completed: {}, in-flight: {}, done: {}", + print_id(_query_id), ctx_id, _pending_tasks.size(), _completed_tasks.size(), + _in_flight_tasks_num, done()); + return nullptr; + } + + // Pop and mark in-flight while holding the same lock used by completion and consumption. + // Thus concurrent Context workers cannot admit the same task or both pass the limit check. + auto scan_task = _pending_tasks.top(); + _pending_tasks.pop(); + // ThreadPool admission bypasses ScannerScheduler::submit(); restart the per-scanner wait + // timer here so it measures admission-to-execution instead of everything since the previous + // attempt paused, which would include time the completed block waited for the operator. + if (auto scanner_delegate = scan_task->scanner.lock()) { + scanner_delegate->_scanner->start_wait_worker_timer(); + } + scan_task->set_state(ScanTask::State::IN_FLIGHT); + ++_in_flight_tasks_num; + VLOG_DEBUG << fmt::format("[{}|{}] admit scanner, pending: {}, completed: {}, in-flight: {}", + print_id(_query_id), ctx_id, _pending_tasks.size(), + _completed_tasks.size(), _in_flight_tasks_num); + return scan_task; +} + void ScannerContext::update_peak_running_scanner(int num) { #ifndef BE_TEST _local_state->_peak_running_scanner->add(num); diff --git a/be/src/exec/scan/scanner_context.h b/be/src/exec/scan/scanner_context.h index 72a557c7461fe7..aaeb36927d9b1c 100644 --- a/be/src/exec/scan/scanner_context.h +++ b/be/src/exec/scan/scanner_context.h @@ -129,7 +129,11 @@ class ScanTask { void set_state(State state) { switch (state) { case State::PENDING: - DCHECK(_state == State::PENDING || _state == State::IN_FLIGHT) << (int)_state; + // A task returns to PENDING after the operator consumes its non-EOS cached block. + // For example, one scanner may produce several blocks, so COMPLETED is not terminal. + DCHECK(_state == State::PENDING || _state == State::IN_FLIGHT || + _state == State::COMPLETED) + << (int)_state; DCHECK(cached_block == nullptr); break; case State::IN_FLIGHT: @@ -207,8 +211,9 @@ class ScannerContext : public std::enable_shared_from_this, // set the `eos` to `ScanTask::eos` if there is no more data in current scanner Status submit_scan_task(std::shared_ptr scan_task, std::unique_lock&); - // Push back a scan task. - void push_back_scan_task(std::shared_ptr scan_task); + // Publish a task whose current scan attempt has completed. The operator consumes its cached + // block and returns a non-EOS task to PENDING for its next scan attempt. + void push_completed_scan_task(std::shared_ptr scan_task); // Return true if this ScannerContext need no more process bool done() const { return _is_finished || _should_stop; } @@ -253,6 +258,39 @@ class ScannerContext : public std::enable_shared_from_this, std::unique_lock& transfer_lock, std::unique_lock& scheduler_lock); + // Context scheduling and operator consumption share this lock so queue-state changes and task + // admission form one atomic decision. For example, two worker callbacks cannot both admit the + // last available concurrency slot. + std::mutex& transfer_lock() { return _transfer_lock; } + + // One Context submission represents many pending scanners in the ThreadPool scheduler. + // Keeping this separate from scanner execution prevents duplicate runnables from accumulating. + bool is_context_queued(const std::unique_lock& transfer_lock) const; + // Transition the Context runnable's queue state. The caller must hold _transfer_lock. + void set_context_queued(bool queued, const std::unique_lock& transfer_lock); + + // Publish a scheduler failure and make the Context terminal. The caller must hold + // _transfer_lock so a retained ThreadPool callback cannot admit another scanner concurrently. + void set_context_failure(const Status& failure, + const std::unique_lock& transfer_lock); + + // Return a scanner to the admission queue after its block is consumed. It may not own a cached + // block and may not be EOS: EOS scanners are terminal and must not run again. + void push_pending_scan_task(std::shared_ptr scan_task, + const std::unique_lock& transfer_lock); + + // Return whether a Context worker can currently admit one pending scanner. This check has no + // side effects, so the scheduler can avoid submitting a runnable that would immediately exit. + // It always admits one scanner when nothing is progressing so the operator can be woken, and it + // holds the Context at max(1, _min_scan_concurrency) while the scheduler pool has no slack, + // like _get_margin() on the TaskExecutor path. The caller must hold _transfer_lock. + bool can_admit_scan_task(const std::unique_lock& transfer_lock) const; + + // Atomically check whether this context can start another scan task, move one task from + // pending to in-flight, and return it. The caller must hold _transfer_lock. + std::shared_ptr try_get_next_scan_task( + const std::unique_lock& transfer_lock); + protected: /// Four criteria to determine whether to increase the parallelism of the scanners /// 1. It ran for at least `SCALE_UP_DURATION` ms after last scale up @@ -293,26 +331,35 @@ class ScannerContext : public std::enable_shared_from_this, // current_concurrency = _completed_tasks.size() + _in_flight_tasks_num // // Lifecycle of a ScanTask: - // _pending_tasks --(submit_scan_task)--> [thread pool] --(push_back_scan_task)--> - // _completed_tasks --(get_block_from_queue)--> operator + // _pending_tasks --(submit_scan_task on the TaskExecutor path, + // try_get_next_scan_task on the ThreadPool path)--> [thread pool] + // --(push_completed_scan_task)--> _completed_tasks --(get_block_from_queue)--> operator // After consumption: non-EOS task goes back to _pending_tasks; EOS increments // _num_finished_scanners. // Completed scan tasks whose cached_block is ready for the operator to consume. - // Protected by _transfer_lock. Written by push_back_scan_task() (scanner thread), + // Protected by _transfer_lock. Written by push_completed_scan_task() (scanner thread), // read/popped by get_block_from_queue() (operator thread). std::list> _completed_tasks; - // Scanners waiting to be submitted to the scheduler thread pool. Stored as a stack - // (LIFO) so that recently-used scanners are re-scheduled first, which is more likely - // to be cache-friendly. Protected by _transfer_lock. Populated in the constructor - // and by schedule_scan_task() when the concurrency limit is reached; drained by - // _pull_next_scan_task() during scheduling. + // Scanners waiting to be admitted for execution. Stored as a stack (LIFO) so that + // recently-used scanners are re-scheduled first, which is more likely to be cache-friendly. + // Protected by _transfer_lock. Populated in the constructor and when an operator returns a + // non-EOS task; drained by try_get_next_scan_task() or the TaskExecutor scheduler. std::stack> _pending_tasks; + // True from the start of one Context submission until its runnable starts. The marker may + // remain true when no runnable was retained: a failed submission makes the Context terminal, + // and a submission that threw inside _run_context() publishes the error through the task that + // was already admitted. In both cases the operator observes _process_status, so no further + // submission is attempted. It does not describe scanners executing on workers. Protected by + // _transfer_lock. + bool _is_context_queued = false; + // 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. + // (i.e. in-flight). Incremented before a task is submitted or directly admitted for + // thread-pool execution, and decremented by push_completed_scan_task() when the worker + // returns it. // 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. diff --git a/be/src/exec/scan/scanner_scheduler.cpp b/be/src/exec/scan/scanner_scheduler.cpp index 3036cbdf6ee099..916078004e1a36 100644 --- a/be/src/exec/scan/scanner_scheduler.cpp +++ b/be/src/exec/scan/scanner_scheduler.cpp @@ -75,17 +75,7 @@ Status ScannerScheduler::submit(std::shared_ptr ctx, 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(); + return execute_scan_task(ctx, scanner_ref); }; SimplifiedScanTask simple_scan_task = {work_func, ctx, scan_task}; return this->submit_scan_task(simple_scan_task); @@ -104,6 +94,22 @@ Status ScannerScheduler::submit(std::shared_ptr ctx, return Status::OK(); } +bool ScannerScheduler::execute_scan_task(const std::shared_ptr& ctx, + const std::shared_ptr& scan_task) { + // Both schedulers admit tasks differently, but exceptions must always become a completed task + // so the operator observes the error and releases the task's in-flight concurrency slot. + auto status = [&] { + RETURN_IF_CATCH_EXCEPTION(_scanner_scan(ctx, scan_task)); + return Status::OK(); + }(); + if (!status.ok()) { + scan_task->set_status(status); + ctx->push_completed_scan_task(scan_task); + return true; + } + return scan_task->is_eos(); +} + void handle_reserve_memory_failure(RuntimeState* state, std::shared_ptr ctx, const Status& st, size_t reserve_size) { ctx->clear_free_blocks(); @@ -308,7 +314,7 @@ void ScannerScheduler::_scanner_scan(std::shared_ptr ctx, "{}, eos: {}, status: {}", ctx->ctx_id, ctx->num_scheduled_scanners(), eos, status.to_string()); - ctx->push_back_scan_task(scan_task); + ctx->push_completed_scan_task(scan_task); } // NOLINTEND(readability-function-cognitive-complexity,readability-function-size) diff --git a/be/src/exec/scan/scanner_scheduler.h b/be/src/exec/scan/scanner_scheduler.h index fa5387d2736a12..273954274d443b 100644 --- a/be/src/exec/scan/scanner_scheduler.h +++ b/be/src/exec/scan/scanner_scheduler.h @@ -138,7 +138,11 @@ class ScannerScheduler { protected: int _min_active_scan_threads; -private: + // Execute one admitted task for both scheduler implementations. The return value is consumed + // by TaskExecutor to distinguish terminal EOS/error tasks from scanners that remain runnable. + static bool execute_scan_task(const std::shared_ptr& ctx, + const std::shared_ptr& scan_task); + static void _scanner_scan(std::shared_ptr ctx, std::shared_ptr scan_task); @@ -240,12 +244,13 @@ class ThreadPoolSimplifiedScanScheduler MOCK_REMOVE(final) : public ScannerSched std::unique_lock& transfer_lock) override; private: + void _run_context(std::shared_ptr scanner_ctx); + std::unique_ptr _scan_thread_pool; std::atomic _is_stop; std::weak_ptr _cgroup_cpu_ctl; std::string _sched_name; std::string _workload_group; - std::shared_mutex _lock; }; class TaskExecutorSimplifiedScanScheduler final : public ScannerScheduler { diff --git a/be/src/exec/scan/simplified_scan_scheduler.cpp b/be/src/exec/scan/simplified_scan_scheduler.cpp index 275461ff1dd409..e3ef29d6f58877 100644 --- a/be/src/exec/scan/simplified_scan_scheduler.cpp +++ b/be/src/exec/scan/simplified_scan_scheduler.cpp @@ -17,8 +17,12 @@ #include +#include "common/exception.h" +#include "common/logging.h" #include "exec/scan/scanner_context.h" #include "exec/scan/scanner_scheduler.h" +#include "runtime/thread_context.h" +#include "util/debug_points.h" namespace doris { class ScannerDelegate; @@ -34,7 +38,104 @@ Status TaskExecutorSimplifiedScanScheduler::schedule_scan_task( Status ThreadPoolSimplifiedScanScheduler::schedule_scan_task( std::shared_ptr scanner_ctx, std::shared_ptr current_scan_task, std::unique_lock& transfer_lock) { - std::unique_lock wl(_lock); - return scanner_ctx->schedule_scan_task(current_scan_task, transfer_lock, wl); + // Unlike TaskExecutor, ThreadPool queues a Context runnable. It later admits one pending task + // under transfer_lock. This bounds queue entries to one per Context even when many scanners + // become runnable together. + DORIS_CHECK(transfer_lock.owns_lock()); + if (current_scan_task != nullptr) { + // The operator has consumed a non-EOS result, making this scanner eligible for another + // scan attempt. Queue the scanner first; the Context runnable chooses it later. + scanner_ctx->push_pending_scan_task(std::move(current_scan_task), transfer_lock); + } + if (scanner_ctx->is_context_queued(transfer_lock)) { + // A queued runnable will see all pending scanners added before it obtains transfer_lock. + // Submitting another runnable would only duplicate work. + return Status::OK(); + } + if (!scanner_ctx->can_admit_scan_task(transfer_lock)) { + // No runnable is needed when the Context has no pending scanner or its concurrency slots + // are occupied. A completion only wakes the operator; the operator's next consumption in + // get_block_from_queue(), or the successor submit in _run_context(), retries scheduling. + return Status::OK(); + } + + if (_is_stop) { + Status failure = Status::InternalError("scanner pool {} is shutdown.", _sched_name); + scanner_ctx->set_context_failure(failure, transfer_lock); + return failure; + } + + // ThreadPool::submit_func() may return an error after retaining the runnable. Set the marker + // before submission so either outcome is safe: a retained callback clears it, while a truly + // rejected callback leaves a terminal Context that no longer needs rescheduling. + scanner_ctx->set_context_queued(true, transfer_lock); + Status status = + _scan_thread_pool->submit_func([this, scanner_ctx] { _run_context(scanner_ctx); }); + if (status.ok()) { + VLOG_DEBUG << "submit context runnable to scanner pool " << _sched_name << ", " + << scanner_ctx->debug_string(); + return Status::OK(); + } + Status failure = + Status::TooManyTasks("Failed to submit scanner context {} to scanner pool, reason: {}", + scanner_ctx->ctx_id, status.msg()); + scanner_ctx->set_context_failure(failure, transfer_lock); + return failure; +} + +void ThreadPoolSimplifiedScanScheduler::_run_context(std::shared_ptr scanner_ctx) { + std::shared_ptr scan_task; + Status admission_status = [&]() -> Status { + std::unique_lock transfer_lock(scanner_ctx->transfer_lock()); + scanner_ctx->set_context_queued(false, transfer_lock); + + auto task_execution_lock = scanner_ctx->task_exec_ctx(); + if (task_execution_lock == nullptr) { + return Status::OK(); + } +#ifndef BE_TEST + // Attach before admission: allocations below (for example the resubmitted + // FunctionRunnable) must charge the query rather than the orphan tracker. Scoped to this + // lambda so it detaches before execute_scan_task(), whose _scanner_scan() attaches again. + SCOPED_ATTACH_TASK(scanner_ctx->state()); +#endif + Status status = [&]() -> Status { + RETURN_IF_CATCH_EXCEPTION({ + // Admission checks completed results, active tasks, and adaptive limits while + // holding transfer_lock. + scan_task = scanner_ctx->try_get_next_scan_task(transfer_lock); + if (scan_task != nullptr) { + DBUG_EXECUTE_IF("ThreadPoolSimplifiedScanScheduler._run_context.inject_failure", + { + throw Exception(ErrorCode::INTERNAL_ERROR, + "injected admission failure"); + }); + // Queue the next Context runnable before executing this task. Holding + // transfer_lock keeps the admission decision atomic. + RETURN_IF_ERROR(schedule_scan_task(scanner_ctx, nullptr, transfer_lock)); + } + }); + return Status::OK(); + }(); + if (!status.ok() && scan_task == nullptr) { + scanner_ctx->set_context_failure(status, transfer_lock); + } + return status; + }(); + if (!admission_status.ok()) [[unlikely]] { + if (scan_task != nullptr) { + // The scanner was admitted before the failure. Publish it to release the in-flight + // slot and make the error observable by the operator. + scan_task->set_status(admission_status); + scanner_ctx->push_completed_scan_task(scan_task); + } + return; + } + if (scan_task == nullptr) { + return; + } + // The scan runs without transfer_lock so the operator and other Context workers can continue + // consuming results and admitting work. Completion reacquires the lock before publishing. + execute_scan_task(scanner_ctx, scan_task); } } // namespace doris diff --git a/be/test/exec/scan/scanner_context_test.cpp b/be/test/exec/scan/scanner_context_test.cpp index ea2382e3c4e2fe..bff24bf4fbcde8 100644 --- a/be/test/exec/scan/scanner_context_test.cpp +++ b/be/test/exec/scan/scanner_context_test.cpp @@ -23,12 +23,16 @@ #include #include +#include +#include #include #include #include #include +#include #include +#include "common/config.h" #include "common/object_pool.h" #include "core/block/block.h" #include "exec/operator/olap_scan_operator.h" @@ -36,16 +40,71 @@ #include "exec/scan/mock_simplified_scan_scheduler.h" #include "exec/scan/olap_scanner.h" #include "exec/scan/scan_node.h" +#include "exec/scan/scanner.h" #include "exec/scan/scanner_scheduler.h" #include "runtime/descriptors.h" #include "runtime/query_context.h" +#include "runtime/task_execution_context.h" #include "storage/options.h" #include "storage/storage_engine.h" #include "storage/tablet/tablet.h" #include "storage/tablet/tablet_meta.h" #include "testutil/mock/mock_runtime_state.h" +#include "util/countdown_latch.h" +#include "util/debug_points.h" +#include "util/defer_op.h" namespace doris { +// A scanner that produces `blocks_per_scanner` one-row blocks and then reports EOS, without any +// tablet or file behind it. It lets the ThreadPool scheduler chain (admit -> execute -> publish -> +// consume -> re-admit) run end to end in a unit test. `overlap` is counted down the first time two +// attempts run concurrently; the first attempt waits for it so the peak concurrency observed by the +// test does not depend on timing. +class ChainMockScanner : public Scanner { +public: + ChainMockScanner(RuntimeState* state, ScanLocalStateBase* local_state, RuntimeProfile* profile, + int blocks_per_scanner, std::atomic* running, + std::atomic* peak_running, CountDownLatch* overlap) + : Scanner(state, local_state, -1, profile), + _blocks_left(blocks_per_scanner), + _running(running), + _peak_running(peak_running), + _overlap(overlap) {} + +protected: + Status _get_block_impl(RuntimeState* /*state*/, Block* block, bool* eof) override { + const int running = ++*_running; + Defer done([&] { --*_running; }); + int peak = _peak_running->load(); + while (running > peak && !_peak_running->compare_exchange_weak(peak, running)) { + } + if (running >= 2) { + _overlap->count_down(); + } else { + // Bounded so a scheduler that never admits a second scanner fails the test instead of + // hanging it. + static_cast(_overlap->wait_for(std::chrono::seconds(5))); + } + if (_blocks_left == 0) { + *eof = true; + return Status::OK(); + } + --_blocks_left; + block->get_by_position(0).column->assert_mutable()->insert_default(); + *eof = false; + return Status::OK(); + } + + // The local state in these tests has no profile counters. + void _collect_profile_before_close() override {} + +private: + int _blocks_left; + std::atomic* _running; + std::atomic* _peak_running; + CountDownLatch* _overlap; +}; + class ScannerContextTest : public testing::Test { public: void SetUp() override { @@ -169,7 +228,10 @@ TEST_F(ScannerContextTest, test_init) { state->set_query_options(query_options); std::unique_ptr scheduler = std::make_unique(cgroup_cpu_ctl); + // init() is invoked twice below, and each invocation performs one initial scheduling attempt. + // Keep this expectation explicit so changing bootstrap scheduling updates this test too. EXPECT_CALL(*scheduler, schedule_scan_task(testing::_, testing::_, testing::_)) + .Times(2) .WillRepeatedly(testing::Return(Status::OK())); scanner_context->_scanner_scheduler = scheduler.get(); @@ -454,7 +516,7 @@ TEST_F(ScannerContextTest, test_max_column_reader_num) { ASSERT_EQ(scanner_context->_max_scan_concurrency, 1); } -TEST_F(ScannerContextTest, test_push_back_scan_task) { +TEST_F(ScannerContextTest, test_push_completed_scan_task) { const int parallel_tasks = 1; auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, parallel_tasks, TQueryCacheParam {}); @@ -486,7 +548,7 @@ TEST_F(ScannerContextTest, test_push_back_scan_task) { for (int i = 0; i < 5; ++i) { auto scan_task = std::make_shared(std::make_shared(scanner)); - scanner_context->push_back_scan_task(scan_task); + scanner_context->push_completed_scan_task(scan_task); ASSERT_EQ(scanner_context->_in_flight_tasks_num, 10 - i); } } @@ -664,6 +726,633 @@ TEST_F(ScannerContextTest, pull_next_scan_task) { EXPECT_NE(pull_scan_task, nullptr); } +TEST_F(ScannerContextTest, thread_pool_admission_state) { + const int parallel_tasks = 1; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = -1; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + std::list> scanners { + std::make_shared(scanner)}; + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, -1, + scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + + // An idle pool: admission is bounded only by the per-Context limit. + std::unique_ptr scheduler = + std::make_unique(cgroup_cpu_ctl); + EXPECT_CALL(*scheduler, get_active_threads()).WillRepeatedly(testing::Return(0)); + EXPECT_CALL(*scheduler, get_queue_size()).WillRepeatedly(testing::Return(0)); + scanner_context->_scanner_scheduler = scheduler.get(); + scanner_context->_min_scan_concurrency_of_scan_scheduler = 20; + + std::unique_lock context_transfer_lock(scanner_context->transfer_lock()); + scanner_context->_pending_tasks = std::stack>(); + scanner_context->_completed_tasks.clear(); + scanner_context->_in_flight_tasks_num = 0; + // Even if the effective limit is temporarily zero, one pending task must run so it can publish + // a block or EOS and prevent the Context from stalling. + scanner_context->_max_scan_concurrency = 0; + + EXPECT_FALSE(scanner_context->can_admit_scan_task(context_transfer_lock)); + + auto completed_task = std::make_shared(std::make_shared(scanner)); + completed_task->set_state(ScanTask::State::IN_FLIGHT); + completed_task->cached_block = Block::create_unique(); + completed_task->set_state(ScanTask::State::COMPLETED); + completed_task->cached_block.reset(); + // A consumed non-EOS result must be eligible for another Context admission. This also covers + // the COMPLETED -> PENDING transition used by ThreadPool scheduling. + scanner_context->push_pending_scan_task(completed_task, context_transfer_lock); + EXPECT_TRUE(scanner_context->can_admit_scan_task(context_transfer_lock)); + + EXPECT_FALSE(scanner_context->is_context_queued(context_transfer_lock)); + scanner_context->set_context_queued(true, context_transfer_lock); + EXPECT_TRUE(scanner_context->is_context_queued(context_transfer_lock)); + scanner_context->set_context_queued(false, context_transfer_lock); + + // The Context can admit exactly one scanner at its configured concurrency limit. + auto admitted_task = scanner_context->try_get_next_scan_task(context_transfer_lock); + EXPECT_EQ(admitted_task, completed_task); + EXPECT_EQ(admitted_task->_state, ScanTask::State::IN_FLIGHT); + EXPECT_EQ(scanner_context->_in_flight_tasks_num, 1); + + auto blocked_task = std::make_shared(std::make_shared(scanner)); + scanner_context->push_pending_scan_task(blocked_task, context_transfer_lock); + EXPECT_FALSE(scanner_context->can_admit_scan_task(context_transfer_lock)); + EXPECT_EQ(scanner_context->try_get_next_scan_task(context_transfer_lock), nullptr); +} + +TEST_F(ScannerContextTest, thread_pool_admission_refreshes_adaptive_limit) { + const int parallel_tasks = 2; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = -1; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + + std::list> scanners; + for (int i = 0; i < 5; ++i) { + scanners.push_back(std::make_shared(scanner)); + } + + TUniqueId query_id = state->get_query_ctx()->query_id(); + const int64_t query_mem_limit = 1024LL * 1024 * 1024; + auto arbitrator = MemShareArbitrator::create_shared(query_id, query_mem_limit, 0.3); + auto limiter = MemLimiter::create_shared(query_id, parallel_tasks, false, + static_cast(query_mem_limit * 0.3)); + // 200MB budget with 100MB estimated blocks: max_count = 2, so instance 1 gets exactly one + // adaptive slot. ins_idx = 1 keeps _available_pickup_scanner_count() away from the + // arbitrator-driven limit adjustment, which would overwrite this deterministic setup. + limiter->update_open_tasks_count(1); + limiter->update_mem_limit(200LL * 1024 * 1024); + limiter->reestimated_block_mem_bytes(100LL * 1024 * 1024); + + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, -1, + scan_dependency, &shared_limit, arbitrator, limiter, 1, true, parallel_tasks); + std::unique_ptr scheduler = + std::make_unique(cgroup_cpu_ctl); + EXPECT_CALL(*scheduler, get_active_threads()).WillRepeatedly(testing::Return(0)); + EXPECT_CALL(*scheduler, get_queue_size()).WillRepeatedly(testing::Return(0)); + scanner_context->_scanner_scheduler = scheduler.get(); + scanner_context->_min_scan_concurrency_of_scan_scheduler = 20; + + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + ASSERT_TRUE(scanner_context->_enable_adaptive_scanners); + ASSERT_LT(1, scanner_context->_max_scan_concurrency); + EXPECT_EQ(scanner_context->_adaptive_processor->expected_scanners, 0); + + // Admission refreshes the adaptive limit. Nothing is progressing yet, so the first scanner is + // admitted regardless, but expected_scanners must no longer stay at its initial zero. + auto first_task = scanner_context->try_get_next_scan_task(transfer_lock); + ASSERT_NE(first_task, nullptr); + EXPECT_EQ(scanner_context->_adaptive_processor->expected_scanners, 1); + + // One task is in flight and the refreshed adaptive limit is one: admission must refuse the + // next scanner even though _max_scan_concurrency would still allow it. + EXPECT_FALSE(scanner_context->can_admit_scan_task(transfer_lock)); + EXPECT_EQ(scanner_context->try_get_next_scan_task(transfer_lock), nullptr); +} + +TEST_F(ScannerContextTest, thread_pool_admission_keeps_zero_adaptive_allocation) { + const int parallel_tasks = 2; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = -1; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + + std::list> scanners; + for (int i = 0; i < 5; ++i) { + scanners.push_back(std::make_shared(scanner)); + } + + TUniqueId query_id = state->get_query_ctx()->query_id(); + const int64_t query_mem_limit = 1024LL * 1024 * 1024; + auto arbitrator = MemShareArbitrator::create_shared(query_id, query_mem_limit, 0.3); + auto limiter = MemLimiter::create_shared(query_id, parallel_tasks, false, + static_cast(query_mem_limit * 0.3)); + // 100MB budget with 100MB estimated blocks: max_count = 1 for two instances, so instance 1 + // is legitimately allocated zero scanners by the node-wide budget. + limiter->update_open_tasks_count(1); + limiter->update_mem_limit(100LL * 1024 * 1024); + limiter->reestimated_block_mem_bytes(100LL * 1024 * 1024); + + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, -1, + scan_dependency, &shared_limit, arbitrator, limiter, 1, true, parallel_tasks); + std::unique_ptr scheduler = + std::make_unique(cgroup_cpu_ctl); + EXPECT_CALL(*scheduler, get_active_threads()).WillRepeatedly(testing::Return(0)); + EXPECT_CALL(*scheduler, get_queue_size()).WillRepeatedly(testing::Return(0)); + scanner_context->_scanner_scheduler = scheduler.get(); + scanner_context->_min_scan_concurrency_of_scan_scheduler = 20; + + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + ASSERT_LT(1, scanner_context->_max_scan_concurrency); + + // Nothing is progressing, so one scanner is admitted even though the refreshed allocation + // is zero. + auto first_task = scanner_context->try_get_next_scan_task(transfer_lock); + ASSERT_NE(first_task, nullptr); + EXPECT_EQ(scanner_context->_adaptive_processor->expected_scanners, 0); + EXPECT_EQ(scanner_context->_in_flight_tasks_num, 1); + + // Zero is the real ceiling: with one task in flight nothing else may be admitted, matching + // the single progress task the TaskExecutor margin keeps. It must not fall back to + // _max_scan_concurrency. + EXPECT_FALSE(scanner_context->can_admit_scan_task(transfer_lock)); + EXPECT_EQ(scanner_context->try_get_next_scan_task(transfer_lock), nullptr); +} + +TEST_F(ScannerContextTest, thread_pool_admission_holds_minimum_when_pool_saturated) { + const int parallel_tasks = 4; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = -1; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + + // A single-worker pool whose worker is parked: active + queued == 1, i.e. the pool has no + // slack once the scheduler-wide budget is 1. + ThreadPoolSimplifiedScanScheduler scheduler("saturated_pool_test", cgroup_cpu_ctl); + ASSERT_TRUE(scheduler.start(1, 1, 4, 1).ok()); + CountDownLatch task_started(1); + CountDownLatch release_task(1); + Defer cleanup = [&] { + release_task.count_down(); + scheduler.stop(); + }; + ASSERT_TRUE(scheduler + .submit_scan_task(SimplifiedScanTask( + [&] { + task_started.count_down(); + release_task.wait(); + return true; + }, + nullptr, nullptr)) + .ok()); + ASSERT_TRUE(task_started.wait_for(std::chrono::seconds(5))); + ASSERT_EQ(scheduler.get_active_threads(), 1); + ASSERT_EQ(scheduler.get_queue_size(), 0); + + // Two Contexts share the saturated pool; each may have up to four scanners outstanding when + // the pool has slack. + std::vector> contexts; + for (int i = 0; i < 2; ++i) { + std::list> scanners; + for (int j = 0; j < 4; ++j) { + scanners.push_back(std::make_shared(scanner)); + } + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, -1, + scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + scanner_context->_scanner_scheduler = &scheduler; + scanner_context->_min_scan_concurrency = 1; + contexts.push_back(scanner_context); + } + + for (const auto& scanner_context : contexts) { + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + ASSERT_EQ(scanner_context->_max_scan_concurrency, parallel_tasks); + + // Saturated: one outstanding scanner is the ceiling, as _get_margin() enforces on the + // TaskExecutor path. + scanner_context->_min_scan_concurrency_of_scan_scheduler = 1; + scanner_context->_in_flight_tasks_num = 0; + EXPECT_TRUE(scanner_context->can_admit_scan_task(transfer_lock)); + scanner_context->_in_flight_tasks_num = 1; + EXPECT_FALSE(scanner_context->can_admit_scan_task(transfer_lock)); + + // A larger minimum raises the saturated ceiling accordingly. + scanner_context->_min_scan_concurrency = 2; + EXPECT_TRUE(scanner_context->can_admit_scan_task(transfer_lock)); + scanner_context->_in_flight_tasks_num = 2; + EXPECT_FALSE(scanner_context->can_admit_scan_task(transfer_lock)); + + // With slack the Context may ramp to its maximum again. + scanner_context->_min_scan_concurrency_of_scan_scheduler = 20; + EXPECT_TRUE(scanner_context->can_admit_scan_task(transfer_lock)); + scanner_context->_in_flight_tasks_num = parallel_tasks; + EXPECT_FALSE(scanner_context->can_admit_scan_task(transfer_lock)); + } +} + +TEST_F(ScannerContextTest, debug_string_reports_distinguishable_fields) { + const int parallel_tasks = 3; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = -1; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + std::list> scanners { + std::make_shared(scanner)}; + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, 7, + scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + + // Every value is distinct so a misplaced placeholder is visible in the output. + shared_limit.store(100); + scanner_context->_in_flight_tasks_num = 2; + scanner_context->_is_context_queued = true; + scanner_context->_num_finished_scanners = 5; + + const std::string debug = scanner_context->debug_string(); + EXPECT_NE(debug.find("limit: 7, remaining_limit: 100, _in_flight_tasks_num: 2, " + "_is_context_queued: true, _num_finished_scanners: 5, " + "_max_scan_concurrency: 3, expected_scanners: -1,"), + std::string::npos) + << debug; +} + +TEST_F(ScannerContextTest, thread_pool_submit_failure_policy) { + const int parallel_tasks = 2; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = -1; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + + std::list> scanners; + for (int i = 0; i < 2; ++i) { + scanners.push_back(std::make_shared(scanner)); + } + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, -1, + scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + + // One worker, zero queue capacity, worker occupied: every submit_func() is rejected. + ThreadPoolSimplifiedScanScheduler scheduler("submit_failure_policy_test", cgroup_cpu_ctl); + ASSERT_TRUE(scheduler.start(1, 1, 0, 1).ok()); + CountDownLatch task_started(1); + CountDownLatch release_task(1); + Defer cleanup = [&] { + release_task.count_down(); + scheduler.stop(); + }; + ASSERT_TRUE(scheduler + .submit_scan_task(SimplifiedScanTask( + [&] { + task_started.count_down(); + release_task.wait(); + return true; + }, + nullptr, nullptr)) + .ok()); + ASSERT_TRUE(task_started.wait_for(std::chrono::seconds(5))); + scanner_context->_scanner_scheduler = &scheduler; + + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + ASSERT_FALSE(scanner_context->_pending_tasks.empty()); + + // Context submission is fail-fast regardless of other progress. Retrying here would couple + // the scanner scheduler to ThreadPool's internal rejection/retention behavior. + Status surfaced = scheduler.schedule_scan_task(scanner_context, nullptr, transfer_lock); + EXPECT_TRUE(surfaced.is()) << surfaced.to_string(); + EXPECT_TRUE(scanner_context->done()); + EXPECT_FALSE(scanner_context->_process_status.ok()); + EXPECT_TRUE(scan_dependency->ready()); + // The marker is set before submit_func(). This rejected runnable was not retained, but the + // terminal Context no longer needs the marker cleared or another submission attempted. + EXPECT_TRUE(scanner_context->is_context_queued(transfer_lock)); +} + +TEST_F(ScannerContextTest, run_context_publishes_admission_failure) { + const bool old_enable_debug_points = config::enable_debug_points; + config::enable_debug_points = true; + DebugPoints::instance()->add("ThreadPoolSimplifiedScanScheduler._run_context.inject_failure"); + Defer cleanup_debug_point = [&] { + DebugPoints::instance()->remove( + "ThreadPoolSimplifiedScanScheduler._run_context.inject_failure"); + config::enable_debug_points = old_enable_debug_points; + }; + + const int parallel_tasks = 2; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = -1; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + + std::list> scanners; + for (int i = 0; i < 2; ++i) { + scanners.push_back(std::make_shared(scanner)); + } + // The worker's task_exec_ctx() must resolve, otherwise _run_context() exits before admission. + // HasTaskExecutionCtx snapshots the weak_ptr at construction, so set it before create_shared. + auto task_execution_context = std::make_shared(); + state->set_task_execution_context(task_execution_context); + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, -1, + scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + + ThreadPoolSimplifiedScanScheduler scheduler("run_context_failure_test", cgroup_cpu_ctl); + ASSERT_TRUE(scheduler.start(1, 1, 1, 1).ok()); + Defer cleanup = [&] { scheduler.stop(); }; + scanner_context->_scanner_scheduler = &scheduler; + + { + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + ASSERT_TRUE(scheduler.schedule_scan_task(scanner_context, nullptr, transfer_lock).ok()); + ASSERT_TRUE(scanner_context->is_context_queued(transfer_lock)); + } + + // The worker admits a scanner and hits the injected exception. It must publish the failure + // as a completed task instead of terminating the process or leaking the in-flight slot. + bool published = false; + for (int i = 0; i < 10000; ++i) { + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + if (!scanner_context->_completed_tasks.empty()) { + published = true; + break; + } + transfer_lock.unlock(); + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + ASSERT_TRUE(published); + + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + ASSERT_EQ(scanner_context->_completed_tasks.size(), 1); + EXPECT_FALSE(scanner_context->_completed_tasks.front()->status_ok()); + EXPECT_EQ(scanner_context->_in_flight_tasks_num, 0); + EXPECT_FALSE(scanner_context->is_context_queued(transfer_lock)); +} + +TEST_F(ScannerContextTest, thread_pool_context_chain_runs_all_scanners) { + const int parallel_tasks = 2; + const int scanner_count = 3; + const int blocks_per_scanner = 4; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + olap_scan_local_state->_parent = 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(); + scan_operator->_should_run_serial = false; + TQueryOptions query_options; + query_options.__set_max_column_reader_num(0); + state->set_query_options(query_options); + + std::atomic running {0}; + std::atomic peak_running {0}; + CountDownLatch overlap(1); + std::list> scanners; + for (int i = 0; i < scanner_count; ++i) { + std::shared_ptr scanner = std::make_shared( + state.get(), olap_scan_local_state.get(), profile.get(), blocks_per_scanner, + &running, &peak_running, &overlap); + scanners.push_back(std::make_shared(scanner)); + } + + // The worker's task_exec_ctx() must resolve, otherwise _run_context() exits before admission. + auto task_execution_context = std::make_shared(); + state->set_task_execution_context(task_execution_context); + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, -1, + scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + scanner_context->_newly_create_free_blocks_num = newly_create_free_blocks_num.get(); + scanner_context->_scanner_memory_used_counter = scanner_memory_used_counter.get(); + + // Two workers so the successor runnable can overlap with the executing scanner. + ThreadPoolSimplifiedScanScheduler scheduler("context_chain_test", cgroup_cpu_ctl); + ASSERT_TRUE(scheduler.start(2, 2, 16, 1).ok()); + Defer cleanup = [&] { scheduler.stop(); }; + scanner_context->_scanner_scheduler = &scheduler; + // The two-thread pool never reaches this budget, so the Context may ramp to its maximum. + scanner_context->_min_scan_concurrency_of_scan_scheduler = 20; + ASSERT_EQ(scanner_context->_max_scan_concurrency, parallel_tasks); + + // init() performs the bootstrap submission of the first Context runnable. + ASSERT_TRUE(scanner_context->init().ok()); + + int64_t rows = 0; + bool eos = false; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(20); + while (!eos) { + ASSERT_LT(std::chrono::steady_clock::now(), deadline) << scanner_context->debug_string(); + // One Context submission represents all pending scanners, so the pool never holds more + // than one runnable for this Context. + EXPECT_LE(scheduler.get_queue_size(), 1); + Block block; + Status st = scanner_context->get_block_from_queue(state.get(), &block, &eos, 0); + ASSERT_TRUE(st.ok()) << st.to_string(); + rows += block.rows(); + if (!eos && block.rows() == 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1)); + } + } + + // Every consumed non-EOS scanner was re-admitted until it reported EOS. + EXPECT_EQ(rows, scanner_count * blocks_per_scanner); + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + EXPECT_EQ(scanner_context->_num_finished_scanners, scanner_count); + EXPECT_EQ(scanner_context->_in_flight_tasks_num, 0); + EXPECT_TRUE(scanner_context->_pending_tasks.empty()); + EXPECT_FALSE(scanner_context->is_context_queued(transfer_lock)); + EXPECT_TRUE(scanner_context->_process_status.ok()); + // The successor runnable ramped concurrency to the per-Context limit, and never beyond it. + EXPECT_EQ(peak_running.load(), parallel_tasks); +} + +TEST_F(ScannerContextTest, thread_pool_context_runnable_is_deduplicated) { + const int parallel_tasks = 2; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = -1; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + + std::list> scanners; + for (int i = 0; i < 3; ++i) { + scanners.push_back(std::make_shared(scanner)); + } + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, -1, + scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + scanner_context->_newly_create_free_blocks_num = newly_create_free_blocks_num.get(); + scanner_context->_scanner_memory_used_counter = scanner_memory_used_counter.get(); + + // One worker, parked, with queue capacity: a submitted Context runnable stays observable in + // the queue instead of being executed or rejected. + ThreadPoolSimplifiedScanScheduler scheduler("context_dedup_test", cgroup_cpu_ctl); + ASSERT_TRUE(scheduler.start(1, 1, 4, 1).ok()); + CountDownLatch task_started(1); + CountDownLatch release_task(1); + Defer cleanup = [&] { + release_task.count_down(); + scheduler.stop(); + }; + ASSERT_TRUE(scheduler + .submit_scan_task(SimplifiedScanTask( + [&] { + task_started.count_down(); + release_task.wait(); + return true; + }, + nullptr, nullptr)) + .ok()); + ASSERT_TRUE(task_started.wait_for(std::chrono::seconds(5))); + scanner_context->_scanner_scheduler = &scheduler; + scanner_context->_min_scan_concurrency_of_scan_scheduler = 20; + + auto completed_task = std::make_shared(scanners.front()); + { + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + ASSERT_TRUE(scheduler.schedule_scan_task(scanner_context, nullptr, transfer_lock).ok()); + EXPECT_TRUE(scanner_context->is_context_queued(transfer_lock)); + EXPECT_EQ(scheduler.get_queue_size(), 1); + + // A second scheduling attempt while a runnable is queued must not add another runnable. + ASSERT_TRUE(scheduler.schedule_scan_task(scanner_context, nullptr, transfer_lock).ok()); + EXPECT_TRUE(scanner_context->is_context_queued(transfer_lock)); + EXPECT_EQ(scheduler.get_queue_size(), 1); + + // Publish a completed non-EOS result so the operator can consume it below. + completed_task->set_state(ScanTask::State::IN_FLIGHT); + completed_task->cached_block = Block::create_unique(); + completed_task->set_state(ScanTask::State::COMPLETED); + scanner_context->_completed_tasks.push_back(completed_task); + scanner_context->_in_flight_tasks_num = 1; + } + + // Consuming a non-EOS result returns the scanner to the admission queue. The queued runnable + // will see it, so no additional runnable is submitted. + Block block; + bool eos = false; + Status st = scanner_context->get_block_from_queue(state.get(), &block, &eos, 0); + ASSERT_TRUE(st.ok()) << st.to_string(); + EXPECT_FALSE(eos); + { + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + EXPECT_EQ(completed_task->_state, ScanTask::State::PENDING); + EXPECT_EQ(completed_task->cached_block, nullptr); + ASSERT_FALSE(scanner_context->_pending_tasks.empty()); + EXPECT_EQ(scanner_context->_pending_tasks.top(), completed_task); + EXPECT_TRUE(scanner_context->is_context_queued(transfer_lock)); + EXPECT_EQ(scheduler.get_queue_size(), 1); + // Cancel the query before the parked worker runs the queued runnable, so it exits without + // touching the OlapScanner that has no tablet behind it. + scanner_context->_should_stop = true; + } +} + +TEST_F(ScannerContextTest, thread_pool_stopped_scheduler_fails_context) { + const int parallel_tasks = 2; + auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, + parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = -1; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + + std::list> scanners { + std::make_shared(scanner)}; + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, -1, + scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + + ThreadPoolSimplifiedScanScheduler scheduler("stopped_scheduler_test", cgroup_cpu_ctl); + ASSERT_TRUE(scheduler.start(1, 1, 1, 1).ok()); + scheduler.stop(); + scanner_context->_scanner_scheduler = &scheduler; + + std::unique_lock transfer_lock(scanner_context->transfer_lock()); + ASSERT_FALSE(scanner_context->_pending_tasks.empty()); + Status surfaced = scheduler.schedule_scan_task(scanner_context, nullptr, transfer_lock); + EXPECT_TRUE(surfaced.is()) << surfaced.to_string(); + // The Context is terminal and the operator is woken to observe the failure. No runnable was + // submitted, so the marker stays clear. + EXPECT_TRUE(scanner_context->done()); + EXPECT_FALSE(scanner_context->_process_status.ok()); + EXPECT_TRUE(scan_dependency->ready()); + EXPECT_FALSE(scanner_context->is_context_queued(transfer_lock)); +} + TEST_F(ScannerContextTest, schedule_scan_task) { const int parallel_tasks = 4; auto scan_operator = std::make_unique(obj_pool.get(), tnode, 0, *descs, @@ -950,6 +1639,7 @@ TEST_F(ScannerContextTest, get_block_from_queue) { std::shared_ptr scanner_context = ScannerContext::create_shared( state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, limit, scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + shared_limit.store(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->_max_bytes_in_queue = 200; @@ -999,6 +1689,81 @@ TEST_F(ScannerContextTest, get_block_from_queue) { EXPECT_EQ(scanner_context->_num_finished_scanners, 1); } +TEST_F(ScannerContextTest, terminal_eos_skips_context_submission) { + ThreadPoolSimplifiedScanScheduler scheduler("terminal_eos_test", cgroup_cpu_ctl); + ASSERT_TRUE(scheduler.start(1, 1, 0, 1).ok()); + CountDownLatch task_started(1); + CountDownLatch release_task(1); + Defer cleanup = [&] { + release_task.count_down(); + scheduler.stop(); + }; + ASSERT_TRUE(scheduler + .submit_scan_task(SimplifiedScanTask( + [&] { + task_started.count_down(); + release_task.wait(); + return true; + }, + nullptr, nullptr)) + .ok()); + ASSERT_TRUE(task_started.wait_for(std::chrono::seconds(5))); + ASSERT_EQ(scheduler.get_active_threads(), 1); + + auto verify_terminal_context = [&](int scanner_count, int64_t remaining_limit) { + const int parallel_tasks = 1; + auto scan_operator = std::make_unique( + obj_pool.get(), tnode, 0, *descs, parallel_tasks, TQueryCacheParam {}); + auto olap_scan_local_state = + OlapScanLocalState::create_unique(state.get(), scan_operator.get()); + + OlapScanner::Params scanner_params; + scanner_params.state = state.get(); + scanner_params.profile = profile.get(); + scanner_params.limit = 100; + scanner_params.key_ranges = std::vector(); + std::shared_ptr scanner = + OlapScanner::create_shared(olap_scan_local_state.get(), std::move(scanner_params)); + + std::list> scanners; + for (int i = 0; i < scanner_count; ++i) { + scanners.push_back(std::make_shared(scanner)); + } + auto scanner_context = ScannerContext::create_shared( + state.get(), olap_scan_local_state.get(), output_tuple_desc, false, scanners, 100, + scan_dependency, &shared_limit, nullptr, nullptr, 0, false, parallel_tasks); + scanner_context->_scanner_scheduler = &scheduler; + + scanner_context->_pending_tasks = std::stack>(); + auto scanner_iter = scanners.begin(); + auto eos_task = std::make_shared(*scanner_iter++); + eos_task->set_state(ScanTask::State::IN_FLIGHT); + eos_task->set_state(ScanTask::State::EOS); + scanner_context->_completed_tasks.push_back(eos_task); + while (scanner_iter != scanners.end()) { + scanner_context->_pending_tasks.push(std::make_shared(*scanner_iter++)); + } + scanner_context->_in_flight_tasks_num = 0; + shared_limit.store(remaining_limit); + + MockRuntimeStateLocal mock_runtime_state; + EXPECT_CALL(mock_runtime_state, is_cancelled()).WillRepeatedly(testing::Return(false)); + Block block; + bool eos = false; + Status status = scanner_context->get_block_from_queue(&mock_runtime_state, &block, &eos, 0); + + EXPECT_TRUE(status.ok()) << status.to_string(); + EXPECT_TRUE(eos); + EXPECT_EQ(scheduler.get_queue_size(), 0); + }; + + // All scanners completed: no runnable is needed even if the pool cannot accept one. + verify_terminal_context(1, 100); + // Shared LIMIT completed the Context while another scanner is pending. Check terminal state + // before rescheduling so the full pool cannot turn successful EOS into TOO_MANY_TASKS. + verify_terminal_context(2, 0); +} + /** MemShareArbitrator Tests (5 tests) - scanner_mem_share_arbitrator_basic: Tests initialization, query_id, memory limits, and initial state