use std::collections::HashMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use crossbeam_channel::Receiver;
use wireshift_core::backend::{BackendCompletion, BackendSubmission};
use wireshift_core::{Error, Result};
use crate::ops::execute_descriptor;
#[derive(Debug)]
pub(crate) enum WorkerMessage {
Submit(BackendSubmission),
Shutdown,
}
#[derive(Debug)]
pub(crate) struct CancelRegistry {
pub(crate) states: Mutex<HashMap<u64, Arc<AtomicBool>>>,
}
impl CancelRegistry {
pub(crate) fn new() -> Self {
Self {
states: Mutex::new(HashMap::new()),
}
}
pub(crate) fn insert(&self, id: u64) -> Result<Arc<AtomicBool>> {
let flag = Arc::new(AtomicBool::new(false));
self.states
.lock()
.map_err(|_| {
Error::completion(
"cancel registry mutex was poisoned",
"avoid panicking while operations are in flight",
)
})?
.insert(id, flag.clone());
Ok(flag)
}
pub(crate) fn remove(&self, id: u64) -> Result<()> {
self.states
.lock()
.map_err(|_| {
Error::completion(
"cancel registry mutex was poisoned",
"avoid panicking while operations are being retired",
)
})?
.remove(&id);
Ok(())
}
pub(crate) fn cancel(&self, id: u64) -> Result<bool> {
let canceled = self
.states
.lock()
.map_err(|_| {
Error::completion(
"cancel registry mutex was poisoned",
"avoid panicking while cancellation is requested",
)
})?
.get(&id)
.cloned();
if let Some(flag) = canceled {
flag.store(true, Ordering::Relaxed);
Ok(true)
} else {
Ok(false)
}
}
}
pub(crate) fn worker_loop(
receiver: Receiver<WorkerMessage>,
completion_tx: std::sync::mpsc::Sender<BackendCompletion>,
cancel_registry: Arc<CancelRegistry>,
) {
loop {
match receiver.recv() {
Ok(WorkerMessage::Submit(submission)) => {
let result = execute_submission(submission, cancel_registry.clone());
let _ = cancel_registry.remove(result.id);
let _ = completion_tx.send(result);
}
Ok(WorkerMessage::Shutdown) | Err(_) => return,
}
}
}
fn execute_submission(
submission: BackendSubmission,
cancel_registry: Arc<CancelRegistry>,
) -> BackendCompletion {
let canceled = cancel_registry
.states
.lock()
.ok()
.and_then(|states| states.get(&submission.id).cloned());
let result = if let Some(flag) = canceled {
if flag.load(Ordering::Relaxed) {
Err(Error::canceled(
format!("request {} was canceled before execution", submission.id),
"stop canceling the request if it still needs to run",
))
} else {
execute_descriptor(submission.descriptor, flag, submission.timeout)
}
} else {
Err(Error::completion(
format!("request {} lost its cancellation state", submission.id),
"keep cancellation bookkeeping alive for the full request lifetime",
))
};
BackendCompletion {
id: submission.id,
result,
}
}