scoped_stream_sink/

scoped_sink.rs

use alloc::boxed::Box;
use core::cell::{Cell, UnsafeCell};
use core::future::Future;
use core::marker::{PhantomData, PhantomPinned};
#[cfg(feature = "std")]
use core::ops::Deref;
use core::ops::DerefMut;
use core::pin::Pin;
use core::ptr::NonNull;
use core::task::{Context, Poll};

use futures_core::Stream;
use futures_sink::Sink;
use pin_project_lite::pin_project;

#[cfg(feature = "std")]
use crate::LocalThread;

#[cfg(feature = "std")]
/// Erased type for the scope function.
///
/// It accepts a [`SinkInner`] reference and returns a future, capturing it's parameter.
///
/// # Examples
///
/// ```
/// # use futures_util::StreamExt;
/// let func: scoped_stream_sink::DynSinkFn<usize, ()> = Box::new(|mut stream| Box::pin(async move {
///     while let Some(v) = stream.next().await {
///         println!("Value: {v}");
///     }
///     Ok(())
/// }));
/// ```
pub type DynSinkFn<'env, T, E> = Box<
    dyn 'env
        + Send
        + for<'scope> FnMut(Pin<&'scope mut SinkInner<'scope, 'env, T>>) -> DynSinkFuture<'scope, E>,
>;

/// Erased type for the scoped future.
#[cfg(feature = "std")]
pub type DynSinkFuture<'scope, E> = Pin<Box<dyn Future<Output = Result<(), E>> + Send + 'scope>>;

#[cfg(feature = "std")]
pin_project! {
    /// Sink with a scoped future.
    ///
    /// It is useful to easily create [`Sink`] type, without
    /// hassle of manually constructing one.
    /// Safety is guaranteed by the inner reference cannot be moved outside the future,
    /// similiar to [`scope`](std::thread::scope).
    #[must_use = "Sink will not do anything if not used"]
    pub struct ScopedSink<'env, T, E> {
        f: DynSinkFn<'env, T, E>,
        inner: Option<DynSinkFuture<'env, E>>,

        data: Pin<Box<SinkInner<'env, 'env, T>>>,
    }
}

struct SinkInnerData<T> {
    data: UnsafeCell<Option<T>>,
    closed: Cell<bool>,

    // Borrow technique from Tokio to pass pesky Miri :table-flip:
    // <https://github.com/rust-lang/rust/pull/82834>
    _pinned: PhantomPinned,
}

// SAFETY: We don't ever use immutable borrow for any of the operations, so it's automatically Sync too.
// Similar to unstable Exclusive struct.
unsafe impl<T: Send> Send for SinkInnerData<T> {}
unsafe impl<T: Send> Sync for SinkInnerData<T> {}

impl<T> SinkInnerData<T> {
    const fn new() -> Self {
        Self {
            data: UnsafeCell::new(None),
            closed: Cell::new(false),
            _pinned: PhantomPinned,
        }
    }
}

#[cfg(feature = "std")]
pin_project! {
    /// Inner type for [`ScopedSink`].
    ///
    /// `'scope` defines the lifetime of it's scope,
    /// and `'env` defines the lifetime of it's environment. Lifetimes are constrained
    /// such that the reference cannot be sent outside it's scope.
    ///
    /// # Note About Thread-safety
    ///
    /// Even though [`SinkInner`] is both [`Send`] and [`Sink`], it's reference
    /// **should** not be sent across thread. This is currently impossible, due to
    /// lack of async version of [`scope`](std::thread::scope).
    /// To future-proof that possibility, any usage of it will panic if called from different
    /// thread than the outer thread. It also may panics outer thread too.
    ///
    /// Also do note that some of the check depends on `debug_assertions` build config
    /// (AKA only on debug builds).
    #[must_use = "SinkInner will not do anything if not used"]
    pub struct SinkInner<'scope, 'env: 'scope, T> {
        #[pin]
        inner: LocalThread<SinkInnerData<T>>,

        phantom: PhantomData<&'scope mut &'env T>,
    }
}

