Skip to content

[fix](be) Account for TaskExecutor admission in scanner scheduling #66838

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
mrhhsg wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

[fix](be) Account for TaskExecutor admission in scanner scheduling #66838

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
186 changes: 110 additions & 76 deletions be/src/exec/scan/scanner_context.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
#include <glog/logging.h>
#include <zconf.h>

#include <algorithm>
#include <chrono>
#include <cstdint>
#include <ctime>
#include <memory>
Expand All @@ -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"
Expand Down Expand Up @@ -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<TaskExecutor>* task_executor,
std::shared_ptr<TaskHandle>* task_handle) {
if (auto* task_executor_scheduler =
dynamic_cast<TaskExecutorSimplifiedScanScheduler*>(scanner_scheduler)) {
std::shared_ptr<TaskExecutor> 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
Expand All @@ -197,20 +224,10 @@ Status ScannerContext::init() {

auto scanner = _all_scanners.front().lock();
DCHECK(scanner != nullptr);

if (auto* task_executor_scheduler =
dynamic_cast<TaskExecutorSimplifiedScanScheduler*>(_scanner_scheduler)) {
std::shared_ptr<TaskExecutor> 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.
Expand Down Expand Up @@ -580,13 +597,19 @@ void ScannerContext::reestimated_block_mem_bytes(int64_t num) {

int32_t ScannerContext::_get_margin(std::unique_lock<std::mutex>& transfer_lock,
std::unique_lock<std::shared_mutex>& scheduler_lock) {
return _get_margin(transfer_lock, scheduler_lock, _in_flight_tasks_num.load());
}

int32_t ScannerContext::_get_margin(std::unique_lock<std::mutex>& transfer_lock,
std::unique_lock<std::shared_mutex>& 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<int32_t>(_completed_tasks.size()) + _in_flight_tasks_num);
int32_t margin_1 =
_min_scan_concurrency - (cast_set<int32_t>(_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 =
Expand All @@ -596,7 +619,7 @@ int32_t ScannerContext::_get_margin(std::unique_lock<std::mutex>& transfer_lock,
// margin_3 is used to respect adaptive max concurrency limit
int32_t margin_3 =
std::max(effective_max_concurrency -
(cast_set<int32_t>(_completed_tasks.size()) + _in_flight_tasks_num),
(cast_set<int32_t>(_completed_tasks.size()) + in_flight_tasks),
1);

if (margin_1 <= 0 && margin_2 <= 0) {
Expand All @@ -611,24 +634,32 @@ int32_t ScannerContext::_get_margin(std::unique_lock<std::mutex>& 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<ScanTask> current_scan_task,
std::unique_lock<std::mutex>& transfer_lock,
std::unique_lock<std::shared_mutex>& scheduler_lock) {
Expand All @@ -639,76 +670,74 @@ Status ScannerContext::schedule_scan_task(std::shared_ptr<ScanTask> current_scan

std::list<std::shared_ptr<ScanTask>> 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<int32_t>(
_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<int32_t>(tasks_to_submit.size()), 0);
while (remaining_submission_margin-- > 0) {
std::shared_ptr<ScanTask> task_to_run;
const int32_t current_concurrency = cast_set<int32_t>(
const auto current_concurrency = cast_set<int32_t>(
_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();
}
Expand All @@ -730,7 +759,8 @@ Status ScannerContext::schedule_scan_task(std::shared_ptr<ScanTask> current_scan
}

std::shared_ptr<ScanTask> ScannerContext::_pull_next_scan_task(
std::shared_ptr<ScanTask> current_scan_task, int32_t current_concurrency) {
std::shared_ptr<ScanTask> 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
Expand All @@ -747,6 +777,7 @@ std::shared_ptr<ScanTask> 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.");
Expand All @@ -755,6 +786,9 @@ std::shared_ptr<ScanTask> 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.
Expand Down
26 changes: 18 additions & 8 deletions be/src/exec/scan/scanner_context.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
};
Expand Down Expand Up @@ -289,10 +289,14 @@ class ScannerContext : public std::enable_shared_from_this<ScannerContext>,
// 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)-->
Expand All @@ -312,9 +316,9 @@ class ScannerContext : public std::enable_shared_from_this<ScannerContext>,
// _pull_next_scan_task() during scheduling.
std::stack<std::shared_ptr<ScanTask>> _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.
Expand Down Expand Up @@ -347,10 +351,16 @@ class ScannerContext : public std::enable_shared_from_this<ScannerContext>,
MOCK_REMOVE(const) int32_t _min_scan_concurrency = 1;

std::shared_ptr<ScanTask> _pull_next_scan_task(std::shared_ptr<ScanTask> current_scan_task,
int32_t current_concurrency);
int32_t current_concurrency,
bool only_existing_split = false);

int32_t _get_margin(std::unique_lock<std::mutex>& transfer_lock,
std::unique_lock<std::shared_mutex>& scheduler_lock);
int32_t _get_margin(std::unique_lock<std::mutex>& transfer_lock,
std::unique_lock<std::shared_mutex>& scheduler_lock,
int32_t in_flight_tasks);

int32_t _admitted_in_flight_tasks() const;

// Memory-aware adaptive scheduling
std::shared_ptr<MemLimiter> _scanner_mem_limiter = nullptr;
Expand Down
2 changes: 2 additions & 0 deletions be/src/exec/scan/task_executor/task_handle.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
5 changes: 5 additions & 0 deletions be/src/exec/scan/task_executor/time_sharing/time_sharing_task_handle.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -133,6 +133,11 @@ int TimeSharingTaskHandle::running_leaf_splits() const {
return static_cast<int>(_running_leaf_splits.size());
}

int TimeSharingTaskHandle::queued_leaf_splits() const {
std::lock_guard<std::mutex> lock(_mutex);
return static_cast<int>(_queued_leaf_splits.size());
}

int64_t TimeSharingTaskHandle::scheduled_nanos() const {
std::lock_guard<std::mutex> lock(_mutex);
return _scheduled_nanos;
Expand Down
Loading
Loading