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