Improve reload experience when using with site generator such as cargo

doc

src/lib.rs

@@ -32,6 +32,7 @@ use watcher::create_watcher;
 
 static ADDR: OnceCell<String> = OnceCell::const_new();
 static ROOT: OnceCell<PathBuf> = OnceCell::const_new();
+static HARD: OnceCell<bool> = OnceCell::const_new();
 static TX: OnceCell<broadcast::Sender<()>> = OnceCell::const_new();
 
 pub struct Listener {
@@ -40,6 +41,7 @@ pub struct Listener {
     root_path: PathBuf,
     debouncer: Debouncer<RecommendedWatcher, FileIdMap>,
     rx: Receiver<Result<Vec<DebouncedEvent>, Vec<notify::Error>>>,
+    hard: bool,
 }
 
 impl Listener {
@@ -53,6 +55,7 @@ impl Listener {
     /// }
     /// ```
     pub async fn start(self) -> Result<(), Box<dyn Error>> {
+        HARD.set(self.hard)?;
         ROOT.set(self.root_path.clone())?;
         let (tx, _) = broadcast::channel(16);
         TX.set(tx)?;
@@ -65,6 +68,19 @@ impl Listener {
         Ok(())
     }
 
+    /// Always hard reload the page instead of hot-reload
+    /// ```
+    /// use live_server::listen;
+    ///
+    /// async fn serve_hard() -> Result<(), Box<dyn std::error::Error>> {
+    ///     listen("127.0.0.1:8080", "./").await?.hard_reload().start().await
+    /// }
+    /// ```
+    pub fn hard_reload(mut self) -> Self {
+        self.hard = true;
+        self
+    }
+
     /// Return the link of the server, like `http://127.0.0.1:8080`.
     ///
     /// ```
@@ -118,5 +134,6 @@ pub async fn listen<A: Into<String>, R: Into<PathBuf>>(
         debouncer,
         root_path,
         rx,
+        hard: false,
     })
 }

src/main.rs

@@ -18,6 +18,11 @@ struct Args {
     /// Open the page in browser automatically
     #[clap(short, long)]
     open: bool,
+    /// Hard reload the page on update instead of hot reload
+    ///
+    /// Try using this if the reload is not working as expected
+    #[clap(long)]
+    hard: bool,
 }
 
 #[tokio::main]
@@ -30,15 +35,20 @@ async fn main() {
         port,
         root,
         open,
+        hard,
     } = Args::parse();
 
     let addr = format!("{}:{}", host, port);
-    let listener = listen(addr, root).await.unwrap();
+    let mut listener = listen(addr, root).await.unwrap();
 
     if open {
         let link = listener.link().unwrap();
         open::that(link).unwrap();
     }
 
+    if hard {
+        listener = listener.hard_reload();
+    }
+
     listener.start().await.unwrap();
 }

src/server.rs

@@ -1,6 +1,7 @@
 use std::io::ErrorKind;
 use std::{fs, net::IpAddr};
 
+use axum::extract::ws::WebSocket;
 use axum::{
     body::Body,
     extract::{ws::Message, Request, WebSocketUpgrade},
@@ -12,7 +13,7 @@ use futures::{sink::SinkExt, stream::StreamExt};
 use local_ip_address::local_ip;
 use tokio::net::TcpListener;
 
-use crate::{ADDR, ROOT, TX};
+use crate::{ADDR, HARD, ROOT, TX};
 
 pub(crate) async fn serve(tcp_listener: TcpListener, router: Router) {
     axum::serve(tcp_listener, router).await.unwrap();
@@ -64,40 +65,49 @@ pub(crate) fn create_server() -> Router {
                 ws.on_failed_upgrade(|error| {
                     log::error!("Failed to upgrade websocket: {}", error);
                 })
-                .on_upgrade(|socket| async move {
-                    let (mut sender, mut receiver) = socket.split();
-                    let tx = TX.get().unwrap();
-                    let mut rx = tx.subscribe();
-                    let mut send_task = tokio::spawn(async move {
-                        while rx.recv().await.is_ok() {
-                            sender.send(Message::Text(String::new())).await.unwrap();
-                        }
-                    });
-                    let mut recv_task =
-                        tokio::spawn(
-                            async move { while let Some(Ok(_)) = receiver.next().await {} },
-                        );
-                    tokio::select! {
-                        _ = (&mut send_task) => recv_task.abort(),
-                        _ = (&mut recv_task) => send_task.abort(),
-                    };
-                })
+                .on_upgrade(on_websocket_upgrade)
             }),
         )
 }
 
