Extract memory tag API to library

13 files changed, 156 insertions, 96 deletions

diff --git a/library/cpp/yt/memory/CMakeLists.darwin-x86_64.txt b/library/cpp/yt/memory/CMakeLists.darwin-x86_64.txt
index fff3bb860b..5e0d60507c 100644
--- a/library/cpp/yt/memory/CMakeLists.darwin-x86_64.txt
+++ b/library/cpp/yt/memory/CMakeLists.darwin-x86_64.txt
@@ -25,6 +25,7 @@ target_sources(cpp-yt-memory PRIVATE
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_memory_pool.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_memory_pool_output.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_output_stream.cpp
+  ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/memory_tag.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/ref.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/ref_tracked.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/safe_memory_reader.cpp
diff --git a/library/cpp/yt/memory/CMakeLists.linux-aarch64.txt b/library/cpp/yt/memory/CMakeLists.linux-aarch64.txt
index 899e0b4d0d..9e5b9c7aed 100644
--- a/library/cpp/yt/memory/CMakeLists.linux-aarch64.txt
+++ b/library/cpp/yt/memory/CMakeLists.linux-aarch64.txt
@@ -26,6 +26,7 @@ target_sources(cpp-yt-memory PRIVATE
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_memory_pool.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_memory_pool_output.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_output_stream.cpp
+  ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/memory_tag.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/ref.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/ref_tracked.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/safe_memory_reader.cpp
diff --git a/library/cpp/yt/memory/CMakeLists.linux-x86_64.txt b/library/cpp/yt/memory/CMakeLists.linux-x86_64.txt
index 899e0b4d0d..9e5b9c7aed 100644
--- a/library/cpp/yt/memory/CMakeLists.linux-x86_64.txt
+++ b/library/cpp/yt/memory/CMakeLists.linux-x86_64.txt
@@ -26,6 +26,7 @@ target_sources(cpp-yt-memory PRIVATE
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_memory_pool.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_memory_pool_output.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_output_stream.cpp
+  ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/memory_tag.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/ref.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/ref_tracked.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/safe_memory_reader.cpp
diff --git a/library/cpp/yt/memory/CMakeLists.windows-x86_64.txt b/library/cpp/yt/memory/CMakeLists.windows-x86_64.txt
index c0ecdaaeb8..9f24448cd5 100644
--- a/library/cpp/yt/memory/CMakeLists.windows-x86_64.txt
+++ b/library/cpp/yt/memory/CMakeLists.windows-x86_64.txt
@@ -22,6 +22,7 @@ target_sources(cpp-yt-memory PRIVATE
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_memory_pool.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_memory_pool_output.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/chunked_output_stream.cpp
+  ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/memory_tag.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/ref.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/ref_tracked.cpp
   ${CMAKE_SOURCE_DIR}/library/cpp/yt/memory/safe_memory_reader.cpp
diff --git a/library/cpp/yt/memory/memory_tag-inl.h b/library/cpp/yt/memory/memory_tag-inl.h
new file mode 100644
index 0000000000..42065b4899
--- /dev/null
+++ b/library/cpp/yt/memory/memory_tag-inl.h
@@ -0,0 +1,44 @@
+#pragma once
+#ifndef MEMORY_TAG_INL_H_
+#error "Direct inclusion of this file is not allowed, include memory_tag.h"
+// For the sake of sane code completion.
+#include "memory_tag.h"
+#endif
+
+#include <util/system/types.h>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+
+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
diff --git a/library/cpp/yt/memory/memory_tag.cpp b/library/cpp/yt/memory/memory_tag.cpp
new file mode 100644
index 0000000000..342be5df34
--- /dev/null
+++ b/library/cpp/yt/memory/memory_tag.cpp
@@ -0,0 +1,28 @@
+#include "memory_tag.h"
+
+#include <util/system/compiler.h>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+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*/)
+{ }
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
+
diff --git a/library/cpp/yt/memory/memory_tag.h b/library/cpp/yt/memory/memory_tag.h
new file mode 100644
index 0000000000..5f154e274e
--- /dev/null
+++ b/library/cpp/yt/memory/memory_tag.h
@@ -0,0 +1,57 @@
+#pragma once
+
+#include "public.h"
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+// 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.
+
+// 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);
+
+////////////////////////////////////////////////////////////////////////////////
+
+//! 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
+
+#define MEMORY_TAG_INL_H_
+#include "memory_tag-inl.h"
+#undef MEMORY_TAG_INL_H_
diff --git a/library/cpp/yt/memory/public.h b/library/cpp/yt/memory/public.h
index f05a6b4569..e348c23439 100644
--- a/library/cpp/yt/memory/public.h
+++ b/library/cpp/yt/memory/public.h
@@ -14,6 +14,10 @@ class TChunkedMemoryPool;
 DECLARE_REFCOUNTED_STRUCT(IMemoryChunkProvider)
 DECLARE_REFCOUNTED_STRUCT(TSharedRangeHolder)
 
+using TMemoryTag = ui32;
+constexpr TMemoryTag NullMemoryTag = 0;
+constexpr TMemoryTag MaxMemoryTag = (1ULL << 22) - 1;
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace NYT
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>