Simplify shutdown with cancellation tokens (#601)

simonft · web-flow · commit 766f3461d31f · 2025-09-18T00:33:44.000+02:00
diff --git a/daemon/src/display/generic_framebuffer.rs b/daemon/src/display/generic_framebuffer.rs
@@ -9,9 +9,7 @@ use rayhunter::analysis::analyzer::EventType;
 
 use log::{error, info};
 use tokio::sync::mpsc::Receiver;
-use tokio::sync::oneshot;
-use tokio::sync::oneshot::error::TryRecvError;
-use tokio_util::task::TaskTracker;
+use tokio_util::{sync::CancellationToken, task::TaskTracker};
 
 use include_dir::{Dir, include_dir};
 
@@ -173,7 +171,7 @@ pub fn update_ui(
     task_tracker: &TaskTracker,
     config: &config::Config,
     mut fb: impl GenericFramebuffer,
-    mut ui_shutdown_rx: oneshot::Receiver<()>,
+    shutdown_token: CancellationToken,
     mut ui_update_rx: Receiver<DisplayState>,
 ) {
     static IMAGE_DIR: Dir<'_> = include_dir!("$CARGO_MANIFEST_DIR/images/");
@@ -204,13 +202,9 @@ pub fn update_ui(
             );
         }
         loop {
-            match ui_shutdown_rx.try_recv() {
-                Ok(_) => {
-                    info!("received UI shutdown");
-                    break;
-                }
-                Err(TryRecvError::Empty) => {}
-                Err(e) => panic!("error receiving shutdown message: {e}"),
+            if shutdown_token.is_cancelled() {
+                info!("received UI shutdown");
+                break;
             }
             match ui_update_rx.try_recv() {
                 Ok(state) => {
diff --git a/daemon/src/display/headless.rs b/daemon/src/display/headless.rs
@@ -1,6 +1,6 @@
 use log::info;
 use tokio::sync::mpsc::Receiver;
-use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
 use crate::config;
@@ -9,7 +9,7 @@ use crate::display::DisplayState;
 pub fn update_ui(
     _task_tracker: &TaskTracker,
     _config: &config::Config,
-    _ui_shutdown_rx: oneshot::Receiver<()>,
+    _shutdown_token: CancellationToken,
     _ui_update_rx: Receiver<DisplayState>,
 ) {
     info!("Headless mode, not spawning UI.");
diff --git a/daemon/src/display/orbic.rs b/daemon/src/display/orbic.rs
@@ -4,7 +4,7 @@ use crate::display::generic_framebuffer::{self, Dimensions, GenericFramebuffer};
 use async_trait::async_trait;
 
 use tokio::sync::mpsc::Receiver;
-use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
 const FB_PATH: &str = "/dev/fb0";
@@ -38,14 +38,14 @@ impl GenericFramebuffer for Framebuffer {
 pub fn update_ui(
     task_tracker: &TaskTracker,
     config: &config::Config,
-    ui_shutdown_rx: oneshot::Receiver<()>,
+    shutdown_token: CancellationToken,
     ui_update_rx: Receiver<DisplayState>,
 ) {
     generic_framebuffer::update_ui(
         task_tracker,
         config,
         Framebuffer,
-        ui_shutdown_rx,
+        shutdown_token,
         ui_update_rx,
     )
 }
diff --git a/daemon/src/display/tmobile.rs b/daemon/src/display/tmobile.rs
@@ -4,7 +4,7 @@
 /// DisplayState::WarningDetected { .. } => Signal LED slowly blinks red.
 use log::{error, info};
 use tokio::sync::mpsc;
-use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
 use std::time::Duration;
@@ -27,7 +27,7 @@ async fn stop_blinking(path: String) {
 pub fn update_ui(
     task_tracker: &TaskTracker,
     config: &config::Config,
-    mut ui_shutdown_rx: oneshot::Receiver<()>,
+    shutdown_token: CancellationToken,
     mut ui_update_rx: mpsc::Receiver<DisplayState>,
 ) {
     let mut invisible: bool = false;
@@ -40,13 +40,9 @@ pub fn update_ui(
         let mut last_state = DisplayState::Paused;
 
         loop {
-            match ui_shutdown_rx.try_recv() {
-                Ok(_) => {
-                    info!("received UI shutdown");
-                    break;
-                }
-                Err(oneshot::error::TryRecvError::Empty) => {}
-                Err(e) => panic!("error receiving shutdown message: {e}"),
+            if shutdown_token.is_cancelled() {
+                info!("received UI shutdown");
+                break;
             }
             match ui_update_rx.try_recv() {
                 Ok(new_state) => state = new_state,
diff --git a/daemon/src/display/tplink.rs b/daemon/src/display/tplink.rs
@@ -1,6 +1,6 @@
 use log::info;
 use tokio::sync::mpsc::Receiver;
-use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
 use crate::config;
@@ -11,7 +11,7 @@ use std::fs;
 pub fn update_ui(
     task_tracker: &TaskTracker,
     config: &config::Config,
-    ui_shutdown_rx: oneshot::Receiver<()>,
+    shutdown_token: CancellationToken,
     ui_update_rx: Receiver<DisplayState>,
 ) {
     let display_level = config.ui_level;
@@ -23,9 +23,9 @@ pub fn update_ui(
     // The alternative would be to make the entire initialization async
     if fs::exists(tplink_onebit::OLED_PATH).unwrap_or_default() {
         info!("detected one-bit display");
-        tplink_onebit::update_ui(task_tracker, config, ui_shutdown_rx, ui_update_rx)
+        tplink_onebit::update_ui(task_tracker, config, shutdown_token, ui_update_rx)
     } else {
         info!("fallback to framebuffer");
-        tplink_framebuffer::update_ui(task_tracker, config, ui_shutdown_rx, ui_update_rx)
+        tplink_framebuffer::update_ui(task_tracker, config, shutdown_token, ui_update_rx)
     }
 }
diff --git a/daemon/src/display/tplink_framebuffer.rs b/daemon/src/display/tplink_framebuffer.rs
@@ -2,13 +2,13 @@ use async_trait::async_trait;
 use std::os::fd::AsRawFd;
 use tokio::fs::OpenOptions;
 use tokio::io::AsyncWriteExt;
+use tokio_util::sync::CancellationToken;
 
 use crate::config;
 use crate::display::DisplayState;
 use crate::display::generic_framebuffer::{self, Dimensions, GenericFramebuffer};
 
 use tokio::sync::mpsc::Receiver;
-use tokio::sync::oneshot;
 use tokio_util::task::TaskTracker;
 
 const FB_PATH: &str = "/dev/fb0";
@@ -80,14 +80,14 @@ impl GenericFramebuffer for Framebuffer {
 pub fn update_ui(
     task_tracker: &TaskTracker,
     config: &config::Config,
-    ui_shutdown_rx: oneshot::Receiver<()>,
+    shutdown_token: CancellationToken,
     ui_update_rx: Receiver<DisplayState>,
 ) {
     generic_framebuffer::update_ui(
         task_tracker,
         config,
         Framebuffer,
-        ui_shutdown_rx,
+        shutdown_token,
         ui_update_rx,
     )
 }
diff --git a/daemon/src/display/tplink_onebit.rs b/daemon/src/display/tplink_onebit.rs
@@ -6,8 +6,7 @@ use crate::display::DisplayState;
 
 use log::{error, info};
 use tokio::sync::mpsc::Receiver;
-use tokio::sync::oneshot;
-use tokio::sync::oneshot::error::TryRecvError;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
 use std::time::Duration;
@@ -112,7 +111,7 @@ const STATUS_WARNING: &[u8] = pixelart! {
 pub fn update_ui(
     task_tracker: &TaskTracker,
     config: &config::Config,
-    mut ui_shutdown_rx: oneshot::Receiver<()>,
+    shutdown_token: CancellationToken,
     mut ui_update_rx: Receiver<DisplayState>,
 ) {
     let display_level = config.ui_level;
@@ -124,13 +123,9 @@ pub fn update_ui(
         let mut pixels = STATUS_SMILING;
 
         loop {
-            match ui_shutdown_rx.try_recv() {
-                Ok(_) => {
-                    info!("received UI shutdown");
-                    break;
-                }
-                Err(TryRecvError::Empty) => {}
-                Err(e) => panic!("error receiving shutdown message: {e}"),
+            if shutdown_token.is_cancelled() {
+                info!("received UI shutdown");
+                break;
             }
 
             match ui_update_rx.try_recv() {
diff --git a/daemon/src/display/uz801.rs b/daemon/src/display/uz801.rs
@@ -4,7 +4,7 @@
 /// DisplayState::WarningDetected => Signal LED is solid red.
 use log::{error, info};
 use tokio::sync::mpsc;
-use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
 use std::time::Duration;
@@ -27,7 +27,7 @@ async fn led_off(path: String) {
 pub fn update_ui(
     task_tracker: &TaskTracker,
     config: &config::Config,
-    mut ui_shutdown_rx: oneshot::Receiver<()>,
+    shutdown_token: CancellationToken,
     mut ui_update_rx: mpsc::Receiver<DisplayState>,
 ) {
     let mut invisible: bool = false;
@@ -41,13 +41,9 @@ pub fn update_ui(
         let mut last_update = std::time::Instant::now();
 
         loop {
-            match ui_shutdown_rx.try_recv() {
-                Ok(_) => {
-                    info!("received UI shutdown");
-                    break;
-                }
-                Err(oneshot::error::TryRecvError::Empty) => {}
-                Err(e) => panic!("error receiving shutdown message: {e}"),
+            if shutdown_token.is_cancelled() {
+                info!("received UI shutdown");
+                break;
             }
             match ui_update_rx.try_recv() {
                 Ok(new_state) => state = new_state,
diff --git a/daemon/src/display/wingtech.rs b/daemon/src/display/wingtech.rs
@@ -10,7 +10,7 @@ use crate::display::generic_framebuffer::{self, Dimensions, GenericFramebuffer};
 use async_trait::async_trait;
 
 use tokio::sync::mpsc::Receiver;
-use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
 const FB_PATH: &str = "/dev/fb0";
@@ -43,14 +43,14 @@ impl GenericFramebuffer for Framebuffer {
 pub fn update_ui(
     task_tracker: &TaskTracker,
     config: &config::Config,
-    ui_shutdown_rx: oneshot::Receiver<()>,
+    shutdown_token: CancellationToken,
     ui_update_rx: Receiver<DisplayState>,
 ) {
     generic_framebuffer::update_ui(
         task_tracker,
         config,
         Framebuffer,
-        ui_shutdown_rx,
+        shutdown_token,
         ui_update_rx,
     )
 }
diff --git a/daemon/src/key_input.rs b/daemon/src/key_input.rs
@@ -3,7 +3,7 @@ use std::time::{Duration, Instant};
 use tokio::fs::File;
 use tokio::io::AsyncReadExt;
 use tokio::sync::mpsc::Sender;
-use tokio::sync::oneshot;
+use tokio_util::sync::CancellationToken;
 use tokio_util::task::TaskTracker;
 
 use crate::config;
@@ -21,7 +21,7 @@ pub fn run_key_input_thread(
     task_tracker: &TaskTracker,
     config: &config::Config,
     diag_tx: Sender<DiagDeviceCtrlMessage>,
-    mut ui_shutdown_rx: oneshot::Receiver<()>,
+    cancellation_token: CancellationToken,
 ) {
     if config.key_input_mode == 0 {
         return;
@@ -43,7 +43,7 @@ pub fn run_key_input_thread(
 
         loop {
             tokio::select! {
-                _ = &mut ui_shutdown_rx => {
+               _ = cancellation_token.cancelled() => {
                     info!("received key input shutdown");
                     return;
                 }
diff --git a/daemon/src/main.rs b/daemon/src/main.rs
