From be3f3f5de0cb46255586c398353e1693b321f54b Mon Sep 17 00:00:00 2001
From: clearloop <tianyi.gc@gmail.com>
Date: Sun, 31 Dec 2023 13:04:35 +0800
Subject: [PATCH] fix(serve): websocket implementation

---
 Cargo.lock                | 48 +++++++++++++++++++++++++++++++++++++++
 Cargo.toml                |  1 +
 blog/templates/layout.hbs |  1 -
 src/cmd/serve.rs          | 10 ++++----
 4 files changed, 55 insertions(+), 5 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 535cd9f..a7177e8 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -95,6 +95,17 @@ version = "1.0.77"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c9d19de80eff169429ac1e9f48fffb163916b448a44e8e046186232046d9e1f9"
 
+[[package]]
+name = "async-lock"
+version = "3.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7125e42787d53db9dd54261812ef17e937c95a51e4d291373b670342fa44310c"
+dependencies = [
+ "event-listener",
+ "event-listener-strategy",
+ "pin-project-lite",
+]
+
 [[package]]
 name = "autocfg"
 version = "1.1.0"
@@ -285,6 +296,15 @@ dependencies = [
  "windows-sys 0.48.0",
 ]
 
+[[package]]
+name = "concurrent-queue"
+version = "2.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d16048cd947b08fa32c24458a22f5dc5e835264f689f4f5653210c69fd107363"
+dependencies = [
+ "crossbeam-utils",
+]
+
 [[package]]
 name = "core-foundation-sys"
 version = "0.8.6"
@@ -334,6 +354,7 @@ name = "cydonia"
 version = "0.0.2"
 dependencies = [
  "anyhow",
+ "async-lock",
  "ccli",
  "chrono",
  "colored",
@@ -391,6 +412,27 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "event-listener"
+version = "4.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "84f2cdcf274580f2d63697192d744727b3198894b1bf02923643bf59e2c26712"
+dependencies = [
+ "concurrent-queue",
+ "parking",
+ "pin-project-lite",
+]
+
+[[package]]
+name = "event-listener-strategy"
+version = "0.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "958e4d70b6d5e81971bebec42271ec641e7ff4e170a6fa605f2b8a8b65cb97d3"
+dependencies = [
+ "event-listener",
+ "pin-project-lite",
+]
+
 [[package]]
 name = "eyre"
 version = "0.6.11"
@@ -954,6 +996,12 @@ version = "3.5.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "c1b04fb49957986fdce4d6ee7a65027d55d4b6d2265e5848bbb507b58ccfdb6f"
 
+[[package]]
+name = "parking"
+version = "2.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "bb813b8af86854136c6922af0598d719255ecb2179515e6e7730d468f05c9cae"
+
 [[package]]
 name = "percent-encoding"
 version = "2.3.1"
diff --git a/Cargo.toml b/Cargo.toml
index e9ab8f5..8812597 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -25,6 +25,7 @@ notify = { version = "6.1.1", optional = true }
 tokio = { version = "1.35.1", features = [ "rt-multi-thread" ], optional = true }
 warp = { version = "0.3.6", optional = true }
 futures = { version = "0.3.30", optional = true }
+async-lock = "3.2.0"
 
 [features]
 default = [ "cli" ]
diff --git a/blog/templates/layout.hbs b/blog/templates/layout.hbs
index 7d52e9b..bb4a93b 100644
--- a/blog/templates/layout.hbs
+++ b/blog/templates/layout.hbs
@@ -14,7 +14,6 @@
       <script>
        const socket = new WebSocket(`ws://${location.host}` + "/" + "{{{ livereload }}}");
        socket.onmessage = function (event) {
-         console.log(event)
          if (event.data === "reload") {
            socket.close();
            location.reload();
diff --git a/src/cmd/serve.rs b/src/cmd/serve.rs
index 6aad91f..7d85135 100644
--- a/src/cmd/serve.rs
+++ b/src/cmd/serve.rs
@@ -2,6 +2,7 @@
 
 use crate::{cmd::Watch, LIVERELOAD_ENDPOINT};
 use anyhow::Result;
+use async_lock::Mutex;
 use ccli::{clap, clap::Parser};
 use futures::{sink::SinkExt, FutureExt, StreamExt};
 use notify::Event;
@@ -9,7 +10,7 @@ use std::{
     net::{Ipv4Addr, TcpListener},
     sync::{
         mpsc::{self, Receiver},
-        Arc, Mutex,
+        Arc,
     },
 };
 use tokio::runtime::Runtime;
@@ -59,9 +60,10 @@ impl Serve {
             .map(|ws: Ws, rx: Arc<Mutex<Receiver<Event>>>| {
                 ws.on_upgrade(move |socket: WebSocket| async move {
                     let (mut tx, _) = socket.split();
-                    if let Ok(rx) = rx.lock() {
-                        if rx.recv().is_ok() {
-                            let _ = tx.send(Message::text("reload"));
+                    let rx = rx.lock().await;
+                    if rx.recv().is_ok() {
+                        if let Err(e) = tx.send(Message::text("reload")).await {
+                            tracing::error!("failed to send reload message: {}", e);
                         }
                     }
                 })