logicaffeine_runtime/
scheduler.rs

1//! The scheduler core + cooperative driver.
2//!
3//! A single, deterministic, seed-driven scheduler that multiplexes [`Task`]s over
4//! channels, `Select`, and timers. The cooperative driver ([`Scheduler::run`])
5//! runs everything on one thread, resolving every nondeterministic decision
6//! through the embedded [`Chooser`] (so a fixed seed ⇒ a fixed run, and a trace
7//! ⇒ an exact replay). The ready-selection discipline is the configured
8//! [`SchedulePolicy`]; determinism holds under all of them.
9
10use std::collections::HashMap;
11
12use crate::channel::{Chan, ChanId};
13use crate::config::{SchedulePolicy, SchedulerConfig};
14use crate::payload::RtPayload;
15use crate::seed::{ChoiceKind, Chooser, SchedSeed, SchedTrace};
16use crate::task::{SelectArm, Task, TaskCtx, TaskId, TaskStep};
17
18/// The result of running a scheduler to quiescence.
19#[derive(Debug, Clone, PartialEq)]
20pub enum RunOutcome {
21    /// Every task finished. Carries the main task's result (or `Nothing`).
22    Done(RtPayload),
23    /// Tasks remain but all are blocked with no timer to fire — a deadlock.
24    Deadlock,
25    /// Tasks remain, none ready and no timer to fire, but at least one is parked on
26    /// external async I/O ([`TaskState::BlockedIo`]) — not a deadlock. The scheduler has no
27    /// reactor; the async drive loop must yield to the host reactor, then [`Scheduler::wake_io`]
28    /// to re-poll the parked tasks. The synchronous `run()` cannot make progress here.
29    WaitingForIo,
30}
31
32enum TaskState {
33    Ready,
34    BlockedRecv,
35    BlockedSend,
36    BlockedSelect,
37    BlockedTimer,
38    BlockedAwait,
39    BlockedIo,
40    Done,
41}
42
43/// A public, read-only projection of a task's scheduling state — what the Studio's
44/// Tasks/Channels strip renders. Mirrors the private [`TaskState`] without exposing the
45/// task body or any mutable handle.
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub enum TaskStateKind {
48    Ready,
49    BlockedRecv,
50    BlockedSend,
51    BlockedSelect,
52    BlockedTimer,
53    BlockedAwait,
54    BlockedIo,
55    Done,
56}
57
58impl TaskState {
59    fn kind(&self) -> TaskStateKind {
60        match self {
61            TaskState::Ready => TaskStateKind::Ready,
62            TaskState::BlockedRecv => TaskStateKind::BlockedRecv,
63            TaskState::BlockedSend => TaskStateKind::BlockedSend,
64            TaskState::BlockedSelect => TaskStateKind::BlockedSelect,
65            TaskState::BlockedTimer => TaskStateKind::BlockedTimer,
66            TaskState::BlockedAwait => TaskStateKind::BlockedAwait,
67            TaskState::BlockedIo => TaskStateKind::BlockedIo,
68            TaskState::Done => TaskStateKind::Done,
69        }
70    }
71}
72
73/// A read-only view of one task for [`Scheduler::snapshot`].
74#[derive(Debug, Clone, Copy, PartialEq, Eq)]
75pub struct TaskView {
76    pub id: TaskId,
77    pub kind: TaskStateKind,
78    pub is_main: bool,
79}
80
81/// A read-only view of one channel for [`Scheduler::snapshot`].
82#[derive(Debug, Clone, Copy, PartialEq, Eq)]
83pub struct ChanView {
84    pub id: ChanId,
85    /// Buffered values waiting to be received.
86    pub depth: usize,
87    /// `None` for an unbounded channel, `Some(0)` for a rendezvous channel.
88    pub capacity: Option<usize>,
89    pub blocked_senders: usize,
90    pub blocked_receivers: usize,
91    pub closed: bool,
92}
93
94/// A deterministic, read-only snapshot of the scheduler between steps — the observability
95/// seam the browser drive loop emits to the Studio so a running concurrent program shows
96/// its live task and channel state instead of freezing the UI.
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct SchedSnapshot {
99    /// The logical clock (the timer wheel's current tick).
100    pub clock: u64,
101    /// How many tasks are ready to run right now.
102    pub ready_len: usize,
103    /// Every task, sorted by id (deterministic order for a stable UI).
104    pub tasks: Vec<TaskView>,
105    /// Every channel, sorted by id.
106    pub channels: Vec<ChanView>,
107}
108
109/// How a finished task finished — delivered to anyone awaiting it.
110#[derive(Debug, Clone)]
111enum AwaitResult {
112    Completed(RtPayload),
113    Aborted,
114}
115
116impl AwaitResult {
117    fn payload(&self) -> RtPayload {
118        match self {
119            AwaitResult::Completed(p) => p.clone(),
120            AwaitResult::Aborted => RtPayload::Nothing,
121        }
122    }
123}
124
125struct TaskSlot<'t> {
126    task: Option<Box<dyn Task<'t> + 't>>,
127    priority: u8,
128    resume: RtPayload,
129    selected_arm: Option<usize>,
130    select_arms: Option<Vec<SelectArm>>,
131    /// Channel id to deliver on the next poll (after a `NewChan`).
132    resume_chan: Option<ChanId>,
133    /// Task id to deliver on the next poll (after a `Spawn`).
134    resume_task: Option<TaskId>,
135    state: TaskState,
136    is_main: bool,
137}
138
139struct TimerEntry {
140    fire_at: u64,
141    tid: TaskId,
142    /// `Some(i)` if this is the timeout arm `i` of a `Select`; `None` for a plain `Sleep`.
143    arm: Option<usize>,
144}
145
146/// The resume state of a ready task, taken by the work-stealing coordinator
147/// ([`crate::executor`]) to deliver to a worker thread. Mirrors the `TaskCtx` the
148/// cooperative `run` builds inline. (Native-only.)
149#[cfg(not(target_arch = "wasm32"))]
150pub(crate) struct WsDispatch {
151    pub tid: TaskId,
152    pub priority: u8,
153    pub selected_arm: Option<usize>,
154    pub resume: RtPayload,
155    pub new_chan: Option<ChanId>,
156    pub spawned: Option<TaskId>,
157}
158
159/// The deterministic task scheduler.
160pub struct Scheduler<'t> {
161    config: SchedulerConfig,
162    chooser: Chooser,
163    tasks: HashMap<TaskId, TaskSlot<'t>>,
164    ready: Vec<TaskId>,
165    chans: HashMap<ChanId, Chan>,
166    timers: Vec<TimerEntry>,
167    clock: u64,
168    next_id: u64,
169    main_result: Option<RtPayload>,
170    results: HashMap<TaskId, AwaitResult>,
171    awaiters: HashMap<TaskId, Vec<TaskId>>,
172}
173
174impl<'t> Scheduler<'t> {
175    /// A fresh scheduler with the given config and decision source.
176    pub fn new(config: SchedulerConfig, chooser: Chooser) -> Self {
177        Scheduler {
178            config,
179            chooser,
180            tasks: HashMap::new(),
181            ready: Vec::new(),
182            chans: HashMap::new(),
183            timers: Vec::new(),
184            clock: 0,
185            next_id: 0,
186            main_result: None,
187            results: HashMap::new(),
188            awaiters: HashMap::new(),
189        }
190    }
191
192    fn bump(&mut self) -> u64 {
193        let v = self.next_id;
194        self.next_id += 1;
195        v
196    }
197
198    /// Create a channel with an explicit capacity (`None` = unbounded, `Some(0)` = rendezvous).
199    pub fn new_chan(&mut self, capacity: Option<usize>) -> ChanId {
200        let id = ChanId(self.bump());
201        self.chans.insert(id, Chan::new(capacity));
202        id
203    }
204
205    /// Create a channel with the config's default capacity.
206    pub fn new_default_chan(&mut self) -> ChanId {
207        let cap = self.config.default_channel_capacity;
208        self.new_chan(Some(cap))
209    }
210
211    /// Spawn a task. Returns its id; it is enqueued ready.
212    pub fn spawn(&mut self, task: Box<dyn Task<'t> + 't>) -> TaskId {
213        let priority = task.priority();
214        self.spawn_inner(Some(task), priority, false)
215    }
216
217    /// Spawn the *main* task — its `Exit` value becomes the run's [`RunOutcome::Done`] payload.
218    pub fn spawn_main(&mut self, task: Box<dyn Task<'t> + 't>) -> TaskId {
219        let priority = task.priority();
220        self.spawn_inner(Some(task), priority, true)
221    }
222
223    fn spawn_inner(
224        &mut self,
225        task: Option<Box<dyn Task<'t> + 't>>,
226        priority: u8,
227        is_main: bool,
228    ) -> TaskId {
229        let id = TaskId(self.bump());
230        self.tasks.insert(
231            id,
232            TaskSlot {
233                task,
234                priority,
235                resume: RtPayload::Nothing,
236                selected_arm: None,
237                select_arms: None,
238                resume_chan: None,
239                resume_task: None,
240                state: TaskState::Ready,
241                is_main,
242            },
243        );
244        self.ready.push(id);
245        id
246    }
247
248    /// Spawn a *metadata-only* slot (no task body) for the work-stealing driver:
249    /// the body lives worker-side, built from a `Send` descriptor. Otherwise
250    /// identical to [`Self::spawn`] — enqueued ready, same id allocation. Returns
251    /// its id. (Native-only: the coordinator that uses it is `cfg(not(wasm32))`.)
252    #[cfg(not(target_arch = "wasm32"))]
253    pub(crate) fn spawn_meta(&mut self, priority: u8, is_main: bool) -> TaskId {
254        self.spawn_inner(None, priority, is_main)
255    }
256
257    /// Consume the scheduler and return the recorded scheduling trace.
258    pub fn into_trace(self) -> SchedTrace {
259        self.chooser.into_trace()
260    }
261
262    /// Run cooperatively to quiescence.
263    pub fn run(&mut self) -> RunOutcome {
264        loop {
265            if let Some(outcome) = self.poll_once() {
266                return outcome;
267            }
268        }
269    }
270
271    /// Advance the scheduler by a single step: dispatch one ready task, or (if none is
272    /// ready) advance the timer wheel. Returns `Some(outcome)` only at quiescence — when no
273    /// task is ready and no timer can fire — and `None` while there is still work to do.
274    /// `run` is `poll_once` to a fixpoint; the browser drive loop calls it in slices,
275    /// yielding a macrotask between them so the UI repaints. Both produce identical
276    /// behavior and identical traces by construction.
277    pub fn poll_once(&mut self) -> Option<RunOutcome> {
278        if let Some(tid) = self.pick_ready() {
279            let step = {
280                let slot = self.tasks.get_mut(&tid).expect("ready task exists");
281                let mut boxed = slot.task.take().expect("task body present");
282                let mut ctx = TaskCtx {
283                    resumed_with: std::mem::replace(&mut slot.resume, RtPayload::Nothing),
284                    selected_arm: slot.selected_arm.take(),
285                    new_chan: slot.resume_chan.take(),
286                    spawned: slot.resume_task.take(),
287                };
288                let step = boxed.poll(&mut ctx);
289                slot.task = Some(boxed);
290                step
291            };
292            self.process_step(tid, step);
293            None
294        } else if !self.timers.is_empty() {
295            self.advance_timers();
296            None
297        } else {
298            let all_done = self.tasks.values().all(|s| matches!(s.state, TaskState::Done));
299            let any_io = self
300                .tasks
301                .values()
302                .any(|s| matches!(s.state, TaskState::BlockedIo));
303            Some(if all_done {
304                RunOutcome::Done(self.main_result.clone().unwrap_or(RtPayload::Nothing))
305            } else if any_io {
306                // Not a deadlock: a task awaits external I/O the async drive loop will service.
307                RunOutcome::WaitingForIo
308            } else {
309                RunOutcome::Deadlock
310            })
311        }
312    }
313
314    /// Run at most `max_steps` steps. Returns `Some(outcome)` if quiescence was reached
315    /// within the budget, or `None` if work remains (the caller should yield and call
316    /// again). A `max_steps` of 0 makes no progress and returns `None`.
317    pub fn run_slice(&mut self, max_steps: usize) -> Option<RunOutcome> {
318        for _ in 0..max_steps {
319            if let Some(outcome) = self.poll_once() {
320                return Some(outcome);
321            }
322        }
323        None
324    }
325
326    /// Re-ready every task parked on external I/O ([`TaskState::BlockedIo`]), in ascending
327    /// id order (deterministic). The async drive loop calls this after yielding to the host
328    /// reactor so the parked tasks re-poll their network futures. Returns whether any task was
329    /// woken. Channel/timer tasks are untouched, so a non-networking program never sees this.
330    pub fn wake_io(&mut self) -> bool {
331        let mut io_tasks: Vec<TaskId> = self
332            .tasks
333            .iter()
334            .filter(|(_, s)| matches!(s.state, TaskState::BlockedIo))
335            .map(|(id, _)| *id)
336            .collect();
337        io_tasks.sort_by_key(|t| t.0);
338        for tid in &io_tasks {
339            let slot = self.tasks.get_mut(tid).unwrap();
340            slot.state = TaskState::Ready;
341            self.ready.push(*tid);
342        }
343        !io_tasks.is_empty()
344    }
345
346    /// A deterministic, read-only snapshot of the current task and channel state — emitted
347    /// between slices to drive the Studio's Tasks/Channels strip. Tasks and channels are
348    /// sorted by id so the UI order is stable across snapshots.
349    pub fn snapshot(&self) -> SchedSnapshot {
350        let mut tasks: Vec<TaskView> = self
351            .tasks
352            .iter()
353            .map(|(id, slot)| TaskView { id: *id, kind: slot.state.kind(), is_main: slot.is_main })
354            .collect();
355        tasks.sort_by_key(|t| t.id.0);
356        let mut channels: Vec<ChanView> = self
357            .chans
358            .iter()
359            .map(|(id, ch)| ChanView {
360                id: *id,
361                depth: ch.queue.len(),
362                capacity: ch.capacity,
363                blocked_senders: ch.blocked_senders.len(),
364                blocked_receivers: ch.blocked_receivers.len(),
365                closed: ch.closed,
366            })
367            .collect();
368        channels.sort_by_key(|c| c.id.0);
369        SchedSnapshot { clock: self.clock, ready_len: self.ready.len(), tasks, channels }
370    }
371
372    // ─── Work-stealing coordinator hooks (native-only) ──────────────────────
373    // The work-stealing driver (`executor.rs`) owns a `Scheduler` with
374    // metadata-only slots and drives it externally: it dispatches ready tasks to
375    // worker threads (each builds + polls the `!Send` body locally) and applies
376    // the workers' `Send` `StepReport`s through the SAME `process_step`/`do_*`
377    // methods — so channel and decision semantics, and the seeded trace, are
378    // identical to the cooperative run.
379
380    /// Spawn a metadata child + ready the parent with the child's handle, exactly
381    /// as the cooperative `Spawn` step does. Returns the child's id (the
382    /// coordinator records its build descriptor under that id).
383    #[cfg(not(target_arch = "wasm32"))]
384    pub(crate) fn ws_spawn_child(&mut self, parent: TaskId, priority: u8) -> TaskId {
385        let child = self.spawn_meta(priority, false);
386        self.tasks.get_mut(&parent).unwrap().resume_task = Some(child);
387        self.mark_ready(parent, RtPayload::Nothing);
388        child
389    }
390
391    /// Pick the next ready task (the same `pick_ready` discipline as cooperative)
392    /// and take its resume state, for the coordinator to deliver to a worker.
393    #[cfg(not(target_arch = "wasm32"))]
394    pub(crate) fn ws_next_dispatch(&mut self) -> Option<WsDispatch> {
395        let tid = self.pick_ready()?;
396        let slot = self.tasks.get_mut(&tid).expect("ready task exists");
397        Some(WsDispatch {
398            tid,
399            priority: slot.priority,
400            selected_arm: slot.selected_arm.take(),
401            resume: std::mem::replace(&mut slot.resume, RtPayload::Nothing),
402            new_chan: slot.resume_chan.take(),
403            spawned: slot.resume_task.take(),
404        })
405    }
406
407    /// Every task finished?
408    #[cfg(not(target_arch = "wasm32"))]
409    pub(crate) fn ws_all_done(&self) -> bool {
410        self.tasks.values().all(|s| matches!(s.state, TaskState::Done))
411    }
412
413    /// Are there pending timers (advance the logical clock when no task is ready)?
414    #[cfg(not(target_arch = "wasm32"))]
415    pub(crate) fn ws_has_timers(&self) -> bool {
416        !self.timers.is_empty()
417    }
418
419    /// Advance the timer wheel (marks timer / select-timeout tasks ready).
420    #[cfg(not(target_arch = "wasm32"))]
421    pub(crate) fn ws_advance_timers(&mut self) {
422        self.advance_timers();
423    }
424
425    /// The main task's result (or `Nothing`).
426    #[cfg(not(target_arch = "wasm32"))]
427    pub(crate) fn ws_main_result(&self) -> RtPayload {
428        self.main_result.clone().unwrap_or(RtPayload::Nothing)
429    }
430
431    fn pick_ready(&mut self) -> Option<TaskId> {
432        if self.ready.is_empty() {
433            return None;
434        }
435        let idx = match self.config.policy {
436            SchedulePolicy::Fifo | SchedulePolicy::RoundRobin => 0,
437            SchedulePolicy::Lifo => self.ready.len() - 1,
438            SchedulePolicy::Random => self.chooser.decide(ChoiceKind::TaskPick, self.ready.len()),
439            SchedulePolicy::Priority => {
440                let mut best = 0usize;
441                let mut best_pri = self.tasks[&self.ready[0]].priority;
442                for i in 1..self.ready.len() {
443                    let p = self.tasks[&self.ready[i]].priority;
444                    if p > best_pri {
445                        best_pri = p;
446                        best = i;
447                    }
448                }
449                best
450            }
451        };
452        Some(self.ready.remove(idx))
453    }
454
455    pub(crate) fn process_step(&mut self, tid: TaskId, step: TaskStep<'t>) {
456        match step {
457            TaskStep::Yield => self.mark_ready(tid, RtPayload::Nothing),
458            TaskStep::Await(target) => self.do_await(tid, target),
459            TaskStep::Abort(target) => {
460                self.do_abort(target);
461                self.mark_ready(tid, RtPayload::Nothing);
462            }
463            TaskStep::Close(ch) => {
464                self.do_close(ch);
465                self.mark_ready(tid, RtPayload::Nothing);
466            }
467            TaskStep::Exit(v) => {
468                if self.tasks[&tid].is_main {
469                    self.main_result = Some(v.clone());
470                }
471                self.tasks.get_mut(&tid).unwrap().state = TaskState::Done;
472                self.complete_task(tid, AwaitResult::Completed(v));
473            }
474            TaskStep::Spawn(child) => {
475                let child_id = self.spawn(child);
476                self.tasks.get_mut(&tid).unwrap().resume_task = Some(child_id);
477                self.mark_ready(tid, RtPayload::Nothing);
478            }
479            // The descriptor form is produced only by work-stealing tasks and is
480            // serviced by the coordinator (which builds the child worker-side),
481            // never the inline `process_step`.
482            TaskStep::SpawnDesc { .. } => {
483                unreachable!("SpawnDesc is serviced by the work-stealing coordinator, not process_step")
484            }
485            TaskStep::NewChan(cap) => {
486                let capacity = Some(cap.unwrap_or(self.config.default_channel_capacity));
487                let id = self.new_chan(capacity);
488                self.tasks.get_mut(&tid).unwrap().resume_chan = Some(id);
489                self.mark_ready(tid, RtPayload::Nothing);
490            }
491            TaskStep::Send(ch, payload) => self.do_send(tid, ch, payload),
492            TaskStep::TrySend(ch, payload) => self.do_try_send(tid, ch, payload),
493            TaskStep::Recv(ch) => self.do_recv(tid, ch),
494            TaskStep::TryRecv(ch) => self.do_try_recv(tid, ch),
495            TaskStep::Sleep(d) => {
496                self.timers.push(TimerEntry { fire_at: self.clock + d, tid, arm: None });
497                self.tasks.get_mut(&tid).unwrap().state = TaskState::BlockedTimer;
498            }
499            TaskStep::Select(arms) => self.do_select(tid, arms),
500            TaskStep::IoPending => {
501                // Parked on external async I/O. The scheduler has no reactor; the async drive
502                // loop re-readies this task via `wake_io` after yielding to the host reactor.
503                self.tasks.get_mut(&tid).unwrap().state = TaskState::BlockedIo;
504            }
505        }
506    }
507
508    fn mark_ready(&mut self, tid: TaskId, value: RtPayload) {
509        let slot = self.tasks.get_mut(&tid).unwrap();
510        slot.resume = value;
511        slot.state = TaskState::Ready;
512        self.ready.push(tid);
513    }
514
515    fn do_send(&mut self, tid: TaskId, ch: ChanId, payload: RtPayload) {
516        let waiting_receiver = self.chans.get_mut(&ch).unwrap().blocked_receivers.pop_front();
517        if let Some(rid) = waiting_receiver {
518            self.deliver_to_receiver(rid, ch, payload);
519            self.mark_ready(tid, RtPayload::Nothing);
520            return;
521        }
522        if self.chans[&ch].has_room() {
523            self.chans.get_mut(&ch).unwrap().queue.push_back(payload);
524            self.mark_ready(tid, RtPayload::Nothing);
525        } else {
526            self.chans.get_mut(&ch).unwrap().blocked_senders.push_back((tid, payload));
527            self.tasks.get_mut(&tid).unwrap().state = TaskState::BlockedSend;
528        }
529    }
530
531    fn do_recv(&mut self, tid: TaskId, ch: ChanId) {
532        let from_queue = self.chans.get_mut(&ch).unwrap().queue.pop_front();
533        if let Some(v) = from_queue {
534            self.mark_ready(tid, v);
535            self.refill_from_blocked_sender(ch);
536            return;
537        }
538        let from_sender = self.chans.get_mut(&ch).unwrap().blocked_senders.pop_front();
539        if let Some((sid, payload)) = from_sender {
540            self.mark_ready(tid, payload);
541            self.mark_ready(sid, RtPayload::Nothing);
542            return;
543        }
544        if self.chans[&ch].closed {
545            self.mark_ready(tid, RtPayload::Nothing);
546            return;
547        }
548        self.chans.get_mut(&ch).unwrap().blocked_receivers.push_back(tid);
549        self.tasks.get_mut(&tid).unwrap().state = TaskState::BlockedRecv;
550    }
551
552    fn do_try_send(&mut self, tid: TaskId, ch: ChanId, payload: RtPayload) {
553        let waiting_receiver = self.chans.get_mut(&ch).unwrap().blocked_receivers.pop_front();
554        if let Some(rid) = waiting_receiver {
555            self.deliver_to_receiver(rid, ch, payload);
556            self.mark_ready(tid, RtPayload::Bool(true));
557            return;
558        }
559        if self.chans[&ch].has_room() {
560            self.chans.get_mut(&ch).unwrap().queue.push_back(payload);
561            self.mark_ready(tid, RtPayload::Bool(true));
562        } else {
563            // Would block — report failure and drop the payload (never park).
564            self.mark_ready(tid, RtPayload::Bool(false));
565        }
566    }
567
568    fn do_try_recv(&mut self, tid: TaskId, ch: ChanId) {
569        let from_queue = self.chans.get_mut(&ch).unwrap().queue.pop_front();
570        if let Some(v) = from_queue {
571            self.mark_ready(tid, v);
572            self.refill_from_blocked_sender(ch);
573            return;
574        }
575        let from_sender = self.chans.get_mut(&ch).unwrap().blocked_senders.pop_front();
576        if let Some((sid, payload)) = from_sender {
577            self.mark_ready(tid, payload);
578            self.mark_ready(sid, RtPayload::Nothing);
579            return;
580        }
581        // Empty (or closed-and-empty) — never park; resume with Nothing.
582        self.mark_ready(tid, RtPayload::Nothing);
583    }
584
585    fn refill_from_blocked_sender(&mut self, ch: ChanId) {
586        let take = {
587            let chan = self.chans.get_mut(&ch).unwrap();
588            if chan.has_room() {
589                chan.blocked_senders.pop_front()
590            } else {
591                None
592            }
593        };
594        if let Some((sid, payload)) = take {
595            self.chans.get_mut(&ch).unwrap().queue.push_back(payload);
596            self.mark_ready(sid, RtPayload::Nothing);
597        }
598    }
599
600    fn deliver_to_receiver(&mut self, rid: TaskId, ch: ChanId, value: RtPayload) {
601        let select_arms = self.tasks.get_mut(&rid).unwrap().select_arms.take();
602        if let Some(arms) = select_arms {
603            let arm_idx = arms
604                .iter()
605                .position(|a| matches!(a, SelectArm::Recv(c) if *c == ch))
606                .expect("delivering channel must be one of the select arms");
607            {
608                let slot = self.tasks.get_mut(&rid).unwrap();
609                slot.resume = value;
610                slot.selected_arm = Some(arm_idx);
611                slot.state = TaskState::Ready;
612            }
613            self.ready.push(rid);
614            self.cancel_select_registrations(rid, Some(ch));
615        } else {
616            self.mark_ready(rid, value);
617        }
618    }
619
620    fn cancel_select_registrations(&mut self, rid: TaskId, except: Option<ChanId>) {
621        for (cid, chan) in self.chans.iter_mut() {
622            if Some(*cid) == except {
623                continue;
624            }
625            chan.blocked_receivers.retain(|t| *t != rid);
626        }
627        self.timers.retain(|t| t.tid != rid);
628    }
629
630    fn do_select(&mut self, tid: TaskId, arms: Vec<SelectArm>) {
631        let ready_arms: Vec<usize> = arms
632            .iter()
633            .enumerate()
634            .filter_map(|(i, a)| match a {
635                SelectArm::Recv(ch) if self.chans[ch].can_recv() => Some(i),
636                _ => None,
637            })
638            .collect();
639
640        if !ready_arms.is_empty() {
641            let w = if ready_arms.len() == 1 {
642                0
643            } else {
644                self.chooser.decide(ChoiceKind::SelectWinner, ready_arms.len())
645            };
646            let arm_idx = ready_arms[w];
647            let ch = match &arms[arm_idx] {
648                SelectArm::Recv(ch) => *ch,
649                SelectArm::Timeout(_) => unreachable!("ready arms are recv arms"),
650            };
651            let v = self.take_one(ch);
652            {
653                let slot = self.tasks.get_mut(&tid).unwrap();
654                slot.resume = v;
655                slot.selected_arm = Some(arm_idx);
656                slot.state = TaskState::Ready;
657            }
658            self.ready.push(tid);
659        } else {
660            for (i, a) in arms.iter().enumerate() {
661                match a {
662                    SelectArm::Recv(ch) => {
663                        self.chans.get_mut(ch).unwrap().blocked_receivers.push_back(tid);
664                    }
665                    SelectArm::Timeout(d) => {
666                        self.timers.push(TimerEntry { fire_at: self.clock + *d, tid, arm: Some(i) });
667                    }
668                }
669            }
670            let slot = self.tasks.get_mut(&tid).unwrap();
671            slot.select_arms = Some(arms);
672            slot.state = TaskState::BlockedSelect;
673        }
674    }
675
676    fn take_one(&mut self, ch: ChanId) -> RtPayload {
677        if let Some(v) = self.chans.get_mut(&ch).unwrap().queue.pop_front() {
678            self.refill_from_blocked_sender(ch);
679            return v;
680        }
681        if let Some((sid, payload)) = self.chans.get_mut(&ch).unwrap().blocked_senders.pop_front() {
682            self.mark_ready(sid, RtPayload::Nothing);
683            return payload;
684        }
685        RtPayload::Nothing
686    }
687
688    fn advance_timers(&mut self) {
689        let earliest = match self.timers.iter().map(|t| t.fire_at).min() {
690            Some(e) => e,
691            None => return,
692        };
693        self.clock = earliest;
694
695        let mut due = Vec::new();
696        let mut rest = Vec::new();
697        for t in self.timers.drain(..) {
698            if t.fire_at == earliest {
699                due.push(t);
700            } else {
701                rest.push(t);
702            }
703        }
704        self.timers = rest;
705
706        // Seeded tie-break order among timers firing on the same tick.
707        let order: Vec<usize> = if due.len() > 1 {
708            let mut pool: Vec<usize> = (0..due.len()).collect();
709            let mut shuffled = Vec::with_capacity(pool.len());
710            while !pool.is_empty() {
711                let k = self.chooser.decide(ChoiceKind::TimerTieBreak, pool.len());
712                shuffled.push(pool.remove(k));
713            }
714            shuffled
715        } else {
716            vec![0]
717        };
718
719        for idx in order {
720            let tid = due[idx].tid;
721            let arm = due[idx].arm;
722            self.fire_timer(tid, arm);
723        }
724    }
725
726    fn fire_timer(&mut self, tid: TaskId, arm: Option<usize>) {
727        let still_blocked = matches!(
728            self.tasks.get(&tid).map(|s| &s.state),
729            Some(TaskState::BlockedTimer) | Some(TaskState::BlockedSelect)
730        );
731        if !still_blocked {
732            return;
733        }
734        match arm {
735            Some(arm_idx) => {
736                let had = self.tasks.get_mut(&tid).unwrap().select_arms.take();
737                {
738                    let slot = self.tasks.get_mut(&tid).unwrap();
739                    slot.selected_arm = Some(arm_idx);
740                    slot.resume = RtPayload::Nothing;
741                    slot.state = TaskState::Ready;
742                }
743                self.ready.push(tid);
744                if had.is_some() {
745                    self.cancel_select_registrations(tid, None);
746                }
747            }
748            None => self.mark_ready(tid, RtPayload::Nothing),
749        }
750    }
751
752    fn complete_task(&mut self, tid: TaskId, result: AwaitResult) {
753        let waiters = self.awaiters.remove(&tid).unwrap_or_default();
754        let payload = result.payload();
755        self.results.insert(tid, result);
756        for w in waiters {
757            self.mark_ready(w, payload.clone());
758        }
759    }
760
761    fn do_await(&mut self, tid: TaskId, target: TaskId) {
762        if let Some(res) = self.results.get(&target) {
763            let payload = res.payload();
764            self.mark_ready(tid, payload);
765        } else if !self.tasks.contains_key(&target) {
766            // Unknown target — treat as already gone.
767            self.mark_ready(tid, RtPayload::Nothing);
768        } else {
769            self.awaiters.entry(target).or_default().push(tid);
770            self.tasks.get_mut(&tid).unwrap().state = TaskState::BlockedAwait;
771        }
772    }
773
774    fn do_abort(&mut self, target: TaskId) {
775        if self.results.contains_key(&target) {
776            return; // already finished
777        }
778        match self.tasks.get(&target) {
779            None => return,
780            Some(s) if matches!(s.state, TaskState::Done) => return,
781            _ => {}
782        }
783        self.ready.retain(|t| *t != target);
784        for chan in self.chans.values_mut() {
785            chan.blocked_receivers.retain(|t| *t != target);
786            chan.blocked_senders.retain(|(t, _)| *t != target);
787        }
788        self.timers.retain(|t| t.tid != target);
789        // Don't let a future completion try to re-wake an aborted task.
790        for list in self.awaiters.values_mut() {
791            list.retain(|t| *t != target);
792        }
793        self.tasks.get_mut(&target).unwrap().state = TaskState::Done;
794        self.complete_task(target, AwaitResult::Aborted);
795    }
796
797    fn do_close(&mut self, ch: ChanId) {
798        let waiters: Vec<TaskId> = {
799            let chan = self.chans.get_mut(&ch).unwrap();
800            chan.closed = true;
801            chan.blocked_receivers.drain(..).collect()
802        };
803        for rid in waiters {
804            self.wake_closed_receiver(rid, ch);
805        }
806    }
807
808    fn wake_closed_receiver(&mut self, rid: TaskId, ch: ChanId) {
809        let select_arms = self.tasks.get_mut(&rid).unwrap().select_arms.take();
810        if let Some(arms) = select_arms {
811            let arm_idx = arms
812                .iter()
813                .position(|a| matches!(a, SelectArm::Recv(c) if *c == ch))
814                .unwrap_or(0);
815            {
816                let slot = self.tasks.get_mut(&rid).unwrap();
817                slot.resume = RtPayload::Nothing;
818                slot.selected_arm = Some(arm_idx);
819                slot.state = TaskState::Ready;
820            }
821            self.ready.push(rid);
822            self.cancel_select_registrations(rid, Some(ch));
823        } else {
824            self.mark_ready(rid, RtPayload::Nothing);
825        }
826    }
827}
828
829/// Run `setup` (which spawns tasks and creates channels) under a fresh recording
830/// scheduler seeded by `seed`, returning the outcome and the scheduling trace.
831pub fn run_with_seed<'t, F>(config: SchedulerConfig, seed: SchedSeed, setup: F) -> (RunOutcome, SchedTrace)
832where
833    F: FnOnce(&mut Scheduler<'t>),
834{
835    let mut sched = Scheduler::new(config, Chooser::record(seed));
836    setup(&mut sched);
837    let outcome = sched.run();
838    (outcome, sched.into_trace())
839}
840
841/// Re-run `setup` under a replaying scheduler driven by `trace`, returning the outcome.
842pub fn run_with_trace<'t, F>(config: SchedulerConfig, trace: SchedTrace, setup: F) -> RunOutcome
843where
844    F: FnOnce(&mut Scheduler<'t>),
845{
846    let mut sched = Scheduler::new(config, Chooser::replay(trace));
847    setup(&mut sched);
848    sched.run()
849}
logicaffeine_runtime/scheduler.rs

logicaffeine_runtime/
scheduler.rs
