Sqlx-core: code for async-global-executor

Author: martin-kolarik

sqlx-core/src/rt/rt_async_global_executor/join_handle.rs

@@ -0,0 +1,30 @@
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+use async_global_executor::Task;
+
+pub struct JoinHandle<T> {
+    pub task: Option<Task<T>>,
+}
+
+impl<T> Drop for JoinHandle<T> {
+    fn drop(&mut self) {
+        if let Some(task) = self.task.take() {
+            task.detach();
+        }
+    }
+}
+
+impl<T> Future for JoinHandle<T> {
+    type Output = T;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        match self.task.as_mut() {
+            Some(task) => Future::poll(Pin::new(task), cx),
+            None => unreachable!("JoinHandle polled after dropping"),
+        }
+    }
+}

sqlx-core/src/rt/rt_async_global_executor/mod.rs

@@ -0,0 +1,10 @@
+mod join_handle;
+pub use join_handle::*;
+
+mod socket;
+
+mod timeout;
+pub use timeout::*;
+
+pub mod yield_now;
+pub use yield_now::*;

sqlx-core/src/rt/rt_async_global_executor/socket.rs

@@ -0,0 +1,55 @@
+use crate::net::Socket;
+
+use std::io;
+use std::io::{Read, Write};
+use std::net::{Shutdown, TcpStream};
+
+use std::task::{Context, Poll};
+
+use crate::io::ReadBuf;
+use async_io_global_executor::Async;
+
+impl Socket for Async<TcpStream> {
+    fn try_read(&mut self, buf: &mut dyn ReadBuf) -> io::Result<usize> {
+        unsafe { self.get_mut().read(buf.init_mut()) }
+    }
+
+    fn try_write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        unsafe { self.get_mut().write(buf) }
+    }
+
+    fn poll_read_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.poll_readable(cx)
+    }
+
+    fn poll_write_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.poll_writable(cx)
+    }
+
+    fn poll_shutdown(&mut self, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        unsafe { Poll::Ready(self.get_mut().shutdown(Shutdown::Both)) }
+    }
+}
+
+#[cfg(unix)]
+impl Socket for Async<std::os::unix::net::UnixStream> {
+    fn try_read(&mut self, buf: &mut dyn ReadBuf) -> io::Result<usize> {
+        unsafe { self.get_mut().read(buf.init_mut()) }
+    }
+
+    fn try_write(&mut self, buf: &[u8]) -> io::Result<usize> {
+        unsafe { self.get_mut().write(buf) }
+    }
+
+    fn poll_read_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.poll_readable(cx)
+    }
+
+    fn poll_write_ready(&mut self, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        self.poll_writable(cx)
+    }
+
+    fn poll_shutdown(&mut self, _cx: &mut Context<'_>) -> Poll<io::Result<()>> {
+        unsafe { Poll::Ready(self.get_mut().shutdown(Shutdown::Both)) }
+    }
+}

sqlx-core/src/rt/rt_async_global_executor/timeout.rs

@@ -0,0 +1,55 @@
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+    time::Duration,
+};
+
+use crate::rt::TimeoutError;
+
+pub async fn sleep(duration: Duration) {
+    timeout_future(duration).await;
+}
+
+pub fn timeout<F: Future>(
+    duration: Duration,
+    future: F,
+) -> impl Future<Output = Result<F::Output, TimeoutError>> {
+    TimeoutFuture::new(future, timeout_future(duration))
+}
+
+fn timeout_future(duration: Duration) -> impl Future {
+    async_io_global_executor::Timer::after(duration)
+}
+
+pub struct TimeoutFuture<F, D> {
+    future: F,
+    delay: D,
+}
+
+impl<F, D> TimeoutFuture<F, D> {
+    fn new(future: F, delay: D) -> TimeoutFuture<F, D> {
+        TimeoutFuture { future, delay }
+    }
+}
+
+impl<F: Future, D: Future> Future for TimeoutFuture<F, D> {
+    type Output = Result<F::Output, TimeoutError>;
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let future_polled = {
+            let future = unsafe { self.as_mut().map_unchecked_mut(|s| &mut s.future) }; // safe, as self is Pin
+            future.poll(cx)
+        };
+        match future_polled {
+            Poll::Ready(v) => Poll::Ready(Ok(v)),
+            Poll::Pending => {
+                let delay = unsafe { self.map_unchecked_mut(|s| &mut s.delay) }; // safe, as self is Pin
+                match delay.poll(cx) {
+                    Poll::Ready(_) => Poll::Ready(Err(TimeoutError)),
+                    Poll::Pending => Poll::Pending,
+                }
+            }
+        }
+    }
+}

sqlx-core/src/rt/rt_async_global_executor/yield_now.rs

@@ -0,0 +1,28 @@
+use std::{
+    future::Future,
+    pin::Pin,
+    task::{Context, Poll},
+};
+
+pub fn yield_now() -> impl Future<Output = ()> {
+    YieldNow(false)
+}
+
+struct YieldNow(bool);
+
+impl Future for YieldNow {
+    type Output = ();
+
+    // The futures executor is implemented as a FIFO queue, so all this future
+    // does is re-schedule the future back to the end of the queue, giving room
+    // for other futures to progress.
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        if !self.0 {
+            self.0 = true;
+            cx.waker().wake_by_ref();
+            Poll::Pending
+        } else {
+            Poll::Ready(())
+        }
+    }
+}