strop_core/
worker.rs

1use parking_lot::Mutex;
2use std::sync::{
3    atomic::{AtomicBool, Ordering},
4    Arc,
5};
6
7mod effect;
8pub use effect::spawn_effect;
9
10#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, serde::Serialize, serde::Deserialize)]
11#[serde(transparent)]
12pub struct WorkerId(u64);
13impl WorkerId {
14    pub const fn new(value: u64) -> Self {
15        Self(value)
16    }
17    pub const fn get(self) -> u64 {
18        self.0
19    }
20}
21#[derive(Debug, Default, Clone, serde::Serialize, serde::Deserialize)]
22#[serde(transparent)]
23pub struct WorkerIds(u64);
24impl WorkerIds {
25    pub fn allocate(&mut self) -> Result<WorkerId, Failure> {
26        let next = self.0.checked_add(1).ok_or_else(|| {
27            Failure::new(
28                FailureKind::IdentityExhausted,
29                "worker request IDs exhausted",
30            )
31        })?;
32        self.0 = next;
33        Ok(WorkerId::new(next))
34    }
35}
36#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
37pub enum FailureKind {
38    IdentityExhausted,
39    ThreadStart,
40    Panic,
41    Io,
42    Spawn,
43    Wait,
44    Exit,
45    InvalidInput,
46    Unavailable,
47    Protocol,
48    Disconnected,
49}
50#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
51pub struct Failure {
52    pub kind: FailureKind,
53    pub message: String,
54}
55impl Failure {
56    pub fn new(kind: FailureKind, message: impl Into<String>) -> Self {
57        Self {
58            kind,
59            message: message.into(),
60        }
61    }
62}
63#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
64pub enum CancelReason {
65    Superseded,
66    OwnerClosed,
67    Dismissed,
68    Shutdown,
69}
70#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
71pub enum Outcome<T> {
72    Success(T),
73    Failed {
74        failure: Failure,
75        partial: Option<T>,
76    },
77    Cancelled(CancelReason),
78}
79impl<T> Outcome<T> {
80    pub fn failed(kind: FailureKind, message: impl Into<String>) -> Self {
81        Self::Failed {
82            failure: Failure::new(kind, message),
83            partial: None,
84        }
85    }
86}
87#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
88pub struct Ticket<K> {
89    pub request: WorkerId,
90    pub key: K,
91}
92#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
93pub struct Completion<K, T> {
94    pub ticket: Ticket<K>,
95    pub outcome: Outcome<T>,
96}
97#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
98pub enum Load<K> {
99    Idle,
100    Running(Ticket<K>),
101    Ready(K),
102    Failed { key: K, failure: Failure },
103    Cancelled { key: K, reason: CancelReason },
104}
105impl<K: PartialEq> Load<K> {
106    pub fn owns(&self, ticket: &Ticket<K>) -> bool {
107        matches!(self, Self::Running(current) if current == ticket)
108    }
109    pub fn covers(&self, key: &K) -> bool {
110        match self {
111            Self::Idle => false,
112            Self::Running(t) => &t.key == key,
113            Self::Ready(k) | Self::Failed { key: k, .. } | Self::Cancelled { key: k, .. } => {
114                k == key
115            }
116        }
117    }
118    pub fn retry_failed(&mut self) {
119        if matches!(self, Self::Failed { .. } | Self::Cancelled { .. }) {
120            *self = Self::Idle;
121        }
122    }
123}
124
125type Resource = Box<dyn FnOnce() -> Result<(), Failure> + Send>;
126#[derive(Default)]
127struct Cancellation {
128    cancelled: AtomicBool,
129    resource: Mutex<Option<Resource>>,
130}
131#[derive(Clone)]
132pub struct CancelToken(Arc<Cancellation>);
133fn invoke(resource: Resource) -> Result<(), Failure> {
134    match std::panic::catch_unwind(std::panic::AssertUnwindSafe(resource)) {
135        Ok(result) => result,
136        Err(_) => Err(Failure::new(
137            FailureKind::Panic,
138            "cancellation resource panicked",
139        )),
140    }
141}
142impl CancelToken {
143    pub fn is_cancelled(&self) -> bool {
144        self.0.cancelled.load(Ordering::Acquire)
145    }
146
147    /// Install the single optional resource hook before acquiring native resources.
148    /// A late installation executes immediately and returns its failure to the
149    /// installer. Hooks must be prompt: they must not join or wait for workers.
150    /// A resource published after cancellation must also be revoked by its owner.
151    pub fn register_cancel_resource(
152        &self,
153        resource: impl FnOnce() -> Result<(), Failure> + Send + 'static,
154    ) -> Result<(), Failure> {
155        let resource: Resource = Box::new(resource);
156        {
157            let mut slot = self.0.resource.lock();
158            if !self.is_cancelled() {
159                if slot.is_some() {
160                    return Err(Failure::new(
161                        FailureKind::Protocol,
162                        "cancellation resource already registered",
163                    ));
164                }
165                *slot = Some(resource);
166                return Ok(());
167            }
168        }
169        invoke(resource)
170    }
171
172    /// Call only after normal resource cleanup. An already detached callback
173    /// may still run; the resource itself must serialize revocation and reuse.
174    pub fn clear_cancel_resource(&self) {
175        let resource = self.0.resource.lock().take();
176        drop(resource);
177    }
178
179    fn cancel_resource(&self) -> Result<(), Failure> {
180        let resource = {
181            let mut slot = self.0.resource.lock();
182            self.0.cancelled.store(true, Ordering::Release);
183            slot.take()
184        };
185        match resource {
186            Some(resource) => invoke(resource),
187            None => Ok(()),
188        }
189    }
190}
191type Emitter<T> = Arc<Mutex<Option<Box<dyn FnOnce(Outcome<T>) + Send>>>>;
192fn finish<T>(emitter: &Emitter<T>, outcome: Outcome<T>) {
193    let emit = emitter.lock().take();
194    if let Some(emit) = emit {
195        emit(outcome);
196    }
197}
198pub struct CancelHandle {
199    cancel: Option<Box<dyn FnOnce(CancelReason) + Send>>,
200}
201impl CancelHandle {
202    pub fn cancel(mut self, reason: CancelReason) {
203        if let Some(cancel) = self.cancel.take() {
204            cancel(reason);
205        }
206    }
207}
208impl Drop for CancelHandle {
209    fn drop(&mut self) {
210        if let Some(cancel) = self.cancel.take() {
211            cancel(CancelReason::OwnerClosed);
212        }
213    }
214}
215/// A registered work envelope without an eagerly spawned thread.
216/// Retain a cancellation handle until it runs or is explicitly cancelled.
217#[must_use]
218pub struct PreparedWork<T> {
219    token: CancelToken,
220    emitter: Emitter<T>,
221}
222
223pub fn prepare<T: Send + 'static>(
224    emit: impl FnOnce(Outcome<T>) + Send + 'static,
225) -> PreparedWork<T> {
226    PreparedWork {
227        token: CancelToken(Arc::new(Cancellation::default())),
228        emitter: Arc::new(Mutex::new(Some(Box::new(emit)))),
229    }
230}
231
232impl<T: Send + 'static> PreparedWork<T> {
233    /// Additional guards share one terminal reservation; dropping any live guard
234    /// cancels the work. A scheduler can retain its own guard while the UI owns one.
235    pub fn cancel_handle(&self) -> CancelHandle {
236        let cancel_token = self.token.clone();
237        let cancel_emitter = self.emitter.clone();
238        CancelHandle {
239            cancel: Some(Box::new(move |reason| {
240                let emit = cancel_emitter.lock().take();
241                if let Some(emit) = emit {
242                    let outcome = match cancel_token.cancel_resource() {
243                        Ok(()) => Outcome::Cancelled(reason),
244                        Err(failure) => Outcome::Failed {
245                            failure,
246                            partial: None,
247                        },
248                    };
249                    emit(outcome);
250                }
251            })),
252        }
253    }
254
255    pub fn run(self, work: impl FnOnce(CancelToken) -> Outcome<T>) {
256        if self.token.is_cancelled() {
257            return;
258        }
259        let outcome = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
260            work(self.token.clone())
261        })) {
262            Ok(outcome) => outcome,
263            Err(_) => match self.token.cancel_resource() {
264                Ok(()) => Outcome::failed(FailureKind::Panic, "worker panicked"),
265                Err(failure) => Outcome::Failed {
266                    failure,
267                    partial: None,
268                },
269            },
270        };
271        self.token.clear_cancel_resource();
272        finish(&self.emitter, outcome);
273    }
274}
275
276pub fn spawn<T: Send + 'static>(
277    name: &'static str,
278    emit: impl FnOnce(Outcome<T>) + Send + 'static,
279    work: impl FnOnce(CancelToken) -> Outcome<T> + Send + 'static,
280) -> CancelHandle {
281    let prepared = prepare(emit);
282    let handle = prepared.cancel_handle();
283    let emitter = prepared.emitter.clone();
284    if let Err(error) = std::thread::Builder::new()
285        .name(name.into())
286        .spawn(move || prepared.run(work))
287    {
288        finish(
289            &emitter,
290            Outcome::failed(FailureKind::ThreadStart, error.to_string()),
291        );
292    }
293    handle
294}
295
296/// Scoped readers use the same envelope, and their scope joins them before a
297/// provider can release physical ownership or start its successor.
298pub fn spawn_scoped<'scope, 'env, T: Send + 'static>(
299    scope: &'scope std::thread::Scope<'scope, 'env>,
300    name: &'static str,
301    emit: impl FnOnce(Outcome<T>) + Send + 'static,
302    work: impl FnOnce(CancelToken) -> Outcome<T> + Send + 'scope,
303) -> CancelHandle {
304    let prepared = prepare(emit);
305    let handle = prepared.cancel_handle();
306    let emitter = prepared.emitter.clone();
307    if let Err(error) = std::thread::Builder::new()
308        .name(name.into())
309        .spawn_scoped(scope, move || prepared.run(work))
310    {
311        finish(
312            &emitter,
313            Outcome::failed(FailureKind::ThreadStart, error.to_string()),
314        );
315    }
316    handle
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322    use std::sync::mpsc::channel;
323
324    #[test]
325    fn cancellation_reserves_terminal_before_callback_and_worker_success() {
326        let (ready_tx, ready_rx) = channel();
327        let (entered_tx, entered_rx) = channel();
328        let (release_tx, release_rx) = channel();
329        let (work_tx, work_rx) = channel();
330        let (done_tx, done_rx) = channel();
331        let (tx, rx) = channel();
332        let handle = spawn(
333            "race",
334            move |result| {
335                tx.send(result).unwrap();
336            },
337            move |token| {
338                token
339                    .register_cancel_resource(move || {
340                        entered_tx.send(()).unwrap();
341                        release_rx.recv().unwrap();
342                        Err(Failure::new(FailureKind::Io, "cleanup failed"))
343                    })
344                    .unwrap();
345                ready_tx.send(()).unwrap();
346                work_rx.recv().unwrap();
347                done_tx.send(()).unwrap();
348                Outcome::Success(())
349            },
350        );
351        ready_rx.recv().unwrap();
352        let cancel = std::thread::spawn(move || handle.cancel(CancelReason::Dismissed));
353        entered_rx.recv().unwrap();
354        work_tx.send(()).unwrap();
355        done_rx.recv().unwrap();
356        assert!(rx.try_recv().is_err());
357        release_tx.send(()).unwrap();
358        cancel.join().unwrap();
359        assert!(
360            matches!(rx.recv().unwrap(), Outcome::Failed { failure, .. } if failure.kind == FailureKind::Io)
361        );
362        assert!(rx.recv().is_err());
363    }
364
365    #[test]
366    fn late_registration_runs_immediately_and_success_wins_when_already_published() {
367        let token = CancelToken(Arc::new(Cancellation::default()));
368        token.cancel_resource().unwrap();
369        let (tx, rx) = channel();
370        token
371            .register_cancel_resource(move || {
372                tx.send(()).unwrap();
373                Ok(())
374            })
375            .unwrap();
376        rx.recv().unwrap();
377        let (tx, rx) = channel();
378        let handle = spawn(
379            "success",
380            move |result| {
381                tx.send(result).unwrap();
382            },
383            |_| Outcome::Success(7),
384        );
385        assert!(matches!(rx.recv().unwrap(), Outcome::Success(7)));
386        drop(handle);
387        assert!(rx.recv().is_err());
388    }
389
390    #[test]
391    fn queued_cancellation_prevents_work_from_starting() {
392        let (tx, rx) = channel();
393        let prepared: PreparedWork<()> = prepare(move |outcome| {
394            tx.send(outcome).unwrap();
395        });
396        prepared.cancel_handle().cancel(CancelReason::Superseded);
397        prepared.run(|_| panic!("cancelled queued work executed"));
398        assert!(matches!(
399            rx.recv().unwrap(),
400            Outcome::Cancelled(CancelReason::Superseded)
401        ));
402        assert!(rx.recv().is_err());
403    }
404
405    #[test]
406    fn scoped_worker_cancellation_does_not_skip_physical_join() {
407        let (control_tx, control_rx) = channel();
408        let (entered_tx, entered_rx) = channel();
409        let (release_tx, release_rx) = channel();
410        let (exited_tx, exited_rx) = channel();
411        let (tx, rx) = channel();
412        let scope_owner = std::thread::spawn(move || {
413            std::thread::scope(|scope| {
414                let handle = spawn_scoped(
415                    scope,
416                    "scoped-reader",
417                    move |outcome| {
418                        tx.send(outcome).unwrap();
419                    },
420                    move |_| {
421                        entered_tx.send(()).unwrap();
422                        release_rx.recv().unwrap();
423                        Outcome::Success(())
424                    },
425                );
426                assert!(control_tx.send(handle).is_ok());
427            });
428            exited_tx.send(()).unwrap();
429        });
430        let handle = control_rx.recv().unwrap();
431        entered_rx.recv().unwrap();
432        handle.cancel(CancelReason::Dismissed);
433        assert!(matches!(
434            rx.recv().unwrap(),
435            Outcome::Cancelled(CancelReason::Dismissed)
436        ));
437        assert!(exited_rx.try_recv().is_err());
438        release_tx.send(()).unwrap();
439        exited_rx.recv().unwrap();
440        scope_owner.join().unwrap();
441    }
442}
strop_core/worker.rs

strop_core/
worker.rs
