YT-19555: Refactoring using Allocation tags for MemoryTag

Fix style issues in unit-tests

Refactoring TraceContext

Minor style fixes

8 files changed, 45 insertions, 50 deletions

diff --git a/yt/yt/core/actions/bind-inl.h b/yt/yt/core/actions/bind-inl.h
index 48a57a4ebdc..0633a5712f1 100644
--- a/yt/yt/core/actions/bind-inl.h
+++ b/yt/yt/core/actions/bind-inl.h
@@ -508,7 +508,7 @@ public:
         auto* volatile unoptimizedState = state;
         Y_UNUSED(unoptimizedState);
 
-        TMemoryTag memoryTag = GetCurrentMemoryTag();
+        auto memoryTag = GetCurrentMemoryTag();
 
         auto propagatingStorageGuard = state->MakePropagatingStorageGuard();
         Y_UNUSED(propagatingStorageGuard);
diff --git a/yt/yt/core/misc/unittests/memory_tag_ut.cpp b/yt/yt/core/misc/unittests/memory_tag_ut.cpp
index 1c80aea4e07..06f2e240f1c 100644
--- a/yt/yt/core/misc/unittests/memory_tag_ut.cpp
+++ b/yt/yt/core/misc/unittests/memory_tag_ut.cpp
@@ -245,7 +245,7 @@ INSTANTIATE_TEST_SUITE_P(MemoryTagTest, TMemoryTagTest, Values(
 
 using namespace NTracing;
 
-TEST(MemoryTagTest, MemoryTagPropagationViaAllocationTags)
+TEST(TMemoryTagTest, MemoryTagPropagationViaAllocationTags)
 {
     auto localContext = CreateTraceContextFromCurrent("MemoryTagPropagation");
     auto localTag = 1u;
@@ -291,7 +291,7 @@ void TestYield(TMemoryTag tag)
     EXPECT_EQ(currentTag, tag);
 }
 
-TEST(MemoryTagTest, MemoryTagWithYieldContextPropagation)
+TEST(TMemoryTagTest, MemoryTagWithYieldContextPropagation)
 {
     auto localContext = CreateTraceContextFromCurrent("MemoryTagSwitchContextPropagation");
     auto localTag = 1u;
diff --git a/yt/yt/core/tracing/allocation_hooks.cpp b/yt/yt/core/tracing/allocation_hooks.cpp
index 0a10e786165..8baa6f49e2c 100644
--- a/yt/yt/core/tracing/allocation_hooks.cpp
+++ b/yt/yt/core/tracing/allocation_hooks.cpp
@@ -26,25 +26,25 @@ void* CreateAllocationTagsData()
     return static_cast<void*>(allocationTags.Release());
 }
 
-void* CopyAllocationTagsData(void* ptr)
+void* CopyAllocationTagsData(void* userData)
 {
-    if (ptr) {
-        auto* allocationTagsPtr = static_cast<TAllocationTags*>(ptr);
+    if (userData) {
+        auto* allocationTagsPtr = static_cast<TAllocationTags*>(userData);
         allocationTagsPtr->Ref();
     }
-    return ptr;
+    return userData;
 }
 
-void DestroyAllocationTagsData(void* ptr)
+void DestroyAllocationTagsData(void* userData)
 {
-    auto* allocationTagsPtr = static_cast<TAllocationTags*>(ptr);
+    auto* allocationTagsPtr = static_cast<TAllocationTags*>(userData);
     // NB. No need to check for nullptr here, because ScheduleFree already does that.
     FreeList->ScheduleFree(allocationTagsPtr);
 }
 
-const TAllocationTags::TTags& ReadAllocationTagsData(void* ptr)
+const TAllocationTags::TTags& ReadAllocationTagsData(void* userData)
 {
-    auto* allocationTagsPtr = static_cast<TAllocationTags*>(ptr);
+    auto* allocationTagsPtr = static_cast<TAllocationTags*>(userData);
     if (!allocationTagsPtr) {
         static TAllocationTags::TTags emptyTags;
         return emptyTags;
diff --git a/yt/yt/core/tracing/allocation_tags.cpp b/yt/yt/core/tracing/allocation_tags.cpp
index 66177fdedbf..791dc7de3fb 100644
--- a/yt/yt/core/tracing/allocation_tags.cpp
+++ b/yt/yt/core/tracing/allocation_tags.cpp
@@ -24,15 +24,11 @@ std::optional<TAllocationTags::TValue> TAllocationTags::FindTagValue(
 {
     std::optional<TAllocationTags::TValue> value;
 
-    auto memoryTagIterator = std::find_if(
-        tags.cbegin(),
-        tags.cend(),
-        [&] (const auto& pair) {
-            return pair.first == key;
-        });
-
-    if (memoryTagIterator != tags.cend()) {
-        value = memoryTagIterator->second;
+    for (const auto& [key_, value_] : tags) {
+        if (key_ == key) {
+            value = value_;
+            break;
+        }
     }
 
     return value;
diff --git a/yt/yt/core/tracing/trace_context-inl.h b/yt/yt/core/tracing/trace_context-inl.h
index 2466564904d..dcb2bb78d75 100644
--- a/yt/yt/core/tracing/trace_context-inl.h
+++ b/yt/yt/core/tracing/trace_context-inl.h
@@ -84,13 +84,13 @@ void TTraceContext::AddTag(const TString& tagName, const T& tagValue)
 template <typename TTag>
 std::optional<TTag> TTraceContext::DoFindAllocationTag(const TString& key) const
 {
-    VERIFY_SPINLOCK_AFFINITY(AllocationTagsRWLock_);
+    VERIFY_SPINLOCK_AFFINITY(AllocationTagsLock_);
 
-    TAllocationTagsPtr tags = nullptr;
+    TAllocationTagsPtr tags;
 
     {
         // Local guard for copy RefCounted AllocationTags_.
-        auto guard = Guard(AllocationTagsAsRefCountedSpinlock_);
+        auto guard = Guard(AllocationTagsAsRefCountedLock_);
         tags = AllocationTags_;
     }
 
@@ -108,14 +108,14 @@ std::optional<TTag> TTraceContext::DoFindAllocationTag(const TString& key) const
 template <typename TTag>
 std::optional<TTag> TTraceContext::FindAllocationTag(const TString& key) const
 {
-    auto readerGuard = ReaderGuard(AllocationTagsRWLock_);
+    auto readerGuard = ReaderGuard(AllocationTagsLock_);
     return DoFindAllocationTag<TTag>(key);
 }
 
 template <typename TTag>
 std::optional<TTag> TTraceContext::RemoveAllocationTag(const TString& key)
 {
-    auto writerGuard = NThreading::WriterGuard(AllocationTagsRWLock_);
+    auto writerGuard = NThreading::WriterGuard(AllocationTagsLock_);
     auto newTags = DoGetAllocationTags();
 
     auto foundTagIt = std::remove_if(
@@ -143,7 +143,7 @@ std::optional<TTag> TTraceContext::SetAllocationTag(const TString& key, TTag new
 {
     auto newTagString = ToString(newTag);
 
-    auto writerGuard = NThreading::WriterGuard(AllocationTagsRWLock_);
+    auto writerGuard = NThreading::WriterGuard(AllocationTagsLock_);
     auto newTags = DoGetAllocationTags();
 
     if (!newTags.empty()) {
diff --git a/yt/yt/core/tracing/trace_context.cpp b/yt/yt/core/tracing/trace_context.cpp
index d813e84db4b..bb88ff60f80 100644
--- a/yt/yt/core/tracing/trace_context.cpp
+++ b/yt/yt/core/tracing/trace_context.cpp
@@ -266,62 +266,61 @@ void TTraceContext::SetLoggingTag(const TString& loggingTag)
 
 void TTraceContext::ClearAllocationTagsPtr() noexcept
 {
-    auto writerGuard = WriterGuard(AllocationTagsRWLock_);
-    auto guard = Guard(AllocationTagsAsRefCountedSpinlock_);
+    auto writerGuard = WriterGuard(AllocationTagsLock_);
+    auto guard = Guard(AllocationTagsAsRefCountedLock_);
     AllocationTags_ = nullptr;
 }
 
 TAllocationTags::TTags TTraceContext::DoGetAllocationTags() const
 {
-    VERIFY_SPINLOCK_AFFINITY(AllocationTagsRWLock_);
+    VERIFY_SPINLOCK_AFFINITY(AllocationTagsLock_);
 
-    TAllocationTagsPtr tags = nullptr;
+    TAllocationTagsPtr tags;
 
     {
         // Local guard for copy RefCounted AllocationTags_.
-        auto guard = Guard(AllocationTagsAsRefCountedSpinlock_);
+        auto guard = Guard(AllocationTagsAsRefCountedLock_);
         tags = AllocationTags_;
     }
 
-    if (tags != nullptr) {
-        return tags->GetTags();
+    if (!tags) {
+        return {};
     }
 
-    return {};
+    return tags->GetTags();
 }
 
 TAllocationTags::TTags TTraceContext::GetAllocationTags() const
 {
-    auto readerGuard = ReaderGuard(AllocationTagsRWLock_);
+    auto readerGuard = ReaderGuard(AllocationTagsLock_);
     return DoGetAllocationTags();
 }
 
 TAllocationTagsPtr TTraceContext::GetAllocationTagsPtr() const noexcept
 {
     // Local guard for copy RefCounted AllocationTags_ for allocator callback CreateAllocationTagsData().
-    auto guard = Guard(AllocationTagsAsRefCountedSpinlock_);
+    auto guard = Guard(AllocationTagsAsRefCountedLock_);
 
-    auto copy = AllocationTags_;
-    return copy;
+    return AllocationTags_;
 }
 
 void TTraceContext::DoSetAllocationTags(TAllocationTags::TTags&& tags)
 {
-    VERIFY_SPINLOCK_AFFINITY(AllocationTagsRWLock_);
+    VERIFY_SPINLOCK_AFFINITY(AllocationTagsLock_);
 
-    TAllocationTagsPtr allocationTagsPtr = nullptr;
+    TAllocationTagsPtr allocationTagsPtr;
     if (!tags.empty()) {
         // Allocation MUST be done BEFORE Guard(AllocationTagsAsRefCountedSpinlock_) to avoid deadlock with CreateAllocationTagsData().
         allocationTagsPtr = New<TAllocationTags>(std::move(tags));
     }
 
-    auto guard = Guard(AllocationTagsAsRefCountedSpinlock_);
-    AllocationTags_ = allocationTagsPtr;
+    auto guard = Guard(AllocationTagsAsRefCountedLock_);
+    AllocationTags_ = std::move(allocationTagsPtr);
 }
 
 void TTraceContext::SetAllocationTags(TAllocationTags::TTags&& tags)
 {
-    auto writerGuard = WriterGuard(AllocationTagsRWLock_);
+    auto writerGuard = WriterGuard(AllocationTagsLock_);
 
     return DoSetAllocationTags(std::move(tags));
 }
diff --git a/yt/yt/core/tracing/trace_context.h b/yt/yt/core/tracing/trace_context.h
index 7276a802f5c..0a6d6ae6cba 100644
--- a/yt/yt/core/tracing/trace_context.h
+++ b/yt/yt/core/tracing/trace_context.h
@@ -1,8 +1,6 @@
 #pragma once
 
 #include "allocation_tags.h"
-#include "library/cpp/yt/threading/rw_spin_lock.h"
-#include "public.h"
 
 #include <yt/yt/library/tracing/public.h>
 
@@ -14,6 +12,7 @@
 
 #include <yt/yt/core/concurrency/public.h>
 
+#include <library/cpp/yt/threading/rw_spin_lock.h>
 #include <library/cpp/yt/threading/spin_lock.h>
 
 #include <atomic>
@@ -243,9 +242,9 @@ private:
 
     std::vector<std::pair<TString, std::variant<TString, i64>>> ProfilingTags_;
 
-    // Must NOT allocate memory on the heap in callbacks with usage of AllocationTags_ to avoid deadlock with allocator.
-    YT_DECLARE_SPIN_LOCK(NThreading::TReaderWriterSpinLock, AllocationTagsRWLock_);
-    YT_DECLARE_SPIN_LOCK(NThreading::TSpinLock, AllocationTagsAsRefCountedSpinlock_);
+    // Must NOT allocate memory on the heap in callbacks with modifying AllocationTags_ to avoid deadlock with allocator.
+    YT_DECLARE_SPIN_LOCK(NThreading::TReaderWriterSpinLock, AllocationTagsLock_);
+    YT_DECLARE_SPIN_LOCK(NThreading::TSpinLock, AllocationTagsAsRefCountedLock_);
     TAllocationTagsPtr AllocationTags_;
 
     TTraceContext(
@@ -265,7 +264,6 @@ private:
 
     template <typename TTag>
     std::optional<TTag> DoFindAllocationTag(const TString& key) const;
-
 };
 
 DEFINE_REFCOUNTED_TYPE(TTraceContext)
diff --git a/yt/yt/core/tracing/unittests/allocation_tags_ut.cpp b/yt/yt/core/tracing/unittests/allocation_tags_ut.cpp
index 8f7e79d15b5..c53533dfe0e 100644
--- a/yt/yt/core/tracing/unittests/allocation_tags_ut.cpp
+++ b/yt/yt/core/tracing/unittests/allocation_tags_ut.cpp
@@ -4,10 +4,11 @@
 #include <yt/yt/core/tracing/trace_context.h>
 
 namespace NYT::NTracing {
+namespace {
 
 ////////////////////////////////////////////////////////////////////////////////
 
-TEST(TestAllocationTags, GetSetAllocationTags)
+TEST(TAllocationTagsTest, GetSetAllocationTags)
 {
     auto traceContext = TTraceContext::NewRoot("Root");
     TTraceContextGuard guard(traceContext);
@@ -41,4 +42,5 @@ TEST(TestAllocationTags, GetSetAllocationTags)
 
 ////////////////////////////////////////////////////////////////////////////////
 
+} // namespace
 } // namespace NYT::NTracing