async_rs/implementors/

tokio.rs

//! tokio implementation of async runtime definition traits

use crate::{AsyncIOHandle, Executor, IOHandle, Reactor, Runtime, RuntimeKit, Task, sys::IO};
use async_trait::async_trait;
use cfg_if::cfg_if;
use futures_core::Stream;
use std::{
    future::Future,
    io,
    net::SocketAddr,
    pin::Pin,
    task::{Context, Poll},
    time::{Duration, Instant},
};
use tokio::{net::TcpStream, runtime::Handle};
use tokio_stream::{StreamExt, wrappers::IntervalStream};
use tokio_util::compat::TokioAsyncReadCompatExt;

/// Type alias for the tokio runtime
pub type TokioRuntime = Runtime<Tokio>;

impl TokioRuntime {
    /// Create a new TokioRuntime and bind it to the current tokio runtime by default.
    pub fn tokio() -> Self {
        Self::new(Tokio::current())
    }

    /// Create a new TokioRuntime and bind it to the tokio runtime associated to this handle by default.
    pub fn tokio_with_handle(handle: Handle) -> Self {
        Self::new(Tokio::default().with_handle(handle))
    }
}

/// Dummy object implementing async common interfaces on top of tokio
#[derive(Default, Debug, Clone)]
pub struct Tokio {
    handle: Option<Handle>,
}

impl Tokio {
    /// Bind to the tokio Runtime associated to this handle by default.
    pub fn with_handle(mut self, handle: Handle) -> Self {
        self.handle = Some(handle);
        self
    }

    /// Bind to the current tokio Runtime by default.
    pub fn current() -> Self {
        Self::default().with_handle(Handle::current())
    }

    pub(crate) fn handle(&self) -> Option<Handle> {
        Handle::try_current().ok().or_else(|| self.handle.clone())
    }
}

struct TTask<T: Send>(Option<tokio::task::JoinHandle<T>>);

impl RuntimeKit for Tokio {}

impl Executor for Tokio {
    fn block_on<T, F: Future<Output = T>>(&self, f: F) -> T {
        if let Some(handle) = self.handle() {
            handle.block_on(f)
        } else {
            Handle::current().block_on(f)
        }
    }

    fn spawn<T: Send + 'static>(
        &self,
        f: impl Future<Output = T> + Send + 'static,
    ) -> impl Task<T> {
        TTask(Some(if let Some(handle) = self.handle() {
            handle.spawn(f)
        } else {
            tokio::task::spawn(f)
        }))
    }

    fn spawn_blocking<F: FnOnce() -> T + Send + 'static, T: Send + 'static>(
        &self,
        f: F,
    ) -> impl Task<T> {
        TTask(Some(if let Some(handle) = self.handle() {
            handle.spawn_blocking(f)
        } else {
            tokio::task::spawn_blocking(f)
        }))
    }
}

#[async_trait(?Send)]
impl<T: Send> Task<T> for TTask<T> {
    async fn cancel(&mut self) -> Option<T> {
        let task = self.0.take()?;
        task.abort();
        task.await.ok()
    }
}

impl<T: Send> Future for TTask<T> {
    type Output = T;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let task = self.0.as_mut().expect("task has been canceled");
        match Pin::new(task).poll(cx) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(res) => Poll::Ready(res.expect("task has been canceled")),
        }
    }
}

#[async_trait]
impl Reactor for Tokio {
    fn register<H: IO + Send + 'static>(
        &self,
        socket: IOHandle<H>,
    ) -> io::Result<impl AsyncIOHandle + Send> {
        let _enter = self.handle().as_ref().map(|handle| handle.enter());
        cfg_if! {
            if #[cfg(unix)] {
                Ok(Box::new(unix::AsyncFdWrapper(
                    tokio::io::unix::AsyncFd::new(socket)?,
                )))
            } else {
                Err::<windows::Dummy, _>(io::Error::other(
                    "Registering FD on tokio reactor is only supported on unix",
                ))
            }
        }
    }

    fn sleep(&self, dur: Duration) -> impl Future<Output = ()> {
        tokio::time::sleep(dur)
    }

    fn interval(&self, dur: Duration) -> impl Stream<Item = Instant> {
        let _enter = self.handle().as_ref().map(|handle| handle.enter());
        Box::new(
            IntervalStream::new(tokio::time::interval(dur)).map(tokio::time::Instant::into_std),
        )
    }

    async fn tcp_connect(&self, addr: SocketAddr) -> io::Result<impl AsyncIOHandle + Send> {
        // We cannot do that as EnterGuard is not Send (which makes sense)
        // let _enter = self.handle().as_ref().map(|handle| handle.enter());
        Ok(TcpStream::connect(addr).await?.compat())
    }
}

