[enhancement](mem-tracker) Use thread local mem tracker to track s3 f…

…ile buffer memory usage
apache · Sep 10, 2024 · 0e4d72e · 0e4d72e
1 parent 5392cb1
commit 0e4d72e
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/be/src/io/fs/s3_file_bufferpool.cpp b/be/src/io/fs/s3_file_bufferpool.cpp
@@ -84,10 +84,11 @@ FileBuffer::FileBuffer(BufferType type, std::function<FileBlocksHolderPtr()> all
           _size(0),
           _state(std::move(state)),
           _inner_data(std::make_unique<FileBuffer::PartData>()),
-          _capacity(_inner_data->size()) {}
+          _capacity(_inner_data->size()),
+          _mem_tracker(doris::thread_context()->thread_mem_tracker_mgr->limiter_mem_tracker()) {}
 
 FileBuffer::~FileBuffer() {
-    SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(ExecEnv::GetInstance()->s3_file_buffer_tracker());
+    SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(_mem_tracker);
     _inner_data.reset();
 }
 
@@ -240,7 +241,6 @@ FileBufferBuilder& FileBufferBuilder::set_allocate_file_blocks_holder(
 }
 
 Status FileBufferBuilder::build(std::shared_ptr<FileBuffer>* buf) {
-    SCOPED_SWITCH_THREAD_MEM_TRACKER_LIMITER(ExecEnv::GetInstance()->s3_file_buffer_tracker());
     OperationState state(_sync_after_complete_task, _is_cancelled);
 
     if (_type == BufferType::UPLOAD) {

diff --git a/be/src/io/fs/s3_file_bufferpool.h b/be/src/io/fs/s3_file_bufferpool.h
@@ -27,6 +27,7 @@
 
 #include "common/status.h"
 #include "io/cache/file_block.h"
+#include "runtime/memory/mem_tracker_limiter.h"
 #include "util/crc32c.h"
 #include "util/slice.h"
 #include "util/threadpool.h"
@@ -127,6 +128,7 @@ struct FileBuffer {
     struct PartData;
     std::unique_ptr<PartData> _inner_data;
     size_t _capacity;
+    std::shared_ptr<MemTrackerLimiter> _mem_tracker;
 };
 
 struct DownloadFileBuffer final : public FileBuffer {