From 54a6030d8a4d829b12b33c0335f38006dbd17114 Mon Sep 17 00:00:00 2001
From: Xinyi Zou <zouxinyi02@gmail.com>
Date: Fri, 29 Mar 2024 22:03:39 +0800
Subject: [PATCH]  [fix](memory) Fix CloudDeltaWriter init attach task in
 thread context (#33048)

---
 be/src/cloud/cloud_delta_writer.cpp | 3 ++-
 be/src/cloud/cloud_delta_writer.h   | 3 +++
 be/src/olap/delta_writer.h          | 1 -
 be/src/olap/storage_engine.cpp      | 1 -
 be/src/olap/storage_engine.h        | 1 -
 5 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/be/src/cloud/cloud_delta_writer.cpp b/be/src/cloud/cloud_delta_writer.cpp
index 72f82e9f4d5ea9..ef5f4dc33f7eb9 100644
--- a/be/src/cloud/cloud_delta_writer.cpp
+++ b/be/src/cloud/cloud_delta_writer.cpp
@@ -29,6 +29,7 @@ CloudDeltaWriter::CloudDeltaWriter(CloudStorageEngine& engine, const WriteReques
                                    RuntimeProfile* profile, const UniqueId& load_id)
         : BaseDeltaWriter(req, profile, load_id), _engine(engine) {
     _rowset_builder = std::make_unique<CloudRowsetBuilder>(engine, req, profile);
+    _query_thread_context.init();
 }
 
 CloudDeltaWriter::~CloudDeltaWriter() = default;
@@ -46,7 +47,7 @@ Status CloudDeltaWriter::batch_init(std::vector<CloudDeltaWriter*> writers) {
         }
 
         tasks.emplace_back([writer] {
-            ThreadLocalHandle::create_thread_local_if_not_exits();
+            SCOPED_ATTACH_TASK(writer->query_thread_context());
             std::lock_guard<bthread::Mutex> lock(writer->_mtx);
             if (writer->_is_init || writer->_is_cancelled) {
                 return Status::OK();
diff --git a/be/src/cloud/cloud_delta_writer.h b/be/src/cloud/cloud_delta_writer.h
index c4153a075db32e..83f2c494708286 100644
--- a/be/src/cloud/cloud_delta_writer.h
+++ b/be/src/cloud/cloud_delta_writer.h
@@ -53,12 +53,15 @@ class CloudDeltaWriter final : public BaseDeltaWriter {
 
     Status set_txn_related_delete_bitmap();
 
+    QueryThreadContext query_thread_context() { return _query_thread_context; }
+
 private:
     // Convert `_rowset_builder` from `BaseRowsetBuilder` to `CloudRowsetBuilder`
     CloudRowsetBuilder* rowset_builder();
 
     bthread::Mutex _mtx;
     CloudStorageEngine& _engine;
+    QueryThreadContext _query_thread_context;
 };
 
 } // namespace doris
diff --git a/be/src/olap/delta_writer.h b/be/src/olap/delta_writer.h
index b0b460dd34fe7d..c7b22c63a5f637 100644
--- a/be/src/olap/delta_writer.h
+++ b/be/src/olap/delta_writer.h
@@ -43,7 +43,6 @@ namespace doris {
 
 class FlushToken;
 class MemTable;
-class MemTracker;
 class StorageEngine;
 class TupleDescriptor;
 class SlotDescriptor;
diff --git a/be/src/olap/storage_engine.cpp b/be/src/olap/storage_engine.cpp
index 3f0347524e6390..d341cf718b0df7 100644
--- a/be/src/olap/storage_engine.cpp
+++ b/be/src/olap/storage_engine.cpp
@@ -75,7 +75,6 @@
 #include "olap/tablet_meta_manager.h"
 #include "olap/task/engine_task.h"
 #include "olap/txn_manager.h"
-#include "runtime/memory/mem_tracker.h"
 #include "runtime/stream_load/stream_load_recorder.h"
 #include "util/doris_metrics.h"
 #include "util/metrics.h"
diff --git a/be/src/olap/storage_engine.h b/be/src/olap/storage_engine.h
index be1d5935a54739..6888af9ee1747e 100644
--- a/be/src/olap/storage_engine.h
+++ b/be/src/olap/storage_engine.h
@@ -60,7 +60,6 @@ class BaseCompaction;
 class CumulativeCompaction;
 class SingleReplicaCompaction;
 class CumulativeCompactionPolicy;
-class MemTracker;
 class StreamLoadRecorder;
 class TCloneReq;
 class TCreateTabletReq;