From 00b933a90de1799d9f0a8074e427491339e7134a Mon Sep 17 00:00:00 2001
From: seth <getchoo@tuta.io>
Date: Tue, 5 Dec 2023 06:33:54 -0500
Subject: feat: cleanly shutdown on ctrl+c & sigterm

---
 src/main.rs | 31 ++++++++++++++++++++++++++-----
 1 file changed, 26 insertions(+), 5 deletions(-)

(limited to 'src/main.rs')

diff --git a/src/main.rs b/src/main.rs
index b16a13d..06e6724 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,11 +1,16 @@
+use std::sync::Arc;
 use std::time::Duration;
 
 use color_eyre::eyre::{eyre, Context as _, Report, Result};
+use color_eyre::owo_colors::OwoColorize;
 use log::*;
-use poise::serenity_prelude as serenity;
+use poise::serenity_prelude::{self as serenity, ShardManager};
 use poise::{EditTracker, Framework, FrameworkOptions, PrefixFrameworkOptions};
 use redis::ConnectionLike;
 use storage::Storage;
+use tokio::signal::ctrl_c;
+use tokio::signal::unix::{signal, SignalKind};
+use tokio::sync::Mutex;
 
 mod api;
 mod colors;
@@ -63,6 +68,12 @@ async fn setup(
 	Ok(data)
 }
 
+async fn handle_shutdown(shard_manager: Arc<Mutex<ShardManager>>, reason: &str) {
+	warn!("{reason}! Shutting down bot...");
+	shard_manager.lock().await.shutdown_all().await;
+	println!("{}", "Everything is shutdown. Goodbye!".green());
+}
+
 #[tokio::main]
 async fn main() -> Result<()> {
 	dotenvy::dotenv().ok();
@@ -100,12 +111,22 @@ async fn main() -> Result<()> {
 		.token(token)
 		.intents(intents)
 		.options(options)
-		.setup(|ctx, ready, framework| Box::pin(setup(ctx, ready, framework)));
+		.setup(|ctx, ready, framework| Box::pin(setup(ctx, ready, framework)))
+		.build()
+		.await
+		.wrap_err_with(|| "Failed to build framework!")?;
+
+	let shard_manager = framework.shard_manager().clone();
+	let mut sigterm = signal(SignalKind::terminate())?;
 
 	tokio::select! {
-		result = framework.run() => result.map_err(Report::from),
-		_ = tokio::signal::ctrl_c() => {
-			info!("Interrupted! Exiting...");
+		result = framework.start() => result.map_err(Report::from),
+		_ = sigterm.recv() => {
+			handle_shutdown(shard_manager, "Recieved SIGTERM").await;
+			std::process::exit(0);
+		},
+		_ = ctrl_c() => {
+			handle_shutdown(shard_manager, "Interrupted").await;
 			std::process::exit(130);
 		}
 	}
-- 
cgit v1.2.3