mirror of
https://github.com/EFForg/rayhunter.git
synced 2026-05-01 09:59:57 -07:00
Simplify shutdown with cancellation tokens (#601)
This commit is contained in:
committed by
GitHub
parent
d30dd6fd9d
commit
766f3461d3
@@ -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) => {
|
||||
|
||||
@@ -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.");
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -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() {
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user