From 9ae8b98e1aa44d8d51419f05ea22c02ced54550a Mon Sep 17 00:00:00 2001 From: seth Date: Wed, 13 Dec 2023 11:40:39 -0500 Subject: [PATCH] feat: cleanly shutdown on ctrl+c & sigterm --- src/main.rs | 34 ++++++++++++++++++++++++++++------ 1 file changed, 28 insertions(+), 6 deletions(-) diff --git a/src/main.rs b/src/main.rs index a90ea10..92544fa 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,12 +1,18 @@ -use std::{sync::Arc, time::Duration}; +use std::sync::Arc; +use std::time::Duration; use color_eyre::eyre::{Context as _, Report, Result}; +use color_eyre::owo_colors::OwoColorize; use config::Config; use log::*; use poise::{ - serenity_prelude as serenity, EditTracker, Framework, FrameworkOptions, PrefixFrameworkOptions, + serenity_prelude::{self as serenity, ShardManager}, + EditTracker, Framework, FrameworkOptions, PrefixFrameworkOptions, }; use storage::Storage; +use tokio::signal::ctrl_c; +use tokio::signal::unix::{signal, SignalKind}; +use tokio::sync::Mutex; mod api; mod commands; @@ -53,6 +59,12 @@ async fn setup( Ok(data) } +async fn handle_shutdown(shard_manager: Arc>, 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(); @@ -91,12 +103,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); } }