#[cfg(feature = "std")]
impl<'env, T: 'env, E: 'env> ScopedSink<'env, T, E> {
    /// Create new [`ScopedSink`] from a [`DynSinkFn`].
    ///
    /// You should probably use [`new`](Self::new) instead.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anyhow::Error;
    /// # use scoped_stream_sink::ScopedSink;
    /// let mut sink: ScopedSink<usize, Error> = ScopedSink::new_dyn(Box::new(|_| {
    ///     Box::pin(async { Ok(()) })
    /// }));
    /// ```
    pub fn new_dyn(f: DynSinkFn<'env, T, E>) -> Self {
        Self {
            data: Box::pin(SinkInner {
                inner: LocalThread::new(SinkInnerData::new()),

                phantom: PhantomData,
            }),

            f,
            inner: None,
        }
    }

    /// Create new [`ScopedSink`].
    ///
    /// Future should be consuming all items of the stream.
    /// If not, the function will be called again to restart itself.
    /// It is guaranteed if the sink is closed, the future will never restarts.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anyhow::Error;
    /// # use futures_util::{SinkExt, StreamExt};
    /// # use scoped_stream_sink::ScopedSink;
    /// # fn main() -> Result<(), Error> {
    /// # tokio::runtime::Builder::new_current_thread().enable_all().build()?.block_on(async {
    /// let mut sink = <ScopedSink<_, Error>>::new(|mut stream| Box::pin(async move {
    ///     // Reads a value. If future returns before sink is closed, it will be restarted.
    ///     if let Some(v) = stream.next().await {
    ///         println!("Value: {v}");
    ///     }
    ///     Ok(())
    /// }));
    ///
    /// // Send a value
    /// sink.send(1).await?;
    ///
    /// // Close the sink
    /// sink.close().await?;
    /// # Ok(()) })}
    /// ```
    pub fn new<F>(f: F) -> Self
    where
        for<'scope> F: 'env
            + Send
            + FnMut(
                Pin<&'scope mut SinkInner<'scope, 'env, T>>,
            ) -> Pin<Box<dyn Future<Output = Result<(), E>> + Send + 'scope>>,
    {
        Self::new_dyn(Box::new(f))
    }
}

impl<T> SinkInnerData<T> {
    fn flush<E, U, F>(
        &self,
        cx: &mut Context<'_>,
        fut: &mut Option<Pin<U>>,
        mut f: F,
    ) -> Poll<Result<(), E>>
    where
        U: DerefMut,
        U::Target: Future<Output = Result<(), E>>,
        F: FnMut() -> Pin<U>,
    {
        loop {
            if unsafe { (*self.data.get()).is_none() } {
                // No need to poll future.
                return Poll::Ready(Ok(()));
            }

            let fp = if let Some(v) = fut {
                v
            } else if self.closed.get() {
                return Poll::Ready(Ok(()));
            } else {
                fut.get_or_insert_with(&mut f)
            };

            let Poll::Ready(v) = fp.as_mut().poll(cx) else {
                break;
            };

            // Dispose future.
            *fut = None;

            if v.is_err() {
                return Poll::Ready(v);
            }

            // We have to repoll the future, otherwise it will never be awoken.
        }

        if unsafe { (*self.data.get()).is_some() } {
            Poll::Pending
        } else {
            Poll::Ready(Ok(()))
        }
    }

    fn send(&self, item: T) {
        if self.closed.get() {
            panic!("Sink is closed!");
        }
        let data = unsafe { &mut *self.data.get() };
        if data.is_some() {
            panic!("poll_ready() is not called yet!");
        }

        *data = Some(item);
    }

    fn close<E, U, F>(
        &self,
        cx: &mut Context<'_>,
        fut: &mut Option<Pin<U>>,
        f: F,
    ) -> Poll<Result<(), E>>
    where
        U: DerefMut,
        U::Target: Future<Output = Result<(), E>>,
        F: FnMut() -> Pin<U>,
    {
        self.closed.set(true);

        // There is still some data
        if unsafe { (*self.data.get()).is_some() } {
            let ret = self.flush(cx, &mut *fut, f);
            if let Poll::Ready(Err(_)) = ret {
                return ret;
            }
            return match fut {
                // Must have been pending then.
                Some(_) => Poll::Pending,
                None => ret,
            };
        }

        let ret = match fut {
            Some(p) => p.as_mut().poll(cx),
            None => return Poll::Ready(Ok(())),
        };
        if ret.is_ready() {
            *fut = None;
        }
        ret
    }

