Use jemalloc to trace thread allocation

contrib/tiflash-proxy-cmake/CMakeLists.txt

@@ -14,10 +14,18 @@
 
 if (CMAKE_BUILD_TYPE_UC STREQUAL "DEBUG" OR SAN_DEBUG)
  set(_TIFLASH_PROXY_BUILD_PROFILE "debug")
- set(_TIFLASH_PROXY_MAKE_COMMAND make debug)
+ if (ENABLE_JEMALLOC)
+ set(_TIFLASH_PROXY_MAKE_COMMAND ENABLE_FEATURES="external-jemalloc" make debug)
+ else()
+ set(_TIFLASH_PROXY_MAKE_COMMAND make debug)
+ endif()
 else()
  set(_TIFLASH_PROXY_BUILD_PROFILE "release")
- set(_TIFLASH_PROXY_MAKE_COMMAND make release)
+ if (ENABLE_JEMALLOC)
+ set(_TIFLASH_PROXY_MAKE_COMMAND ENABLE_FEATURES="external-jemalloc" make release)
+ else()
+ set(_TIFLASH_PROXY_MAKE_COMMAND make release)
+ endif()
 endif()
 
 set(_TIFLASH_PROXY_SOURCE_DIR "${TiFlash_SOURCE_DIR}/contrib/tiflash-proxy")

dbms/src/Common/MemoryTrace.cpp

@@ -0,0 +1,37 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed 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 <Common/MemoryTrace.h>
+
+#ifdef USE_JEMALLOC
+#include <jemalloc/jemalloc.h>
+#endif
+
+namespace DB
+{
+std::tuple<uint64_t *, uint64_t *> getAllocDeallocPtr()
+{
+#ifdef USE_JEMALLOC
+ uint64_t * ptr1 = nullptr;
+ uint64_t size1 = sizeof ptr1;
+ je_mallctl("thread.allocatedp", (void *)&ptr1, &size1, nullptr, 0);
+ uint64_t * ptr2 = nullptr;
+ uint64_t size2 = sizeof ptr2;
+ je_mallctl("thread.deallocatedp", (void *)&ptr2, &size2, nullptr, 0);
+ return std::make_tuple(ptr1, ptr2);
+#else
+ return std::make_tuple(nullptr, nullptr);
+#endif
+}
+} // namespace DB

dbms/src/Common/MemoryTrace.h

@@ -0,0 +1,22 @@
+// Copyright 2024 PingCAP, Inc.
+//
+// Licensed 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.
+
+#pragma once
+
+#include <tuple>
+
+namespace DB
+{
+std::tuple<uint64_t *, uint64_t *> getAllocDeallocPtr();
+} // namespace DB

dbms/src/Common/TiFlashMetrics.cpp

@@ -58,6 +58,8 @@ TiFlashMetrics::TiFlashMetrics()
  .Name("tiflash_storage_sync_replica_ru")
  .Help("RU for synchronous replica of keyspace")
  .Register(*registry);
+ registered_raft_proxy_thread_memory_usage_family
+ = &prometheus::BuildGauge().Name(raft_proxy_thread_memory_usage).Help("").Register(*registry);
 }
 
 void TiFlashMetrics::addReplicaSyncRU(UInt32 keyspace_id, UInt64 ru)
@@ -96,4 +98,35 @@ void TiFlashMetrics::removeReplicaSyncRUCounter(UInt32 keyspace_id)
  registered_keyspace_sync_replica_ru_family->Remove(itr->second);
  registered_keyspace_sync_replica_ru.erase(itr);
 }
-} // namespace DB
+
+double TiFlashMetrics::getProxyThreadMemory(const std::string & k)
+{
+ std::shared_lock lock(proxy_thread_report_mtx);
+ auto it = registered_raft_proxy_thread_memory_usage_metrics.find(k);
+ RUNTIME_CHECK(it != registered_raft_proxy_thread_memory_usage_metrics.end(), k);
+ return it->second->Value();
+}
+
+void TiFlashMetrics::setProxyThreadMemory(const std::string & k, Int64 v)
+{
+ std::shared_lock lock(proxy_thread_report_mtx);
+ if unlikely (!registered_raft_proxy_thread_memory_usage_metrics.count(k))
+ {
+ // New metrics added through `Reset`.
+ return;
+ }
+ registered_raft_proxy_thread_memory_usage_metrics[k]->Set(v);
+}
+
+void TiFlashMetrics::registerProxyThreadMemory(const std::string & k)
+{
+ std::unique_lock lock(proxy_thread_report_mtx);
+ if unlikely (!registered_raft_proxy_thread_memory_usage_metrics.count(k))
+ {
+ registered_raft_proxy_thread_memory_usage_metrics.emplace(
+ k,
+ &registered_raft_proxy_thread_memory_usage_family->Add({{"type", k}}));
+ }
+}
+
+} // namespace DB

dbms/src/Common/TiFlashMetrics.h

@@ -1088,6 +1088,9 @@ class TiFlashMetrics
 
  void addReplicaSyncRU(UInt32 keyspace_id, UInt64 ru);
  UInt64 debugQueryReplicaSyncRU(UInt32 keyspace_id);
+ void setProxyThreadMemory(const std::string & k, Int64 v);
+ double getProxyThreadMemory(const std::string & k);
+ void registerProxyThreadMemory(const std::string & k);
 
 private:
  TiFlashMetrics();
