Revision control
Copy as Markdown
Other Tools
use alloc::sync::{Arc, Weak};
use core::cell::UnsafeCell;
use core::sync::atomic::Ordering::{self, Relaxed, SeqCst};
use core::sync::atomic::{AtomicBool, AtomicPtr};
use super::abort::abort;
use super::ReadyToRunQueue;
use crate::task::{waker_ref, ArcWake, WakerRef};
pub(super) struct Task<Fut> {
// The future
pub(super) future: UnsafeCell<Option<Fut>>,
// Next pointer for linked list tracking all active tasks (use
// `spin_next_all` to read when access is shared across threads)
pub(super) next_all: AtomicPtr<Task<Fut>>,
// Previous task in linked list tracking all active tasks
pub(super) prev_all: UnsafeCell<*const Task<Fut>>,
// Length of the linked list tracking all active tasks when this node was
// inserted (use `spin_next_all` to synchronize before reading when access
// is shared across threads)
pub(super) len_all: UnsafeCell<usize>,
// Next pointer in ready to run queue
pub(super) next_ready_to_run: AtomicPtr<Task<Fut>>,
// Queue that we'll be enqueued to when woken
pub(super) ready_to_run_queue: Weak<ReadyToRunQueue<Fut>>,
// Whether or not this task is currently in the ready to run queue
pub(super) queued: AtomicBool,
// Whether the future was awoken during polling
// It is possible for this flag to be set to true after the polling,
// but it will be ignored.
pub(super) woken: AtomicBool,
}
// `Task` can be sent across threads safely because it ensures that
// the underlying `Fut` type isn't touched from any of its methods.
//
// The parent (`super`) module is trusted not to access `future`
// across different threads.
unsafe impl<Fut> Send for Task<Fut> {}
unsafe impl<Fut> Sync for Task<Fut> {}
impl<Fut> ArcWake for Task<Fut> {
fn wake_by_ref(arc_self: &Arc<Self>) {
let inner = match arc_self.ready_to_run_queue.upgrade() {
Some(inner) => inner,
None => return,
};
arc_self.woken.store(true, Relaxed);
// It's our job to enqueue this task it into the ready to run queue. To
// do this we set the `queued` flag, and if successful we then do the
// actual queueing operation, ensuring that we're only queued once.
//
// Once the task is inserted call `wake` to notify the parent task,
// as it'll want to come along and run our task later.
//
// Note that we don't change the reference count of the task here,
// we merely enqueue the raw pointer. The `FuturesUnordered`
// implementation guarantees that if we set the `queued` flag that
// there's a reference count held by the main `FuturesUnordered` queue
// still.
let prev = arc_self.queued.swap(true, SeqCst);
if !prev {
inner.enqueue(Arc::as_ptr(arc_self));
inner.waker.wake();
}
}
}
impl<Fut> Task<Fut> {
/// Returns a waker reference for this task without cloning the Arc.
pub(super) fn waker_ref(this: &Arc<Self>) -> WakerRef<'_> {
waker_ref(this)
}
/// Spins until `next_all` is no longer set to `pending_next_all`.
///
/// The temporary `pending_next_all` value is typically overwritten fairly
/// quickly after a node is inserted into the list of all futures, so this
/// should rarely spin much.
///
/// When it returns, the correct `next_all` value is returned.
///
/// `Relaxed` or `Acquire` ordering can be used. `Acquire` ordering must be
/// used before `len_all` can be safely read.
#[inline]
pub(super) fn spin_next_all(
&self,
pending_next_all: *mut Self,
ordering: Ordering,
) -> *const Self {
loop {
let next = self.next_all.load(ordering);
if next != pending_next_all {
return next;
}
}
}
}
impl<Fut> Drop for Task<Fut> {
fn drop(&mut self) {
// Since `Task<Fut>` is sent across all threads for any lifetime,
// regardless of `Fut`, we, to guarantee memory safety, can't actually
// touch `Fut` at any time except when we have a reference to the
// `FuturesUnordered` itself .
//
// Consequently it *should* be the case that we always drop futures from
// the `FuturesUnordered` instance. This is a bomb, just in case there's
// a bug in that logic.
unsafe {
if (*self.future.get()).is_some() {
abort("future still here when dropping");
}
}
}
}