dasboard cannot block shutdown

1 files changed, 55 insertions, 24 deletions

diff --git a/src/api/dashboard/mod.rs b/src/api/dashboard/mod.rs
index 6c7f1c5..92271ef 100644
--- a/src/api/dashboard/mod.rs
+++ b/src/api/dashboard/mod.rs
@@ -18,7 +18,7 @@ use std::convert::Infallible;
 use std::sync::{Arc, Mutex};
 use std::time::Duration;
 
-use axum::http::{HeaderValue, header::CONTENT_TYPE};
+use axum::http::{HeaderValue, StatusCode, header::CONTENT_TYPE};
 use axum::response::sse::{Event, KeepAlive, Sse};
 use axum::response::{Html, IntoResponse, Response};
 use axum::routing::get;
@@ -57,6 +57,9 @@ struct Hub {
     /// Each patch, to the clients connected.
     patches: broadcast::Sender<Arc<Published>>,
     control: Mutex<Control>,
+    /// True once the node stops: every client gets an `offline` event, and its stream
+    /// ends (so the HTTP server needn't wait for clients to go).
+    closing: watch::Receiver<bool>,
 }
 
 struct Control {
@@ -66,7 +69,7 @@ struct Control {
 }
 
 impl Hub {
-    fn new(state: AppState) -> Arc<Self> {
+    fn new(state: AppState, closing: watch::Receiver<bool>) -> Arc<Self> {
         let (latest, _) = watch::channel(None);
         let (patches, _) = broadcast::channel(64);
         Arc::new(Self {
@@ -77,6 +80,7 @@ impl Hub {
                 running: false,
                 wanted_until: Instant::now(),
             }),
+            closing,
         })
     }
 
@@ -84,7 +88,7 @@ impl Hub {
     fn want(self: &Arc<Self>) {
         let mut control = self.control.lock_or_recover();
         control.wanted_until = Instant::now() + LINGER;
-        if !control.running {
+        if !control.running && !*self.closing.borrow() {
             control.running = true;
             // What it published before it stopped is stale.
             self.latest.send_replace(None);
@@ -92,24 +96,25 @@ impl Hub {
         }
     }
 
-    /// The latest snapshot, once there is one.
-    async fn current(&self) -> Arc<Published> {
+    /// The latest snapshot, once there is one; None if the node stops first.
+    async fn current(&self) -> Option<Arc<Published>> {
         let mut latest = self.latest.subscribe();
-        loop {
-            // The borrow (a lock guard, not Send) ends here, before any later await.
-            let published = latest
-                .wait_for(Option::is_some)
-                .await
-                .ok()
-                .and_then(|current| current.clone());
-            if let Some(published) = published {
-                return published;
+        let mut closing = self.closing.clone();
+        // The borrows (lock guards, not Send) end within the arms.
+        tokio::select! {
+            published = latest.wait_for(Option::is_some) => {
+                published.ok().and_then(|current| current.clone())
             }
-            // The sender lives as long as the hub: not reached.
-            tokio::time::sleep(TICK).await;
+            _ = closing.wait_for(|closing| *closing) => None,
         }
     }
 
+    /// Resolves once the node stops (or its signal is gone).
+    async fn closed(&self) {
+        let mut closing = self.closing.clone();
+        let _ = closing.wait_for(|closing| *closing).await;
+    }
+
     async fn run(self: Arc<Self>) {
         let mut last: Option<JsonValue> = None;
         let mut seq = 0u64;
@@ -143,7 +148,9 @@ impl Hub {
             tokio::time::sleep(TICK).await;
             {
                 let mut control = self.control.lock_or_recover();
-                if self.patches.receiver_count() == 0 && Instant::now() >= control.wanted_until {
+                let unwanted =
+                    self.patches.receiver_count() == 0 && Instant::now() >= control.wanted_until;
+                if unwanted || *self.closing.borrow() {
                     control.running = false;
                     return;
                 }
@@ -301,12 +308,13 @@ fn merge_patch(old: &JsonValue, new: &JsonValue) -> Option<JsonValue> {
 }
 
 /// The dashboard's routes (authenticated, as the rest of `/admin`), with their hub.
-pub(crate) fn router(state: AppState) -> Router<AppState> {
+/// `closing` turns true when the node stops (`Hub::closing`).
+pub(crate) fn router(state: AppState, closing: watch::Receiver<bool>) -> Router<AppState> {
     Router::new()
         .route("/admin/dashboard", get(page))
         .route("/admin/dashboard/state", get(current_state))
         .route("/admin/dashboard/events", get(events))
-        .layer(Extension(Hub::new(state)))
+        .layer(Extension(Hub::new(state, closing)))
 }
 
 /// `GET /admin/dashboard`: the page.
@@ -317,7 +325,9 @@ async fn page() -> Html<&'static str> {
 /// `GET /admin/dashboard/state`: the latest snapshot, for a client that wants it whole.
 async fn current_state(Extension(hub): Extension<Arc<Hub>>) -> Response {
     hub.want();
-    let current = hub.current().await;
+    let Some(current) = hub.current().await else {
+        return StatusCode::SERVICE_UNAVAILABLE.into_response();
+    };
     (
         [(CONTENT_TYPE, HeaderValue::from_static("application/json"))],
         current.json.clone(),
@@ -327,17 +337,30 @@ async fn current_state(Extension(hub): Extension<Arc<Hub>>) -> Response {
 
 /// `GET /admin/dashboard/events`: Server-Sent Events: a `snapshot` first, then a
 /// `patch` whenever something changed (each with its number as the event's ID; a
-/// client that fell behind gets a `snapshot` again).
+/// client that fell behind gets a `snapshot` again), and `offline` when the node
+/// stops, which ends the stream.
 async fn events(Extension(hub): Extension<Arc<Hub>>) -> impl IntoResponse {
     // Subscribed before the snapshot is read: no patch after it is missed.
     let mut patches = hub.patches.subscribe();
     hub.want();
     let first = hub.current().await;
     let stream = async_stream::stream! {
+        let Some(first) = first else {
+            yield Ok::<_, Infallible>(offline_event());
+            return;
+        };
         let mut seen = first.seq;
-        yield Ok::<_, Infallible>(snapshot_event(&first));
+        yield Ok(snapshot_event(&first));
         loop {
-            match patches.recv().await {
+            let received = tokio::select! {
+                received = patches.recv() => Some(received),
+                () = hub.closed() => None,
+            };
+            let Some(received) = received else {
+                yield Ok(offline_event());
+                break;
+            };
+            match received {
                 Ok(patch) if patch.seq <= seen => {}
                 Ok(patch) => {
                     seen = patch.seq;
@@ -347,7 +370,10 @@ async fn events(Extension(hub): Extension<Arc<Hub>>) -> impl IntoResponse {
                         .data(patch.json.clone()));
                 }
                 Err(broadcast::error::RecvError::Lagged(_)) => {
-                    let current = hub.current().await;
+                    let Some(current) = hub.current().await else {
+                        yield Ok(offline_event());
+                        break;
+                    };
                     seen = current.seq;
                     yield Ok(snapshot_event(&current));
                 }
@@ -358,6 +384,11 @@ async fn events(Extension(hub): Extension<Arc<Hub>>) -> impl IntoResponse {
     Sse::new(stream).keep_alive(KeepAlive::default())
 }
 
+/// The node stops: the client may reconnect once it is back.
+fn offline_event() -> Event {
+    Event::default().event("offline").data("{}")
+}
+
 fn snapshot_event(published: &Published) -> Event {
     Event::default()
         .event("snapshot")