@@ -1098,6 +1101,7 @@ class TiFlashMetrics
  static constexpr auto profile_events_prefix = "tiflash_system_profile_event_";
  static constexpr auto current_metrics_prefix = "tiflash_system_current_metric_";
  static constexpr auto async_metrics_prefix = "tiflash_system_asynchronous_metric_";
+ static constexpr auto raft_proxy_thread_memory_usage = "tiflash_raft_proxy_thread_memory_usage";
 
  std::shared_ptr<prometheus::Registry> registry = std::make_shared<prometheus::Registry>();
  // Here we add a ProcessCollector to collect cpu/rss/vsize/start_time information.
@@ -1118,6 +1122,11 @@ class TiFlashMetrics
  std::mutex replica_sync_ru_mtx;
  std::unordered_map<KeyspaceID, prometheus::Counter *> registered_keyspace_sync_replica_ru;
 
+ // TODO: Use CAS+HazPtr to remove proxy_thread_report_mtx, or hash some slots here.
+ prometheus::Family<prometheus::Gauge> * registered_raft_proxy_thread_memory_usage_family;
+ std::shared_mutex proxy_thread_report_mtx;
+ std::unordered_map<std::string, prometheus::Gauge *> registered_raft_proxy_thread_memory_usage_metrics;
+
 public:
 #define MAKE_METRIC_MEMBER_M(family_name, help, type, ...) \
  MetricFamily<prometheus::type> family_name \

dbms/src/Server/Server.cpp

@@ -1029,6 +1029,12 @@ int Server::main(const std::vector<std::string> & /*args*/)
  LOG_INFO(log, "UniPS is not enabled for proxy, page_version={}", STORAGE_FORMAT_CURRENT.page);
  }
 
+#ifdef USE_JEMALLOC
+ LOG_INFO(log, "Using Jemalloc for TiFlash");
+#else
+ LOG_INFO(log, "Not using Jemalloc for TiFlash");
+#endif
+
  RaftStoreProxyRunner proxy_runner(RaftStoreProxyRunner::RunRaftStoreProxyParms{&helper, proxy_conf}, log);
 
  if (proxy_conf.is_proxy_runnable)
@@ -1708,6 +1714,8 @@ int Server::main(const std::vector<std::string> & /*args*/)
  LOG_ERROR(log, "Current status of engine-store is NOT Running, should not happen");
  exit(-1);
  }
+ LOG_INFO(log, "Stop collecting thread alloc metrics");
+ tmt_context.getKVStore()->stopThreadAllocInfo();
  LOG_INFO(log, "Set store context status Stopping");
  tmt_context.setStatusStopping();
  {

dbms/src/Storages/KVStore/FFI/ProxyFFI.cpp

@@ -1030,4 +1030,44 @@ BaseBuffView GetLockByKey(const EngineStoreServerWrap * server, uint64_t region_
  }
 }
 
+void ReportThreadAllocateInfo(
+ EngineStoreServerWrap * server,
+ uint64_t tid,
+ BaseBuffView name,
+ ReportThreadAllocateInfoType type,
+ uint64_t value)
+{
+ try
+ {
+ UNUSED(tid);
+ if (!server || !server->tmt || !server->tmt->getKVStore())
+ return;
+ server->tmt->getKVStore()->reportThreadAllocInfo(buffToStrView(name), type, value);
+ }
+ catch (...)
+ {
+ tryLogCurrentFatalException(__PRETTY_FUNCTION__);
+ exit(-1);
+ }
+}
+
+void ReportThreadAllocateBatch(
+ EngineStoreServerWrap * server,
+ uint64_t tid,
+ BaseBuffView name,
+ ReportThreadAllocateInfoBatch data)
+{
+ try
+ {
+ UNUSED(server);
+ UNUSED(tid);
+ KVStore::reportThreadAllocBatch(buffToStrView(name), data);
+ }
+ catch (...)
+ {
+ tryLogCurrentFatalException(__PRETTY_FUNCTION__);
+ exit(-1);
+ }
+}
+
 } // namespace DB

dbms/src/Storages/KVStore/FFI/ProxyFFI.h

@@ -206,6 +206,17 @@ FapSnapshotState QueryFapSnapshotState(
  uint64_t term);
 void ClearFapSnapshot(EngineStoreServerWrap * server, uint64_t region_id);
 bool KvstoreRegionExists(EngineStoreServerWrap * server, uint64_t region_id);
+void ReportThreadAllocateInfo(
+ EngineStoreServerWrap *,
+ uint64_t tid,
+ BaseBuffView name,
+ ReportThreadAllocateInfoType type,
+ uint64_t value);
+void ReportThreadAllocateBatch(
+ EngineStoreServerWrap *,
+ uint64_t tid,
+ BaseBuffView name,
+ ReportThreadAllocateInfoBatch data);
 }
 
 inline EngineStoreServerHelper GetEngineStoreServerHelper(EngineStoreServerWrap * tiflash_instance_wrap)
@@ -257,7 +268,8 @@ inline EngineStoreServerHelper GetEngineStoreServerHelper(EngineStoreServerWrap
  .fn_fast_add_peer = FastAddPeer,
  .fn_query_fap_snapshot_state = QueryFapSnapshotState,
  .fn_clear_fap_snapshot = ClearFapSnapshot,
- .fn_kvstore_region_exists = KvstoreRegionExists,
+ .fn_report_thread_allocate_info = ReportThreadAllocateInfo,
+ .fn_report_thread_allocate_batch = ReportThreadAllocateBatch,
  };
 }