Repos / s4g / df315f13be
commit df315f13be65c391e35da8931835ef106d03fce8
Author: Nhân <hi@imnhan.com>
Date:   Sun Jul 9 01:10:17 2023 +0700

    make livereload work with multiple browser tabs

diff --git a/livereload/livereload.go b/livereload/livereload.go
index d440790..e7de673 100644
--- a/livereload/livereload.go
+++ b/livereload/livereload.go
@@ -12,6 +12,7 @@
 )
 
 const endpoint = "/_livereload"
+const clientIdHeader = "Client-Id"
 
 //go:embed livereload.html
 var lrScript []byte
@@ -20,13 +21,26 @@
 var dontReload = []byte("0")
 
 var state struct {
-	shouldReload bool
-	mut          sync.RWMutex
+	// Maps each client ID to whether they should reload on next ajax request.
+	//
+	// Client IDs are generated on client side so that an open tab's
+	// livereload feature keeps working even when the server is restarted.
+	clients map[string]bool
+	mut     sync.RWMutex
 }
 
 func init() {
-	lrScript = bytes.ReplaceAll(lrScript, []byte("{{LR_ENDPOINT}}"), []byte(endpoint))
-	lrScript = bytes.ReplaceAll(lrScript, []byte("{{SHOULD_RELOAD}}"), pleaseReload)
+	state.clients = make(map[string]bool)
+
+	lrScript = bytes.ReplaceAll(
+		lrScript, []byte("{{LR_ENDPOINT}}"), []byte(endpoint),
+	)
+	lrScript = bytes.ReplaceAll(
+		lrScript, []byte("{{PLEASE_RELOAD}}"), pleaseReload,
+	)
+	lrScript = bytes.ReplaceAll(
+		lrScript, []byte("{{CLIENT_ID_HEADER}}"), []byte(clientIdHeader),
+	)
 }
 
 // For html pages, insert a script tag to enable livereload
@@ -36,14 +50,26 @@ func Middleware(fsys writablefs.FS, f http.Handler) http.Handler {
 
 		// Handle AJAX endpoint
 		if path == endpoint {
+			clientId := r.Header.Get(clientIdHeader)
 			state.mut.RLock()
-			shouldReload := state.shouldReload
+			shouldReload, ok := state.clients[clientId]
 			state.mut.RUnlock()
 
+			// New client: add client to state, don't reload
+			if !ok {
+				//fmt.Println("New livereload client:", clientId)
+				state.mut.Lock()
+				state.clients[clientId] = false
+				state.mut.Unlock()
+				w.Write(dontReload)
+				return
+			}
+
+			// Existing client:
 			if shouldReload {
 				w.Write(pleaseReload)
 				state.mut.Lock()
-				state.shouldReload = false
+				state.clients[clientId] = false
 				state.mut.Unlock()
 			} else {
 				w.Write(dontReload)
@@ -76,8 +102,10 @@ func Middleware(fsys writablefs.FS, f http.Handler) http.Handler {
 
 func Trigger() {
 	state.mut.Lock()
-	state.shouldReload = true
-	state.mut.Unlock()
+	defer state.mut.Unlock()
+	for k := range state.clients {
+		state.clients[k] = true
+	}
 }
 
 func withLiveReload(original []byte) []byte {
diff --git a/livereload/livereload.html b/livereload/livereload.html
index 8f35245..d3feba0 100644
--- a/livereload/livereload.html
+++ b/livereload/livereload.html
@@ -1,9 +1,12 @@
 <script>
+  const clientId =
+    Date.now().toString(36) + Math.random().toString(36).substr(2);
+
   setInterval(() => {
-    fetch("{{LR_ENDPOINT}}")
+    fetch("{{LR_ENDPOINT}}", { headers: { "{{CLIENT_ID_HEADER}}": clientId } })
       .then((resp) => resp.text())
       .then((text) => {
-        if (text === "{{SHOULD_RELOAD}}") {
+        if (text === "{{PLEASE_RELOAD}}") {
           location.reload();
         }
       });