From 6e666e5f1a017c362b2dd6f5b77fa914043c8e44 Mon Sep 17 00:00:00 2001 From: Coca Date: Mon, 19 Dec 2022 03:28:23 +0000 Subject: [PATCH] graceful shutdowns --- Cargo.lock | 10 +++++++ Cargo.toml | 2 +- src/commands.rs | 2 +- src/commands/help.rs | 2 +- src/commands/ping.rs | 2 +- src/main.rs | 63 +++++++++++++++++++++++++++++++------------- 6 files changed, 59 insertions(+), 22 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2495e22..d256a7c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1366,6 +1366,15 @@ dependencies = [ "lazy_static", ] +[[package]] +name = "signal-hook-registry" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e51e73328dc4ac0c7ccbda3a494dfa03df1de2f46018127f60c693f2648455b0" +dependencies = [ + "libc", +] + [[package]] name = "slab" version = "0.4.7" @@ -1630,6 +1639,7 @@ dependencies = [ "mio", "num_cpus", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys", diff --git a/Cargo.toml b/Cargo.toml index 1a91b75..1d46b15 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,7 +7,7 @@ edition = "2021" [dependencies] poise = "0.5.2" -tokio = "1.23.0" +tokio = { version = "1.23.0", features = [ "signal" ] } dotenvy = "0.15.6" tracing = "0.1.37" tracing-subscriber = { version = "0.3.16", features = [ "env-filter" ] } diff --git a/src/commands.rs b/src/commands.rs index 9e4026a..11709c0 100644 --- a/src/commands.rs +++ b/src/commands.rs @@ -1,2 +1,2 @@ pub mod help; -pub mod ping; \ No newline at end of file +pub mod ping; diff --git a/src/commands/help.rs b/src/commands/help.rs index 44d62ab..9f05e06 100644 --- a/src/commands/help.rs +++ b/src/commands/help.rs @@ -14,4 +14,4 @@ You can edit most command messages and the bot will edit its response.", }; poise::builtins::help(ctx, command.as_deref(), config).await?; Ok(()) -} \ No newline at end of file +} diff --git a/src/commands/ping.rs b/src/commands/ping.rs index 0311c17..638bf7d 100644 --- a/src/commands/ping.rs +++ b/src/commands/ping.rs @@ -5,4 +5,4 @@ use crate::{Context, Error}; pub async fn pong(ctx: Context<'_>) -> Result<(), Error> { ctx.say("pong!").await?; Ok(()) -} \ No newline at end of file +} diff --git a/src/main.rs b/src/main.rs index db76a0c..9bfbdc9 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,9 +1,9 @@ use std::env; use std::env::VarError; +use commands::{help::help, ping::pong}; use dotenvy::dotenv; -use sqlx::{PgPool, postgres::PgPoolOptions}; -use commands::{ping::pong, help::help}; +use sqlx::{postgres::PgPoolOptions, PgPool}; use poise::{serenity_prelude as serenity, Prefix}; use serenity::GatewayIntents; @@ -11,14 +11,15 @@ use serenity::GatewayIntents; mod commands; // You might want to change this to include more privileged intents or to make it not be so broad -const INTENTS: GatewayIntents = GatewayIntents::non_privileged().union(serenity::GatewayIntents::MESSAGE_CONTENT); +const INTENTS: GatewayIntents = + GatewayIntents::non_privileged().union(serenity::GatewayIntents::MESSAGE_CONTENT); pub type Context<'a> = poise::Context<'a, Data, Error>; pub type Error = Box; // Data shared across commands and events pub struct Data { - pub db: PgPool + pub db: PgPool, } #[tokio::main] @@ -39,7 +40,8 @@ async fn main() { // These are done at runtime so changes can be made when running the bot without the need of a recompilation let token = env::var("DISCORD_TOKEN").expect("No discord token found in environment variables"); - let database_url = env::var("DATABASE_URL").expect("No database url found in environment variables"); + let database_url = + env::var("DATABASE_URL").expect("No database url found in environment variables"); let (primary_prefix, addition_prefixes) = parse_prefixes(); // Setting up database connections @@ -54,14 +56,16 @@ async fn main() { .await .expect("Unable to apply migrations!"); - let data = Data { db }; + let data = Data { db: db.clone() }; - let framework = poise::Framework::builder() + let framework_builder = poise::Framework::builder() .options(poise::FrameworkOptions { prefix_options: poise::PrefixFrameworkOptions { prefix: primary_prefix, additional_prefixes: addition_prefixes, - edit_tracker: Some(poise::EditTracker::for_timespan(std::time::Duration::from_secs(120))), + edit_tracker: Some(poise::EditTracker::for_timespan( + std::time::Duration::from_secs(120), + )), ..Default::default() }, commands, @@ -69,15 +73,34 @@ async fn main() { }) .token(token) .intents(INTENTS) - .setup(|ctx, _ready, framework| - Box::pin(async move { + .setup(|ctx, _ready, framework| { + Box::pin(async move { poise::builtins::register_globally(ctx, &framework.options().commands).await?; - Ok(data) - } - )); + Ok(data) + }) + }); - framework.run().await.unwrap(); + // Build the framework + let framework = framework_builder + .build() + .await + .expect("Cannot build the bot framework!"); + + // ctrl+c handler for graceful shutdowns + let shard_handler = framework.shard_manager().clone(); + tokio::spawn(async move { + tokio::signal::ctrl_c() + .await + .expect("Cannot register a ctrl+c handler!"); + + tracing::info!("Shutting down the bot!"); + shard_handler.lock().await.shutdown_all().await; + db.close().await; + }); + + tracing::info!("Starting the bot!"); + framework.start().await.unwrap(); } fn not_using_dotenv() -> bool { @@ -96,15 +119,19 @@ fn parse_prefixes() -> (Option, Vec) { Ok(unparsed) => unparsed, // The defaults for prefix & additional_prefixes is these Err(VarError::NotPresent) => return (None, Vec::new()), - _ => panic!("Could not handle the environment variable for prefixes") + _ => panic!("Could not handle the environment variable for prefixes"), }; let mut split = unparsed.split(' ').map(|x| x.to_string()); - let first = split.next().expect("Could not parse prefixes from environment variables"); + let first = split + .next() + .expect("Could not parse prefixes from environment variables"); // We need to leak these strings since `Prefix::Literal` only accepts `&'static str` for some reason - let split = split.map(|x| Box::leak(Box::new(x))).map(|x| Prefix::Literal(x)); + let split = split + .map(|x| Box::leak(Box::new(x))) + .map(|x| Prefix::Literal(x)); (Some(first), split.collect()) -} \ No newline at end of file +}