Block shutdown until all pending messages are sent (Fixes #27).

examples/zeroconf_test.py

@@ -14,8 +14,12 @@ def has_found(self, name):
         return name in self.found
 
     def add_service(self, zeroconf, type, name):
+        print(f"  Service {name} added")
         self.found.append(name.replace("." + TYPE, ""))
 
+    def remove_service(self, zeroconf, type, name):
+        print(f"  Service {name} removed")
+
 
 zeroconf = Zeroconf()
 listener = MyListener()

src/fsm.rs

@@ -222,9 +222,13 @@ impl<AF: Unpin + AddressFamily> Future for FSM<AF> {
     type Output = ();
     fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<()> {
         let pinned = Pin::get_mut(self);
+        let mut shutdown = false;
         while let Poll::Ready(cmd) = Pin::new(&mut pinned.commands).poll_recv(cx) {
             match cmd {
-                Some(Command::Shutdown) => return Poll::Ready(()),
+                Some(Command::Shutdown) => {
+                    shutdown = true;
+                    break;
+                }
                 Some(Command::SendUnsolicited {
                     svc,
                     ttl,
@@ -239,9 +243,11 @@ impl<AF: Unpin + AddressFamily> Future for FSM<AF> {
             }
         }
 
-        match pinned.recv_packets(cx) {
-            Ok(_) => (),
-            Err(e) => error!("ResponderRecvPacket Error: {:?}", e),
+        if !shutdown {
+            match pinned.recv_packets(cx) {
+                Ok(_) => (),
+                Err(e) => error!("ResponderRecvPacket Error: {:?}", e),
+            }
         }
 
         while let Some((ref response, addr)) = pinned.outgoing.pop_front() {
@@ -256,6 +262,11 @@ impl<AF: Unpin + AddressFamily> Future for FSM<AF> {
             }
         }
 
+        if shutdown {
+            trace!("shutting down {:?}", std::any::type_name::<AF>());
+            return Poll::Ready(());
+        }
+
         Poll::Pending
     }
 }

src/lib.rs

@@ -1,5 +1,5 @@
 use futures_util::{future, future::FutureExt};
-use log::warn;
+use log::{trace, warn};
 use std::cell::RefCell;
 use std::future::Future;
 use std::io;
@@ -26,14 +26,13 @@ const MDNS_PORT: u16 = 5353;
 pub struct Responder {
     services: Services,
     commands: RefCell<CommandSender>,
-    shutdown: Arc<Shutdown>,
+    shutdown: Shutdown,
 }
 
 pub struct Service {
     id: usize,
     services: Services,
     commands: CommandSender,
-    _shutdown: Arc<Shutdown>,
 }
 
 type ResponderTask = Box<dyn Future<Output = ()> + Send + Unpin>;
@@ -42,7 +41,7 @@ impl Responder {
     /// Spawn a responder task on an os thread.
     pub fn new() -> io::Result<Responder> {
         let (tx, rx) = std::sync::mpsc::sync_channel(0);
-        thread::Builder::new()
+        let join_handle = thread::Builder::new()
             .name("mdns-responder".to_owned())
             .spawn(move || {
                 let rt = tokio::runtime::Builder::new_current_thread()
@@ -59,7 +58,9 @@ impl Responder {
                     }
                 })
             })?;
-        rx.recv().expect("rx responder channel closed")
+        let mut responder = rx.recv().expect("rx responder channel closed")?;
+        responder.shutdown.1 = Some(join_handle);
+        Ok(responder)
     }
 
     /// Spawn a `Responder` with the provided tokio `Handle`.
@@ -115,7 +116,7 @@ impl Responder {
         let responder = Responder {
             services: services,
             commands: RefCell::new(commands.clone()),
-            shutdown: Arc::new(Shutdown(commands)),
+            shutdown: Shutdown(commands, None),
         };
 
         Ok((responder, task))
@@ -177,7 +178,6 @@ impl Responder {
             id: id,
             commands: self.commands.borrow().clone(),
             services: self.services.clone(),
-            _shutdown: self.shutdown.clone(),
         }
     }
 }
@@ -189,12 +189,16 @@ impl Drop for Service {
     }
 }
 
-struct Shutdown(CommandSender);
+struct Shutdown(CommandSender, Option<thread::JoinHandle<()>>);
 
 impl Drop for Shutdown {
     fn drop(&mut self) {
+        trace!("Shutting down...");
         self.0.send_shutdown();
-        // TODO wait for tasks to shutdown
+
+        if let Some(handle) = self.1.take() {
+            handle.join().expect("failed to join thread");
+        }
     }
 }