Add monitoring to tokio tasks

1 files changed, 29 insertions, 8 deletions

diff --git a/src/main.rs b/src/main.rs
index dc16023..fdcdd47 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -8,6 +8,7 @@ use std::fs;
 use std::sync::Arc;
 use std::time::Duration;
 use tokio::sync::Mutex;
+use tokio::task::{Id, JoinSet};
 
 type StateMap = Arc<Mutex<HashMap<String, PrinterState>>>;
 
@@ -18,8 +19,11 @@ async fn main() {
     let content = fs::read_to_string("config.toml").expect("Couldn't read config.toml");
     let config = Config::load(&content);
     let state: StateMap = Arc::new(Mutex::new(HashMap::new()));
+    let mut tasks = JoinSet::new();
+    let mut task_names: HashMap<Id, String> = HashMap::new();
 
     for printer in config.printers {
+        let state = state.clone();
         match printer {
             Printer::Prusa {
                 name,
@@ -27,7 +31,9 @@ async fn main() {
                 api_key,
             } => {
                 tracing::info!(name, host, "Found Prusa");
-                tokio::spawn(poll_prusa(name, host, api_key, state.clone()));
+                let task_name = name.clone();
+                let handle = tasks.spawn(poll_prusa(name, host, api_key, state));
+                task_names.insert(handle.id(), task_name);
             }
             Printer::Bambu {
                 name,
@@ -36,17 +42,32 @@ async fn main() {
                 access_code,
             } => {
                 tracing::info!(name, host, "Found Bambu");
-                tokio::spawn(poll_bambu(
-                    name,
-                    host,
-                    serial_number,
-                    access_code,
-                    state.clone(),
-                ));
+                let task_name = name.clone();
+                let handle = tasks.spawn(poll_bambu(name, host, serial_number, access_code, state));
+                task_names.insert(handle.id(), task_name);
             }
         }
     }
 
+    tokio::spawn(async move {
+        while let Some(result) = tasks.join_next_with_id().await {
+            match result {
+                Ok((id, ())) => {
+                    let name = task_names.get(&id).map(String::as_str).unwrap_or("unknown");
+                    tracing::warn!(name, "Printer polling task exited unexpectedly");
+                }
+                Err(e) => {
+                    let name = task_names
+                        .get(&e.id())
+                        .map(String::as_str)
+                        .unwrap_or("unknown");
+                    tracing::error!(name, error = ?e, "A printer polling task panicked");
+                }
+            }
+        }
+        tracing::warn!("All printer polling tasks have exited");
+    });
+
     let app = Router::new().route("/", get(root)).with_state(state);
 
     let listener = tokio::net::TcpListener::bind("0.0.0.0:3000").await.unwrap();