handle SIGTERM

Author: Martin Špirk

.gitignore

@@ -1,2 +1,3 @@
 /target
 /.idea
+/.qtcreator

Cargo.lock

@@ -30,6 +30,7 @@ futures-rustls = "0.26.0"
 rustls-pemfile = "2.2.0"
 rustls-platform-verifier = "0.7.0"
 async-channel = "2.3.1"
+async-signal = "0.2.13"
 serialport = { version = "4.7.1", optional = true }
 
 # For local development

Cargo.toml

@@ -584,7 +584,7 @@ fn timeout_param_to_duration(timeout_ms: u64) -> Option<Duration> {
     }
 }
 
-async fn make_burst_call(opts: &Opts) -> Result {
+async fn make_burst_call(opts: &Opts, shutdown_receiver: Receiver<()>) -> Result {
     if opts.method.is_none() {
         return Err("--method parameter missing".into());
     }
@@ -610,16 +610,41 @@ async fn make_burst_call(opts: &Opts) -> Result {
         taskno: i32,
         count: i32,
         timeout: Option<Duration>,
-        user_agent: String
+        user_agent: String,
+        shutdown_receiver: Receiver<()>,
     ) {
         println!("Starting burst task #{taskno}, {count} calls of {path}:{method}");
-        let (mut frame_reader, mut frame_writer) = login(&url, user_agent).await.unwrap();
+        let (mut frame_reader, mut frame_writer) = match login(&url, user_agent).await {
+            Ok(conn) => conn,
+            Err(err) => {
+                error!(target: "Burst", "Burst task #{taskno} failed to login: {err}");
+                return;
+            }
+        };
         for _ in 0..count {
+            if shutdown_receiver.try_recv().is_ok() {
+                info!(target: "Burst", "Shutdown requested, stopping burst task #{taskno}");
+                return;
+            }
             let rqid = frame_writer
                 .send_request(&path, &method, param.clone())
                 .await
                 .unwrap();
-            receive_response(&mut frame_reader, rqid, timeout).await.unwrap();
+            let response_fut = receive_response(&mut frame_reader, rqid, timeout).fuse();
+            let shutdown_fut = shutdown_receiver.recv().fuse();
+            futures::pin_mut!(response_fut, shutdown_fut);
+            match select(response_fut, shutdown_fut).await {
+                futures::future::Either::Left((response, _)) => {
+                    if let Err(err) = response {
+                        error!(target: "Burst", "Burst task #{taskno} request failed: {err}");
+                        return;
+                    }
+                }
+                futures::future::Either::Right((_, _)) => {
+                    info!(target: "Burst", "Shutdown requested, stopping burst task #{taskno}");
+                    return;
+                }
+            }
         }
         println!("Burst task #{taskno} finished, after {count} calls made successfully.");
     }
@@ -635,7 +660,8 @@ async fn make_burst_call(opts: &Opts) -> Result {
                 taskno + 1,
                 nmsg,
                 timeout_param_to_duration(opts.timeout),
-                opts.extract_user_agent()
+                opts.extract_user_agent(),
+                shutdown_receiver.clone(),
             ))
         })
         .collect::<FuturesUnordered<_>>()
@@ -660,14 +686,48 @@ fn split_quoted(s: &str) -> Vec<&str> {
 }
 #[derive(Debug)]
 struct Tunnel {
+    tunid: Option<u64>,
     create_rqid: RqId,
     write_rqid: RqId,
+    close_rqid: Option<RqId>,
     frame_sender: Sender<RpcFrame>,
 }
