ixa/
plan_queue.rs

1//! Implementation details for Ixa's plan queues.
2//!
3//! [`PlanId`] is the only public item in this module. The queues store regular
4//! time-ordered plans and shutdown-time plans, sharing a single `PlanId`
5//! allocator and cancellation map.
6
7use std::cmp::Ordering;
8use std::collections::BinaryHeap;
9
10use crate::context::{Context, ExecutionPhase};
11use crate::{trace, HashMap, HashMapExt};
12
13type Callback = dyn FnOnce(&mut Context);
14type BoxedCallback = Box<Callback>;
15
16struct QueuedPlan {
17    callback: BoxedCallback,
18    is_passive: bool,
19}
20
21/// A priority queue that stores scheduled plans.
22///
23/// Regular plans are ordered by simulation time, execution phase, and plan ID.
24/// Shutdown-time plans are ordered by execution phase and plan ID; their stored
25/// time is only an internal constant and has no simulation-time meaning.
26pub(crate) struct PlanQueue {
27    queue: BinaryHeap<PlanSchedule>,
28    shutdown_queue: BinaryHeap<PlanSchedule>,
29    data_map: HashMap<u64, QueuedPlan>,
30    /// Count of scheduled plans excluding passive and shutdown-time plans.
31    regular_plan_count: usize,
32    /// The next plan ID that will be issued.
33    next_plan_id: u64,
34    /// Tracks the high water mark of plans in flight (scheduled but not yet executed).
35    /// This is the max of the two heap lengths, not of `self.data_map.len()`.
36    #[cfg(feature = "profiling")]
37    pub(crate) max_plans_in_flight: u64,
38    #[cfg(feature = "profiling")]
39    pub(crate) max_memory_in_use: u64,
40}
41
42impl PlanQueue {
43    /// Create a new empty `PlanQueue`.
44    #[must_use]
45    pub(crate) fn new() -> PlanQueue {
46        PlanQueue {
47            queue: BinaryHeap::new(),
48            shutdown_queue: BinaryHeap::new(),
49            data_map: HashMap::new(),
50            regular_plan_count: 0,
51            next_plan_id: 0,
52            #[cfg(feature = "profiling")]
53            max_plans_in_flight: 0,
54            #[cfg(feature = "profiling")]
55            max_memory_in_use: 0,
56        }
57    }
58
59    /// Add a regular plan to the queue at the specified time.
60    ///
61    /// Returns a [`PlanId`] for the newly-added plan that can be used to cancel it
62    /// if needed.
63    pub(crate) fn add_plan(
64        &mut self,
65        time: f64,
66        callback: BoxedCallback,
67        phase: ExecutionPhase,
68        is_passive: bool,
69    ) -> PlanId {
70        trace!("adding plan at {time}");
71        let plan_id = self.next_plan_id;
72        self.queue.push(PlanSchedule {
73            plan_id,
74            time,
75            phase,
76        });
77        self.data_map.insert(
78            plan_id,
79            QueuedPlan {
80                callback,
81                is_passive,
82            },
83        );
84        if !is_passive {
85            self.regular_plan_count += 1;
86        }
87        self.next_plan_id += 1;
88        self.update_profiling_high_water_marks();
89
90        PlanId(plan_id)
91    }
92
93    /// Add a shutdown-time plan.
94    ///
95    /// Shutdown-time plans have no simulation time. They are ordered by phase and
96    /// plan ID.
97    pub(crate) fn add_shutdown_plan(
98        &mut self,
99        callback: BoxedCallback,
100        phase: ExecutionPhase,
101    ) -> PlanId {
102        trace!("adding shutdown-time plan");
103        let plan_id = self.next_plan_id;
104        self.shutdown_queue.push(PlanSchedule {
105            plan_id,
106            time: 0.0,
107            phase,
108        });
109        self.data_map.insert(
110            plan_id,
111            QueuedPlan {
112                callback,
113                is_passive: true,
114            },
115        );
116        self.next_plan_id += 1;
117        self.update_profiling_high_water_marks();
118
119        PlanId(plan_id)
120    }
121
122    /// Cancel a plan that has been added to either queue.
123    pub(crate) fn cancel_plan(&mut self, plan_id: &PlanId) -> Option<BoxedCallback> {
124        trace!("cancel plan {plan_id:?}");
125        // Delete the plan from the map, but leave in the heap. It will be skipped
126        // when its heap entry reaches the root.
127        self.data_map.remove(&plan_id.0).map(|queued_plan| {
128            if !queued_plan.is_passive {
129                self.regular_plan_count -= 1;
130            }
131            queued_plan.callback
132        })
133    }
134
135    /// Return the time the next plan is scheduled for, if there is one.
136    #[must_use]
137    pub(crate) fn next_time(&mut self) -> Option<f64> {
138        while let Some(entry) = self.queue.peek() {
139            // We only want to report the time if the plan has not been canceled.
140            if self.data_map.contains_key(&entry.plan_id) {
141                return Some(entry.time);
142            }
143            // Trim the canceled plan.
144            self.queue.pop();
145        }
146        None
147    }
148
149    /// Completely empties the queue, including the plans scheduled at shutdown time.
150    #[allow(dead_code)]
151    pub(crate) fn clear(&mut self) {
152        self.data_map.clear();
153        self.queue.clear();
154        self.shutdown_queue.clear();
155        self.regular_plan_count = 0;
156        self.next_plan_id = 0;
157    }
158
159    /// Retrieve the earliest regular plan only if the regular queue is active.
160    ///
161    /// The queue is active while at least one non-passive regular plan is
162    /// scheduled. The returned plan is simply the next regular plan by time,
163    /// phase, and plan ID; it may be passive. If no live non-passive regular
164    /// plan is scheduled, this returns `None` without removing passive plans
165    /// from the regular queue.
166    pub(crate) fn pop_next_if_active(&mut self) -> Option<Plan> {
167        if self.regular_plan_count == 0 {
168            return None;
169        }
170
171        trace!("getting next plan");
172        loop {
173            // The `pop` should be infallible when the plan count is positive unless the
174            // queue invariants have been violated.
175            let entry = self
176                .queue
177                .pop()
178                .expect("plan count was positive but no plan was available");
179
180            // Discard any cancelled plans we encounter.
181            if let Some(queued_plan) = self.data_map.remove(&entry.plan_id) {
182                if !queued_plan.is_passive {
183                    self.regular_plan_count -= 1;
184                }
185                return Some(Plan {
186                    time: entry.time,
187                    data: queued_plan.callback,
188                });
189            }
190        }
191    }
192
193    /// Retrieve the earliest regular plan only if it is scheduled at `time`.
194    ///
195    /// Returns `None` without removing a future plan if the next regular plan is
196    /// later than `time`.
197    pub(crate) fn pop_next_at(&mut self, time: f64) -> Option<Plan> {
198        loop {
199            match self.queue.peek() {
200                // Trim any cancelled plans
201                Some(entry) if !self.data_map.contains_key(&entry.plan_id) => {
202                    self.queue.pop();
203                }
204
205                // Return only if the plan is scheduled for the given time
206                Some(entry) if entry.time == time => {
207                    // Pop is infallible here.
208                    let entry = self.queue.pop().unwrap();
209                    let queued_plan = self
210                        .data_map
211                        .remove(&entry.plan_id)
212                        .expect("live plan must have callback");
213                    if !queued_plan.is_passive {
214                        self.regular_plan_count -= 1;
215                    }
216                    return Some(Plan {
217                        time,
218                        data: queued_plan.callback,
219                    });
220                }
221
222                // There are no plans scheduled at the given time
223                _ => return None,
224            }
225        }
226    }
227
228    /// Retrieve the next shutdown-time plan.
229    ///
230    /// Returns the next shutdown-time plan if it exists or else `None` if the
231    /// shutdown-time queue is empty.
232    pub(crate) fn pop_next_shutdown(&mut self) -> Option<Plan> {
233        trace!("getting next shutdown-time plan");
234        std::iter::from_fn(|| self.shutdown_queue.pop()).find_map(|entry| {
235            // If there's no `data_map` entry, the plan has been canceled, so discard
236            // and pop another plan.
237            let queued_plan = self.data_map.remove(&entry.plan_id)?;
238            if !queued_plan.is_passive {
239                self.regular_plan_count -= 1;
240            }
241            Some(Plan {
242                time: entry.time,
243                data: queued_plan.callback,
244            })
245        })
246    }
247
248    fn update_profiling_high_water_marks(&mut self) {
249        #[cfg(feature = "profiling")]
250        {
251            let plans_in_flight = self.queue.len() + self.shutdown_queue.len();
252            self.max_plans_in_flight = self.max_plans_in_flight.max(plans_in_flight as u64);
253            self.max_memory_in_use = self
254                .max_memory_in_use
255                .max(self.estimated_memory_in_use() as u64);
256        }
257    }
258
259    #[cfg(feature = "profiling")]
260    fn estimated_memory_in_use(&self) -> usize {
261        let queue_bytes =
262            (self.queue.capacity() + self.shutdown_queue.capacity()) * size_of::<PlanSchedule>();
263
264        let map_entry_bytes = self.data_map.capacity() * size_of::<(u64, QueuedPlan)>();
265
266        queue_bytes + map_entry_bytes
267    }
268}
269
270impl Default for PlanQueue {
271    fn default() -> Self {
272        Self::new()
273    }
274}
275
276/// A time, id, and phase object used to order plans in a [`PlanQueue`].
277///
278/// Regular [`PlanSchedule`] objects are sorted in increasing order of time,
279/// phase, and then plan id. Shutdown-time schedules all have the same internal
280/// time and are therefore sorted by phase and then plan id.
281#[derive(PartialEq, Debug, Clone, Copy)]
282pub(crate) struct PlanSchedule {
283    pub plan_id: u64,
284    pub time: f64,
285    pub phase: ExecutionPhase,
286}
287
288impl Eq for PlanSchedule {}
289
290impl PartialOrd for PlanSchedule {
291    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
292        Some(self.cmp(other))
293    }
294}
295
296/// Entry objects are ordered in increasing order by time, phase, and then plan id.
297impl Ord for PlanSchedule {
298    fn cmp(&self, other: &Self) -> Ordering {
299        let time_ordering = self.time.partial_cmp(&other.time).unwrap().reverse();
300        match time_ordering {
301            Ordering::Equal => {
302                let phase_ordering = self.phase.partial_cmp(&other.phase).unwrap().reverse();
303                match phase_ordering {
304                    Ordering::Equal => self.plan_id.cmp(&other.plan_id).reverse(),
305                    _ => phase_ordering,
306                }
307            }
308            _ => time_ordering,
309        }
310    }
311}
312
313/// A unique identifier for a plan scheduled on a [`Context`].
314///
315/// `PlanId` values are returned by [`Context::add_plan`] and related scheduling
316/// methods, and can be passed to [`Context::cancel_plan`].
317///
318/// # Examples
319///
320/// ```
321/// use ixa::{Context, PlanId};
322///
323/// let mut context = Context::new();
324/// let plan_id: PlanId = context.add_plan(1.0, |_| {});
325/// context.cancel_plan(&plan_id);
326/// ```
327#[derive(Clone, Copy, Debug, Hash, Eq, PartialEq)]
328pub struct PlanId(pub(crate) u64);
329
330/// A plan that holds a callback intended to be executed at the specified time.
331pub(crate) struct Plan {
332    pub time: f64,
333    pub data: BoxedCallback,
334}
335
336#[cfg(test)]
337#[allow(clippy::float_cmp)]
338mod tests {
339    use std::cell::RefCell;
340    use std::rc::Rc;
341
342    use super::PlanQueue;
343    use crate::context::{Context, ExecutionPhase};
344
345    fn callback(value: u32, observed: Rc<RefCell<Vec<u32>>>) -> Box<dyn FnOnce(&mut Context)> {
346        Box::new(move |_| observed.borrow_mut().push(value))
347    }
348
349    fn run_plan(plan: super::Plan, context: &mut Context) {
350        (plan.data)(context);
351    }
352
353    #[test]
354    fn empty_queue() {
355        let mut plan_queue = PlanQueue::new();
356        assert!(plan_queue.pop_next_if_active().is_none());
357    }
358
359    #[test]
360    fn add_plans() {
361        let observed = Rc::new(RefCell::new(Vec::new()));
362        let mut context = Context::new();
363        let mut plan_queue = PlanQueue::new();
364        plan_queue.add_plan(
365            1.0,
366            callback(1, Rc::clone(&observed)),
367            ExecutionPhase::Normal,
368            false,
369        );
370        plan_queue.add_plan(
371            3.0,
372            callback(3, Rc::clone(&observed)),
373            ExecutionPhase::Normal,
374            false,
375        );
376        plan_queue.add_plan(
377            2.0,
378            callback(2, Rc::clone(&observed)),
379            ExecutionPhase::Normal,
380            false,
381        );
382        assert_eq!(plan_queue.next_time(), Some(1.0));
383
384        let next_plan = plan_queue.pop_next_if_active().unwrap();
385        assert_eq!(next_plan.time, 1.0);
386        run_plan(next_plan, &mut context);
387
388        assert_eq!(plan_queue.next_time(), Some(2.0));
389        let next_plan = plan_queue.pop_next_if_active().unwrap();
390        assert_eq!(next_plan.time, 2.0);
391        run_plan(next_plan, &mut context);
392
393        assert_eq!(plan_queue.next_time(), Some(3.0));
394        let next_plan = plan_queue.pop_next_if_active().unwrap();
395        assert_eq!(next_plan.time, 3.0);
396        run_plan(next_plan, &mut context);
397
398        assert!(plan_queue.pop_next_if_active().is_none());
399        assert_eq!(*observed.borrow(), vec![1, 2, 3]);
400    }
401
402    #[test]
403    fn add_plans_at_same_time_with_same_phase() {
404        let observed = Rc::new(RefCell::new(Vec::new()));
405        let mut context = Context::new();
406        let mut plan_queue = PlanQueue::new();
407        plan_queue.add_plan(
408            1.0,
409            callback(1, Rc::clone(&observed)),
410            ExecutionPhase::Normal,
411            false,
412        );
413        plan_queue.add_plan(
414            1.0,
415            callback(2, Rc::clone(&observed)),
416            ExecutionPhase::Normal,
417            false,
418        );
419
420        let next_plan = plan_queue.pop_next_if_active().unwrap();
421        assert_eq!(next_plan.time, 1.0);
422        run_plan(next_plan, &mut context);
423        let next_plan = plan_queue.pop_next_if_active().unwrap();
424        assert_eq!(next_plan.time, 1.0);
425        run_plan(next_plan, &mut context);
426
427        assert!(plan_queue.pop_next_if_active().is_none());
428        assert_eq!(*observed.borrow(), vec![1, 2]);
429    }
430
431    #[test]
432    fn add_plans_at_same_time_with_different_phase() {
433        let observed = Rc::new(RefCell::new(Vec::new()));
434        let mut context = Context::new();
435        let mut plan_queue = PlanQueue::new();
436        plan_queue.add_plan(
437            1.0,
438            callback(1, Rc::clone(&observed)),
439            ExecutionPhase::Normal,
440            false,
441        );
442        plan_queue.add_plan(
443            1.0,
444            callback(2, Rc::clone(&observed)),
445            ExecutionPhase::First,
446            false,
447        );
448
449        let next_plan = plan_queue.pop_next_if_active().unwrap();
450        assert_eq!(next_plan.time, 1.0);
451        run_plan(next_plan, &mut context);
452        let next_plan = plan_queue.pop_next_if_active().unwrap();
453        assert_eq!(next_plan.time, 1.0);
454        run_plan(next_plan, &mut context);
455
456        assert!(plan_queue.pop_next_if_active().is_none());
457        assert_eq!(*observed.borrow(), vec![2, 1]);
458    }
459
460    #[test]
461    fn cancel_plan() {
462        let observed = Rc::new(RefCell::new(Vec::new()));
463        let mut context = Context::new();
464        let mut plan_queue = PlanQueue::new();
465        plan_queue.add_plan(
466            1.0,
467            callback(1, Rc::clone(&observed)),
468            ExecutionPhase::Normal,
469            false,
470        );
471        let plan_to_cancel = plan_queue.add_plan(
472            2.0,
473            callback(2, Rc::clone(&observed)),
474            ExecutionPhase::Normal,
475            false,
476        );
477        plan_queue.add_plan(
478            3.0,
479            callback(3, Rc::clone(&observed)),
480            ExecutionPhase::Normal,
481            false,
482        );
483        plan_queue.cancel_plan(&plan_to_cancel);
484
485        let next_plan = plan_queue.pop_next_if_active().unwrap();
486        assert_eq!(next_plan.time, 1.0);
487        run_plan(next_plan, &mut context);
488
489        let next_plan = plan_queue.pop_next_if_active().unwrap();
490        assert_eq!(next_plan.time, 3.0);
491        run_plan(next_plan, &mut context);
492
493        assert!(plan_queue.pop_next_if_active().is_none());
494        assert_eq!(*observed.borrow(), vec![1, 3]);
495    }
496
497    #[test]
498    fn passive_only_plans_do_not_pop_during_active_execution() {
499        let observed = Rc::new(RefCell::new(Vec::new()));
500        let mut plan_queue = PlanQueue::new();
501        plan_queue.add_plan(
502            1.0,
503            callback(1, Rc::clone(&observed)),
504            ExecutionPhase::Normal,
505            true,
506        );
507
508        assert_eq!(plan_queue.regular_plan_count, 0);
509        assert!(plan_queue.pop_next_if_active().is_none());
510        assert_eq!(plan_queue.next_time(), Some(1.0));
511        assert!(observed.borrow().is_empty());
512    }
513
514    #[test]
515    fn passive_plan_can_pop_before_later_non_passive_plan() {
516        let observed = Rc::new(RefCell::new(Vec::new()));
517        let mut context = Context::new();
518        let mut plan_queue = PlanQueue::new();
519        plan_queue.add_plan(
520            1.0,
521            callback(1, Rc::clone(&observed)),
522            ExecutionPhase::Normal,
523            true,
524        );
525        plan_queue.add_plan(
526            2.0,
527            callback(2, Rc::clone(&observed)),
528            ExecutionPhase::Normal,
529            false,
530        );
531
532        assert_eq!(plan_queue.regular_plan_count, 1);
533        let next_plan = plan_queue.pop_next_if_active().unwrap();
534        assert_eq!(next_plan.time, 1.0);
535        run_plan(next_plan, &mut context);
536        assert_eq!(plan_queue.regular_plan_count, 1);
537
538        let next_plan = plan_queue.pop_next_if_active().unwrap();
539        assert_eq!(next_plan.time, 2.0);
540        run_plan(next_plan, &mut context);
541        assert_eq!(plan_queue.regular_plan_count, 0);
542
543        assert_eq!(*observed.borrow(), vec![1, 2]);
544    }
545
546    #[test]
547    fn canceling_non_passive_plan_decrements_regular_count() {
548        let observed = Rc::new(RefCell::new(Vec::new()));
549        let mut plan_queue = PlanQueue::new();
550        let non_passive = plan_queue.add_plan(
551            1.0,
552            callback(1, Rc::clone(&observed)),
553            ExecutionPhase::Normal,
554            false,
555        );
556        let passive = plan_queue.add_plan(
557            2.0,
558            callback(2, Rc::clone(&observed)),
559            ExecutionPhase::Normal,
560            true,
561        );
562
563        assert_eq!(plan_queue.regular_plan_count, 1);
564        plan_queue.cancel_plan(&passive);
565        assert_eq!(plan_queue.regular_plan_count, 1);
566        plan_queue.cancel_plan(&non_passive);
567        assert_eq!(plan_queue.regular_plan_count, 0);
568    }
569
570    #[test]
571    fn shutdown_plans_do_not_affect_regular_count() {
572        let observed = Rc::new(RefCell::new(Vec::new()));
573        let mut context = Context::new();
574        let mut plan_queue = PlanQueue::new();
575        plan_queue.add_shutdown_plan(callback(1, Rc::clone(&observed)), ExecutionPhase::Normal);
576
577        assert_eq!(plan_queue.regular_plan_count, 0);
578        let next_plan = plan_queue.pop_next_shutdown().unwrap();
579        run_plan(next_plan, &mut context);
580        assert_eq!(plan_queue.regular_plan_count, 0);
581        assert_eq!(*observed.borrow(), vec![1]);
582    }
583
584    #[test]
585    fn next_time_ignores_canceled_root() {
586        let observed = Rc::new(RefCell::new(Vec::new()));
587        let mut plan_queue = PlanQueue::new();
588        let plan_to_cancel = plan_queue.add_plan(
589            1.0,
590            callback(1, Rc::clone(&observed)),
591            ExecutionPhase::Normal,
592            false,
593        );
594        plan_queue.add_plan(
595            2.0,
596            callback(2, Rc::clone(&observed)),
597            ExecutionPhase::Normal,
598            false,
599        );
600
601        plan_queue.cancel_plan(&plan_to_cancel);
602
603        assert_eq!(plan_queue.next_time(), Some(2.0));
604    }
605
606    #[test]
607    fn pop_next_at_leaves_future_plan_in_queue() {
608        let observed = Rc::new(RefCell::new(Vec::new()));
609        let mut context = Context::new();
610        let mut plan_queue = PlanQueue::new();
611        plan_queue.add_plan(
612            2.0,
613            callback(2, Rc::clone(&observed)),
614            ExecutionPhase::Normal,
615            false,
616        );
617
618        assert!(plan_queue.pop_next_at(1.0).is_none());
619
620        let next_plan = plan_queue.pop_next_if_active().unwrap();
621        assert_eq!(next_plan.time, 2.0);
622        run_plan(next_plan, &mut context);
623        assert_eq!(*observed.borrow(), vec![2]);
624    }
625
626    #[test]
627    fn shutdown_plans_use_phase_and_fifo_order() {
628        let observed = Rc::new(RefCell::new(Vec::new()));
629        let mut context = Context::new();
630        let mut plan_queue = PlanQueue::new();
631        plan_queue.add_shutdown_plan(callback(3, Rc::clone(&observed)), ExecutionPhase::Last);
632        plan_queue.add_shutdown_plan(callback(1, Rc::clone(&observed)), ExecutionPhase::First);
633        plan_queue.add_shutdown_plan(callback(2, Rc::clone(&observed)), ExecutionPhase::Normal);
634        plan_queue.add_shutdown_plan(callback(4, Rc::clone(&observed)), ExecutionPhase::Last);
635
636        while let Some(plan) = plan_queue.pop_next_shutdown() {
637            run_plan(plan, &mut context);
638        }
639
640        assert_eq!(*observed.borrow(), vec![1, 2, 3, 4]);
641    }
642
643    #[test]
644    fn plan_ids_are_shared_between_regular_and_shutdown_queues() {
645        let observed = Rc::new(RefCell::new(Vec::new()));
646        let mut plan_queue = PlanQueue::new();
647        let regular_id = plan_queue.add_plan(
648            1.0,
649            callback(1, Rc::clone(&observed)),
650            ExecutionPhase::Normal,
651            false,
652        );
653        let shutdown_id =
654            plan_queue.add_shutdown_plan(callback(2, Rc::clone(&observed)), ExecutionPhase::Normal);
655
656        assert_ne!(regular_id, shutdown_id);
657        assert_eq!(regular_id.0, 0);
658        assert_eq!(shutdown_id.0, 1);
659    }
660}
ixa/plan_queue.rs

ixa/
plan_queue.rs
