YT-20220: Fix error due allocation at __atexit(...)

Fix error due allocation at __atexit(...)

4 files changed, 25 insertions, 15 deletions

diff --git a/yt/yt/core/tracing/allocation_hooks.cpp b/yt/yt/core/tracing/allocation_hooks.cpp
index 8baa6f49e2..36be840cfa 100644
--- a/yt/yt/core/tracing/allocation_hooks.cpp
+++ b/yt/yt/core/tracing/allocation_hooks.cpp
@@ -42,14 +42,14 @@ void DestroyAllocationTagsData(void* userData)
     FreeList->ScheduleFree(allocationTagsPtr);
 }
 
-const TAllocationTags::TTags& ReadAllocationTagsData(void* userData)
+const TAllocationTags::TTags* ReadAllocationTagsData(void* userData)
 {
-    auto* allocationTagsPtr = static_cast<TAllocationTags*>(userData);
-    if (!allocationTagsPtr) {
-        static TAllocationTags::TTags emptyTags;
-        return emptyTags;
+    if (!userData) {
+        return nullptr;
     }
-    return allocationTagsPtr->GetTags();
+
+    const auto* allocationTagsPtr = static_cast<TAllocationTags*>(userData);
+    return allocationTagsPtr->GetTagsPtr();
 }
 
 std::optional<TString> FindTagValue(
diff --git a/yt/yt/core/tracing/allocation_tags.cpp b/yt/yt/core/tracing/allocation_tags.cpp
index 791dc7de3f..e38154db17 100644
--- a/yt/yt/core/tracing/allocation_tags.cpp
+++ b/yt/yt/core/tracing/allocation_tags.cpp
@@ -13,6 +13,11 @@ const TAllocationTags::TTags& TAllocationTags::GetTags() const noexcept
     return Tags_;
 }
 
+const TAllocationTags::TTags* TAllocationTags::GetTagsPtr() const noexcept
+{
+    return &Tags_;
+}
+
 std::optional<TAllocationTags::TValue> TAllocationTags::FindTagValue(const TKey& key) const
 {
     return FindTagValue(Tags_, key);
diff --git a/yt/yt/core/tracing/allocation_tags.h b/yt/yt/core/tracing/allocation_tags.h
index 549bdc30bc..ea7b5a6447 100644
--- a/yt/yt/core/tracing/allocation_tags.h
+++ b/yt/yt/core/tracing/allocation_tags.h
@@ -19,6 +19,8 @@ public:
 
     const TTags& GetTags() const noexcept;
 
+    const TTags* GetTagsPtr() const noexcept;
+
     std::optional<TValue> FindTagValue(const TKey& key) const;
 
     static std::optional<TValue> FindTagValue(
diff --git a/yt/yt/library/ytprof/heap_profiler.cpp b/yt/yt/library/ytprof/heap_profiler.cpp
index e04a8a1759..70f5d9ece6 100644
--- a/yt/yt/library/ytprof/heap_profiler.cpp
+++ b/yt/yt/library/ytprof/heap_profiler.cpp
@@ -34,10 +34,9 @@ Y_WEAK void* CopyAllocationTagsData(void* userData)
 Y_WEAK void DestroyAllocationTagsData(void* /*userData*/)
 { }
 
-Y_WEAK const std::vector<std::pair<TString, TString>>& ReadAllocationTagsData(void* /*userData*/)
+Y_WEAK const std::vector<std::pair<TString, TString>>* ReadAllocationTagsData(void* /*userData*/)
 {
-    static const std::vector<std::pair<TString, TString>> emptyTags;
-    return emptyTags;
+    return nullptr;
 }
 
 Y_WEAK std::optional<TString> FindTagValue(
@@ -132,10 +131,12 @@ NProto::Profile ConvertAllocationProfile(const tcmalloc::Profile& snapshot)
         }
 
         // TODO(gepardo): Deduplicate values in string table
-        for (const auto& [key, value] : ReadAllocationTagsData(sample.user_data)) {
-            auto label = sampleProto->add_label();
-            label->set_key(addString(key));
-            label->set_str(addString(value));
+        if (const auto* data = ReadAllocationTagsData(sample.user_data)) {
+            for (const auto& [key, value] : *data) {
+                auto label = sampleProto->add_label();
+                label->set_key(addString(key));
+                label->set_str(addString(value));
+            }
         }
     });
 
@@ -193,8 +194,10 @@ TMemoryUsageSnapshotPtr CollectMemoryUsageSnapshot()
 
     auto snapshot = tcmalloc::MallocExtension::SnapshotCurrent(tcmalloc::ProfileType::kHeap);
     snapshot.Iterate([&] (const tcmalloc::Profile::Sample& sample) {
-        for (const auto& [tagName, tag] : ReadAllocationTagsData(sample.user_data)) {
-            usage[tagName][tag] += sample.sum;
+        if (const auto* data = ReadAllocationTagsData(sample.user_data)) {
+            for (const auto& [tagName, tag] : *data) {
+                usage[tagName][tag] += sample.sum;
+            }
         }
     });