From fbb9d6a44eacc978a82ec421c62fc8cec6c4bf98 Mon Sep 17 00:00:00 2001 From: tison Date: Fri, 11 Sep 2026 02:30:36 +0800 Subject: [PATCH] feat: add runtime-agnostic select macro --- CHANGELOG.md | 1 + Cargo.lock | 7 + Cargo.toml | 1 + README.md | 3 +- asyncband/Cargo.toml | 2 + asyncband/src/lib.rs | 7 + asyncband/src/select.rs | 233 ++++++++++++++++ examples/Cargo.toml | 7 + examples/src/select_messages.rs | 73 +++++ tests-integration/Cargo.toml | 1 + tests-integration/tests/select_test.rs | 362 +++++++++++++++++++++++++ 11 files changed, 696 insertions(+), 1 deletion(-) create mode 100644 asyncband/src/select.rs create mode 100644 examples/src/select_messages.rs create mode 100644 tests-integration/tests/select_test.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index ea2f1e3c..643749ed 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,7 @@ All notable changes to this project will be documented in this file. ### New features +* Add the opt-in `select!` macro to wait for one of several `IntoFuture` branches, with random starting order, explicit `biased;` priority, branch conditions, and an all-disabled `else` fallback; owned futures are dropped before the selected handler runs, while borrowed futures can be retained across selections. * Add bounded MPSC `reserve` and `try_reserve` methods returning a `Permit`, allowing callers to wait for capacity before constructing a message; pending sends and reservations receive capacity in wait-queue order, and unused permits release capacity without claiming message order. ### Bug fixes diff --git a/Cargo.lock b/Cargo.lock index c6baa36c..5a1c90cd 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -80,6 +80,7 @@ dependencies = [ name = "asyncband" version = "0.7.2" dependencies = [ + "fastrand", "hashbrown", "tokio", ] @@ -362,6 +363,12 @@ dependencies = [ "tokio", ] +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "find-msvc-tools" version = "0.1.9" diff --git a/Cargo.toml b/Cargo.toml index 05e963ce..9d88226e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -32,6 +32,7 @@ rust-version = "1.86.0" asyncband = { path = "asyncband" } # Optional runtime dependencies +fastrand = { version = "2.5.0" } hashbrown = { version = "0.17.1", default-features = false } # Dev dependencies diff --git a/README.md b/README.md index 1203ae01..b28e028f 100644 --- a/README.md +++ b/README.md @@ -58,7 +58,7 @@ Public paths stay direct—such as `asyncband::mutex`, `asyncband::pool`, and `a ## Examples -Runnable examples live in the [`examples`](examples) workspace crate. They demonstrate how to choose and compose Asyncband primitives in complete programs. +Runnable examples live in the [`examples`](examples) workspace crate. They demonstrate how to choose and compose Asyncband primitives in complete programs. The [`select_messages`](examples/src/select_messages.rs) example combines message reception, a retained stop notification, and a caller-owned deadline using `select!`. ## API map @@ -84,6 +84,7 @@ Runnable examples live in the [`examples`](examples) workspace crate. They demon | | [`broadcast`](https://docs.rs/asyncband/*/asyncband/broadcast/) | `broadcast` | Deliver every value to receivers active at send time; retain an unbounded backlog until each consumes or drops. | | | [`watch`](https://docs.rs/asyncband/*/asyncband/watch/) | `watch` | Publish cloneable latest state from one or more senders; receivers independently coalesce intermediate updates. | | Object reuse | [`pool`](https://docs.rs/asyncband/*/asyncband/pool/) | `pool` | Reuse objects through bounded or unbounded pool variants. | +| Future composition | [`select!`](https://docs.rs/asyncband/*/asyncband/macro.select.html) | `select` | Wait for one asynchronous branch, with explicit cancellation and polling order. | | Sync interop | [`FutureExt`](https://docs.rs/asyncband/*/asyncband/blocking/trait.FutureExt.html) | `blocking` | Drive one runtime-agnostic future from a blocking thread. | ## Synchronous interoperability diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index d2d74f45..f5770f66 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -60,6 +60,7 @@ once-map = ["dep:hashbrown", "once-cell"] oneshot = [] pool = ["semaphore"] rwlock = [] +select = ["dep:fastrand"] semaphore = [] shutdown = ["latch", "waitgroup"] singleflight = ["dep:hashbrown", "once-cell"] @@ -67,6 +68,7 @@ waitgroup = [] watch = [] [dependencies] +fastrand = { workspace = true, optional = true } hashbrown = { workspace = true, default-features = false, features = [ "inline-more", ], optional = true } diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index 93ce485e..3573e7c5 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -76,6 +76,7 @@ //! | | [`broadcast`] | `broadcast` | Deliver every value to receivers active at send time; retain an unbounded backlog until each consumes or drops. | //! | | [`watch`] | `watch` | Publish cloneable latest state from one or more senders; receivers independently coalesce intermediate updates. | //! | Object reuse | [`pool`] | `pool` | Reuse objects through bounded or unbounded pool variants. | +//! | Future composition | [`select!`] | `select` | Wait for one asynchronous branch, with explicit cancellation and polling order. | //! | Sync interop | [`FutureExt`](blocking::FutureExt) | `blocking` | Drive one runtime-agnostic future from a blocking thread. | //! //! # Scope and runtime model @@ -118,6 +119,12 @@ //! code, it does indicate that the project has yet to be fully endorsed by the ASF. mod internal; +#[cfg(feature = "select")] +mod select; +#[cfg(feature = "select")] +#[doc(hidden)] +pub use select::random_start as __select_random_start; + #[cfg(feature = "barrier")] pub mod barrier; #[cfg(feature = "blocking")] diff --git a/asyncband/src/select.rs b/asyncband/src/select.rs new file mode 100644 index 00000000..2da0e7fc --- /dev/null +++ b/asyncband/src/select.rs @@ -0,0 +1,233 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +/// Implementation support for the exported macro. +#[doc(hidden)] +pub fn random_start(branches: usize) -> usize { + fastrand::usize(..branches) +} + +/// Waits for one of several asynchronous branches to complete. +/// +/// Available with the `select` feature. Each branch has the form +/// `binding = expression, if condition => handler`, with an optional condition. Expressions +/// implement [`IntoFuture`](std::future::IntoFuture), and their outputs may have different types. +/// Every handler must produce the same result type, which becomes the value of the selection. +/// Branches are separated by commas; the last comma is optional. Up to 32 branches are supported. +/// +/// Bindings must be irrefutable patterns, such as `value`, `_`, or `(left, right)`. Handle errors, +/// channel disconnection, and other alternatives explicitly inside the handler. Unlike selectors +/// that filter outputs with refutable patterns, this macro never discards a completed result to +/// try another branch. +/// +/// ``` +/// # async fn example() { +/// let value = asyncband::select! { +/// value = async { 42 } => value, +/// _ = std::future::pending::<()>() => 0, +/// }; +/// assert_eq!(value, 42); +/// # } +/// ``` +/// +/// # Execution and polling order +/// +/// The macro awaits its branches on the current task, without spawning tasks or allocating storage +/// for them on the heap. Branches need not be `Send`, `'static`, or `Unpin`. A branch that blocks +/// inside `poll` prevents every other branch from progressing. A ready error is a completed +/// result, just like a ready success. +/// +/// Each poll starts at a randomly chosen branch and scans circularly, stopping at the first ready +/// enabled branch. This reduces fixed-order bias across repeated selections; it does not guarantee +/// equal probabilities among ready branches, starvation freedom, or fairness between tasks. Add +/// `biased;` to poll in source order instead. Put higher-priority branches first in that mode. +/// +/// ``` +/// # async fn example() { +/// let value = asyncband::select! { +/// biased; +/// value = async { 1 } => value, +/// value = async { 2 } => value, +/// }; +/// assert_eq!(value, 1); +/// # } +/// ``` +/// +/// # Conditions and fallback +/// +/// All conditions are evaluated once in source order, before constructing any branch. Then every +/// branch expression is evaluated and converted with `IntoFuture::into_future` in source order, +/// including disabled branches. A false condition prevents polling, not construction or ownership +/// transfer. Conditions remain fixed for this selection, even after a wakeup. +/// +/// An optional final `else => handler` runs when all conditions are false. It does not run when +/// enabled branches are merely pending. Without an `else`, an all-disabled selection panics. At +/// least one asynchronous branch is required; an empty selection is a compile error. +/// +/// # Cancellation and reuse +/// +/// Before running the selected handler, the macro drops all branch futures it owns, including the +/// completed one. This ends their borrows and runs their cancellation cleanup. Handlers may await, +/// use `?`, return from the enclosing function, or break and continue enclosing loops. +/// +/// Dropping a future does not undo work already performed by its polls. For example, cancelling a +/// pending send may drop its message, cancelling a lock request loses its queue position, and +/// cancelling a barrier wait does not retract its arrival. Dropping a task handle also does not +/// necessarily stop the underlying task. Consult each operation's cancellation contract. +/// +/// To preserve an unfinished operation across selections, create and pin it outside the selection, +/// then pass `future.as_mut()`. Only that temporary borrow is dropped when another branch wins; +/// the underlying future remains available to poll again. Disable or remove it after completion: +/// this macro does not fuse futures across separate selections. +/// +/// ``` +/// # async fn example() { +/// let mut operation = std::pin::pin!(async { String::from("done") }); +/// let mut completed = false; +/// loop { +/// asyncband::select! { +/// result = operation.as_mut(), if !completed => { +/// assert_eq!(result, "done"); +/// completed = true; +/// }, +/// else => break, +/// } +/// } +/// # } +/// ``` +/// +/// Timers are ordinary caller-provided branches. Keep a deadline future outside a loop when its +/// deadline must survive other branches winning. The macro has no built-in timer or nonblocking +/// `default` branch. +/// +/// ```compile_fail +/// # async fn example() { +/// asyncband::select! { +/// Ok(value) = async { Ok::<_, ()>(1) } => value, +/// }; +/// # } +/// ``` +#[macro_export] +macro_rules! select { + (@condition) => { true }; + (@condition $condition:expr) => { $condition }; + (@start biased $count:expr) => { 0usize }; + (@start random $count:expr) => { $crate::__select_random_start($count) }; + + (@build $mode:ident [ + $(($index:tt $variant:ident ($binding:pat) ($future:expr) ($condition:expr) ($handler:expr)))+ + ] ($fallback:expr)) => {{ + enum __SelectOutput<$($variant),+> { + $($variant($variant),)+ + Disabled, + } + + let output = { + let enabled = [$($condition,)+]; + // Pin each future separately so tuple projection needs no unsafe code. This scope also + // drops the underlying futures, rather than just their pins, before invoking a handler. + let mut futures = ($( + ::core::pin::pin!(::core::future::IntoFuture::into_future($future)), + )+); + + ::core::future::poll_fn(|cx| { + if !enabled.iter().any(|enabled| *enabled) { + return ::core::task::Poll::Ready(__SelectOutput::Disabled); + } + let start = $crate::select!(@start $mode enabled.len()); + for offset in 0..enabled.len() { + let index = (start + offset) % enabled.len(); + if !enabled[index] { + continue; + } + match index { + $( + $index => { + if let ::core::task::Poll::Ready(value) = + ::core::future::Future::poll(futures.$index.as_mut(), cx) + { + return ::core::task::Poll::Ready(__SelectOutput::$variant(value)); + } + } + )+ + _ => ::core::unreachable!(), + } + } + ::core::task::Poll::Pending + }).await + }; + + match output { + $( + __SelectOutput::$variant(value) => { + let $binding = value; + $handler + } + )+ + __SelectOutput::Disabled => $fallback, + } + }}; + + (@collect $mode:ident [] [$($slots:tt)*]; $(,)?) => { + ::core::compile_error!("select! requires at least one asynchronous branch") + }; + (@collect $mode:ident [] [$($slots:tt)*]; else => $fallback:expr $(,)?) => { + ::core::compile_error!("select! requires at least one asynchronous branch") + }; + (@collect $mode:ident [$($branches:tt)+] [$($slots:tt)*]; $(,)?) => { + $crate::select!(@build $mode [$($branches)+] + (::core::panic!("select! has no enabled branches"))) + }; + (@collect $mode:ident [$($branches:tt)+] [$($slots:tt)*]; else => $fallback:expr $(,)?) => { + $crate::select!(@build $mode [$($branches)+] ($fallback)) + }; + (@collect $mode:ident [$($branches:tt)*] [($index:tt $variant:ident) $($slots:tt)*]; + $binding:pat = $future:expr $(, if $condition:expr)? => $handler:expr, $($rest:tt)* + ) => { + $crate::select!(@collect $mode [ + $($branches)* + ($index $variant ($binding) ($future) + ($crate::select!(@condition $($condition)?)) ($handler)) + ] [$($slots)*]; $($rest)*) + }; + (@collect $mode:ident [$($branches:tt)*] [$($slots:tt)*]; + $binding:pat = $future:expr $(, if $condition:expr)? => $handler:expr + ) => { + $crate::select!(@collect $mode [$($branches)*] [$($slots)*]; + $binding = $future $(, if $condition)? => $handler,) + }; + (@collect $mode:ident [$($branches:tt)*] []; $($rest:tt)+) => { + ::core::compile_error!("select! supports at most 32 asynchronous branches") + }; + (@collect $($invalid:tt)*) => { + ::core::compile_error!("expected `binding = future, if condition => handler,` or a final `else => handler`") + }; + (@init $mode:ident; $($branches:tt)*) => { + $crate::select!(@collect $mode [] [ + (0 V0) (1 V1) (2 V2) (3 V3) (4 V4) (5 V5) (6 V6) (7 V7) + (8 V8) (9 V9) (10 V10) (11 V11) (12 V12) (13 V13) (14 V14) (15 V15) + (16 V16) (17 V17) (18 V18) (19 V19) (20 V20) (21 V21) (22 V22) (23 V23) + (24 V24) (25 V25) (26 V26) (27 V27) (28 V28) (29 V29) (30 V30) (31 V31) + ]; $($branches)*) + }; + (biased; $($branches:tt)*) => { + $crate::select!(@init biased; $($branches)*) + }; + ($($branches:tt)*) => { + $crate::select!(@init random; $($branches)*) + }; +} diff --git a/examples/Cargo.toml b/examples/Cargo.toml index 6e079a89..fab171a3 100644 --- a/examples/Cargo.toml +++ b/examples/Cargo.toml @@ -29,7 +29,10 @@ release = false asyncband = { workspace = true, features = [ "completion", "lazy-cell", + "mpsc", "once-cell", + "oneshot", + "select", "shutdown", ] } tokio = { workspace = true, features = [ @@ -58,3 +61,7 @@ path = "src/graceful_shutdown.rs" [[example]] name = "shared_completion" path = "src/shared_completion.rs" + +[[example]] +name = "select_messages" +path = "src/select_messages.rs" diff --git a/examples/src/select_messages.rs b/examples/src/select_messages.rs new file mode 100644 index 00000000..37a7b91b --- /dev/null +++ b/examples/src/select_messages.rs @@ -0,0 +1,73 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Receive messages while retaining a stop notification and one deadline across selections. +//! +//! Asyncband supplies the selection and channels. The application supplies Tokio's timer and task +//! execution. Once the message channel closes, its branch is disabled so its ready error cannot +//! turn the loop into a busy loop. + +use std::future::IntoFuture; +use std::pin::pin; +use std::time::Duration; + +use asyncband::mpsc; +use asyncband::oneshot; + +#[tokio::main(flavor = "current_thread")] +async fn main() { + let (sender, mut receiver) = mpsc::bounded(1); + let (stop_sender, stop_receiver) = oneshot::channel(); + let producer = tokio::spawn(async move { + for index in 1..=3 { + if sender.send(format!("message {index}")).await.is_err() { + return; + } + } + drop(sender); + tokio::time::sleep(Duration::from_millis(50)).await; + let _ = stop_sender.send(()); + }); + + // Recreating these operations inside the loop would discard the oneshot receiver or restart + // the timeout each time a message wins. Passing pinned borrows preserves both operations. + let mut stop = pin!(stop_receiver.into_future()); + let mut deadline = pin!(tokio::time::sleep(Duration::from_secs(1))); + let mut messages_open = true; + loop { + asyncband::select! { + biased; + result = stop.as_mut() => { + println!("producer finished: {result:?}"); + break; + }, + _ = deadline.as_mut() => { + println!("deadline reached"); + break; + }, + result = receiver.recv(), if messages_open => match result { + Ok(message) => println!("received {message}"), + Err(_) => { + messages_open = false; + println!("message channel drained; waiting for stop"); + } + }, + } + } + drop(receiver); + producer.await.unwrap(); +} diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 04df5894..6839482c 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -43,6 +43,7 @@ asyncband = { workspace = true, features = [ "oneshot", "pool", "rwlock", + "select", "semaphore", "shutdown", "singleflight", diff --git a/tests-integration/tests/select_test.rs b/tests-integration/tests/select_test.rs new file mode 100644 index 00000000..5864942a --- /dev/null +++ b/tests-integration/tests/select_test.rs @@ -0,0 +1,362 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cell::Cell; +use std::future::IntoFuture; +use std::future::pending; +use std::future::poll_fn; +use std::future::ready; +use std::pin::pin; +use std::rc::Rc; +use std::task::Poll; + +use asyncband::barrier::Barrier; +use asyncband::blocking::FutureExt as _; +use asyncband::mpsc; +use asyncband::oneshot; +use tests_integration::WakeCounter; +use tests_integration::expect_ready; +use tests_integration::poll_once; +use tests_integration::poll_with; + +#[tokio::test] +async fn selection_remains_send_and_works_on_an_executor() { + let (sender, receiver) = oneshot::channel(); + let task = tokio::spawn(async move { + asyncband::select! { + result = receiver => result.unwrap(), + _ = pending::<()>() => unreachable!(), + } + }); + sender.send(String::from("ready")).unwrap(); + assert_eq!(task.await.unwrap(), "ready"); +} + +#[test] +fn selection_consumes_only_one_ready_message() { + let (left_tx, mut left_rx) = mpsc::unbounded(); + let (right_tx, mut right_rx) = mpsc::unbounded(); + left_tx.send(String::from("left")).unwrap(); + right_tx.send(String::from("right")).unwrap(); + + let selected_left = async { + asyncband::select! { + value = left_rx.recv() => { + assert_eq!(value.unwrap(), "left"); + true + }, + value = right_rx.recv() => { + assert_eq!(value.unwrap(), "right"); + false + }, + } + } + .block_on(); + + if selected_left { + assert!(left_rx.try_recv().is_err()); + assert_eq!(right_rx.try_recv().unwrap(), "right"); + } else { + assert_eq!(left_rx.try_recv().unwrap(), "left"); + assert!(right_rx.try_recv().is_err()); + } +} + +#[test] +fn ready_errors_are_delivered_without_polling_lower_priority_branches() { + let (sender, receiver) = oneshot::channel::(); + drop(sender); + async { + asyncband::select! { + biased; + result = receiver => assert!(result.is_err()), + _ = poll_fn(|_| -> Poll<()> { panic!("polled after choosing a result") }) => {}, + } + } + .block_on(); +} + +#[test] +fn a_later_branch_can_wake_an_earlier_pending_branch() { + let signalled = Cell::new(false); + let mut selection = pin!(async { + asyncband::select! { + biased; + value = poll_fn(|_| { + if signalled.get() { Poll::Ready(42) } else { Poll::Pending } + }) => value, + _ = poll_fn(|cx| { + if !signalled.replace(true) { + cx.waker().wake_by_ref(); + } + Poll::<()>::Pending + }) => unreachable!(), + } + }); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_with(selection.as_mut(), &waker).is_pending()); + assert_eq!(wakes.take(), 1); + assert_eq!(expect_ready(poll_with(selection.as_mut(), &waker)), 42); +} + +#[test] +fn pending_branches_receive_the_latest_parent_waker() { + let (left_tx, left_rx) = oneshot::channel::(); + let (right_tx, right_rx) = oneshot::channel::(); + let mut selection = pin!(async { + asyncband::select! { + result = left_rx => result.unwrap(), + result = right_rx => result.unwrap().len() as i32, + } + }); + let (old_waker, old_wakes) = WakeCounter::new(); + let (new_waker, new_wakes) = WakeCounter::new(); + assert!(poll_with(selection.as_mut(), &old_waker).is_pending()); + assert!(poll_with(selection.as_mut(), &new_waker).is_pending()); + left_tx.send(7).unwrap(); + assert_eq!(old_wakes.count(), 0); + assert_eq!(new_wakes.count(), 1); + assert_eq!(expect_ready(poll_with(selection.as_mut(), &new_waker)), 7); + assert!(right_tx.send(String::from("cancelled")).is_err()); +} + +#[test] +fn dropping_the_selection_cancels_all_owned_branches() { + let (left_tx, left_rx) = oneshot::channel::<()>(); + let (right_tx, right_rx) = oneshot::channel::<()>(); + { + let mut selection = pin!(async { + asyncband::select! { + result = left_rx => result, + result = right_rx => result, + } + }); + assert!(poll_once(selection.as_mut()).is_pending()); + } + assert!(left_tx.send(()).is_err()); + assert!(right_tx.send(()).is_err()); +} + +#[test] +fn cancelling_a_granted_reservation_releases_capacity_before_the_handler() { + let (sender, mut receiver) = mpsc::bounded(1); + let mut held = Some(sender.try_reserve().unwrap()); + async { + asyncband::select! { + biased; + _ = sender.reserve() => panic!("capacity is initially occupied"), + _ = poll_fn(|_| { + // The earlier branch receives a capacity grant after it returned Pending. + drop(held.take()); + Poll::Ready(()) + }) => sender.try_send(7).unwrap(), + } + } + .block_on(); + assert_eq!(receiver.try_recv().unwrap(), 7); +} + +#[test] +fn borrowed_barrier_arrival_survives_another_branch_winning() { + let barrier = Barrier::new(2); + let mut arrival = pin!(barrier.wait()); + async { + asyncband::select! { + biased; + _ = arrival.as_mut() => panic!("second participant has not arrived"), + _ = ready(()) => {}, + } + } + .block_on(); + + assert!(barrier.wait().block_on().is_leader()); + assert!(!arrival.block_on().is_leader()); +} + +struct DropFlag(Rc>); + +impl Drop for DropFlag { + fn drop(&mut self) { + self.0.set(true); + } +} + +struct OwnedReady(DropFlag); + +impl Future for OwnedReady { + type Output = usize; + + fn poll(self: std::pin::Pin<&mut Self>, _: &mut std::task::Context<'_>) -> Poll { + assert!(!self.0.0.get()); + Poll::Ready(42) + } +} + +#[test] +fn completed_future_is_dropped_before_the_handler() { + let dropped = Rc::new(Cell::new(false)); + let future = OwnedReady(DropFlag(dropped.clone())); + async { + asyncband::select! { + value = future => { + assert!(dropped.get()); + assert_eq!(value, 42); + }, + } + } + .block_on(); +} + +struct DisabledAwaitable { + converted: Rc>, + drop_flag: DropFlag, +} + +impl IntoFuture for DisabledAwaitable { + type Output = DropFlag; + type IntoFuture = std::future::Ready; + + fn into_future(self) -> Self::IntoFuture { + self.converted.set(true); + ready(self.drop_flag) + } +} + +#[test] +fn disabled_branches_are_constructed_converted_and_dropped_before_fallback() { + let converted = Rc::new(Cell::new(false)); + let dropped = Rc::new(Cell::new(false)); + let conditions = Cell::new(0); + async { + asyncband::select! { + _ = { + assert_eq!(conditions.get(), 2); + DisabledAwaitable { + converted: converted.clone(), + drop_flag: DropFlag(dropped.clone()), + } + }, if { conditions.set(conditions.get() + 1); false } => panic!("disabled branch ran"), + _ = async { panic!("disabled branch was polled") }, + if { conditions.set(conditions.get() + 1); false } => {}, + else => { + assert!(converted.get()); + assert!(dropped.get()); + }, + } + } + .block_on(); +} + +#[test] +fn conditions_are_not_reevaluated_and_fallback_does_not_replace_pending() { + let condition_checks = Cell::new(0); + let (sender, receiver) = oneshot::channel(); + let mut selection = pin!(async { + asyncband::select! { + value = receiver, if { condition_checks.set(condition_checks.get() + 1); true } => value, + else => panic!("enabled branch is pending"), + } + }); + assert!(poll_once(selection.as_mut()).is_pending()); + assert!(poll_once(selection.as_mut()).is_pending()); + assert_eq!(condition_checks.get(), 1); + sender.send(9).unwrap(); + assert_eq!(expect_ready(poll_once(selection.as_mut())), Ok(9)); +} + +#[test] +#[should_panic(expected = "select! has no enabled branches")] +fn all_disabled_without_fallback_panics() { + async { + asyncband::select! { + _ = pending::<()>(), if false => {}, + } + } + .block_on(); +} + +#[test] +fn handlers_release_borrows_and_preserve_enclosing_control_flow() { + async fn run() -> Result, ()> { + let local = Rc::new(String::from("local")); + let mut values = Vec::new(); + loop { + asyncband::select! { + biased; + result = async { + pending::<()>().await; + values.push(String::from("unfinished")); + } => result, + (mut value, marker) = async { (local.to_string(), true) } => { + ready(Ok::<_, ()>(())).await?; + value.push('!'); + values.push(value); + assert!(marker); + if values.len() == 1 { + continue; + } + break; + }, + } + } + asyncband::select! { + should_return = ready(true) => { + if should_return { + return Ok(values); + } + }, + } + Err(()) + } + + assert_eq!(run().block_on().unwrap(), ["local!", "local!"]); +} + +#[test] +fn retained_oneshot_is_disabled_after_completion() { + let (sender, receiver) = oneshot::channel(); + sender.send(String::from("done")).unwrap(); + let mut receiver = pin!(receiver.into_future()); + let mut output = None; + async { + loop { + asyncband::select! { + result = receiver.as_mut(), if output.is_none() => output = Some(result.unwrap()), + else => break, + } + } + } + .block_on(); + assert_eq!(output.as_deref(), Some("done")); +} + +#[test] +fn panicking_poll_drops_previously_registered_branches() { + let (sender, receiver) = oneshot::channel::<()>(); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + async { + asyncband::select! { + biased; + _ = receiver => {}, + _ = poll_fn(|_| -> Poll<()> { panic!("poll failed") }) => {}, + } + } + .block_on(); + })); + assert!(result.is_err()); + assert!(sender.send(()).is_err()); +}