Simplify shutdown with cancellation tokens (#601)

This commit is contained in:
Simon Fondrie-Teitler
2025-09-17 18:33:44 -04:00
committed by GitHub
parent d30dd6fd9d
commit 766f3461d3
11 changed files with 48 additions and 89 deletions

View File

@@ -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) => {

View File

@@ -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.");

View File

@@ -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,
)
}

View File

@@ -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,

View File

@@ -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)
}
}

View File

@@ -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,
)
}

View File

@@ -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() {

View File

@@ -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,

View File

@@ -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,
)
}