refactor: gracefully shutdown

Cargo.toml

@@ -12,15 +12,14 @@ 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"
 directories = "4.0.1"
 ouroboros = "0.15.5"
 thiserror = "1.0.38"
 once_cell = "1.17.0"
-ctrlc = "3.2.5"
 
 [features]
 default = []

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,32 +30,50 @@ async fn run_tcp_forever(bind_addr: SocketAddr) -> tokio::io::Result<()> {
     }
 }
 
+async fn run(mut shutdown: Receiver<()>) -> Result<(), Box<dyn std::error::Error>> {
+    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 <bind_addr>]")
 }
 
 #[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();
     }
 }