coven_foundation/
bounded_workers.rs

1//! Owned synchronous workers with bounded, cancellation-aware FIFO admission.
2
3use std::num::NonZeroUsize;
4use std::sync::{Arc, Mutex};
5use tokio::sync::{mpsc, oneshot};
6
7type Job<State> = Box<dyn FnOnce(&mut State) + Send>;
8
9/// Each worker retains one state. All workers dequeue from the same bounded
10/// channel, so work is never stranded behind a busy worker's private queue.
11pub struct BoundedWorkers<State> {
12    inner: Arc<Workers<State>>,
13}
14
15impl<State> Clone for BoundedWorkers<State> {
16    fn clone(&self) -> Self {
17        Self {
18            inner: self.inner.clone(),
19        }
20    }
21}
22
23impl<State: Send + 'static> BoundedWorkers<State> {
24    /// Start all workers before returning a capability that can admit work.
25    pub fn start(states: Vec<State>, capacity: NonZeroUsize, name: &str) -> std::io::Result<Self> {
26        assert!(
27            !states.is_empty(),
28            "a worker pool needs at least one worker"
29        );
30        let (sender, receiver) = mpsc::channel::<Job<State>>(capacity.get());
31        let receiver = Arc::new(Mutex::new(receiver));
32        let mut joins = Vec::with_capacity(states.len());
33        for (index, mut state) in states.into_iter().enumerate() {
34            let receiver = receiver.clone();
35            joins.push(
36                std::thread::Builder::new()
37                    .name(format!("{name}-{index}"))
38                    .spawn(move || {
39                        loop {
40                            // Release this lock before running the job. It protects only
41                            // dequeue; every worker executes against its own state.
42                            let job = receiver
43                                .lock()
44                                .expect("worker queue mutex poisoned")
45                                .blocking_recv();
46                            match job {
47                                Some(job) => job(&mut state),
48                                None => break,
49                            }
50                        }
51                    })?,
52            );
53        }
54        Ok(Self {
55            inner: Arc::new(Workers {
56                sender: Some(sender),
57                joins,
58            }),
59        })
60    }
61
62    /// Wait for bounded admission and completion. Cancelling before execution
63    /// discards the closure; running work finishes, but its reply is discarded.
64    /// A closure panic resumes on the caller and leaves the worker available.
65    pub async fn call<F, R>(&self, operation: F) -> R
66    where
67        F: FnOnce(&mut State) -> R + Send + 'static,
68        R: Send + 'static,
69    {
70        let (reply, result) = oneshot::channel();
71        let job: Job<State> = Box::new(move |state| {
72            if reply.is_closed() {
73                tracing::debug!("discarding cancelled queued work");
74                return;
75            }
76            let outcome =
77                std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| operation(state)));
78            // A cancelled caller deliberately abandons its reply.
79            if reply.send(outcome).is_err() {
80                tracing::debug!("discarding completed work for a cancelled caller");
81            }
82        });
83        if self
84            .inner
85            .sender
86            .as_ref()
87            .expect("live workers retain their sender")
88            .send(job)
89            .await
90            .is_err()
91        {
92            panic!("worker pool stopped before admitting a call");
93        }
94        match result.await {
95            Ok(Ok(value)) => value,
96            Ok(Err(panic)) => std::panic::resume_unwind(panic),
97            Err(_) => panic!("worker pool dropped a call without responding"),
98        }
99    }
100}
101
102struct Workers<State> {
103    sender: Option<mpsc::Sender<Job<State>>>,
104    joins: Vec<std::thread::JoinHandle<()>>,
105}
106
107impl<State> Drop for Workers<State> {
108    fn drop(&mut self) {
109        // Closing admission lets workers drain cancelled jobs and drop their
110        // retained state on their own thread. Never block an async executor.
111        drop(self.sender.take());
112        let current_thread = std::thread::current().id();
113        let on_owned_worker = self
114            .joins
115            .iter()
116            .any(|join| join.thread().id() == current_thread);
117        if !on_owned_worker && tokio::runtime::Handle::try_current().is_err() {
118            for join in self.joins.drain(..) {
119                if join.join().is_err() {
120                    tracing::error!("bounded worker panicked");
121                }
122            }
123        }
124    }
125}
126
127#[cfg(test)]
128#[path = "bounded_workers_tests.rs"]
129mod tests;
coven_foundation/bounded_workers.rs

coven_foundation/
bounded_workers.rs
