join_me_maybe 0.6.0

use futures::stream::FuturesUnordered;
use futures::{Stream, StreamExt, stream};
use join_me_maybe::join;
use std::cell::Cell;
use std::future::ready;
use std::hash::Hash;
use std::pin::Pin;
use std::task::{Context, Poll, Waker};
use tokio::time::{Duration, sleep};
use tokio_stream::StreamMap;

#[tokio::test]
async fn test_maybe() {
    let ret = join!(
        maybe ready(1),
        ready(2),
        ready(3),
        maybe ready(4),
    );
    assert_eq!(ret, (Some(1), 2, 3, None));
}

#[tokio::test]
async fn test_cancel() {
    let ret = join!(
        maybe ready(0),
        ready(1),
        foo: async {
            sleep(Duration::from_secs(1_000_000)).await;
            2 // we'll never get here
        },
        bar: maybe async {
            sleep(Duration::from_secs(1_000_000)).await;
            3 // we'll never get here
        },
        maybe async {
            foo.cancel();
            4
        },
        async {
            bar.cancel();
            5
        },
        baz: maybe ready(6),
    );
    assert_eq!(ret, (Some(0), 1, None, None, Some(4), 5, None));
}

#[tokio::test]
async fn test_early_exit() {
    let ret = join!(
        maybe async {
            foo.cancel();
            // Because of this yield, we'll never get to the return value in this arm.
            sleep(Duration::from_secs(0)).await;
            0
        },
        foo: ready(1),
    );
    assert_eq!(ret, (None, None));
}

#[tokio::test]
async fn test_cancel_already_finished() {
    let ret = join!(
        foo: ready(0),
        async {
            // `foo` will have already finished above by the time we try to cancel it here. This is
            // testing that we don't screw up the count.
            foo.cancel();
            1
        },
        async {
            // Hypothetically if we screwed up the count, we might skip this arm.
            2
        },
    );
    assert_eq!(ret, (Some(0), 1, 2));
}

#[tokio::test]
async fn test_drop_promptly() {
    let mutex1 = tokio::sync::Mutex::new(());
    let mutex2 = tokio::sync::Mutex::new(());
    let mutex3 = tokio::sync::Mutex::new(());
    let ret = join!(
        foo: async {
            // Polling order is (currently) deterministic, so this arm definitely gets the lock
            // here. If that ever changes we could acquire the guard above and move it in here.
            let _guard = mutex1.lock().await;
            // This arm tries to sleep "forever" while holding the lock. The other arm isn't happy
            // about that.
            sleep(Duration::from_secs(1_000_000)).await;
        },
        // These maybe arms should get cancelled automatically when the last arm reaches its body.
        // Test both an unlabeled one and a labeled one.
        maybe async {
            let _guard = mutex2.lock().await;
            sleep(Duration::from_secs(1_000_000)).await;
        },
        bar: maybe async {
            let _guard = mutex3.lock().await;
            sleep(Duration::from_secs(1_000_000)).await;
        },
        _ = async {
            // Take the lock from foo...by force!
            foo.cancel();
            // If cancelling `foo` doesn't drop its future promptly, this will deadlock.
            let _guard = mutex1.lock().await;
        } => {
            // Now all the definitely scrutinees are finished, so the maybe arms should also get
            // cancelled and promptly dropped.
            let _guard = mutex2.lock().await;
            let _guard = mutex3.lock().await;
        }
    );
    assert_eq!(ret, (None, None, None, ()));
}

#[tokio::test]
async fn test_future_items_drop_promptly() {
    let mutex = tokio::sync::Mutex::new(());
    let mut foo_guard_acquired = false;
    let mut foo_body_ran = false;
    join!(
        maybe _ = ready(()) => {
            // This maybe body runs forever, preventing the `foo` body from starting.
            sleep(Duration::from_secs(1_000_000)).await;
        },
        // `_guard` gets acquired immediately here...
        foo: _guard = async {
            let guard = mutex.lock().await;
            foo_guard_acquired = true;
            guard
        } => {
            // ...but we cancel this body before it ever gets a chance to run.
            foo_body_ran = true;
        },
        async {
            foo.cancel();
            // The `foo` body will never run, which means we need to promptly drop the guard
            // sitting in its "item" slot. If we forget to do that, this deadlocks.
            let _guard = mutex.lock().await;
        }
    );
    assert!(foo_guard_acquired);
    assert!(!foo_body_ran);
}

