diff --git a/Cargo.toml b/Cargo.toml index 5e274d7..04faa6b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,6 +13,7 @@ license = "AGPL-3.0" aho-corasick = "0.7.14" base64 = "0.12.3" blake3 = "0.3.7" +ctrlc = { version = "3.1.7", features = [ "termination" ] } dashmap = "3.11.10" futures = "0.3.7" handlebars = "3.5.1" diff --git a/Dockerfile b/Dockerfile index 3e85827..006c8f6 100644 --- a/Dockerfile +++ b/Dockerfile @@ -21,7 +21,7 @@ ARG gitcommithash=HEAD RUN git checkout "$gitcommithash" \ && git reset --hard \ -&& echo "pub const GIT_VERSION: Option <&str> = Some (\"$(git-rev-parse HEAD)\");" > src/git_version.rs \ +&& echo "pub const GIT_VERSION: Option <&str> = Some (\"$(git rev-parse HEAD)\");" > src/git_version.rs \ && cargo test --release \ && cargo build --release --bin ptth_relay @@ -29,12 +29,13 @@ FROM debian:buster-slim as deploy RUN apt-get update \ && apt-get install -y libssl1.1 ca-certificates \ -&& apt-get upgrade +&& apt-get upgrade -y COPY --from=build /usr/src/target/release/ptth_relay /root -COPY ptth_relay.toml /root/ptth_relay.toml +COPY --from=build /usr/src/src/git_version.rs /root/git_version.rs +COPY --from=build /usr/src/ptth_handlebars /root/ptth_handlebars -COPY --from=build /usr/src/git_version.rs /root/git_version.rs +COPY ptth_relay.toml /root/ptth_relay.toml WORKDIR /root diff --git a/src/bin/ptth_relay.rs b/src/bin/ptth_relay.rs index 71314da..392e90b 100644 --- a/src/bin/ptth_relay.rs +++ b/src/bin/ptth_relay.rs @@ -1,8 +1,14 @@ use std::{ error::Error, fs::File, + sync::Arc, }; +use tokio::sync::oneshot; + +use ptth::relay; +use ptth::relay::RelayState; + #[tokio::main] async fn main () -> Result <(), Box > { use std::io::Read; @@ -21,5 +27,29 @@ async fn main () -> Result <(), Box > { eprintln! ("ptth_relay Git version: {:?}", ptth::git_version::GIT_VERSION); - ptth::relay::main (config_file).await + 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)), + Some (rx) + ).await } diff --git a/src/lib.rs b/src/lib.rs index dddcb52..2e1c271 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -100,7 +100,7 @@ mod tests { let relay_state_2 = relay_state.clone (); spawn (async move { - relay::run_relay (relay_state_2).await.unwrap (); + relay::run_relay (relay_state_2, None).await.unwrap (); }); assert! (relay_state.list_servers ().await.is_empty ()); diff --git a/src/relay/mod.rs b/src/relay/mod.rs index 8f4d602..cb84e4f 100644 --- a/src/relay/mod.rs +++ b/src/relay/mod.rs @@ -388,7 +388,8 @@ pub fn load_templates () } pub async fn run_relay ( - state: Arc + state: Arc , + shutdown_oneshot: Option > ) -> Result <(), Box > { @@ -421,21 +422,19 @@ pub async fn run_relay ( } }); - let server = Server::bind (&addr).serve (make_svc); + let server = Server::bind (&addr) + .serve (make_svc); - server.await?; + match shutdown_oneshot { + Some (rx) => server.with_graceful_shutdown (async { + rx.await.ok (); + }).await?, + None => server.await?, + }; Ok (()) } -pub async fn main (config_file: ConfigFile) --> Result <(), Box > -{ - let state = Arc::new (RelayState::from (&config_file)); - - run_relay (state).await -} - #[cfg (test)] mod tests {