Add Ctrl-C handler

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"

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
 

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 <dyn Error>> {
 	use std::io::Read;
@@ -21,5 +27,29 @@ async fn main () -> Result <(), Box <dyn Error>> {
 	
 	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
 }

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 ());

src/relay/mod.rs

@@ -388,7 +388,8 @@ pub fn load_templates ()
 }
 
 pub async fn run_relay (
-	state: Arc <RelayState>
+	state: Arc <RelayState>,
+	shutdown_oneshot: Option <tokio::sync::oneshot::Receiver <()>>
 ) 
 	-> Result <(), Box <dyn Error>> 
 {
@@ -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 <dyn Error>> 
-{
-	let state = Arc::new (RelayState::from (&config_file));
-	
-	run_relay (state).await
-}
-
 #[cfg (test)]
 mod tests {