From 90cd078973b23b2291cf156e46729842f33c1806 Mon Sep 17 00:00:00 2001
From: Jake Stanger <mail@jstanger.dev>
Date: Sat, 28 Jan 2023 14:40:12 +0000
Subject: [PATCH] fix(mpd): stops working if connection lost

The client will now attempt to reconnect when a connection loss is detected.

Fixes #21.
---
 src/clients/music/mod.rs   |  6 ++++-
 src/clients/music/mpd.rs   | 35 +++++++++++++++++++++++----
 src/clients/music/mpris.rs |  2 +-
 src/modules/music.rs       | 49 ++++++++++++++++++++++++--------------
 4 files changed, 67 insertions(+), 25 deletions(-)

diff --git a/src/clients/music/mod.rs b/src/clients/music/mod.rs
index 6c3edd2..554dd7c 100644
--- a/src/clients/music/mod.rs
+++ b/src/clients/music/mod.rs
@@ -7,7 +7,11 @@ use tokio::sync::broadcast;
 pub mod mpd;
 pub mod mpris;
 
-pub type PlayerUpdate = (Option<Track>, Status);
+#[derive(Clone, Debug)]
+pub enum PlayerUpdate {
+    Update(Box<Option<Track>>, Status),
+    Disconnect,
+}
 
 #[derive(Clone, Debug)]
 pub struct Track {
diff --git a/src/clients/music/mpd.rs b/src/clients/music/mpd.rs
index 770d8b1..3dfd48d 100644
--- a/src/clients/music/mpd.rs
+++ b/src/clients/music/mpd.rs
@@ -19,7 +19,7 @@ use tokio::spawn;
 use tokio::sync::broadcast::{channel, error::SendError, Receiver, Sender};
 use tokio::sync::Mutex;
 use tokio::time::sleep;
-use tracing::{debug, error};
+use tracing::{debug, error, info};
 
 lazy_static! {
     static ref CONNECTIONS: Arc<Mutex<HashMap<String, Arc<MpdClient>>>> =
@@ -72,7 +72,9 @@ impl MpdClient {
                         Subsystem::Player | Subsystem::Queue | Subsystem::Mixer,
                     ) = change
                     {
-                        Self::send_update(&client, &tx, &music_dir).await?;
+                        Self::send_update(&client, &tx, &music_dir)
+                            .await
+                            .expect("Failed to send update");
                     }
                 }
 
@@ -92,7 +94,7 @@ impl MpdClient {
         client: &Client,
         tx: &Sender<PlayerUpdate>,
         music_dir: &Path,
-    ) -> Result<(), SendError<(Option<Track>, Status)>> {
+    ) -> Result<(), SendError<PlayerUpdate>> {
         let current_song = client.command(commands::CurrentSong).await;
         let status = client.command(commands::Status).await;
 
@@ -100,12 +102,22 @@ impl MpdClient {
             let track = current_song.map(|s| Self::convert_song(&s.song, music_dir));
             let status = Status::from(status);
 
-            tx.send((track, status))?;
+            tx.send(PlayerUpdate::Update(Box::new(track), status))?;
         }
 
         Ok(())
     }
 
+    fn is_connected(&self) -> bool {
+        !self.client.is_connection_closed()
+    }
+
+    fn send_disconnect_update(&self) -> Result<(), SendError<PlayerUpdate>> {
+        info!("Connection to MPD server lost");
+        self.tx.send(PlayerUpdate::Disconnect)?;
+        Ok(())
+    }
+
     fn convert_song(song: &Song, music_dir: &Path) -> Track {
         let (track, disc) = song.number();
 
@@ -189,7 +201,20 @@ pub async fn get_client(
             connections.insert(host.to_string(), Arc::clone(&client));
             Ok(client)
         }
-        Some(client) => Ok(Arc::clone(client)),
+        Some(client) => {
+            if client.is_connected() {
+                Ok(Arc::clone(client))
+            } else {
+                client
+                    .send_disconnect_update()
+                    .expect("Failed to send disconnect update");
+
+                let client = MpdClient::new(host, music_dir).await?;
+                let client = Arc::new(client);
+                connections.insert(host.to_string(), Arc::clone(&client));
+                Ok(client)
+            }
+        }
     }
 }
 
diff --git a/src/clients/music/mpris.rs b/src/clients/music/mpris.rs
index 80c4ca3..7675106 100644
--- a/src/clients/music/mpris.rs
+++ b/src/clients/music/mpris.rs
@@ -170,7 +170,7 @@ impl Client {
 
         let track = Track::from(metadata);
 
-        let player_update: PlayerUpdate = (Some(track), status);
+        let player_update = PlayerUpdate::Update(Box::new(Some(track)), status);
 
         tx.send(player_update)
             .expect("Failed to send player update");
diff --git a/src/modules/music.rs b/src/modules/music.rs
index a30a9b9..dcb300a 100644
--- a/src/modules/music.rs
+++ b/src/modules/music.rs
@@ -178,30 +178,43 @@ impl Module<Button> for MusicModule {
             let music_dir = self.music_dir.clone();
 
             spawn(async move {
-                let mut rx = {
-                    let client = get_client(player_type, &host, music_dir).await;
-                    client.subscribe_change()
-                };
+                loop {
+                    let mut rx = {
+                        let client = get_client(player_type, &host, music_dir.clone()).await;
+                        client.subscribe_change()
+                    };
 
-                while let Ok((track, status)) = rx.recv().await {
-                    match track {
-                        Some(track) => {
-                            let display_string =
-                                replace_tokens(format.as_str(), &tokens, &track, &status, &icons);
+                    while let Ok(update) = rx.recv().await {
+                        match update {
+                            PlayerUpdate::Update(track, status) => match *track {
+                                Some(track) => {
+                                    let display_string = replace_tokens(
+                                        format.as_str(),
+                                        &tokens,
+                                        &track,
+                                        &status,
+                                        &icons,
+                                    );
 
-                            let update = SongUpdate {
-                                song: track,
-                                status,
-                                display_string,
-                            };
+                                    let update = SongUpdate {
+                                        song: track,
+                                        status,
+                                        display_string,
+                                    };
 
-                            tx.send(ModuleUpdateEvent::Update(Some(update))).await?;
+                                    tx.send(ModuleUpdateEvent::Update(Some(update)))
+                                        .await
+                                        .expect(ERR_CHANNEL_SEND);
+                                }
+                                None => tx
+                                    .send(ModuleUpdateEvent::Update(None))
+                                    .await
+                                    .expect(ERR_CHANNEL_SEND),
+                            },
+                            PlayerUpdate::Disconnect => break,
                         }
-                        None => tx.send(ModuleUpdateEvent::Update(None)).await?,
                     }
                 }
-
-                Ok::<(), mpsc::error::SendError<ModuleUpdateEvent<Self::SendMessage>>>(())
             });
         }