medi_rs/adapters/
lifecycle.rs

1//! Completion tracking for generated mediator workers and runtime tasks.
2
3use core::future::{Future, poll_fn};
4use core::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
5
6use crate::StartError;
7use core::task::Poll;
8use futures::task::AtomicWaker;
9
10/// Tracks work started by a generated mediator.
11///
12/// Generated code registers every event worker and `#[medi_task]` before it is
13/// spawned, and marks it complete when its future returns. [`Self::wait`] is
14/// therefore the completion guarantee used by mediator shutdown.
15pub struct Lifecycle {
16    started: AtomicBool,
17    running: AtomicUsize,
18    shutdown_requested: AtomicBool,
19    waker: AtomicWaker,
20}
21
22impl Lifecycle {
23    /// Create an idle lifecycle tracker.
24    pub const fn new() -> Self {
25        Self {
26            started: AtomicBool::new(false),
27            running: AtomicUsize::new(0),
28            shutdown_requested: AtomicBool::new(false),
29            waker: AtomicWaker::new(),
30        }
31    }
32
33    /// Mark the mediator as started.
34    ///
35    /// Returns [`StartError::AlreadyStarted`] if another caller has already
36    /// started it. This operation is atomic across concurrent callers.
37    pub fn start(&self) -> core::result::Result<(), StartError> {
38        self.started
39            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
40            .map(|_| ())
41            .map_err(|_| StartError::AlreadyStarted)
42    }
43
44    /// Return whether the mediator has been started.
45    pub fn is_started(&self) -> bool {
46        self.started.load(Ordering::Acquire)
47    }
48
49    /// Request shutdown, returning `true` only for the caller that initiated it.
50    pub fn request_shutdown(&self) -> bool {
51        self.shutdown_requested
52            .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
53            .is_ok()
54    }
55
56    /// Register work immediately before spawning it.
57    pub fn begin(&self) {
58        self.running.fetch_add(1, Ordering::AcqRel);
59    }
60
61    /// Mark previously registered work as complete.
62    pub fn finish(&self) {
63        if self.running.fetch_sub(1, Ordering::AcqRel) == 1 {
64            self.waker.wake();
65        }
66    }
67
68    /// Wait until all registered workers and tasks have returned.
69    pub fn wait(&self) -> impl Future<Output = ()> + '_ {
70        poll_fn(|cx| {
71            if self.running.load(Ordering::Acquire) == 0 {
72                return Poll::Ready(());
73            }
74            self.waker.register(cx.waker());
75            if self.running.load(Ordering::Acquire) == 0 {
76                Poll::Ready(())
77            } else {
78                Poll::Pending
79            }
80        })
81    }
82}
83
84impl Default for Lifecycle {
85    fn default() -> Self {
86        Self::new()
87    }
88}
medi_rs/adapters/lifecycle.rs

medi_rs/adapters/
lifecycle.rs
