Extract memory tag API to library

5 files changed, 19 insertions, 96 deletions

diff --git a/library/cpp/ytalloc/api/fallback.cpp b/library/cpp/ytalloc/api/fallback.cpp
index 5880ede439..d094cf2e93 100644
--- a/library/cpp/ytalloc/api/fallback.cpp
+++ b/library/cpp/ytalloc/api/fallback.cpp
@@ -56,24 +56,6 @@ Y_WEAK size_t GetAllocationSize(size_t size)
 
 ////////////////////////////////////////////////////////////////////////////////
 
-Y_WEAK TMemoryTag GetCurrentMemoryTag()
-{
-    return NullMemoryTag;
-}
-
-Y_WEAK void SetCurrentMemoryTag(TMemoryTag /*tag*/)
-{ }
-
-Y_WEAK size_t GetMemoryUsageForTag(TMemoryTag /*tag*/)
-{
-    return 0;
-}
-
-Y_WEAK void GetMemoryUsageForTags(const TMemoryTag* /*tags*/, size_t /*count*/, size_t* /*results*/)
-{ }
-
-////////////////////////////////////////////////////////////////////////////////
-
 Y_WEAK void SetCurrentMemoryZone(EMemoryZone /*zone*/)
 { }
 
diff --git a/library/cpp/ytalloc/api/ytalloc-inl.h b/library/cpp/ytalloc/api/ytalloc-inl.h
index 263108423d..997b230aae 100644
--- a/library/cpp/ytalloc/api/ytalloc-inl.h
+++ b/library/cpp/ytalloc/api/ytalloc-inl.h
@@ -73,33 +73,4 @@ void* AllocateConstSize()
 
 ////////////////////////////////////////////////////////////////////////////////
 
-inline TMemoryTagGuard::TMemoryTagGuard()
-    : Active_(false)
-    , PreviousTag_(NullMemoryTag)
-{ }
-
-inline TMemoryTagGuard::TMemoryTagGuard(TMemoryTag tag)
-    : Active_(true)
-    , PreviousTag_(GetCurrentMemoryTag())
-{
-    SetCurrentMemoryTag(tag);
-}
-
-inline TMemoryTagGuard::TMemoryTagGuard(TMemoryTagGuard&& other)
-    : Active_(other.Active_)
-    , PreviousTag_(other.PreviousTag_)
-{
-    other.Active_ = false;
-    other.PreviousTag_ = NullMemoryTag;
-}
-
-inline TMemoryTagGuard::~TMemoryTagGuard()
-{
-    if (Active_) {
-        SetCurrentMemoryTag(PreviousTag_);
-    }
-}
-
-////////////////////////////////////////////////////////////////////////////////
-
 } // namespace NYT::NYTAlloc
diff --git a/library/cpp/ytalloc/api/ytalloc.h b/library/cpp/ytalloc/api/ytalloc.h
index d942dde638..219814def2 100644
--- a/library/cpp/ytalloc/api/ytalloc.h
+++ b/library/cpp/ytalloc/api/ytalloc.h
@@ -71,35 +71,6 @@ size_t GetAllocationSize(const void* ptr);
 size_t GetAllocationSize(size_t size);
 
 ////////////////////////////////////////////////////////////////////////////////
-// Memory tagging API
-//
-// Each allocation can be tagged with a number (from 1 to MaxMemoryTag).
-// Setting this to NullMemoryTag disables tagging.
-// Internally, YTAlloc tracks the number of bytes used by each tag.
-//
-// Tagged allocations are somewhat slower. Others (large and huge) are not affected
-// (but for these performance implications are negligible anyway).
-//
-// The current memory tag used for allocations is stored in TLS.
-
-using TMemoryTag = ui32;
-constexpr TMemoryTag NullMemoryTag = 0;
-constexpr TMemoryTag MaxMemoryTag = (1ULL << 22) - 1;
-
-// Updates the current tag value in TLS.
-void SetCurrentMemoryTag(TMemoryTag tag);
-
-// Returns the current tag value from TLS.
-TMemoryTag GetCurrentMemoryTag();
-
-// Returns the memory usage for a given tag.
-// The value is somewhat approxiate and racy.
-size_t GetMemoryUsageForTag(TMemoryTag tag);
-
-// A batched version of GetMemoryUsageForTag.
-void GetMemoryUsageForTags(const TMemoryTag* tags, size_t count, size_t* results);
-
-////////////////////////////////////////////////////////////////////////////////
 // Memory zone API
 //
 // Each allocation is either in the "normal zone" or "undumpable zone".
@@ -390,25 +361,6 @@ std::vector<TProfiledAllocation> GetProfiledAllocationStatistics();
 
 ////////////////////////////////////////////////////////////////////////////////
 
-//! An RAII guard for setting the current memory tag in a scope.
-class TMemoryTagGuard
-{
-public:
-    TMemoryTagGuard();
-    explicit TMemoryTagGuard(TMemoryTag tag);
-
-    TMemoryTagGuard(const TMemoryTagGuard& other) = delete;
-    TMemoryTagGuard(TMemoryTagGuard&& other);
-
-    ~TMemoryTagGuard();
-
-private:
-    bool Active_;
-    TMemoryTag PreviousTag_;
-};
-
-////////////////////////////////////////////////////////////////////////////////
-
 } // namespace NYT::NYTAlloc
 
 #define YT_ALLOC_INL_H_
diff --git a/library/cpp/ytalloc/impl/bridge.cpp b/library/cpp/ytalloc/impl/bridge.cpp
index c8287d80de..5444a8d48b 100644
--- a/library/cpp/ytalloc/impl/bridge.cpp
+++ b/library/cpp/ytalloc/impl/bridge.cpp
@@ -4,6 +4,8 @@
 
 #include <library/cpp/malloc/api/malloc.h>
 
+#include <library/cpp/yt/memory/memory_tag.h>
+
 namespace NYT::NYTAlloc {
 
 ////////////////////////////////////////////////////////////////////////////////
@@ -49,11 +51,19 @@ size_t GetAllocationSize(size_t size)
 #endif
 
 ////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NYTAlloc
+
+namespace NYT {
+
+using namespace NYTAlloc;
+
+////////////////////////////////////////////////////////////////////////////////
 // Memory tags API bridge
 
 TMemoryTag GetCurrentMemoryTag()
 {
-    return TThreadManager::GetCurrentMemoryTag();
+    return NYTAlloc::TThreadManager::GetCurrentMemoryTag();
 }
 
 void SetCurrentMemoryTag(TMemoryTag tag)
@@ -75,6 +85,12 @@ size_t GetMemoryUsageForTag(TMemoryTag tag)
 }
 
 ////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+
+namespace NYT::NYTAlloc {
+
+////////////////////////////////////////////////////////////////////////////////
 // Memory zone API bridge
 
 void SetCurrentMemoryZone(EMemoryZone zone)
diff --git a/library/cpp/ytalloc/impl/core-inl.h b/library/cpp/ytalloc/impl/core-inl.h
index e8e5d25442..babf58aaa8 100644
--- a/library/cpp/ytalloc/impl/core-inl.h
+++ b/library/cpp/ytalloc/impl/core-inl.h
@@ -7,6 +7,8 @@
 
 #include <library/cpp/yt/containers/intrusive_linked_list.h>
 
+#include <library/cpp/yt/memory/memory_tag.h>
+
 #include <library/cpp/yt/threading/at_fork.h>
 #include <library/cpp/yt/threading/fork_aware_spin_lock.h>