#[tokio::test]
async fn test_stream_items_drop_promptly() {
    let mutex = tokio::sync::Mutex::new(());
    let mut foo_guard_acquired = false;
    let mut foo_body_ran = false;
    join!(
        maybe _ = ready(()) => {
            sleep(Duration::from_secs(1_000_000)).await;
        },
        // Just like `test_future_items_drop_promptly` above, but using a stream here instead.
        foo: _guard in stream::once(async {
            let guard = mutex.lock().await;
            foo_guard_acquired = true;
            guard
        }) => {
            foo_body_ran = true;
        },
        async {
            foo.cancel();
            let _guard = mutex.lock().await;
        }
    );
    assert!(foo_guard_acquired);
    assert!(!foo_body_ran);
}

#[tokio::test]
async fn test_drop_promptly_within_future_body() {
    let mutex = tokio::sync::Mutex::new(());
    let ret = join!(
        foo: async {
            let _guard = mutex.lock().await;
            sleep(Duration::from_secs(1_000_000)).await;
        },
        _ = ready(()) => {
            // Like `test_drop_promptly` above, but from within a future arm body.
            foo.cancel();
            let _guard = mutex.lock().await;
            assert!(foo.with_pin_mut(|foo| foo.is_none()));
            42
        }
    );
    assert_eq!(ret, (None, 42));
}

#[tokio::test]
async fn test_drop_promptly_within_stream_body() {
    let mutex = tokio::sync::Mutex::new(());
    let mut evidence = None;
    join!(
        foo: async {
            let _guard = mutex.lock().await;
            sleep(Duration::from_secs(1_000_000)).await;
        },
        _ in stream::iter([()]) => {
            // Like `test_drop_promptly` above, but from within a stream arm body.
            foo.cancel();
            let _guard = mutex.lock().await;
            assert!(foo.with_pin_mut(|foo| foo.is_none()));
            evidence = Some(42);
        }
    );
    assert_eq!(evidence, Some(42));
}

#[tokio::test]
async fn test_drop_promptly_within_stream_finally() {
    let mutex = tokio::sync::Mutex::new(());
    let ret = join!(
        foo: async {
            let _guard = mutex.lock().await;
            sleep(Duration::from_secs(1_000_000)).await;
        },
        _ in stream::iter(0..3) => {} finally {
            // Like `test_drop_promptly` above, but from within a stream finally.
            foo.cancel();
            let _guard = mutex.lock().await;
            assert!(foo.with_pin_mut(|foo| foo.is_none()));
            42
        }
    );
    assert_eq!(ret, (None, 42));
}

// Most of the cases above rely on simple `ready` futures, but here we do at least one case that
// actually returns `Pending` before eventually returning `ready`.
#[tokio::test]
async fn test_nontrivial_futures() {
    let ret = join!(
        maybe async {
            sleep(Duration::from_millis(1)).await;
            1
        },
        async {
            sleep(Duration::from_millis(10)).await;
            2
        },
    );
    assert_eq!(ret, (Some(1), 2));
}

#[tokio::test]
async fn test_future_arms_with_bodies() {
    let mut counter = 0;
    // Note that all of the arm bodies here can mutate `counter`.
    let ret = join!(
        maybe x = ready(1) => {
            assert_eq!(x, 1);
            counter += 1;
            "hello"
        },
        y = ready(2) => {
            counter += 1;
            10 * y
        },
        _ = sleep(Duration::from_millis(1)) => counter += 1,
        // This arm gets cancelled.
        maybe _ = sleep(Duration::from_millis(10)) => counter += 1,
    );
    assert_eq!(ret, (Some("hello"), 20, (), None));
    assert_eq!(counter, 3);
}