    fn next(&self) -> Poll<Option<T>> {
        match unsafe { (*self.data.get()).take() } {
            v @ Some(_) => Poll::Ready(v),
            None if self.closed.get() => Poll::Ready(None),
            None => Poll::Pending,
        }
    }
}

unsafe fn make_future<'a, T: 'a, R, F>(mut ptr: NonNull<T>, mut f: F) -> impl FnMut() -> R
where
    F: FnMut(Pin<&'a mut T>) -> R,
{
    move || f(Pin::new_unchecked(ptr.as_mut()))
}

#[cfg(feature = "std")]
impl<'env, T: 'env, E: 'env> ScopedSink<'env, T, E> {
    fn future_wrapper(
        self: Pin<&mut Self>,
    ) -> (
        impl Deref<Target = SinkInnerData<T>> + '_,
        &mut Option<DynSinkFuture<'env, E>>,
        impl FnMut() -> DynSinkFuture<'env, E> + '_,
    ) {
        let this = self.project();
        // SAFETY: We constrained data lifetime to be 'scope.
        // Since 'scope is contained within self, it is safe to extend it.
        let f = unsafe {
            make_future(
                NonNull::from(this.data.as_mut().get_unchecked_mut()),
                this.f,
            )
        };

        (
            this.data.as_ref().get_ref().inner.set_inner_ctx(),
            this.inner,
            f,
        )
    }
}

#[cfg(feature = "std")]
impl<'env, T: 'env, E: 'env> Sink<T> for ScopedSink<'env, T, E> {
    type Error = E;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), E>> {
        self.poll_flush(cx)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), E>> {
        let (data, fut, f) = self.future_wrapper();

        data.flush(cx, fut, f)
    }

    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), E> {
        self.data.inner.set_inner_ctx().send(item);
        Ok(())
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), E>> {
        let (data, fut, f) = self.future_wrapper();

        data.close(cx, fut, f)
    }
}

#[cfg(feature = "std")]
impl<'scope, 'env, T> Stream for SinkInner<'scope, 'env, T> {
    type Item = T;

    fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.get_inner().next()
    }
}

/// Erased type for the local scope function.
///
/// It accepts a [`LocalSinkInner`] reference and returns a future, capturing it's parameter.
///
/// # Examples
///
/// ```
/// # use futures_util::StreamExt;
/// let func: scoped_stream_sink::DynLocalSinkFn<usize, ()> = Box::new(|mut stream| Box::pin(async move {
///     while let Some(v) = stream.next().await {
///         println!("Value: {v}");
///     }
///     Ok(())
/// }));
/// ```
pub type DynLocalSinkFn<'env, T, E> = Box<
    dyn 'env
        + for<'scope> FnMut(
            Pin<&'scope mut LocalSinkInner<'scope, 'env, T>>,
        ) -> DynLocalSinkFuture<'scope, E>,
>;

/// Erased type for the locally scoped future.
pub type DynLocalSinkFuture<'scope, E> = Pin<Box<dyn Future<Output = Result<(), E>> + 'scope>>;

pin_project! {
    /// Local sink with a scoped future.
    ///
    /// Unlike [`ScopedSink`] it is not [`Send`], so it can work in no-std environment.
    #[must_use = "Sink will not do anything if not used"]
    pub struct LocalScopedSink<'env, T, E> {
        f: DynLocalSinkFn<'env, T, E>,
        inner: Option<DynLocalSinkFuture<'env, E>>,

        data: Pin<Box<LocalSinkInner<'env, 'env, T>>>,
    }
}

