From a2be43bf27f7e6d615c83fe74287eac083c252f6 Mon Sep 17 00:00:00 2001
From: Ilia Shakhov <pixcc@ydb.tech>
Date: Mon, 7 Oct 2024 17:14:05 +0300
Subject: [PATCH] 24-3: Add min delay before shutdown (#9688) (#10087)

---
 ydb/core/driver_lib/run/run.cpp | 12 +++++++++++-
 ydb/core/driver_lib/run/run.h   |  1 +
 ydb/core/protos/config.proto    |  5 +++++
 3 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/ydb/core/driver_lib/run/run.cpp b/ydb/core/driver_lib/run/run.cpp
index e1e7ef77d0f3..5413287f4b96 100644
--- a/ydb/core/driver_lib/run/run.cpp
+++ b/ydb/core/driver_lib/run/run.cpp
@@ -495,8 +495,11 @@ static TString ReadFile(const TString& fileName) {
 }
 
 void TKikimrRunner::InitializeGracefulShutdown(const TKikimrRunConfig& runConfig) {
-    Y_UNUSED(runConfig);
     GracefulShutdownSupported = true;
+    const auto& config = runConfig.AppConfig.GetShutdownConfig();
+    if (config.HasMinDelayBeforeShutdownSeconds()) {
+        MinDelayBeforeShutdown = TDuration::Seconds(config.GetMinDelayBeforeShutdownSeconds());
+    }
 }
 
 void TKikimrRunner::InitializeKqpController(const TKikimrRunConfig& runConfig) {
@@ -1708,6 +1711,7 @@ void TKikimrRunner::KikimrStop(bool graceful) {
         ActorSystem->Send(new IEventHandle(NGRpcService::CreateGrpcPublisherServiceActorId(), {}, new TEvents::TEvPoisonPill));
     }
 
+    THPTimer timer;
     TIntrusivePtr<TDrainProgress> drainProgress(new TDrainProgress());
     if (AppData->FeatureFlags.GetEnableDrainOnShutdown() && GracefulShutdownSupported && ActorSystem) {
         drainProgress->OnSend();
@@ -1741,6 +1745,12 @@ void TKikimrRunner::KikimrStop(bool graceful) {
         }
     }
 
+    // Wait for a minimum delay to make sure that clients forget about this node
+    auto passedTime = TDuration::Seconds(timer.Passed());
+    if (MinDelayBeforeShutdown > passedTime) {
+        Sleep(MinDelayBeforeShutdown - passedTime);
+    }
+
     if (ActorSystem) {
         ActorSystem->BroadcastToProxies([](const TActorId& proxyId) {
             return new IEventHandle(proxyId, {}, new TEvInterconnect::TEvTerminate);
diff --git a/ydb/core/driver_lib/run/run.h b/ydb/core/driver_lib/run/run.h
index a6c9283fc43d..eaa4ff1f91a6 100644
--- a/ydb/core/driver_lib/run/run.h
+++ b/ydb/core/driver_lib/run/run.h
@@ -42,6 +42,7 @@ class TKikimrRunner : public virtual TThrRefBase, private IGlobalObjectStorage {
 
     bool EnabledGrpcService = false;
     bool GracefulShutdownSupported = false;
+    TDuration MinDelayBeforeShutdown;
     THolder<NSQS::TAsyncHttpServer> SqsHttp;
 
     THolder<NYdb::TDriver> YdbDriver;
diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto
index 1bf74c929823..5757e561370e 100644
--- a/ydb/core/protos/config.proto
+++ b/ydb/core/protos/config.proto
@@ -1853,6 +1853,10 @@ message TMetadataCacheConfig {
     optional uint64 RefreshPeriodMs = 1 [default = 15000];
 }
 
+message TShutdownConfig {
+    optional uint32 MinDelayBeforeShutdownSeconds = 1;
+}
+
 message TLabel {
     optional string Name = 1;
     optional string Value = 2;
@@ -1937,6 +1941,7 @@ message TAppConfig {
     //optional TMemoryControllerConfig MemoryControllerConfig = 81; NB. exist in main
 	optional TGroupedMemoryLimiterConfig GroupedMemoryLimiterConfig = 82;
     optional NKikimrReplication.TReplicationDefaults ReplicationConfig = 83;
+    optional TShutdownConfig ShutdownConfig = 84;
 
     repeated TNamedConfig NamedConfigs = 100;
     optional string ClusterYamlConfig = 101;