Add feature to automatically leave when the voicechannel is empty

Signed-off-by: trivernis <trivernis@protonmail.com>

src/commands/music/mod.rs

@@ -30,6 +30,7 @@ use crate::providers::music::{
 use crate::utils::error::{BotError, BotResult};
 use crate::utils::store::Store;
 use regex::Regex;
+use std::sync::atomic::{AtomicIsize, AtomicUsize, Ordering};
 use std::time::Duration;
 
 mod clear;
@@ -50,6 +51,68 @@ mod skip;
 #[prefix("m")]
 pub struct Music;
 
+struct SongEndNotifier {
+    channel_id: ChannelId,
+    http: Arc<Http>,
+    queue: Arc<Mutex<MusicQueue>>,
+    handler: Arc<Mutex<Call>>,
+}
+
+#[async_trait]
+impl VoiceEventHandler for SongEndNotifier {
+    async fn act(&self, _ctx: &EventContext<'_>) -> Option<Event> {
+        log::debug!("Song ended in {}. Playing next one", self.channel_id);
+        while !play_next_in_queue(&self.http, &self.channel_id, &self.queue, &self.handler).await {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+        }
+
+        None
+    }
+}
+
+struct ChannelDurationNotifier {
+    channel_id: ChannelId,
+    guild_id: GuildId,
+    count: Arc<AtomicUsize>,
+    queue: Arc<Mutex<MusicQueue>>,
+    leave_in: Arc<AtomicIsize>,
+    handler: Arc<Mutex<Call>>,
+    manager: Arc<Songbird>,
+}
+
+#[async_trait]
+impl VoiceEventHandler for ChannelDurationNotifier {
+    async fn act(&self, _ctx: &EventContext<'_>) -> Option<Event> {
+        let count_before = self.count.fetch_add(1, Ordering::Relaxed);
+        log::debug!(
+            "Playing in channel {} for {} minutes",
+            self.channel_id,
+            count_before
+        );
+        let queue_lock = self.queue.lock().await;
+        if queue_lock.leave_flag {
+            log::debug!("Waiting to leave");
+            if self.leave_in.fetch_sub(1, Ordering::Relaxed) <= 0 {
+                log::debug!("Leaving voice channel");
+                {
+                    let mut handler_lock = self.handler.lock().await;
+                    handler_lock.remove_all_global_events();
+                }
+                if let Some(current) = queue_lock.current() {
+                    let _ = current.stop();
+                }
+                let _ = self.manager.remove(self.guild_id).await;
+                log::debug!("Left the voice channel");
+            }
+        } else {
+            log::debug!("Resetting leave value");
+            self.leave_in.store(5, Ordering::Relaxed)
+        }
+
+        None
+    }
+}
+
 /// Joins a voice channel
 async fn join_channel(ctx: &Context, channel_id: ChannelId, guild_id: GuildId) -> Arc<Mutex<Call>> {
     log::debug!(
@@ -76,12 +139,25 @@ async fn join_channel(ctx: &Context, channel_id: ChannelId, guild_id: GuildId) -
         handler_lock.add_global_event(
             Event::Track(TrackEvent::End),
             SongEndNotifier {
-                channel_id,
+                channel_id: channel_id.clone(),
                 http: ctx.http.clone(),
                 queue: Arc::clone(&queue),
                 handler: handler.clone(),
             },
         );
+
+        handler_lock.add_global_event(
+            Event::Periodic(Duration::from_secs(60), None),
+            ChannelDurationNotifier {
+                channel_id,
+                guild_id,
+                count: Default::default(),
+                queue: Arc::clone(&queue),
+                handler: handler.clone(),
+                leave_in: Arc::new(AtomicIsize::new(5)),
+                manager: manager.clone(),
+            },
+        );
     }
 
     handler
@@ -105,7 +181,7 @@ async fn get_voice_manager(ctx: &Context) -> Arc<Songbird> {
 }
 
 /// Returns a reference to a guilds music queue
-async fn get_queue_for_guild(
+pub(crate) async fn get_queue_for_guild(
     ctx: &Context,
     guild_id: &GuildId,
 ) -> BotResult<Arc<Mutex<MusicQueue>>> {
@@ -120,25 +196,6 @@ async fn get_queue_for_guild(
     Ok(queue)
 }
 
-struct SongEndNotifier {
-    channel_id: ChannelId,
-    http: Arc<Http>,
-    queue: Arc<Mutex<MusicQueue>>,
-    handler: Arc<Mutex<Call>>,
-}
-
-#[async_trait]
-impl VoiceEventHandler for SongEndNotifier {
-    async fn act(&self, _ctx: &EventContext<'_>) -> Option<Event> {
-        log::debug!("Song ended in {}. Playing next one", self.channel_id);
-        while !play_next_in_queue(&self.http, &self.channel_id, &self.queue, &self.handler).await {
-            tokio::time::sleep(Duration::from_millis(100)).await;
-        }
-
-        None
-    }
-}
-
 /// Plays the next song in the queue
 async fn play_next_in_queue(
     http: &Arc<Http>,

src/commands/music/play.rs

@@ -58,7 +58,7 @@ async fn play(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
 
     if play_first {
         log::debug!("Playing first song in queue");
-        play_next_in_queue(&ctx.http, &msg.channel_id, &queue, &handler_lock).await;
+        while !play_next_in_queue(&ctx.http, &msg.channel_id, &queue, &handler_lock).await {}
     }
 
     Ok(())

src/commands/music/play_next.rs

@@ -48,7 +48,7 @@ async fn play_next(ctx: &Context, msg: &Message, args: Args) -> CommandResult {
     };
 
     if play_first {
-        play_next_in_queue(&ctx.http, &msg.channel_id, &queue, &handler).await;
+        while !play_next_in_queue(&ctx.http, &msg.channel_id, &queue, &handler).await {}
     }
 
     Ok(())

src/handler.rs

@@ -1,7 +1,11 @@
+use crate::commands::music::get_queue_for_guild;
 use serenity::async_trait;
 use serenity::client::Context;
 use serenity::model::event::ResumedEvent;
 use serenity::model::gateway::{Activity, Ready};
+use serenity::model::guild::Member;
+use serenity::model::id::{ChannelId, GuildId};
+use serenity::model::voice::VoiceState;
 use serenity::prelude::*;
 
 pub(crate) struct Handler;
@@ -18,4 +22,57 @@ impl EventHandler for Handler {
     async fn resume(&self, _: Context, _: ResumedEvent) {
         log::info!("Reconnected to gateway")
     }
+
+    async fn voice_state_update(
+        &self,
+        ctx: Context,
+        guild_id: Option<GuildId>,
+        old_state: Option<VoiceState>,
+        new_state: VoiceState,
+    ) {
+        let mut member_count = None;
+
+        let guild_id = if let Some(gid) = guild_id {
+            gid
+        } else {
+            return;
+        };
+
+        if let Some(old_id) = old_state.and_then(|c| c.channel_id) {
+            member_count = get_own_channel_member_count(&ctx, &old_id).await;
+        }
+        if member_count.is_none() {
+            if let Some(new_id) = new_state.channel_id {
+                member_count = get_own_channel_member_count(&ctx, &new_id).await;
+            }
+        }
+
+        if let Some(count) = member_count {
+            log::debug!("{} Members in channel", count);
+            let queue = get_queue_for_guild(&ctx, &guild_id).await.unwrap();
+            let mut queue_lock = queue.lock().await;
+            log::debug!("Setting leave flag to {}", count == 0);
+            queue_lock.leave_flag = count == 0;
+        }
+    }
+}
+
+/// Returns the number of members in the channel if it's the bots voice channel
+async fn get_own_channel_member_count(ctx: &Context, channel_id: &ChannelId) -> Option<usize> {
+    let channel = ctx.http.get_channel(channel_id.0).await.ok()?;
+    let guild_channel = channel.guild()?;
+    let current_user = ctx.http.get_current_user().await.ok()?;
+
+    let members = guild_channel.members(&ctx).await.ok()?;
+    let own_channel = members
+        .iter()
+        .find(|m| m.user.id == current_user.id)
+        .is_some();
+
+    if !own_channel {
+        return None;
+    }
+    let members: Vec<Member> = members.into_iter().filter(|m| !m.user.bot).collect();
+
+    Some(members.len())
 }

src/providers/music/queue.rs

@@ -12,6 +12,7 @@ pub struct MusicQueue {
     inner: VecDeque<Song>,
     current: Option<TrackHandle>,
     paused: bool,
+    pub leave_flag: bool,
 }
 
 impl MusicQueue {
@@ -20,6 +21,7 @@ impl MusicQueue {
             inner: VecDeque::new(),
             current: None,
             paused: false,
+            leave_flag: false,
         }
     }