diff --git a/forest/Cargo.toml b/forest/Cargo.toml index c9b190ef708e..efe91aa0f001 100644 --- a/forest/Cargo.toml +++ b/forest/Cargo.toml @@ -16,3 +16,4 @@ async-log = "2.0.0" async-std = { version = "1.4.0", features = ["attributes"] } serde = { version = "1.0", features = ["derive"] } pretty_env_logger = "0.3" +ctrlc = "3.1.4" diff --git a/forest/src/main.rs b/forest/src/main.rs index 55f22b2121d5..b26ada310511 100644 --- a/forest/src/main.rs +++ b/forest/src/main.rs @@ -8,8 +8,35 @@ use async_std::task; use forest_libp2p::{get_keypair, Libp2pService}; use libp2p::identity::{ed25519, Keypair}; use log::{info, trace}; +use std::cell::RefCell; +use std::process; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; use utils::write_to_file; +// Blocks current thread until ctrl-c is received +fn block_until_sigint() { + let (ctrlc_send, ctrlc_oneshot) = futures::channel::oneshot::channel(); + let ctrlc_send_c = RefCell::new(Some(ctrlc_send)); + + let running = Arc::new(AtomicUsize::new(0)); + ctrlc::set_handler(move || { + let prev = running.fetch_add(1, Ordering::SeqCst); + if prev == 0 { + println!("Got interrupt, shutting down..."); + // Send sig int in channel to blocking task + if let Some(ctrlc_send) = ctrlc_send_c.try_borrow_mut().unwrap().take() { + ctrlc_send.send(()).expect("Error sending ctrl-c message"); + } + } else { + process::exit(0); + } + }) + .expect("Error setting Ctrl-C handler"); + + task::block_on(ctrlc_oneshot).unwrap(); +} + #[async_std::main] async fn main() { logger::setup_logger(); @@ -33,11 +60,17 @@ async fn main() { } }; - let lp2p_service = Libp2pService::new(&config.network, net_keypair); - - task::block_on(async move { - lp2p_service.run().await; + // Start libp2p service + let p2p_service = Libp2pService::new(&config.network, net_keypair); + let p2p_thread = task::spawn(async { + p2p_service.run().await; }); + // Block until ctrl-c is hit + block_until_sigint(); + + // Drop threads + drop(p2p_thread); + info!("Forest finish shutdown"); }