#[tokio::test]
async fn test_stream_arms() {
    let mut elements1 = Vec::new();
    let mut elements2 = Vec::new();
    let mut counter = 0;
    let ret = join!(
        x in stream::iter(0..5) => {
            elements1.push(x);
            counter+= 1;
        },
        x in stream::iter(5..8) => {
            elements2.push(x);
            counter+= 1;
        },
        _ = ready(()) => counter += 100,
    );
    assert_eq!(elements1, [0, 1, 2, 3, 4]);
    assert_eq!(elements2, [5, 6, 7]);
    assert_eq!(counter, 108);
    assert_eq!(ret, ((), (), ()));
}

#[tokio::test]
async fn test_stream_body_sleep_does_not_miss_items() {
    let mut elements = Vec::new();
    join!(
        x in stream::iter(0..3) => {
            sleep(Duration::from_millis(1)).await;
            elements.push(x);
        },
    );
    assert_eq!(elements, [0, 1, 2]);
}

#[tokio::test]
async fn test_stream_body_sleep_does_not_miss_other_stream_items() {
    let mut slow_elements = Vec::new();
    let mut fast_elements = Vec::new();
    join!(
        x in stream::iter(0..3) => {
            sleep(Duration::from_millis(1)).await;
            slow_elements.push(x);
        },
        x in stream::iter(0..3) => {
            fast_elements.push(x);
        },
    );
    assert_eq!(slow_elements, [0, 1, 2]);
    assert_eq!(fast_elements, [0, 1, 2]);
}

#[tokio::test]
async fn test_future_body_sleep_does_not_miss_stream_items() {
    let mut elements = Vec::new();
    join!(
        _ = ready(()) => {
            sleep(Duration::from_millis(1)).await;
        },
        x in stream::iter(0..3) => {
            elements.push(x);
        },
    );
    assert_eq!(elements, [0, 1, 2]);
}

#[tokio::test]
async fn test_finally_sleep_does_not_miss_stream_items() {
    let mut elements = Vec::new();
    join!(
        _ in stream::empty::<()>() => {} finally {
            sleep(Duration::from_millis(1)).await;
        },
        x in stream::iter(0..3) => {
            elements.push(x);
        },
    );
    assert_eq!(elements, [0, 1, 2]);
}

fn resuming<S>(stream: S) -> ResumingStream<S> {
    ResumingStream {
        stream: Box::pin(stream),
        waker: None,
        ended: false,
    }
}

struct ResumingStream<S> {
    stream: Pin<Box<S>>,
    waker: Option<Waker>,
    ended: bool,
}

impl<S> ResumingStream<S> {
    fn inner(self: Pin<&mut Self>) -> Pin<&mut S> {
        if let Some(waker) = &self.waker {
            // Mutating the stream might mean it needs to be polled again.
            waker.wake_by_ref();
        }
        self.get_mut().stream.as_mut()
    }

    fn end(&mut self) {
        self.ended = true;
    }
}

impl<S: Stream> Stream for ResumingStream<S> {
    type Item = S::Item;

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        match self.stream.as_mut().poll_next(cx) {
            // Once `self.ended` is set, we can let the caller observe end-of-stream.
            Poll::Ready(None) if self.ended => Poll::Ready(None),
            // If not `self.ended`, refuse to allow the underlying stream to report that it's done.
            // Of course this causes us to poll the underlying stream again after it *tried* to
            // report that it's done, which isn't generally allowed, but `FuturesUnordered` expects
            // it.
            Poll::Pending | Poll::Ready(None) => {
                // Stash the waker so that we can request a re-poll if the inner stream is mutated.
                self.waker = Some(cx.waker().clone());
                Poll::Pending
            }
            Poll::Ready(Some(item)) => {
                // Wakeups are not registered unless Pending is returned. Clear the waker.
                self.waker = None;
                Poll::Ready(Some(item))
            }
        }
    }
}

