glory_cli/service/
reload.rs

1use std::sync::Arc;
2use std::{fmt::Display, net::SocketAddr};
3
4use once_cell::sync::Lazy;
5use salvo::prelude::*;
6use salvo::websocket::{Message, WebSocket};
7use serde::Serialize;
8use tokio::{net::TcpStream, select, sync::RwLock, task::JoinHandle};
9
10use crate::config::Project;
11use crate::ext::sync::wait_for_socket;
12use crate::logger::GRAY;
13use crate::signal::Interrupt;
14use crate::signal::{ReloadSignal, ReloadType};
15
16static SITE_ADDR: Lazy<RwLock<SocketAddr>> = Lazy::new(|| RwLock::new(SocketAddr::new([127, 0, 0, 1].into(), 8000)));
17static CSS_LINK: Lazy<RwLock<String>> = Lazy::new(|| RwLock::new(String::default()));
18
19pub async fn spawn(proj: &Arc<Project>) -> JoinHandle<()> {
20    let proj = proj.clone();
21
22    let mut site_addr = SITE_ADDR.write().await;
23    *site_addr = proj.site.addr;
24    if let Some(file) = &proj.style.file {
25        let mut css_link = CSS_LINK.write().await;
26        // Always use `/` as separator in links
27        *css_link = file.site.components().map(|c| c.as_str()).collect::<Vec<_>>().join("/");
28    }
29
30    tokio::spawn(async move {
31        let _change = ReloadSignal::subscribe();
32
33        let reload_addr = proj.site.reload;
34
35        if TcpStream::connect(&reload_addr).await.is_ok() {
36            log::error!("Reload TCP port {reload_addr} already in use. You can set the port in the server integration's RenderOptions reload_port");
37            Interrupt::request_shutdown().await;
38
39            return;
40        }
41
42        log::debug!("Reload server started {}", GRAY.paint(reload_addr.to_string()));
43
44        let router = Router::with_path("/live_reload").get(live_reload);
45        let acceptor = TcpListener::new(reload_addr).bind().await;
46        match salvo::Server::new(acceptor).try_serve(router).await {
47            Ok(_) => log::debug!("Reload server stopped"),
48            Err(e) => log::error!("Reload {e}"),
49        }
50    })
51}
52
53#[handler]
54async fn live_reload(req: &mut Request, res: &mut Response) -> Result<(), StatusError> {
55    WebSocketUpgrade::new().upgrade(req, res, handle_socket).await
56}
57
58async fn handle_socket(mut stream: WebSocket) {
59    let mut rx = ReloadSignal::subscribe();
60    let mut int = Interrupt::subscribe_any();
61
62    log::trace!("Reload websocket connected");
63    tokio::spawn(async move {
64        loop {
65            select! {
66                res = rx.recv() =>{
67                    match res {
68                        Ok(ReloadType::Full) => {
69                            send_and_close(stream, BrowserMessage::all()).await;
70                            return
71                        }
72                        Ok(ReloadType::Style) => {
73                            send(&mut stream, BrowserMessage::css().await).await;
74                        },
75                        Ok(ReloadType::ViewPatches(data)) => {
76                            send(&mut stream, BrowserMessage::view(data)).await;
77                        }
78                        Err(e) => log::debug!("Reload recive error {e}")
79                    }
80                }
81                _ = int.recv(), if Interrupt::is_shutdown_requested().await => {
82                    log::trace!("Reload websocket closed");
83                    return
84                },
85            }
86        }
87    });
88}
89
90async fn send(stream: &mut WebSocket, msg: BrowserMessage) {
91    let site_addr = *SITE_ADDR.read().await;
92    if !wait_for_socket("Reload", site_addr).await {
93        log::warn!(r#"Reload could not send "{msg}" to websocket"#);
94    }
95
96    let text = serde_json::to_string(&msg).unwrap();
97    match stream.send(Message::text(text)).await {
98        Err(e) => {
99            log::debug!("Reload could not send {msg} due to {e}");
100        }
101        Ok(_) => {
102            log::debug!(r#"Reload sent "{msg}" to browser"#);
103        }
104    }
105}
106
107async fn send_and_close(mut stream: WebSocket, msg: BrowserMessage) {
108    send(&mut stream, msg).await;
109    let _ = stream.close().await;
110    log::trace!("Reload websocket closed");
111}
112
113#[derive(Serialize)]
114struct BrowserMessage {
115    css: Option<String>,
116    view: Option<String>,
117    all: bool,
118}
119
120impl BrowserMessage {
121    async fn css() -> Self {
122        let link = CSS_LINK.read().await.clone();
123        if link.is_empty() {
124            log::error!("Reload internal error: sending css reload but no css file is set.");
125        }
126        Self {
127            css: Some(link),
128            view: None,
129            all: false,
130        }
131    }
132
133    fn view(data: String) -> Self {
134        Self {
135            css: None,
136            view: Some(data),
137            all: false,
138        }
139    }
140
141    fn all() -> Self {
142        Self {
143            css: None,
144            view: None,
145            all: true,
146        }
147    }
148}
149
150impl Display for BrowserMessage {
151    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
152        if let Some(css) = &self.css {
153            write!(f, "reload {}", css)
154        } else {
155            write!(f, "reload all")
156        }
157    }
158}