From 56eeca89025cbe1574d76e44b622c3ffc3e34511 Mon Sep 17 00:00:00 2001
From: Erb3 <49862976+Erb3@users.noreply.github.com>
Date: Sat, 3 Aug 2024 13:55:13 +0200
Subject: [PATCH] feat: termination kick cli option

---
 src/cli.rs    |  9 +++++++++
 src/main.rs   |  2 ++
 src/server.rs | 37 ++++++++++++++++++-------------------
 3 files changed, 29 insertions(+), 19 deletions(-)

diff --git a/src/cli.rs b/src/cli.rs
index 7961b58..4f8a931 100644
--- a/src/cli.rs
+++ b/src/cli.rs
@@ -42,6 +42,15 @@ pub(crate) struct Cli {
 	/// guessed. Default is 3s
 	#[arg(long, env = "SVEIO_SHOWCASE_TIME")]
 	pub(crate) showcase_time: Option<u64>,
+
+	/// Optional boolean deciding if players should be removed to the game
+	/// when the server receives a termination signal. With this enabled,
+	/// shutting down the server will take an additional five seconds
+	/// so that clients can load the error message. For debugging
+	/// reasons this is disabled by default, if you are not running in release.
+	/// If you are running in release, this will be enabled by default.
+	#[arg(long, env = "SVEIO_TERMINATION_KICK")]
+	pub(crate) termination_kick: Option<bool>,
 }
 
 pub(crate) fn get_settings() -> Cli {
diff --git a/src/main.rs b/src/main.rs
index 9cb7b40..60661ba 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -25,6 +25,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 			showcase_time: settings.showcase_time.unwrap_or(3),
 		},
 		port: Some(settings.port.unwrap_or(8085)),
+		server_termination_kick: settings.termination_kick.unwrap_or(!cfg!(debug_assertions)),
 	})
 	.await;
 
@@ -42,6 +43,7 @@ async fn main() -> shuttle_axum::ShuttleAxum {
 			showcase_time: 3,
 		},
 		port: None,
+		server_termination_kick: true,
 	})
 	.await
 	.unwrap()
diff --git a/src/server.rs b/src/server.rs
index 9d1548c..0031749 100644
--- a/src/server.rs
+++ b/src/server.rs
@@ -13,6 +13,7 @@ use tracing::info;
 pub(crate) struct ServerOptions {
 	pub(crate) game: game::GameOptions,
 	pub(crate) port: Option<u32>,
+	pub(crate) server_termination_kick: bool,
 }
 
 pub(crate) async fn create_server(opts: ServerOptions) -> Option<axum::Router> {
@@ -67,7 +68,7 @@ pub(crate) async fn create_server(opts: ServerOptions) -> Option<axum::Router> {
 
 		info!("✅ Listening on http://{}", listener.local_addr().unwrap());
 		axum::serve(listener, app)
-			.with_graceful_shutdown(shutdown_signal(shutdown_io))
+			.with_graceful_shutdown(shutdown_signal(shutdown_io, opts.server_termination_kick))
 			.await
 			.unwrap();
 
@@ -77,7 +78,7 @@ pub(crate) async fn create_server(opts: ServerOptions) -> Option<axum::Router> {
 	Some(app)
 }
 
-async fn shutdown_signal(io: Arc<SocketIo>) {
+async fn shutdown_signal(io: Arc<SocketIo>, should_kick: bool) {
 	let ctrl_c = async {
 		signal::ctrl_c()
 			.await
@@ -100,24 +101,22 @@ async fn shutdown_signal(io: Arc<SocketIo>) {
 		_ = terminate => {},
 	}
 
-	info!("Termination signal received, starting graceful shutdown. Exiting soon");
-	for socket in io.sockets().unwrap() {
-		socket
-			.emit(
-				"kick",
-				packets::DisconnectPacket {
-					message: "Server going down".to_string(),
-				},
-			)
-			.unwrap();
-		socket.disconnect().unwrap();
+	if should_kick {
+		info!("Termination signal received, starting graceful shutdown. Exiting soon");
+		for socket in io.sockets().unwrap() {
+			socket
+				.emit(
+					"kick",
+					packets::DisconnectPacket {
+						message: "Server going down".to_string(),
+					},
+				)
+				.unwrap();
+			socket.disconnect().unwrap();
+		}
+
+		tokio::time::sleep(Duration::from_secs(5)).await;
 	}
 
-	tokio::time::sleep(Duration::from_secs(if cfg!(debug_assertions) {
-		0
-	} else {
-		5
-	}))
-	.await;
 	info!("Exit imminent")
 }