update

Author: Gabriel39

be/src/pipeline/dependency.cpp

@@ -70,7 +70,7 @@ void Dependency::set_ready() {
     for (auto task : local_block_task) {
         if (auto t = task.lock()) {
             std::unique_lock<std::mutex> lc(_task_lock);
-            THROW_IF_ERROR(t->wake_up(this));
+            t->wake_up();
         }
     }
 }
@@ -109,7 +109,7 @@ std::string RuntimeFilterDependency::debug_string(int indentation_level) {
     return fmt::to_string(debug_string_buffer);
 }
 
-Dependency* RuntimeFilterDependency::is_blocked_by(PipelineTask* task) {
+Dependency* RuntimeFilterDependency::is_blocked_by(std::shared_ptr<PipelineTask> task) {
     std::unique_lock<std::mutex> lc(_task_lock);
     auto ready = _ready.load();
     if (!ready && task) {

be/src/pipeline/dependency.h

@@ -282,7 +282,7 @@ class RuntimeFilterDependency final : public Dependency {
             : Dependency(id, node_id, name), _runtime_filter(runtime_filter) {}
     std::string debug_string(int indentation_level = 0) override;
 
-    Dependency* is_blocked_by(PipelineTask* task) override;
+    Dependency* is_blocked_by(std::shared_ptr<PipelineTask> task) override;
 
 private:
     const IRuntimeFilter* _runtime_filter = nullptr;

be/src/pipeline/pipeline.h

@@ -135,6 +135,7 @@ class Pipeline : public std::enable_shared_from_this<Pipeline> {
     }
 
     int num_tasks_of_parent() const { return _num_tasks_of_parent; }
+    std::string& name() { return _name; }
 
 private:
     void _init_profile();

be/src/pipeline/pipeline_fragment_context.cpp

@@ -153,7 +153,6 @@ PipelineFragmentContext::~PipelineFragmentContext() {
     _runtime_filter_states.clear();
     _runtime_filter_mgr_map.clear();
     _op_id_to_le_state.clear();
-    _op_id_to_shared_state.clear();
     _query_ctx.reset();
 }
 
@@ -447,7 +446,8 @@ Status PipelineFragmentContext::_build_pipeline_tasks(const doris::TPipelineFrag
                 task_runtime_state->set_task_id(cur_task_id);
                 task_runtime_state->set_task_num(pipeline->num_tasks());
                 auto task = std::make_shared<PipelineTask>(
-                        pipeline, cur_task_id, task_runtime_state.get(), shared_from_this(),
+                        pipeline, cur_task_id, task_runtime_state.get(),
+                        std::dynamic_pointer_cast<PipelineFragmentContext>(shared_from_this()),
                         pipeline_id_to_profile[pip_idx].get(), get_local_exchange_state(pipeline),
                         i);
                 pipeline->incr_created_tasks(i, task.get());

be/src/pipeline/pipeline_task.cpp

@@ -59,6 +59,7 @@ PipelineTask::PipelineTask(
 #else
           _query_id(fragment_context->get_query_id()),
 #endif
+          _pip_id(pipeline->id()),
           _index(task_id),
           _pipeline(pipeline),
           _opened(false),
@@ -71,7 +72,8 @@ PipelineTask::PipelineTask(
           _sink(pipeline->sink_shared_pointer()),
           _le_state_map(std::move(le_state_map)),
           _task_idx(task_idx),
-          _execution_dep(state->get_query_ctx()->get_execution_dependency()) {
+          _execution_dep(state->get_query_ctx()->get_execution_dependency()),
+          _pipeline_name(_pipeline->name()){
     _pipeline_task_watcher.start();
 
     auto shared_state = _sink->create_shared_state();
@@ -314,11 +316,11 @@ Status PipelineTask::execute(bool* eos) {
         }
         int64_t delta_cpu_time = cpu_time_stop_watch.elapsed_time();
         _task_cpu_timer->update(delta_cpu_time);
-        auto cpu_qs = query_context()->get_cpu_statistics();
+        auto cpu_qs = fragment_context->get_query_ctx()->get_cpu_statistics();
         if (cpu_qs) {
             cpu_qs->add_cpu_nanos(delta_cpu_time);
         }
-        query_context()->update_cpu_time(delta_cpu_time);
+        fragment_context->get_query_ctx()->update_cpu_time(delta_cpu_time);
     }};
     if (_wait_to_start()) {
         if (config::enable_prefetch_tablet) {
@@ -502,13 +504,18 @@ bool PipelineTask::should_revoke_memory(RuntimeState* state, int64_t revocable_m
 void PipelineTask::finalize() {
     auto fragment = _fragment_context.lock();
     if (!fragment) {
-        return Status::OK();
+        return;
     }
+    SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(fragment->get_query_ctx()->query_mem_tracker);
     std::unique_lock<std::mutex> lc(_dependency_lock);
     _finalized = true;
     _sink_shared_state.reset();
     _op_shared_states.clear();
     _le_state_map.clear();
+    _block.reset();
+    _operators.clear();
+    _sink.reset();
+    _pipeline.reset();
 }
 
 Status PipelineTask::close(Status exec_status, bool close_sink) {
@@ -543,7 +550,6 @@ Status PipelineTask::close(Status exec_status, bool close_sink) {
 }
 
 std::string PipelineTask::debug_string() {
-    std::unique_lock<std::mutex> lc(_dependency_lock);
     fmt::memory_buffer debug_string_buffer;
 
     fmt::format_to(debug_string_buffer, "QueryId: {}\n", print_id(_query_id));
@@ -560,6 +566,7 @@ std::string PipelineTask::debug_string() {
     auto* cur_blocked_dep = _blocked_dep;
     auto fragment = _fragment_context.lock();
     if (is_finalized() || !fragment) {
+        fmt::format_to(debug_string_buffer, " pipeline name = {}", _pipeline_name);
         return fmt::to_string(debug_string_buffer);
     }
     auto elapsed = fragment->elapsed_time() / NANOS_PER_SEC;

be/src/pipeline/pipeline_task.h

@@ -48,7 +48,8 @@ class Dependency;
 class PipelineTask : public std::enable_shared_from_this<PipelineTask> {
 public:
     PipelineTask(PipelinePtr& pipeline, uint32_t task_id, RuntimeState* state,
-                 std::shared_ptr<PipelineFragmentContext> fragment_context, RuntimeProfile* parent_profile,
+                 std::shared_ptr<PipelineFragmentContext> fragment_context,
+                 RuntimeProfile* parent_profile,
                  std::map<int, std::pair<std::shared_ptr<LocalExchangeSharedState>,
                                          std::shared_ptr<Dependency>>>
                          le_state_map,
@@ -224,6 +225,8 @@ class PipelineTask : public std::enable_shared_from_this<PipelineTask> {
 
     RuntimeState* runtime_state() const { return _state; }
 
+    RuntimeProfile* get_task_profile() const { return _task_profile.get(); }
+
     std::string task_name() const { return fmt::format("task{}({})", _index, _pipeline->_name); }
 
     void stop_if_finished() {
@@ -235,7 +238,7 @@ class PipelineTask : public std::enable_shared_from_this<PipelineTask> {
         }
     }
 
-    PipelineId pipeline_id() const { return _pipeline->id(); }
+    PipelineId pipeline_id() const { return _pip_id; }
 
     bool wake_up_early() const { return _wake_up_early; }
 
@@ -250,6 +253,7 @@ class PipelineTask : public std::enable_shared_from_this<PipelineTask> {
     Status _open();
 
     const TUniqueId _query_id;
+    const PipelineId _pip_id;
     const uint32_t _index;
     PipelinePtr _pipeline;
     bool _has_exceed_timeout = false;
@@ -258,7 +262,7 @@ class PipelineTask : public std::enable_shared_from_this<PipelineTask> {
     int _previous_schedule_id = -1;
     uint32_t _schedule_time = 0;
     std::unique_ptr<doris::vectorized::Block> _block;
-    std::weak_ptr<PipelineFragmentContext> _fragment_context = nullptr;
+    std::weak_ptr<PipelineFragmentContext> _fragment_context;
     TaskQueue* _task_queue = nullptr;
 
     // used for priority queue
@@ -321,6 +325,7 @@ class PipelineTask : public std::enable_shared_from_this<PipelineTask> {
     std::atomic<bool> _running = false;
     std::atomic<bool> _eos = false;
     std::atomic<bool> _wake_up_early = false;
+    const std::string _pipeline_name;
 };
 
 using PipelineTaskSPtr = std::shared_ptr<PipelineTask>;

be/src/pipeline/task_queue.h

@@ -40,13 +40,13 @@ 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;
+    virtual Status push_back(PipelineTaskSPtr task) = 0;
 
     // push from worker
-    virtual Status push_back(PipelineTask* task, int core_id) = 0;
+    virtual Status push_back(PipelineTaskSPtr task, int core_id) = 0;
 
     virtual void update_statistics(PipelineTask* task, int64_t time_spent) {}
 
@@ -133,7 +133,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(PipelineTaskSPtr task) override;

be/src/pipeline/task_scheduler.cpp

@@ -92,7 +92,7 @@ void _close_task(PipelineTask* task, Status exec_status, PipelineFragmentContext
 
 void TaskScheduler::_do_work(size_t index) {
     while (_markers[index]) {
-        auto task = _task_queue.take(index);
+        auto task = _task_queue->take(index);
         if (!task) {
             continue;
         }
@@ -122,7 +122,7 @@ void TaskScheduler::_do_work(size_t index) {
             // If pipeline is canceled, it will report after pipeline closed, and will propagate
             // errors to downstream through exchange. So, here we needn't send_report.
             // fragment_ctx->send_report(true);
-            _close_task(task, fragment_context->get_query_ctx()->exec_status(),
+            _close_task(task.get(), fragment_context->get_query_ctx()->exec_status(),
                         fragment_context.get());
             continue;
         }
@@ -166,7 +166,7 @@ void TaskScheduler::_do_work(size_t index) {
             LOG(WARNING) << fmt::format("Pipeline task failed. query_id: {} reason: {}",
                                         print_id(fragment_context->get_query_ctx()->query_id()),
                                         status.to_string());
-            _close_task(task, status, fragment_context.get());
+            _close_task(task.get(), status, fragment_context.get());
             continue;
         }
         fragment_context->trigger_report_if_necessary();
@@ -179,7 +179,7 @@ void TaskScheduler::_do_work(size_t index) {
                 task->set_running(false);
             } else {
                 Status exec_status = fragment_context->get_query_ctx()->exec_status();
-                _close_task(task, exec_status, fragment_context.get());
+                _close_task(task.get(), exec_status, fragment_context.get());
             }
             continue;
         }