pin_project! {
    /// Inner type for [`LocalScopedSink`].
    ///
    /// Similiar to [`SinkInner`], but not [`Send`].
    pub struct LocalSinkInner<'scope, 'env: 'scope, T> {
        #[pin]
        inner: SinkInnerData<T>,

        phantom: PhantomData<(&'scope mut &'env T, *mut u8)>,
    }
}

impl<'env, T: 'env, E: 'env> LocalScopedSink<'env, T, E> {
    /// Create new [`LocalScopedSink`] from a [`DynLocalSinkFn`].
    ///
    /// You should probably use [`new`](Self::new) instead.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anyhow::Error;
    /// # use scoped_stream_sink::LocalScopedSink;
    /// let mut sink: LocalScopedSink<usize, Error> = LocalScopedSink::new_dyn(Box::new(|_| {
    ///     Box::pin(async { Ok(()) })
    /// }));
    /// ```
    pub fn new_dyn(f: DynLocalSinkFn<'env, T, E>) -> Self {
        Self {
            data: Box::pin(LocalSinkInner {
                inner: SinkInnerData::new(),

                phantom: PhantomData,
            }),

            f,
            inner: None,
        }
    }

    /// Create new [`LocalScopedSink`].
    ///
    /// Future should be consuming all items of the stream.
    /// If not, the function will be called again to restart itself.
    /// It is guaranteed if the sink is closed, the future will never restarts.
    ///
    /// # Examples
    ///
    /// ```
    /// # use anyhow::Error;
    /// # use futures_util::{SinkExt, StreamExt};
    /// # use scoped_stream_sink::LocalScopedSink;
    /// # fn main() -> Result<(), Error> {
    /// # tokio::runtime::Builder::new_current_thread().enable_all().build()?.block_on(async {
    /// let mut sink = <LocalScopedSink<_, Error>>::new(|mut stream| Box::pin(async move {
    ///     // Reads a value. If future returns before sink is closed, it will be restarted.
    ///     if let Some(v) = stream.next().await {
    ///         println!("Value: {v}");
    ///     }
    ///     Ok(())
    /// }));
    ///
    /// // Send a value
    /// sink.send(1).await?;
    ///
    /// // Close the sink
    /// sink.close().await?;
    /// # Ok(()) })}
    /// ```
    pub fn new<F>(f: F) -> Self
    where
        for<'scope> F: 'env
            + FnMut(Pin<&'scope mut LocalSinkInner<'scope, 'env, T>>) -> DynLocalSinkFuture<'scope, E>,
    {
        Self::new_dyn(Box::new(f))
    }
}

impl<'env, T: 'env, E: 'env> LocalScopedSink<'env, T, E> {
    fn future_wrapper(
        self: Pin<&mut Self>,
    ) -> (
        &SinkInnerData<T>,
        &mut Option<DynLocalSinkFuture<'env, E>>,
        impl FnMut() -> DynLocalSinkFuture<'env, E> + '_,
    ) {
        let this = self.project();
        // SAFETY: We constrained data lifetime to be 'scope.
        // Since 'scope is contained within self, it is safe to extend it.
        let f = unsafe {
            make_future(
                NonNull::from(this.data.as_mut().get_unchecked_mut()),
                this.f,
            )
        };

        (&this.data.inner, this.inner, f)
    }
}

impl<'env, T: 'env, E: 'env> Sink<T> for LocalScopedSink<'env, T, E> {
    type Error = E;

    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), E>> {
        self.poll_flush(cx)
    }

    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), E>> {
        let (data, fut, f) = self.future_wrapper();

        data.flush(cx, fut, f)
    }

    fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), E> {
        self.data.inner.send(item);
        Ok(())
    }

    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), E>> {
        let (data, fut, f) = self.future_wrapper();

        data.close(cx, fut, f)
    }
}

impl<'scope, 'env, T> Stream for LocalSinkInner<'scope, 'env, T> {
    type Item = T;

    fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        self.inner.next()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    use std::prelude::rust_2021::*;
    use std::sync::atomic::{AtomicU8, Ordering};
    use std::sync::Arc;
    use std::time::Duration;

