Live reload

Cargo.lock

@@ -351,6 +351,12 @@ dependencies = [
  "typenum",
 ]
 
+[[package]]
+name = "data-encoding"
+version = "2.6.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8566979429cf69b49a5c740c60791108e86440e8be149bbea4fe54d2c32d6e2"
+
 [[package]]
 name = "debounced"
 version = "0.2.0"
@@ -765,6 +771,21 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "hyper-tungstenite"
+version = "0.17.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0110a0487cbc65c3d1f38c2ef851dbf8bee8c2761e5a96be6a59ba84412b4752"
+dependencies = [
+ "http-body-util",
+ "hyper",
+ "hyper-util",
+ "pin-project-lite",
+ "tokio",
+ "tokio-tungstenite",
+ "tungstenite",
+]
+
 [[package]]
 name = "hyper-util"
 version = "0.1.10"
@@ -1597,6 +1618,17 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "sha1"
+version = "0.10.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba"
+dependencies = [
+ "cfg-if",
+ "cpufeatures",
+ "digest",
+]
+
 [[package]]
 name = "sha2"
 version = "0.10.8"
@@ -1713,6 +1745,12 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "sync_wrapper"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263"
+
 [[package]]
 name = "synstructure"
 version = "0.13.1"
@@ -1842,6 +1880,18 @@ dependencies = [
  "tokio",
 ]
 
+[[package]]
+name = "tokio-tungstenite"
+version = "0.26.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "be4bf6fecd69fcdede0ec680aaf474cdab988f9de6bc73d3758f0160e3b7025a"
+dependencies = [
+ "futures-util",
+ "log",
+ "tokio",
+ "tungstenite",
+]
+
 [[package]]
 name = "tokio-util"
 version = "0.7.13"
@@ -1895,6 +1945,10 @@ version = "0.5.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d039ad9159c98b70ecfd540b2573b97f7f52c3e8d9f8ad57a24b916a536975f9"
 dependencies = [
+ "futures-core",
+ "futures-util",
+ "pin-project-lite",
+ "sync_wrapper",
  "tower-layer",
  "tower-service",
 ]
@@ -1952,6 +2006,24 @@ version = "0.1.33"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "e672c95779cf947c5311f83787af4fa8fffd12fb27e4993211a84bdfd9610f9c"
 
+[[package]]
+name = "tungstenite"
+version = "0.26.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "413083a99c579593656008130e29255e54dcaae495be556cc26888f211648c24"
+dependencies = [
+ "byteorder",
+ "bytes",
+ "data-encoding",
+ "http",
+ "httparse",
+ "log",
+ "rand",
+ "sha1",
+ "thiserror",
+ "utf-8",
+]
+
 [[package]]
 name = "typenum"
 version = "1.17.0"
@@ -2091,7 +2163,11 @@ dependencies = [
  "grass",
  "grass_compiler",
  "html5ever",
+ "http",
+ "http-body",
+ "http-body-util",
  "hyper",
+ "hyper-tungstenite",
  "hyper-util",
  "log",
  "markup5ever_rcdom",

Cargo.toml

@@ -29,7 +29,11 @@ grass_compiler = { version = "0.13.4", features = [
     "custom-builtin-fns",
 ], git = "https://git.shadowfacts.net/shadowfacts/grass.git", branch = "custom-global-variables" }
 html5ever = "0.27.0"
+http = "1.2.0"
+http-body = "1.0.1"
+http-body-util = "0.1.2"
 hyper = { version = "1.5.2", features = ["server", "http1"] }
+hyper-tungstenite = "0.17.0"
 hyper-util = { version = "0.1.10", features = ["tokio", "service"] }
 log = "0.4.22"
 markup5ever_rcdom = "0.3.0"
@@ -42,7 +46,7 @@ tera = "1.20.0"
 tokio = { version = "1.42.0", features = ["full"] }
 tokio-stream = "0.1.17"
 toml = "0.8.19"
-tower = "0.5.2"
+tower = { version = "0.5.2", features = ["steer", "util"] }
 tower-http = { version = "0.6.2", features = ["fs"] }
 unicode-normalization = "0.1.24"
 url = "2.5.4"

site_test/layout/default.html

@@ -36,5 +36,17 @@
     {% block content %}{% endblock %}
 
     <script data-goatcounter="https://shadowfacts.goatcounter.com/count" async src="//gc.zgo.at/count.v3.js" crossorigin="anonymous"></script>
+
+    {% if _development %}
+    <script>
+    let ws = new WebSocket("/_dev/live_reload");
+    ws.onmessage = (event) => {
+      if (event.data == "regenerated") {
+        ws.close();
+        window.location.reload();
+      }
+    };
+    </script>
+    {% endif %}
 </body>
 </html>

src/generator/templates.rs

@@ -96,6 +96,7 @@ impl<T: 'static, F: Fn(&T, &mut Context) -> () + 'static> Rule for BuildTemplate
         context.insert("_domain", &*DOMAIN);
         context.insert("_permalink", &self.permalink);
         context.insert("_stylesheet_cache_buster", &*CB);
+        context.insert("_development", &cfg!(debug_assertions));
         context
     }
 }

