From 00128f32a2a25a332c7a5c7a465fdb9aeaccfce0 Mon Sep 17 00:00:00 2001 From: smarm Date: Wed, 10 Jun 2026 07:45:00 +0000 Subject: [PATCH] =?UTF-8?q?feat(channel):=20select=20=E2=80=94=20ready-ind?= =?UTF-8?q?ex=20wait=20over=20multiple=20receivers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit select(&[&dyn Selectable]) -> usize registers (pid, epoch) in every arm under one wait epoch, parks once, and returns the first ready index in documented priority order (BEAM-style; no fairness promise). A closed arm counts as ready — and stays ready forever, so callers drop it from the set once observed. Losing arms need no cancellation pass: the winning wake consumed the epoch, so their registrations die at their wakers' failed CAS or are overwritten by the receiver's next wait on that channel (the single-receiver debug_asserts relax to 'none or own pid' accordingly). The one genuinely new protocol piece is the no-park exit: returning with an arm ready at registration time leaves earlier arms holding LIVE-epoch registrations, whose wakes could fault a later one-shot park as a pending notification. scheduler::retire_wait closes it — bump the epoch (in-flight wakes die at their CAS), eat a notification that already landed (StateWord::clear_notify), then re-observe the stop flag, in that order. Proved by the new loom theorem retire_eats_late_arm_notification; the integration probes in tests/select.rs fire stale loser-arm wakes and assert a subsequent sleep's one-shot park holds its full duration. --- src/channel.rs | 122 ++++++++++++++++++++++- src/lib.rs | 2 +- src/runtime.rs | 10 ++ src/scheduler.rs | 16 +++ src/slot_state.rs | 75 ++++++++++++++ tests/select.rs | 245 ++++++++++++++++++++++++++++++++++++++++++++++ 6 files changed, 466 insertions(+), 4 deletions(-) create mode 100644 tests/select.rs diff --git a/src/channel.rs b/src/channel.rs index 43b07be..80f8a20 100644 --- a/src/channel.rs +++ b/src/channel.rs @@ -170,7 +170,7 @@ impl Receiver { let me = crate::actor::current_pid() .expect("recv() called outside an actor"); debug_assert!( - g.parked_receiver.is_none(), + g.parked_receiver.is_none_or(|(p, _)| p == me), "channel has more than one receiver" ); // begin_wait is lock-free — legal under the Channel lock; @@ -227,7 +227,7 @@ impl Receiver { return Err(RecvTimeoutError::Disconnected); } debug_assert!( - g.parked_receiver.is_none(), + g.parked_receiver.is_none_or(|(p, _)| p == me), "channel has more than one receiver" ); epoch = crate::scheduler::begin_wait(); @@ -283,7 +283,7 @@ impl Receiver { let me = crate::actor::current_pid() .expect("recv_match() called outside an actor"); debug_assert!( - g.parked_receiver.is_none(), + g.parked_receiver.is_none_or(|(p, _)| p == me), "channel has more than one receiver" ); g.parked_receiver = Some((me, crate::scheduler::begin_wait())); @@ -356,3 +356,119 @@ impl crate::timer::TimerTarget for RawMutex> { } } } + +// --------------------------------------------------------------------------- +// select — ready-index wait over multiple receivers +// --------------------------------------------------------------------------- + +mod sealed { + pub trait Sealed {} +} +impl sealed::Sealed for Receiver {} + +/// An arm of a [`select`]. Implemented by [`Receiver`]; sealed, because the +/// registration contract below is part of the runtime's wake protocol. +/// +/// Contract (all under the arm's own lock): `sel_register` checks-or- +/// registers atomically — if the arm is ready it does NOT register and +/// returns `false`; otherwise it publishes `(pid, epoch)` where its wakers +/// will find it. "Ready" means a receive would not park: a message is +/// queued, or the arm is closed. +pub trait Selectable: sealed::Sealed { + #[doc(hidden)] + fn sel_register(&self, pid: Pid, epoch: u32) -> bool; + #[doc(hidden)] + fn sel_ready(&self) -> bool; +} + +impl Selectable for Receiver { + fn sel_register(&self, pid: Pid, epoch: u32) -> bool { + let mut g = self.inner.lock(); + if !g.queue.is_empty() || g.senders == 0 { + return false; + } + debug_assert!( + g.parked_receiver.is_none_or(|(p, _)| p == pid), + "channel has more than one receiver" + ); + g.parked_receiver = Some((pid, epoch)); + true + } + + fn sel_ready(&self) -> bool { + let g = self.inner.lock(); + !g.queue.is_empty() || g.senders == 0 + } +} + +/// Park on every arm at once; return the index of the first ready one. +/// +/// "Ready" means a receive on that arm would not park: a message is queued, +/// or the arm is **closed** (so the caller's `try_recv` observes the +/// disconnect — a dead arm is an event, not a hang). The caller consumes the +/// arm itself, typically via [`Receiver::try_recv`]; single-receiver +/// channels guarantee nothing can steal the message in between. +/// +/// A closed arm stays ready *forever*: once its disconnect has been +/// observed, drop it from the arm set — under priority order it would +/// otherwise win every subsequent call and starve every higher-indexed arm. +/// +/// Arms are scanned **in order**: index 0 is the highest priority, both on +/// the immediate-ready path and after a wake. This is a documented +/// guarantee (compose like BEAM receive clauses: put control channels +/// first), not an accident — and therefore there is NO fairness promise; a +/// saturated arm 0 starves arm 1 by design. +/// +/// One actor may select on a channel and later `recv` on it (or select on +/// overlapping sets) freely. What stays illegal is what was always illegal: +/// two *different* actors receiving on one channel. +/// +/// Built on the consuming-wake protocol (see slot_state.rs): all arms are +/// registered under one wait epoch; the winning wake consumes it, so losing +/// arms' registrations are inert and need no cancellation pass — they +/// self-clean at their wakers' failed CAS, or get overwritten by this +/// receiver's next wait on that channel. +/// +/// Panics if `arms` is empty, or when called outside an actor. +pub fn select(arms: &[&dyn Selectable]) -> usize { + assert!(!arms.is_empty(), "select() on an empty arm list"); + let me = crate::actor::current_pid().expect("select() called outside an actor"); + loop { + let epoch = crate::scheduler::begin_wait(); + + // Check-or-register each arm, in priority order, each atomically + // under its own lock. Cross-arm atomicity is unnecessary: an arm + // becoming ready right after its registration wakes us through the + // protocol (the prep-to-park window is closed by RunningNotified). + let mut ready = None; + for (i, arm) in arms.iter().enumerate() { + if !arm.sel_register(me, epoch) { + ready = Some(i); + break; + } + } + if let Some(i) = ready { + // Returning WITHOUT parking while earlier arms hold live-epoch + // registrations: retire the wait (bump the epoch, eat a + // notification that already landed, re-observe a pending stop) + // so no stale arm wake can fault a later one-shot park. + crate::scheduler::retire_wait(); + return i; + } + + crate::scheduler::park_current(); + // Woken precisely: an arm's send (message) or last-sender drop + // (closure) consumed our epoch, and both leave their arm ready — + // return the first one, in priority order (which may be a + // different, higher-priority arm than the one that woke us; its + // message stays queued and re-reports ready on the next call). + for (i, arm) in arms.iter().enumerate() { + if arm.sel_ready() { + return i; + } + } + // Unreachable by protocol (a stop wake unwinds out of + // park_current). Defensive: re-open the wait and re-register — + // stale own-registrations are overwritten. + } +} diff --git a/src/lib.rs b/src/lib.rs index 8ea33b5..f42b926 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -45,7 +45,7 @@ static ALLOCATOR: preempt::PreemptingAllocator = preempt::PreemptingAllocator; // Public API re-exports // --------------------------------------------------------------------------- -pub use channel::{channel, Receiver, RecvError, RecvTimeoutError, Sender}; +pub use channel::{channel, select, Receiver, RecvError, RecvTimeoutError, Selectable, Sender}; pub use gen_server::{CallError, CallTimeoutError, CastError, GenServer, ServerRef}; pub use link::{link, trap_exit, unlink, ExitSignal}; pub use monitor::{demonitor, monitor, Down, DownReason, Monitor, MonitorId}; diff --git a/src/runtime.rs b/src/runtime.rs index 7807bb8..2b428d2 100644 --- a/src/runtime.rs +++ b/src/runtime.rs @@ -553,6 +553,16 @@ impl RuntimeInner { slot.word.begin_wait(pid.generation()) } + /// Retire the calling actor's current wait without parking on it: bump + /// the epoch (invalidating every in-flight registration-based wake), + /// then eat a notification that already landed. The caller MUST + /// re-check its stop flag afterwards — see `StateWord::clear_notify`. + pub(crate) fn retire_wait(&self, pid: Pid) { + let slot = self.slot_at(pid).expect("retire_wait: own slot vanished"); + let _ = slot.word.begin_wait(pid.generation()); + slot.word.clear_notify(pid.generation()); + } + fn unpark_inner(&self, pid: Pid, want: Option) { if let Some(slot) = self.slot_at(pid) { match slot.word.unpark(pid.generation(), want) { diff --git a/src/scheduler.rs b/src/scheduler.rs index 5c4295b..28f9c7f 100644 --- a/src/scheduler.rs +++ b/src/scheduler.rs @@ -246,6 +246,22 @@ pub(crate) fn begin_wait() -> u32 { with_runtime(|inner| inner.begin_wait(me)) } +/// Close the current actor's wait WITHOUT parking on it. For the no-park +/// exit of multi-registration waits (`select` finding an arm ready at +/// registration time): bumps the epoch so in-flight wakes die at their CAS, +/// eats a notification that already landed, then observes a pending stop — +/// in that order. After this, leftover registrations are stale-epoch and +/// self-clean at their wakers' failed CAS; nothing can fault the actor's +/// next one-shot park. +pub(crate) fn retire_wait() { + let me = current_pid().expect("retire_wait() called outside an actor"); + with_runtime(|inner| inner.retire_wait(me)); + // A request_stop that fired before the clear had its notification + // eaten, but it set the stop flag first — observe it here and unwind. + // One that fires after re-notifies a Running word as usual. + crate::preempt::check_cancelled(); +} + /// Request cooperative cancellation of `pid`. /// /// Sets the actor's stop flag and wakes it so it observes the flag promptly: diff --git a/src/slot_state.rs b/src/slot_state.rs index 8ab37b4..01f65a5 100644 --- a/src/slot_state.rs +++ b/src/slot_state.rs @@ -321,6 +321,49 @@ impl StateWord { } } + /// Eat a pending notification: RunningNotified → Running, epoch + /// preserved; no-op on Running. Called by the RUNNING actor itself, on + /// the no-park exit of a wait it registered for but never parked on + /// (`select` returning a ready arm at registration time), AFTER bumping + /// the epoch and BEFORE re-checking its stop flag: + /// + /// - post-bump, the only wakers that can have set RunningNotified are + /// ones stamped with the just-retired epoch (a select arm) or a + /// terminal wildcard (`request_stop`); + /// - the caller's stop-flag check AFTER the clear catches the terminal + /// case (the flag is set before the wake fires), so eating its + /// notification loses nothing — and a stop arriving later re-notifies + /// a Running word as usual; + /// - what remains eaten is exactly the stale arm wake that would + /// otherwise fault the actor's next one-shot park. + /// + /// Returns whether a notification was eaten. + pub(crate) fn clear_notify(&self, gen: u32) -> bool { + loop { + let w = self.load(); + debug_assert!( + matches!(word_state(w), ST_RUNNING | ST_RUNNING_NOTIFIED) + && word_gen(w) == gen, + "clear_notify from invalid word {w:#x}" + ); + if word_state(w) != ST_RUNNING_NOTIFIED { + return false; + } + if self + .0 + .compare_exchange( + w, + pack(gen, word_epoch(w), ST_RUNNING), + Ordering::AcqRel, + Ordering::Acquire, + ) + .is_ok() + { + return true; + } + } + } + /// Finalize: Running | RunningNotified → Done, epoch zeroed (wait /// identity never crosses an occupancy). Called by the scheduler that /// just ran the actor to completion (so those are the only legal prior @@ -499,6 +542,38 @@ mod loom_tests { }); } + /// The retire theorem — `select`'s no-park exit. An actor opens a wait + /// and registers, then finds an arm ready and returns WITHOUT parking; + /// a loser arm's waker fires concurrently, stamped with the live epoch. + /// The exit retires the wait (bump, then eat): in every interleaving + /// the run ends on a clean Running word — no pending notification + /// survives to fault the actor's next one-shot park — and the waker + /// never enqueues. + #[test] + fn retire_eats_late_arm_notification() { + loom::model(|| { + let word = Arc::new(StateWord::new()); + word.publish_queued(0); + assert!(word.try_claim(0)); + let epoch = word.begin_wait(0); // select opens + registers + + let w = word.clone(); + let waker = thread::spawn(move || w.unpark(0, Some(epoch))); + + // No-park exit: bump (invalidates in-flight wakes), then eat + // (consumes one that already landed). + let _ = word.begin_wait(0); + word.clear_notify(0); + + assert_ne!(waker.join().unwrap(), Unpark::Enqueue); + assert_eq!( + word_state(word.load()), + ST_RUNNING, + "stale arm wake survived the retire" + ); + }); + } + /// The ABA theorem: a stale-generation unpark racing reclaim + reuse can /// never touch the slot's new occupant. #[test] diff --git a/tests/select.rs b/tests/select.rs new file mode 100644 index 0000000..d51f4ed --- /dev/null +++ b/tests/select.rs @@ -0,0 +1,245 @@ +//! select tests. Beyond the functional cases, the *_stays_precise tests are +//! soundness probes for the consuming-wake protocol: they fire stale loser- +//! arm wakes at an actor and then run a one-shot park (`sleep`, whose early +//! return would be observable as a short elapsed time) to prove the stale +//! wake died at its epoch CAS instead of corrupting the next wait. + +use smarm::{channel, run, select, spawn}; +use std::sync::atomic::{AtomicI64, Ordering}; +use std::time::{Duration, Instant}; + +static OUT: AtomicI64 = AtomicI64::new(0); + +#[test] +fn ready_arm_returns_immediately_without_parking() { + OUT.store(0, Ordering::SeqCst); + run(|| { + let (txa, rxa) = channel::(); + let (_txb, rxb) = channel::(); + txa.send(42).unwrap(); + let i = select(&[&rxb, &rxa]); + assert_eq!(i, 1); + OUT.store(rxa.try_recv().unwrap().expect("ready arm must hold a message"), Ordering::SeqCst); + }); + assert_eq!(OUT.load(Ordering::SeqCst), 42); +} + +#[test] +fn lower_index_wins_when_several_arms_are_ready() { + run(|| { + let (txa, rxa) = channel::(); + let (txb, rxb) = channel::(); + txa.send(1).unwrap(); + txb.send(2).unwrap(); + // Documented priority order: index 0 first, regardless of send order. + assert_eq!(select(&[&rxa, &rxb]), 0); + assert_eq!(select(&[&rxb, &rxa]), 0); + }); +} + +#[test] +fn parks_until_any_arm_fires() { + OUT.store(0, Ordering::SeqCst); + run(|| { + let (txa, _rxa_keepalive) = (channel::().0, ()); + let _hold = txa; // arm a: sender alive, never sends + let (txa, rxa) = channel::(); + let (txb, rxb) = channel::(); + let _keep_a_open = txa; + let h = spawn(move || { + smarm::sleep(Duration::from_millis(10)); + txb.send(7).unwrap(); + }); + let i = select(&[&rxa, &rxb]); + assert_eq!(i, 1); + OUT.store(rxb.try_recv().unwrap().unwrap(), Ordering::SeqCst); + h.join().unwrap(); + }); + assert_eq!(OUT.load(Ordering::SeqCst), 7); +} + +#[test] +fn closed_arm_counts_as_ready() { + run(|| { + let (_keep, rxa) = channel::(); + let (txb, rxb) = channel::(); + drop(txb); + let i = select(&[&rxa, &rxb]); + assert_eq!(i, 1); + // The caller observes the closure on the arm itself. + assert!(rxb.try_recv().is_err()); + }); +} + +#[test] +fn closure_while_parked_wakes_the_select() { + run(|| { + let (_keep, rxa) = channel::(); + let (txb, rxb) = channel::(); + let h = spawn(move || { + smarm::sleep(Duration::from_millis(10)); + drop(txb); + }); + assert_eq!(select(&[&rxa, &rxb]), 1); + assert!(rxb.try_recv().is_err()); + h.join().unwrap(); + }); +} + +#[test] +fn loser_arm_wake_after_parked_select_stays_precise() { + // Both arms registered; arm 0 wins the wake; arm 1's sender then fires + // a wake stamped with the consumed epoch. The subsequent sleep is a + // one-shot park: an early return would mean the stale wake landed. + run(|| { + let (txa, rxa) = channel::(); + let (txb, rxb) = channel::(); + let h = spawn(move || { + smarm::sleep(Duration::from_millis(10)); + txa.send(1).unwrap(); // wins + txb.send(2).unwrap(); // loser arm: stale-epoch wake + }); + assert_eq!(select(&[&rxa, &rxb]), 0); + assert_eq!(rxa.try_recv().unwrap(), Some(1)); + let t0 = Instant::now(); + smarm::sleep(Duration::from_millis(40)); + assert!( + t0.elapsed() >= Duration::from_millis(40), + "one-shot park returned early: a stale loser-arm wake landed" + ); + // The loser's message was never lost. + assert_eq!(select(&[&rxa, &rxb]), 1); + assert_eq!(rxb.try_recv().unwrap(), Some(2)); + h.join().unwrap(); + }); +} + +#[test] +fn no_park_exit_retires_the_wait() { + // Arm 1 is ready at registration time, so select returns WITHOUT + // parking while arm 0 holds a live-epoch registration. Arm 0's sender + // then fires. retire_wait must have invalidated/eaten that wake; + // the sleep proves it. + run(|| { + let (txa, rxa) = channel::(); + let (txb, rxb) = channel::(); + txb.send(2).unwrap(); + assert_eq!(select(&[&rxa, &rxb]), 1); + assert_eq!(rxb.try_recv().unwrap(), Some(2)); + + let h = spawn(move || { + txa.send(1).unwrap(); // fires at the (retired) select epoch + }); + h.join().unwrap(); + + let t0 = Instant::now(); + smarm::sleep(Duration::from_millis(40)); + assert!( + t0.elapsed() >= Duration::from_millis(40), + "one-shot park returned early: the no-park exit leaked a wake" + ); + assert_eq!(rxa.try_recv().unwrap(), Some(1)); + }); +} + +#[test] +fn select_then_plain_recv_on_a_loser_arm() { + // A leftover own-registration from a select must not trip the next + // direct wait on that channel (it is overwritten, not asserted away). + run(|| { + let (txa, rxa) = channel::(); + let (txb, rxb) = channel::(); + txa.send(1).unwrap(); + assert_eq!(select(&[&rxa, &rxb]), 0); // rxb may keep a registration + assert_eq!(rxa.try_recv().unwrap(), Some(1)); + let h = spawn(move || { + smarm::sleep(Duration::from_millis(5)); + txb.send(9).unwrap(); + }); + assert_eq!(rxb.recv().unwrap(), 9); + h.join().unwrap(); + }); +} + +#[test] +fn select_loop_drains_two_producers_completely() { + const N: i64 = 200; + OUT.store(0, Ordering::SeqCst); + run(|| { + let (txa, rxa) = channel::(); + let (txb, rxb) = channel::(); + let ha = spawn(move || { + for v in 0..N { + txa.send(v).unwrap(); + if v % 7 == 0 { + smarm::yield_now(); + } + } + }); + let hb = spawn(move || { + for v in 0..N { + txb.send(v).unwrap(); + if v % 5 == 0 { + smarm::yield_now(); + } + } + }); + // A closed arm stays permanently ready, so it must leave the arm + // set once observed (keeping it would starve the other arm at + // priority order) — select until the FIRST closure, then drain the + // survivor with plain recv. + let mut sum = 0i64; + let survivor: &smarm::Receiver = loop { + let i = select(&[&rxa, &rxb]); + let arm: &smarm::Receiver = if i == 0 { &rxa } else { &rxb }; + match arm.try_recv() { + Ok(Some(v)) => sum += v, + Ok(None) => {} // defensive: ready arm already drained + Err(_) => break if i == 0 { &rxb } else { &rxa }, + } + }; + loop { + match survivor.recv() { + Ok(v) => sum += v, + Err(_) => break, + } + } + OUT.store(sum, Ordering::SeqCst); + ha.join().unwrap(); + hb.join().unwrap(); + }); + assert_eq!(OUT.load(Ordering::SeqCst), 2 * (0..200i64).sum::()); +} + +#[test] +fn motivating_pattern_inbox_plus_monitor_down() { + // The gen_server/handle_info shape select unlocks: a server waiting on + // its request inbox AND a monitor Down channel in one park. + use smarm::monitor; + run(|| { + let (req_tx, req_rx) = channel::(); + let worker = spawn(|| { + smarm::sleep(Duration::from_millis(10)); + // worker exits -> Down fires + }); + let m = monitor(worker.pid()); + let _keep_inbox_open = req_tx; + + let mut served = 0; + loop { + match select(&[&m.rx, &req_rx]) { + 0 => { + let _down = m.rx.try_recv().unwrap().expect("Down message"); + break; + } + _ => { + if let Ok(Some(_)) = req_rx.try_recv() { + served += 1; + } + } + } + } + assert_eq!(served, 0); + worker.join().unwrap(); + }); +}