diff --git a/be/src/exec/scan/scanner_context.cpp b/be/src/exec/scan/scanner_context.cpp index 3371af5bf8e829..279040c5168279 100644 --- a/be/src/exec/scan/scanner_context.cpp +++ b/be/src/exec/scan/scanner_context.cpp @@ -40,6 +40,7 @@ #include "exec/operator/scan_operator.h" #include "exec/scan/scan_node.h" #include "exec/scan/scanner_scheduler.h" +#include "exec/scan/task_executor/task_executor.h" #include "runtime/descriptors.h" #include "runtime/exec_env.h" #include "runtime/runtime_profile.h" @@ -123,6 +124,7 @@ Status ScannerContext::init() { if (auto* task_executor_scheduler = dynamic_cast(_scanner_scheduler)) { std::shared_ptr task_executor = task_executor_scheduler->task_executor(); + _task_executor = task_executor; TaskId task_id(fmt::format("{}-{}", print_id(_state->query_id()), ctx_id)); _task_handle = DORIS_TRY(task_executor->create_task( task_id, []() { return 0.0; }, @@ -187,11 +189,11 @@ ScannerContext::~ScannerContext() { block.reset(); DorisMetrics::instance()->scanner_ctx_cnt->increment(-1); if (_task_handle) { - if (auto* task_executor_scheduler = - dynamic_cast(_scanner_scheduler)) { - static_cast(task_executor_scheduler->task_executor()->remove_task(_task_handle)); + if (auto task_executor = _task_executor.lock()) { + static_cast(task_executor->remove_task(_task_handle)); } _task_handle = nullptr; + _task_executor.reset(); } } @@ -389,11 +391,11 @@ void ScannerContext::stop_scanners(RuntimeState* state) { } _tasks_queue.clear(); if (_task_handle) { - if (auto* task_executor_scheduler = - dynamic_cast(_scanner_scheduler)) { - static_cast(task_executor_scheduler->task_executor()->remove_task(_task_handle)); + if (auto task_executor = _task_executor.lock()) { + static_cast(task_executor->remove_task(_task_handle)); } _task_handle = nullptr; + _task_executor.reset(); } // TODO yiguolei, call mark close to scanners if (state->enable_profile()) { diff --git a/be/src/exec/scan/scanner_context.h b/be/src/exec/scan/scanner_context.h index 857dfb736e91bd..ba1c323566062b 100644 --- a/be/src/exec/scan/scanner_context.h +++ b/be/src/exec/scan/scanner_context.h @@ -239,6 +239,7 @@ class ScannerContext : public std::enable_shared_from_this, std::shared_ptr _resource_ctx; std::shared_ptr _dependency = nullptr; std::shared_ptr _task_handle; + std::weak_ptr _task_executor; std::atomic _block_memory_usage = 0; diff --git a/be/src/exec/scan/scanner_scheduler.h b/be/src/exec/scan/scanner_scheduler.h index f1939409bda579..66d5fd55f65cb7 100644 --- a/be/src/exec/scan/scanner_scheduler.h +++ b/be/src/exec/scan/scanner_scheduler.h @@ -351,10 +351,14 @@ class TaskExecutorSimplifiedScanScheduler final : public ScannerScheduler { : std::max(48, CpuInfo::num_cores() * 2), std::chrono::milliseconds(100), std::nullopt)); - auto wrapped_scan_func = [this, task_handle, scan_func = scan_task.scan_func]() { + std::weak_ptr task_executor = _task_executor; + auto wrapped_scan_func = [task_executor, task_handle, + scan_func = scan_task.scan_func]() { bool result = scan_func(); if (result) { - static_cast(_task_executor->remove_task(task_handle)); + if (auto executor = task_executor.lock()) { + static_cast(executor->remove_task(task_handle)); + } } return result; };