wireshift_fallback/
pool.rs

1use std::collections::HashMap;
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::sync::{Arc, Mutex};
4
5use crossbeam_channel::Receiver;
6use wireshift_core::backend::{BackendCompletion, BackendSubmission};
7use wireshift_core::{Error, Result};
8
9use crate::ops::execute_descriptor;
10
11#[derive(Debug)]
12pub(crate) enum WorkerMessage {
13    Submit(BackendSubmission),
14    Shutdown,
15}
16
17#[derive(Debug)]
18pub(crate) struct CancelRegistry {
19    pub(crate) states: Mutex<HashMap<u64, Arc<AtomicBool>>>,
20}
21
22impl CancelRegistry {
23    pub(crate) fn new() -> Self {
24        Self {
25            states: Mutex::new(HashMap::new()),
26        }
27    }
28
29    pub(crate) fn insert(&self, id: u64) -> Result<Arc<AtomicBool>> {
30        let flag = Arc::new(AtomicBool::new(false));
31        self.states
32            .lock()
33            .map_err(|_| {
34                Error::completion(
35                    "cancel registry mutex was poisoned",
36                    "avoid panicking while operations are in flight",
37                )
38            })?
39            .insert(id, flag.clone());
40        Ok(flag)
41    }
42
43    pub(crate) fn remove(&self, id: u64) -> Result<()> {
44        self.states
45            .lock()
46            .map_err(|_| {
47                Error::completion(
48                    "cancel registry mutex was poisoned",
49                    "avoid panicking while operations are being retired",
50                )
51            })?
52            .remove(&id);
53        Ok(())
54    }
55
56    pub(crate) fn cancel(&self, id: u64) -> Result<bool> {
57        let canceled = self
58            .states
59            .lock()
60            .map_err(|_| {
61                Error::completion(
62                    "cancel registry mutex was poisoned",
63                    "avoid panicking while cancellation is requested",
64                )
65            })?
66            .get(&id)
67            .cloned();
68        if let Some(flag) = canceled {
69            flag.store(true, Ordering::Relaxed);
70            Ok(true)
71        } else {
72            Ok(false)
73        }
74    }
75}
76
77pub(crate) fn worker_loop(
78    receiver: Receiver<WorkerMessage>,
79    completion_tx: std::sync::mpsc::Sender<BackendCompletion>,
80    cancel_registry: Arc<CancelRegistry>,
81) {
82    loop {
83        match receiver.recv() {
84            Ok(WorkerMessage::Submit(submission)) => {
85                let result = execute_submission(submission, cancel_registry.clone());
86                let _ = cancel_registry.remove(result.id);
87                let _ = completion_tx.send(result);
88            }
89            Ok(WorkerMessage::Shutdown) | Err(_) => return,
90        }
91    }
92}
93
94fn execute_submission(
95    submission: BackendSubmission,
96    cancel_registry: Arc<CancelRegistry>,
97) -> BackendCompletion {
98    let canceled = cancel_registry
99        .states
100        .lock()
101        .ok()
102        .and_then(|states| states.get(&submission.id).cloned());
103    let result = if let Some(flag) = canceled {
104        if flag.load(Ordering::Relaxed) {
105            Err(Error::canceled(
106                format!("request {} was canceled before execution", submission.id),
107                "stop canceling the request if it still needs to run",
108            ))
109        } else {
110            execute_descriptor(submission.descriptor, flag, submission.timeout)
111        }
112    } else {
113        Err(Error::completion(
114            format!("request {} lost its cancellation state", submission.id),
115            "keep cancellation bookkeeping alive for the full request lifetime",
116        ))
117    };
118    BackendCompletion {
119        id: submission.id,
120        result,
121    }
122}
wireshift_fallback/pool.rs

wireshift_fallback/
pool.rs