+async fn on_websocket_upgrade(socket: WebSocket) {
+    let (mut sender, mut receiver) = socket.split();
+    let tx = TX.get().unwrap();
+    let mut rx = tx.subscribe();
+    let mut send_task = tokio::spawn(async move {
+        while rx.recv().await.is_ok() {
+            sender.send(Message::Text(String::new())).await.unwrap();
+        }
+    });
+    let mut recv_task =
+        tokio::spawn(async move { while let Some(Ok(_)) = receiver.next().await {} });
+    tokio::select! {
+        _ = (&mut send_task) => recv_task.abort(),
+        _ = (&mut recv_task) => send_task.abort(),
+    };
+}
+
 async fn static_assets(req: Request<Body>) -> (StatusCode, HeaderMap, Body) {
     let addr = ADDR.get().unwrap();
     let root = ROOT.get().unwrap();
 
+    let is_reload = req.uri().query().is_some_and(|x| x == "reload");
+
     // Get the path and mime of the static file.
-    let mut path = req.uri().path().to_string();
-    path.remove(0);
-    let mut path = root.join(path);
+    let uri_path = req.uri().path();
+    let mut path = root.join(&uri_path[1..]);
     if path.is_dir() {
+        if !uri_path.ends_with('/') {
+            // redirect so parent links work correctly
+            let redirect = format!("{}/", uri_path);
+            let mut headers = HeaderMap::new();
+            headers.append(header::LOCATION, HeaderValue::from_str(&redirect).unwrap());
+            return (StatusCode::TEMPORARY_REDIRECT, headers, Body::empty());
+        }
         path.push("index.html");
     }
     let mime = mime_guess::from_path(&path).first_or_text_plain();
+
     let mut headers = HeaderMap::new();
     headers.append(
         header::CONTENT_TYPE,
@@ -117,7 +127,7 @@ async fn static_assets(req: Request<Body>) -> (StatusCode, HeaderMap, Body) {
                 _ => StatusCode::INTERNAL_SERVER_ERROR,
             };
             if mime == "text/html" {
-                let script = format!(include_str!("templates/websocket.html"), addr);
+                let script = format_script(addr, is_reload, true);
                 let html = format!(include_str!("templates/error.html"), script, err);
                 let body = Body::from(html);
 
@@ -137,9 +147,45 @@ async fn static_assets(req: Request<Body>) -> (StatusCode, HeaderMap, Body) {
                 return (StatusCode::INTERNAL_SERVER_ERROR, headers, body);
             }
         };
-        let script = format!(include_str!("templates/websocket.html"), addr);
+        let script = format_script(addr, is_reload, false);
         file = format!("{text}{script}").into_bytes();
+    } else if !HARD.get().copied().unwrap_or(false) {
+        // allow client to cache assets for a smoother reload.
+        // client handles preloading to refresh cache before reloading.
+        headers.append(
+            header::CACHE_CONTROL,
+            HeaderValue::from_str("max-age=30").unwrap(),
+        );
     }
 
     (StatusCode::OK, headers, Body::from(file))
 }
+
+/// JS script containing a function that takes in the address and connects to the websocket.
+const WEBSOCKET_FUNCTION: &str = include_str!("templates/websocket.js");
+
+/// JS script to inject to the HTML on reload so the client
+/// knows it's a successful reload.
+const RELOAD_PAYLOAD: &str = include_str!("templates/reload.js");
+
+/// Inject the address into the websocket script and wrap it in a script tag
+fn format_script(addr: &str, is_reload: bool, is_error: bool) -> String {
+    match (is_reload, is_error) {
+        // successful reload, inject the reload payload
+        (true, false) => format!("<script>{}</script>", RELOAD_PAYLOAD),
+        // failed reload, don't inject anything so the client polls again
+        (true, true) => String::new(),
+        // normal connection, inject the websocket client
+        _ => {
+            let hard = if HARD.get().copied().unwrap_or(false) {
+                "true"
+            } else {
+                "false"
+            };
+            format!(
+                r#"<script>{}("{}", {})</script>"#,
+                WEBSOCKET_FUNCTION, addr, hard
+            )
+        }
+    }
+}

