UDP seems to work now

src/shutdown.rs

@@ -1,56 +1,60 @@
 use tokio::sync::broadcast;
 
-/// Listens for the server shutdown signal.
-///
-/// Shutdown is signalled using a `broadcast::Receiver`. Only a single value is
-/// ever sent. Once a value has been sent via the broadcast channel, the server
-/// should shutdown.
-///
-/// The `Shutdown` struct listens for the signal and tracks that the signal has
-/// been received. Callers may query for whether the shutdown signal has been
-/// received or not.
-#[derive(Debug)]
 pub(crate) struct Shutdown {
-    /// `true` if the shutdown signal has been received
     shutdown: bool,
-
-    /// The receive half of the channel used to listen for shutdown.
     notify: broadcast::Receiver<()>,
+    sender: broadcast::Sender<()>,
 }
 
 impl Shutdown {
-    /// Create a new `Shutdown` backed by the given `broadcast::Receiver`.
-    pub(crate) fn new(notify: broadcast::Receiver<()>) -> Shutdown {
+    pub(crate) fn new() -> Shutdown {
+        let (sender, notify) = broadcast::channel(1);
         Shutdown {
             shutdown: false,
             notify,
+            sender,
+        }
+    }
+
+    pub(crate) fn shutdown(&mut self) {
+        if self.shutdown {
+            return;
+        }
+        let _ = self.sender.send(());
+        self.shutdown = true;
+    }
+
+    pub(crate) fn receiver(&self) -> ShutdownReceiver {
+        ShutdownReceiver::new(self.notify.resubscribe())
+    }
+}
+
+#[derive(Debug)]
+pub(crate) struct ShutdownReceiver {
+    shutdown: bool,
+    notify: broadcast::Receiver<()>,
+}
+
+impl ShutdownReceiver {
+    pub(crate) fn new(notify: broadcast::Receiver<()>) -> ShutdownReceiver {
+        ShutdownReceiver {
+            shutdown: false,
+            notify,
         }
     }
 
-    /// Returns `true` if the shutdown signal has been received.
-    // pub(crate) fn is_shutdown(&self) -> bool {
-    //    self.shutdown
-    // }
-
-    /// Receive the shutdown notice, waiting if necessary.
     pub(crate) async fn recv(&mut self) {
-        // If the shutdown signal has already been received, then return
-        // immediately.
         if self.shutdown {
             return;
         }
-
-        // Cannot receive a "lag error" as only one value is ever sent.
         let _ = self.notify.recv().await;
-
-        // Remember that the signal has been received.
         self.shutdown = true;
     }
 }
 
-impl Clone for Shutdown {
+impl Clone for ShutdownReceiver {
     fn clone(&self) -> Self {
-        Shutdown {
+        ShutdownReceiver {
             shutdown: self.shutdown,
             notify: self.notify.resubscribe(),
         }