Extract memory tag API to library

8 files changed, 137 insertions, 0 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