From 2cf2897d31066a6962eff94e9c82163b833c97c6 Mon Sep 17 00:00:00 2001
From: "James O. D. Hunt" <james.o.hunt@intel.com>
Date: Mon, 22 Feb 2021 15:39:33 +0000
Subject: [PATCH] main: Use task list for stopping tasks

Maintain a list of tasks and wait on them all before main returns.

This is preparatory work for the agent shutdown: all tasks that are
started need to be added to the list. This aggregation makes it easier
to identify what needs to stop before the agent can exit cleanly.

Signed-off-by: James O. D. Hunt <james.o.hunt@intel.com>
---
 src/agent/src/main.rs | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/src/agent/src/main.rs b/src/agent/src/main.rs
index efac66d5cf..0aa027e38b 100644
--- a/src/agent/src/main.rs
+++ b/src/agent/src/main.rs
@@ -64,11 +64,13 @@ use uevent::watch_uevents;
 
 use std::sync::Mutex as SyncMutex;
 
+use futures::future::join_all;
 use futures::StreamExt as _;
 use rustjail::pipestream::PipeStream;
 use tokio::{
     signal::unix::{signal, SignalKind},
     sync::{oneshot::Sender, Mutex, RwLock},
+    task::JoinHandle,
 };
 use tokio_vsock::{Incoming, VsockListener, VsockStream};
 
@@ -124,6 +126,9 @@ async fn get_vsock_stream(fd: RawFd) -> Result<VsockStream> {
 async fn real_main() -> std::result::Result<(), Box<dyn std::error::Error>> {
     env::set_var("RUST_BACKTRACE", "full");
 
+    // List of tasks that need to be stopped for a clean shutdown
+    let mut tasks: Vec<JoinHandle<()>> = vec![];
+
     lazy_static::initialize(&SHELLS);
 
     lazy_static::initialize(&AGENT_CONFIG);
@@ -197,6 +202,8 @@ async fn real_main() -> std::result::Result<(), Box<dyn std::error::Error>> {
         let _ = tokio::io::copy(&mut reader, &mut stdout_writer).await;
     });
 
+    tasks.push(log_handle);
+
     let writer = unsafe { File::from_raw_fd(wfd) };
 
     // Recreate a logger with the log level get from "/proc/cmdline".
@@ -218,7 +225,13 @@ async fn real_main() -> std::result::Result<(), Box<dyn std::error::Error>> {
 
     start_sandbox(&logger, &config, init_mode).await?;
 
-    let _ = log_handle.await.unwrap();
+    let results = join_all(tasks).await;
+
+    for result in results {
+        if let Err(e) = result {
+            return Err(anyhow!(e).into());
+        }
+    }
 
     Ok(())
 }