bamboo_engine/runtime/execution/
event_publication.rs1use 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 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 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}