use crate::{
queue::{Queue, QueueEntry, Queueable},
task::{self, Builder, Task},
};
use nix::{
fcntl::OFlag,
sys::epoll::{Epoll, EpollCreateFlags, EpollEvent, EpollFlags, EpollTimeout},
unistd::{pipe2, read, write},
};
use pin_project_lite::pin_project;
use std::{
cell::{Cell, RefCell, UnsafeCell},
io::Error,
marker::PhantomData,
os::fd::{AsFd, AsRawFd, BorrowedFd, OwnedFd, RawFd},
pin::Pin,
sync::{Arc, Weak},
task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
};
thread_local! {
pub(crate) static EXECUTOR: RefCell<Pin<Box<Executor>>> = RefCell::new(Executor::new().unwrap());
}
pub(crate) fn exec<R, F: FnOnce(Pin<&mut Executor>) -> R>(f: F) -> R {
EXECUTOR.with_borrow_mut(|e| f(e.as_mut()))
}
pub(crate) struct ExecutorTask {
queue_entry: UnsafeCell<QueueEntry>,
pipe_wr: RawFd,
waker: Waker,
pub(crate) handle_waker: Cell<Waker>,
priority: Priority,
pub(crate) task: Pin<Box<RefCell<Task>>>,
has_completed: Cell<bool>,
name: std::borrow::Cow<'static, str>,
#[cfg(feature = "tracing")]
tracing_span: Option<tracing::Id>,
_pinned: std::marker::PhantomPinned,
}
pub(crate) type TaskRef = Pin<Arc<ExecutorTask>>;
impl Queueable for ExecutorTask {
fn with_entry<R, F: FnMut(Pin<&mut QueueEntry>) -> R>(self: &Pin<Arc<Self>>, mut f: F) -> R {
f(unsafe { Pin::new_unchecked(&mut *self.queue_entry.get()) })
}
fn from_entry(entry: Pin<&QueueEntry>) -> Pin<Arc<Self>> {
let cell_inner_offset = {
let cell = UnsafeCell::new(QueueEntry::new());
let inner_p = cell.get().cast::<u8>();
let cell_p = std::ptr::from_ref(&cell).cast::<u8>();
unsafe { inner_p.offset_from_unsigned(cell_p) }
};
let queue_entry_offset = std::mem::offset_of!(Self, queue_entry) + cell_inner_offset;
let task = unsafe {
std::ptr::from_ref(&*entry.as_ref())
.cast::<u8>()
.sub(queue_entry_offset)
.cast()
};
unsafe { Pin::new_unchecked(Arc::from_raw(task)) }
}
}
impl ExecutorTask {
fn is_queued(self: &Pin<Arc<Self>>) -> bool {
self.with_entry(|e| e.is_queued())
}
pub(crate) fn name(&self) -> &str {
&self.name
}
}
fn create_epoll() -> Result<Epoll, Error> {
Ok(Epoll::new(EpollCreateFlags::EPOLL_CLOEXEC)?)
}
pin_project! {
pub(crate) struct Executor {
epoll: Epoll,
#[pin]
runq_high: Queue<ExecutorTask>,
#[pin]
runq_normal: Queue<ExecutorTask>,
#[pin]
runq_low: Queue<ExecutorTask>,
#[pin]
sleepq: Queue<ExecutorTask>,
pipe_rd: OwnedFd,
pipe_wr: OwnedFd,
events: Vec<EpollEvent>,
shutdown: bool,
_not_send_not_sync: PhantomData<*mut ()>,
}
}
impl Executor {
fn new() -> Result<Pin<Box<Self>>, Error> {
let (pipe_rd, pipe_wr) = pipe2(OFlag::O_CLOEXEC | OFlag::O_NONBLOCK)?;
let mut s = Box::pin(Self {
epoll: create_epoll()?,
runq_high: unsafe { Queue::new() },
runq_normal: unsafe { Queue::new() },
runq_low: unsafe { Queue::new() },
sleepq: unsafe { Queue::new() },
pipe_rd,
pipe_wr,
events: Vec::new(),
shutdown: false,
_not_send_not_sync: PhantomData,
});
let this = s.as_mut().project();
unsafe {
this.runq_high.init();
this.runq_normal.init();
this.runq_low.init();
this.sleepq.init();
}
this.events.push(EpollEvent::empty());
this.epoll
.add(this.pipe_rd, EpollEvent::new(EpollFlags::EPOLLIN, 0))
.unwrap();
Ok(s)
}
fn spawn<F: Future + 'static>(
self: Pin<&mut Self>,
future: F,
priority: Priority,
metadata: task::Metadata,
) -> crate::task::Handle<F::Output> {
let task = Arc::new_cyclic(|me: &Weak<ExecutorTask>| {
#[cfg(feature = "tracing")]
let (tracing_span, future) =
task::instrument_future(future, &metadata, me.as_ptr().addr());
ExecutorTask {
queue_entry: UnsafeCell::new(QueueEntry::new()),
pipe_wr: self.pipe_wr.as_raw_fd(),
waker: task_waker::build(me),
handle_waker: Cell::new(Waker::noop().clone()),
priority,
task: Task::new_pinned_boxed_refcell(future),
has_completed: Cell::new(false),
name: metadata.name,
#[cfg(feature = "tracing")]
tracing_span,
_pinned: std::marker::PhantomPinned,
}
});
let task = unsafe { Pin::new_unchecked(task) };
self.enqueue(task.clone());
crate::task::Handle::new(task)
}
pub(crate) fn abort(self: Pin<&mut Self>, task: TaskRef) {
let this = self.project();
this.sleepq.release(task);
}
pub(crate) fn epoll_add(
self: Pin<&mut Self>,
fd: impl AsFd,
events: EpollFlags,
ew: Pin<&RefCell<EpollWaker>>,
) -> Result<(), Error> {
let this = self.project();
this.events.push(EpollEvent::empty());
this.epoll.add(
fd,
EpollEvent::new(events, std::ptr::from_ref(Pin::into_inner(ew)) as _),
)?;
Ok(())
}
pub(crate) fn epoll_del(self: Pin<&mut Self>, fd: impl AsFd) -> Result<(), Error> {
let this = self.project();
this.events.pop();
this.epoll.delete(fd)?;
Ok(())
}
fn runq(self: Pin<&mut Self>, priority: Priority) -> Pin<&mut Queue<ExecutorTask>> {
let this = self.project();
match priority {
Priority::High => this.runq_high,
Priority::Normal => this.runq_normal,
Priority::Low => this.runq_low,
}
}
fn epoll_wait(mut self: Pin<&mut Self>, timeout: EpollTimeout) -> Result<(), Error> {
if self.events.is_empty() {
return Ok(());
}
let n_events = {
let this = self.as_mut().project();
this.epoll
.wait(this.events.as_mut_slice(), timeout)
.or_else(|e| {
if e == nix::Error::EINTR {
Ok(0)
} else {
Err(e)
}
})
}?;
for ei in 0..n_events {
let e = self.events[ei];
if e.data() == 0 {
self.as_mut().wake_from_pipe()?;
continue;
}
let ew = Pin::new(unsafe { &*(e.data() as *mut RefCell<EpollWaker>) });
let mut waker = ew.borrow_mut();
if waker.events.contains(e.events()) {
continue;
}
waker.events |= e.events();
if *waker.waker.vtable() == task_waker::VTABLE {
let weak = weak_from_raw(waker.waker.data());
if let Some(task) = weak.upgrade() {
let task = unsafe { Pin::new_unchecked(task) };
self.as_mut().enqueue(task);
}
let _ = Weak::into_raw(weak);
}
}
Ok(())
}
fn enqueue(self: Pin<&mut Self>, task: TaskRef) {
if !task.has_completed.get() {
self.runq(task.priority).push(task);
}
}
fn sleep(self: Pin<&mut Self>, task: TaskRef) {
let this = self.project();
this.sleepq.push(task);
}
fn highest_priority_runnable_task(self: Pin<&mut Self>) -> Option<TaskRef> {
let this = self.project();
this.runq_high
.pop()
.or_else(|| this.runq_normal.pop())
.or_else(|| this.runq_low.pop())
}
fn shutdown(self: Pin<&mut Self>) {
*self.project().shutdown = true;
}
fn wake_from_pipe(mut self: Pin<&mut Self>) -> Result<(), Error> {
let mut buf = std::mem::MaybeUninit::<*const ()>::uninit();
read(&self.pipe_rd, unsafe {
std::slice::from_raw_parts_mut(
buf.as_mut_ptr().cast::<u8>(),
std::mem::size_of_val(&buf),
)
})?;
let p = unsafe { buf.assume_init() };
if p.is_null() {
return Ok(());
}
let weak = weak_from_raw(p);
if let Some(task) = weak.upgrade() {
let task = unsafe { Pin::new_unchecked(task) };
self.as_mut().enqueue(task);
}
Ok(())
}
}
fn weak_from_raw(p: *const ()) -> Weak<ExecutorTask> {
unsafe { Weak::from_raw(p.cast::<ExecutorTask>()) }
}
mod task_waker {
use super::{ExecutorTask, Weak, exec, weak_from_raw};
use nix::sys::epoll::EpollTimeout;
use nix::unistd::write;
use std::{
os::fd::{AsRawFd, BorrowedFd},
pin::Pin,
sync::Arc,
task::{RawWaker, RawWakerVTable, Waker},
};
pub(super) fn build(task: &Weak<ExecutorTask>) -> Waker {
unsafe { Waker::from_raw(RawWaker::new(Weak::into_raw(task.clone()).cast(), &VTABLE)) }
}
pub(super) const VTABLE: RawWakerVTable =
RawWakerVTable::new(clone, wake_by_val, wake_by_ref, drop);
fn wake(weak: &Weak<ExecutorTask>) {
if let Some(task) = weak.upgrade() {
exec(|mut e| {
if task.pipe_wr == e.pipe_wr.as_raw_fd() {
if task.task.try_borrow().is_err() {
let _ = e.as_mut().epoll_wait(EpollTimeout::ZERO);
}
let task = unsafe { Pin::new_unchecked(task) };
e.enqueue(task);
} else {
let p = Weak::into_raw(weak.clone());
let bytes = (p as usize).to_ne_bytes();
match write(unsafe { BorrowedFd::borrow_raw(task.pipe_wr) }, &bytes) {
Ok(_) | Err(nix::errno::Errno::EBADF) => {}
Err(e) => panic!("error writing to task wake pipe: {e:?}"),
}
}
});
}
}
unsafe fn clone(p: *const ()) -> RawWaker {
let weak = weak_from_raw(p);
let clone = if let Some(task) = weak.upgrade() {
#[cfg(feature = "tracing")]
if let Some(task_span) = &task.tracing_span {
tracing::trace!(
target: "runtime::waker",
op = "waker.clone",
task.id = task_span.into_u64()
);
}
RawWaker::new(Weak::into_raw(Arc::downgrade(&task)).cast(), &VTABLE)
} else {
RawWaker::new(Waker::noop().data(), Waker::noop().vtable())
};
let _ = Weak::into_raw(weak);
clone
}
unsafe fn wake_by_val(p: *const ()) {
let weak = weak_from_raw(p);
#[cfg(feature = "tracing")]
if let Some(task) = weak.upgrade() {
if let Some(task_span) = &task.tracing_span {
tracing::trace!(
target: "runtime::waker",
op = "waker.wake_by_val",
task.id = task_span.into_u64()
);
}
}
wake(&weak);
}
unsafe fn wake_by_ref(p: *const ()) {
let weak = weak_from_raw(p);
#[cfg(feature = "tracing")]
if let Some(task) = weak.upgrade() {
if let Some(task_span) = &task.tracing_span {
tracing::trace!(
target: "runtime::waker",
op = "waker.wake_by_ref",
task.id = task_span.into_u64()
);
}
}
wake(&weak);
let _ = Weak::into_raw(weak);
}
unsafe fn drop(p: *const ()) {
let weak = weak_from_raw(p);
#[cfg(feature = "tracing")]
if let Some(task) = weak.upgrade() {
if let Some(task_span) = &task.tracing_span {
tracing::trace!(
target: "runtime::waker",
op = "waker.drop",
task.id = task_span.into_u64()
);
}
}
#[cfg(not(feature = "tracing"))]
let _ = weak;
}
}
pub(crate) struct EpollWaker {
waker: Waker,
events: EpollFlags,
burst: usize,
_not_send_not_sync: core::marker::PhantomData<*mut ()>,
}
#[must_use]
pub(crate) enum ControlFlow<T> {
Normal(Poll<T>),
Yield,
}
impl EpollWaker {
pub(crate) fn new(events: EpollFlags) -> Self {
Self {
waker: Waker::noop().clone(),
events,
burst: 0,
_not_send_not_sync: PhantomData,
}
}
pub(crate) fn poll_with<T, F: FnMut(EpollFlags) -> Result<T, std::io::Error>>(
&mut self,
cx: &std::task::Context<'_>,
mut func: F,
) -> ControlFlow<Result<T, std::io::Error>> {
if self.events.is_empty() {
self.burst = 0;
self.waker.clone_from(cx.waker());
return ControlFlow::Normal(Poll::Pending);
}
const BURST_LIMIT: usize = 20;
if self.burst >= BURST_LIMIT {
self.burst = 0;
return ControlFlow::Yield;
}
loop {
match func(self.events) {
Err(e) if e.kind() == std::io::ErrorKind::Interrupted => { }
Err(e) if e.kind() == std::io::ErrorKind::WouldBlock => {
self.burst = 0;
self.waker.clone_from(cx.waker());
self.events = EpollFlags::empty();
break ControlFlow::Normal(Poll::Pending);
}
res => {
self.burst += 1;
break ControlFlow::Normal(Poll::Ready(res));
}
}
}
}
}
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
pub enum Priority {
High,
Normal,
Low,
}
#[track_caller]
pub fn spawn_checked<T: 'static, E: 'static, F: Future<Output = Result<T, E>> + 'static>(
future: F,
) -> task::Handle<Result<T, E>> {
Builder::new(future).spawn_checked()
}
#[track_caller]
pub fn spawn_checked_with_priority<
T: 'static,
E: 'static,
F: Future<Output = Result<T, E>> + 'static,
>(
future: F,
priority: Priority,
) -> task::Handle<Result<T, E>> {
Builder::new(future).priority(priority).spawn_checked()
}
#[track_caller]
pub fn spawn<F: Future + 'static>(future: F) -> task::Handle<F::Output> {
Builder::new(future).spawn()
}
#[track_caller]
pub fn spawn_with_priority<F: Future + 'static>(
future: F,
priority: Priority,
) -> task::Handle<F::Output> {
Builder::new(future).priority(priority).spawn()
}
impl<T: 'static, E: 'static, F: Future<Output = Result<T, E>> + 'static> Builder<F> {
pub fn spawn_checked(self) -> task::Handle<F::Output> {
let (future, meta, priority) = self.into_parts();
let future = async move {
future
.await
.inspect_err(|_| unsafe { shutdown_executor_unchecked() })
};
exec(|e| e.spawn(future, priority, meta))
}
}
impl<F: Future + 'static> Builder<F> {
pub fn spawn(self) -> task::Handle<F::Output> {
let (future, meta, priority) = self.into_parts();
exec(|e| e.spawn(future, priority, meta))
}
fn into_parts(self) -> (F, task::Metadata, Priority) {
let Self {
future,
priority,
spawn_location,
name,
} = self;
let meta = task::Metadata {
spawn_location,
priority,
spawn_checked: true,
name: name.unwrap_or_default(),
};
(future, meta, priority)
}
}
fn run_task(t: TaskRef) {
let pending = {
let mut cx = Context::from_waker(&t.waker);
let mut ti = t.task.borrow_mut();
let ti = unsafe { Pin::new_unchecked(&mut *ti) };
match ti.poll(&mut cx) {
Poll::Pending => true,
Poll::Ready(()) => {
t.handle_waker.replace(Waker::noop().clone()).wake_by_ref();
t.has_completed.replace(true);
false
}
}
};
if pending && !t.is_queued() {
exec(|e| e.sleep(t));
}
}
pub fn run() -> Result<(), std::io::Error> {
'outer: loop {
while let Some(runnable_task) = exec(|e| e.highest_priority_runnable_task()) {
run_task(runnable_task);
if exec(|e| e.shutdown) {
break 'outer;
}
exec(|e| e.epoll_wait(EpollTimeout::ZERO))?;
}
if exec(|e| e.sleepq.is_empty()) {
break;
}
exec(|e| e.epoll_wait(EpollTimeout::NONE))?;
}
let _ = EXECUTOR.replace(Executor::new()?);
Ok(())
}
#[allow(
clippy::must_use_candidate,
reason = "we want to inherit the must_use of F::Output"
)]
#[track_caller]
pub fn block_on<F: Future>(future: F) -> F::Output {
#[cfg(feature = "tracing")]
let future = {
use tracing::Instrument;
let location = std::panic::Location::caller();
let id = {
use std::hash::{Hash, Hasher};
let mut hasher = std::hash::DefaultHasher::new();
location.hash(&mut hasher);
hasher.finish()
};
let span = tracing::trace_span!(
target: "epox::task",
parent: None,
"runtime.spawn",
kind = "blocking",
task.id = id,
loc.file = location.file(),
loc.line = location.line(),
loc.col = location.column(),
);
future.instrument(span)
};
let block_on_waker = {
fn wake(p: *const ()) {
let waker_pipe_wr = p as i32;
exec(|e| {
if waker_pipe_wr == e.pipe_wr.as_raw_fd() {
return;
}
let bytes = 0usize.to_ne_bytes();
match write(unsafe { BorrowedFd::borrow_raw(waker_pipe_wr) }, &bytes) {
Ok(_) | Err(nix::errno::Errno::EBADF) => {}
Err(e) => panic!("error writing to task wake pipe: {e:?}"),
}
});
}
const VTABLE: RawWakerVTable = RawWakerVTable::new(
|p| RawWaker::new(p, &VTABLE),
|p| wake(p),
|p| wake(p),
|_| {},
);
RawWaker::new(exec(|e| e.pipe_wr.as_raw_fd()) as *const _, &VTABLE)
};
let waker = unsafe { std::task::Waker::from_raw(block_on_waker) };
let mut cx = std::task::Context::from_waker(&waker);
let mut future = std::pin::pin!(future);
loop {
if let Poll::Ready(v) = future.as_mut().poll(&mut cx) {
return v;
}
if let Some(runnable_task) = exec(|e| e.highest_priority_runnable_task()) {
run_task(runnable_task);
exec(|e| e.epoll_wait(EpollTimeout::ZERO)).unwrap();
} else {
exec(|e| e.epoll_wait(EpollTimeout::NONE)).unwrap();
}
}
}
pub unsafe fn shutdown_executor_unchecked() {
exec(|e| e.shutdown());
}
pub const fn shutdown() -> Shutdown {
Shutdown { _private: () }
}
#[must_use]
pub struct Shutdown {
_private: (),
}
impl Future for Shutdown {
#[cfg(feature = "nightly")]
type Output = !;
#[cfg(not(feature = "nightly"))]
type Output = std::convert::Infallible;
fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<Self::Output> {
unsafe { shutdown_executor_unchecked() };
Poll::Pending
}
}