#[cfg(any(unix, windows))]
use std::path::Path;
#[cfg(unix)]
use std::process::{Command, ExitStatus, Stdio};
use std::{
future::Future,
io,
net::SocketAddr,
pin::Pin,
task::{Context, Poll},
time::Duration,
};
#[cfg(windows)]
use socket2::{Domain, SockRef};
#[cfg(unix)]
use tokio::io::unix::AsyncFd;
use tokio::runtime::Handle;
#[cfg(unix)]
use super::io::unix_socket_address;
use super::{Interest, IoSource, traits};
#[cfg(target_os = "wasi")]
use crate::log::warn;
#[derive(Clone, Debug)]
pub(crate) struct Tokio {
handle: Handle,
}
impl Tokio {
pub(crate) fn current() -> Option<Self> {
Handle::try_current().ok().map(|handle| Self { handle })
}
}
impl traits::Runtime for Tokio {
type RegisteredIoSource = Registration;
type Sleep = tokio::time::Sleep;
type Task<T>
= TokioTask<T>
where
T: Send + 'static;
#[cfg(unix)]
fn register_io_source(&self, source: IoSource) -> io::Result<Registration> {
let _guard = self.handle.enter();
let async_fd = unsafe { AsyncFd::register(source) }?;
Ok(Registration(async_fd))
}
#[cfg(windows)]
fn register_io_source(&self, source: IoSource) -> io::Result<Registration> {
let socket = SockRef::from(&source);
if socket.local_addr()?.domain() == Domain::UNIX {
return Err(io::Error::from(io::ErrorKind::Unsupported));
}
let stream = std::net::TcpStream::from(socket.try_clone()?);
let _guard = self.handle.enter();
tokio::net::TcpStream::from_std(stream).map(Registration)
}
#[cfg(target_os = "wasi")]
fn register_io_source(&self, source: IoSource) -> io::Result<Registration> {
use std::os::fd::{AsRawFd, FromRawFd};
let stream = unsafe { std::net::TcpStream::from_raw_fd(source.as_raw_fd()) };
let _guard = self.handle.enter();
let stream = match tokio::net::TcpStream::from_std(stream) {
Ok(stream) => stream,
Err(e) => {
std::mem::forget(source);
return Err(e);
}
};
Ok(Registration {
stream: Some(stream),
source: Some(source),
})
}
fn sleep(&self, duration: Duration) -> tokio::time::Sleep {
let _guard = self.handle.enter();
tokio::time::sleep(duration)
}
fn spawn<T>(&self, name: &str, future: impl Future<Output = T> + Send + 'static) -> TokioTask<T>
where
T: Send + 'static,
{
TokioTask::new(spawn_named(&self.handle, name, future))
}
#[cfg(not(all(target_os = "wasi", target_env = "p1")))]
fn connect_tcp(
&self,
address: SocketAddr,
) -> impl Future<Output = io::Result<IoSource>> + Send {
entered(&self.handle, async move {
let stream = tokio::net::TcpStream::connect(address).await?.into_std()?;
Ok(IoSource::new(stream.into()))
})
}
#[cfg(all(target_os = "wasi", target_env = "p1"))]
fn connect_tcp(
&self,
_address: SocketAddr,
) -> impl Future<Output = io::Result<IoSource>> + Send {
std::future::ready(Err(io::Error::from(io::ErrorKind::Unsupported)))
}
#[cfg(unix)]
fn connect_unix(&self, path: &Path) -> impl Future<Output = io::Result<IoSource>> + Send {
let address = unix_socket_address(path).map(tokio::net::unix::SocketAddr::from);
entered(&self.handle, async move {
let address = address?;
let stream = loop {
match tokio::net::UnixStream::connect_addr(&address).await {
Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
tokio::time::sleep(BACKLOG_INTERVAL).await;
}
result => break result?,
}
};
Ok(IoSource::new(stream.into_std()?.into()))
})
}
#[cfg(windows)]
fn connect_unix(&self, _path: &Path) -> impl Future<Output = io::Result<IoSource>> + Send {
std::future::ready(Err(io::Error::from(io::ErrorKind::Unsupported)))
}
#[cfg(all(unix, any(feature = "unixexec", feature = "ibus", target_os = "macos")))]
fn spawn_process(
&self,
mut command: Command,
stdin: Stdio,
stdout: Stdio,
stderr: Stdio,
) -> io::Result<Pin<Box<dyn Future<Output = io::Result<ExitStatus>> + Send + 'static>>> {
command.stdin(stdin).stdout(stdout).stderr(stderr);
let mut command = tokio::process::Command::from(command);
let child = {
let _guard = self.handle.enter();
command.spawn()
};
drop(command);
let mut child = child?;
let handle = self.handle.clone();
Ok(Box::pin(
async move { entered(&handle, child.wait()).await },
))
}
#[cfg(all(
unix,
not(any(feature = "unixexec", feature = "ibus", target_os = "macos"))
))]
fn spawn_process(
&self,
_command: Command,
_stdin: Stdio,
_stdout: Stdio,
_stderr: Stdio,
) -> io::Result<Pin<Box<dyn Future<Output = io::Result<ExitStatus>> + Send + 'static>>> {
Err(io::Error::from(io::ErrorKind::Unsupported))
}
#[cfg(not(target_os = "wasi"))]
fn spawn_blocking<T>(
&self,
work: impl FnOnce() -> T + Send + 'static,
) -> Pin<Box<dyn Future<Output = T> + Send + 'static>>
where
T: Send + 'static,
{
let task = self.handle.spawn_blocking(work);
Box::pin(async move {
match task.await {
Ok(value) => value,
Err(e) if e.is_panic() => std::panic::resume_unwind(e.into_panic()),
Err(_) => {
panic!("the Tokio runtime shut down while zbus was waiting on blocking work")
}
}
})
}
}
#[cfg(unix)]
#[derive(Debug)]
pub(crate) struct Registration(AsyncFd<IoSource>);
#[cfg(unix)]
impl traits::PollIo for Registration {
fn poll_io<T>(
&self,
cx: &mut Context<'_>,
interest: Interest,
mut operation: impl FnMut() -> io::Result<T>,
) -> Poll<io::Result<T>> {
loop {
let mut guard = match interest {
Interest::Readable => std::task::ready!(self.0.poll_read_ready(cx))?,
Interest::Writable => std::task::ready!(self.0.poll_write_ready(cx))?,
};
match operation() {
Err(e) if e.kind() == io::ErrorKind::WouldBlock => guard.clear_ready(),
result => return Poll::Ready(result),
}
}
}
}
#[cfg(windows)]
#[derive(Debug)]
pub(crate) struct Registration(tokio::net::TcpStream);
#[cfg(windows)]
impl traits::PollIo for Registration {
fn poll_io<T>(
&self,
cx: &mut Context<'_>,
interest: Interest,
operation: impl FnMut() -> io::Result<T>,
) -> Poll<io::Result<T>> {
poll_stream_io(&self.0, cx, interest, operation)
}
}
#[cfg(target_os = "wasi")]
#[derive(Debug)]
pub(crate) struct Registration {
stream: Option<tokio::net::TcpStream>,
source: Option<IoSource>,
}
#[cfg(target_os = "wasi")]
impl traits::PollIo for Registration {
fn poll_io<T>(
&self,
cx: &mut Context<'_>,
interest: Interest,
operation: impl FnMut() -> io::Result<T>,
) -> Poll<io::Result<T>> {
let stream = self
.stream
.as_ref()
.expect("a registration's stream is only taken as it is dropped");
poll_stream_io(stream, cx, interest, operation)
}
}
#[cfg(target_os = "wasi")]
impl Drop for Registration {
fn drop(&mut self) {
use std::os::fd::IntoRawFd;
let (Some(stream), Some(source)) = (self.stream.take(), self.source.take()) else {
return;
};
match stream.into_std() {
Ok(stream) => {
let _ = stream.into_raw_fd();
}
Err(e) => {
warn!("Tokio failed to stop watching a socket: {}", e);
std::mem::forget(source);
}
}
}
}
#[cfg(any(windows, target_os = "wasi"))]
fn poll_stream_io<T>(
stream: &tokio::net::TcpStream,
cx: &mut Context<'_>,
interest: Interest,
mut operation: impl FnMut() -> io::Result<T>,
) -> Poll<io::Result<T>> {
let tokio_interest = match interest {
Interest::Readable => tokio::io::Interest::READABLE,
Interest::Writable => tokio::io::Interest::WRITABLE,
};
loop {
std::task::ready!(match interest {
Interest::Readable => stream.poll_read_ready(cx),
Interest::Writable => stream.poll_write_ready(cx),
})?;
match stream.try_io(tokio_interest, &mut operation) {
Err(e) if e.kind() == io::ErrorKind::WouldBlock => continue,
result => return Poll::Ready(result),
}
}
}
#[cfg(not(all(target_os = "wasi", target_env = "p1")))]
async fn entered<F>(handle: &Handle, future: F) -> F::Output
where
F: Future,
{
use std::{future::poll_fn, pin::pin};
let mut future = pin!(future);
poll_fn(|cx| {
let _guard = handle.enter();
future.as_mut().poll(cx)
})
.await
}
#[cfg(unix)]
const BACKLOG_INTERVAL: Duration = Duration::from_millis(20);
fn spawn_named<F>(handle: &Handle, name: &str, future: F) -> tokio::task::JoinHandle<F::Output>
where
F: Future + Send + 'static,
F::Output: Send + 'static,
{
#[cfg(all(tokio_unstable, feature = "tracing"))]
{
tokio::task::Builder::new()
.name(name)
.spawn_on(future, handle)
.expect("`Builder::spawn_on` never returns `Err`")
}
#[cfg(not(all(tokio_unstable, feature = "tracing")))]
{
let _ = name;
handle.spawn(future)
}
}
#[derive(Debug)]
pub(crate) struct TokioTask<T>(Option<tokio::task::JoinHandle<T>>);
impl<T> TokioTask<T> {
pub(super) fn new(handle: tokio::task::JoinHandle<T>) -> Self {
Self(Some(handle))
}
}
impl<T> traits::TaskHandle<T> for TokioTask<T>
where
T: Send + 'static,
{
fn detach(mut self) {
drop(self.0.take());
}
}
impl<T> Drop for TokioTask<T> {
fn drop(&mut self) {
if let Some(handle) = self.0.take() {
handle.abort();
}
}
}
impl<T> Future for TokioTask<T> {
type Output = io::Result<T>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let handle = self
.get_mut()
.0
.as_mut()
.expect("a `TokioTask` is only polled before it is detached or dropped");
Pin::new(handle).poll(cx).map(|r| match r {
Ok(v) => Ok(v),
Err(e) => {
if e.is_cancelled() {
Err(io::Error::other("tokio::task cancelled"))
} else {
panic!("tokio::task::JoinHandle error: {e}")
}
}
})
}
}
#[cfg(test)]
mod tests {
use ntest::timeout;
use super::*;
use crate::runtime::traits::Runtime as _;
#[tokio::test]
#[timeout(15000)]
async fn a_spawned_task_hands_its_output_back() {
let runtime = Tokio::current().expect("a Tokio runtime is current");
let task = runtime.spawn("a task with an output", async { 42 });
assert_eq!(task.await.unwrap(), 42);
}
#[test]
#[timeout(15000)]
fn a_tcp_connect_polled_outside_the_runtime_is_made_on_it() {
use std::net::{Ipv4Addr, TcpListener};
use socket2::SockRef;
let (_tokio, runtime) = runtime_and_backend();
let listener = TcpListener::bind((Ipv4Addr::LOCALHOST, 0)).unwrap();
let source =
futures_lite::future::block_on(runtime.connect_tcp(listener.local_addr().unwrap()))
.unwrap();
let (_accepted, peer) = listener.accept().unwrap();
assert_eq!(
SockRef::from(&source).local_addr().unwrap().as_socket(),
Some(peer),
);
}
#[cfg(unix)]
#[test]
#[timeout(15000)]
fn a_unix_connect_polled_outside_the_runtime_is_made_on_it() {
use std::{io::Write, mem::MaybeUninit, os::unix::net::UnixListener};
use socket2::SockRef;
let (_tokio, runtime) = runtime_and_backend();
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("socket");
let listener = UnixListener::bind(&path).unwrap();
let source = futures_lite::future::block_on(runtime.connect_unix(&path)).unwrap();
let (mut accepted, _) = listener.accept().unwrap();
accepted.write_all(b"!").unwrap();
let mut byte = [MaybeUninit::uninit(); 1];
assert_eq!(SockRef::from(&source).recv(&mut byte).unwrap(), 1);
}
#[cfg(any(target_os = "android", target_os = "linux"))]
#[test]
#[timeout(15000)]
fn a_unix_connect_to_a_full_backlog_polled_outside_the_runtime_waits_on_it() {
use socket2::{Domain, SockAddr, SockRef, Socket, Type};
let (_tokio, runtime) = runtime_and_backend();
let directory = tempfile::tempdir().unwrap();
let path = directory.path().join("socket");
let listener = Socket::new(Domain::UNIX, Type::STREAM, None).unwrap();
listener.bind(&SockAddr::unix(&path).unwrap()).unwrap();
listener.listen(0).unwrap();
let _queued = futures_lite::future::block_on(runtime.connect_unix(&path)).unwrap();
let mut pending = std::pin::pin!(runtime.connect_unix(&path));
assert!(
futures_lite::future::block_on(futures_lite::future::poll_once(pending.as_mut()))
.is_none(),
"the connection was made before the listener had room for it",
);
let _accepted = listener.accept().unwrap();
let source = futures_lite::future::block_on(pending).unwrap();
assert!(SockRef::from(&source).peer_addr().is_ok());
}
#[cfg(all(unix, any(feature = "unixexec", feature = "ibus", target_os = "macos")))]
#[test]
#[timeout(15000)]
fn a_process_spawned_and_waited_for_outside_the_runtime_is_made_on_it() {
use std::process::{Command, Stdio};
let (_tokio, runtime) = runtime_and_backend();
let mut command = Command::new("sh");
command.args(["-c", "sleep 0.2; exit 3"]);
let exit = runtime
.spawn_process(command, Stdio::null(), Stdio::null(), Stdio::null())
.unwrap();
let status = futures_lite::future::block_on(exit).unwrap();
assert_eq!(status.code(), Some(3));
}
fn runtime_and_backend() -> (tokio::runtime::Runtime, Tokio) {
let tokio = tokio::runtime::Runtime::new().unwrap();
let backend =
tokio.block_on(async { Tokio::current().expect("a Tokio runtime is current") });
assert!(
Tokio::current().is_none(),
"the test's thread is inside a Tokio runtime",
);
(tokio, backend)
}
}