Add a Dream.livereload middleware

The code is adapted from dream-livereload, itself adapted from the w-live-reload example. We update the w-live-reload to use the newly introduced middleware.
aantron · May 13, 2023 · d5125a4 · d5125a4
1 parent e14bd91
commit d5125a4
Show file tree

Hide file tree

Showing 7 changed files with 103 additions and 143 deletions.
diff --git a/dream.opam b/dream.opam
@@ -60,11 +60,13 @@ depends: [
   "fmt" {>= "0.8.7"}  # `Italic.
   "graphql_parser"
   "graphql-lwt"
+  "lambdasoup" {>= "0.6.1"}
   "lwt"
   "lwt_ppx" {>= "1.2.2"}
   "lwt_ssl"
   "logs" {>= "0.5.0"}
   "magic-mime"
+  "markup" {>= "1.0.2"}
   "mirage-clock" {>= "3.0.0"}  # now_d_ps : unit -> int * int64.
   "mirage-crypto" {>= "0.8.1"}  # AES-256-GCM.
   "mirage-crypto-rng"
@@ -85,7 +87,6 @@ depends: [
   "crunch" {with-test}
   "js_of_ocaml" {with-test}
   "js_of_ocaml-ppx" {with-test}
-  "lambdasoup" {with-test}
   "ppx_expect" {with-test & >= "v0.15.0"}  # Formatting changes.
   "ppx_yojson_conv" {with-test}
   "reason" {with-test}

diff --git a/example/w-live-reload/README.md b/example/w-live-reload/README.md
@@ -2,83 +2,20 @@
 
 <br>
 
-This example shows a simple live reloading setup. It works by injecting a script
-into the `<head>` of HTML documents. The script opens a WebSocket back to the
-server. If the WebSocket gets closed, the script tries to reconnect. When the
-server comes back up, the client is able to reconnect, and reloads itself.
+This example shows a simple live reloading setup using the `Dream.livereload`
+middleware. It works by injecting a script into the `<head>` of HTML documents.
+The script opens a WebSocket back to the server. If the WebSocket gets closed,
+the script tries to reconnect. When the server comes back up, the client is able
+to reconnect and reloads itself.
 
-```js
-var socketUrl = "ws://" + location.host + "/_live-reload"
-var socket = new WebSocket(socketUrl);
-
-socket.onclose = function(event) {
-  const intervalMs = 100;
-  const attempts = 100;
-  let attempt = 0;
-
-  function reload() {
-    ++attempt;
-
-    if(attempt > attempts) {
-      console.error("Could not reconnect to server");
-      return;
-    }
-
-    reconnectSocket = new WebSocket(socketUrl);
-
-    reconnectSocket.onerror = function(event) {
-      setTimeout(reload, intervalMs);
-    };
-
-    reconnectSocket.onopen = function(event) {
-      location.reload();
-    };
-  };
-
-  reload();
-};
-```
-
-The injection is done by a small middleware:
-
-```ocaml
-
-let inject_live_reload_script inner_handler request =
-  let%lwt response = inner_handler request in
-
-  match Dream.header "Content-Type" response with
-  | Some "text/html; charset=utf-8" ->
-    let%lwt body = Dream.body response in
-    let soup =
-      Markup.string body
-      |> Markup.parse_html ~context:`Document
-      |> Markup.signals
-      |> Soup.from_signals
-    in
-
-    begin match Soup.Infix.(soup $? "head") with
-    | None ->
-      Lwt.return response
-    | Some head ->
-      Soup.create_element "script" ~inner_text:live_reload_script
-      |> Soup.append_child head;
-      response
-      |> Dream.with_body (Soup.to_string soup)
-      |> Lwt.return
-    end
-
-  | _ ->
-    Lwt.return response
-```
-
-The example server just wraps a single page at `/` with the middleware. The page
+The example server just wraps a single page at `/` with the `Dream.livereload` middleware. The page
 displays a tag that changes each time it is loaded:
 
 ```ocaml
 let () =
   Dream.run
   @@ Dream.logger
-  @@ inject_live_reload_script
+  @@ Dream.livereload
   @@ Dream.router [
 
     Dream.get "/" (fun _ ->
@@ -87,11 +24,6 @@ let () =
       |> Printf.sprintf "Good morning, world! Random tag: %s"
       |> Dream.html);
 
-    Dream.get "/_live-reload" (fun _ ->
-      Dream.websocket (fun socket ->
-        let%lwt _ = Dream.receive socket in
-        Dream.close_websocket socket));
-
   ]
 ```
 
@@ -129,7 +61,7 @@ changes.
 **See also:**
 
 - [**`k-websocket`**](../k-websocket#files) introduces WebSockets.
-- [**`w-fswatch`**](../w-fswatch#files) rebuilds and restarts a server each
+- [**`w-watch`**](../w-watch#files) rebuilds and restarts a server each
   time its source code changes.
 
 <br>

diff --git a/example/w-live-reload/live_reload.ml b/example/w-live-reload/live_reload.ml
@@ -1,67 +1,7 @@
-let live_reload_script = {js|
-
-var socketUrl = "ws://" + location.host + "/_live-reload"
-var socket = new WebSocket(socketUrl);
-
-socket.onclose = function(event) {
-  const intervalMs = 100;
-  const attempts = 100;
-  let attempt = 0;
-
-  function reload() {
-    ++attempt;
-
-    if(attempt > attempts) {
-      console.error("Could not reconnect to server");
-      return;
-    }
-
-    reconnectSocket = new WebSocket(socketUrl);
-
-    reconnectSocket.onerror = function(event) {
-      setTimeout(reload, intervalMs);
-    };
-
-    reconnectSocket.onopen = function(event) {
-      location.reload();
-    };
-  };
-
-  reload();
-};
-
-|js}
-
-let inject_live_reload_script inner_handler request =
-  let%lwt response = inner_handler request in
-
-  match Dream.header response "Content-Type" with
-  | Some "text/html; charset=utf-8" ->
-    let%lwt body = Dream.body response in
-    let soup =
-      Markup.string body
-      |> Markup.parse_html ~context:`Document
-      |> Markup.signals
-      |> Soup.from_signals
-    in
-
-    begin match Soup.Infix.(soup $? "head") with
-    | None ->
-      Lwt.return response
-    | Some head ->
-      Soup.create_element "script" ~inner_text:live_reload_script
-      |> Soup.append_child head;
-      Dream.set_body response (Soup.to_string soup);
-      Lwt.return response
-    end
-
-  | _ ->
-    Lwt.return response
-
 let () =
   Dream.run
   @@ Dream.logger
-  @@ inject_live_reload_script
+  @@ Dream.livereload ()
   @@ Dream.router [
 
     Dream.get "/" (fun _ ->
@@ -70,9 +10,4 @@ let () =
       |> Printf.sprintf "Good morning, world! Random tag: %s"
       |> Dream.html);
 
-    Dream.get "/_live-reload" (fun _ ->
-      Dream.websocket (fun socket ->
-        let%lwt _ = Dream.receive socket in
-        Dream.close_websocket socket));
-
   ]
diff --git a/src/dream.ml b/src/dream.ml
@@ -17,6 +17,7 @@ module Formats = Dream_pure.Formats
 module Graphql = Dream__graphql.Graphql
 module Helpers = Dream__server.Helpers
 module Http = Dream__http.Http
+module Livereload = Dream__server.Livereload
 module Message = Dream_pure.Message
 module Method = Dream_pure.Method
 module Origin_referrer_check = Dream__server.Origin_referrer_check
@@ -216,6 +217,7 @@ let csrf_tag = Tag.csrf_tag ~now
 
 let no_middleware = Message.no_middleware
 let pipeline = Message.pipeline
+let livereload = Livereload.livereload
 
 
 

diff --git a/src/dream.mli b/src/dream.mli
@@ -1304,6 +1304,17 @@ val no_middleware : middleware
         Dream.no_middleware
     ]} *)
 
