Gracefully exit if Ctrl-C is pressed (#293)

* Feature: Gracefully exit if Ctrl-C is pressed
* Bugfix: Return error code on exit as set by application
* changelog

5 files changed, 54 insertions, 6 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 2c9aa09db..6d64fd347 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,4 +1,9 @@
 ##
+- Feature: Gracefully exit if Ctrl-C is pressed
+- Bugfix: Return error code on exit as set by application
+- Bugfix: Fix crash at startup if dead process handles are detected in ZenServerState
+
+## 0.2.10
 - Feature: zenserver now writes a state_marker file in the root of the data directory. Deleting this file will cause zenserver to exit. This is used to detect if user is deleting the data folder while zenserver is running
 - Feature: Disk writes are now blocked early and return an insufficient storage error if free disk space falls below the `--low-diskspace-threshold` value
 - Feature: zenserver: Add command line option `--sentry-allow-personal-info` to allow personally identifiable information in sentry reports, disabled by default
@@ -8,7 +13,6 @@
 - Bugfix: Close down http server gracefully when exiting even while requests are still being processed
 - Bugfix: Flush snapshot for filecase on flush/exit
 - Bugfix: Fix log of size found when scanning for files in filecas
-- Bugfix: Fix crash at startup if dead process handles are detected in ZenServerState
 
 ## 0.2.9
 - Bugfix: Treat reading outside of block store file as a not found error. We may encounter truncated blocks due to earlier abnormal termination of zenserver or disk failures
diff --git a/src/zencore/include/zencore/zencore.h b/src/zencore/include/zencore/zencore.h
index 5bcd77239..d4173e128 100644
--- a/src/zencore/include/zencore/zencore.h
+++ b/src/zencore/include/zencore/zencore.h
@@ -329,6 +329,7 @@ namespace zen {
 
 ZENCORE_API bool IsApplicationExitRequested();
 ZENCORE_API void RequestApplicationExit(int ExitCode);
+ZENCORE_API int	 ApplicationExitCode();
 ZENCORE_API bool IsDebuggerPresent();
 ZENCORE_API void SetIsInteractiveSession(bool Value);
 ZENCORE_API bool IsInteractiveSession();
diff --git a/src/zencore/zencore.cpp b/src/zencore/zencore.cpp
index 96d61a0ab..ae3d3934d 100644
--- a/src/zencore/zencore.cpp
+++ b/src/zencore/zencore.cpp
@@ -99,6 +99,12 @@ RequestApplicationExit(int ExitCode)
 	s_ApplicationExitRequested = true;
 }
 
+int
+ApplicationExitCode()
+{
+	return s_ApplicationExitCode;
+}
+
 #if ZEN_WITH_TESTS
 void
 zencore_forcelinktests()
diff --git a/src/zenserver/windows/service.cpp b/src/zenserver/windows/service.cpp
index 89bacab0b..21116f1c4 100644
--- a/src/zenserver/windows/service.cpp
+++ b/src/zenserver/windows/service.cpp
@@ -179,7 +179,7 @@ WindowsService::ServiceMain()
 
 	zen::SetIsInteractiveSession(false);
 
-	return 0;
+	return zen::ApplicationExitCode();
 }
 
 int
diff --git a/src/zenserver/zenserver.cpp b/src/zenserver/zenserver.cpp
index f9221ed0f..ce5e4e964 100644
--- a/src/zenserver/zenserver.cpp
+++ b/src/zenserver/zenserver.cpp
@@ -130,6 +130,16 @@ namespace zen {
 using namespace std::literals;
 
 namespace utils {
+	static std::atomic_uint32_t SignalCounter[NSIG] = {0};
+
+	static void SignalCallbackHandler(int SigNum)
+	{
+		if (SigNum >= 0 && SigNum < NSIG)
+		{
+			SignalCounter[SigNum].fetch_add(1);
+		}
+	}
+
 #if ZEN_USE_SENTRY
 	class sentry_sink final : public spdlog::sinks::base_sink<spdlog::details::null_mutex>
 	{
@@ -520,17 +530,36 @@ public:
 		EnsureIoRunner();
 	}
 
+	void EnqueueSigIntTimer()
+	{
+		m_SigIntTimer.expires_after(std::chrono::milliseconds(500));
+		m_SigIntTimer.async_wait([this](const asio::error_code&) { CheckSigInt(); });
+		EnsureIoRunner();
+	}
+
 	void CheckStateMarker()
 	{
 		std::filesystem::path StateMarkerPath = m_DataRoot / "state_marker";
 		if (!std::filesystem::exists(StateMarkerPath))
 		{
 			ZEN_ERROR("state marker at {} has been deleted, exiting", StateMarkerPath);
-			RequestExit(0);
+			RequestExit(1);
+			return;
 		}
 		EnqueueStateMarkerTimer();
 	}
 
+	void CheckSigInt()
+	{
+		if (utils::SignalCounter[SIGINT] > 0)
+		{
+			ZEN_INFO("SIGINT triggered (Ctrl+C), exiting");
+			RequestExit(128 + SIGINT);
+			return;
+		}
+		EnqueueSigIntTimer();
+	}
+
 	void CheckOwnerPid()
 	{
 		// Pick up any new "owner" processes
@@ -615,6 +644,7 @@ private:
 	asio::io_context				m_IoContext;
 	asio::steady_timer				m_PidCheckTimer{m_IoContext};
 	asio::steady_timer				m_StateMakerTimer{m_IoContext};
+	asio::steady_timer				m_SigIntTimer{m_IoContext};
 	zen::ProcessMonitor				m_ProcessMonitor;
 	zen::NamedMutex					m_ServerMutex;
 
@@ -686,6 +716,8 @@ ZenServer::OnReady()
 void
 ZenServer::InitializeState(const ZenServerOptions& ServerOptions)
 {
+	EnqueueSigIntTimer();
+
 	// Check root manifest to deal with schema versioning
 
 	bool		WipeState  = false;
@@ -1198,8 +1230,11 @@ ZenEntryPoint::Run()
 			ZEN_INFO("shutdown monitor thread waiting for shutdown signal '{}'", ShutdownEventName);
 			if (ShutdownEvent->Wait())
 			{
-				ZEN_INFO("shutdown signal received");
-				Server.RequestExit(0);
+				if (!IsApplicationExitRequested())
+				{
+					ZEN_INFO("shutdown signal received");
+					Server.RequestExit(0);
+				}
 			}
 			else
 			{
@@ -1235,7 +1270,7 @@ ZenEntryPoint::Run()
 
 	ShutdownLogging();
 
-	return 0;
+	return ApplicationExitCode();
 }
 
 }  // namespace zen
@@ -1306,6 +1341,8 @@ main(int argc, char* argv[])
 	}
 #endif
 
+	signal(SIGINT, utils::SignalCallbackHandler);
+
 	try
 	{
 		ZenServerOptions ServerOptions;