Files
loki/src/utils/abort_signal.rs
2025-10-07 10:45:42 -06:00

89 lines
2.1 KiB
Rust

use anyhow::Result;
use crossterm::event::{self, Event, KeyCode, KeyModifiers};
use std::{
sync::{
atomic::{AtomicBool, Ordering},
Arc,
},
time::Duration,
};
pub type AbortSignal = Arc<AbortSignalInner>;
pub struct AbortSignalInner {
ctrlc: AtomicBool,
ctrld: AtomicBool,
}
pub fn create_abort_signal() -> AbortSignal {
AbortSignalInner::new()
}
impl AbortSignalInner {
pub fn new() -> AbortSignal {
Arc::new(Self {
ctrlc: AtomicBool::new(false),
ctrld: AtomicBool::new(false),
})
}
pub fn aborted(&self) -> bool {
if self.aborted_ctrlc() {
return true;
}
if self.aborted_ctrld() {
return true;
}
false
}
pub fn aborted_ctrlc(&self) -> bool {
self.ctrlc.load(Ordering::SeqCst)
}
pub fn aborted_ctrld(&self) -> bool {
self.ctrld.load(Ordering::SeqCst)
}
pub fn reset(&self) {
self.ctrlc.store(false, Ordering::SeqCst);
self.ctrld.store(false, Ordering::SeqCst);
}
pub fn set_ctrlc(&self) {
self.ctrlc.store(true, Ordering::SeqCst);
}
pub fn set_ctrld(&self) {
self.ctrld.store(true, Ordering::SeqCst);
}
}
pub async fn wait_abort_signal(abort_signal: &AbortSignal) {
loop {
if abort_signal.aborted() {
break;
}
tokio::time::sleep(Duration::from_millis(25)).await;
}
}
pub fn poll_abort_signal(abort_signal: &AbortSignal) -> Result<bool> {
if event::poll(Duration::from_millis(25))? {
if let Event::Key(key) = event::read()? {
match key.code {
KeyCode::Char('c') if key.modifiers == KeyModifiers::CONTROL => {
abort_signal.set_ctrlc();
return Ok(true);
}
KeyCode::Char('d') if key.modifiers == KeyModifiers::CONTROL => {
abort_signal.set_ctrld();
return Ok(true);
}
_ => {}
}
}
}
Ok(false)
}