From ce0d2fd0ce92fcefc182385fef1bf6ceb8b2ef58 Mon Sep 17 00:00:00 2001
From: DaXcess <daxcess@daxcess.io>
Date: Sat, 10 Dec 2022 22:58:29 +0100
Subject: [PATCH] Fix double disconnect/corrupt state on auto leave

---
 src/main.rs            |  2 +-
 src/session/manager.rs | 13 ++++++-------
 src/session/mod.rs     | 26 +++++++++++++++++++-------
 3 files changed, 26 insertions(+), 15 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 894f1da..6858ed4 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -25,7 +25,7 @@ mod session;
 mod stats;
 mod utils;
 
-#[tokio::main]
+#[tokio::main(flavor = "current_thread")]
 async fn main() {
   if std::env::var("RUST_LOG").is_err() {
     #[cfg(debug_assertions)]
diff --git a/src/session/manager.rs b/src/session/manager.rs
index b574f37..fa4c265 100644
--- a/src/session/manager.rs
+++ b/src/session/manager.rs
@@ -69,11 +69,10 @@ impl InnerSessionManager {
   }
 
   /// Remove a session
-  pub async fn remove_session(&mut self, guild_id: GuildId) {
-    if let Some(session) = self.sessions.get(&guild_id) {
-      if let Some(owner) = session.owner().await {
-        self.owner_map.remove(&owner);
-      }
+  pub async fn remove_session(&mut self, guild_id: GuildId, owner: Option<UserId>) {
+    // Remove the owner from the owner map (if it exists)
+    if let Some(owner) = owner {
+      self.owner_map.remove(&owner);
     }
 
     self.sessions.remove(&guild_id);
@@ -149,8 +148,8 @@ impl SessionManager {
   }
 
   /// Remove a session
-  pub async fn remove_session(&self, guild_id: GuildId) {
-    self.0.write().await.remove_session(guild_id).await;
+  pub async fn remove_session(&self, guild_id: GuildId, owner: Option<UserId>) {
+    self.0.write().await.remove_session(guild_id, owner).await;
   }
 
   /// Remove owner from owner map.
diff --git a/src/session/mod.rs b/src/session/mod.rs
index 23978ee..e03c2b4 100644
--- a/src/session/mod.rs
+++ b/src/session/mod.rs
@@ -54,6 +54,10 @@ struct InnerSpoticordSession {
   disconnect_handle: Option<tokio::task::JoinHandle<()>>,
 
   client: Option<Client>,
+
+  /// Whether the session has been disconnected
+  /// If this is true then this instance should no longer be used and dropped
+  disconnected: bool,
 }
 
 impl SpoticordSession {
@@ -90,6 +94,7 @@ impl SpoticordSession {
       playback_info: None,
       disconnect_handle: None,
       client: None,
+      disconnected: false,
     };
 
     let mut instance = Self(Arc::new(RwLock::new(inner)));
@@ -278,6 +283,7 @@ impl SpoticordSession {
                   message.embed(|embed| {
                     embed.title("Failed to connect to Spotify");
                     embed.description(why);
+                    embed.footer(|footer| footer.text("Please try again"));
                     embed.color(Status::Error as u64);
 
                     embed
@@ -290,9 +296,6 @@ impl SpoticordSession {
                 error!("Failed to send error message: {:?}", why);
               }
 
-              // Clean up session
-              instance.player_stopped().await;
-
               break;
             }
 
@@ -382,6 +385,11 @@ impl SpoticordSession {
             _ => {}
           }
         }
+
+        // Clean up session
+        if !inner.read().await.disconnected {
+          instance.player_stopped().await;
+        }
       }
     });
 
@@ -515,7 +523,7 @@ impl SpoticordSession {
     //  read lock to read the current owner.
     // This would deadlock if we have an active write lock
     {
-      let inner = self.0.read().await;
+      let mut inner = self.0.write().await;
       inner.disconnect_no_abort().await;
     }
 
@@ -599,7 +607,7 @@ impl SpoticordSession {
 
   pub async fn disconnect_with_message(&self, content: &str) {
     {
-      let inner = self.0.read().await;
+      let mut inner = self.0.write().await;
 
       // Firstly we disconnect
       inner.disconnect_no_abort().await;
@@ -674,8 +682,12 @@ impl SpoticordSession {
 
 impl InnerSpoticordSession {
   /// Internal version of disconnect, which does not abort the disconnect timer
-  async fn disconnect_no_abort(&self) {
-    self.session_manager.remove_session(self.guild_id).await;
+  async fn disconnect_no_abort(&mut self) {
+    self.disconnected = true;
+    self
+      .session_manager
+      .remove_session(self.guild_id, self.owner)
+      .await;
 
     let mut call = self.call.lock().await;