transmit: Add SIGINT shutdown hook (re)

srt-transmit/Cargo.toml

@@ -20,6 +20,7 @@ bytes = "1.0"
 anyhow = "1"
 pretty_env_logger = { version = "0.4", default-features = false }
 futures = { version = "0.3", default-features = false, features = ["std", "async-await"] }
+signal-hook = "0.3.12"
 
 [dependencies.tokio]
 version = "1"
@@ -33,6 +34,7 @@ features = ["full"]
 [dev-dependencies]
 rand = "0.8"
 tokio-stream = "0.1"
+nix = "0.23"
 
 [features]
 default = []

srt-transmit/src/main.rs

@@ -8,6 +8,10 @@ use std::{
     path::Path,
     pin::Pin,
     process::exit,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
     task::{Context, Poll},
     time::{Duration, Instant},
 };
@@ -16,20 +20,23 @@ use anyhow::{anyhow, bail, format_err, Error};
 use bytes::Bytes;
 use clap::{App, Arg};
 use log::info;
+use signal_hook::{consts::SIGINT, flag::register};
 use url::{Host, Url};
 
 use futures::{
     future,
     prelude::*,
     ready,
     stream::{self, once, unfold, BoxStream},
-    try_join,
 };
 use tokio::{
     io::{AsyncRead, AsyncReadExt},
     net::TcpListener,
     net::TcpStream,
     net::UdpSocket,
+    select,
+    time::interval,
+    try_join,
 };
 use tokio_util::{codec::BytesCodec, codec::Framed, codec::FramedWrite, udp::UdpFramed};
 
@@ -580,14 +587,31 @@ impl Sink<Bytes> for MultiSinkFlatten {
     }
 }
 
+// A ZST error which represents a user interrupt via SIGINT.
+#[derive(Debug)]
+struct UserInterruptError;
+
+impl std::fmt::Display for UserInterruptError {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
+        "Interrupted by user via SIGINT.".fmt(f)
+    }
+}
+
+impl std::error::Error for UserInterruptError {}
+
 #[tokio::main]
 async fn main() {
     if let Err(e) = run().await {
         eprintln!(
             "Invalid settings detected: {}\n\nSee srt-transmit --help for more info",
             e
         );
-        exit(1);
+
+        if e.is::<UserInterruptError>() {
+            exit(130); // by convention, 128 + signal where SIGINT = 2
+        } else {
+            exit(1);
+        }
     }
 }
 
@@ -637,14 +661,37 @@ async fn run() -> Result<(), Error> {
 
     let mut sinks = MultiSinkFlatten::new(sink_streams.drain(..));
 
-    // poll sink and stream in parallel, only yielding when there is something ready for the sink and the stream is good.
-    while let (_, Some(stream)) = try_join!(
-        future::poll_fn(|cx| Pin::new(&mut sinks).poll_ready(cx)),
-        stream_stream.try_next()
-    )? {
-        // let a: () = &mut *stream;
-        sinks.send_all(&mut stream.map(Ok)).await?;
-    }
+    // setup SIGINT handling prerequisites
+    let mut sig_intval = interval(Duration::from_millis(500));
+    let sig_flg = Arc::new(AtomicBool::new(false));
+    register(SIGINT, sig_flg.clone()).unwrap();
+
+    // a future that only completes with an error when we get a sigint
+    // the flag is only checked every 500ms as to not hinder performance or spin.
+    let sig_fut = async {
+        while !sig_flg.load(Ordering::Relaxed) {
+            sig_intval.tick().await;
+        }
+        Err(UserInterruptError)
+    };
+
+    // a future that only completes once the transmit has finished or an error occurred.
+    // polls sink and stream in parallel, only yielding when there is something ready for the sink and the stream is good
+    let transmit_fut = async {
+        while let (_, Some(stream)) = try_join!(
+            future::poll_fn(|cx| Pin::new(&mut sinks).poll_ready(cx)),
+            stream_stream.try_next()
+        )? {
+            // let a: () = &mut *stream;
+            sinks.send_all(&mut stream.map(Ok)).await?;
+        }
+        Ok::<(), Error>(())
+    };
+
+    select! {
+        sig_res = sig_fut => sig_res?,
+        transmit_res = transmit_fut => transmit_res?,
+    };
 
     sinks.close().await?;
     Ok(())

srt-transmit/tests/timeout.rs

@@ -1,5 +1,11 @@
 use std::{env, path::PathBuf, process::Stdio, time::Instant};
 
+#[cfg(not(windows))]
+use nix::{
+    sys::signal::{kill, Signal},
+    unistd::Pid,
+};
+
 use bytes::Bytes;
 use futures::prelude::*;
 use tokio::{
@@ -109,3 +115,78 @@ async fn sender_timeout() {
     };
     futures::join!(sender, recvr);
 }
+
+// There doesn't seem to exist any crates for programmatically sending Ctrl+C events to a child process on Windows
+// within Rust. One avenue to explore would be the raw winapi binding - GenerateConsoleCtrlEvent.
+
+#[cfg(not(windows))]
+#[tokio::test]
+async fn sigint_termination_idle() {
+    let _ = pretty_env_logger::try_init();
+
+    let stranmsit_rs = find_stransmit_rs();
+    let mut a = Command::new(&stranmsit_rs)
+        .args(&["-", "-"])
+        .stdout(Stdio::piped())
+        .stdin(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()
+        .unwrap();
+
+    sleep(Duration::from_millis(500)).await;
+    let pid = Pid::from_raw(a.id().unwrap() as i32);
+    kill(pid, Signal::SIGINT).unwrap();
+
+    let out = a.wait().await.unwrap();
+    assert_eq!(out.code().unwrap(), 130);
+}
+
+#[cfg(not(windows))]
+#[tokio::test]
+async fn sigint_termination_work() {
+    let _ = pretty_env_logger::try_init();
+
+    let b = SrtSocket::builder().listen_on(1880);
+
+    let stranmsit_rs = find_stransmit_rs();
+    let mut a = Command::new(&stranmsit_rs)
+        .args(&["srt://localhost:1880", "-"])
+        .stdout(Stdio::piped())
+        .stderr(Stdio::piped())
+        .spawn()
+        .unwrap();
+
+    let sender = async move {
+        let mut b = b.await.unwrap();
+
+        let mut got_done = false;
+        for _ in 0..200 {
+            if b.send((Instant::now(), Bytes::from_static(b"asdf\n")))
+                .await
+                .is_err()
+            {
+                got_done = true;
+                break;
+            }
+            sleep(Duration::from_millis(100)).await;
+        }
+        assert!(got_done);
+    };
+
+    let recvr = async move {
+        let mut out = BufReader::new(a.stdout.as_mut().unwrap());
+        let mut line = String::new();
+
+        for _ in 0..10 {
+            out.read_line(&mut line).await.unwrap();
+        }
+
+        let pid = Pid::from_raw(a.id().unwrap() as i32);
+        kill(pid, Signal::SIGINT).unwrap();
+
+        let out = a.wait().await.unwrap();
+        assert_eq!(out.code().unwrap(), 130);
+    };
+
+    futures::join!(sender, recvr);
+}