fix: crash and graceful shutdown improvements

Cargo.lock

@@ -1777,11 +1777,13 @@ dependencies = [
  "extension-traits",
  "futures",
  "hound",
+ "humantime",
  "itertools",
  "moka",
  "opendal",
  "opus2",
  "patricia_tree 0.10.1",
+ "rayon",
  "rhai",
  "rustls 0.23.40",
  "secrecy 0.10.3",
@@ -2517,6 +2519,12 @@ version = "1.0.3"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9"
 
+[[package]]
+name = "humantime"
+version = "2.3.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "135b12329e5e3ce057a9f972339ea52bc954fe1e9358ef27f95e89716fbc5424"
+
 [[package]]
 name = "hyper"
 version = "1.9.0"

Cargo.toml

@@ -14,6 +14,7 @@ dashmap = "6.1.0"
 extension-traits = "2.0.2"
 futures = "0.3.32"
 hound = "3.5.1"
+humantime = "2.3.0"
 itertools = "0.14.0"
 moka = { version = "0.12.15", features = ["future"] }
 opendal = { git = "https://github.com/apache/opendal", rev = "ecf840b04afd2be109830b9978ba89759adfee79", features = [
@@ -55,6 +56,7 @@ opendal = { git = "https://github.com/apache/opendal", rev = "ecf840b04afd2be109
 ] }
 opus2 = "0.4.0"
 patricia_tree = "0.10.1"
+rayon = "1.12"
 rhai = { version = "1.23.6", features = ["sync"] }
 rustls = "0.23"
 secrecy = { version = "0.10.3", features = ["serde"] }

build.rs

@@ -1,6 +1,9 @@
 use shadow_rs::{BuildPattern, ShadowBuilder};
 
 fn main() {
+    println!("cargo::rerun-if-changed=bot.capnp");
+    println!("cargo::rerun-if-changed=user.capnp");
+
     capnpc::CompilerCommand::new()
         .file("bot.capnp")
         .file("user.capnp")

src/main.rs

@@ -4,10 +4,19 @@ use fomo_reducer::{
     RecordingManager, RenderManager, State, Storage, UserManager, VCsSender, all_commands, command,
     heat_seek, initialize_vcs, update_vcs,
 };
+use futures::{StreamExt as _, stream::FuturesUnordered};
+use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
 use secrecy::{ExposeSecret, SecretString};
 use snafu::{OptionExt, ResultExt, Snafu};
 use songbird::{Config, Songbird, driver::DecodeConfig, shards::TwilightMap};
-use std::{collections::BTreeMap, fmt::Debug, str::FromStr, sync::Arc, time::Duration};
+use std::{
+    collections::BTreeMap,
+    fmt::{Debug, Display},
+    num::NonZero,
+    str::FromStr,
+    sync::Arc,
+    time::Duration,
+};
 use tokio::{select, signal::ctrl_c, task::JoinSet};
 use tokio_util::{sync::CancellationToken, time::FutureExt as _};
 use tracing::Level;
@@ -15,7 +24,7 @@ use tracing_subscriber::{
     EnvFilter,
     fmt::{format::FmtSpan, writer::MakeWriterExt},
 };
-use twilight_gateway::{Event, EventTypeFlags, Intents, Shard, StreamExt};
+use twilight_gateway::{Event, EventTypeFlags, Intents, Shard, StreamExt as _};
 use twilight_model::{
     application::interaction::InteractionData,
     gateway::{
@@ -68,6 +77,29 @@ fn parse_guild_vc_to_text_channel(
     Ok((guild, voice_channel, text_channel))
 }
 
+#[derive(Clone)]
+struct HumanDuration(Duration);
+
+impl FromStr for HumanDuration {
+    type Err = humantime::DurationError;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        humantime::parse_duration(s).map(Self)
+    }
+}
+
+impl Debug for HumanDuration {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        self.0.fmt(f)
+    }
+}
+
+impl Display for HumanDuration {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}", humantime::format_duration(self.0))
+    }
+}
+
 #[derive(Debug, Parser)]
 struct AppArgs {
     #[arg(long, env)]
@@ -80,7 +112,7 @@ struct AppArgs {
     discord_nickname: Option<Arc<str>>,
 
     #[arg(long, env)]
-    discord_status: Option<Arc<str>>,
+    discord_status: Option<String>,
 
     #[arg(long, env, value_parser = parse_guild_vc_to_text_channel)]
     discord_voice_channel_corresponding_text_channel:
@@ -89,7 +121,7 @@ struct AppArgs {
     #[arg(long, env, default_value_t = AudioChannels::Mono)]
     audio_channels: AudioChannels,
 
-    #[arg(long, env, default_value_t = AudioSampleRate::Hz12000)]
+    #[arg(long, env, default_value_t = AudioSampleRate::Hz24000)]
     audio_sample_rate: AudioSampleRate,
 
     #[arg(long, env)]
@@ -103,6 +135,15 @@ struct AppArgs {
 
     #[arg(long, env)]
     render_data: Storage,
+
+    #[arg(long, env, default_value_t = HumanDuration(Duration::from_secs(120)))]
+    watchdog_warmup: HumanDuration,
+
+    #[arg(long, env, default_value_t = HumanDuration(Duration::from_secs(5)))]
+    watchdog_frequency: HumanDuration,
+
+    #[arg(long, env, default_value_t = 8.try_into().unwrap())]
+    watchdog_channel_size: NonZero<usize>,
 }
 
 #[derive(Parser)]
@@ -166,6 +207,9 @@ async fn main() -> Result<(), MainError> {
         user_data,
         recording_data,
         render_data,
+        watchdog_warmup: HumanDuration(watchdog_warmup),
+        watchdog_frequency: HumanDuration(watchdog_frequency),
+        watchdog_channel_size,
     } = app_args;
 
     let cancellation_token = CancellationToken::new();
@@ -216,32 +260,6 @@ async fn main() -> Result<(), MainError> {
 
     let discord_application_id = current_application.id;
 
-    let intents = Intents::GUILD_VOICE_STATES;
-    let config = twilight_gateway::Config::new(discord_token.expose_secret().to_owned(), intents);
-
-    let shards = twilight_gateway::create_recommended(&discord_client, config, |_id, builder| {
-        builder.build()
-    })
-    .await
-    .expect("TODO");
-    let shards = Vec::from_iter(shards);
-
-    let senders = TwilightMap::new(
-        shards
-            .iter()
-            .map(|shard| (shard.id().number(), shard.sender()))
-            .collect(),
-    );
-
-    let senders = Arc::new(senders);
-    let songbird = Songbird::twilight(senders, discord_user_id);
-    songbird.set_config(
-        Config::default().decode_mode(songbird::driver::DecodeMode::Decode(DecodeConfig::new(
-            audio_channels.into(),
-            audio_sample_rate.into(),
-        ))),
-    );
-
     let interaction_client = discord_client.interaction(discord_application_id);
 
     let commands = all_commands();
@@ -281,19 +299,16 @@ async fn main() -> Result<(), MainError> {
     let discord_opt_in_command_id = discord_opt_in_command.id.expect("TODO");
     let discord_opt_out_command_id = discord_opt_out_command.id.expect("TODO");
 
-    let discord_info_command_name = discord_info_command.name.into();
-    let discord_opt_in_command_name = discord_opt_in_command.name.into();
-    let discord_opt_out_command_name = discord_opt_out_command.name.into();
+    let discord_info_command_name: Arc<str> = discord_info_command.name.into();
+    let discord_opt_in_command_name: Arc<str> = discord_opt_in_command.name.into();
+    let discord_opt_out_command_name: Arc<str> = discord_opt_out_command.name.into();
 
     let command_router = CommandRouter::from_iter(commands);
     let command_router = Arc::new(command_router);
 
     let discord_client = Arc::new(discord_client);
-    let songbird = Arc::new(songbird);
     let vcs_sender = VCsSender::new(Default::default());
 
-    let initializing_vcs = initialize_vcs(&vcs_sender, &discord_client);
-
     let bot_data = bot_data.into_inner();
     let recording_data = recording_data.into_inner();
     let render_data = render_data.into_inner();
@@ -320,39 +335,117 @@ async fn main() -> Result<(), MainError> {
     let discord_voice_channel_corresponding_text_channel =
         Arc::new(discord_voice_channel_corresponding_text_channel);
 
-    let state = State {
-        audio_channels,
-        audio_sample_rate,
-        bot_manager,
-        cancellation_token: cancellation_token.clone(),
-        discord_application_id,
-        discord_bot_owner_user_id,
-        discord_client: discord_client.clone(),
-        discord_info_command_id,
-        discord_info_command_name,
-        discord_opt_in_command_id,
-        discord_opt_in_command_name,
-        discord_opt_out_command_id,
-        discord_opt_out_command_name,
-        discord_user_id,
-        discord_voice_channel_corresponding_text_channel,
-        recording_manager,
-        render_manager,
-        songbird,
-        user_manager,
-        vcs_sender: vcs_sender.clone(),
-    };
+    tokio::spawn({
+        let cancellation_token = cancellation_token.clone();
+        async move {
+            match ctrl_c().await {
+                Ok(()) => cancellation_token.cancel(),
+                Err(error) => tracing::error!(?error, "failed to listen for interrupt signal"),
+            }
+        }
+    });
 
-    let heat_seeking = tokio::spawn(heat_seek(state.clone()));
+    let (mut watchdog_tx, mut watchdog_rx) =
+        futures::channel::mpsc::channel(watchdog_channel_size.get());
 
-    if let Some(discord_status) = discord_status {
+    std::thread::spawn({
+        let discord_voice_channel_corresponding_text_channel =
+            discord_voice_channel_corresponding_text_channel.clone();
+        let discord_client = discord_client.clone();
+        let vcs_watcher = vcs_sender.subscribe();
+
+        move || {
+            std::thread::sleep(watchdog_warmup);
+
+            loop {
+                tracing::debug!("waiting to send check-in");
+
+                if watchdog_tx.try_send(()).is_err() {
+                    tracing::error!("tokio runtime deadlocked");
+
+                    vcs_watcher
+                        .borrow()
+                        .par_iter()
+                        .for_each(|(&guild_id, vcs_in_guild)| {
+                            if let Some(&voice_channel_id) =
+                                vcs_in_guild.get_left_for(&discord_user_id)
+                            {
+                                let text_channel_id =
+                                    discord_voice_channel_corresponding_text_channel
+                                        .get(&guild_id)
+                                        .and_then(|guild_mappings| {
+                                            guild_mappings.get_right_for(&voice_channel_id).copied()
+                                        })
+                                        .unwrap_or(voice_channel_id);
+
+                                    tokio::runtime::Runtime::new().unwrap().block_on(discord_client.create_message(text_channel_id).content("so sorry I died, I'm in purgatory now, I don't like it here.\nbut I will be back in 5-20 minutes (even if it says I'm still there, I'm not currently recording and will be disconnected soon before later reconnecting and announcing recording again)").into_future());
+                            }
+                        });
+
+                    std::process::exit(1);
+                }
+
+                std::thread::sleep(watchdog_frequency);
+            }
+        }
+    });
+
+    tokio::spawn(async move {
+        tokio::time::sleep(watchdog_warmup).await;
+
+        loop {
+            tracing::debug!("waiting to acknowledge the watchdog");
+
+            if watchdog_rx.recv().await.is_err() {
+                tracing::error!("watchdog died (this should be impossible)");
+                std::process::exit(1);
+            }
+        }
+    });
+
+    loop {
+        tokio::spawn({
+            let vcs_sender = vcs_sender.clone();
+            let discord_client = discord_client.clone();
+
+            async move { initialize_vcs(&vcs_sender, &discord_client).await }
+        });
+
+        let intents = Intents::GUILD_VOICE_STATES;
+        let config =
+            twilight_gateway::Config::new(discord_token.expose_secret().to_owned(), intents);
+
+        let shards =
+            twilight_gateway::create_recommended(&discord_client, config, |_id, builder| {
+                builder.build()
+            })
+            .await
+            .expect("TODO");
+        let shards = Vec::from_iter(shards);
+
+        let senders = TwilightMap::new(
+            shards
+                .iter()
+                .map(|shard| (shard.id().number(), shard.sender()))
+                .collect(),
+        );
+
+        let senders = Arc::new(senders);
+        let songbird = Songbird::twilight(senders, discord_user_id);
+        songbird.set_config(
+            Config::default().decode_mode(songbird::driver::DecodeMode::Decode(DecodeConfig::new(
+                audio_channels.into(),
+                audio_sample_rate.into(),
+            ))),
+        );
+        if let Some(discord_status) = &discord_status {
             shards.iter().for_each(|shard| {
                 shard.command(
                     &UpdatePresence::new(
                         vec![
                             MinimalActivity {
                                 kind: ActivityType::Listening,
-                            name: (*discord_status).to_owned(),
+                                name: discord_status.clone(),
                                 url: None,
                             }
                             .into(),
@@ -366,49 +459,82 @@ async fn main() -> Result<(), MainError> {
             });
         }
 
+        let songbird = Arc::new(songbird);
+
+        let state = State {
+            audio_channels,
+            audio_sample_rate,
+            bot_manager: bot_manager.clone(),
+            cancellation_token: cancellation_token.clone(),
+            discord_application_id,
+            discord_bot_owner_user_id,
+            discord_client: discord_client.clone(),
+            discord_info_command_id,
+            discord_info_command_name: discord_info_command_name.clone(),
+            discord_opt_in_command_id,
+            discord_opt_in_command_name: discord_opt_in_command_name.clone(),
+            discord_opt_out_command_id,
+            discord_opt_out_command_name: discord_opt_out_command_name.clone(),
+            discord_user_id,
+            discord_voice_channel_corresponding_text_channel:
+                discord_voice_channel_corresponding_text_channel.clone(),
+            recording_manager: recording_manager.clone(),
+            render_manager: render_manager.clone(),
+            songbird,
+            user_manager: user_manager.clone(),
+            vcs_sender: vcs_sender.clone(),
+        };
+
+        let mut heat_seeking = tokio::spawn(heat_seek(state.clone()));
+
         let run_shards = shards
             .into_iter()
             .map(|shard| handle_events(command_router.clone(), state.clone(), shard));
-    let run_shards = JoinSet::from_iter(run_shards);
-    let run_shards = run_shards.join_all();
-
-    tokio::spawn({
-        let cancellation_token = cancellation_token.clone();
-        async move {
-            match ctrl_c().await {
-                Ok(()) => cancellation_token.cancel(),
-                Err(error) => tracing::error!(?error, "failed to listen for interrupt signal"),
-            }
-        }
-    });
-
-    tokio::spawn(async {
-        let duration = Duration::from_secs(120);
-        let mut interval = tokio::time::interval(duration);
-
-        loop {
-            interval.tick().await;
-
-            tracing::debug!("this process is still alive");
-        }
-    });
-
-    let finished_naturally = async move {
-        initializing_vcs.await;
-        heat_seeking.await.unwrap();
-        run_shards.await;
-    };
-    tokio::pin!(finished_naturally);
+        let mut run_shards = JoinSet::from_iter(run_shards);
 
         select! {
-        _ = &mut finished_naturally => {
-            Ok(())
+            _heat_seeking_exited = &mut heat_seeking => {
+                tracing::warn!("heat seeking exited, which shouldn't happen. let's try again");
+                continue;
+            }
+            _first_shard_exited = run_shards.join_next() => {
+                tracing::warn!("a shard exited when it's not supposed to, let's try reconnecting them all");
+                heat_seeking.abort();
+                continue;
             }
             () = cancellation_token.cancelled() => {
-            tracing::warn!("waiting for tasks to gracefully shut down");
-            finished_naturally.await;
+                tracing::warn!("gracefully shutting down");
 
-            Err(MainError::Cancelled)
+                FuturesUnordered::from_iter(
+                    vcs_sender
+                        .borrow()
+                        .iter()
+                        .map(|(&guild_id, vcs_in_guild)| {
+                            let discord_client = discord_client.clone();
+                            let discord_voice_channel_corresponding_text_channel = discord_voice_channel_corresponding_text_channel.clone();
+
+                            async move {
+                                if let Some(&voice_channel_id) =
+                                    vcs_in_guild.get_left_for(&discord_user_id)
+                                {
+                                    let text_channel_id =
+                                        discord_voice_channel_corresponding_text_channel
+                                            .get(&guild_id)
+                                            .and_then(|guild_mappings| {
+                                                guild_mappings.get_right_for(&voice_channel_id).copied()
+                                            })
+                                            .unwrap_or(voice_channel_id);
+
+                                        let _ = discord_client.create_message(text_channel_id).content("probably about to be updated, back in 5-20 minutes").await;
+                                }
+                            }
+                        })
+                    ).collect::<()>().await;
+                heat_seeking.await.unwrap();
+                run_shards.join_all().await;
+
+                return Err(MainError::Cancelled);
+            }
         }
     }
 }