Add toggle to disable autoshutdown

2 files changed, 23 insertions, 2 deletions

diff --git a/yt/yt/core/misc/shutdown.cpp b/yt/yt/core/misc/shutdown.cpp
index e029e96f304..0d46fcfdb8a 100644
--- a/yt/yt/core/misc/shutdown.cpp
+++ b/yt/yt/core/misc/shutdown.cpp
@@ -56,7 +56,7 @@ public:
         return registeredCallback;
     }
 
-    void Shutdown(const TShutdownOptions& options)
+    void Shutdown(const TShutdownOptions& options = {})
     {
         std::vector<TRegisteredCallback> registeredCallbacks;
 
@@ -123,11 +123,23 @@ public:
         }
     }
 
+    void AutoShutdown()
+    {
+        if (AutoShutdownEnabled_.load()) {
+            Shutdown();
+        }
+    }
+
     bool IsShutdownStarted()
     {
         return ShutdownStarted_.load();
     }
 
+    void SetAutoShutdownEnabled(bool enabled)
+    {
+        AutoShutdownEnabled_.store(enabled);
+    }
+
     void EnableShutdownLoggingToStderr()
     {
         ShutdownLogFile_.store(stderr);
@@ -180,6 +192,7 @@ private:
 
     std::unordered_set<TRefCountedRegisteredCallback*> RegisteredCallbacks_;
     std::atomic<bool> ShutdownStarted_ = false;
+    std::atomic<bool> AutoShutdownEnabled_ = true;
     std::atomic<size_t> ShutdownThreadId_ = 0;
 
 
@@ -227,6 +240,11 @@ bool IsShutdownStarted()
     return TShutdownManager::Get()->IsShutdownStarted();
 }
 
+void SetAutoShutdownEnabled(bool enabled)
+{
+    TShutdownManager::Get()->SetAutoShutdownEnabled(enabled);
+}
+
 void EnableShutdownLoggingToStderr()
 {
     TShutdownManager::Get()->EnableShutdownLoggingToStderr();
@@ -258,7 +276,7 @@ static const void* ShutdownGuardInitializer = [] {
             if (auto* logFile = TShutdownManager::Get()->TryGetShutdownLogFile()) {
                 fprintf(logFile, "*** Shutdown guard destructed\n");
             }
-            Shutdown();
+            TShutdownManager::Get()->AutoShutdown();
         }
     };
 
diff --git a/yt/yt/core/misc/shutdown.h b/yt/yt/core/misc/shutdown.h
index 59f261a1566..9fdb86f5476 100644
--- a/yt/yt/core/misc/shutdown.h
+++ b/yt/yt/core/misc/shutdown.h
@@ -60,6 +60,9 @@ void Shutdown(const TShutdownOptions& options = {});
 //! (and is possibly already completed).
 bool IsShutdownStarted();
 
+//! Controls if shutdown must be invoked automatically on process teardown.
+void SetAutoShutdownEnabled(bool enabled);
+
 //! Enables logging shutdown messages to stderr.
 void EnableShutdownLoggingToStderr();