dfir_rs/scheduled/
context.rs

1//! Module for the user-facing [`Context`] object.
2//!
3//! Provides APIs for state and scheduling.
4
5use 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
24/// The main state and scheduler of the runtime instance. Provided as the `context` API to each
25/// subgraph/operator as it is run.
26///
27/// Each instance stores eactly one Context inline. Before the `Context` is provided to
28/// a running operator, the `subgraph_id` field must be updated.
29pub struct Context {
30    /// Storage for the user-facing State API.
31    states: SlotVec<StateTag, StateData>,
32
33    /// Priority stack for handling strata within loops. Prioritized by loop depth.
34    pub(super) stratum_stack: PriorityStack<usize>,
35
36    /// Stack of loop nonces. Used to identify when a new loop iteration begins.
37    pub(super) loop_nonce_stack: Vec<usize>,
38
39    /// TODO(mingwei):
40    /// used for loop iteration scheduling.
41    pub(super) schedule_deferred: Vec<SubgraphId>,
42
43    /// TODO(mingwei): separate scheduler into its own struct/trait?
44    /// Index is stratum, value is FIFO queue for that stratum.
45    pub(super) stratum_queues: Vec<VecDeque<SubgraphId>>,
46
47    /// Receive events, if second arg indicates if it is an external "important" event (true).
48    pub(super) event_queue_recv: UnboundedReceiver<(SubgraphId, bool)>,
49    /// If external events or data can justify starting the next tick.
50    pub(super) can_start_tick: bool,
51    /// If the events have been received for this tick.
52    pub(super) events_received_tick: bool,
53
54    // TODO(mingwei): as long as this is here, it's impossible to know when all work is done.
55    // Second field (bool) is for if the event is an external "important" event (true).
56    pub(super) event_queue_send: UnboundedSender<(SubgraphId, bool)>,
57
58    /// If the current subgraph wants to reschedule the current loop block (in the current tick).
59    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    /// Depth of loop (zero for top-level).
70    pub(super) loop_depth: SlotVec<LoopTag, usize>,
71    /// For each loop, state which needs to be reset between loop executions.
72    loop_states: SecondarySlotVec<LoopTag, Vec<StateId>>,
73    /// Used to differentiate between loop executions. Incremented at the start of each loop execution.
74    pub(super) loop_nonce: usize,
75
76    /// For each subgraph, state which needs to be reset between executions.
77    subgraph_states: SecondarySlotVec<SubgraphTag, Vec<StateId>>,
78
79    /// The SubgraphId of the currently running operator. When this context is
80    /// not being forwarded to a running operator, this field is meaningless.
81    pub(super) subgraph_id: SubgraphId,
82
83    tasks_to_spawn: Vec<Pin<Box<dyn Future<Output = ()> + 'static>>>,
84    /// Join handles for spawned tasks.
85    task_join_handles: Vec<JoinHandle<()>>,
86}
87/// Public APIs.
88impl Context {
89    /// Gets the current tick (local time) count.
90    pub fn current_tick(&self) -> TickInstant {
91        self.current_tick
92    }
93
94    /// Gets the timestamp of the beginning of the current tick.
95    pub fn current_tick_start(&self) -> SystemTime {
96        self.current_tick_start
97    }
98
99    /// Gets whether this is the first time this subgraph is being scheduled for this tick
100    pub fn is_first_run_this_tick(&self) -> bool {
101        self.is_first_run_this_tick
102    }
103
104    /// Gets the current loop iteration count.
105    pub fn loop_iter_count(&self) -> usize {
106        self.loop_iter_count
107    }
108
109    /// Gets the current stratum nubmer.
110    pub fn current_stratum(&self) -> usize {
111        self.current_stratum
112    }
113
114    /// Gets the ID of the current subgraph.
115    pub fn current_subgraph(&self) -> SubgraphId {
116        self.subgraph_id
117    }
118
119    /// Schedules a subgraph for the next tick.
120    ///
121    /// If `is_external` is `true`, the scheduling will trigger the next tick to begin. If it is
122    /// `false` then scheduling will be lazy and the next tick will not begin unless there is other
123    /// reason to.
124    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    /// Schedules the current loop block to be run again (_in this tick_).
129    pub fn reschedule_loop_block(&self) {
130        self.reschedule_loop_block.set(true);
131    }
132
133    /// Allow another iteration of the loop, if more data comes.
134    pub fn allow_another_iteration(&self) {
135        self.allow_another_iteration.set(true);
136    }
137
138    /// Returns a `Waker` for interacting with async Rust.
139    /// Waker events are considered to be extenral.
140    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    /// Returns a shared reference to the state.
162    ///
163    /// # Safety
164    /// `StateHandle<T>` must be from _this_ instance, created via [`Self::add_state`].
165    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            // SAFETY: `handle` is from this instance.
180            // TODO(shadaj): replace with `downcast_ref_unchecked` when it's stabilized
181            &*(state as *const dyn Any as *const T)
182        }
183    }
184
185    /// Returns a shared reference to the state.
186    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    /// Returns an exclusive reference to the state.
199    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    /// Adds state to the context and returns the handle.
212    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    /// Sets a hook to modify the state at the end of each tick, using the supplied closure.
230    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                // Already included in `run_state_hooks_tick`.
260            }
261            StateLifespan::Static => {
262                // Never resets.
263            }
264        }
265    }
266
267    /// Prepares an async task to be launched by [`Self::spawn_tasks`].
268    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    /// Launches all tasks requested with [`Self::request_task`] on the internal Tokio executor.
276    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    /// Aborts all tasks spawned with [`Self::spawn_tasks`].
283    pub fn abort_tasks(&mut self) {
284        for task in self.task_join_handles.drain(..) {
285            task.abort();
286        }
287    }
288
289    /// Waits for all tasks spawned with [`Self::spawn_tasks`] to complete.
290    ///
291    /// Will probably just hang.
292    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()]; // Always initialize stratum #0.
300        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            // Will be re-set before use.
334            subgraph_id: SubgraphId::from_raw(0),
335
336            tasks_to_spawn: Vec::new(),
337            task_join_handles: Vec::new(),
338        }
339    }
340}
341/// Internal APIs.
342impl Context {
343    /// Makes sure stratum STRATUM is initialized.
344    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    /// Call this at the end of a tick,
352    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    // Run the state hooks for each state in the loop.
386    // Call at the end of each loop execution.
387    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
409/// Internal struct containing a pointer to instance-owned state.
410struct StateData {
411    state: Box<dyn Any>,
412    lifespan_hook_fn: Option<LifespanResetFn>, // TODO(mingwei): replace with trait?
413    /// `None` for static.
414    lifespan: Option<StateLifespan>,
415}
416type LifespanResetFn = Box<dyn FnMut(&mut dyn Any)>;