#[tokio::test]
async fn test_canceller_mut_futuresunordered() {
    let inputs = futures::stream::iter(0..5).then(|i| async move {
        sleep(Duration::from_millis(1)).await;
        i
    });
    let mut outputs = Vec::new();
    join!(
        i in inputs => {
            unordered.with_pin_mut(|unordered| {
                unordered.unwrap().inner().push(async move { i });
            });
        // `ResumingStream` and `FuturesUnordered` are `Unpin`, so we can test `with_mut` here.
        } finally unordered.with_mut(|unordered| unordered.unwrap().end()),
        unordered: i in resuming(FuturesUnordered::new()) => outputs.push(i),
    );
    outputs.sort();
    assert_eq!(outputs, [0, 1, 2, 3, 4]);
}

// Similar to `ResumingStream` above, but more tailored to `StreamMap` specifically.
struct WellBehavedStreamMap<K, V> {
    map: StreamMap<K, V>,
    waker: Option<Waker>,
    drain: bool,
}

impl<K, V> WellBehavedStreamMap<K, V> {
    fn new() -> Self {
        Self {
            map: StreamMap::new(),
            waker: None,
            drain: false,
        }
    }

    fn start_drain(&mut self) {
        self.drain = true;
    }
}

impl<K: Hash + Eq, V: Stream> WellBehavedStreamMap<K, V> {
    fn insert(&mut self, key: K, stream: V) {
        assert!(!self.drain, "already draining");
        self.map.insert(key, stream);
        if let Some(waker) = &self.waker {
            waker.wake_by_ref();
        }
    }
}

impl<K: Clone + Unpin, V: Stream + Unpin> Stream for WellBehavedStreamMap<K, V> {
    type Item = (K, V::Item);

    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        match Pin::new(&mut self.map).poll_next(cx) {
            Poll::Ready(None) if self.drain => {
                // Once drain is set, we can let the caller observe end-of-stream.
                Poll::Ready(None)
            }
            Poll::Pending | Poll::Ready(None) => {
                self.waker = Some(cx.waker().clone());
                Poll::Pending
            }
            Poll::Ready(Some(item)) => {
                self.waker = None;
                Poll::Ready(Some(item))
            }
        }
    }
}

#[tokio::test]
async fn test_canceller_mut_streammap() {
    let inputs = futures::stream::iter(0..5).then(|i| async move {
        sleep(Duration::from_millis(1)).await;
        i
    });
    let mut outputs = Vec::new();
    join!(
        i in inputs => {
            stream_map.with_pin_mut(|stream_map| {
                stream_map.unwrap().insert(i, futures::stream::iter(vec![i; i]));
            });
        } finally {
            stream_map.with_pin_mut(|stream_map| {
                stream_map.unwrap().start_drain();
            });
        },
        stream_map: (_k, v) in WellBehavedStreamMap::new() => outputs.push(v),
    );
    outputs.sort();
    assert_eq!(outputs, [1, 2, 2, 3, 3, 3, 4, 4, 4, 4]);
}

#[tokio::test]
async fn test_potentially_ambiguous_colons() {
    // All of these examples need to parse cleanly.
    let ret = join!(
        ready(0),
        core::future::ready(1),
        ::core::future::ready(2),
        label3: ready(3),
        label4: core::future::ready(4),
        label5: ::core::future::ready(5),
    );
    assert_eq!(ret, (0, 1, 2, Some(3), Some(4), Some(5)));
}

#[tokio::test]
async fn test_finally_values() {
    // Streams with a finally block return a value.
    let ret = join!(
        maybe _ in futures::stream::iter([()]) => {},
        // This stream will complete before the "definitely" streams below, just because it comes
        // first, but its finally expression will get cancelled.
        maybe _ in futures::stream::iter([()]) => {} finally 99,
        _ in futures::stream::iter([()]) => {},
        _ in futures::stream::iter([()]) => {} finally 42,
        maybe _ in futures::stream::iter([()]) => {} finally 1_000_000,
    );
    assert_eq!(ret, ((), None, (), 42, None));
}

