use clap::Parser;
use log::LevelFilter;
use nix::{
sys::signal::{self, Signal},
unistd::Pid,
};
use notify::{RecommendedWatcher, RecursiveMode, Watcher};
use serde::Serialize;
use signal_hook::{
consts::signal::{SIGHUP, SIGUSR1},
iterator::Signals,
};
use std::{
env,
io::{self, Write},
path::PathBuf,
process::Command,
sync::{Arc, Mutex, mpsc},
thread,
time::Duration,
};
#[derive(Parser, Debug)]
pub struct DevServer {
#[arg(short, long, env, default_value = "src")]
watch: Option<Vec<PathBuf>>,
#[arg(short, long, env)]
bin: Option<PathBuf>,
#[arg(short, long)]
cwd: Option<PathBuf>,
#[arg(short, long)]
release: bool,
#[arg(short, long)]
example: Option<String>,
#[arg(short, long, default_value = "SIGTERM")]
signal: Signal,
}
#[derive(Debug, Clone, Serialize)]
#[serde(tag = "type")]
pub enum Event {
BinaryChanged,
Rebuild,
Restarted,
BuildSuccess,
CompileError { error: String },
}
impl DevServer {
fn determine_bin(&self) -> PathBuf {
if let Some(ref bin) = self.bin {
bin.canonicalize().unwrap()
} else {
let metadata = Command::new("cargo")
.current_dir(self.cwd.clone().unwrap())
.args(["metadata", "--format-version", "1"])
.output()
.unwrap();
let value: serde_json::Value = serde_json::from_slice(&metadata.stdout).unwrap();
let target_dir =
PathBuf::from(value.get("target_directory").unwrap().as_str().unwrap());
let root = value
.get("resolve")
.unwrap()
.get("root")
.unwrap()
.as_str()
.unwrap()
.split(' ')
.next()
.unwrap();
let target_dir = target_dir.join(if self.release { "release" } else { "debug" });
let target_dir = if let Some(example) = &self.example {
target_dir.join("examples").join(example)
} else {
target_dir.join(root)
};
target_dir.canonicalize().unwrap()
}
}
pub fn run(mut self) {
env_logger::Builder::new()
.filter_level(LevelFilter::Debug)
.init();
let cwd = self
.cwd
.get_or_insert_with(|| env::current_dir().unwrap())
.clone();
let bin = self.determine_bin();
let mut run = Command::new(&bin);
run.current_dir(&cwd);
let mut build = Command::new("cargo");
let mut args = vec!["build", "--color=always"];
if self.release {
args.push("--release");
}
let signal = self.signal;
if let Some(example) = &self.example {
args.push("--example");
args.push(example);
self.watch
.get_or_insert_with(Vec::new)
.push(cwd.join("examples"));
}
build.args(&args[..]);
build.current_dir(&cwd);
let mut child = run.spawn().unwrap();
let child_id = Arc::new(Mutex::new(child.id()));
let (tx, rx) = mpsc::channel();
let (broadcast_tx, broadcast_rx) = async_broadcast::broadcast(10);
{
let tx = tx.clone();
thread::spawn(move || {
let mut signals = Signals::new([SIGHUP, SIGUSR1]).unwrap();
loop {
for signal in signals.pending() {
if let SIGHUP = signal as libc::c_int {
tx.send(Event::BinaryChanged).unwrap();
}
}
}
});
}
thread::spawn(move || {
let (t, r) = mpsc::channel();
let mut watcher = RecommendedWatcher::new(t, Default::default()).unwrap();
if let Some(watches) = self.watch {
for watch in watches {
let watch = if watch.is_relative() {
cwd.join(watch)
} else {
watch
};
let watch = watch.canonicalize().unwrap();
log::info!("watching {:?}", &watch);
watcher.watch(&watch, RecursiveMode::Recursive).unwrap();
}
}
log::info!("watching {:?}", &bin);
watcher.watch(&bin, RecursiveMode::NonRecursive).unwrap();
while let Ok(Ok(m)) = r.recv() {
for path in m.paths {
if let Ok(path) = path.canonicalize() {
if path == bin {
tx.send(Event::BinaryChanged).unwrap();
} else {
tx.send(Event::Rebuild).unwrap();
}
}
}
}
});
{
let child_id = child_id.clone();
let broadcaster = broadcast_tx.clone();
thread::spawn(move || {
loop {
child.wait().unwrap();
log::info!("shut down, restarting");
child = run.spawn().unwrap();
*child_id.lock().unwrap() = child.id();
thread::sleep(Duration::from_millis(500));
async_io::block_on(broadcaster.broadcast_direct(Event::Restarted)).unwrap();
}
});
}
{
let broadcaster = broadcast_tx.clone();
thread::spawn(move || {
loop {
let event = rx.recv().unwrap();
async_io::block_on(broadcaster.broadcast_direct(event.clone())).unwrap();
match event {
Event::BinaryChanged => {
log::info!("attempting to send {}", &signal);
signal::kill(
Pid::from_raw((*child_id.lock().unwrap()).try_into().unwrap()),
signal,
)
.unwrap();
}
Event::Rebuild => {
log::info!("building...");
let output = build.output();
match output {
Ok(ok) => {
if ok.status.success() {
log::debug!("{}", String::from_utf8_lossy(&ok.stdout[..]));
async_io::block_on(
broadcaster.broadcast_direct(Event::BuildSuccess),
)
.ok();
} else {
io::stderr().write_all(&ok.stderr).unwrap();
async_io::block_on(
broadcaster.broadcast_direct(Event::CompileError {
error: ansi_to_html::convert(
&String::from_utf8_lossy(&ok.stderr),
)
.unwrap(),
}),
)
.ok();
}
}
Err(e) => {
eprintln!("{:?}", e);
}
}
}
_ => {}
}
}
});
}
proxy_app::run(format!("http://{}:{}", "localhost", "8080"), broadcast_rx);
}
}
mod proxy_app {
use super::Event;
use trillium::{Conn, KnownHeaderName, State};
use trillium_client::Client;
use trillium_html_rewriter::{
HtmlRewriter,
html::{Settings, element, html_content::ContentType},
};
use trillium_proxy::Proxy;
use trillium_router::Router;
use trillium_smol::ClientConfig;
use trillium_websockets::{WebSocket, WebSocketConn};
pub fn run(proxy: String, rx: async_broadcast::Receiver<Event>) {
static PORT: u16 = 8082;
let client = Client::new(ClientConfig::default().with_nodelay(true));
trillium_smol::config()
.without_signals()
.with_port(PORT)
.run((
Router::new()
.get("/_dev_server.js", |conn: Conn| async move {
conn.with_response_header(
KnownHeaderName::ContentType,
"application/javascript",
)
.ok(include_str!("./dev_server.js"))
})
.get(
"/_dev_server.ws",
(
State::new(rx),
WebSocket::new(|mut conn: WebSocketConn| async move {
let mut rx = conn
.take_state::<async_broadcast::Receiver<Event>>()
.unwrap();
while let Ok(message) = rx.recv_direct().await {
if let Err(e) = conn.send_json(&message).await {
log::error!("{:?}", e);
return;
}
}
}),
),
),
Proxy::new(client, &*proxy),
HtmlRewriter::new(|| Settings {
element_content_handlers: vec![element!("body", |el| {
el.append(
r#"<script src="/_dev_server.js"></script>"#,
ContentType::Html,
);
Ok(())
})],
..Settings::new_send()
}),
));
}
}