src/templates/reload.js

@@ -0,0 +1,5 @@
+const meta = document.createElement("meta");
+meta.name = "live-server";
+meta.content = "reload";
+document.head.appendChild(meta);
+

src/templates/websocket.js

@@ -0,0 +1,107 @@
+(async (addr, hard) => {
+  addr = `ws://${addr}/live-server-ws`;
+  const sleep = (x) => new Promise((r) => setTimeout(r, x));
+  const preload = async (url, requireSuccess) => {
+    const resp = await fetch(url, { cache: "reload" }); // reset cache
+    if (requireSuccess && (!resp.ok || resp.status !== 200)) {
+      throw new Error();
+    }
+  };
+  /** Reset cache in link.href and strip scripts */
+  const preloadNode = (n, ps) => {
+    if (n.tagName === "SCRIPT" && n.src) {
+      ps.push(preload(n.src, false));
+      return;
+    }
+    if (n.tagName === "LINK" && n.href) {
+      ps.push(preload(n.href, false));
+      return;
+    }
+    let c = n.firstChild;
+    while (c) {
+      const nc = c.nextSibling;
+      preloadNode(c, ps);
+      c = nc;
+    }
+  };
+  let reloading = false; // if the page is currently being reloaded
+  let scheduled = false; // if another reload is scheduled while the page is being reloaded
+  async function reload() {
+    // schedule the reload for later if it's already reloading
+    if (reloading) {
+      scheduled = true;
+      return;
+    }
+    let ifr;
+    reloading = true;
+    while (true) {
+      scheduled = false;
+      const url = location.origin + location.pathname;
+      const promises = [];
+      preloadNode(document.head, promises);
+      preloadNode(document.body, promises);
+      await Promise.allSettled(promises);
+      try {
+        await new Promise((resolve) => {
+          ifr = document.createElement("iframe");
+          ifr.src = url + "?reload";
+          ifr.style.display = "none";
+          ifr.onload = resolve;
+          document.body.appendChild(ifr);
+        });
+      } catch {}
+      // reload only if the iframe loaded successfully
+      // with the reload payload. If the reload payload
+      // is absent, it probably means the server responded
+      // with a 404 page
+      const meta = ifr.contentDocument.head.lastChild;
+      if (
+        meta &&
+        meta.tagName === "META" &&
+        meta.name === "live-server" &&
+        meta.content === "reload"
+      ) {
+        // do reload if there's no further scheduled reload
+        // otherwise, let the next scheduled reload do the job
+        if (!scheduled) {
+          if (hard) {
+            location.reload();
+          } else {
+            reloading = false;
+            document.head.replaceWith(ifr.contentDocument.head);
+            document.body.replaceWith(ifr.contentDocument.body);
+            ifr.remove();
+            console.log("[Live Server] Reloaded");
+          }
+          return;
+        }
+      }
+      if (ifr) {
+        ifr.remove();
+      }
+      // wait for some time before trying again
+      await sleep(500);
+    }
+  }
+  let connectedInterrupted = false; // track if it's the first connection or a reconnection
+  while (true) {
+    try {
+      await new Promise((resolve) => {
+        const ws = new WebSocket(addr);
+        ws.onopen = () => {
+          console.log("[Live Server] Connection Established");
+          // on reconnection, refresh the page
+          if (connectedInterrupted) {
+            reload();
+          }
+        };
+        ws.onmessage = reload;
+        ws.onerror = () => ws.close();
+        ws.onclose = resolve;
+      });
+    } catch {}
+    connectedInterrupted = true;
+    await sleep(3000);
+    console.log("[Live Server] Reconnecting...");
+  }
+})