#[cfg(unix)]
mod unix {
    use super::*;
    use futures_io::{AsyncRead, AsyncWrite};
    use std::io::{IoSlice, IoSliceMut, Read, Write};
    use tokio::io::unix::AsyncFd;

    pub(super) struct AsyncFdWrapper<H: IO + Send + 'static>(pub(super) AsyncFd<IOHandle<H>>);

    impl<H: IO + Send + 'static> AsyncFdWrapper<H> {
        fn read<F: FnOnce(&mut AsyncFd<IOHandle<H>>) -> io::Result<usize>>(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            f: F,
        ) -> Option<Poll<io::Result<usize>>> {
            Some(match self.0.poll_read_ready_mut(cx) {
                Poll::Pending => Poll::Pending,
                Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
                Poll::Ready(Ok(mut guard)) => match guard.try_io(f) {
                    Ok(res) => Poll::Ready(res),
                    Err(_) => return None,
                },
            })
        }

        fn write<R, F: FnOnce(&mut AsyncFd<IOHandle<H>>) -> io::Result<R>>(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            f: F,
        ) -> Option<Poll<io::Result<R>>> {
            Some(match self.0.poll_write_ready_mut(cx) {
                Poll::Pending => Poll::Pending,
                Poll::Ready(Err(e)) => Poll::Ready(Err(e)),
                Poll::Ready(Ok(mut guard)) => match guard.try_io(f) {
                    Ok(res) => Poll::Ready(res),
                    Err(_) => return None,
                },
            })
        }
    }

    impl<H: IO + Send + 'static> Unpin for AsyncFdWrapper<H> {}

    impl<H: IO + Send + 'static> AsyncRead for AsyncFdWrapper<H> {
        fn poll_read(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &mut [u8],
        ) -> Poll<io::Result<usize>> {
            loop {
                if let Some(res) = self.as_mut().read(cx, |socket| socket.get_mut().read(buf)) {
                    return res;
                }
            }
        }

        fn poll_read_vectored(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            bufs: &mut [IoSliceMut<'_>],
        ) -> Poll<io::Result<usize>> {
            loop {
                if let Some(res) = self
                    .as_mut()
                    .read(cx, |socket| socket.get_mut().read_vectored(bufs))
                {
                    return res;
                }
            }
        }
    }

    impl<H: IO + Send + 'static> AsyncWrite for AsyncFdWrapper<H> {
        fn poll_write(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<io::Result<usize>> {
            loop {
                if let Some(res) = self
                    .as_mut()
                    .write(cx, |socket| socket.get_mut().write(buf))
                {
                    return res;
                }
            }
        }

        fn poll_write_vectored(
            mut self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            bufs: &[IoSlice<'_>],
        ) -> Poll<io::Result<usize>> {
            loop {
                if let Some(res) = self
                    .as_mut()
                    .write(cx, |socket| socket.get_mut().write_vectored(bufs))
                {
                    return res;
                }
            }
        }

        fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
            loop {
                if let Some(res) = self.as_mut().write(cx, |socket| socket.get_mut().flush()) {
                    return res;
                }
            }
        }

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

#[cfg(windows)]
mod windows {
    use super::*;
    use futures_io::{AsyncRead, AsyncWrite};

    pub(super) struct Dummy;

    impl AsyncRead for Dummy {
        fn poll_read(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &mut [u8],
        ) -> Poll<io::Result<usize>> {
            Poll::Pending
        }
    }

    impl AsyncWrite for Dummy {
        fn poll_write(
            self: Pin<&mut Self>,
            cx: &mut Context<'_>,
            buf: &[u8],
        ) -> Poll<io::Result<usize>> {
            Poll::Pending
        }

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

        fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
            Poll::Pending
        }
    }
}

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

    #[test]
    fn dyn_compat() {
        struct Test {
            _executor: Box<dyn Executor>,
            _reactor: Box<dyn Reactor>,
            _kit: Box<dyn RuntimeKit>,
            _task: Box<dyn Task<String>>,
        }

        let _ = Test {
            _executor: Box::new(Tokio::default()),
            _reactor: Box::new(Tokio::default()),
            _kit: Box::new(Tokio::default()),
            _task: Box::new(TTask(None)),
        };
    }
}