[improvement](be) Refactor thread-pool scan scheduling

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<ScanTask> scan_task,
                                         std::unique_lock<std::mutex>& /*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<ScanTask> scan_task) {
+void ScannerContext::push_completed_scan_task(std::shared_ptr<ScanTask> 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<ScanTask> 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<std::mutex> 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<std::mutex>& transfer_lock) const {
+    DORIS_CHECK(transfer_lock.owns_lock());
+    return _is_context_queued;
+}
+
+void ScannerContext::set_context_queued(bool queued,
+                                        const std::unique_lock<std::mutex>& 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<std::mutex>& 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<ScanTask> scan_task,
+                                            const std::unique_lock<std::mutex>& 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<std::mutex>& 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<int32_t>(_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<ScanTask> ScannerContext::try_get_next_scan_task(
+        const std::unique_lock<std::mutex>& 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<void>(_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);

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<ScannerContext>,
     // set the `eos` to `ScanTask::eos` if there is no more data in current scanner
     Status submit_scan_task(std::shared_ptr<ScanTask> scan_task, std::unique_lock<std::mutex>&);
 
-    // Push back a scan task.
-    void push_back_scan_task(std::shared_ptr<ScanTask> 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<ScanTask> 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<ScannerContext>,
                               std::unique_lock<std::mutex>& transfer_lock,
                               std::unique_lock<std::shared_mutex>& 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<std::mutex>& 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<std::mutex>& 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<std::mutex>& 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<ScanTask> scan_task,
+                                const std::unique_lock<std::mutex>& 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<std::mutex>& 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<ScanTask> try_get_next_scan_task(
+            const std::unique_lock<std::mutex>& 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<ScannerContext>,
     //   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<std::shared_ptr<ScanTask>> _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<std::shared_ptr<ScanTask>> _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.

be/src/exec/scan/scanner_scheduler.cpp

@@ -75,17 +75,7 @@ Status ScannerScheduler::submit(std::shared_ptr<ScannerContext> 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<ScannerContext> ctx,
     return Status::OK();
 }
 
+bool ScannerScheduler::execute_scan_task(const std::shared_ptr<ScannerContext>& ctx,
+                                         const std::shared_ptr<ScanTask>& 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<ScannerContext> ctx,
                                    const Status& st, size_t reserve_size) {
     ctx->clear_free_blocks();
@@ -308,7 +314,7 @@ void ScannerScheduler::_scanner_scan(std::shared_ptr<ScannerContext> 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)