intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

8 files changed, 356 insertions, 0 deletions

diff --git a/library/cpp/actors/prof/tag.cpp b/library/cpp/actors/prof/tag.cpp
new file mode 100644
index 0000000000..9ccf03e1a9
--- /dev/null
+++ b/library/cpp/actors/prof/tag.cpp
@@ -0,0 +1,119 @@
+#include "tag.h"
+#include "tcmalloc.h"
+
+#include <library/cpp/charset/ci_string.h>
+#include <library/cpp/containers/atomizer/atomizer.h>
+#include <library/cpp/malloc/api/malloc.h>
+
+#if defined(PROFILE_MEMORY_ALLOCATIONS)
+#include <library/cpp/lfalloc/dbg_info/dbg_info.h>
+#include <library/cpp/ytalloc/api/ytalloc.h>
+#endif
+
+#include <util/generic/singleton.h>
+#include <util/generic/string.h>
+#include <util/generic/vector.h>
+#include <util/system/mutex.h>
+
+namespace NProfiling {
+    class TStringAtoms {
+    private:
+        TMutex Mutex;
+        atomizer<ci_hash, ci_equal_to> Tags;
+
+    public:
+        static TStringAtoms& Instance() {
+            return *Singleton<TStringAtoms>();
+        }
+
+        ui32 MakeTag(const char* s) {
+            Y_VERIFY(s);
+            with_lock (Mutex) {
+                return Tags.string_to_atom(s);
+            }
+        }
+
+        ui32 MakeTags(const TVector<const char*>& ss) {
+            Y_VERIFY(ss);
+            with_lock (Mutex) {
+                ui32 baseTag = Tags.string_to_atom(ss[0]);
+                ui32 nextTag = baseTag + 1;
+                for (auto i = ss.begin() + 1; i != ss.end(); ++i, ++nextTag) {
+                    Y_VERIFY(*i);
+                    ui32 ctag = Tags.string_to_atom(*i);
+                    Y_VERIFY(ctag == nextTag);
+                }
+                return baseTag;
+            }
+        }
+
+        const char* GetTag(ui32 tag) const {
+            with_lock (Mutex) {
+                return Tags.get_atom_name(tag);
+            }
+        }
+
+        size_t GetTagsCount() const {
+            with_lock (Mutex) {
+                return Tags.size();
+            }
+        }
+    };
+
+    ui32 MakeTag(const char* s) {
+        return TStringAtoms::Instance().MakeTag(s);
+    }
+
+    ui32 MakeTags(const TVector<const char*>& ss) {
+        return TStringAtoms::Instance().MakeTags(ss);
+    }
+
+    const char* GetTag(ui32 tag) {
+        return TStringAtoms::Instance().GetTag(tag);
+    }
+
+    size_t GetTagsCount() {
+        return TStringAtoms::Instance().GetTagsCount();
+    }
+
+    static ui32 SetThreadAllocTag_Default(ui32 tag) {
+        Y_UNUSED(tag);
+        return 0;
+    }
+
+#if defined(PROFILE_MEMORY_ALLOCATIONS)
+    static ui32 SetThreadAllocTag_YT(ui32 tag) {
+        auto prev = NYT::NYTAlloc::GetCurrentMemoryTag();
+        NYT::NYTAlloc::SetCurrentMemoryTag(tag);
+        return prev;
+    }
+
+    static TSetThreadAllocTag* SetThreadAllocTagFn() {
+        const auto& info = NMalloc::MallocInfo();
+
+        TStringBuf name(info.Name);
+        if (name.StartsWith("lf")) {
+            return (TSetThreadAllocTag*)NAllocDbg::SetThreadAllocTag;
+        } else if (name.StartsWith("yt")) {
+            return SetThreadAllocTag_YT;
+        } else if (name.StartsWith("tc")) {
+            return SetTCMallocThreadAllocTag;
+        } else {
+            return SetThreadAllocTag_Default;
+        }
+    }
+#else
+    static TSetThreadAllocTag* SetThreadAllocTagFn() {
+        const auto& info = NMalloc::MallocInfo();
+
+        TStringBuf name(info.Name);
+        if (name.StartsWith("tc")) {
+            return SetTCMallocThreadAllocTag;
+        } else {
+            return SetThreadAllocTag_Default;
+        }
+    }
+#endif
+
+    TSetThreadAllocTag* SetThreadAllocTag = SetThreadAllocTagFn();
+}
diff --git a/library/cpp/actors/prof/tag.h b/library/cpp/actors/prof/tag.h
new file mode 100644
index 0000000000..357e264a22
--- /dev/null
+++ b/library/cpp/actors/prof/tag.h
@@ -0,0 +1,73 @@
+#pragma once
+
+#include <util/generic/fwd.h>
+
+/*
+  Common registry for tagging memory profiler.
+  Register a new tag with MakeTag using a unique string.
+  Use registered tags with SetThreadAllocTag function in allocator API.
+*/
+
+namespace NProfiling {
+    ui32 MakeTag(const char* s);
+
+    // Make only unique tags. Y_VERIFY inside.
+    ui32 MakeTags(const TVector<const char*>& ss);
+
+    const char* GetTag(ui32 tag);
+    size_t GetTagsCount();
+
+    using TSetThreadAllocTag = ui32(ui32 tag);
+    extern TSetThreadAllocTag* SetThreadAllocTag;
+
+    class TMemoryTagScope {
+    public:
+        explicit TMemoryTagScope(ui32 tag)
+            : RestoreTag(SetThreadAllocTag(tag))
+        {
+        }
+
+        explicit TMemoryTagScope(const char* tagName) {
+            ui32 newTag = MakeTag(tagName);
+            RestoreTag = SetThreadAllocTag(newTag);
+        }
+
+        TMemoryTagScope(TMemoryTagScope&& move)
+            : RestoreTag(move.RestoreTag)
+            , Released(move.Released)
+        {
+            move.Released = true;
+        }
+
+        TMemoryTagScope& operator=(TMemoryTagScope&& move) {
+            RestoreTag = move.RestoreTag;
+            Released = move.Released;
+            move.Released = true;
+            return *this;
+        }
+
+        static void Reset(ui32 tag) {
+            SetThreadAllocTag(tag);
+        }
+
+        void Release() {
+            if (!Released) {
+                SetThreadAllocTag(RestoreTag);
+                Released = true;
+            }
+        }
+
+        ~TMemoryTagScope() {
+            if (!Released) {
+                SetThreadAllocTag(RestoreTag);
+            }
+        }
+
+    protected:
+        TMemoryTagScope(const TMemoryTagScope&) = delete;
+        void operator=(const TMemoryTagScope&) = delete;
+
+        ui32 RestoreTag = 0;
+        bool Released = false;
+    };
+}
diff --git a/library/cpp/actors/prof/tcmalloc.cpp b/library/cpp/actors/prof/tcmalloc.cpp
new file mode 100644
index 0000000000..3d4f203dbb
--- /dev/null
+++ b/library/cpp/actors/prof/tcmalloc.cpp
@@ -0,0 +1,32 @@
+#include "tcmalloc.h"
+
+#include <contrib/libs/tcmalloc/tcmalloc/malloc_extension.h>
+
+namespace NProfiling {
+
+static thread_local ui32 AllocationTag = 0;
+
+static struct TInitTCMallocCallbacks {
+    static void* CreateTag() {
+        return reinterpret_cast<void*>(AllocationTag);
+    }
+    static void* CopyTag(void* tag) {
+        return tag;
+    }
+    static void DestroyTag(void* tag) {
+        Y_UNUSED(tag);
+    }
+
+    TInitTCMallocCallbacks() {
+        tcmalloc::MallocExtension::SetSampleUserDataCallbacks(
+            CreateTag, CopyTag, DestroyTag);
+    }
+} InitTCMallocCallbacks;
+
+ui32 SetTCMallocThreadAllocTag(ui32 tag) {
+    ui32 prev = AllocationTag;
+    AllocationTag = tag;
+    return prev;
+}
+
+}
diff --git a/library/cpp/actors/prof/tcmalloc.h b/library/cpp/actors/prof/tcmalloc.h
new file mode 100644
index 0000000000..659fb4eaf3
--- /dev/null
+++ b/library/cpp/actors/prof/tcmalloc.h
@@ -0,0 +1,9 @@
+#pragma once
+
+#include <util/generic/fwd.h>
+
+namespace NProfiling {
+
+ui32 SetTCMallocThreadAllocTag(ui32 tag);
+
+}
diff --git a/library/cpp/actors/prof/tcmalloc_null.cpp b/library/cpp/actors/prof/tcmalloc_null.cpp
new file mode 100644
index 0000000000..75c0013154
--- /dev/null
+++ b/library/cpp/actors/prof/tcmalloc_null.cpp
@@ -0,0 +1,10 @@
+#include "tcmalloc.h"
+
+namespace NProfiling {
+
+ui32 SetTCMallocThreadAllocTag(ui32 tag) {
+    Y_UNUSED(tag);
+    return 0;
+}
+
+}
diff --git a/library/cpp/actors/prof/ut/tag_ut.cpp b/library/cpp/actors/prof/ut/tag_ut.cpp
new file mode 100644
index 0000000000..accf3921ab
--- /dev/null
+++ b/library/cpp/actors/prof/ut/tag_ut.cpp
@@ -0,0 +1,68 @@
+#include "tag.h"
+
+#include <library/cpp/testing/unittest/registar.h>
+
+using namespace NProfiling;
+
+class TAtomTagsTest: public TTestBase {
+private:
+    UNIT_TEST_SUITE(TAtomTagsTest);
+    UNIT_TEST(Test_MakeTag);
+    UNIT_TEST(Test_Make2Tags);
+    UNIT_TEST(Test_MakeTagTwice);
+
+    UNIT_TEST(Test_MakeAndGetTag);
+
+    UNIT_TEST(Test_MakeVector);
+    UNIT_TEST_SUITE_END();
+
+public:
+    void Test_MakeTag();
+    void Test_Make2Tags();
+    void Test_MakeTagTwice();
+    void Test_MakeAndGetTag();
+    void Test_MakeVector();
+};
+
+UNIT_TEST_SUITE_REGISTRATION(TAtomTagsTest);
+
+void TAtomTagsTest::Test_MakeTag() {
+    ui32 tag = MakeTag("a tag");
+    UNIT_ASSERT(tag != 0);
+}
+
+void TAtomTagsTest::Test_Make2Tags() {
+    ui32 tag1 = MakeTag("a tag 1");
+    ui32 tag2 = MakeTag("a tag 2");
+    UNIT_ASSERT(tag1 != 0);
+    UNIT_ASSERT(tag2 != 0);
+    UNIT_ASSERT(tag1 != tag2);
+}
+
+void TAtomTagsTest::Test_MakeTagTwice() {
+    ui32 tag1 = MakeTag("a tag twice");
+    ui32 tag2 = MakeTag("a tag twice");
+    UNIT_ASSERT(tag1 != 0);
+    UNIT_ASSERT(tag1 == tag2);
+}
+
+void TAtomTagsTest::Test_MakeAndGetTag() {
+    const char* makeStr = "tag to get";
+    ui32 tag = MakeTag(makeStr);
+    const char* tagStr = GetTag(tag);
+    UNIT_ASSERT_STRINGS_EQUAL(makeStr, tagStr);
+}
+
+void TAtomTagsTest::Test_MakeVector() {
+    TVector<const char*> strs = {
+        "vector tag 0",
+        "vector tag 1",
+        "vector tag 3",
+        "vector tag 4"};
+    ui32 baseTag = MakeTags(strs);
+    UNIT_ASSERT(baseTag != 0);
+    for (ui32 i = 0; i < strs.size(); ++i) {
+        const char* str = GetTag(baseTag + i);
+        UNIT_ASSERT_STRINGS_EQUAL(str, strs[i]);
+    }
+}
diff --git a/library/cpp/actors/prof/ut/ya.make b/library/cpp/actors/prof/ut/ya.make
new file mode 100644
index 0000000000..47c58a8fb7
--- /dev/null
+++ b/library/cpp/actors/prof/ut/ya.make
@@ -0,0 +1,12 @@
+UNITTEST_FOR(library/cpp/actors/prof)
+
+OWNER(
+    agri
+    g:kikimr
+)
+
+SRCS(
+    tag_ut.cpp
+)
+
+END()
diff --git a/library/cpp/actors/prof/ya.make b/library/cpp/actors/prof/ya.make
new file mode 100644
index 0000000000..b5e2497563
--- /dev/null
+++ b/library/cpp/actors/prof/ya.make
@@ -0,0 +1,33 @@
+LIBRARY()
+
+OWNER(
+    agri
+    g:kikimr
+)
+
+SRCS(
+    tag.cpp
+)
+
+PEERDIR(
+    library/cpp/charset
+    library/cpp/containers/atomizer
+)
+
+IF (PROFILE_MEMORY_ALLOCATIONS)
+    CFLAGS(-DPROFILE_MEMORY_ALLOCATIONS)
+    PEERDIR(
+        library/cpp/malloc/api
+        library/cpp/lfalloc/dbg_info
+        library/cpp/ytalloc/api
+    )
+ENDIF()
+
+IF(ALLOCATOR == "TCMALLOC_256K")
+    SRCS(tcmalloc.cpp)
+    PEERDIR(contrib/libs/tcmalloc)
+ELSE()
+    SRCS(tcmalloc_null.cpp)
+ENDIF()
+
+END()