+
+enum TunnelEvent {
+    ClientConnectionClosed(u64),
+}
+
+async fn send_tunnel_close(tunnel_path: &str, tunid: u64, write_frame_sender: &mut Sender<RpcFrame>) -> Option<RqId> {
+    let mut rq = RpcMessage::new_request(format!("{tunnel_path}/{tunid}"), "close");
+    let rqid = RpcMessage::next_request_id();
+    rq.set_request_id(rqid);
+    match rq.to_frame() {
+        Ok(frame) => {
+            if let Err(err) = write_frame_sender.send(frame).await {
+                error!(target: "Tunnel", "Failed to send tunnel close request: {err}");
+                None
+            } else {
+                Some(rqid)
+            }
+        }
+        Err(err) => {
+            error!(target: "Tunnel", "Failed to build tunnel close request: {err}");
+            None
+        }
+    }
+}
+
+fn remove_pre_tunid_closed_tunnels(tunnels: &mut Vec<Tunnel>) {
+    tunnels.extract_if(.., |tunnel| tunnel.tunid.is_none() && tunnel.frame_sender.is_closed()).for_each(|tunnel| {
+        debug!(target: "Tunnel", "Removing tunnel task that finished before tunid assignment: {:?}", tunnel);
+    });
+}
+
 async fn start_tunnel_server(
     mut broker_frame_reader: BoxedFrameReader,
     mut broker_frame_writer: BoxedFrameWriter,
     opts: &Opts,
+    shutdown_receiver: Receiver<()>,
 ) -> Result {
     if opts.tunnel_path.is_none() {
         warn!("Using default .app/tunnel endpoint. This is usually not what you want. Set tunnel path to the broker you want to create a tunnel to.");
@@ -701,30 +761,48 @@ async fn start_tunnel_server(
     let local_host = local_host.to_owned();
 
     let mut tunnels: Vec<Tunnel> = Vec::new();
+    let mut tunnel_tasks: Vec<smol::Task<()>> = Vec::new();
+    let mut shutting_down = false;
 
     debug!(target: "Tunnel", "Starting TCP server on {local_host}:{local_port}");
     let listener = TcpListener::bind(format!("{local_host}:{local_port}")).await?;
     let mut incoming = listener.incoming();
 
-    let (write_frame_sender, write_frame_receiver) = async_channel::unbounded();
+    let (tunnel_event_sender, tunnel_event_receiver) = async_channel::unbounded();
+    let (mut write_frame_sender, write_frame_receiver) = async_channel::unbounded();
+    let (shutdown_timeout_sender, shutdown_timeout_receiver) = async_channel::bounded::<()>(1);
     loop {
+        remove_pre_tunid_closed_tunnels(&mut tunnels);
+        tunnel_tasks.retain(|task| !task.is_finished());
+        if shutting_down && tunnels.iter().all(|tunnel| tunnel.close_rqid.is_none()) {
+            info!(target: "Tunnel", "All tunnel close responses received, finishing tunnel loop");
+            break;
+        }
         select! {
             stream = incoming.next().fuse() => {
                 if let Some(stream) = stream {
+                    if shutting_down {
+                        drop(stream?);
+                        continue;
+                    }
                     let stream = stream?;
                     debug!(target: "Tunnel", "New connection from {:?}", stream.local_addr());
                     let create_rqid = RpcMessage::next_request_id();
                     let write_rqid = RpcMessage::next_request_id();
                     let (read_frame_sender, read_frame_receiver) = async_channel::unbounded();
-                    let tunnel = Tunnel {create_rqid, write_rqid, frame_sender: read_frame_sender};
+                    let tunnel = Tunnel {tunid: None, create_rqid, write_rqid, close_rqid: None, frame_sender: read_frame_sender};
                     tunnels.push(tunnel);
                     let read_frame_receiver = read_frame_receiver.clone();
                     let write_frame_sender = write_frame_sender.clone();
                     let remote_host_port = remote_host_port.clone();
                     let tunnel_path = tunnel_path.clone();
-                    spawn_and_log_error(async move {
-                        handle_tunnel_socket(stream, remote_host_port, tunnel_path, create_rqid, write_rqid, read_frame_receiver, write_frame_sender.clone()).await.map_err(|e | e.to_string())
+                    let tunnel_event_sender = tunnel_event_sender.clone();
+                    let task = smol::spawn(async move {
+                        if let Err(err) = handle_tunnel_socket(stream, remote_host_port, tunnel_path, create_rqid, write_rqid, read_frame_receiver, write_frame_sender.clone(), tunnel_event_sender).await {
+                            error!(target: "Tunnel", "Tunnel task finished with error: {err}");
+                        }
                     });
+                    tunnel_tasks.push(task);
                 } else {
                     break;
                 }
@@ -733,22 +811,88 @@ async fn start_tunnel_server(
                 match frame {
                     Ok(frame) => {
                         let rqid = frame.request_id().unwrap_or(0);
-                        for tunnel in &tunnels {
+                        if tunnels.iter().any(|tunnel| tunnel.close_rqid == Some(rqid)) {
+                            tunnels.extract_if(.., |tunnel| tunnel.close_rqid == Some(rqid)).for_each(|tunnel| {
+                                debug!(target: "Tunnel", "Tunnel close ACK received, removing tunnel {:?}", tunnel);
+                            });
+                            if shutting_down && tunnels.iter().all(|tunnel| tunnel.close_rqid.is_none()) {
+                                info!(target: "Tunnel", "All tunnel close responses received, finishing tunnel loop");
+                                break;
+                            }
+                            continue;
+                        }
+                        for tunnel in &mut tunnels {
                             if tunnel.write_rqid == rqid || tunnel.create_rqid == rqid {
+                                if tunnel.create_rqid == rqid && tunnel.tunid.is_none() {
+                                    if let Ok(rpcmsg) = frame.to_rpcmesage()
+                                        && let Ok(shvrpc::rpcmessage::Response::Success(val)) = rpcmsg.response()
+                                        && let Ok(tunid) = val.as_str().parse::<u64>()
+                                    {
+                                        tunnel.tunid = Some(tunid);
+                                    }
+                                }
                                 tunnel.frame_sender.send(frame).await?;
                                 break;
                             }
                         }
-                        tunnels.extract_if(.., |tunnel| tunnel.frame_sender.is_closed()).for_each(|tunnel| {
-                            debug!(target: "Tunnel", "Removing closed tunnel {:?}", tunnel);
-                        });
                     }
                     Err(e) => {
                         error!("Get response receiver error: {e}");
                         break;
                     }
                 }
             }
+            event = tunnel_event_receiver.recv().fuse() => {
+                match event {
+                    Ok(TunnelEvent::ClientConnectionClosed(tunid)) => {
+                        if let Some(tunnel) = tunnels.iter_mut().find(|tunnel| tunnel.tunid == Some(tunid)) {
+                            if tunnel.close_rqid.is_none() {
+                                if let Some(close_rqid) = send_tunnel_close(&tunnel_path, tunid, &mut write_frame_sender).await {
+                                    tunnel.close_rqid = Some(close_rqid);
+                                } else {
+                                    tunnels.extract_if(.., |candidate| candidate.tunid == Some(tunid)).for_each(|removed| {
+                                        debug!(target: "Tunnel", "Tunnel close send failed, removing tunnel {:?}", removed);
+                                    });
+                                }
+                            }
+                        }
+                    }
+                    Err(err) => {
+                        debug!(target: "Tunnel", "Tunnel event receiver closed: {err}");
+                    }
+                }
+            }
+            shutdown = shutdown_receiver.recv().fuse() => {
+                if shutdown.is_ok() && !shutting_down {
+                    info!(target: "Tunnel", "Received shutdown signal, shutting down TCP tunnel server");
+                    shutting_down = true;
+                    let shutdown_timeout_sender = shutdown_timeout_sender.clone();
+                    smol::spawn(async move {
+                        smol::Timer::after(std::time::Duration::from_secs(2)).await;
+                        let _ = shutdown_timeout_sender.send(()).await;
+                    }).detach();
+                    if tunnels.is_empty() {
+                        info!(target: "Tunnel", "No active tunnels, finishing tunnel loop");
+                        break;
+                    }
+                    for tunnel in &mut tunnels {
+                        if tunnel.close_rqid.is_some() {
+                            continue;
+                        }
+                        if let Some(tunid) = tunnel.tunid {
+                            if let Some(rqid) = send_tunnel_close(&tunnel_path, tunid, &mut write_frame_sender).await {
+                                tunnel.close_rqid = Some(rqid);
+                            }
+                        }
+                    }
+                }
+            }
+            timeout = shutdown_timeout_receiver.recv().fuse() => {
+                if shutting_down && timeout.is_ok() {
+                    info!(target: "Tunnel", "Shutdown deadline reached, finishing tunnel loop");
+                    break;
+                }
+            }
             frame = write_frame_receiver.recv().fuse() => {
                 match frame {
                     Ok(frame) => {
@@ -762,10 +906,17 @@ async fn start_tunnel_server(
             }
         }
     }
+
+    for task in tunnel_tasks.drain(..) {
+        if !task.is_finished() {
+            task.cancel().await;
+        }
+    }
+
     Ok(())
 }
 
-async fn handle_tunnel_socket(stream: TcpStream, remote_host_port: String, tunnel_path: String, create_rqid: RqId, write_rqid: RqId, read_frame_receiver: Receiver<RpcFrame>, mut write_frame_sender: Sender<RpcFrame>) -> Result {
+async fn handle_tunnel_socket(stream: TcpStream, remote_host_port: String, tunnel_path: String, create_rqid: RqId, write_rqid: RqId, read_frame_receiver: Receiver<RpcFrame>, mut write_frame_sender: Sender<RpcFrame>, tunnel_event_sender: Sender<TunnelEvent>) -> Result {
     let tunid = {
         debug!(target: "Tunnel", "Creating tunnel");
         let tun_opts = Map::from([("host".into(), (remote_host_port).into())]);
@@ -799,20 +950,30 @@ async fn handle_tunnel_socket(stream: TcpStream, remote_host_port: String, tunne
         rq.set_seqno(seqno_to_write);
         seqno_to_write += 1;
         debug!(target: "Tunnel", "Starting data exchange");
-        write_frame_sender.send(rq.to_frame()?).await?;
+        if let Err(err) = write_frame_sender.send(rq.to_frame()?).await {
+            return Err(err.into());
+        }
     };
+
     let (mut sock_reader, mut sock_writer) = stream.split();
     let mut sock_read_buff: [u8; 1024] = [0; 1024];
     loop {
         select! {
             n = sock_reader.read(&mut sock_read_buff).fuse() => {
-                let n = n?;
-                if n == 0 {
-                    debug!(target: "Tunnel", "Tunnel client socket closed");
-                    break;
+                match n {
+                    Ok(n) => {
+                        if n == 0 {
+                            debug!(target: "Tunnel", "Tunnel client socket closed");
+                            let _ = tunnel_event_sender.send(TunnelEvent::ClientConnectionClosed(tunid)).await;
+                            return Ok(());
+                        }
+                        let data = &sock_read_buff[0 .. n];
+                        seqno_to_write = process_socket_to_broker_data(&tunnel_path, tunid, seqno_to_write, write_rqid, data, &mut write_frame_sender).await?;
+                    }
+                    Err(err) => {
+                        return Err(err.into());
+                    }
                 }
-                let data = &sock_read_buff[0 .. n];
-                seqno_to_write = process_socket_to_broker_data(&tunnel_path, tunid, seqno_to_write, write_rqid, data, &mut write_frame_sender).await?;
             }
             frame = read_frame_receiver.recv().fuse() => {
                 match frame {
@@ -866,14 +1027,40 @@ async fn process_socket_to_broker_data(tunnel_path: &str, tunid: u64, seqno_to_w
 }
 
 pub async fn try_main(opts: Opts) -> Result {
+    let (shutdown_sender, shutdown_receiver) = async_channel::bounded::<()>(1);
+    smol::spawn(async move {
+        match async_signal::Signals::new(&[
+            async_signal::Signal::Term,
+            async_signal::Signal::Int,
+        ]) {
+            Ok(mut signals) => {
+                if signals.next().await.is_some() {
+                    let _ = shutdown_sender.send(()).await;
+                }
+            }
+            Err(err) => {
+                error!("Failed to initialize signal handling: {err}");
+            }
+        }
+    }).detach();
+
     if opts.burst.is_some() {
-        return make_burst_call(&opts).await;
+        return make_burst_call(&opts, shutdown_receiver).await;
     }
     let (frame_reader, frame_writer) = login(&opts.url, opts.extract_user_agent()).await?;
     let res = if opts.tunnel.is_some() {
-        start_tunnel_server(frame_reader, frame_writer, &opts).await
+        start_tunnel_server(frame_reader, frame_writer, &opts, shutdown_receiver).await
     } else {
-        make_call(frame_reader, frame_writer, &opts).await
+        let call_fut = make_call(frame_reader, frame_writer, &opts).fuse();
+        let shutdown_fut = shutdown_receiver.recv().fuse();
+        futures::pin_mut!(call_fut, shutdown_fut);
+        match select(call_fut, shutdown_fut).await {
+            futures::future::Either::Left((call_res, _)) => call_res,
+            futures::future::Either::Right((_, _)) => {
+                info!("Received shutdown signal, exiting");
+                Ok(())
+            }
+        }
     };
     match res {
         Ok(_) => Ok(()),