#[tokio::test]
async fn test_await_in_bodies() {
    let mut x = 0;
    let ret = join!(
        y = ready(1) => {
            x += y;
            ready(42).await
        },
        z in stream::iter([10]) => {
            x += z;
            sleep(Duration::from_millis(1)).await;
        } finally {
            x += 100;
            ready(99).await;
            x
        },
    );
    assert_eq!(ret, (42, 111));
}

#[tokio::test]
async fn test_return_in_bodies() {
    fn io_fail() -> std::io::Result<i32> {
        // Normally I'd just try to open /nonexistent.txt or something, but this plays more nicely
        // with Miri.
        Err(std::io::Error::new(std::io::ErrorKind::NotFound, "fail!!!"))
    }
    async fn foo() -> std::io::Result<()> {
        join!(
            sleep(Duration::from_secs(1_000_000)),
            _ = sleep(Duration::from_millis(10)) => {
                io_fail()?;
            }
        );
        Ok(())
    }
    assert!(foo().await.is_err());
}

#[tokio::test]
#[should_panic = "already mutably borrowed"]
async fn test_with_pin_mut_panic() {
    join!(
        foo: sleep(Duration::from_secs(1_000_000)),
        _ = async {} => {
            foo.with_pin_mut(|_| {
                foo.with_pin_mut(|_| {});
            });
        },
    );
}

// This test looks really stupid, but it actually fails to compile ("borrowed value does not live
// long enough") if canceller references don't outlive the futures that capture them. The problem
// is that non-move async blocks will automatically capture `Copy` types by value *except* when
// they're used in a generic context. It's pretty confusing.
#[tokio::test]
async fn test_manipulate_canceller_generically() {
    fn do_something<T>(_: T) {}
    join!(
        foo: sleep(Duration::from_millis(1)),
        async {
            do_something(foo);
        },
    );
}

#[tokio::test]
async fn test_canceller_send_sync_when_future_is_send() {
    fn assert_send_sync<T: Send + Sync>(_: &T) {}
    join!(
        foo: async {
            // `RefCell` makes this future `!Sync`, but its cancellers are always `Send`.
            let _x = std::cell::RefCell::new(());
            sleep(Duration::from_millis(1)).await;
            drop(_x);
        },
        _ = async {
            assert_send_sync::<join_me_maybe::Canceller>(foo);
        } => {
            assert_send_sync::<join_me_maybe::CancellerMut<_>>(foo);
        },
    );
}

#[tokio::test]
async fn test_canceller_send_sync_when_future_is_not_send() {
    fn assert_send_sync<T: Send + Sync>(_: &T) {}
    join!(
        foo: async {
            // `Rc` makes this future `!Send`. The `Canceller` will still be `Send+Sync`, but not
            // the `CancellerMut`. See `tests/ui/non_send_canceller.rs`.
            let _x = std::rc::Rc::new(());
            sleep(Duration::from_millis(1)).await;
            drop(_x);
        },
        async {
            assert_send_sync::<join_me_maybe::Canceller>(foo);
        },
    );
}

// This test covers some behavior, but it also just checks that this compiles and doesn't hit
// "infinite size type" errors.
#[tokio::test]
async fn test_cancel_self() {
    let mut did_write = false;
    join!(
        foo: async {
            foo.cancel();
            // Control still reaches here.
            did_write = true;
            sleep(Duration::from_secs(1_000_000)).await;
        },
    );
    assert!(did_write);
}

// This test covers some behavior, but it also just checks that this compiles and doesn't hit
// "infinite size type" errors.
#[tokio::test]
async fn test_circular_cancellation() {
    let mut did_write1 = false;
    let mut did_write2 = false;
    join!(
        foo: async {
            bar.cancel();
            sleep(Duration::from_millis(1)).await;
            // Control still reaches here.
            did_write1 = true;
        },
        bar: async {
            // Control does not reach here.
            foo.cancel();
            did_write2 = true;
        },
    );
    assert!(did_write1);
    assert!(!did_write2);
}

