Intermediate changes

5 files changed, 168 insertions, 0 deletions

diff --git a/yt/yt/library/ytprof/allocation_tag_profiler/allocation_tag_profiler.cpp b/yt/yt/library/ytprof/allocation_tag_profiler/allocation_tag_profiler.cpp
new file mode 100644
index 0000000000..de5e7988bc
--- /dev/null
+++ b/yt/yt/library/ytprof/allocation_tag_profiler/allocation_tag_profiler.cpp
@@ -0,0 +1,84 @@
+#include "allocation_tag_profiler.h"
+
+#include <yt/yt/core/concurrency/periodic_executor.h>
+#include <yt/yt/library/ytprof/heap_profiler.h>
+
+namespace NYT::NYTProf {
+
+using namespace NProfiling;
+using namespace NConcurrency;
+
+////////////////////////////////////////////////////////////////////////////////
+
+THeapUsageProfiler::THeapUsageProfiler(
+    std::vector<TString> tags,
+    IInvokerPtr invoker,
+    std::optional<TDuration> updatePeriod,
+    std::optional<i64> samplingRate,
+    NProfiling::TProfiler profiler)
+    : Profiler_(std::move(profiler))
+    , TagTypes_(std::move(tags))
+    , UpdateExecutor_(New<TPeriodicExecutor>(
+        std::move(invoker),
+        BIND(&THeapUsageProfiler::UpdateGauges, MakeWeak(this)),
+        std::move(updatePeriod)))
+{
+    if (samplingRate) {
+        tcmalloc::MallocExtension::SetProfileSamplingRate(*samplingRate);
+    }
+
+    UpdateExecutor_->Start();
+}
+
+void THeapUsageProfiler::UpdateGauges()
+{
+    const auto memorySnapshot = GetMemoryUsageSnapshot();
+    YT_VERIFY(memorySnapshot);
+
+    for (const auto& tagType : TagTypes_) {
+        auto& heapUsageMap = HeapUsageByType_.emplace(tagType, THashMap<TString, TGauge>{}).first->second;
+        const auto& snapshotSlice = memorySnapshot->GetUsage(tagType);
+
+        for (auto &[tag, gauge] : heapUsageMap) {
+            if (const auto& iter = snapshotSlice.find(tag)) {
+                gauge.Update(iter->second);
+            } else {
+                gauge.Update(0.0);
+            }
+        }
+
+        for (const auto& [tag, usage] : snapshotSlice) {
+            auto gauge = heapUsageMap.find(tag);
+
+            if (gauge.IsEnd()) {
+                gauge = heapUsageMap.emplace(tag, Profiler_
+                    .WithTag(tagType, tag)
+                    .Gauge(tagType))
+                    .first;
+                gauge->second.Update(usage);
+            }
+        }
+    }
+}
+
+///////////////////////////////////////////////////////////////////
+
+THeapUsageProfilerPtr CreateHeapProfilerWithTags(
+    std::vector<TString>&& tags,
+    IInvokerPtr invoker,
+    std::optional<TDuration> updatePeriod,
+    std::optional<i64> samplingRate,
+    NYT::NProfiling::TProfiler profiler)
+{
+    return New<THeapUsageProfiler>(
+        std::move(tags),
+        std::move(invoker),
+        std::move(updatePeriod),
+        std::move(samplingRate),
+        std::move(profiler));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NYTProf
+
diff --git a/yt/yt/library/ytprof/allocation_tag_profiler/allocation_tag_profiler.h b/yt/yt/library/ytprof/allocation_tag_profiler/allocation_tag_profiler.h
new file mode 100644
index 0000000000..fcf7be4cfe
--- /dev/null
+++ b/yt/yt/library/ytprof/allocation_tag_profiler/allocation_tag_profiler.h
@@ -0,0 +1,53 @@
+#pragma once
+
+#include "public.h"
+
+#include <yt/yt/library/profiling/sensor.h>
+
+#include <yt/yt/core/actions/public.h>
+
+#include <yt/yt/core/concurrency/public.h>
+
+#include <yt/yt/core/tracing/public.h>
+
+#include <yt/yt/core/profiling/public.h>
+
+namespace NYT::NYTProf {
+
+////////////////////////////////////////////////////////////////////////////////
+
+class THeapUsageProfiler
+    : public TRefCounted
+{
+public:
+    THeapUsageProfiler(
+        std::vector<TString> tags,
+        IInvokerPtr invoker,
+        std::optional<TDuration> updatePeriod,
+        std::optional<i64> samplingRate,
+        NProfiling::TProfiler profiler);
+
+private:
+    NProfiling::TProfiler Profiler_;
+    const std::vector<TString> TagTypes_;
+    THashMap<TString, THashMap<TString, NProfiling::TGauge>> HeapUsageByType_;
+
+    const NConcurrency::TPeriodicExecutorPtr UpdateExecutor_;
+
+    void UpdateGauges();
+};
+
+DEFINE_REFCOUNTED_TYPE(THeapUsageProfiler)
+
+////////////////////////////////////////////////////////////////////////////////
+
+THeapUsageProfilerPtr CreateHeapProfilerWithTags(
+    std::vector<TString>&& tags,
+    IInvokerPtr invoker,
+    std::optional<TDuration> updatePeriod,
+    std::optional<i64> samplingRate,
+    NProfiling::TProfiler profiler = NProfiling::TProfiler{"/heap_usage/"});
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NYTProf
diff --git a/yt/yt/library/ytprof/allocation_tag_profiler/public.h b/yt/yt/library/ytprof/allocation_tag_profiler/public.h
new file mode 100644
index 0000000000..b5d8d9f983
--- /dev/null
+++ b/yt/yt/library/ytprof/allocation_tag_profiler/public.h
@@ -0,0 +1,13 @@
+#pragma once
+
+#include <library/cpp/yt/memory/ref_counted.h>
+
+namespace NYT::NYTProf {
+
+////////////////////////////////////////////////////////////////////////////////
+
+DECLARE_REFCOUNTED_CLASS(THeapUsageProfiler);
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT::NYTProf
diff --git a/yt/yt/library/ytprof/allocation_tag_profiler/ya.make b/yt/yt/library/ytprof/allocation_tag_profiler/ya.make
new file mode 100644
index 0000000000..24aa93e440
--- /dev/null
+++ b/yt/yt/library/ytprof/allocation_tag_profiler/ya.make
@@ -0,0 +1,17 @@
+LIBRARY()
+
+INCLUDE(${ARCADIA_ROOT}/yt/ya_cpp.make.inc)
+
+SRCS(
+    allocation_tag_profiler.cpp
+)
+
+PEERDIR(
+    yt/yt/library/profiling
+    yt/yt/library/ytprof
+    yt/yt/core
+)
+
+END()
+
+
diff --git a/yt/yt/library/ytprof/ya.make b/yt/yt/library/ytprof/ya.make
index 3ceb7aeda9..f093c2f7ad 100644
--- a/yt/yt/library/ytprof/ya.make
+++ b/yt/yt/library/ytprof/ya.make
@@ -51,6 +51,7 @@ CXXFLAGS(-DYTPROF_BUILD_TYPE='\"${BUILD_TYPE}\"')
 END()
 
 RECURSE(
+    allocation_tag_profiler
     http
     example
     bundle