Skip to main content

bamboo_engine/runtime/execution/
event_publication.rs

1//! Per-run atomic publication fence. Token traffic never acquires the shared
2//! runner registry. Replacement closes the old fence and drains synchronous
3//! publications before the successor can emit Started on the same channel.
4
5use chrono::{DateTime, Utc};
6use std::sync::atomic::{AtomicI64, AtomicUsize, Ordering};
7
8const CLOSED: usize = 1 << (usize::BITS - 1);
9
10#[derive(Debug, Default)]
11pub struct EventPublication {
12    state: AtomicUsize,
13    last_event_millis: AtomicI64,
14}
15
16struct PublicationGuard<'a>(&'a AtomicUsize);
17impl Drop for PublicationGuard<'_> {
18    fn drop(&mut self) {
19        self.0.fetch_sub(1, Ordering::Release);
20    }
21}
22
23impl EventPublication {
24    /// Only synchronous sends belong in this closure. An async producer cannot
25    /// keep a publication permit across suspension and delay a successor.
26    pub fn publish(&self, send: impl FnOnce()) -> bool {
27        let mut state = self.state.load(Ordering::Acquire);
28        loop {
29            if state & CLOSED != 0 {
30                return false;
31            }
32            match self.state.compare_exchange_weak(
33                state,
34                state + 1,
35                Ordering::AcqRel,
36                Ordering::Acquire,
37            ) {
38                Ok(_) => break,
39                Err(current) => state = current,
40            }
41        }
42        let _permit = PublicationGuard(&self.state);
43        self.touch();
44        send();
45        true
46    }
47
48    pub fn touch(&self) {
49        self.last_event_millis
50            .fetch_max(Utc::now().timestamp_millis(), Ordering::Relaxed);
51    }
52
53    pub fn last_event_at(&self) -> Option<DateTime<Utc>> {
54        let millis = self.last_event_millis.load(Ordering::Relaxed);
55        (millis != 0)
56            .then(|| DateTime::from_timestamp_millis(millis))
57            .flatten()
58    }
59
60    pub async fn retire(&self) {
61        self.state.fetch_or(CLOSED, Ordering::AcqRel);
62        while self.state.load(Ordering::Acquire) & !CLOSED != 0 {
63            tokio::task::yield_now().await;
64        }
65    }
66
67    /// Maintenance removes an entry only once admitted publications have exited.
68    pub fn retire_if_idle(&self) -> bool {
69        self.state.fetch_or(CLOSED, Ordering::AcqRel) & !CLOSED == 0
70    }
71}
72
73#[cfg(test)]
74mod tests {
75    use super::*;
76    use std::sync::{Arc, Barrier};
77
78    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
79    async fn replacement_drains_admitted_frame_and_rejects_late_frames() {
80        let gate = Arc::new(EventPublication::default());
81        let admitted = Arc::new(Barrier::new(2));
82        let release = Arc::new(Barrier::new(2));
83        let publisher = {
84            let gate = gate.clone();
85            let admitted = admitted.clone();
86            let release = release.clone();
87            std::thread::spawn(move || {
88                gate.publish(|| {
89                    admitted.wait();
90                    release.wait();
91                })
92            })
93        };
94        admitted.wait();
95        assert!(!gate.retire_if_idle());
96        assert!(!gate.publish(|| panic!("old generation published")));
97        let retiring = {
98            let gate = gate.clone();
99            tokio::spawn(async move { gate.retire().await })
100        };
101        tokio::task::yield_now().await;
102        assert!(!retiring.is_finished());
103        release.wait();
104        assert!(publisher.join().unwrap());
105        retiring.await.unwrap();
106        assert!(gate.retire_if_idle());
107    }
108
109    #[test]
110    fn panic_releases_publication_permit() {
111        let gate = EventPublication::default();
112        let _ = std::panic::catch_unwind(|| gate.publish(|| panic!("sink panic")));
113        assert!(gate.retire_if_idle());
114    }
115}