// This is mainly to check that these cases don't warn about unreachable code.
#[tokio::test]
#[should_panic]
async fn test_panic_in_body() {
    join!(
        _ = ready(2) => todo!(),
        _ in stream::iter([1, 2, 3]) => todo!() finally todo!(),
    );
}

#[tokio::test]
async fn test_type_annotation_needed() {
    join!(
        x = ready(Some(true)) => {
            // TODO: Currently this type annotation is needed. See ui/type_annotation_needed.rs for
            // the error output without it.
            let x: Option<bool> = x;
            x.unwrap();
        }
    );
}

#[tokio::test]
async fn test_maybe_bodies_cancelled() {
    let mut definitely_finished = false;
    let mut definitely_body_finished = false;
    let mut maybe_started = false;
    let mut maybe_finished = false;
    join!(
        _ = async {
            sleep(Duration::from_millis(10)).await;
            definitely_finished = true;
        } => {
            // This body runs after cancelling the maybe body.
            definitely_body_finished = true;
        },
        maybe _ = ready(()) => {
            // This body starts immediately, but because it's a `maybe` arm, it gets cancelled
            // after the sleep finishes.
            maybe_started = true;
            sleep(Duration::from_secs(1_000_000)).await;
            maybe_finished = true;
        },
    );
    assert!(definitely_finished);
    assert!(definitely_body_finished);
    assert!(maybe_started);
    assert!(!maybe_finished);
}

#[tokio::test]
async fn test_cancelled_bodies_cancelled() {
    let mut labeled_started = false;
    let mut labeled_finished = false;
    join!(
        async {
            sleep(Duration::from_millis(10)).await;
            labeled.cancel();
        },
        labeled: _ = ready(()) => {
            // This body starts immediately, but it gets cancelled.
            labeled_started = true;
            sleep(Duration::from_secs(1_000_000)).await;
            labeled_finished = true;
        },
    );
    assert!(labeled_started);
    assert!(!labeled_finished);
}

#[tokio::test]
async fn test_maybe_stream_body_cancelled() {
    let mut body_started = false;
    let mut body_finished = false;
    join!(
        sleep(Duration::from_millis(10)),
        maybe x in stream::iter([1]) => {
            _ = x;
            body_started = true;
            sleep(Duration::from_secs(1_000_000)).await;
            body_finished = true;
        },
    );
    assert!(body_started);
    assert!(!body_finished);
}

#[tokio::test]
async fn test_maybe_finally_cancelled_mid_execution() {
    let mut finally_started = false;
    let mut finally_finished = false;
    join!(
        sleep(Duration::from_millis(10)),
        maybe _ in stream::iter([()]) => {} finally {
            finally_started = true;
            sleep(Duration::from_secs(1_000_000)).await;
            finally_finished = true;
            42
        },
    );
    assert!(finally_started);
    assert!(!finally_finished);
}

#[tokio::test]
async fn test_labeled_maybe_body_cancelled_explicitly() {
    let mut started = false;
    let mut finished = false;
    join!(
        async {
            sleep(Duration::from_millis(10)).await;
            labeled.cancel();
        },
        labeled: maybe _ = ready(()) => {
            started = true;
            sleep(Duration::from_secs(1_000_000)).await;
            finished = true;
        },
    );
    assert!(started);
    assert!(!finished);
}

#[tokio::test]
async fn test_body_self_cancellation() {
    let mut started = false;
    let mut finished = false;
    join!(
        foo: _ = ready(()) => {
            started = true;
            foo.cancel();
            sleep(Duration::from_secs(1_000_000)).await;
            finished = true;
        },
    );
    assert!(started);
    assert!(!finished);
}

