Implement process cleanup for the broker

sidebus-broker/src/main.rs

@@ -4,9 +4,10 @@ mod vsock;
 
 use bus::SharedHostedBus;
 use clap::Parser;
-use futures::TryFutureExt;
+use futures::{TryFutureExt, stream::FuturesUnordered};
 use std::{path::PathBuf, sync::Arc};
 use tokio::{net::UnixListener, process::Command, sync::Mutex};
+use tokio_stream::StreamExt as _;
 use tracing::{Instrument, debug, error, info_span};
 use zbus::names::WellKnownName;
 
@@ -66,6 +67,7 @@ async fn main() -> eyre::Result<()> {
     let (priv_bus, _, mut priv_lst) = new_hosted_bus().await?;
 
     let mut server_tasks = tokio::task::JoinSet::new();
+    let mut child_procs = Vec::new();
 
     if let Some(dir_path) = cli.debug_access {
         if !dir_path.is_dir() {
@@ -89,33 +91,38 @@ async fn main() -> eyre::Result<()> {
 
     std::fs::create_dir_all(&cli.runtime_dir)?;
 
-    let _xps = priv_bus
+    child_procs.push(
+        priv_bus
             .clone()
             .spawn_external_client(
                 Command::new(env!("BIN_XDG_PERMISSION_STORE"))
                     .env("XDG_RUNTIME_DIR", cli.runtime_dir.as_os_str())
                     .kill_on_drop(true),
             )
-        .await?;
+            .await?,
+    );
 
     let impl_permission_store =
         WellKnownName::from_static_str("org.freedesktop.impl.portal.PermissionStore")?.into();
     priv_lst.wait_for_acquisition(impl_permission_store).await?;
 
-    let _xdp = priv_bus
+    child_procs.push(
+        priv_bus
             .clone()
             .spawn_external_client(
                 Command::new(env!("BIN_XDG_DOCUMENT_PORTAL"))
                     .env("XDG_RUNTIME_DIR", cli.runtime_dir.as_os_str())
                     .kill_on_drop(true),
             )
-        .await?;
+            .await?,
+    );
 
     let portal_documents =
         WellKnownName::from_static_str("org.freedesktop.portal.Documents")?.into();
     priv_lst.wait_for_acquisition(portal_documents).await?;
 
-    let _vfs = Command::new(env!("BIN_VIRTIOFSD"))
+    child_procs.push(
+        Command::new(env!("BIN_VIRTIOFSD"))
             .args(&[
                 "--shared-dir",
                 cli.runtime_dir.join("doc").to_str().unwrap(),
@@ -130,8 +137,8 @@ async fn main() -> eyre::Result<()> {
             ])
             .env("XDG_RUNTIME_DIR", cli.runtime_dir.as_os_str())
             .kill_on_drop(true)
-        .spawn();
+            .spawn()?,
-    // TODO: die when it exits
+    );
 
     let vm_bus_conn = vm_bus.lock().await.connect_channel(false).await?;
     let priv_bus_conn = priv_bus.lock().await.connect_channel(false).await?;
@@ -206,6 +213,21 @@ async fn main() -> eyre::Result<()> {
         );
     }
 
-    let _ = server_tasks.join_all().await;
+    let mut waiter = child_procs
+        .iter_mut()
+        .map(|child| child.wait())
+        .collect::<FuturesUnordered<_>>();
+    debug!("starting..");
+    tokio::select! {
+        _ = server_tasks.join_all() => debug!("server tasks ended"),
+        res = waiter.next() => debug!(?res, "child process terminated"),
+        _ = tokio::signal::ctrl_c() => debug!("interrupt signal"),
+    };
+    drop(waiter);
+    for mut child in child_procs {
+        if let Err(e) = child.kill().await {
+            error!(?e, "could not kill process");
+        }
+    }
     Ok(())
 }