[feature-wip] (memory tracker) (step5) Fix track bthread, fix track vectorized query

be/src/common/utils.h

@@ -21,7 +21,9 @@
 
 namespace doris {
 
+#ifndef ARRAY_SIZE
 #define ARRAY_SIZE(a) (sizeof(a)/sizeof((a)[0]))
+#endif
 
 struct AuthInfo {
  std::string user;

be/src/exec/exchange_node.cpp

@@ -80,6 +80,7 @@ Status ExchangeNode::prepare(RuntimeState* state) {
 Status ExchangeNode::open(RuntimeState* state) {
  SCOPED_TIMER(_runtime_profile->total_time_counter());
  SCOPED_SWITCH_TASK_THREAD_LOCAL_MEM_TRACKER(mem_tracker());
+ ADD_THREAD_LOCAL_MEM_TRACKER(_stream_recvr->mem_tracker());
  RETURN_IF_ERROR(ExecNode::open(state));
  if (_is_merging) {
  RETURN_IF_ERROR(_sort_exec_exprs.open(state));
@@ -215,7 +216,6 @@ Status ExchangeNode::get_next_merging(RuntimeState* state, RowBatch* output_batc
  RETURN_IF_CANCELLED(state);
  RETURN_IF_ERROR(state->check_query_state("Exchange, while merging next."));
 
- ADD_THREAD_LOCAL_MEM_TRACKER(_stream_recvr->mem_tracker());
  RETURN_IF_ERROR(_stream_recvr->get_next(output_batch, eos));
  while ((_num_rows_skipped < _offset)) {
  _num_rows_skipped += output_batch->num_rows();

be/src/exec/olap_scanner.cpp

@@ -52,7 +52,7 @@ OlapScanner::OlapScanner(RuntimeState* runtime_state, OlapScanNode* parent, bool
  _version(-1),
  _mem_tracker(MemTracker::create_tracker(
  tracker->limit(),
- tracker->label() + ":OlapScanner:" + thread_local_ctx.get()->thread_id_str(),
+ tracker->label() + ":OlapScanner:" + tls_ctx()->thread_id_str(),
  tracker)) {}
 
 Status OlapScanner::prepare(

be/src/exec/tablet_sink.cpp

@@ -50,7 +50,7 @@ NodeChannel::NodeChannel(OlapTableSink* parent, IndexChannel* index_channel, int
  _tuple_data_buffer_ptr = &_tuple_data_buffer;
  }
  _node_channel_tracker =
- MemTracker::create_tracker(-1, "NodeChannel" + thread_local_ctx.get()->thread_id_str());
+ MemTracker::create_tracker(-1, "NodeChannel" + tls_ctx()->thread_id_str());
 }
 
 NodeChannel::~NodeChannel() noexcept {
@@ -654,6 +654,7 @@ void IndexChannel::add_row(BlockRow& block_row, int64_t tablet_id) {
 
 void IndexChannel::mark_as_failed(int64_t node_id, const std::string& host, const std::string& err,
  int64_t tablet_id) {
+ SCOPED_SWITCH_THREAD_LOCAL_MEM_TRACKER(_index_channel_tracker);
  const auto& it = _tablets_by_channel.find(node_id);
  if (it == _tablets_by_channel.end()) {
  return;

be/src/exec/tablet_sink.h

@@ -33,6 +33,7 @@
 #include "exec/tablet_info.h"
 #include "gen_cpp/Types_types.h"
 #include "gen_cpp/internal_service.pb.h"
+#include "runtime/thread_context.h"
 #include "util/bitmap.h"
 #include "util/countdown_latch.h"
 #include "util/ref_count_closure.h"
@@ -325,6 +326,7 @@ class IndexChannel {
 
  void for_each_node_channel(
  const std::function<void(const std::shared_ptr<NodeChannel>&)>& func) {
+ SCOPED_SWITCH_THREAD_LOCAL_MEM_TRACKER(_index_channel_tracker);
  for (auto& it : _node_channels) {
  func(it.second);
  }
@@ -365,7 +367,7 @@ class IndexChannel {
  std::unordered_map<int64_t, std::string> _failed_channels_msgs;
  Status _intolerable_failure_status = Status::OK();
 
- std::shared_ptr<MemTracker> _index_channel_tracker; // TODO(zxy) use after
+ std::shared_ptr<MemTracker> _index_channel_tracker;
 };
 
 // Write data to Olap Table.

be/src/olap/byte_buffer.cpp

@@ -20,6 +20,7 @@
 #include <sys/mman.h>
 
 #include "olap/utils.h"
+#include "runtime/thread_context.h"
 
 namespace doris {
 
@@ -42,6 +43,8 @@ void StorageByteBuffer::BufDeleter::operator()(char* p) {
  if (0 != munmap(p, _mmap_length)) {
  LOG(FATAL) << "fail to munmap: mem=" << p << ", len=" << _mmap_length
  << ", errno=" << Errno::no() << ", errno_str=" << Errno::str();
+ } else {
+ RELEASE_THREAD_LOCAL_MEM_TRACKER(_mmap_length);
  }
  } else {
  delete[] p;
@@ -93,10 +96,12 @@ StorageByteBuffer* StorageByteBuffer::reference_buffer(StorageByteBuffer* refere
 
 StorageByteBuffer* StorageByteBuffer::mmap(void* start, uint64_t length, int prot, int flags,
  int fd, uint64_t offset) {
+ CONSUME_THREAD_LOCAL_MEM_TRACKER(length);
  char* memory = (char*)::mmap(start, length, prot, flags, fd, offset);
 
  if (MAP_FAILED == memory) {
  OLAP_LOG_WARNING("fail to mmap. [errno='%d' errno_str='%s']", Errno::no(), Errno::str());
+ RELEASE_THREAD_LOCAL_MEM_TRACKER(length);
  return nullptr;
  }
 
@@ -108,6 +113,7 @@ StorageByteBuffer* StorageByteBuffer::mmap(void* start, uint64_t length, int pro
  if (nullptr == buf) {
  deleter(memory);
  OLAP_LOG_WARNING("fail to allocate StorageByteBuffer.");
+ RELEASE_THREAD_LOCAL_MEM_TRACKER(length);
  return nullptr;
  }
 
@@ -128,10 +134,12 @@ StorageByteBuffer* StorageByteBuffer::mmap(FileHandler* handler, uint64_t offset
 
  size_t length = handler->length();
  int fd = handler->fd();
+ CONSUME_THREAD_LOCAL_MEM_TRACKER(length);
  char* memory = (char*)::mmap(nullptr, length, prot, flags, fd, offset);
 
  if (MAP_FAILED == memory) {
  OLAP_LOG_WARNING("fail to mmap. [errno='%d' errno_str='%s']", Errno::no(), Errno::str());
+ RELEASE_THREAD_LOCAL_MEM_TRACKER(length);
  return nullptr;
  }
 
@@ -143,6 +151,7 @@ StorageByteBuffer* StorageByteBuffer::mmap(FileHandler* handler, uint64_t offset
  if (nullptr == buf) {
  deleter(memory);
  OLAP_LOG_WARNING("fail to allocate StorageByteBuffer.");
+ RELEASE_THREAD_LOCAL_MEM_TRACKER(length);
  return nullptr;
  }
 
@@ -173,7 +182,7 @@ Status StorageByteBuffer::put(uint64_t index, char src) {
 }
 
 Status StorageByteBuffer::put(const char* src, uint64_t src_size, uint64_t offset,
-uint64_t length) {
+ uint64_t length) {
  //没有足够的空间可以写
  if (length > remaining()) {
  return Status::OLAPInternalError(OLAP_ERR_BUFFER_OVERFLOW);

be/src/olap/lru_cache.cpp

@@ -475,8 +475,7 @@ Cache::Handle* ShardedLRUCache::insert(const CacheKey& key, void* value, size_t
  CachePriority priority) {
  // The memory of the parameter value should be recorded in the tls mem tracker,
  // transfer the memory ownership of the value to ShardedLRUCache::_mem_tracker.
- thread_local_ctx.get()->_thread_mem_tracker_mgr->mem_tracker()->transfer_to(_mem_tracker.get(),
- charge);
+ tls_ctx()->_thread_mem_tracker_mgr->mem_tracker()->transfer_to(_mem_tracker.get(), charge);
  SCOPED_SWITCH_THREAD_LOCAL_MEM_TRACKER(_mem_tracker);
  const uint32_t hash = _hash_slice(key);
  return _shards[_shard(hash)]->insert(key, hash, value, charge, deleter, priority);

be/src/olap/out_stream.cpp

@@ -25,7 +25,7 @@
 namespace doris {
 
 OutStreamFactory::OutStreamFactory(CompressKind compress_kind, uint32_t stream_buffer_size)
- : _compress_kind(compress_kind), _stream_buffer_size(stream_buffer_size) {
+ : _stream_buffer_size(stream_buffer_size) {
  switch (compress_kind) {
  case COMPRESS_NONE:
  _compressor = nullptr;

be/src/olap/out_stream.h

@@ -141,7 +141,6 @@ class OutStreamFactory {
 
 private:
  std::map<StreamName, OutStream*> _streams; // All created streams
- CompressKind _compress_kind;
  Compressor _compressor;
  uint32_t _stream_buffer_size;
 

be/src/olap/rowset/column_writer.cpp

@@ -488,8 +488,7 @@ void ByteColumnWriter::record_position() {
 
 IntegerColumnWriter::IntegerColumnWriter(uint32_t column_id, uint32_t unique_column_id,
  OutStreamFactory* stream_factory, bool is_singed)
- : _column_id(column_id),
- _unique_column_id(unique_column_id),
+ : _unique_column_id(unique_column_id),
  _stream_factory(stream_factory),
  _writer(nullptr),
  _is_signed(is_singed) {}

be/src/olap/rowset/column_writer.h

@@ -179,7 +179,6 @@ class IntegerColumnWriter {
  Status flush() { return _writer->flush(); }
 
 private:
- uint32_t _column_id;
  uint32_t _unique_column_id;
  OutStreamFactory* _stream_factory;
  RunLengthIntegerWriter* _writer;

be/src/olap/rowset/segment_v2/segment.cpp

@@ -51,10 +51,9 @@ Segment::Segment(const FilePathDesc& path_desc, uint32_t segment_id,
  const TabletSchema* tablet_schema)
  : _path_desc(path_desc), _segment_id(segment_id), _tablet_schema(tablet_schema) {
 #ifndef BE_TEST
- _mem_tracker = MemTracker::create_virtual_tracker(
- -1, "Segment", StorageEngine::instance()->tablet_mem_tracker());
+ _mem_tracker = StorageEngine::instance()->tablet_mem_tracker();
 #else
- _mem_tracker = MemTracker::create_virtual_tracker(-1, "Segment");
+ _mem_tracker = MemTracker::get_process_tracker();
 #endif
 }
 

be/src/runtime/bufferpool/system_allocator.cc

@@ -22,6 +22,7 @@
 
 #include "common/config.h"
 #include "gutil/strings/substitute.h"
+#include "runtime/thread_context.h"
 #include "util/bit_util.h"
 #include "util/error_util.h"
 
@@ -75,9 +76,11 @@ Status SystemAllocator::AllocateViaMMap(int64_t len, uint8_t** buffer_mem) {
  // Map an extra huge page so we can fix up the alignment if needed.
  map_len += HUGE_PAGE_SIZE;
  }
+ CONSUME_THREAD_LOCAL_MEM_TRACKER(map_len);
  uint8_t* mem = reinterpret_cast<uint8_t*>(
  mmap(nullptr, map_len, PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_PRIVATE, -1, 0));
  if (mem == MAP_FAILED) {
+ RELEASE_THREAD_LOCAL_MEM_TRACKER(map_len);
  return Status::BufferAllocFailed("mmap failed");
  }
 
@@ -89,10 +92,12 @@ Status SystemAllocator::AllocateViaMMap(int64_t len, uint8_t** buffer_mem) {
  if (misalignment != 0) {
  uintptr_t fixup = HUGE_PAGE_SIZE - misalignment;
  munmap(mem, fixup);
+ RELEASE_THREAD_LOCAL_MEM_TRACKER(fixup);
  mem += fixup;
  map_len -= fixup;
  }
  munmap(mem + len, map_len - len);
+ RELEASE_THREAD_LOCAL_MEM_TRACKER(map_len - len);
  DCHECK_EQ(reinterpret_cast<uintptr_t>(mem) % HUGE_PAGE_SIZE, 0) << mem;
  // Mark the buffer as a candidate for promotion to huge pages. The Linux Transparent
  // Huge Pages implementation will try to back the memory with a huge page if it is
@@ -142,6 +147,7 @@ Status SystemAllocator::AllocateViaMalloc(int64_t len, uint8_t** buffer_mem) {
 void SystemAllocator::Free(BufferPool::BufferHandle&& buffer) {
  if (config::mmap_buffers) {
  int rc = munmap(buffer.data(), buffer.len());
+ RELEASE_THREAD_LOCAL_MEM_TRACKER(buffer.len());
  DCHECK_EQ(rc, 0) << "Unexpected munmap() error: " << errno;
  } else {
  bool use_huge_pages = buffer.len() % HUGE_PAGE_SIZE == 0 && config::madvise_huge_pages;

be/src/runtime/disk_io_mgr.cc

@@ -220,7 +220,7 @@ void DiskIoMgr::BufferDescriptor::reset(RequestContext* reader, ScanRange* range
  _eosr = false;
  _status = Status::OK();
  // Consume in the tls mem tracker when the buffer is allocated.
- _buffer_mem_tracker = thread_local_ctx.get()->_thread_mem_tracker_mgr->mem_tracker().get();
+ _buffer_mem_tracker = tls_ctx()->_thread_mem_tracker_mgr->mem_tracker().get();
 }
 
 void DiskIoMgr::BufferDescriptor::return_buffer() {
@@ -739,7 +739,7 @@ char* DiskIoMgr::get_free_buffer(int64_t* buffer_size) {
  buffer = new char[*buffer_size];
  } else {
  // This means the buffer's memory ownership is transferred from DiskIoMgr to tls tracker.
- _mem_tracker->transfer_to(thread_local_ctx.get()->_thread_mem_tracker_mgr->mem_tracker().get(), *buffer_size);
+ _mem_tracker->transfer_to(tls_ctx()->_thread_mem_tracker_mgr->mem_tracker().get(), *buffer_size);
  buffer = _free_buffers[idx].front();
  _free_buffers[idx].pop_front();
  }
@@ -767,7 +767,7 @@ void DiskIoMgr::gc_io_buffers(int64_t bytes_to_free) {
  // The deleted buffer is released in the tls mem tracker, the deleted buffer belongs to DiskIoMgr,
  // so the freed memory should be recorded in the DiskIoMgr mem tracker. So if the tls mem tracker
  // and the DiskIoMgr tracker are different, transfer memory ownership.
- _mem_tracker->transfer_to(thread_local_ctx.get()->_thread_mem_tracker_mgr->mem_tracker().get(), bytes_freed);
+ _mem_tracker->transfer_to(tls_ctx()->_thread_mem_tracker_mgr->mem_tracker().get(), bytes_freed);
 }
 
 void DiskIoMgr::return_free_buffer(BufferDescriptor* desc) {
@@ -793,7 +793,7 @@ void DiskIoMgr::return_free_buffer(char* buffer, int64_t buffer_size, MemTracker
  // The deleted buffer is released in the tls mem tracker. When the buffer was allocated,
  // it was consumed in BufferDescriptor->buffer_mem_tracker, so if the tls mem tracker and
  // the tracker in the parameters are different, transfer memory ownership.
- tracker->transfer_to(thread_local_ctx.get()->_thread_mem_tracker_mgr->mem_tracker().get(), buffer_size);
+ tracker->transfer_to(tls_ctx()->_thread_mem_tracker_mgr->mem_tracker().get(), buffer_size);
  }
 }
 

be/src/runtime/fold_constant_executor.cpp

@@ -44,6 +44,7 @@ TUniqueId FoldConstantExecutor::_dummy_id;
 
 Status FoldConstantExecutor::fold_constant_expr(
  const TFoldConstantParams& params, PConstantExprResult* response) {
+ SCOPED_SWITCH_THREAD_LOCAL_MEM_TRACKER(_mem_tracker);
  const auto& expr_map = params.expr_map;
  auto expr_result_map = response->mutable_expr_result_map();
 
@@ -53,7 +54,6 @@ Status FoldConstantExecutor::fold_constant_expr(
  if (UNLIKELY(!status.ok())) {
  return status;
  }
- SCOPED_SWITCH_THREAD_LOCAL_MEM_TRACKER(_mem_tracker);
 
  for (const auto& m : expr_map) {
  PExprResultMap pexpr_result_map;

be/src/runtime/fragment_mgr.cpp

@@ -467,8 +467,7 @@ void FragmentMgr::_exec_actual(std::shared_ptr<FragmentExecState> exec_state, Fi
  .instance_id(exec_state->fragment_instance_id())
  .tag("pthread_id", std::to_string((uintptr_t)pthread_self()));
 #ifndef BE_TEST
- SCOPED_ATTACH_TASK_THREAD(exec_state->executor()->runtime_state()->query_type(),
- print_id(exec_state->query_id()), exec_state->fragment_instance_id(),
+ SCOPED_ATTACH_TASK_THREAD(exec_state->executor()->runtime_state(),
  exec_state->executor()->runtime_state()->instance_mem_tracker());
 #endif
  exec_state->execute();

be/src/runtime/mem_pool.cpp

@@ -65,7 +65,7 @@ MemPool::MemPool()
  total_allocated_bytes_(0),
  total_reserved_bytes_(0),
  peak_allocated_bytes_(0),
- _mem_tracker(thread_local_ctx.get()->_thread_mem_tracker_mgr->mem_tracker().get()) {}
+ _mem_tracker(tls_ctx()->_thread_mem_tracker_mgr->mem_tracker().get()) {}
 
 MemPool::ChunkInfo::ChunkInfo(const Chunk& chunk_) : chunk(chunk_), allocated_bytes(0) {
  DorisMetrics::instance()->memory_pool_bytes_total->increment(chunk.size);

be/src/runtime/mem_tracker.cpp

@@ -60,6 +60,19 @@ MemTracker* MemTracker::get_raw_process_tracker() {
  return raw_process_tracker;
 }
 
+// Track memory for all brpc server responses.
+static std::shared_ptr<MemTracker> brpc_server_tracker;
+static GoogleOnceType brpc_server_tracker_once = GOOGLE_ONCE_INIT;
+
+void MemTracker::create_brpc_server_tracker() {
+ brpc_server_tracker = MemTracker::create_tracker(-1, "Brpc", get_process_tracker(), MemTrackerLevel::OVERVIEW);
+}
+
+std::shared_ptr<MemTracker> MemTracker::get_brpc_server_tracker() {
+ GoogleOnceInit(&brpc_server_tracker_once, &MemTracker::create_brpc_server_tracker);
+ return brpc_server_tracker;
+}
+
 void MemTracker::list_process_trackers(std::vector<std::shared_ptr<MemTracker>>* trackers) {
  trackers->clear();
  std::deque<std::shared_ptr<MemTracker>> to_process;
@@ -88,7 +101,8 @@ std::shared_ptr<MemTracker> MemTracker::create_tracker(int64_t byte_limit, const
  const std::shared_ptr<MemTracker>& parent,
  MemTrackerLevel level,
  RuntimeProfile* profile) {
- std::shared_ptr<MemTracker> reset_parent = parent ? parent : thread_local_ctx.get()->_thread_mem_tracker_mgr->mem_tracker();
+ std::shared_ptr<MemTracker> reset_parent =
+ parent ? parent : tls_ctx()->_thread_mem_tracker_mgr->mem_tracker();
  DCHECK(reset_parent);
 
  std::shared_ptr<MemTracker> tracker(
@@ -102,7 +116,8 @@ std::shared_ptr<MemTracker> MemTracker::create_tracker(int64_t byte_limit, const
 std::shared_ptr<MemTracker> MemTracker::create_virtual_tracker(
  int64_t byte_limit, const std::string& label, const std::shared_ptr<MemTracker>& parent,
  MemTrackerLevel level) {
- std::shared_ptr<MemTracker> reset_parent = parent ? parent : thread_local_ctx.get()->_thread_mem_tracker_mgr->mem_tracker();
+ std::shared_ptr<MemTracker> reset_parent =
+ parent ? parent : tls_ctx()->_thread_mem_tracker_mgr->mem_tracker();
  DCHECK(reset_parent);
 
  std::shared_ptr<MemTracker> tracker(
@@ -121,6 +136,7 @@ MemTracker::MemTracker(int64_t byte_limit, const std::string& label,
  RuntimeProfile* profile)
  : _limit(byte_limit),
  _label(label),
+ // Not 100% sure the id is unique. This is generated because it is faster than converting to int after hash.
  _id((GetCurrentTimeMicros() % 1000000) * 100 + _label.length()),
  _parent(parent),
  _level(level) {

be/src/runtime/mem_tracker.h

@@ -97,6 +97,8 @@ class MemTracker {
  // Gets a shared_ptr to the "process" tracker, creating it if necessary.
  static std::shared_ptr<MemTracker> get_process_tracker();
  static MemTracker* get_raw_process_tracker();
+ // Gets a shared_ptr to the "brpc server" tracker, creating it if necessary.
+ static std::shared_ptr<MemTracker> get_brpc_server_tracker();
 
  Status check_sys_mem_info(int64_t bytes) {
  if (MemInfo::initialized() && MemInfo::current_mem() + bytes >= MemInfo::mem_limit()) {
@@ -464,6 +466,8 @@ class MemTracker {
 
  // Creates the process tracker.
  static void create_process_tracker();
+ // Creates the brpc server tracker.
+ static void create_brpc_server_tracker();
 
  // Limit on memory consumption, in bytes. If limit_ == -1, there is no consumption limit.
  int64_t _limit;