diff --git a/Cargo.toml b/Cargo.toml index 084908e..03242b7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ librime-sys = { version = "0.1.0", git = "https://github.com/lotem/librime-sys" dashmap = "5.4.0" regex = "1.7.1" ropey = "1.5.1" -tokio = { version = "1", features = ["io-util", "io-std", "macros", "rt-multi-thread", "net"] } +tokio = { version = "1", features = ["io-util", "io-std", "macros", "rt-multi-thread", "net", "signal"] } tower-lsp = "0.17.0" serde = { version = "1.0", features = ["derive"] } serde_json = "1.0.91" @@ -20,7 +20,6 @@ directories = "4.0.1" ouroboros = "0.15.5" thiserror = "1.0.38" once_cell = "1.17.0" -ctrlc = "3.2.5" [features] default = [] diff --git a/src/main.rs b/src/main.rs index d398422..f67c951 100644 --- a/src/main.rs +++ b/src/main.rs @@ -3,6 +3,7 @@ use std::{net::SocketAddr, str::FromStr}; use rime_ls::lsp::Backend; use rime_ls::rime::Rime; use tokio::net::{TcpListener, TcpStream}; +use tokio::sync::broadcast::{self, Receiver}; use tower_lsp::{LspService, Server}; async fn run_stdio() { @@ -29,6 +30,28 @@ async fn run_tcp_forever(bind_addr: SocketAddr) -> tokio::io::Result<()> { } } +async fn run(mut shutdown: Receiver<()>) -> Result<(), Box> { + let mut args = std::env::args(); + match args.nth(1).as_deref() { + None => { + tokio::select! { + _ = shutdown.recv() => (), + _ = run_stdio() => () + } + } + Some("--listen") => { + let addr = args.next().unwrap_or("127.0.0.1:9257".to_owned()); + let addr = SocketAddr::from_str(&addr)?; + tokio::select! { + _ = shutdown.recv() => Ok(()), + Err(e) = run_tcp_forever(addr) => Err(e), + }? + } + _ => usage(), + } + Ok(()) +} + fn usage() { println!("rime_ls v{}", env!("CARGO_PKG_VERSION")); println!("Usage: rime_ls [--listen ]") @@ -36,25 +59,21 @@ fn usage() { #[tokio::main] async fn main() { - // set handler to finalize rime - // TODO: it is ugly - ctrlc::set_handler(move || { + // tell things to shutdown + let (tx, rx) = broadcast::channel(1); + // waiting for ctrl-c + tokio::spawn(async move { + tokio::signal::ctrl_c().await.unwrap(); println!("Ctrl-C pressed."); - if Rime::is_initialized() { - Rime::global().destroy(); - } - std::process::exit(0); // 0? - }) - .expect("Error setting Ctrl-C handler"); - - let mut args = std::env::args(); - match args.nth(1).as_deref() { - None => run_stdio().await, - Some("--listen") => { - let addr = args.next().unwrap_or("127.0.0.1:9257".to_owned()); - let addr = SocketAddr::from_str(&addr).unwrap(); - run_tcp_forever(addr).await.unwrap(); - } - _ => usage(), + tx.send(()).unwrap(); + Ok::<(), tokio::io::Error>(()) + }); + // run + if let Err(e) = run(rx).await { + eprintln!("{e}"); + } + // finalize rime if necessary + if Rime::is_initialized() { + Rime::global().destroy(); } }