+val livereload : ?script:string -> ?path:string -> unit -> middleware
+(** Adds live reloading to your Dream application.
+
+    It works by injecting a script in the HTML pages sent to clients that will
+    initiate a WebSocket.
+
+    When the server restarts, the WebSocket connection is lost, at which point,
+    the client will try to reconnect every 500ms for 5s. If within these 5s the
+    client is able to reconnect to the server, it will trigger a reload of the
+    page. *)
+
 val pipeline : middleware list -> middleware
 (** Combines a sequence of middlewares into one, such that these two lines are
     equivalent:

diff --git a/src/server/dune b/src/server/dune
@@ -6,6 +6,7 @@
   dream.cipher
   dream-pure
   fmt
+  lambdasoup
   logs
   lwt
   magic-mime

diff --git a/src/server/livereload.ml b/src/server/livereload.ml
@@ -0,0 +1,78 @@
+(* This file is part of Dream, released under the MIT license. See LICENSE.md
+   for details, or visit https://github.com/aantron/dream.
+
+   Copyright 2021 Anton Bachin *)
+
+module Message = Dream_pure.Message
+
+let default_script ?(retry_interval_ms = 500) ?(max_retry_ms = 5000)
+    ?(route = "/_livereload") () =
+  Printf.sprintf
+    {js|
+var socketUrl = "ws://" + location.host + "%s"
+var s = new WebSocket(socketUrl);
+
+s.onopen = function(even) {
+  console.log("WebSocket connection open.");
+};
+
+s.onclose = function(even) {
+  console.log("WebSocket connection closed.");
+  const innerMs = %i;
+  const maxMs = %i;
+  const maxAttempts = Math.round(maxMs / innerMs);
+  let attempts = 0;
+  function reload() {
+    attempts++;
+    if(attempts > maxAttempts) {
+      console.error("Could not reconnect to dev server.");
+      return;
+    }
+
+    s2 = new WebSocket(socketUrl);
+
+    s2.onerror = function(event) {
+      setTimeout(reload, innerMs);
+    };
+
+    s2.onopen = function(event) {
+      location.reload();
+    };
+  };
+  reload();
+};
+
+s.onerror = function(event) {
+  console.error("WebSocket error observed:", event);
+};
+|js}
+    route retry_interval_ms max_retry_ms
+
+let livereload ?(script = default_script ()) ?(path = "/_livereload") ()
+    (next_handler : Message.request -> Message.response Lwt.t)
+    (request : Message.request) : Message.response Lwt.t =
+  match Message.target request with
+  | target when target = path ->
+    Helpers.websocket (fun socket ->
+        Lwt.bind (Helpers.receive socket) (fun _ ->
+            Message.close_websocket socket))
+  | _ -> (
+    let%lwt response = next_handler request in
+    match Message.header response "Content-Type" with
+    | Some "text/html" | Some "text/html; charset=utf-8" -> (
+      let%lwt body = Message.body response in
+      let soup =
+        Markup.string body
+        |> Markup.parse_html ~context:`Document
+        |> Markup.signals
+        |> Soup.from_signals
+      in
+      let open Soup.Infix in
+      match soup $? "head" with
+      | None -> Lwt.return response
+      | Some head ->
+        Soup.create_element "script" ~inner_text:script
+        |> Soup.append_child head;
+        Message.set_body response (Soup.to_string soup);
+        Lwt.return response)
+    | _ -> Lwt.return response)