diff --git a/src/bin/ptth_relay.rs b/src/bin/ptth_relay.rs index a0ab66c..dd3ebd2 100644 --- a/src/bin/ptth_relay.rs +++ b/src/bin/ptth_relay.rs @@ -3,42 +3,17 @@ use std::{ sync::Arc, }; -use tokio::sync::oneshot; - -mod load_toml; - use ptth::relay; use ptth::relay::RelayState; #[tokio::main] async fn main () -> Result <(), Box > { - let config_file = load_toml::load ("config/ptth_relay.toml"); + let config_file = ptth::load_toml::load ("config/ptth_relay.toml"); eprintln! ("ptth_relay Git version: {:?}", ptth::git_version::GIT_VERSION); - let rx = { - let (tx, rx) = oneshot::channel::<()> (); - - // I have to put the tx into a Cell here so that if Ctrl-C gets - // called multiple times, we won't send multiple shutdowns to the - // oneshot channel. (Which would be a compile error) - - let tx = Some (tx); - let tx = std::cell::Cell::new (tx); - - ctrlc::set_handler (move ||{ - let tx = tx.replace (None); - - if let Some (tx) = tx { - tx.send (()).unwrap (); - } - }).expect ("Error setting Ctrl-C handler"); - - rx - }; - relay::run_relay ( Arc::new (RelayState::from (&config_file)), - rx + ptth::graceful_shutdown::init () ).await } diff --git a/src/bin/ptth_server.rs b/src/bin/ptth_server.rs index 495c4b6..790dc97 100644 --- a/src/bin/ptth_server.rs +++ b/src/bin/ptth_server.rs @@ -4,9 +4,6 @@ use std::{ }; use structopt::StructOpt; -use tokio::sync::oneshot; - -mod load_toml; #[derive (Debug, StructOpt)] struct Opt { @@ -18,31 +15,10 @@ struct Opt { async fn main () -> Result <(), Box > { tracing_subscriber::fmt::init (); - let config_file = load_toml::load ("config/ptth_server.toml"); - - let rx = { - let (tx, rx) = oneshot::channel::<()> (); - - // I have to put the tx into a Cell here so that if Ctrl-C gets - // called multiple times, we won't send multiple shutdowns to the - // oneshot channel. (Which would be a compile error) - - let tx = Some (tx); - let tx = std::cell::Cell::new (tx); - - ctrlc::set_handler (move ||{ - let tx = tx.replace (None); - - if let Some (tx) = tx { - tx.send (()).unwrap (); - } - }).expect ("Error setting Ctrl-C handler"); - - rx - }; + let config_file = ptth::load_toml::load ("config/ptth_server.toml"); ptth::server::run_server ( config_file, - rx + ptth::graceful_shutdown::init () ).await } diff --git a/src/graceful_shutdown.rs b/src/graceful_shutdown.rs new file mode 100644 index 0000000..63de1b3 --- /dev/null +++ b/src/graceful_shutdown.rs @@ -0,0 +1,23 @@ +use std::cell::Cell; +use tokio::sync::oneshot; + +pub fn init () -> oneshot::Receiver <()> { + let (tx, rx) = oneshot::channel::<()> (); + + // I have to put the tx into a Cell here so that if Ctrl-C gets + // called multiple times, we won't send multiple shutdowns to the + // oneshot channel. (Which would be a compile error) + + let tx = Some (tx); + let tx = Cell::new (tx); + + ctrlc::set_handler (move ||{ + let tx = tx.replace (None); + + if let Some (tx) = tx { + tx.send (()).unwrap (); + } + }).expect ("Error setting Ctrl-C handler"); + + rx +} diff --git a/src/lib.rs b/src/lib.rs index 82724e3..16d239c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -11,6 +11,8 @@ pub const PTTH_MAGIC_HEADER: &str = "X-PTTH-2LJYXWC4"; // test stuff like spawn them both in the same process pub mod git_version; +pub mod graceful_shutdown; +pub mod load_toml; pub mod relay; pub mod server; diff --git a/src/bin/load_toml.rs b/src/load_toml.rs similarity index 100% rename from src/bin/load_toml.rs rename to src/load_toml.rs