1use 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#[derive(ScheduleLabel, Debug, Clone, Copy, PartialEq, Eq, Hash)]
36pub struct RigSchedule;
37
38#[derive(ScheduleLabel, Debug, Clone, Copy, PartialEq, Eq, Hash)]
41pub struct RigEnd;
42
43#[derive(SystemSet, Debug, Clone, Copy, PartialEq, Eq, Hash)]
46pub enum BusSet {
47 Gate,
49 Dispatch,
51 Collect,
53 Judge,
55}
56
57#[derive(Resource, Debug, Clone, Copy, PartialEq, Eq)]
59pub struct Policy(pub ServingPolicy);
60
61#[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 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 pub fn generation(&self) -> u64 {
94 self.0.generation.load(Ordering::Acquire)
95 }
96
97 pub fn take(&self) -> bool {
99 self.0.raised.swap(false, Ordering::AcqRel)
100 }
101
102 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 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
143pub 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#[derive(Debug, Clone)]
169pub struct BusPlugin {
170 pub policy: ServingPolicy,
174 pub ambiguity: LogLevel,
176 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 pub fn with_policy(policy: ServingPolicy) -> Self {
193 Self {
194 policy,
195 ..Self::default()
196 }
197 }
198
199 #[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 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 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;