capture signal handling in a function

This commit is contained in:
2024-02-16 14:47:32 +01:00
parent bd33f22b30
commit 5560d31638
2 changed files with 25 additions and 21 deletions

View File

@@ -2,17 +2,15 @@ mod handler;
use handler::Handler; use handler::Handler;
mod commands; mod commands;
mod metadata; mod metadata;
mod signal;
mod state; mod state;
use dotenv::dotenv; use dotenv::dotenv;
use futures::StreamExt; use futures::StreamExt;
use signal::signal_handler;
use songbird::{shards::TwilightMap, Songbird}; use songbird::{shards::TwilightMap, Songbird};
use state::StateRef; use state::StateRef;
use std::{env, error::Error, sync::Arc}; use std::{env, error::Error, sync::Arc};
use tokio::{ use tokio::select;
select,
signal::unix::{signal, SignalKind},
sync::watch,
};
use tracing::{debug, info}; use tracing::{debug, info};
use twilight_cache_inmemory::InMemoryCache; use twilight_cache_inmemory::InMemoryCache;
use twilight_gateway::{ use twilight_gateway::{
@@ -32,20 +30,6 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
// Initialize the tracing subscriber. // Initialize the tracing subscriber.
tracing_subscriber::fmt::init(); tracing_subscriber::fmt::init();
let (stop_tx, mut stop_rx) = watch::channel(());
tokio::spawn(async move {
let mut sigterm = signal(SignalKind::terminate()).unwrap();
let mut sigint = signal(SignalKind::interrupt()).unwrap();
loop {
select! {
_ = sigterm.recv() => println!("Receive SIGTERM"),
_ = sigint.recv() => println!("Receive SIGTERM"),
};
stop_tx.send(()).unwrap();
}
});
let (mut shards, state) = { let (mut shards, state) = {
let token = env::var("DISCORD_TOKEN")?; let token = env::var("DISCORD_TOKEN")?;
let app_id = env::var("DISCORD_APP_ID")?.parse()?; let app_id = env::var("DISCORD_APP_ID")?.parse()?;
@@ -91,6 +75,7 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
}; };
let mut handler = Handler::new(Arc::clone(&state)); let mut handler = Handler::new(Arc::clone(&state));
let mut stop_rx = signal_handler();
let mut stream = ShardEventStream::new(shards.iter_mut()); let mut stream = ShardEventStream::new(shards.iter_mut());
loop { loop {
select! { select! {
@@ -109,11 +94,9 @@ async fn main() -> Result<(), Box<dyn Error + Send + Sync + 'static>> {
Some((_, Ok(event))) => event, Some((_, Ok(event))) => event,
Some((_, Err(source))) => { Some((_, Err(source))) => {
tracing::warn!(?source, "error receiving event"); tracing::warn!(?source, "error receiving event");
if source.is_fatal() { if source.is_fatal() {
break; break;
} }
continue; continue;
} }
None => break, None => break,

21
src/signal.rs Normal file
View File

@@ -0,0 +1,21 @@
use tokio::{
select,
signal::unix::{signal, SignalKind},
sync::watch,
};
pub(crate) fn signal_handler() -> watch::Receiver<()> {
let (stop_tx, stop_rx) = watch::channel(());
tokio::spawn(async move {
let mut sigterm = signal(SignalKind::terminate()).unwrap();
let mut sigint = signal(SignalKind::interrupt()).unwrap();
loop {
select! {
_ = sigterm.recv() => println!("Receive SIGTERM"),
_ = sigint.recv() => println!("Receive SIGTERM"),
};
stop_tx.send(()).unwrap();
}
});
stop_rx
}