Intermediate changes

10 files changed, 501 insertions, 1 deletions

diff --git a/yt/yt/library/oom/oom.cpp b/yt/yt/library/oom/oom.cpp
new file mode 100644
index 0000000000..648485c77b
--- /dev/null
+++ b/yt/yt/library/oom/oom.cpp
@@ -0,0 +1,138 @@
+#include "oom.h"
+
+#include <thread>
+#include <mutex>
+
+#include <yt/yt/core/misc/proc.h>
+#include <yt/yt/core/misc/ref_counted_tracker.h>
+
+#include <library/cpp/yt/assert/assert.h>
+#include <library/cpp/yt/logging/logger.h>
+
+#include <yt/yt/library/ytprof/heap_profiler.h>
+#include <yt/yt/library/ytprof/profile.h>
+
+#include <util/datetime/base.h>
+#include <util/system/file.h>
+#include <util/stream/output.h>
+#include <util/stream/file.h>
+#include <util/string/split.h>
+#include <util/system/fs.h>
+
+namespace NYT {
+
+static NYT::NLogging::TLogger Logger{"OOM"};
+
+////////////////////////////////////////////////////////////////////////////////
+
+static const char* TCMallocStats[] = {
+    "tcmalloc.per_cpu_caches_active",
+    "generic.virtual_memory_used",
+    "generic.physical_memory_used",
+    "generic.bytes_in_use_by_app",
+    "generic.heap_size",
+    "tcmalloc.central_cache_free",
+    "tcmalloc.cpu_free",
+    "tcmalloc.page_heap_free",
+    "tcmalloc.page_heap_unmapped",
+    "tcmalloc.page_algorithm",
+    "tcmalloc.max_total_thread_cache_bytes",
+    "tcmalloc.thread_cache_free",
+    "tcmalloc.thread_cache_count",
+    "tcmalloc.local_bytes",
+    "tcmalloc.external_fragmentation_bytes",
+    "tcmalloc.metadata_bytes",
+    "tcmalloc.transfer_cache_free",
+    "tcmalloc.hard_usage_limit_bytes",
+    "tcmalloc.desired_usage_limit_bytes",
+    "tcmalloc.required_bytes"
+};
+
+void OOMWatchdog(TOOMOptions options)
+{
+    while (true) {
+        auto rss = GetProcessMemoryUsage().Rss;
+
+        if (options.MemoryLimit && static_cast<i64>(rss) > *options.MemoryLimit) {
+            auto profile = NYTProf::ReadHeapProfile(tcmalloc::ProfileType::kHeap);
+
+            TFileOutput output(options.HeapDumpPath);
+            NYTProf::WriteProfile(&output, profile);
+            output.Finish();
+
+            auto rctDump = TRefCountedTracker::Get()->GetDebugInfo();
+            for (const auto& line : StringSplitter(rctDump).Split('\n')) {
+                YT_LOG_DEBUG("RCT %v", line.Token());
+            }
+
+            auto parseMemoryAmount = [] (const TStringBuf strValue) {
+                const TStringBuf kbSuffix = " kB";
+                YT_VERIFY(strValue.EndsWith(kbSuffix));
+                auto startPos = strValue.find_first_not_of(' ');
+                auto valueString = strValue.substr(
+                    startPos,
+                    strValue.size() - kbSuffix.size() - startPos);
+                return FromString<ui64>(valueString) * 1_KB;
+            };
+
+            ui64 rssAnon = 0;
+            ui64 rssFile = 0;
+            ui64 rssShmem = 0;
+
+            TFileInput statusFile(Format("/proc/self/status"));
+            TString line;
+            while (statusFile.ReadLine(line)) {
+                const TStringBuf rssAnonHeader = "RssAnon:\t";
+                if (line.StartsWith(rssAnonHeader)) {
+                    rssAnon = parseMemoryAmount(line.substr(rssAnonHeader.size()));
+                    continue;
+                }
+
+                const TStringBuf rssFileHeader = "RssFile:\t";
+                if (line.StartsWith(rssFileHeader)) {
+                    rssFile = parseMemoryAmount(line.substr(rssFileHeader.size()));
+                    continue;
+                }
+
+                const TStringBuf rssShmemHeader = "RssShmem:\t";
+                if (line.StartsWith(rssShmemHeader)) {
+                    rssShmem = parseMemoryAmount(line.substr(rssShmemHeader.size()));
+                    continue;
+                }
+            }
+
+            YT_LOG_DEBUG("Memory statistis (RssTotal: %v, RssAnon: %v, RssFile %v, RssShmem: %v, TCMalloc: %v)",
+                rss,
+                rssAnon,
+                rssFile,
+                rssShmem,
+                MakeFormattableView(
+                    MakeRange(TCMallocStats),
+                    [&] (auto* builder, auto metric) {
+                        auto value = tcmalloc::MallocExtension::GetNumericProperty(metric);
+                        builder->AppendFormat("%v: %v", metric, value);
+                    }));
+
+            YT_LOG_FATAL("Early OOM triggered (MemoryUsage: %v, MemoryLimit: %v, HeapDump: %v, CurrentWorkingDirectory: %v)",
+                rss,
+                *options.MemoryLimit,
+                options.HeapDumpPath,
+                NFs::CurrentWorkingDirectory());
+        }
+
+        Sleep(TDuration::MilliSeconds(10));
+    }
+}
+
+void EnableEarlyOOMWatchdog(TOOMOptions options)
+{
+    static std::once_flag onceFlag;
+
+    std::call_once(onceFlag, [options] {
+        std::thread(OOMWatchdog, options).detach();
+    });
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/yt/yt/library/oom/oom.h b/yt/yt/library/oom/oom.h
new file mode 100644
index 0000000000..aa85cc678b
--- /dev/null
+++ b/yt/yt/library/oom/oom.h
@@ -0,0 +1,34 @@
+#pragma once
+
+#include <optional>
+
+#include <util/system/types.h>
+#include <util/generic/string.h>
+#include <util/datetime/base.h>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+struct TOOMOptions
+{
+    std::optional<i64> MemoryLimit;
+    TString HeapDumpPath = "oom.pb.gz";
+};
+
+void EnableEarlyOOMWatchdog(TOOMOptions options);
+
+////////////////////////////////////////////////////////////////////////////////
+
+struct TTCMallocLimitHandlerOptions
+{
+    TString HeapDumpDirectory;
+    TDuration Timeout = TDuration::Seconds(300);
+};
+
+void EnableTCMallocLimitHandler(TTCMallocLimitHandlerOptions options);
+void DisableTCMallocLimitHandler();
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace NYT
diff --git a/yt/yt/library/oom/tcmalloc_memory_limit_handler.cpp b/yt/yt/library/oom/tcmalloc_memory_limit_handler.cpp
new file mode 100644
index 0000000000..48520bd873
--- /dev/null
+++ b/yt/yt/library/oom/tcmalloc_memory_limit_handler.cpp
@@ -0,0 +1,213 @@
+#include "oom.h"
+
+#include <yt/yt/library/ytprof/external_pprof.h>
+#include <yt/yt/library/ytprof/heap_profiler.h>
+#include <yt/yt/library/ytprof/profile.h>
+#include <yt/yt/library/ytprof/spinlock_profiler.h>
+#include <yt/yt/library/ytprof/symbolize.h>
+
+#include <yt/yt/core/misc/crash_handler.h>
+#include <yt/yt/core/misc/error.h>
+
+#include <library/cpp/yt/string/format.h>
+
+#include <library/cpp/yt/memory/atomic_intrusive_ptr.h>
+
+#include <util/datetime/base.h>
+#include <util/stream/file.h>
+#include <util/stream/output.h>
+#include <util/string/split.h>
+#include <util/system/env.h>
+#include <util/system/file.h>
+#include <util/system/fs.h>
+#include <util/system/shellcommand.h>
+
+#include <mutex>
+#include <thread>
+
+namespace NYT {
+
+////////////////////////////////////////////////////////////////////////////////
+
+void CollectAndDumpMemoryProfile(const TString& memoryProfilePath)
+{
+    auto profile = NYTProf::ReadHeapProfile(tcmalloc::ProfileType::kHeap);
+    SymbolizeByExternalPProf(&profile, NYTProf::TSymbolizationOptions{
+        .RunTool = [] (const std::vector<TString>& args) {
+            TShellCommand command{args[0], TList<TString>{args.begin()+1, args.end()}};
+            command.Run();
+        },
+    });
+
+    TFileOutput output(memoryProfilePath);
+    NYTProf::WriteProfile(&output, profile);
+    output.Finish();
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+void MemoryProfileTimeoutHandler(int /*signal*/)
+{
+    WriteToStderr("*** Process hung during dumping heap profile ***\n");
+    ::_exit(1);
+}
+
+void SetupMemoryProfileTimeout(int timeout)
+{
+    ::signal(SIGALRM, &MemoryProfileTimeoutHandler);
+    ::alarm(timeout);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+class TTCMallocLimitHandler
+    : public TRefCounted
+{
+public:
+    explicit TTCMallocLimitHandler(TTCMallocLimitHandlerOptions options)
+        : Options_(options)
+    {
+        Thread_ = std::thread([this] {
+            Handle();
+        });
+    }
+
+    ~TTCMallocLimitHandler()
+    {
+        {
+            std::unique_lock<std::mutex> lock(Mutex_);
+            Fired_ = true;
+            CV_.notify_all();
+        }
+
+        Thread_.join();
+    }
+
+    void Fire()
+    {
+        std::unique_lock<std::mutex> lock(Mutex_);
+        Fired_ = true;
+        NeedToHandle_ = true;
+        CV_.notify_all();
+    }
+
+private:
+    const TTCMallocLimitHandlerOptions Options_;
+    bool Fired_ = false;
+    bool NeedToHandle_ = false;
+    std::mutex Mutex_;
+    std::condition_variable CV_;
+    std::thread Thread_;
+
+
+    void Handle()
+    {
+        std::unique_lock<std::mutex> lock(Mutex_);
+        CV_.wait(lock, [&] {
+            return Fired_;
+        });
+
+        if (!NeedToHandle_) {
+            return;
+        }
+
+        auto heapDumpPath = GetHeapDumpPath();
+        Cerr << "TTCMallocLimitHandler: Fork process to write heap profile: "
+            << heapDumpPath
+            << Endl;
+
+        SetupMemoryProfileTimeout(Options_.Timeout.Seconds());
+        auto childPid = fork();
+
+        if (childPid == 0) {
+            SetupMemoryProfileTimeout(Options_.Timeout.Seconds());
+            CollectAndDumpMemoryProfile(heapDumpPath);
+
+            Cerr << "TTCMallocLimitHandler: Heap profile written." << Endl;
+            ::_exit(0);
+        }
+
+        if (childPid < 0) {
+            Cerr << "TTCMallocLimitHandler: fork failed with code:" << LastSystemErrorText() << Endl;
+            ::_exit(1);
+        }
+
+        ExecWaitForChild(childPid);
+        ::_exit(0);
+    }
+
+    TString GetHeapDumpPath() const
+    {
+        return Format(
+            "%v/heap_%v.pb.gz",
+            Options_.HeapDumpDirectory,
+            TInstant::Now().FormatLocalTime("%Y%m%dT%H%M%S"));
+    }
+
+    void ExecWaitForChild(int pid)
+    {
+        Cerr << "TTCMallocLimitHandler: Before waiting for child" << Endl;
+
+        auto command = Format("while [ -e /proc/%v ]; do sleep 1; done;", pid);
+        execl("/bin/bash", "/bin/bash", "-c",  command.c_str(), (void*)nullptr);
+
+        Cerr << "TTCMallocLimitHandler: Failed to switch main process to dummy child waiter: "
+            << LastSystemErrorText() << Endl;
+    }
+};
+
+DEFINE_REFCOUNTED_TYPE(TTCMallocLimitHandler);
+
+////////////////////////////////////////////////////////////////////////////////
+
+template <class TMallocExtension>
+concept CSupportsLimitHandler = requires (TMallocExtension extension)
+{
+    { extension.GetSoftMemoryLimitHandler() };
+};
+
+template <typename TMallocExtension, typename THandler>
+void SetSoftMemoryLimitHandler(THandler)
+{
+    WriteToStderr("TCMalloc does not support memory limit handler\n");
+}
+
+template <CSupportsLimitHandler TMallocExtension, typename THandler>
+void SetSoftMemoryLimitHandler(THandler handler)
+{
+    TMallocExtension::SetSoftMemoryLimitHandler(handler);
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+static TAtomicIntrusivePtr<TTCMallocLimitHandler> LimitHandler_;
+
+////////////////////////////////////////////////////////////////////////////////
+
+void HandleTCMallocLimit()
+{
+    if (auto handler = LimitHandler_.Acquire()) {
+        handler->Fire();
+    }
+}
+
+void EnableTCMallocLimitHandler(TTCMallocLimitHandlerOptions options)
+{
+    {
+        if (LimitHandler_.Acquire()) {
+            return;
+        }
+
+        TAtomicIntrusivePtr<TTCMallocLimitHandler>::TRawPtr expected{nullptr};
+        LimitHandler_.CompareAndSwap(expected, New<TTCMallocLimitHandler>(options));
+    }
+
+    SetSoftMemoryLimitHandler<tcmalloc::MallocExtension>(&HandleTCMallocLimit);
+}
+
+void DisableTCMallocLimitHandler()
+{
+    LimitHandler_.Reset();
+}
+
+} // namespace NYT
diff --git a/yt/yt/library/oom/unittests/oom_ut.cpp b/yt/yt/library/oom/unittests/oom_ut.cpp
new file mode 100644
index 0000000000..6c61f456f9
--- /dev/null
+++ b/yt/yt/library/oom/unittests/oom_ut.cpp
@@ -0,0 +1,41 @@
+#include <gtest/gtest.h>
+
+#include <yt/yt/library/oom/oom.h>
+
+#include <util/datetime/base.h>
+#include <util/system/fs.h>
+#include <util/generic/size_literals.h>
+
+namespace NYT {
+namespace {
+
+////////////////////////////////////////////////////////////////////////////////
+
+TEST(EarlyOOM, Crash)
+{
+    auto checkOOM = [] {
+        EnableEarlyOOMWatchdog(TOOMOptions{
+            .MemoryLimit = 0,
+        });
+
+        Sleep(TDuration::Seconds(5));
+    };
+
+    ASSERT_DEATH(checkOOM(), "");
+
+    ASSERT_TRUE(NFs::Exists("oom.pb.gz"));
+}
+
+TEST(EarlyOOM, NoCrash)
+{
+    EnableEarlyOOMWatchdog(TOOMOptions{
+        .MemoryLimit = 1_GB,
+    });
+
+    Sleep(TDuration::Seconds(5));
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+} // namespace
+} // namespace NYT
diff --git a/yt/yt/library/oom/unittests/ya.make b/yt/yt/library/oom/unittests/ya.make
new file mode 100644
index 0000000000..23392352b9
--- /dev/null
+++ b/yt/yt/library/oom/unittests/ya.make
@@ -0,0 +1,17 @@
+GTEST()
+
+INCLUDE(${ARCADIA_ROOT}/yt/ya_cpp.make.inc)
+
+ALLOCATOR(TCMALLOC)
+
+SRCS(
+    oom_ut.cpp
+)
+
+INCLUDE(${ARCADIA_ROOT}/yt/opensource.inc)
+
+PEERDIR(
+    yt/yt/library/oom
+)
+
+END()
diff --git a/yt/yt/library/oom/ya.make b/yt/yt/library/oom/ya.make
new file mode 100644
index 0000000000..5ad7018578
--- /dev/null
+++ b/yt/yt/library/oom/ya.make
@@ -0,0 +1,22 @@
+LIBRARY()
+
+INCLUDE(${ARCADIA_ROOT}/yt/ya_cpp.make.inc)
+
+SRCS(
+    oom.cpp
+    tcmalloc_memory_limit_handler.cpp
+)
+
+PEERDIR(
+    yt/yt/core
+    yt/yt/library/ytprof
+    library/cpp/yt/logging
+)
+
+END()
+
+IF (OS_LINUX AND NOT SANITIZER_TYPE)
+    RECURSE(
+        unittests
+    )
+ENDIF()
diff --git a/yt/yt/library/program/config.cpp b/yt/yt/library/program/config.cpp
index 1b903185d7..b44abdd6e9 100644
--- a/yt/yt/library/program/config.cpp
+++ b/yt/yt/library/program/config.cpp
@@ -14,6 +14,12 @@ void THeapSizeLimit::Register(TRegistrar registrar)
         .Optional();
     registrar.Parameter("hard", &TThis::Hard)
         .Default(false);
+    registrar.Parameter("dump_memory_profile_on_violation", &TThis::DumpMemoryProfileOnViolation)
+        .Default(false);
+    registrar.Parameter("dump_memory_profile_timeout", &TThis::DumpMemoryProfileTimeout)
+        .Default(TDuration::Minutes(10));
+    registrar.Parameter("dump_memory_profile_path", &TThis::DumpMemoryProfilePath)
+        .Default();
 }
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/yt/yt/library/program/config.h b/yt/yt/library/program/config.h
index 78c63c21e1..96f70cc339 100644
--- a/yt/yt/library/program/config.h
+++ b/yt/yt/library/program/config.h
@@ -46,6 +46,10 @@ public:
     //! If true tcmalloc crashes when system allocates more memory than #ContainerMemoryRatio.
     bool Hard;
 
+    bool DumpMemoryProfileOnViolation;
+    TDuration DumpMemoryProfileTimeout;
+    TString DumpMemoryProfilePath;
+
     REGISTER_YSON_STRUCT(THeapSizeLimit);
 
     static void Register(TRegistrar registrar);
diff --git a/yt/yt/library/program/helpers.cpp b/yt/yt/library/program/helpers.cpp
index 51b66b7568..606c1cc406 100644
--- a/yt/yt/library/program/helpers.cpp
+++ b/yt/yt/library/program/helpers.cpp
@@ -10,6 +10,8 @@
 
 #include <yt/yt/core/bus/tcp/dispatcher.h>
 
+#include <yt/yt/library/oom/oom.h>
+
 #include <yt/yt/library/tracing/jaeger/tracer.h>
 
 #include <yt/yt/library/profiling/perf/counters.h>
@@ -63,6 +65,7 @@ public:
         i64 totalMemory = GetContainerMemoryLimit();
         AdjustPageHeapLimit(totalMemory, config);
         AdjustAggressiveReleaseThreshold(totalMemory, config);
+        SetupMemoryLimitHandler(config);
     }
 
     i64 GetAggressiveReleaseThreshold()
@@ -103,6 +106,20 @@ private:
         }
     }
 
+    void SetupMemoryLimitHandler(const TTCMallocConfigPtr& config)
+    {
+        TTCMallocLimitHandlerOptions handlerOptions {
+            .HeapDumpDirectory = config->HeapSizeLimit->DumpMemoryProfilePath,
+            .Timeout = config->HeapSizeLimit->DumpMemoryProfileTimeout,
+        };
+
+        if (config->HeapSizeLimit->DumpMemoryProfileOnViolation) {
+            EnableTCMallocLimitHandler(handlerOptions);
+        } else {
+            DisableTCMallocLimitHandler();
+        }
+    }
+
     i64 GetContainerMemoryLimit() const
     {
         auto resourceTracker = NProfiling::GetResourceTracker();
@@ -237,6 +254,13 @@ void ConfigureSingletons(const TSingletonsConfigPtr& config)
     NYson::SetProtobufInteropConfig(config->ProtobufInterop);
 }
 
+TTCMallocConfigPtr MergeTCMallocDynamicConfig(const TTCMallocConfigPtr& staticConfig, const TTCMallocConfigPtr& dynamicConfig)
+{
+    auto mergedConfig = CloneYsonStruct(dynamicConfig);
+    mergedConfig->HeapSizeLimit->DumpMemoryProfilePath = staticConfig->HeapSizeLimit->DumpMemoryProfilePath;
+    return mergedConfig;
+}
+
 void ReconfigureSingletons(const TSingletonsConfigPtr& config, const TSingletonsDynamicConfigPtr& dynamicConfig)
 {
     SetSpinWaitSlowPathLoggingThreshold(dynamicConfig->SpinWaitSlowPathLoggingThreshold.value_or(config->SpinWaitSlowPathLoggingThreshold));
@@ -271,7 +295,7 @@ void ReconfigureSingletons(const TSingletonsConfigPtr& config, const TSingletons
     }
 
     if (dynamicConfig->TCMalloc) {
-        ConfigureTCMalloc(dynamicConfig->TCMalloc);
+        ConfigureTCMalloc(MergeTCMallocDynamicConfig(config->TCMalloc, dynamicConfig->TCMalloc));
     } else if (config->TCMalloc) {
         ConfigureTCMalloc(config->TCMalloc);
     }
diff --git a/yt/yt/library/program/ya.make b/yt/yt/library/program/ya.make
index 02da9f2e57..4c0604f04f 100644
--- a/yt/yt/library/program/ya.make
+++ b/yt/yt/library/program/ya.make
@@ -16,6 +16,7 @@ PEERDIR(
     yt/yt/core
     yt/yt/core/service_discovery/yp
     yt/yt/library/monitoring
+    yt/yt/library/oom
     yt/yt/library/profiling/solomon
     yt/yt/library/profiling/tcmalloc
     yt/yt/library/profiling/perf