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,
     };
 }