ssg/server/
hmr.rs

1// Copyright © 2023 - 2026 Static Site Generator (SSG). All rights reserved.
2// SPDX-License-Identifier: Apache-2.0 OR MIT
3
4//! Hot-module-reload protocol for `ssg dev` (issue #526).
5//!
6//! Defines the WebSocket message format the dev server pushes to every
7//! connected browser tab and the [`HmrBroadcaster`] that fans messages
8//! out across N tabs.
9//!
10//! # Protocol (AC5)
11//!
12//! Each frame is a JSON object:
13//!
14//! ```json
15//! { "type": "hmr-css",  "paths": ["assets/style.css"],  "sha": "abc..." }
16//! { "type": "hmr-html", "paths": ["/blog/foo/"],         "sha": "def..." }
17//! { "type": "reload",   "paths": [],                     "sha": "" }
18//! ```
19//!
20//! * `type` — one of `hmr-css`, `hmr-html`, `reload`.
21//!   * `hmr-css`: client swaps `<link>` href in-place (no reload, scroll
22//!     preserved, no FOUC). AC2.
23//!   * `hmr-html`: client fetches each path and swaps the `<main>` body
24//!     (or `<body>` if no `<main>`) — scroll position preserved. AC3/AC4.
25//!   * `reload`: client does a full `location.reload()` — used for
26//!     `<head>` / config / build-error recovery.
27//! * `paths` — affected paths. For `hmr-html` these are the rebuilt
28//!   page URLs the dep graph reported invalidated.
29//! * `sha` — short content hash so the client can de-dupe redundant
30//!   frames from rapid saves. Optional; empty string is fine.
31//!
32//! Frames are pushed as Text WebSocket messages from
33//! [`HmrBroadcaster::broadcast`]; the JS client in
34//! `crate::livereload` parses them.
35//!
36//! # Architecture
37//!
38//! ```text
39//! caller ──▶ broadcast() ──▶ for each registered Sender:
40//!                              try_send(json) ──▶ tab thread ──▶ ws.send()
41//! ```
42//!
43//! Each connected tab owns a `Sender<String>`; the broadcaster holds the
44//! receiving side in a `Mutex<Vec<Sender>>`. Send failures evict the
45//! tab from the list — there is no explicit unregister API, dead tabs
46//! garbage-collect on the next broadcast.
47
48use std::sync::Mutex;
49
50use serde::{Deserialize, Serialize};
51
52/// Message types pushed to the browser HMR client.
53#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
54#[serde(rename_all = "kebab-case")]
55pub enum HmrType {
56    /// Stylesheet swap — no reload, no FOUC.
57    HmrCss,
58    /// Partial page replacement — swap `<main>` body, preserve scroll.
59    HmrHtml,
60    /// Full `location.reload()`.
61    Reload,
62}
63
64impl HmrType {
65    /// The wire string the JS client matches on (`msg.type`).
66    ///
67    /// # Examples
68    ///
69    /// ```
70    /// use ssg::hmr::HmrType;
71    /// assert_eq!(HmrType::HmrCss.wire(), "hmr-css");
72    /// assert_eq!(HmrType::Reload.wire(), "reload");
73    /// ```
74    #[must_use]
75    pub const fn wire(self) -> &'static str {
76        match self {
77            Self::HmrCss => "hmr-css",
78            Self::HmrHtml => "hmr-html",
79            Self::Reload => "reload",
80        }
81    }
82}
83
84/// A single HMR frame.
85#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
86pub struct HmrMessage {
87    /// Frame type.
88    #[serde(rename = "type")]
89    pub kind: HmrType,
90    /// Paths affected. For `hmr-html` these are page URLs; for
91    /// `hmr-css` these are asset URLs; for `reload` this may be empty.
92    pub paths: Vec<String>,
93    /// Short content hash. Optional — empty string is fine.
94    #[serde(default)]
95    pub sha: String,
96}
97
98impl HmrMessage {
99    /// Build a CSS-swap frame for one or more stylesheet paths.
100    ///
101    /// # Examples
102    ///
103    /// ```
104    /// use ssg::hmr::{HmrMessage, HmrType};
105    /// let m = HmrMessage::css(vec!["x.css".into()]);
106    /// assert_eq!(m.kind, HmrType::HmrCss);
107    /// assert_eq!(m.paths, vec!["x.css"]);
108    /// ```
109    #[must_use]
110    pub const fn css(paths: Vec<String>) -> Self {
111        Self {
112            kind: HmrType::HmrCss,
113            paths,
114            sha: String::new(),
115        }
116    }
117
118    /// Build an HTML-partial frame for one or more page URLs.
119    ///
120    /// # Examples
121    ///
122    /// ```
123    /// use ssg::hmr::{HmrMessage, HmrType};
124    /// let m = HmrMessage::html(vec!["/a/".into()]);
125    /// assert_eq!(m.kind, HmrType::HmrHtml);
126    /// ```
127    #[must_use]
128    pub const fn html(paths: Vec<String>) -> Self {
129        Self {
130            kind: HmrType::HmrHtml,
131            paths,
132            sha: String::new(),
133        }
134    }
135
136    /// Build a full-page-reload frame.
137    ///
138    /// # Examples
139    ///
140    /// ```
141    /// use ssg::hmr::{HmrMessage, HmrType};
142    /// let m = HmrMessage::reload();
143    /// assert_eq!(m.kind, HmrType::Reload);
144    /// assert!(m.paths.is_empty());
145    /// ```
146    #[must_use]
147    pub const fn reload() -> Self {
148        Self {
149            kind: HmrType::Reload,
150            paths: Vec::new(),
151            sha: String::new(),
152        }
153    }
154
155    /// Attach a content hash for client-side de-duplication.
156    ///
157    /// # Examples
158    ///
159    /// ```
160    /// use ssg::hmr::HmrMessage;
161    /// let m = HmrMessage::reload().with_sha("abc123");
162    /// assert_eq!(m.sha, "abc123");
163    /// ```
164    #[must_use]
165    pub fn with_sha(mut self, sha: impl Into<String>) -> Self {
166        self.sha = sha.into();
167        self
168    }
169
170    /// Serialise to the JSON wire format. Infallible — the struct is
171    /// always serialisable.
172    ///
173    /// # Panics
174    ///
175    /// Only if `serde_json` cannot serialise a primitive `Vec<String>`,
176    /// which is impossible.
177    ///
178    /// # Examples
179    ///
180    /// ```
181    /// use ssg::hmr::HmrMessage;
182    /// let m = HmrMessage::css(vec!["a.css".into()]);
183    /// let json = m.to_json();
184    /// assert!(json.contains("\"type\":\"hmr-css\""));
185    /// ```
186    #[must_use]
187    pub fn to_json(&self) -> String {
188        serde_json::to_string(self).unwrap_or_else(|_| {
189            // Defensive: serialise a static reload as a last-ditch frame.
190            r#"{"type":"reload","paths":[],"sha":""}"#.to_string()
191        })
192    }
193}
194
195/// Per-tab callback that emits one HMR frame.
196///
197/// The broadcaster keeps a vector of these and invokes each on every
198/// `broadcast()`. Callbacks return `Ok(())` on a successful WebSocket
199/// send and `Err(())` to be evicted (closed socket, write error).
200pub type HmrSink = Box<dyn Fn(&str) -> Result<(), ()> + Send + Sync>;
201
202/// Fan-out HMR sender shared across the dev server.
203///
204/// Construct one per `ssg dev` process. Each connected tab registers
205/// an [`HmrSink`] via [`Self::subscribe`]; [`Self::broadcast`] pushes
206/// the same frame to every live tab and evicts any sink that errored.
207#[derive(Default)]
208pub struct HmrBroadcaster {
209    sinks: Mutex<Vec<HmrSink>>,
210}
211
212impl std::fmt::Debug for HmrBroadcaster {
213    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
214        f.debug_struct("HmrBroadcaster")
215            .field("subscribers", &self.sinks.lock().map_or(0, |v| v.len()))
216            .finish()
217    }
218}
219
220impl HmrBroadcaster {
221    /// Construct an empty broadcaster.
222    ///
223    /// # Examples
224    ///
225    /// ```
226    /// use ssg::hmr::HmrBroadcaster;
227    /// let b = HmrBroadcaster::new();
228    /// assert_eq!(b.subscriber_count(), 0);
229    /// ```
230    #[must_use]
231    pub fn new() -> Self {
232        Self {
233            sinks: Mutex::new(Vec::new()),
234        }
235    }
236
237    /// Register a new tab. The sink is invoked on every subsequent
238    /// [`Self::broadcast`] until it returns `Err(())`.
239    ///
240    /// # Examples
241    ///
242    /// ```
243    /// use ssg::hmr::HmrBroadcaster;
244    /// let b = HmrBroadcaster::new();
245    /// b.subscribe(Box::new(|_| Ok(())));
246    /// assert_eq!(b.subscriber_count(), 1);
247    /// ```
248    pub fn subscribe(&self, sink: HmrSink) {
249        if let Ok(mut g) = self.sinks.lock() {
250            g.push(sink);
251        }
252    }
253
254    /// Returns the current number of subscribed tabs. Useful for tests
255    /// and the dev-server status banner.
256    ///
257    /// # Examples
258    ///
259    /// ```
260    /// use ssg::hmr::HmrBroadcaster;
261    /// let b = HmrBroadcaster::new();
262    /// assert_eq!(b.subscriber_count(), 0);
263    /// ```
264    #[must_use]
265    pub fn subscriber_count(&self) -> usize {
266        self.sinks.lock().map_or(0, |g| g.len())
267    }
268
269    /// Push `msg` to every subscriber. Sinks that return `Err(())` are
270    /// evicted in-place.
271    ///
272    /// Returns the number of tabs that successfully received the frame.
273    ///
274    /// # Examples
275    ///
276    /// ```
277    /// use ssg::hmr::{HmrBroadcaster, HmrMessage};
278    /// let b = HmrBroadcaster::new();
279    /// b.subscribe(Box::new(|_| Ok(())));
280    /// assert_eq!(b.broadcast(&HmrMessage::reload()), 1);
281    /// ```
282    pub fn broadcast(&self, msg: &HmrMessage) -> usize {
283        let payload = msg.to_json();
284        let Ok(mut sinks) = self.sinks.lock() else {
285            return 0;
286        };
287        let mut delivered = 0usize;
288        // Walk in reverse so swap_remove preserves the unvisited prefix.
289        let mut i = sinks.len();
290        while i > 0 {
291            i -= 1;
292            match (sinks[i])(&payload) {
293                Ok(()) => delivered += 1,
294                Err(()) => {
295                    let _ = sinks.swap_remove(i);
296                }
297            }
298        }
299        delivered
300    }
301}
302
303#[cfg(test)]
304mod tests {
305    use super::*;
306    use std::sync::{
307        atomic::{AtomicUsize, Ordering},
308        Arc,
309    };
310
311    #[test]
312    fn hmr_type_wire_strings_match_spec() {
313        assert_eq!(HmrType::HmrCss.wire(), "hmr-css");
314        assert_eq!(HmrType::HmrHtml.wire(), "hmr-html");
315        assert_eq!(HmrType::Reload.wire(), "reload");
316    }
317
318    #[test]
319    fn hmr_message_css_serialises_to_kebab_type() {
320        let m = HmrMessage::css(vec!["assets/style.css".into()]);
321        let json = m.to_json();
322        assert!(json.contains(r#""type":"hmr-css""#));
323        assert!(json.contains(r#""paths":["assets/style.css"]"#));
324        assert!(json.contains(r#""sha":"""#));
325    }
326
327    #[test]
328    fn hmr_message_html_serialises_paths() {
329        let m = HmrMessage::html(vec!["/a/".into(), "/b/".into()]);
330        let json = m.to_json();
331        assert!(json.contains(r#""type":"hmr-html""#));
332        assert!(json.contains(r#""/a/""#));
333        assert!(json.contains(r#""/b/""#));
334    }
335
336    #[test]
337    fn hmr_message_reload_has_empty_paths() {
338        let m = HmrMessage::reload();
339        let json = m.to_json();
340        assert!(json.contains(r#""type":"reload""#));
341        assert!(json.contains(r#""paths":[]"#));
342    }
343
344    #[test]
345    fn hmr_message_with_sha_round_trips() {
346        let m = HmrMessage::html(vec!["/x/".into()]).with_sha("deadbeef");
347        let json = m.to_json();
348        let back: HmrMessage = serde_json::from_str(&json).unwrap();
349        assert_eq!(back, m);
350        assert_eq!(back.sha, "deadbeef");
351    }
352
353    #[test]
354    fn hmr_message_default_sha_omittable() {
355        // Without an explicit sha field on the wire, deserialise should
356        // still succeed thanks to `#[serde(default)]`.
357        let json = r#"{"type":"hmr-css","paths":["a.css"]}"#;
358        let m: HmrMessage = serde_json::from_str(json).unwrap();
359        assert_eq!(m.kind, HmrType::HmrCss);
360        assert_eq!(m.sha, "");
361    }
362
363    #[test]
364    fn broadcaster_new_has_zero_subscribers() {
365        let b = HmrBroadcaster::new();
366        assert_eq!(b.subscriber_count(), 0);
367    }
368
369    #[test]
370    fn broadcaster_default_has_zero_subscribers() {
371        let b = HmrBroadcaster::default();
372        assert_eq!(b.subscriber_count(), 0);
373    }
374
375    #[test]
376    fn subscribe_increments_count() {
377        let b = HmrBroadcaster::new();
378        b.subscribe(Box::new(|_| Ok(())));
379        b.subscribe(Box::new(|_| Ok(())));
380        assert_eq!(b.subscriber_count(), 2);
381        // Deliver one frame so both registered sinks actually run.
382        assert_eq!(b.broadcast(&HmrMessage::reload()), 2);
383    }
384
385    #[test]
386    fn broadcast_delivers_to_every_subscriber() {
387        let b = HmrBroadcaster::new();
388        let calls = Arc::new(AtomicUsize::new(0));
389        for _ in 0..3 {
390            let calls = Arc::clone(&calls);
391            b.subscribe(Box::new(move |_| {
392                let _ = calls.fetch_add(1, Ordering::SeqCst);
393                Ok(())
394            }));
395        }
396        let delivered = b.broadcast(&HmrMessage::reload());
397        assert_eq!(delivered, 3);
398        assert_eq!(calls.load(Ordering::SeqCst), 3);
399    }
400
401    #[test]
402    fn broadcast_evicts_failed_sinks() {
403        let b = HmrBroadcaster::new();
404        b.subscribe(Box::new(|_| Err(())));
405        b.subscribe(Box::new(|_| Ok(())));
406        b.subscribe(Box::new(|_| Err(())));
407
408        let delivered = b.broadcast(&HmrMessage::reload());
409        assert_eq!(delivered, 1);
410        // The two erroring sinks should have been evicted.
411        assert_eq!(b.subscriber_count(), 1);
412
413        // Second broadcast still sees the surviving subscriber.
414        let delivered2 = b.broadcast(&HmrMessage::reload());
415        assert_eq!(delivered2, 1);
416        assert_eq!(b.subscriber_count(), 1);
417    }
418
419    #[test]
420    fn broadcast_passes_json_payload_to_sinks() {
421        let b = HmrBroadcaster::new();
422        let seen: Arc<Mutex<Option<String>>> = Arc::new(Mutex::new(None));
423        let seen_clone = Arc::clone(&seen);
424        b.subscribe(Box::new(move |s| {
425            *seen_clone.lock().unwrap() = Some(s.to_string());
426            Ok(())
427        }));
428
429        let _ = b.broadcast(&HmrMessage::css(vec!["x.css".into()]));
430        let payload = seen.lock().unwrap().clone().unwrap();
431        assert!(payload.contains(r#""type":"hmr-css""#));
432        assert!(payload.contains("x.css"));
433    }
434
435    #[test]
436    fn broadcast_with_no_subscribers_is_noop() {
437        let b = HmrBroadcaster::new();
438        assert_eq!(b.broadcast(&HmrMessage::reload()), 0);
439    }
440
441    #[test]
442    fn broadcaster_debug_format_includes_count() {
443        let b = HmrBroadcaster::new();
444        b.subscribe(Box::new(|_| Ok(())));
445        let d = format!("{b:?}");
446        assert!(d.contains("HmrBroadcaster"));
447        assert!(d.contains('1'));
448        // Run the registered sink once so the closure body executes.
449        assert_eq!(b.broadcast(&HmrMessage::reload()), 1);
450    }
451
452    #[test]
453    fn broadcaster_recovers_from_poisoned_sink_lock() {
454        // Poison the internal sinks mutex, then exercise every lock
455        // site: subscribe silently no-ops, broadcast returns 0, the
456        // counters fall back to 0.
457        let b = HmrBroadcaster::new();
458        b.subscribe(Box::new(|_| Ok(())));
459        let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
460            let _guard = b.sinks.lock().unwrap();
461            panic!("poison sinks");
462        }));
463
464        b.subscribe(Box::new(|_| Ok(()))); // if-let Ok fails: no-op
465        assert_eq!(b.subscriber_count(), 0, "poisoned lock reads as empty");
466        assert_eq!(b.broadcast(&HmrMessage::reload()), 0);
467        let d = format!("{b:?}");
468        assert!(d.contains('0'), "Debug falls back to 0: {d}");
469    }
470
471    #[test]
472    fn hmr_message_to_json_is_valid_json() {
473        let m =
474            HmrMessage::html(vec!["/a/".into(), "/b/".into()]).with_sha("h");
475        let v: serde_json::Value = serde_json::from_str(&m.to_json()).unwrap();
476        assert_eq!(v["type"], "hmr-html");
477        assert_eq!(v["paths"][0], "/a/");
478        assert_eq!(v["paths"][1], "/b/");
479        assert_eq!(v["sha"], "h");
480    }
481
482    #[test]
483    fn protocol_three_message_types_distinct() {
484        let css = HmrMessage::css(vec![]).to_json();
485        let html = HmrMessage::html(vec![]).to_json();
486        let reload = HmrMessage::reload().to_json();
487        assert!(css.contains("hmr-css"));
488        assert!(html.contains("hmr-html"));
489        assert!(reload.contains("reload"));
490        assert!(!reload.contains("hmr-"));
491    }
492
493    #[test]
494    fn hmr_type_is_copy() {
495        let t = HmrType::HmrHtml;
496        let _copy = t;
497        assert_eq!(t, HmrType::HmrHtml);
498    }
499
500    #[test]
501    fn hmr_type_clone_and_debug() {
502        // `assert_eq!` only invokes `Debug::fmt` on a *failing*
503        // comparison, so the derived `Clone`/`Debug` impls need an
504        // explicit exercise to be reached at all.
505        let t = HmrType::HmrCss;
506        #[allow(clippy::clone_on_copy)]
507        let cloned = t.clone();
508        assert_eq!(cloned, t);
509        let dbg = format!("{t:?}");
510        assert!(dbg.contains("HmrCss"), "unexpected Debug output: {dbg}");
511    }
512
513    #[test]
514    fn hmr_message_clone_and_debug() {
515        let m = HmrMessage::html(vec!["/a/".into()]).with_sha("abc");
516        let cloned = m.clone();
517        assert_eq!(cloned, m);
518        let dbg = format!("{m:?}");
519        assert!(dbg.contains("HmrHtml"), "unexpected Debug output: {dbg}");
520        assert!(dbg.contains("abc"), "unexpected Debug output: {dbg}");
521    }
522}
ssg/server/hmr.rs

ssg/server/
hmr.rs
