core/test/rust/runtime/async_stream.rs

use dever_runtime::{
    async_stream::{AsyncStream, Error},
    channel::Channel,
    task::{Limits, Runtime},
};
use futures_util::stream;
use std::{
    future::{Future, poll_fn},
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    task::Poll,
};

fn runtime() -> Runtime {
    Runtime::new(Limits {
        threads: 2,
        workers: 1,
        tasks: 8,
    })
    .unwrap()
}

async fn pending(future: std::pin::Pin<&mut impl Future>) {
    let mut future = future;
    poll_fn(|context| {
        assert!(future.as_mut().poll(context).is_pending());
        Poll::Ready(())
    })
    .await;
}

struct Notice(Arc<AtomicUsize>);
impl Drop for Notice {
    fn drop(&mut self) {
        self.0.fetch_add(1, Ordering::SeqCst);
    }
}

#[test]
fn values_are_lazy_and_aliases_share_cursor_and_close() {
    let calls = Arc::new(AtomicUsize::new(0));
    let observed = calls.clone();
    let values = AsyncStream::values((1..=3).inspect(move |_| {
        calls.fetch_add(1, Ordering::SeqCst);
    }));
    assert_eq!(observed.load(Ordering::SeqCst), 0);
    let alias = values.clone();
    runtime()
        .run(async |_| {
            assert_eq!(values.next().await, Ok(Some(1)));
            assert_eq!(alias.next().await, Ok(Some(2)));
            assert_eq!(observed.load(Ordering::SeqCst), 2);
            alias.close();
            assert_eq!(values.next().await, Ok(None));
            assert_eq!(observed.load(Ordering::SeqCst), 2);
        })
        .unwrap();
}

#[test]
fn channel_stream_preserves_null_items_and_external_channel_alias() {
    runtime()
        .run(async |_| {
            let channel = Channel::new(2).unwrap();
            channel.send(None).await.unwrap();
            channel.send(Some(2)).await.unwrap();
            let values = AsyncStream::channel(channel.clone());
            assert_eq!(values.next().await, Ok(Some(None)));
            values.close();
            assert_eq!(values.next().await, Ok(None));
            assert_eq!(channel.receive().await, Ok(Some(Some(2))));
            channel.send(Some(3)).await.unwrap();
            channel.close().await.unwrap();
            let remaining = AsyncStream::channel(channel);
            assert_eq!(remaining.next().await, Ok(Some(Some(3))));
            assert_eq!(remaining.next().await, Ok(None));
        })
        .unwrap();
}

#[test]
fn close_wakes_current_and_queued_pulls() {
    runtime()
        .run(async |_| {
            let values = AsyncStream::channel(Channel::<i64>::new(1).unwrap());
            let mut first = std::pin::pin!(values.next());
            let mut second = std::pin::pin!(values.next());
            pending(first.as_mut()).await;
            pending(second.as_mut()).await;
            values.close();
            assert_eq!(first.await, Ok(None));
            assert_eq!(second.await, Ok(None));
        })
        .unwrap();
}

#[test]
fn cancellation_closes_aliases_and_releases_the_producer() {
    let drops = Arc::new(AtomicUsize::new(0));
    let observed = drops.clone();
    let values = AsyncStream::<i64>::new(stream::unfold(Notice(drops), |owner| async move {
        std::future::pending::<()>().await;
        Some((Ok(1), owner))
    }));
    let alias = values.clone();
    runtime()
        .run(async |scope| {
            let (started, ready) = tokio::sync::oneshot::channel();
            let pending = scope
                .spawn_leaf(move || async move {
                    let _owner = values.close_on_drop();
                    let mut next = std::pin::pin!(values.next());
                    pending(next.as_mut()).await;
                    started.send(()).unwrap();
                    next.await
                })
                .unwrap();
            ready.await.unwrap();
            pending.stop().await.unwrap();
            assert_eq!(observed.load(Ordering::SeqCst), 1);
            assert_eq!(alias.next().await, Ok(None));
        })
        .unwrap();
}

#[test]
fn terminal_fault_releases_source_before_delivery_and_is_not_repeated() {
    let drops = Arc::new(AtomicUsize::new(0));
    let observed = drops.clone();
    let values = AsyncStream::<i64>::new(stream::unfold(Notice(drops), |owner| async move {
        Some((Err(Error::Overflow), owner))
    }));
    runtime()
        .run(async |_| {
            assert_eq!(values.next().await, Err(Error::Overflow));
            assert_eq!(observed.load(Ordering::SeqCst), 1);
            assert_eq!(values.next().await, Ok(None));
        })
        .unwrap();
}

#[test]
fn ticks_are_created_without_runtime_and_closed_without_waiting_for_deadline() {
    assert!(matches!(AsyncStream::ticks(0), Err(Error::Period)));
    assert!(matches!(AsyncStream::ticks(-1), Err(Error::Period)));
    let values = AsyncStream::ticks(60000).unwrap();
    runtime()
        .run(async |_| {
            let mut first = std::pin::pin!(values.next());
            pending(first.as_mut()).await;
            values.close();
            assert_eq!(first.await, Ok(None));
            let values = AsyncStream::ticks(1).unwrap();
            assert_eq!(values.next().await, Ok(Some(1)));
            assert_eq!(values.next().await, Ok(Some(2)));
            values.close();
        })
        .unwrap();
}