    use anyhow::{bail, Error as AnyError, Result as AnyResult};
    use futures_util::{SinkExt, StreamExt};
    use num_integer::Roots as _;
    use tokio::spawn;
    use tokio::sync::mpsc::channel;
    use tokio::task::yield_now;
    use tokio::time::timeout;

    async fn test_helper<F>(f: F) -> AnyResult<()>
    where
        F: Future<Output = AnyResult<()>> + Send,
    {
        match timeout(Duration::from_secs(5), f).await {
            Ok(v) => v,
            Err(_) => bail!("Time ran out"),
        }
    }
    /*
    #[tokio::test]
    async fn test_simple_fail() -> AnyResult<()> {
        use std::pin::pin;

        let mut sink: Pin<&mut ScopedSink<'static, usize, AnyError>> =
            pin!(ScopedSink::new(|mut src| Box::pin(async move {
                tokio::spawn(async move {
                    println!("{:?}", src.next().await);
                    println!("{:?}", src.next().await);
                })
                .await?;
                Ok(())
            })));

        test_helper(async move {
            sink.send(1).await?;
            drop(sink);

            Ok(())
        })
        .await
    }
    */
    #[tokio::test]
    async fn test_simple() -> AnyResult<()> {
        let mut sink: ScopedSink<usize, AnyError> = ScopedSink::new(|_| Box::pin(async { Ok(()) }));

        test_helper(async move {
            println!("Closing");
            sink.close().await?;

            Ok(())
        })
        .await
    }

    #[tokio::test]
    async fn test_send_one() -> AnyResult<()> {
        let mut sink: ScopedSink<usize, AnyError> = ScopedSink::new(|mut src| {
            Box::pin(async move {
                println!("Starting sink");
                while let Some(v) = src.next().await {
                    println!("Value: {v}");
                }
                println!("Stopping sink");

                Ok(())
            })
        });

        test_helper(async move {
            sink.feed(1).await?;

            println!("Closing");
            sink.close().await?;

            Ok(())
        })
        .await
    }

    #[tokio::test]
    async fn test_send_many() -> AnyResult<()> {
        let mut sink: ScopedSink<usize, AnyError> = ScopedSink::new(|mut src| {
            Box::pin(async move {
                println!("Starting sink");
                while let Some(v) = src.next().await {
                    println!("Value: {v}");
                }
                println!("Stopping sink");

                Ok(())
            })
        });

        test_helper(async move {
            for i in 0..10 {
                println!("Sending: {i}");
                sink.feed(i).await?;
            }

            println!("Closing");
            sink.close().await?;

            Ok(())
        })
        .await
    }

    #[tokio::test]
    async fn test_send_yield() -> AnyResult<()> {
        let mut sink: ScopedSink<usize, AnyError> = ScopedSink::new(|mut src| {
            Box::pin(async move {
                println!("Starting sink");
                while let Some(v) = src.next().await {
                    println!("Value: {v}");
                    for _ in 0..5 {
                        yield_now().await;
                    }
                }
                println!("Stopping sink");

                Ok(())
            })
        });

        test_helper(async move {
            for i in 0..10 {
                println!("Sending: {i}");
                sink.feed(i).await?;
            }

            println!("Closing");
            sink.close().await?;

            Ok(())
        })
        .await
    }

    #[tokio::test]
    async fn test_send_yield2() -> AnyResult<()> {
        let mut sink: ScopedSink<usize, AnyError> = ScopedSink::new(|mut src| {
            Box::pin(async move {
                println!("Starting sink");
                while let Some(v) = src.next().await {
                    println!("Value: {v}");
                    for _ in 0..3 {
                        yield_now().await;
                    }
                }
                println!("Stopping sink");

                Ok(())
            })
        });

        test_helper(async move {
            for i in 0..10 {
                println!("Sending: {i}");
                sink.feed(i).await?;

                for _ in 0..5 {
                    yield_now().await;
                }
            }

            println!("Closing");
            sink.close().await?;

            Ok(())
        })
        .await
    }

