Skip to main content

rig_ecs/bus/
plugin.rs

1//! Bus scheduling, serving policy, and wake-driven app execution.
2//!
3//! ```
4//! let wake = rig_ecs::bus::Wake::default();
5//! wake.signal();
6//! assert!(wake.take());
7//! ```
8
9use std::{
10    sync::{
11        Arc, Condvar, Mutex,
12        atomic::{AtomicBool, AtomicU64, Ordering},
13    },
14    time::Duration,
15};
16
17use bevy_app::{App, AppExit, MainScheduleOrder, Plugin, PluginsState, Update};
18use bevy_ecs::{
19    prelude::*,
20    schedule::{LogLevel, ScheduleBuildSettings, ScheduleLabel},
21};
22use bevy_tasks::{IoTaskPool, TaskPool};
23use rig_core::serve::ServingPolicy;
24
25use super::{
26    collect::{collect_streams, collect_tasks, settle},
27    dispatch::dispatch,
28    effect::{IdCounter, SeqCounter, WorldOutcomeCounter},
29    handlers::{HandlerIndex, WorldKinds},
30    record::{DeliveryBatch, begin_delivery_pass, record_bound, record_cancelled, record_outcome},
31};
32
33/// The schedule the bus runs in, once per app update, after `Update`.
34/// Users add their systems here, ordered against [`BusSet`]s.
35#[derive(ScheduleLabel, Debug, Clone, Copy, PartialEq, Eq, Hash)]
36pub struct RigSchedule;
37
38/// Runs after [`RigSchedule`] and its deferred commands in every app update.
39/// Replay adapters can install idle-pass diagnosis here.
40#[derive(ScheduleLabel, Debug, Clone, Copy, PartialEq, Eq, Hash)]
41pub struct RigEnd;
42
43/// Ordered sets of [`RigSchedule`]: gate pending effects, dispatch, collect and
44/// record answers, then apply policy replacements.
45#[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
46pub enum BusSet {
47    /// Before dispatch: patch, deny, hold.
48    Gate,
49    /// The plugin takes pending effects.
50    Dispatch,
51    /// The plugin lands what finished.
52    Collect,
53    /// After the record: replace.
54    Judge,
55}
56
57/// The world's [`ServingPolicy`], as a resource.
58#[derive(Resource, Debug, Clone, Copy, PartialEq, Eq)]
59pub struct Policy(pub ServingPolicy);
60
61/// A shared signal requesting another app update, raised by completed tasks,
62/// stream delivery, or host systems needing a later pass.
63/// Each signal advances [`Wake::generation`] and wakes blocking and async waiters;
64/// signalling requires only shared access to the resource.
65#[derive(Resource, Clone, Default)]
66pub struct Wake(Arc<WakeInner>);
67
68#[derive(Default)]
69struct WakeInner {
70    raised: AtomicBool,
71    generation: AtomicU64,
72    waker: futures::task::AtomicWaker,
73    lock: Mutex<()>,
74    condvar: Condvar,
75}
76
77impl Wake {
78    /// Raise the signal.
79    pub fn signal(&self) {
80        self.0.generation.fetch_add(1, Ordering::AcqRel);
81        self.0.raised.store(true, Ordering::Release);
82        self.0.waker.wake();
83        let guard = self
84            .0
85            .lock
86            .lock()
87            .unwrap_or_else(|poisoned| poisoned.into_inner());
88        self.0.condvar.notify_all();
89        drop(guard);
90    }
91
92    /// Return the wrapping count of signals raised.
93    pub fn generation(&self) -> u64 {
94        self.0.generation.load(Ordering::Acquire)
95    }
96
97    /// Take the signal: whether it was raised since the last take.
98    pub fn take(&self) -> bool {
99        self.0.raised.swap(false, Ordering::AcqRel)
100    }
101
102    /// Block until the signal is raised or `timeout` elapses, then take it.
103    /// Returns whether it was raised.
104    pub fn wait(&self, timeout: Duration) -> bool {
105        let deadline = std::time::Instant::now() + timeout;
106        let mut guard = self
107            .0
108            .lock
109            .lock()
110            .unwrap_or_else(|poisoned| poisoned.into_inner());
111        while !self.0.raised.load(Ordering::Acquire) {
112            let now = std::time::Instant::now();
113            if now >= deadline {
114                break;
115            }
116            let (next, _) = self
117                .0
118                .condvar
119                .wait_timeout(guard, deadline - now)
120                .unwrap_or_else(|poisoned| poisoned.into_inner());
121            guard = next;
122        }
123        drop(guard);
124        self.take()
125    }
126
127    /// Resolve when the signal is raised, taking it.
128    pub fn woken(&self) -> impl Future<Output = ()> + Send + '_ {
129        std::future::poll_fn(move |cx| {
130            if self.take() {
131                return std::task::Poll::Ready(());
132            }
133            self.0.waker.register(cx.waker());
134            if self.take() {
135                std::task::Poll::Ready(())
136            } else {
137                std::task::Poll::Pending
138            }
139        })
140    }
141}
142
143/// A runner that updates the app whenever [`Wake`] is raised, or every
144/// `idle` at most, until the app asks to exit. Install with
145/// `app.set_runner(woken_runner(..))`; [`BusPlugin`] installs it by default.
146pub fn woken_runner(idle: Duration) -> impl FnOnce(App) -> AppExit {
147    move |mut app: App| {
148        if app.plugins_state() == PluginsState::Ready {
149            app.finish();
150            app.cleanup();
151        }
152        let wake = app.world().resource::<Wake>().clone();
153        loop {
154            app.update();
155            if let Some(exit) = app.should_exit() {
156                return exit;
157            }
158            wake.wait(idle);
159        }
160    }
161}
162
163/// The bus: [`RigSchedule`] with its sets and systems after `Update`, the
164/// counters, the handler table, the observers, the [`Wake`] and its runner.
165///
166/// The task pool: `build` calls `IoTaskPool::get_or_init`, so the bus works
167/// with or without a host that initialises the pools itself.
168#[derive(Debug, Clone)]
169pub struct BusPlugin {
170    /// The serving policy: intake per tick, serial keys and bounded delivery.
171    /// `stream_capacity` supplies shared queue slots, clamped to at least one;
172    /// the single sender has one additional reserved slot.
173    pub policy: ServingPolicy,
174    /// Schedule ambiguity detection level, defaulting to `Warn`.
175    pub ambiguity: LogLevel,
176    /// How long the default runner waits for a wake before updating anyway.
177    pub idle: Duration,
178}
179
180impl Default for BusPlugin {
181    fn default() -> Self {
182        Self {
183            policy: ServingPolicy::default(),
184            ambiguity: LogLevel::Warn,
185            idle: Duration::from_millis(100),
186        }
187    }
188}
189
190impl BusPlugin {
191    /// The bus under `policy`.
192    pub fn with_policy(policy: ServingPolicy) -> Self {
193        Self {
194            policy,
195            ..Self::default()
196        }
197    }
198
199    /// Build the schedule with ambiguity detection at `level`.
200    #[must_use = "the setting applies to the returned value"]
201    pub fn ambiguity_detection(mut self, level: LogLevel) -> Self {
202        self.ambiguity = level;
203        self
204    }
205
206    /// Install resources, observers, [`RigSchedule`], and [`RigEnd`] in `world`.
207    /// Panics if the bus is already installed. Does not configure an app runner;
208    /// callers driving a bare world must run the schedules themselves.
209    pub fn install(&self, world: &mut World) {
210        assert!(
211            !world.contains_resource::<Policy>(),
212            "the bus is already installed in this world"
213        );
214        IoTaskPool::get_or_init(TaskPool::default);
215        // A zero intake bound would leave pending effects permanently undispatched.
216        world.insert_resource(Policy(ServingPolicy {
217            command_capacity: self.policy.command_capacity.max(1),
218            ..self.policy
219        }));
220        world.init_resource::<Wake>();
221        world.init_resource::<SeqCounter>();
222        world.init_resource::<IdCounter>();
223        world.init_resource::<DeliveryBatch>();
224        world.init_resource::<WorldOutcomeCounter>();
225        world.init_resource::<HandlerIndex>();
226        world.init_resource::<WorldKinds>();
227        world.init_non_send::<super::effect::Executions>();
228        world.init_non_send::<super::handlers::HandlerTable>();
229        world.add_observer(super::effect::drop_execution);
230        world.add_observer(super::handlers::unbound);
231        world.add_observer(record_outcome);
232        world.add_observer(record_cancelled);
233        world.add_observer(super::record::witness_cancelled);
234        world.add_observer(record_bound);
235        let mut schedule = Schedule::new(RigSchedule);
236        schedule.set_build_settings(ScheduleBuildSettings {
237            ambiguity_detection: self.ambiguity,
238            ..Default::default()
239        });
240        schedule.configure_sets(
241            (
242                BusSet::Gate,
243                BusSet::Dispatch,
244                BusSet::Collect,
245                BusSet::Judge,
246            )
247                .chain(),
248        );
249        schedule.add_systems(dispatch.in_set(BusSet::Dispatch));
250        schedule.add_systems(begin_delivery_pass.before(BusSet::Gate));
251        schedule.add_systems(
252            (
253                collect_tasks,
254                collect_streams,
255                super::collect::collect_world,
256                settle,
257            )
258                .chain()
259                .in_set(BusSet::Collect),
260        );
261        world.init_resource::<Schedules>();
262        world.resource_mut::<Schedules>().insert(schedule);
263        world
264            .resource_mut::<Schedules>()
265            .insert(Schedule::new(RigEnd));
266    }
267}
268
269impl Plugin for BusPlugin {
270    fn build(&self, app: &mut App) {
271        self.install(app.world_mut());
272        let mut order = app.world_mut().resource_mut::<MainScheduleOrder>();
273        order.insert_after(Update, RigSchedule);
274        order.insert_after(RigSchedule, RigEnd);
275        app.set_runner(woken_runner(self.idle));
276    }
277}
278
279#[cfg(test)]
280mod tests;