From c7e85f08b3286e55192b68ae0eb41982f4657578 Mon Sep 17 00:00:00 2001 From: Hu Shenggang Date: Mon, 3 Aug 2026 19:43:32 +0800 Subject: [PATCH] [fix](pipeline) Prevent finalized tasks from being rescheduled ### What problem does this PR solve? Issue Number: None Problem Summary: A delayed dependency wake-up can leave a terminal pipeline task in the runnable queue. The stale task may be executed again after its shared operator state has been released, causing invalid virtual calls or use-after-free. Keep queued tasks and their fragment contexts alive, store dependency waiters weakly, atomically claim PipelineX tasks before execution, and drop terminal tasks before accessing their execution state. Drain pending queue entries safely during shutdown. ### Release note None ### Check List (For Author) - Test: - Unit Test: `./run-be-ut.sh --run --filter='TaskQueueTest.*:TaskSchedulerTest.*' -j 24` - Regression test: Not required for this scheduler lifecycle fix - Behavior changed: No - Does this need documentation: No --- be/src/pipeline/pipeline_task.cpp | 5 + be/src/pipeline/pipeline_task.h | 7 +- be/src/pipeline/pipeline_x/dependency.cpp | 13 +- be/src/pipeline/pipeline_x/dependency.h | 2 +- .../pipeline/pipeline_x/pipeline_x_task.cpp | 3 + be/src/pipeline/task_queue.cpp | 54 +++-- be/src/pipeline/task_queue.h | 31 ++- be/src/pipeline/task_scheduler.cpp | 15 +- be/test/pipeline/task_queue_test.cpp | 186 ++++++++++++++++++ 9 files changed, 277 insertions(+), 39 deletions(-) create mode 100644 be/test/pipeline/task_queue_test.cpp diff --git a/be/src/pipeline/pipeline_task.cpp b/be/src/pipeline/pipeline_task.cpp index 242837da6b986b..6be5d792c16ba0 100644 --- a/be/src/pipeline/pipeline_task.cpp +++ b/be/src/pipeline/pipeline_task.cpp @@ -208,6 +208,11 @@ void PipelineTask::set_task_queue(TaskQueue* task_queue) { _task_queue = task_queue; } +std::shared_ptr PipelineTask::get_task_holder() { + auto context_holder = _fragment_context->shared_from_this(); + return std::shared_ptr(context_holder, this); +} + Status PipelineTask::execute(bool* eos) { SCOPED_TIMER(_task_profile->total_time_counter()); SCOPED_TIMER(_exec_timer); diff --git a/be/src/pipeline/pipeline_task.h b/be/src/pipeline/pipeline_task.h index 08fdb32d7fe34e..c9dc518cc28f8e 100644 --- a/be/src/pipeline/pipeline_task.h +++ b/be/src/pipeline/pipeline_task.h @@ -259,7 +259,12 @@ class PipelineTask { virtual bool is_pipelineX() const { return false; } bool is_running() { return _running.load(); } - void set_running(bool running) { _running = running; } + // Return the previous state so a scheduler can atomically claim the task. + bool set_running(bool running) { return _running.exchange(running); } + + // Pipeline tasks are owned by their fragment context. The aliasing shared pointer keeps the + // context (and therefore this task) alive while the task is waiting in an asynchronous queue. + std::shared_ptr get_task_holder(); bool is_exceed_debug_timeout() { if (_has_exceed_timeout) { diff --git a/be/src/pipeline/pipeline_x/dependency.cpp b/be/src/pipeline/pipeline_x/dependency.cpp index 66aa14e64b029b..3008960c117506 100644 --- a/be/src/pipeline/pipeline_x/dependency.cpp +++ b/be/src/pipeline/pipeline_x/dependency.cpp @@ -44,9 +44,10 @@ Dependency* BasicSharedState::create_sink_dependency(int dest_id, int node_id, s } void Dependency::_add_block_task(PipelineXTask* task) { - DCHECK(_blocked_task.empty() || _blocked_task[_blocked_task.size() - 1] != task) + auto task_holder = std::static_pointer_cast(task->get_task_holder()); + DCHECK(_blocked_task.empty() || _blocked_task.back().lock().get() != task) << "Duplicate task: " << task->debug_string(); - _blocked_task.push_back(task); + _blocked_task.emplace_back(task_holder); } void Dependency::set_ready() { @@ -54,7 +55,7 @@ void Dependency::set_ready() { return; } _watcher.stop(); - std::vector local_block_task {}; + std::vector> local_block_task {}; { std::unique_lock lc(_task_lock); if (_ready) { @@ -63,8 +64,10 @@ void Dependency::set_ready() { _ready = true; local_block_task.swap(_blocked_task); } - for (auto* task : local_block_task) { - task->wake_up(); + for (auto& task : local_block_task) { + if (auto task_holder = task.lock()) { + task_holder->wake_up(); + } } } diff --git a/be/src/pipeline/pipeline_x/dependency.h b/be/src/pipeline/pipeline_x/dependency.h index be2876927248b8..88bc592a92a217 100644 --- a/be/src/pipeline/pipeline_x/dependency.h +++ b/be/src/pipeline/pipeline_x/dependency.h @@ -162,7 +162,7 @@ class Dependency : public std::enable_shared_from_this { MonotonicStopWatch _watcher; std::mutex _task_lock; - std::vector _blocked_task; + std::vector> _blocked_task; // If `_always_ready` is true, `block()` will never block tasks. std::atomic _always_ready = false; diff --git a/be/src/pipeline/pipeline_x/pipeline_x_task.cpp b/be/src/pipeline/pipeline_x/pipeline_x_task.cpp index d158057173ee6a..043b9a231c0c19 100644 --- a/be/src/pipeline/pipeline_x/pipeline_x_task.cpp +++ b/be/src/pipeline/pipeline_x/pipeline_x_task.cpp @@ -547,6 +547,9 @@ std::string PipelineXTask::debug_string() { void PipelineXTask::wake_up() { // call by dependency + if (is_finished()) { + return; + } static_cast(get_task_queue()->push_back(this)); } } // namespace doris::pipeline diff --git a/be/src/pipeline/task_queue.cpp b/be/src/pipeline/task_queue.cpp index 4ed2faf0ca6028..3256cee03ea4a1 100644 --- a/be/src/pipeline/task_queue.cpp +++ b/be/src/pipeline/task_queue.cpp @@ -29,11 +29,11 @@ namespace pipeline { TaskQueue::~TaskQueue() = default; -PipelineTask* SubTaskQueue::try_take(bool is_steal) { +PipelineTaskSPtr SubTaskQueue::try_take(bool is_steal) { if (_queue.empty()) { return nullptr; } - auto task = _queue.front(); + auto task = std::move(_queue.front()); _queue.pop(); return task; } @@ -49,13 +49,30 @@ PriorityTaskQueue::PriorityTaskQueue() : _closed(false) { } void PriorityTaskQueue::close() { - std::unique_lock lock(_work_size_mutex); - _closed = true; - _wait_task.notify_all(); - DorisMetrics::instance()->pipeline_task_queue_size->increment(-_total_task_size); + std::vector pending_tasks; + { + std::unique_lock lock(_work_size_mutex); + if (_closed) { + return; + } + _closed = true; + _wait_task.notify_all(); + const auto pending_task_size = _total_task_size.exchange(0); + DorisMetrics::instance()->pipeline_task_queue_size->increment( + -static_cast(pending_task_size)); + pending_tasks.reserve(pending_task_size); + for (auto& queue : _sub_queues) { + queue.drain(&pending_tasks); + } + } + // Releasing a task may destroy its fragment context, so do it outside the queue lock. + for (const auto& task : pending_tasks) { + task->pop_out_runnable_queue(); + } + pending_tasks.clear(); } -PipelineTask* PriorityTaskQueue::_try_take_unprotected(bool is_steal) { +PipelineTaskSPtr PriorityTaskQueue::_try_take_unprotected(bool is_steal) { if (_total_task_size == 0 || _closed) { return nullptr; } @@ -92,13 +109,13 @@ int PriorityTaskQueue::_compute_level(uint64_t runtime) { return SUB_QUEUE_LEVEL - 1; } -PipelineTask* PriorityTaskQueue::try_take(bool is_steal) { +PipelineTaskSPtr PriorityTaskQueue::try_take(bool is_steal) { // TODO other efficient lock? e.g. if get lock fail, return null_ptr std::unique_lock lock(_work_size_mutex); return _try_take_unprotected(is_steal); } -PipelineTask* PriorityTaskQueue::take(uint32_t timeout_ms) { +PipelineTaskSPtr PriorityTaskQueue::take(uint32_t timeout_ms) { std::unique_lock lock(_work_size_mutex); auto task = _try_take_unprotected(false); if (task) { @@ -113,12 +130,12 @@ PipelineTask* PriorityTaskQueue::take(uint32_t timeout_ms) { } } -Status PriorityTaskQueue::push(PipelineTask* task) { +Status PriorityTaskQueue::push(PipelineTaskSPtr task) { + auto level = _compute_level(task->get_runtime_ns()); + std::unique_lock lock(_work_size_mutex); if (_closed) { return Status::InternalError("WorkTaskQueue closed"); } - auto level = _compute_level(task->get_runtime_ns()); - std::unique_lock lock(_work_size_mutex); // update empty queue's runtime, to avoid too high priority if (_sub_queues[level].empty() && @@ -126,7 +143,8 @@ Status PriorityTaskQueue::push(PipelineTask* task) { _sub_queues[level].adjust_runtime(_queue_level_min_vruntime); } - _sub_queues[level].push_back(task); + task->put_in_runnable_queue(); + _sub_queues[level].push_back(std::move(task)); _total_task_size++; DorisMetrics::instance()->pipeline_task_queue_size->increment(1); _wait_task.notify_one(); @@ -148,8 +166,8 @@ void MultiCoreTaskQueue::close() { [](auto& prio_task_queue) { prio_task_queue.close(); }); } -PipelineTask* MultiCoreTaskQueue::take(int core_id) { - PipelineTask* task = nullptr; +PipelineTaskSPtr MultiCoreTaskQueue::take(int core_id) { + PipelineTaskSPtr task = nullptr; while (!_closed) { DCHECK(_prio_task_queue_list.size() > core_id) << " list size: " << _prio_task_queue_list.size() << " core_id: " << core_id @@ -175,7 +193,7 @@ PipelineTask* MultiCoreTaskQueue::take(int core_id) { return task; } -PipelineTask* MultiCoreTaskQueue::_steal_take(int core_id) { +PipelineTaskSPtr MultiCoreTaskQueue::_steal_take(int core_id) { DCHECK(core_id < _core_size); int next_id = core_id; for (int i = 1; i < _core_size; ++i) { @@ -203,8 +221,8 @@ Status MultiCoreTaskQueue::push_back(PipelineTask* task) { Status MultiCoreTaskQueue::push_back(PipelineTask* task, int core_id) { DCHECK(core_id < _core_size); - task->put_in_runnable_queue(); - return _prio_task_queue_list[core_id].push(task); + auto task_holder = task->get_task_holder(); + return _prio_task_queue_list[core_id].push(std::move(task_holder)); } } // namespace pipeline diff --git a/be/src/pipeline/task_queue.h b/be/src/pipeline/task_queue.h index 1c6d2be4929167..d0dd5cef424312 100644 --- a/be/src/pipeline/task_queue.h +++ b/be/src/pipeline/task_queue.h @@ -27,6 +27,8 @@ #include #include #include +#include +#include #include "common/status.h" #include "pipeline_task.h" @@ -35,6 +37,8 @@ namespace doris { namespace pipeline { +using PipelineTaskSPtr = std::shared_ptr; + class TaskQueue { public: TaskQueue(int core_size) : _core_size(core_size) {} @@ -42,7 +46,7 @@ class TaskQueue { virtual void close() = 0; // Get the task by core id. // TODO: To think the logic is useful? - virtual PipelineTask* take(int core_id) = 0; + virtual PipelineTaskSPtr take(int core_id) = 0; // push from scheduler virtual Status push_back(PipelineTask* task) = 0; @@ -63,9 +67,16 @@ class SubTaskQueue { friend class PriorityTaskQueue; public: - void push_back(PipelineTask* task) { _queue.emplace(task); } + void push_back(PipelineTaskSPtr task) { _queue.emplace(std::move(task)); } + + PipelineTaskSPtr try_take(bool is_steal); - PipelineTask* try_take(bool is_steal); + void drain(std::vector* tasks) { + while (!_queue.empty()) { + tasks->emplace_back(std::move(_queue.front())); + _queue.pop(); + } + } void set_level_factor(double level_factor) { _level_factor = level_factor; } @@ -81,7 +92,7 @@ class SubTaskQueue { bool empty() { return _queue.empty(); } private: - std::queue _queue; + std::queue _queue; // depends on LEVEL_QUEUE_TIME_FACTOR double _level_factor = 1; @@ -95,18 +106,18 @@ class PriorityTaskQueue { void close(); - PipelineTask* try_take(bool is_steal); + PipelineTaskSPtr try_take(bool is_steal); - PipelineTask* take(uint32_t timeout_ms = 0); + PipelineTaskSPtr take(uint32_t timeout_ms = 0); - Status push(PipelineTask* task); + Status push(PipelineTaskSPtr task); void inc_sub_queue_runtime(int level, uint64_t runtime) { _sub_queues[level].inc_runtime(runtime); } private: - PipelineTask* _try_take_unprotected(bool is_steal); + PipelineTaskSPtr _try_take_unprotected(bool is_steal); static constexpr auto LEVEL_QUEUE_TIME_FACTOR = 2; static constexpr size_t SUB_QUEUE_LEVEL = 6; SubTaskQueue _sub_queues[SUB_QUEUE_LEVEL]; @@ -135,7 +146,7 @@ class MultiCoreTaskQueue : public TaskQueue { void close() override; // Get the task by core id. - PipelineTask* take(int core_id) override; + PipelineTaskSPtr take(int core_id) override; // TODO combine these methods to `push_back(task, core_id = -1)` Status push_back(PipelineTask* task) override; @@ -149,7 +160,7 @@ class MultiCoreTaskQueue : public TaskQueue { } private: - PipelineTask* _steal_take(int core_id); + PipelineTaskSPtr _steal_take(int core_id); std::vector _prio_task_queue_list; std::atomic _next_core = 0; diff --git a/be/src/pipeline/task_scheduler.cpp b/be/src/pipeline/task_scheduler.cpp index 93cf2d9dd993b6..fb18a29bf365ff 100644 --- a/be/src/pipeline/task_scheduler.cpp +++ b/be/src/pipeline/task_scheduler.cpp @@ -264,16 +264,23 @@ void _close_task(PipelineTask* task, PipelineTaskState state, Status exec_status void TaskScheduler::_do_work(size_t index) { const auto& marker = _markers[index]; while (*marker) { - auto* task = _task_queue->take(index); - if (!task) { + auto task_holder = _task_queue->take(index); + if (!task_holder) { continue; } - if (task->is_pipelineX() && task->is_running()) { + auto* task = task_holder.get(); + if (task->is_pipelineX() && task->set_running(true)) { static_cast(_task_queue->push_back(task, index)); continue; } + if (task->is_finished()) { + task->set_running(false); + continue; + } task->log_detail_if_need(); - task->set_running(true); + if (!task->is_pipelineX()) { + task->set_running(true); + } task->set_task_queue(_task_queue.get()); auto* fragment_ctx = task->fragment_context(); bool canceled = fragment_ctx->is_canceled(); diff --git a/be/test/pipeline/task_queue_test.cpp b/be/test/pipeline/task_queue_test.cpp new file mode 100644 index 00000000000000..958cab7450649e --- /dev/null +++ b/be/test/pipeline/task_queue_test.cpp @@ -0,0 +1,186 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#include "pipeline/task_queue.h" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "pipeline/pipeline.h" +#include "pipeline/task_scheduler.h" + +namespace doris::pipeline { +namespace { + +class TestPipelineTask : public PipelineTask { +public: + explicit TestPipelineTask(PipelinePtr& pipeline) + : PipelineTask(pipeline, 0, nullptr, nullptr, nullptr) {} +}; + +struct TaskOwner { + explicit TaskOwner(std::atomic* destroyed) + : pipeline(std::make_shared(0, 1, std::weak_ptr())), + task(pipeline), + _destroyed(destroyed) {} + + ~TaskOwner() { _destroyed->store(true); } + + PipelinePtr pipeline; + TestPipelineTask task; + std::atomic* _destroyed; +}; + +TEST(TaskQueueTest, SubTaskQueueKeepsTaskOwnerAlive) { + std::atomic destroyed = false; + auto owner = std::make_shared(&destroyed); + std::weak_ptr weak_owner = owner; + PipelineTaskSPtr task_holder(owner, &owner->task); + + SubTaskQueue queue; + queue.push_back(task_holder); + task_holder.reset(); + owner.reset(); + EXPECT_FALSE(weak_owner.expired()); + + auto taken = queue.try_take(false); + ASSERT_NE(taken, nullptr); + EXPECT_FALSE(weak_owner.expired()); + + taken.reset(); + EXPECT_TRUE(weak_owner.expired()); + EXPECT_TRUE(destroyed.load()); +} + +TEST(TaskQueueTest, CloseReleasesPendingTaskOwner) { + std::atomic destroyed = false; + auto owner = std::make_shared(&destroyed); + std::weak_ptr weak_owner = owner; + PipelineTaskSPtr task_holder(owner, &owner->task); + + PriorityTaskQueue queue; + ASSERT_TRUE(queue.push(task_holder).ok()); + task_holder.reset(); + owner.reset(); + EXPECT_FALSE(weak_owner.expired()); + + queue.close(); + queue.close(); + EXPECT_TRUE(weak_owner.expired()); + EXPECT_TRUE(destroyed.load()); +} + +class FinishedPipelineTask final : public TestPipelineTask { +public: + FinishedPipelineTask(PipelinePtr& pipeline, std::promise* finished_checked, + std::atomic* execute_calls) + : TestPipelineTask(pipeline), + _finished_checked(finished_checked), + _execute_calls(execute_calls) {} + + bool is_pipelineX() const override { return true; } + + bool is_finished() const override { + if (!_check_reported.exchange(true)) { + _finished_checked->set_value(); + } + return true; + } + + Status execute(bool* /*eos*/) override { + _execute_calls->fetch_add(1); + return Status::OK(); + } + +private: + std::promise* _finished_checked; + std::atomic* _execute_calls; + mutable std::atomic _check_reported = false; +}; + +struct FinishedTaskOwner { + FinishedTaskOwner(std::promise* finished_checked, std::atomic* execute_calls) + : pipeline(std::make_shared(0, 1, std::weak_ptr())), + task(pipeline, finished_checked, execute_calls) {} + + PipelinePtr pipeline; + FinishedPipelineTask task; +}; + +class OneShotTaskQueue final : public TaskQueue { +public: + explicit OneShotTaskQueue(PipelineTaskSPtr task) : TaskQueue(1), _task(std::move(task)) {} + + void close() override { + { + std::lock_guard lock(_mutex); + _closed = true; + } + _closed_cv.notify_all(); + } + + PipelineTaskSPtr take(int /*core_id*/) override { + std::unique_lock lock(_mutex); + if (_task) { + return std::move(_task); + } + _closed_cv.wait(lock, [this] { return _closed; }); + return nullptr; + } + + Status push_back(PipelineTask* /*task*/) override { + return Status::InternalError("finished task should not be resubmitted"); + } + + Status push_back(PipelineTask* /*task*/, int /*core_id*/) override { + return Status::InternalError("finished task should not be resubmitted"); + } + +private: + std::mutex _mutex; + std::condition_variable _closed_cv; + PipelineTaskSPtr _task; + bool _closed = false; +}; + +TEST(TaskSchedulerTest, SkipsFinishedPipelineTask) { + std::promise finished_checked; + auto checked_future = finished_checked.get_future(); + std::atomic execute_calls = 0; + auto owner = std::make_shared(&finished_checked, &execute_calls); + PipelineTaskSPtr task_holder(owner, &owner->task); + auto task_queue = std::make_shared(std::move(task_holder)); + + TaskScheduler scheduler(nullptr, nullptr, task_queue, "terminal-task-test", nullptr); + ASSERT_TRUE(scheduler.start().ok()); + ASSERT_EQ(checked_future.wait_for(std::chrono::seconds(5)), std::future_status::ready); + scheduler.stop(); + + EXPECT_EQ(execute_calls.load(), 0); + EXPECT_FALSE(owner->task.is_running()); +} + +} // namespace +} // namespace doris::pipeline