    #[tokio::test]
    async fn test_send_many_flush() -> AnyResult<()> {
        let mut sink: ScopedSink<usize, AnyError> = ScopedSink::new(|mut src| {
            Box::pin(async move {
                println!("Starting sink");
                while let Some(v) = src.next().await {
                    println!("Value: {v}");
                }
                println!("Stopping sink");

                Ok(())
            })
        });

        test_helper(async move {
            for i in 0..10 {
                println!("Sending: {i}");
                sink.feed(i).await?;
            }

            println!("Flushing");
            sink.flush().await?;

            for i in 10..20 {
                println!("Sending: {i}");
                sink.feed(i).await?;
            }

            println!("Closing");
            sink.close().await?;

            Ok(())
        })
        .await
    }

    #[tokio::test]
    async fn test_return_then_receive() -> AnyResult<()> {
        let v = Arc::new(AtomicU8::new(0));
        let mut sink: ScopedSink<usize, AnyError> = ScopedSink::new(move |mut src| {
            let v = v.clone();
            Box::pin(async move {
                let mut v_ = v.load(Ordering::SeqCst);
                v_ = if v_ == 8 {
                    // Should never receive empty as any closure will be returned early.
                    assert_eq!(src.next().await, Some(1));
                    0
                } else {
                    v_ + 1
                };
                v.store(v_, Ordering::SeqCst);

                Ok(())
            })
        });

        test_helper(async move {
            for _ in 0..10 {
                println!("Sending");
                sink.feed(1).await?;
            }

            println!("Closing");
            sink.close().await?;

            Ok(())
        })
        .await
    }

    #[tokio::test]
    async fn test_double_scoped() -> AnyResult<()> {
        let mut sink: ScopedSink<usize, AnyError> = ScopedSink::new(|mut src| {
            Box::pin(async move {
                // Unfortunately, you can't use src inside inner sink yet
                // (due to issues with lifetimes of captured values.
                let mut sink2: ScopedSink<usize, AnyError> = ScopedSink::new(|mut src| {
                    Box::pin(async move {
                        println!("Value: {}", src.next().await.unwrap());
                        Ok(())
                    })
                });

                while let Some(v) = src.next().await {
                    sink2.feed(v).await?;
                    sink2.feed(v).await?;
                }
                sink2.close().await?;

                Ok(())
            })
        });

        test_helper(async move {
            for i in 0..10 {
                sink.feed(i).await?;
            }
            sink.close().await?;

            Ok(())
        })
        .await
    }

    #[tokio::test]
    async fn test_spawn_mpsc() -> AnyResult<()> {
        fn is_prime(v: u64) -> bool {
            (2..v.sqrt()).all(|i| v % i != 0)
        }

        let (s1, mut r1) = channel::<u64>(4);
        let (s2, mut r2) = channel::<u64>(4);

        let mut sink = ScopedSink::new(move |mut stream| {
            let s1 = s1.clone();
            let s2 = s2.clone();

            Box::pin(async move {
                let Some(v) = stream.next().await else {
                    return Ok(());
                };

                if is_prime(v) {
                    s1.send(v).await
                } else {
                    s2.send(v).await
                }
            })
        });

        let mut handles = Vec::new();

        handles.push(spawn(test_helper(async move {
            while let Some(v) = r1.recv().await {
                assert!(is_prime(v));
                for _ in 0..v.sqrt() {
                    yield_now().await
                }
            }

            Ok(())
        })));

        handles.push(spawn(test_helper(async move {
            while let Some(v) = r2.recv().await {
                assert!(!is_prime(v));
                for _ in 0..v.sqrt() {
                    yield_now().await
                }
            }

            Ok(())
        })));

        handles.push(spawn(test_helper(async move {
            for i in 1..1000 {
                sink.feed(i).await?;
            }
            sink.close().await?;

            Ok(())
        })));

        let mut has_error = false;
        for f in handles {
            if let Err(e) = f.await? {
                eprintln!("{e:?}");
                has_error = true;
            }
        }

        if has_error {
            bail!("Some error has happened");
        }

        Ok(())
    }
}