#[tokio::test]
async fn test_multiple_maybe_bodies_all_cancelled() {
    let mut started1 = false;
    let mut finished1 = false;
    let mut started2 = false;
    let mut finished2 = false;
    join!(
        sleep(Duration::from_millis(10)),
        maybe _ = ready(()) => {
            started1 = true;
            sleep(Duration::from_secs(1_000_000)).await;
            finished1 = true;
        },
        maybe _ = ready(()) => {
            started2 = true;
            sleep(Duration::from_secs(1_000_000)).await;
            finished2 = true;
        },
    );
    assert!(started1);
    assert!(!started2);
    assert!(!finished1);
    assert!(!finished2);
}

#[tokio::test]
async fn test_stream_body_cancelled_finally_skipped() {
    let mut body_started = false;
    let mut finally_ran = false;
    let ret = join!(
        sleep(Duration::from_millis(10)),
        maybe _ in stream::iter([()]) => {
            body_started = true;
            sleep(Duration::from_secs(1_000_000)).await;
        } finally {
            finally_ran = true;
            99
        },
    );
    assert!(body_started);
    assert!(!finally_ran);
    assert_eq!(ret, ((), None));
}

#[tokio::test]
async fn test_maybe_body_skipped() {
    let mut maybe_body_ran = false;
    let mut definitely_body_ran = false;
    let ret = join!(
        // This maybe arm finishes first, but the definitely arm below finishes immediately after,
        // so the maybe arm's body does not run.
        maybe _ = ready(()) => {
            maybe_body_ran = true;
        },
        _ = ready(()) => {
            definitely_body_ran = true;
        },
    );
    assert!(!maybe_body_ran);
    assert!(definitely_body_ran);
    assert_eq!(ret, (None, ()));
}

#[tokio::test]
async fn test_maybe_body_cancels_last_definitely_arm() {
    let mut maybe_finished = false;
    join!(
        foo: sleep(Duration::from_secs(1_000_000)),
        maybe _ = ready(()) => {
            foo.cancel();
            sleep(Duration::from_millis(10)).await;
            maybe_finished = true;
        },
    );
    assert!(!maybe_finished);
}

#[tokio::test]
async fn test_self_cancel_immediately_ready() {
    let ret = join!(
        foo: _ = ready(()) => {
            foo.cancel();
            42
        },
    );
    assert_eq!(ret, (Some(42),));
}

#[tokio::test]
async fn test_labeled_stream_body_self_cancellation() {
    join!(
        foo: _ in stream::iter([1]) => {
            foo.cancel();
            sleep(Duration::from_secs(1_000_000)).await;
        },
    );
}

#[tokio::test]
async fn test_cancel_from_finally_block() {
    let ret = join!(
        bar: sleep(Duration::from_secs(1_000_000)),
        _ in stream::iter([()]) => {} finally bar.cancel(),
    );
    assert_eq!(ret, (None, ()));
}

#[tokio::test]
async fn test_labeled_definitely_stream_body_cancelled() {
    let mut body_started = false;
    let mut body_finished = false;
    let mut finally_ran = false;
    let ret = join!(
        async {
            sleep(Duration::from_millis(10)).await;
            labeled.cancel();
        },
        labeled: _ in stream::iter([1]) => {
            body_started = true;
            sleep(Duration::from_secs(1_000_000)).await;
            body_finished = true;
        } finally {
            finally_ran = true;
            99
        },
    );
    assert!(body_started);
    assert!(!body_finished);
    assert!(!finally_ran);
    assert_eq!(ret, ((), None));
}

#[tokio::test]
async fn test_dont_poll_streams_concurrently_with_their_own_bodies() {
    let foo_running = Cell::new(false);
    let foo = async || {
        assert!(!foo_running.get(), "only one call to `foo` at a time");
        foo_running.set(true);
        sleep(Duration::from_millis(10)).await;
        foo_running.set(false);
        42
    };
    let foo_stream = stream::iter(0..3).then(|_| foo());
    join!(
        x in foo_stream => {
            // Sanity check: we're actually running `foo` in the stream and not yielding futures.
            assert_eq!(x, 42);
            // The real test: If we drive the stream concurrently with its body, we'll fail the
            // assert in `foo`.
            foo().await;
        },
    );
}