From 7cdad883278f9764b6888ec2d7d7dc5d7967f75a Mon Sep 17 00:00:00 2001
From: Antonio Andelic <antonio@clickhouse.com>
Date: Sat, 30 Jul 2022 11:34:17 +0000
Subject: [PATCH] Create snapshot on shutdown

---
 src/Coordination/KeeperContext.h        | 3 ++-
 src/Coordination/KeeperServer.cpp       | 9 ++++++++-
 src/Coordination/KeeperServer.h         | 2 ++
 src/Coordination/KeeperStateMachine.cpp | 8 +++++++-
 4 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/src/Coordination/KeeperContext.h b/src/Coordination/KeeperContext.h
index 84ec65cecde..64fa8cea6ec 100644
--- a/src/Coordination/KeeperContext.h
+++ b/src/Coordination/KeeperContext.h
@@ -8,7 +8,8 @@ struct KeeperContext
     enum class Phase : uint8_t
     {
         INIT,
-        RUNNING
+        RUNNING,
+        SHUTDOWN
     };
 
     Phase server_state{Phase::INIT};
diff --git a/src/Coordination/KeeperServer.cpp b/src/Coordination/KeeperServer.cpp
index 587ab9c8f66..20ce7e42afc 100644
--- a/src/Coordination/KeeperServer.cpp
+++ b/src/Coordination/KeeperServer.cpp
@@ -107,8 +107,9 @@ KeeperServer::KeeperServer(
     : server_id(configuration_and_settings_->server_id)
     , coordination_settings(configuration_and_settings_->coordination_settings)
     , log(&Poco::Logger::get("KeeperServer"))
-    , is_recovering(config.has("keeper_server.force_recovery") && config.getBool("keeper_server.force_recovery"))
+    , is_recovering(config.getBool("keeper_server.force_recovery", false))
     , keeper_context{std::make_shared<KeeperContext>()}
+    , create_snapshot_on_exit(config.getBool("keeper_server.create_snapshot_on_exit", true))
 {
     if (coordination_settings->quorum_reads)
         LOG_WARNING(log, "Quorum reads enabled, Keeper will work slower.");
@@ -367,6 +368,12 @@ void KeeperServer::shutdownRaftServer()
     }
 
     raft_instance->shutdown();
+
+    keeper_context->server_state = KeeperContext::Phase::SHUTDOWN;
+
+    if (create_snapshot_on_exit)
+        raft_instance->create_snapshot();
+
     raft_instance.reset();
 
     if (asio_listener)
diff --git a/src/Coordination/KeeperServer.h b/src/Coordination/KeeperServer.h
index 74dd05631f0..1fb3e579214 100644
--- a/src/Coordination/KeeperServer.h
+++ b/src/Coordination/KeeperServer.h
@@ -64,6 +64,8 @@ private:
 
     std::shared_ptr<KeeperContext> keeper_context;
 
+    const bool create_snapshot_on_exit;
+
 public:
     KeeperServer(
         const KeeperConfigurationAndSettingsPtr & settings_,
diff --git a/src/Coordination/KeeperStateMachine.cpp b/src/Coordination/KeeperStateMachine.cpp
index f43a3dbb319..a55acaf9b91 100644
--- a/src/Coordination/KeeperStateMachine.cpp
+++ b/src/Coordination/KeeperStateMachine.cpp
@@ -383,7 +383,13 @@ void KeeperStateMachine::create_snapshot(nuraft::snapshot & s, nuraft::async_res
     };
 
 
-    LOG_DEBUG(log, "In memory snapshot {} created, queueing task to flash to disk", s.get_last_log_idx());
+    if (keeper_context->server_state == KeeperContext::Phase::SHUTDOWN)
+    {
+        snapshot_task.create_snapshot(std::move(snapshot_task.snapshot));
+        return;
+    }
+
+    LOG_DEBUG(log, "In memory snapshot {} created, queueing task to flush to disk", s.get_last_log_idx());
     /// Flush snapshot to disk in a separate thread.
     if (!snapshots_queue.push(std::move(snapshot_task)))
         LOG_WARNING(log, "Cannot push snapshot task into queue");