discovery: revise shutdown handling after review

librespot-org · Oct 21, 2024 · 878d64d · 878d64d
1 parent 1a88825
commit 878d64d
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 11 deletions.
diff --git a/discovery/src/lib.rs b/discovery/src/lib.rs
@@ -12,7 +12,6 @@ mod server;
 
 use std::{
     borrow::Cow,
-    convert::Infallible,
     error::Error as StdError,
     pin::Pin,
     task::{Context, Poll},
@@ -39,9 +38,13 @@ pub enum DiscoveryEvent {
     ZeroconfError(DiscoveryError),
 }
 
+enum ZeroconfCmd {
+    Shutdown,
+}
+
 pub struct DnsSdHandle {
     task_handle: tokio::task::JoinHandle<()>,
-    shutdown_tx: oneshot::Sender<Infallible>,
+    shutdown_tx: oneshot::Sender<ZeroconfCmd>,
 }
 
 impl DnsSdHandle {
@@ -51,9 +54,12 @@ impl DnsSdHandle {
             task_handle,
             shutdown_tx,
         } = self;
-        std::mem::drop(shutdown_tx);
-        let _ = task_handle.await;
-        log::debug!("Zeroconf responder stopped");
+        if shutdown_tx.send(ZeroconfCmd::Shutdown).is_err() {
+            log::warn!("Zeroconf responder unexpectedly disappeared");
+        } else {
+            let _ = task_handle.await;
+            log::debug!("Zeroconf responder stopped");
+        }
     }
 }
 

diff --git a/discovery/src/server.rs b/discovery/src/server.rs
@@ -1,7 +1,6 @@
 use std::{
     borrow::Cow,
     collections::BTreeMap,
-    convert::Infallible,
     net::{Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener},
     sync::{Arc, Mutex},
 };
@@ -252,8 +251,12 @@ impl RequestHandler {
     }
 }
 
+pub(crate) enum DiscoveryServerCmd {
+    Shutdown,
+}
+
 pub struct DiscoveryServer {
-    close_tx: oneshot::Sender<Infallible>,
+    close_tx: oneshot::Sender<DiscoveryServerCmd>,
     task_handle: tokio::task::JoinHandle<()>,
 }
 
@@ -324,14 +327,12 @@ impl DiscoveryServer {
                         });
                     }
                     _ = &mut close_rx => {
-                        debug!("Shutting down discovery server");
                         break;
                     }
                 }
             }
 
             graceful.shutdown().await;
-            debug!("Discovery server stopped");
         });
 
         Ok(Self {
@@ -346,7 +347,12 @@ impl DiscoveryServer {
             task_handle,
             ..
         } = self;
-        std::mem::drop(close_tx);
-        let _ = task_handle.await;
+        log::debug!("Shutting down discovery server");
+        if close_tx.send(DiscoveryServerCmd::Shutdown).is_err() {
+            log::warn!("Discovery server unexpectedly disappeared");
+        } else {
+            let _ = task_handle.await;
+            log::debug!("Discovery server stopped");
+        }
     }
 }