From a22965a3be93409df5f48121fd2f1cbe9fbd04e8 Mon Sep 17 00:00:00 2001 From: Jacob Date: Fri, 12 Jun 2026 00:01:04 -0400 Subject: [PATCH] feat: make reconnecting after discord disconnect faster (not obligate an entire process restart) --- src/main.rs | 203 +++++++++++++++++++++++++++------------------------- 1 file changed, 105 insertions(+), 98 deletions(-) diff --git a/src/main.rs b/src/main.rs index 397f119..fb614ba 100644 --- a/src/main.rs +++ b/src/main.rs @@ -80,7 +80,7 @@ struct AppArgs { discord_nickname: Option>, #[arg(long, env)] - discord_status: Option>, + discord_status: Option, #[arg(long, env, value_parser = parse_guild_vc_to_text_channel)] discord_voice_channel_corresponding_text_channel: @@ -216,32 +216,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 +255,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 = discord_info_command.name.into(); + let discord_opt_in_command_name: Arc = discord_opt_in_command.name.into(); + let discord_opt_out_command_name: Arc = 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,58 +291,6 @@ 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(), - }; - - let heat_seeking = tokio::spawn(heat_seek(state.clone())); - - 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(), - url: None, - } - .into(), - ], - false, - None, - Status::Idle, - ) - .expect("TODO"), - ) - }); - } - - 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 { @@ -393,22 +312,110 @@ async fn main() -> Result<(), MainError> { } }); - let finished_naturally = async move { - initializing_vcs.await; - heat_seeking.await.unwrap(); - run_shards.await; - }; - tokio::pin!(finished_naturally); + loop { + tokio::spawn({ + let vcs_sender = vcs_sender.clone(); + let discord_client = discord_client.clone(); - select! { - _ = &mut finished_naturally => { - Ok(()) + 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.clone(), + url: None, + } + .into(), + ], + false, + None, + Status::Idle, + ) + .expect("TODO"), + ) + }); } - () = cancellation_token.cancelled() => { - tracing::warn!("waiting for tasks to gracefully shut down"); - finished_naturally.await; - Err(MainError::Cancelled) + 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 mut run_shards = JoinSet::from_iter(run_shards); + + select! { + _heat_seeking_exited = &mut heat_seeking => { + // this shouldn't happen, but let's try again + continue; + } + _first_shard_exited = run_shards.join_next() => { + heat_seeking.abort(); + continue; + } + () = cancellation_token.cancelled() => { + heat_seeking.await.unwrap(); + run_shards.join_all().await; + + return Err(MainError::Cancelled); + } } } }