src/live_reload.rs

@@ -0,0 +1,65 @@
+use std::convert::Infallible;
+
+use anyhow::anyhow;
+use futures::{SinkExt, stream::FusedStream};
+use http::{Request, Response};
+use http_body_util::{BodyExt, combinators::UnsyncBoxBody};
+use hyper::body::Bytes;
+use hyper_tungstenite::{HyperWebsocket, tungstenite::Message};
+use log::error;
+use tokio::sync::broadcast::error::RecvError;
+use tower::Service;
+use tower_http::{
+    services::{ServeDir, ServeFile},
+    set_status::SetStatus,
+};
+
+use crate::CloneableReceiver;
+
+pub async fn handle(
+    mut request: Request<hyper::body::Incoming>,
+    mut fallback: ServeDir<SetStatus<ServeFile>>,
+    regen_complete_rx: CloneableReceiver<()>,
+) -> anyhow::Result<Response<UnsyncBoxBody<Bytes, std::io::Error>>> {
+    if request.uri().path() == "/_dev/live_reload"
+        && hyper_tungstenite::is_upgrade_request(&request)
+    {
+        let (response, websocket) = hyper_tungstenite::upgrade(&mut request, None)?;
+
+        tokio::spawn(async move {
+            if let Err(e) = serve_websocket(websocket, regen_complete_rx).await {
+                error!("Error serving websocket: {e:?}");
+            }
+        });
+
+        Ok(response.map(|b| b.map_err(|_: Infallible| unreachable!()).boxed_unsync()))
+    } else {
+        let fallback_resp = fallback.call(request).await?;
+        Ok(Response::new(fallback_resp.boxed_unsync()))
+    }
+}
+
+async fn serve_websocket(
+    websocket: HyperWebsocket,
+    mut regen_complete_rx: CloneableReceiver<()>,
+) -> anyhow::Result<()> {
+    let mut websocket = websocket.await?;
+    'outer: loop {
+        match regen_complete_rx.0.recv().await {
+            Err(RecvError::Closed) => {
+                break 'outer;
+            }
+            Err(_) => (),
+            Ok(_) => {
+                if websocket.is_terminated() {
+                    return Ok(());
+                }
+                let result = websocket.send(Message::text("regenerated")).await;
+                if let Err(e) = result {
+                    return Err(anyhow!(e));
+                }
+            }
+        }
+    }
+    Ok(())
+}

src/main.rs

@@ -1,6 +1,7 @@
 #![feature(let_chains)]
 
 mod generator;
+mod live_reload;
 
 use clap::{Arg, ArgAction, ArgMatches, Command, command};
 use compute_graph::AsyncGraph;
@@ -15,8 +16,10 @@ use std::net::SocketAddr;
 use std::rc::Rc;
 use std::time::Duration;
 use tokio::net::TcpListener;
