1use std::any::Any;
6use std::cell::Cell;
7use std::collections::VecDeque;
8use std::future::Future;
9use std::marker::PhantomData;
10use std::ops::DerefMut;
11use std::pin::Pin;
12
13use tokio::sync::mpsc::{self, UnboundedReceiver, UnboundedSender};
14use tokio::task::JoinHandle;
15use web_time::SystemTime;
16
17use super::graph::StateLifespan;
18use super::state::StateHandle;
19use super::{LoopId, LoopTag, StateId, StateTag, SubgraphId, SubgraphTag};
20use crate::scheduled::ticks::TickInstant;
21use crate::util::priority_stack::PriorityStack;
22use crate::util::slot_vec::{SecondarySlotVec, SlotVec};
23
24pub struct Context {
30 states: SlotVec<StateTag, StateData>,
32
33 pub(super) stratum_stack: PriorityStack<usize>,
35
36 pub(super) loop_nonce_stack: Vec<usize>,
38
39 pub(super) schedule_deferred: Vec<SubgraphId>,
42
43 pub(super) stratum_queues: Vec<VecDeque<SubgraphId>>,
46
47 pub(super) event_queue_recv: UnboundedReceiver<(SubgraphId, bool)>,
49 pub(super) can_start_tick: bool,
51 pub(super) events_received_tick: bool,
53
54 pub(super) event_queue_send: UnboundedSender<(SubgraphId, bool)>,
57
58 pub(super) reschedule_loop_block: Cell<bool>,
60 pub(super) allow_another_iteration: Cell<bool>,
61
62 pub(super) current_tick: TickInstant,
63 pub(super) current_stratum: usize,
64
65 pub(super) current_tick_start: SystemTime,
66 pub(super) is_first_run_this_tick: bool,
67 pub(super) loop_iter_count: usize,
68
69 pub(super) loop_depth: SlotVec<LoopTag, usize>,
71 loop_states: SecondarySlotVec<LoopTag, Vec<StateId>>,
73 pub(super) loop_nonce: usize,
75
76 subgraph_states: SecondarySlotVec<SubgraphTag, Vec<StateId>>,
78
79 pub(super) subgraph_id: SubgraphId,
82
83 tasks_to_spawn: Vec<Pin<Box<dyn Future<Output = ()> + 'static>>>,
84 task_join_handles: Vec<JoinHandle<()>>,
86}
87impl Context {
89 pub fn current_tick(&self) -> TickInstant {
91 self.current_tick
92 }
93
94 pub fn current_tick_start(&self) -> SystemTime {
96 self.current_tick_start
97 }
98
99 pub fn is_first_run_this_tick(&self) -> bool {
101 self.is_first_run_this_tick
102 }
103
104 pub fn loop_iter_count(&self) -> usize {
106 self.loop_iter_count
107 }
108
109 pub fn current_stratum(&self) -> usize {
111 self.current_stratum
112 }
113
114 pub fn current_subgraph(&self) -> SubgraphId {
116 self.subgraph_id
117 }
118
119 pub fn schedule_subgraph(&self, sg_id: SubgraphId, is_external: bool) {
125 self.event_queue_send.send((sg_id, is_external)).unwrap()
126 }
127
128 pub fn reschedule_loop_block(&self) {
130 self.reschedule_loop_block.set(true);
131 }
132
133 pub fn allow_another_iteration(&self) {
135 self.allow_another_iteration.set(true);
136 }
137
138 pub fn waker(&self) -> std::task::Waker {
141 use std::sync::Arc;
142
143 struct ContextWaker {
144 subgraph_id: SubgraphId,
145 event_queue_send: UnboundedSender<(SubgraphId, bool)>,
146 }
147 impl futures::task::ArcWake for ContextWaker {
148 fn wake_by_ref(arc_self: &Arc<Self>) {
149 let _recv_closed_error =
150 arc_self.event_queue_send.send((arc_self.subgraph_id, true));
151 }
152 }
153
154 let context_waker = ContextWaker {
155 subgraph_id: self.subgraph_id,
156 event_queue_send: self.event_queue_send.clone(),
157 };
158 futures::task::waker(Arc::new(context_waker))
159 }
160
161 pub unsafe fn state_ref_unchecked<T>(&self, handle: StateHandle<T>) -> &'_ T
166 where
167 T: Any,
168 {
169 let state = self
170 .states
171 .get(handle.state_id)
172 .expect("Failed to find state with given handle.")
173 .state
174 .as_ref();
175
176 debug_assert!(state.is::<T>());
177
178 unsafe {
179 &*(state as *const dyn Any as *const T)
182 }
183 }
184
185 pub fn state_ref<T>(&self, handle: StateHandle<T>) -> &'_ T
187 where
188 T: Any,
189 {
190 self.states
191 .get(handle.state_id)
192 .expect("Failed to find state with given handle.")
193 .state
194 .downcast_ref()
195 .expect("StateHandle wrong type T for casting.")
196 }
197
198 pub fn state_mut<T>(&mut self, handle: StateHandle<T>) -> &'_ mut T
200 where
201 T: Any,
202 {
203 self.states
204 .get_mut(handle.state_id)
205 .expect("Failed to find state with given handle.")
206 .state
207 .downcast_mut()
208 .expect("StateHandle wrong type T for casting.")
209 }
210
211 pub fn add_state<T>(&mut self, state: T) -> StateHandle<T>
213 where
214 T: Any,
215 {
216 let state_data = StateData {
217 state: Box::new(state),
218 lifespan_hook_fn: None,
219 lifespan: None,
220 };
221 let state_id = self.states.insert(state_data);
222
223 StateHandle {
224 state_id,
225 _phantom: PhantomData,
226 }
227 }
228
229 pub fn set_state_lifespan_hook<T>(
231 &mut self,
232 handle: StateHandle<T>,
233 lifespan: StateLifespan,
234 mut hook_fn: impl 'static + FnMut(&mut T),
235 ) where
236 T: Any,
237 {
238 let state_data = self
239 .states
240 .get_mut(handle.state_id)
241 .expect("Failed to find state with given handle.");
242 state_data.lifespan_hook_fn = Some(Box::new(move |state| {
243 (hook_fn)(state.downcast_mut::<T>().unwrap());
244 }));
245 state_data.lifespan = Some(lifespan);
246
247 match lifespan {
248 StateLifespan::Subgraph(key) => {
249 self.subgraph_states
250 .get_or_insert_with(key, Vec::new)
251 .push(handle.state_id);
252 }
253 StateLifespan::Loop(loop_id) => {
254 self.loop_states
255 .get_or_insert_with(loop_id, Vec::new)
256 .push(handle.state_id);
257 }
258 StateLifespan::Tick => {
259 }
261 StateLifespan::Static => {
262 }
264 }
265 }
266
267 pub fn request_task<Fut>(&mut self, future: Fut)
269 where
270 Fut: Future<Output = ()> + 'static,
271 {
272 self.tasks_to_spawn.push(Box::pin(future));
273 }
274
275 pub fn spawn_tasks(&mut self) {
277 for task in self.tasks_to_spawn.drain(..) {
278 self.task_join_handles.push(tokio::task::spawn_local(task));
279 }
280 }
281
282 pub fn abort_tasks(&mut self) {
284 for task in self.task_join_handles.drain(..) {
285 task.abort();
286 }
287 }
288
289 pub async fn join_tasks(&mut self) {
293 futures::future::join_all(self.task_join_handles.drain(..)).await;
294 }
295}
296
297impl Default for Context {
298 fn default() -> Self {
299 let stratum_queues = vec![Default::default()]; let (event_queue_send, event_queue_recv) = mpsc::unbounded_channel();
301 let (stratum_stack, loop_depth) = Default::default();
302 Self {
303 states: SlotVec::new(),
304
305 stratum_stack,
306
307 loop_nonce_stack: Vec::new(),
308
309 schedule_deferred: Vec::new(),
310
311 stratum_queues,
312 event_queue_recv,
313 can_start_tick: false,
314 events_received_tick: false,
315
316 event_queue_send,
317 reschedule_loop_block: Cell::new(false),
318 allow_another_iteration: Cell::new(false),
319
320 current_stratum: 0,
321 current_tick: TickInstant::default(),
322
323 current_tick_start: SystemTime::now(),
324 is_first_run_this_tick: false,
325 loop_iter_count: 0,
326
327 loop_depth,
328 loop_states: SecondarySlotVec::new(),
329 loop_nonce: 0,
330
331 subgraph_states: SecondarySlotVec::new(),
332
333 subgraph_id: SubgraphId::from_raw(0),
335
336 tasks_to_spawn: Vec::new(),
337 task_join_handles: Vec::new(),
338 }
339 }
340}
341impl Context {
343 pub(super) fn init_stratum(&mut self, stratum: usize) {
345 if self.stratum_queues.len() <= stratum {
346 self.stratum_queues
347 .resize_with(stratum + 1, Default::default);
348 }
349 }
350
351 pub(super) fn run_state_hooks_tick(&mut self) {
353 tracing::trace!("Running state hooks for tick.");
354 for state_data in self.states.values_mut() {
355 let StateData {
356 state,
357 lifespan_hook_fn: Some(lifespan_hook_fn),
358 lifespan: Some(StateLifespan::Tick),
359 } = state_data
360 else {
361 continue;
362 };
363 (lifespan_hook_fn)(Box::deref_mut(state));
364 }
365 }
366
367 pub(super) fn run_state_hooks_subgraph(&mut self, subgraph_id: SubgraphId) {
368 tracing::trace!("Running state hooks for subgraph.");
369 for state_id in self.subgraph_states.get(subgraph_id).into_iter().flatten() {
370 let StateData {
371 state,
372 lifespan_hook_fn,
373 lifespan: _,
374 } = self
375 .states
376 .get_mut(*state_id)
377 .expect("Failed to find state with given ID.");
378
379 if let Some(lifespan_hook_fn) = lifespan_hook_fn {
380 (lifespan_hook_fn)(Box::deref_mut(state));
381 }
382 }
383 }
384
385 pub(super) fn run_state_hooks_loop(&mut self, loop_id: LoopId) {
388 tracing::trace!(
389 loop_id = loop_id.to_string(),
390 "Running state hooks for loop."
391 );
392 for state_id in self.loop_states.get(loop_id).into_iter().flatten() {
393 let StateData {
394 state,
395 lifespan_hook_fn,
396 lifespan: _,
397 } = self
398 .states
399 .get_mut(*state_id)
400 .expect("Failed to find state with given ID.");
401
402 if let Some(lifespan_hook_fn) = lifespan_hook_fn {
403 (lifespan_hook_fn)(Box::deref_mut(state));
404 }
405 }
406 }
407}
408
409struct StateData {
411 state: Box<dyn Any>,
412 lifespan_hook_fn: Option<LifespanResetFn>, lifespan: Option<StateLifespan>,
415}
416type LifespanResetFn = Box<dyn FnMut(&mut dyn Any)>;