diff --git a/Cargo.lock b/Cargo.lock index 28413c71..0810cc08 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2625,7 +2625,9 @@ dependencies = [ "k8s-openapi", "kube", "logging", + "mockall", "serenity", + "serenity_utils", "tokio", "tracing", ] @@ -2890,6 +2892,7 @@ dependencies = [ "mockall", "ordered-float 5.0.0", "regex-lite", + "restarter", "serde", "serde_json", "serenity", @@ -3083,6 +3086,14 @@ dependencies = [ "serde_repr", ] +[[package]] +name = "serenity_utils" +version = "0.1.0" +dependencies = [ + "serenity", + "tracing", +] + [[package]] name = "sha1" version = "0.10.6" diff --git a/Cargo.toml b/Cargo.toml index 1d861ab9..af21a42f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,11 +1,24 @@ [workspace] -members = ["crates/database", "crates/logging", "crates/soundboard", "crates/voicevox", "restarter", "seitai"] +members = [ + "crates/database", + "crates/logging", + "crates/restarter", + "crates/serenity-utils", + "crates/soundboard", + "crates/voicevox", + "seitai", +] default-members = ["seitai"] resolver = "3" [workspace.dependencies.anyhow] version = "1.0.96" +[workspace.dependencies.futures] +version = "0.3.31" +default-features =false +features = ["std"] + [workspace.dependencies.http-body-util] version = "0.1.2" @@ -20,6 +33,9 @@ features = ["client", "client-legacy", "http1", "tokio"] [workspace.dependencies.logging] path = "crates/logging" +[workspace.dependencies.restarter] +path = "crates/restarter" + [workspace.dependencies.serde] version = "1.0.218" features = ["derive"] @@ -31,6 +47,9 @@ version = "1.0.139" version = "0.12.4" default-features = false +[workspace.dependencies.serenity_utils] +path = "crates/serenity-utils" + [workspace.dependencies.tokio] version = "1.43.0" features = ["macros", "net", "rt-multi-thread", "signal"] diff --git a/Dockerfile b/Dockerfile index 95ec207b..2eb50b2e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -13,28 +13,13 @@ FROM runtime AS development FROM runtime AS builder RUN --mount=type=bind,source=crates,target=crates \ - --mount=type=bind,source=restarter,target=restarter \ --mount=type=bind,source=seitai,target=seitai \ --mount=type=bind,source=Cargo.toml,target=Cargo.toml \ --mount=type=bind,source=Cargo.lock,target=Cargo.lock \ --mount=type=cache,target=/usr/src/myapp/target \ cargo build --release --workspace \ - && cp target/release/restarter /restarter \ && cp target/release/seitai /seitai -FROM scratch AS restarter -LABEL io.github.hexium310.seitai.app=restarter -LABEL org.opencontainers.image.source=https://github.com/hexium310/seitai -COPY --from=runtime /etc/ssl/certs/ /etc/ssl/certs/ -COPY --from=runtime /lib/x86_64-linux-gnu/libc.so* /lib/x86_64-linux-gnu/ -COPY --from=runtime /lib/x86_64-linux-gnu/libcrypto.so* /lib/x86_64-linux-gnu/ -COPY --from=runtime /lib/x86_64-linux-gnu/libgcc_s.so* /lib/x86_64-linux-gnu/ -COPY --from=runtime /lib/x86_64-linux-gnu/libm.so* /lib/x86_64-linux-gnu/ -COPY --from=runtime /lib/x86_64-linux-gnu/libssl.so* /lib/x86_64-linux-gnu/ -COPY --from=runtime /lib64/ld-linux-x86-64.so* /lib64/ -COPY --from=builder /restarter / -CMD ["/restarter"] - FROM scratch AS seitai LABEL io.github.hexium310.seitai.app=seitai LABEL org.opencontainers.image.source=https://github.com/hexium310/seitai diff --git a/compose.yaml b/compose.yaml index a70829dd..89ae43f5 100644 --- a/compose.yaml +++ b/compose.yaml @@ -22,9 +22,6 @@ services: - type: bind source: seitai target: /usr/src/myapp/seitai - - type: bind - source: restarter - target: /usr/src/myapp/restarter command: /bin/sh -c 'cargo run -p seitai' environment: DISCORD_TOKEN: @@ -53,10 +50,7 @@ services: - type: bind source: seitai target: /usr/src/myapp/seitai - - type: bind - source: restarter - target: /usr/src/myapp/restarter - command: /bin/sh -c 'cargo run -p restarter' + command: /bin/sh -c 'cargo run -- restarter' environment: DISCORD_TOKEN: diff --git a/restarter/Cargo.toml b/crates/restarter/Cargo.toml similarity index 70% rename from restarter/Cargo.toml rename to crates/restarter/Cargo.toml index 78b44317..4c338708 100644 --- a/restarter/Cargo.toml +++ b/crates/restarter/Cargo.toml @@ -3,11 +3,16 @@ name = "restarter" version = "0.1.7" edition = "2024" +[lib] +doctest = false + +[dependencies] + [dependencies.anyhow] workspace = true [dependencies.futures] -version = "0.3.31" +workspace = true [dependencies.k8s-openapi] version = "0.24.0" @@ -30,3 +35,13 @@ workspace = true [dependencies.serenity] workspace = true features = ["cache", "client", "gateway", "model", "native_tls_backend"] + +[dependencies.serenity_utils] +workspace = true + +[dev-dependencies.mockall] +version = "0.13.1" + +[dev-dependencies.tokio] +workspace = true +features = ["test-util"] diff --git a/crates/restarter/src/client.rs b/crates/restarter/src/client.rs new file mode 100644 index 00000000..5d44e26c --- /dev/null +++ b/crates/restarter/src/client.rs @@ -0,0 +1,51 @@ +use std::{collections::HashMap, sync::Arc, time::Duration}; + +use anyhow::Result; +use futures::lock::Mutex; +use serenity::{all::GatewayIntents, Client as SerenityClient}; +use tokio::{signal::unix::{self, SignalKind}, task::JoinHandle}; + +use crate::{event_handler::Handler, restarter::{KubeRestarter, Restarter}}; + +pub struct Client; + +impl Client { + #[tracing::instrument(skip_all)] + pub async fn start(token: String, restart_duration: u64) -> Result<()> { + enable_graceful_shutdown(); + + let intents = GatewayIntents::GUILDS | GatewayIntents::GUILD_VOICE_STATES; + let mut client = SerenityClient::builder(token, intents) + .event_handler(Handler { + connected_channels: Arc::new(Mutex::new(HashMap::new())), + restarter: Restarter::new(Duration::from_secs(restart_duration), KubeRestarter), + }) + .await?; + + tracing::debug!("restarter client starts"); + if let Err(err) = client.start().await { + tracing::error!("failed to start client\nError: {err:?}"); + return Err(err.into()); + } + + Ok(()) + } +} + +fn enable_graceful_shutdown() -> JoinHandle> { + tokio::spawn(async move { + let mut sigint = unix::signal(SignalKind::interrupt())?; + let mut sigterm = unix::signal(SignalKind::terminate())?; + + tokio::select! { + _ = sigint.recv() => { + tracing::info!("received SIGINT, shutting down"); + std::process::exit(130); + }, + _ = sigterm.recv() => { + tracing::info!("received SIGTERM, shutting down"); + std::process::exit(143); + }, + } + }) +} diff --git a/crates/restarter/src/event_handler.rs b/crates/restarter/src/event_handler.rs new file mode 100644 index 00000000..00bb5f7b --- /dev/null +++ b/crates/restarter/src/event_handler.rs @@ -0,0 +1,52 @@ +use std::{collections::HashMap, pin::Pin, sync::Arc}; + +use futures::{lock::Mutex, FutureExt}; +use serenity::{ + all::{ChannelId, GuildId, VoiceState}, + client::{Context, EventHandler}, + model::gateway::Ready, +}; +use tracing::instrument; + +use crate::{event_handler, restarter::Restarter}; + +mod ready; +mod voice_state_update; + +pub(crate) struct Handler { + pub(crate) connected_channels: Arc>>, + pub(crate) restarter: Restarter, +} + +impl EventHandler for Handler { + #[instrument(skip_all)] + fn ready<'s, 'async_trait>(&'s self, ctx: Context, ready: Ready) -> Pin + Send + 'async_trait)>> + where + Self: 'async_trait, + 's: 'async_trait, + { + async move { + if let Err(err) = event_handler::ready::handle(self, ctx, ready).await { + tracing::error!("failed to handle ready event\nError: {err:?}"); + } + }.boxed() + } + + #[instrument(skip_all)] + fn voice_state_update<'s, 'async_trait>( + &'s self, + context: Context, + old_state: Option, + new_state: VoiceState, + ) -> Pin + Send + 'async_trait)>> + where + Self: 'async_trait, + 's: 'async_trait, + { + async move { + if let Err(err) = event_handler::voice_state_update::handle(self, context, old_state, new_state).await { + tracing::error!("failed to handle voice state update event\nError: {err:?}"); + } + }.boxed() + } +} diff --git a/crates/restarter/src/event_handler/ready.rs b/crates/restarter/src/event_handler/ready.rs new file mode 100644 index 00000000..f887f4c1 --- /dev/null +++ b/crates/restarter/src/event_handler/ready.rs @@ -0,0 +1,10 @@ +use anyhow::Result; +use serenity::all::{Context, Ready}; + +use super::Handler; + +pub(crate) async fn handle(_handler: &Handler, _ctx: Context, ready: Ready) -> Result<()> { + tracing::info!("{} is ready", ready.user.name); + + Ok(()) +} diff --git a/crates/restarter/src/event_handler/voice_state_update.rs b/crates/restarter/src/event_handler/voice_state_update.rs new file mode 100644 index 00000000..864a95ed --- /dev/null +++ b/crates/restarter/src/event_handler/voice_state_update.rs @@ -0,0 +1,41 @@ +use anyhow::Result; +use serenity::all::{Context, VoiceState}; +use serenity_utils::voice_state::{VoiceStateAction, VoiceStateConnection}; + +use super::Handler; + +pub(crate) async fn handle(handler: &Handler, ctx: Context, old_state: Option, new_state: VoiceState) -> Result<()> { + let Some(guild_id) = new_state.guild_id else { + return Ok(()); + }; + + let bot_id = ctx.http.get_current_user().await?.id; + let action = VoiceStateAction::new(old_state, new_state); + + if !action.is_bot_action(bot_id) { + return Ok(()); + } + + match action.connection() { + VoiceStateConnection::Joined(channel_id) => { + let mut connected_channels = handler.connected_channels.lock().await; + connected_channels.insert(guild_id, channel_id); + + handler.restarter.set_connection_count(connected_channels.len()).await; + }, + VoiceStateConnection::Left(_) => { + let mut connected_channels = handler.connected_channels.lock().await; + connected_channels.remove(&guild_id); + + let count = connected_channels.len(); + handler.restarter.set_connection_count(count).await; + + if count == 0 { + handler.restarter.schedule_restart().await; + } + }, + VoiceStateConnection::Moved(..) | VoiceStateConnection::NoAction => (), + } + + Ok(()) +} diff --git a/crates/restarter/src/lib.rs b/crates/restarter/src/lib.rs new file mode 100644 index 00000000..d099f828 --- /dev/null +++ b/crates/restarter/src/lib.rs @@ -0,0 +1,5 @@ +pub use client::Client; + +mod event_handler; +mod client; +mod restarter; diff --git a/crates/restarter/src/restarter.rs b/crates/restarter/src/restarter.rs new file mode 100644 index 00000000..68bec921 --- /dev/null +++ b/crates/restarter/src/restarter.rs @@ -0,0 +1,105 @@ +use std::{fmt::Debug, sync::Arc, time::Duration}; + +use anyhow::Result; +use futures::lock::Mutex; +use k8s_openapi::api::apps::v1::StatefulSet; +use kube::{Api, Client, Config}; +use tokio::{sync::oneshot, task::JoinHandle}; +use tracing::Instrument; + +#[derive(Debug, Clone)] +pub(crate) struct Restarter { + duration: Duration, + waiting: Arc>, + restart: Restart, + connection_count: Arc>, +} + +#[derive(Debug, Clone)] +pub(crate) struct KubeRestarter; + +pub(crate) trait Restart: Send { + fn restart(&self) -> impl Future> + Send; +} + +impl Restart for KubeRestarter { + #[tracing::instrument] + async fn restart(&self) -> Result<()> { + let config = match Config::incluster() { + Ok(config) => config, + Err(_) => { + tracing::warn!("This app doesn't running in Kubernetes, so it didn't restart voicevox."); + return Ok(()) + }, + }; + let client = Client::try_from(config)?; + let stateful_sets: Api = Api::default_namespaced(client); + + stateful_sets.restart("voicevox").await?; + + tracing::info!("succeeded in restarting statefulsets/voicevox"); + + Ok(()) + } +} + +impl Restarter { + pub(crate) fn new(duration: Duration, restart: Restart) -> Self { + Self { + duration, + waiting: Arc::new(Mutex::new(false)), + restart, + connection_count: Arc::new(Mutex::new(0)), + } + } + + #[allow(clippy::async_yields_async)] + #[tracing::instrument(skip(self), fields(self.duration = ?self.duration, self.restart = ?self.restart))] + pub(crate) async fn schedule_restart(&self) -> JoinHandle<()> { + let waiting = self.waiting.clone(); + + if *waiting.lock().await { + return tokio::spawn(async {}); + } + + let duration = self.duration; + let restart = self.restart.clone(); + let connection_count = self.connection_count.clone(); + + tracing::info!("statefulsets/voicevox is going to restart in {} secs", duration.as_secs()); + + let (tx, rx) = oneshot::channel(); + + let handle = tokio::spawn(async move { + tx.send(()).expect("failed to send notice that timer task spawned"); + *waiting.lock().await = true; + + tokio::time::sleep(duration).await; + + *waiting.lock().await = false; + + { + let connection_count = *connection_count.lock().await; + if connection_count != 0 { + tracing::info!("statefulsets/voicevox wasn't restart becase {connection_count} clients are connected"); + + return; + } + } + + if let Err(err) = restart.restart().await { + tracing::error!("failed to restart statefulsets/voicevox\nError: {err:?}"); + } + }.in_current_span()); + + rx.await.expect("failed to receive notice that timer task spawned"); + handle + } + + pub(crate) async fn set_connection_count(&self, count: usize) { + *self.connection_count.lock().await = count; + } +} + +#[cfg(test)] +mod tests; diff --git a/crates/restarter/src/restarter/tests.rs b/crates/restarter/src/restarter/tests.rs new file mode 100644 index 00000000..3dbb7617 --- /dev/null +++ b/crates/restarter/src/restarter/tests.rs @@ -0,0 +1,66 @@ +use std::time::Duration; + +use anyhow::Result; +use futures::FutureExt; + +use crate::restarter::{Restart, Restarter}; + +mockall::mock! { + #[derive(Debug)] + Restart {} + impl Clone for Restart { + fn clone(&self) -> Self; + } + impl Restart for Restart { + fn restart(&self) -> impl Future> + Send; + } +} + +#[tokio::test(start_paused = true)] +async fn restart_as_scheduled() { + let mut mock_restart = MockRestart::new(); + mock_restart + .expect_clone() + .once() + .return_once(move || { + let mut mock_restart = MockRestart::new(); + mock_restart + .expect_restart() + .once() + .returning(|| async { Ok(()) }.boxed()); + + mock_restart + }); + + let restarter = Restarter::new(Duration::from_secs(300), mock_restart); + + let wait = restarter.schedule_restart().await; + + restarter.set_connection_count(0).await; + + assert!(wait.await.is_ok()); +} + +#[tokio::test(start_paused = true)] +async fn cancel_restart() { + let mut mock_restart = MockRestart::new(); + mock_restart + .expect_clone() + .once() + .return_once(move || { + let mut mock_restart = MockRestart::new(); + mock_restart + .expect_restart() + .never(); + + mock_restart + }); + + let restarter = Restarter::new(Duration::from_secs(300), mock_restart); + + let wait = restarter.schedule_restart().await; + + restarter.set_connection_count(1).await; + + assert!(wait.await.is_ok()); +} diff --git a/crates/serenity-utils/Cargo.toml b/crates/serenity-utils/Cargo.toml new file mode 100644 index 00000000..8ec088f6 --- /dev/null +++ b/crates/serenity-utils/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "serenity_utils" +version = "0.1.0" +edition = "2024" + +[dependencies.serenity] +workspace = true + +[dependencies.tracing] +workspace = true diff --git a/crates/serenity-utils/src/lib.rs b/crates/serenity-utils/src/lib.rs new file mode 100644 index 00000000..cad7be11 --- /dev/null +++ b/crates/serenity-utils/src/lib.rs @@ -0,0 +1 @@ +pub mod voice_state; diff --git a/crates/serenity-utils/src/voice_state.rs b/crates/serenity-utils/src/voice_state.rs new file mode 100644 index 00000000..8040c9a6 --- /dev/null +++ b/crates/serenity-utils/src/voice_state.rs @@ -0,0 +1,50 @@ +use serenity::all::{ChannelId, UserId, VoiceState}; + +#[derive(Debug)] +pub struct VoiceStateAction { + old_state: Option, + new_state: VoiceState, +} + +pub enum VoiceStateConnection { + Joined(ChannelId), + Left(ChannelId), + Moved(ChannelId, ChannelId), + NoAction, +} + +impl VoiceStateAction { + pub fn new(old_state: Option, new_state: VoiceState) -> Self { + Self { + old_state, + new_state, + } + } + + pub fn connection(&self) -> VoiceStateConnection { + match &self.old_state { + Some(old_state) => match (old_state.channel_id, self.new_state.channel_id) { + (Some(old_channel_id), Some(new_channel_id)) if old_channel_id != new_channel_id => { + tracing::debug!("moved voice channel from {old_channel_id} to {new_channel_id}"); + VoiceStateConnection::Moved(old_channel_id, new_channel_id) + }, + (Some(channel_id), None) => { + tracing::debug!("left voice channel from {channel_id}"); + VoiceStateConnection::Left(channel_id) + }, + _ => VoiceStateConnection::NoAction, + }, + None => match self.new_state.channel_id { + Some(channel_id) => { + tracing::debug!("joined voice channel to {channel_id}"); + VoiceStateConnection::Joined(channel_id) + }, + None => VoiceStateConnection::NoAction, + }, + } + } + + pub fn is_bot_action(&self, bot_id: UserId) -> bool { + self.new_state.user_id == bot_id + } +} diff --git a/docker-bake.hcl b/docker-bake.hcl index 88fb6260..0746266a 100644 --- a/docker-bake.hcl +++ b/docker-bake.hcl @@ -1,19 +1,11 @@ group "default" { - targets = ["restarter", "seitai"] + targets = ["seitai"] } variable "VERSION" { default = "latest" } -target "restarter" { - target = "restarter" - tags = [ - "ghcr.io/hexium310/restarter:latest", - "ghcr.io/hexium310/restarter:${VERSION}", - ] -} - target "seitai" { target = "seitai" tags = [ diff --git a/manifests/statefulset-seitai.yaml b/manifests/statefulset-seitai.yaml index 1d817aa0..ece3b4b5 100644 --- a/manifests/statefulset-seitai.yaml +++ b/manifests/statefulset-seitai.yaml @@ -16,8 +16,9 @@ spec: - name: seitai-migration image: ghcr.io/hexium310/seitai imagePullPolicy: Always - args: + command: - /seitai + args: - migration - apply env: @@ -64,8 +65,12 @@ spec: name: seitai.seitai-database.credentials.postgresql.acid.zalan.do key: password - name: restarter - image: ghcr.io/hexium310/restarter + image: ghcr.io/hexium310/seitai imagePullPolicy: Always + command: + - /seitai + args: + - restarter env: - name: NO_COLOR value: '1' diff --git a/restarter/src/event_handler.rs b/restarter/src/event_handler.rs deleted file mode 100644 index c3d2489f..00000000 --- a/restarter/src/event_handler.rs +++ /dev/null @@ -1,121 +0,0 @@ -use std::{pin::Pin, sync::Arc}; - -use anyhow::{bail, Result}; -use futures::lock::Mutex; -use k8s_openapi::api::apps::v1::StatefulSet; -use kube::{Api, Client, Config}; -use serenity::{ - all::VoiceState, - client::{Context, EventHandler}, - model::gateway::Ready, -}; -use tokio::sync::Notify; -use tracing::instrument; - -use crate::Data; - -pub struct Handler; - -impl EventHandler for Handler { - #[instrument(skip(self, context))] - fn ready<'s, 'async_trait>( - &'s self, - context: Context, - ready: Ready, - ) -> Pin + Send + 'async_trait)>> - where - Self: 'async_trait, - 's: 'async_trait, - { - tracing::info!("{} is ready", ready.user.name); - - Box::pin(async move { - let Some(data) = get_data(&context).await else { - tracing::error!("failed to get data"); - return; - }; - data.lock().await.bot_id = ready.user.id; - }) - } - - fn voice_state_update<'s, 'async_trait>( - &'s self, - context: Context, - _old: Option, - new: VoiceState, - ) -> Pin + Send + 'async_trait)>> - where - Self: 'async_trait, - 's: 'async_trait, - { - Box::pin(async move { - let Some(data) = get_data(&context).await else { - tracing::error!("failed to get data"); - return; - }; - let mut data = data.lock().await; - let bot_id = data.bot_id; - - if new.user_id == bot_id { - // When bot joined a voice channel - if let (Some(channel_id), Some(guild_id)) = (new.channel_id, new.guild_id) { - data.connected_channels.insert(guild_id, channel_id); - data.cancellation.notify_one(); - // When bot left a voice channel - } else if let (None, Some(guild_id)) = (new.channel_id, new.guild_id) { - data.connected_channels.remove(&guild_id); - if data.connected_channels.is_empty() { - wait_restart(&context).await; - } - }; - } - }) - } -} - -async fn get_data(context: &Context) -> Option>> { - let data = context.data.read().await; - data.get::().cloned() -} - -async fn wait_restart(context: &Context) { - let Some(data) = get_data(context).await else { - tracing::error!("failed to get data"); - return; - }; - - tokio::spawn(async move { - let cancellation = { - let mut data = data.lock().await; - data.cancellation = Arc::new(Notify::new()); - data.cancellation.clone() - }; - - tokio::select! { - _ = tokio::time::sleep(tokio::time::Duration::from_secs(300)) => { - if let Err(error) = restart().await { - tracing::error!("failed to restart statefulsets/voicevox\nError: {error:?}"); - return; - } - tracing::info!("succeeded in restarting statefulsets/voicevox"); - }, - _ = cancellation.notified() => { - tracing::info!("canceled restarting statefulsets/voicevox"); - }, - } - }); -} - -async fn restart() -> Result<()> { - let config = match Config::incluster() { - Ok(config) => config, - Err(_) => { - bail!("this app is not running in cluster of Kubernetes"); - }, - }; - let client = Client::try_from(config)?; - let stateful_sets: Api = Api::default_namespaced(client); - stateful_sets.restart("voicevox").await?; - - Ok(()) -} diff --git a/restarter/src/main.rs b/restarter/src/main.rs deleted file mode 100644 index 25670ef0..00000000 --- a/restarter/src/main.rs +++ /dev/null @@ -1,82 +0,0 @@ -use std::{collections::HashMap, env, process::exit, sync::Arc}; - -use futures::lock::Mutex; -use logging::initialize_logging; -use serenity::{ - all::{ChannelId, GuildId, UserId}, - client::Client, - model::gateway::GatewayIntents, - prelude::TypeMapKey, -}; -use tokio::{ - signal::unix::{signal, SignalKind}, - sync::Notify, -}; - -mod event_handler; - -struct Data { - bot_id: UserId, - connected_channels: HashMap, - cancellation: Arc, -} - -impl TypeMapKey for Data { - type Value = Arc>; -} - -#[tokio::main] -async fn main() { - initialize_logging(); - - let token = match env::var("DISCORD_TOKEN") { - Ok(token) => token, - Err(error) => { - tracing::error!("failed to fetch environment variable DISCORD_TOKEN\nError: {error:?}"); - exit(1); - }, - }; - - let intents = GatewayIntents::non_privileged() | GatewayIntents::MESSAGE_CONTENT; - let mut client = match Client::builder(token, intents) - .event_handler(event_handler::Handler) - .await - { - Ok(client) => client, - Err(error) => { - tracing::error!("failed to build serenity client\nError: {error:?}"); - exit(1); - }, - }; - - { - let mut data = client.data.write().await; - - data.insert::(Arc::new(Mutex::new(Data { - bot_id: UserId::default(), - connected_channels: HashMap::new(), - cancellation: Arc::new(Notify::default()), - }))); - } - - tokio::spawn(async move { - if let Err(error) = client.start().await { - tracing::error!("failed to start client\nError: {error:?}"); - exit(1); - } - }); - - let mut sigint = signal(SignalKind::interrupt()).unwrap(); - let mut sigterm = signal(SignalKind::terminate()).unwrap(); - - tokio::select! { - _ = sigint.recv() => { - tracing::info!("received SIGINT, shutting down"); - exit(130); - }, - _ = sigterm.recv() => { - tracing::info!("received SIGTERM, shutting down"); - exit(143); - }, - } -} diff --git a/seitai/Cargo.toml b/seitai/Cargo.toml index afbfb27b..a9000758 100644 --- a/seitai/Cargo.toml +++ b/seitai/Cargo.toml @@ -8,13 +8,13 @@ workspace = true [dependencies.clap] version = "4.5.31" -features = ["derive"] +features = ["derive", "env"] [dependencies.database] path = "../crates/database" [dependencies.futures] -version = "0.3.31" +workspace = true [dependencies.hashbrown] version = "0.15.2" @@ -45,6 +45,9 @@ version = "5.0.0" [dependencies.regex-lite] version = "0.1.6" +[dependencies.restarter] +workspace = true + [dependencies.serde] workspace = true diff --git a/seitai/src/cli.rs b/seitai/src/cli.rs index 08761a2b..697e82b2 100644 --- a/seitai/src/cli.rs +++ b/seitai/src/cli.rs @@ -1,10 +1,11 @@ -use std::process; - use anyhow::Result; use clap::{error::ErrorKind, Parser}; -use database::migrations::{MigrationCommand, Migrator}; +use subcommands::Subcommand; + +use crate::start_bot; -use crate::{set_up_database, start_bot}; +mod args; +mod subcommands; pub struct Application; @@ -16,11 +17,6 @@ pub struct Cli { subcommand: Subcommand, } -#[derive(clap::Subcommand)] -enum Subcommand { - Migration(MigrationCommand) -} - impl Application { pub async fn start() -> Result<()> { let cli = match Cli::try_parse() { @@ -30,18 +26,14 @@ impl Application { return Ok(()); }, Err(help) => { - println!("{help}"); + help.print()?; return Ok(()); }, }; - let migrator = Migrator::new(); - let pgpool = set_up_database().await?; - - #[allow(irrefutable_let_patterns)] - if let Subcommand::Migration(migration) = cli.subcommand { - migration.run(&mut *pgpool.acquire().await?, migrator.into_boxed_inner()).await?; - process::exit(0); + match cli.subcommand { + Subcommand::Migration(migration) => migration.run().await?, + Subcommand::Restarter(restarter) => restarter.run().await?, } Ok(()) diff --git a/seitai/src/cli/args.rs b/seitai/src/cli/args.rs new file mode 100644 index 00000000..3228a62a --- /dev/null +++ b/seitai/src/cli/args.rs @@ -0,0 +1 @@ +pub mod discord; diff --git a/seitai/src/cli/args/discord.rs b/seitai/src/cli/args/discord.rs new file mode 100644 index 00000000..b1e4fd96 --- /dev/null +++ b/seitai/src/cli/args/discord.rs @@ -0,0 +1,7 @@ +use clap::Parser; + +#[derive(Debug, Parser)] +pub struct DiscordConfig { + #[arg(long, env, hide_env_values = true)] + pub discord_token: String, +} diff --git a/seitai/src/cli/subcommands.rs b/seitai/src/cli/subcommands.rs new file mode 100644 index 00000000..66ad2a9c --- /dev/null +++ b/seitai/src/cli/subcommands.rs @@ -0,0 +1,11 @@ +use migration::Migration; +use restarter::Restarter; + +mod migration; +mod restarter; + +#[derive(clap::Subcommand)] +pub enum Subcommand { + Migration(Migration), + Restarter(Restarter), +} diff --git a/seitai/src/cli/subcommands/migration.rs b/seitai/src/cli/subcommands/migration.rs new file mode 100644 index 00000000..4d1b74c2 --- /dev/null +++ b/seitai/src/cli/subcommands/migration.rs @@ -0,0 +1,21 @@ +use anyhow::Result; +use clap::Parser; +use database::migrations::{MigrationCommand, Migrator}; + +use crate::set_up_database; + +#[derive(Debug, Parser)] +pub struct Migration { + #[command(flatten)] + pub command: MigrationCommand, +} + +impl Migration { + pub async fn run(&self) -> Result<()> { + let migrator = Migrator::new(); + let pgpool = set_up_database().await?; + self.command.run(&mut *pgpool.acquire().await?, migrator.into_boxed_inner()).await?; + + Ok(()) + } +} diff --git a/seitai/src/cli/subcommands/restarter.rs b/seitai/src/cli/subcommands/restarter.rs new file mode 100644 index 00000000..0290f564 --- /dev/null +++ b/seitai/src/cli/subcommands/restarter.rs @@ -0,0 +1,24 @@ +use anyhow::Result; +use clap::Parser; + +use crate::cli::args::discord::DiscordConfig; + +#[derive(Debug, Parser)] +#[command(about = "Launches client to restart voicevox statefulset")] +pub struct Restarter { + #[command(flatten)] + pub discord_args: DiscordConfig, + + #[arg(long, default_value_t = 300, help = "secs until this client restarts voicevox after leaving from voice channel")] + pub duration: u64, +} + +impl Restarter { + pub async fn run(&self) -> Result<()> { + let token = self.discord_args.discord_token.clone(); + + restarter::Client::start(token, self.duration).await?; + + Ok(()) + } +}