+use tokio::sync::broadcast;
 use tokio_stream::StreamExt;
 use tokio_stream::wrappers::UnboundedReceiverStream;
+use tower::service_fn;
 use tower_http::services::{ServeDir, ServeFile};
 
 #[tokio::main]
@@ -58,15 +61,26 @@ async fn main() {
         println!("{}", graph.as_dot_string());
     }
 
-    let serve_future = serve(&matches);
+    let (regen_complete_tx, regen_complete_rx) = tokio::sync::broadcast::channel::<()>(4);
-    let watch_future = watch(&matches, graph, watcher);
+    let serve_future = serve(&matches, regen_complete_rx);
+    let watch_future = watch(&matches, graph, watcher, regen_complete_tx);
 
     join!(serve_future, watch_future);
 }
 
-async fn serve(matches: &ArgMatches) {
+struct CloneableReceiver<T>(pub broadcast::Receiver<T>);
+impl<T: Clone> Clone for CloneableReceiver<T> {
+    fn clone(&self) -> Self {
+        Self(self.0.resubscribe())
+    }
+}
+
+async fn serve(matches: &ArgMatches, regen_complete_rx: broadcast::Receiver<()>) {
     if let Some(("serve", _)) = matches.subcommand() {
-        let service = ServeDir::new("out/").not_found_service(ServeFile::new("out/404.html"));
+        let serve_dir = ServeDir::new("out").not_found_service(ServeFile::new("out/404.html"));
+        let receiver = CloneableReceiver(regen_complete_rx);
+        let service =
+            service_fn(move |req| live_reload::handle(req, serve_dir.clone(), receiver.clone()));
 
         let addr = SocketAddr::from((
             if cfg!(debug_assertions) {
@@ -84,7 +98,10 @@ async fn serve(matches: &ArgMatches) {
             let io = TokioIo::new(tcp);
             let service = TowerToHyperService::new(service.clone());
             tokio::task::spawn(async move {
-                let result = http1::Builder::new().serve_connection(io, service).await;
+                let result = http1::Builder::new()
+                    .serve_connection(io, service)
+                    .with_upgrades()
+                    .await;
                 if let Err(e) = result {
                     error!("Error handling connection: {e:?}");
                 }
@@ -93,15 +110,23 @@ async fn serve(matches: &ArgMatches) {
     }
 }
 
-async fn watch(matches: &ArgMatches, mut graph: AsyncGraph<()>, watcher: Rc<RefCell<FileWatcher>>) {
+async fn watch(
+    matches: &ArgMatches,
+    mut graph: AsyncGraph<()>,
+    watcher: Rc<RefCell<FileWatcher>>,
+    regen_complete_tx: broadcast::Sender<()>,
+) {
     if let Some(true) = matches.get_one("watch") {
-        let (tx, rx) = tokio::sync::mpsc::unbounded_channel::<()>();
+        let (regen_tx, regen_rx) = tokio::sync::mpsc::unbounded_channel::<()>();
 
         watcher.borrow_mut().watch(content_base_path(), move || {
-            tx.send(()).expect("sending regenerate signal");
+            regen_tx.send(()).expect("sending regenerate signal");
         });
 
-        let mut debounced = debounced(UnboundedReceiverStream::new(rx), Duration::from_millis(100));
+        let mut debounced = debounced(
+            UnboundedReceiverStream::new(regen_rx),
+            Duration::from_millis(100),
+        );
 
         let watch = async move { FileWatcher::start(watcher).await.expect("watching files") };
 
@@ -109,6 +134,12 @@ async fn watch(matches: &ArgMatches, mut graph: AsyncGraph<()>, watcher: Rc<RefC
             while let Some(_) = debounced.next().await {
                 info!("Regenerating");
                 graph.evaluate_async().await;
+
+                if regen_complete_tx.receiver_count() > 0 {
+                    regen_complete_tx
+                        .send(())
+                        .expect("sending regen complete signal");
+                }
             }
         };