odem_rs_core/continuation/

mod.rs

//! The `continuation` module provides the implementation of a type-erased,
//! intrusively linked futures abstraction used in the scheduler of the
//! simulation library.
//!
//! Continuations are held by a calendar, allowing them to be scheduled and
//! managed based on model time.
//!
//! The [`Continuation`] type represents a continuation that can be scheduled
//! and executed within the simulation library. It is designed to handle
//! type-erased futures and supports branding to tie [token witnesses] to their
//! runtime-state, enabling statically checked validity of state transitions.
//!
//! [token witnesses]: token::State

use core::{
	any::Any,
	cell::Cell,
	fmt,
	hint::unreachable_unchecked,
	panic::Location,
	pin::Pin,
	ptr::NonNull,
	sync::atomic::{AtomicUsize, Ordering},
	task::{Context, Poll},
};
use intrusive_collections::{LinkedList, RBTreeLink};

pub use adapter::{Adapter, PointerOps};
pub use puck::Puck;
pub use share::{Label, Share};

use crate::{
	Dispatch, ExitStatus,
	calendar::{PlanState, Scheduler},
	config::Config,
	error::NotIdle,
	fsm::*,
	ptr::{IntrusivelyCounted, Irc, IrcBox, IrcBoxed},
	simulator::{Mark, Prec},
};

mod adapter;
mod puck;
mod share;

/* ********************************************************************* Continuation */

/// The type of intrusive link stored in every [Continuation].
pub type Link = RBTreeLink;

/// A structure of type-erased, intrusively linked [futures].
///
/// Pointers to instances of this type are used in the scheduler of the
/// simulation library. They are oblivious to lifetime-restrictions and
/// erase the concrete types of the future and the process instance.
/// The latter types are recovered during dynamic lookup, but the
/// lifetime-restrictions have to be observed by the abstractions provided
/// in the simulation library.
///
/// Continuations may be *branded* using an additional lifetime that ties
/// [Token] to specific instances. This allows composition and chaining of
/// method calls that depend on specific [States] without having to
/// continuously test the current states within the method implementations.
///
/// [futures]: Future
/// [Token]: token
/// [States]: State
pub struct Continuation<'brand, C: ?Sized + Config> {
	/// An intrusive link for insertion of a continuation into the calendar.
	hook: Link,
	/// Intrusive reference counter and raw pointer to the most specialized
	/// version of the type containing this continuation.
	///
	/// This is needed to provide run-time polymorphism over the type
	/// of future being executed as well as what code is run on dropping. See
	/// [Dispatch] for further information.
	task_box: IrcBox<ContBox>,
	/// Pointer to data shared across different continuations.
	share: NonNull<Share<C>>,
	/// A list of continuations waiting for this continuation to terminate.
	pending: Cell<LinkedList<Adapter<C>>>,
	/// Precedence of this continuation relative to other continuations from the
	/// same agent.
	prec: Cell<Prec>,
	/// Current state of this continuation instance.
	state: StateMachine<'brand, State<C>>,
	/// Span information for the current continuation.
	#[cfg(feature = "tracing")]
	span: tracing::Span,
}

impl<C: ?Sized + Config> Continuation<'static, C> {
	/// Creates a new continuation initialized in [State::Born] with specific
	/// location information attached.
	///
	/// [State::Born]: erased::State::Born
	pub(crate) fn new(prec: Prec, location: &'static Location<'static>) -> Self {
		Self {
			hook: Link::new(),
			task_box: IrcBox::with_location(ContBox::new(), location),
			share: NonNull::dangling(),
			pending: Cell::new(LinkedList::new(Adapter::NEW)),
			prec: Cell::new(prec),
			state: StateMachine::default(),
			#[cfg(feature = "tracing")]
			span: tracing::Span::current(),
		}
	}

	/// Returns the exit code of the continuation, or `None` if it hasn't
	/// terminated.
	pub(crate) fn result(&self) -> Option<ExitStatus> {
		self.brand(|task, once| {
			task.token(once)
				.into_term()
				.map(|state| task.branded_result(&state))
				.ok()
		})
	}

	/// Reactivates the continuation at the current model-time.
	pub(crate) fn wake(this: Irc<Self>) -> Result<(), NotIdle> {
		this.clone().brand(|task, once| {
			let idle = task.token(once).into_idle()?;

			task.branded_share(&idle)
				.sim()
				.calendar()
				.activate(task.clone(), idle);

			Ok(())
		})
	}

	/// Returns the next model time this continuation is scheduled for, or
	/// `None` if it isn't scheduled.
	pub(crate) fn time(&self) -> Option<C::Time> {
		self.brand(|task, once| {
			task.token(once)
				.into_next()
				.ok()
				.map(|next| task.next_time(&next))
		})
	}
}

impl<'brand, C: ?Sized + Config> Continuation<'brand, C> {
	/// Enters the [`Span`] associated with the agent owning the shared data.
	///
	/// [`Span`]: tracing::Span
	#[cfg(feature = "tracing")]
	pub fn enter_span(&self) -> tracing::span::Entered<'_> {
		self.span.enter()
	}

	#[cfg(not(feature = "tracing"))]
	pub const fn enter_span(&self) {}

	/// Sets a new virtual function pointer for this continuation.
	///
	/// # Safety
	/// The caller is responsible to ensure that the pointer stays valid for
	/// the duration of the continuation's life.
	pub(crate) unsafe fn set_vptr(&self, vptr: NonNull<dyn Dispatch + '_>) {
		unsafe {
			self.task_box.set_vptr(vptr);
		}
	}

	/// Clears the virtual function pointer for this continuation, preventing any
	/// methods from being called.
	///
	/// This can be used to prevent [`Dispatch::reclaim`] from being called,
	/// even if the reference counter reaches zero.
	pub(crate) fn clear_vptr(&self) {
		self.task_box.clear_vptr();
	}

	/// Returns the number of references to this continuation.
	pub(crate) fn use_count(&self) -> usize {
		self.task_box.refs.load(Ordering::Relaxed)
	}

	/// Returns a copy of the internal [State](erased::State).
	pub(crate) fn state(&self) -> &StateMachine<'brand, State<C>> {
		&self.state
	}

	/// Returns the current [Prec] of this continuation.
	pub(crate) fn prec(&self) -> Prec {
		self.prec.get()
	}

	/// Sets the new [Prec] of this continuation.
	pub(crate) fn set_prec(&self, prec: Prec) {
		self.prec.set(prec);
	}

	/// Adds another continuation to the list of continuations to be awoken upon terminating.
	#[inline]
	pub(crate) fn insert_pending(&self, other: Irc<Continuation<'static, C>>) {
		let mut list = self.pending.take();
		list.push_back(other);
		self.pending.set(list);
	}

	/// Removes a previously added continuation from the list of pending continuations.
	///
	/// # Safety
	/// It is the caller's responsibility to ensure that the continuation had been
	/// added previously via [Self::insert_pending].
	pub(crate) unsafe fn remove_pending(&self, other: &Continuation<'static, C>) {
		// Only unlink if the other continuation is actually linked right now.
		// This can happen during panics, when the `drop` impl of a `Join`
		// attempts to remove a stored continuation from the pending list, but
		// it has already been removed by a prior reactivation.
		if other.hook.is_linked() {
			let mut list = self.pending.take();

			unsafe {
				list.cursor_mut_from_ptr(other.detach()).remove();
			}

			self.pending.set(list);
		}
	}

	/// Awakens all pending continuations and clears the list.
	pub(crate) fn wake_pending(&self) {
		for task in self.pending.take() {
			Continuation::wake(task).ok();
		}
	}

	/// Converts the specific brand into a generic brand, breaking the
	/// connection with the equally branded token.
	pub(crate) fn detach(&self) -> &Continuation<'static, C> {
		unsafe { core::mem::transmute(self) }
	}

	/// Returns the [`Location`] information for this `Continuation`.
	pub(crate) fn location(&self) -> &'static Location<'static> {
		IrcBox::location(&self.task_box)
	}

	/// Performs a runtime-check if this continuation has been dereferenced on the
	/// same thread as the one the executor is running and panicks if that is
	/// not the case.
	///
	/// # Safety
	/// This method can only be called after the `Continuation` has been
	/// activated. Calling it is thread-safe.
	unsafe fn is_same_thread(&self) -> bool {
		// Extract the shared data from the task.
		let share = unsafe { self.share.as_ref() };

		// Compare the pointer-address of this continuation's simulation context
		// to the pointer-address of the thread-local simulation context.
		crate::erased::with(|sim| core::ptr::addr_eq(&**share.sim(), sim)).unwrap_or(false)
	}

	/// Returns an enumeration copy of the internal [State](token::State).
	pub(crate) fn token(&self, once: Ephemeral<'brand>) -> token::State<'brand> {
		self.state.token(once)
	}

	/// Binds a continuation to [shared data].
	///
	/// # Safety
	/// The caller is responsible to ensure that the shared-data-reference
	/// outlives the (active) part of the continuation's life.
	///
	/// [shared data]: Share
	pub(crate) unsafe fn bind(
		mut self: Pin<&mut Self>,
		born: token::Born<'brand>,
		share: &Share<C>,
	) -> token::Idle<'brand> {
		// assign the reference to the shared data
		self.share = NonNull::from(share);

		// enter our `Span` to properly record the transition
		let _span = self.enter_span();

		// transition into state `Idle`
		self.state.transition(born, ())
	}

	/// Purges the continuation from the calendar.
	pub(crate) fn deschedule(&self, next: token::Next<'brand>) -> token::Idle<'brand> {
		let share = self.branded_share(&next);

		// remove the continuation from the calendar
		share.sim().calendar().remove(self, next)
	}

	/// Removes the active state from the continuation.
	pub(crate) fn deactivate(&self, busy: token::Busy<'brand>) -> token::Idle<'brand> {
		let share = self.branded_share(&busy);

		// deregister the continuation from the active cell
		share.sim().unslot(self, busy)
	}

	/// Polls the underlying future of this continuation.
	pub(crate) fn poll(&self, busy: token::Busy<'brand>, cx: &mut Context<'_>) -> Poll<()> {
		// read the virtual function from the table

		// SAFETY: the busy-token testifies that the continuation has been
		// bound, which ensures that the vptr to the virtual function table has
		// been set; pinning is ensured by the binding routine requiring it
		let vtab = unsafe { Pin::new_unchecked(self.task_box.vptr.get().unwrap().as_ref()) };

		// temporarily escape the branding to allow `Future::poll()` to rebrand
		// without accidentally creating two tokens for the same instance
		let (once, res) = self.state.debrand(busy, move |_| vtab.poll(cx));

		// analyze the resulting state
		match self.token(once).into_busy() {
			Ok(busy) => match res {
				Poll::Ready(result) => {
					// busy -> done
					let _: token::Done<'_> = self.state.transition(busy, result);

					// reactivate pending continuations on completion
					self.wake_pending();

					Poll::Ready(())
				}
				Poll::Pending => {
					// busy -> idle
					let _: token::Idle<'_> = self.state.transition(busy, ());
					Poll::Pending
				}
			},
			Err(_err) => {
				debug_assert!(
					res.is_pending(),
					"task in state `{:?}` should not have been able to terminate",
					_err.0
				);
				Poll::Pending
			}
		}
	}

	/// Returns a reference to the shared data for this continuation.
	pub(crate) fn branded_share<'s, I>(&'s self, init: &I) -> &'s Share<C>
	where
		I: Into<token::Init<'brand>>,
	{
		let _ = init;

		// SAFETY: the token witness testifies that the shared data is
		// initialized, which only happens during binding
		unsafe { self.share.as_ref() }
	}

	/// Returns a reference to the shared data if initialization of this continuation
	/// has been completed.
	pub(crate) fn share(&self) -> Option<&Share<C>> {
		if self.state().erased().is_init() {
			// SAFETY: once bound, the pointer stays valid
			Some(unsafe { self.share.as_ref() })
		} else {
			None
		}
	}

	/// Grants access to the calendar state if in state [`Next`](State::Next).
	pub(crate) fn next_state<F, R>(&self, next: &token::Next<'brand>, f: F) -> R
	where
		F: FnOnce(&<C::Plan as Scheduler>::State) -> R,
	{
		let _ = next;

		// SAFETY: the `Next` token guarantees that this continuation is in the
		// correct state
		match &*self.state.borrow() {
			State::Next(state) => f(state),
			_ => unsafe { unreachable_unchecked() },
		}
	}

	/// Returns the model time that the continuation will be activated.
	pub(crate) fn next_time(&self, next: &token::Next<'brand>) -> C::Time {
		self.next_state(next, |s| s.time())
	}

	/// Returns the [Cell] containing the current mark of the shared data
	/// associated with this continuation.
	///
	/// This value is used to organize different continuations with identical shared
	/// data in the calendar such that they are executed in a contiguous batch.
	pub(crate) fn mark<'s, I>(&'s self, init: &I) -> &'s Cell<Mark>
	where
		I: Into<token::Init<'brand>>,
	{
		self.branded_share(init).mark()
	}

	/// Returns the exit code of this continuation.
	pub(crate) fn branded_result<T>(&self, _: &T) -> ExitStatus
	where
		T: Into<token::Term<'brand>>,
	{
		// SAFETY: the continuation is in state `Done` or `Gone` per the token witness
		match &*self.state.borrow() {
			State::Done(rc) | State::Gone(rc) => *rc,
			_ => unsafe { unreachable_unchecked() },
		}
	}
}

// Continuations offer unique methods for Branded variants
impl<'b, C: ?Sized + Config> Stateful for Continuation<'b, C> {
	type Brand = &'b ();

	unsafe fn enter(&self) {
		unsafe {
			self.state.enter();
		}
	}

	unsafe fn leave(&self) {
		unsafe {
			self.state.leave();
		}
	}
}

impl<'b, C: ?Sized + Config> Rebrand<'b> for Continuation<'b, C> {
	type Kind<'a> = Continuation<'a, C>;
}

// Continuations only contain pointers to pinned data and are not pinned themselves
impl<C: ?Sized + Config> Unpin for Continuation<'_, C> {}

impl<C: ?Sized + Config> fmt::Debug for Continuation<'_, C> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let mut s = f.debug_struct("Continuation");

		s.field("state", &self.state.borrow());

		self.brand(|task, once| {
			struct PrettyName(Label);

			impl fmt::Debug for PrettyName {
				fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
					write!(f, "\"{}\"", &self.0)
				}
			}

			let state = task.token(once);
			let init = state.as_init()?;

			s.field("label", &PrettyName(task.branded_share(init).label()))
				.field("rank", &task.branded_share(init).rank());

			None::<()>
		});

		s.field("prec", &self.prec().float_range())
			.field("refs", &self.use_count())
			.finish()
	}
}

impl<C: Config> crate::erased::Continuation for Continuation<'static, C> {
	fn subject(&self) -> &dyn Any {
		self.share().map_or(&(), |shared| shared.subject())
	}

	fn label(&self) -> Label {
		self.share()
			.map_or(Label::default(), |shared| shared.label())
	}

	fn prec(&self) -> Prec {
		self.prec()
	}

	fn state(&self) -> erased::State {
		self.state().borrow().erased()
	}
}

unsafe impl<C: ?Sized + Config> IntrusivelyCounted for Continuation<'_, C> {
	fn irc_box(&self) -> &IrcBox<dyn IrcBoxed> {
		&self.task_box
	}
}

/* ************************************************************** Irc Support */

/// Container type used to support [intrusive reference counting](Irc).
pub struct ContBox {
	/// Contains a counter for the number of references to the outer
	/// [`Continuation`].
	refs: AtomicUsize,
	/// Pointer to the most-specialized version of the instance.
	/// Used to recover type information during dynamic dispatch.
	vptr: Cell<Option<NonNull<dyn Dispatch>>>,
}

impl ContBox {
	/// Creates a new irc box from a location with zero references and no
	/// dispatcher.
	const fn new() -> Self {
		ContBox {
			refs: AtomicUsize::new(0),
			vptr: Cell::new(None),
		}
	}

	/// Sets the pointer to the vtable.
	///
	/// # Safety
	/// It is the callers' responsibility to ensure that the `vptr` pointer to
	/// the `dyn Dispatch` object outlives the continuation box.
	unsafe fn set_vptr(&self, vptr: NonNull<dyn Dispatch + '_>) {
		unsafe {
			use core::mem::transmute;

			self.vptr.set(Some(transmute::<
				NonNull<dyn Dispatch + '_>,
				NonNull<dyn Dispatch + 'static>,
			>(vptr)));
		}
	}

	/// Clears the pointer to the vtable.
	///
	/// This can be useful to prevent [`IrcBoxed::release`] from returning
	/// a destructor function in case the reference counter reaches zero.
	fn clear_vptr(&self) {
		self.vptr.set(None);
	}
}

unsafe impl IrcBoxed for ContBox {
	fn ref_count(&self) -> usize {
		self.refs.load(Ordering::SeqCst)
	}

	fn acquire(&self, _: crate::ptr::Private) {
		self.refs.fetch_add(1, Ordering::SeqCst);
	}

	fn release(&self, _: crate::ptr::Private) {
		self.refs.fetch_sub(1, Ordering::SeqCst);
	}

	fn reclaim(&self, _: crate::ptr::Private) -> Option<fn(NonNull<dyn IrcBoxed>)> {
		self.vptr.get().is_some().then_some(
			// return a function reclaiming the outer type if no references point
			// to the Continuation anymore; this separation is necessary to prevent
			// overlapping references to this Continuation from self and (indirectly) from
			// the inner dyn object
			|this| unsafe {
				// restore the pointer to the virtual table
				let this = this.cast::<Self>().as_ref();
				let mut vptr = this.vptr.get().unwrap_unchecked();

				// call the reclaim-method on an exclusive reference to the dyn
				// object; the temporary reference to the continuation has been dropped
				// previously, and no other active references exist to the continuation
				// which allows us to take this exclusive reference
				Pin::new_unchecked(vptr.as_mut()).reclaim();
			},
		)
	}
}

/* ************************************************************** Continuation States */

fsm! {
	/// Represents the state of a [`Continuation`] throughout its lifecycle.
	///
	/// A `Continuation` starts in the [`Born`] state without references to
	/// shared data or its future, as provided by its owning [Job] or [Agent].
	/// After binding pinned references to it, the `Continuation` moves to the
	/// [`Idle`] state, indicating it is unscheduled.
	///
	/// From the [`Idle`] state, the `Continuation` can be scheduled by
	/// inserting it into the event calendar at a specific model time,
	/// transitioning it to the [`Next`] state. When the model time reaches this
	/// point and the `Continuation` is activated, it enters the [`Busy`] state.
	///
	/// Upon completion, the `Continuation` transitions to the [`Done`] state,
	/// indicating that a return value is available for extraction. At any point
	/// before normal termination, the `Continuation` can be aborted, moving it
	/// to the [`Gone`] state and dropping its bound future. The [`Gone`] state
	/// is the final state in a `Continuation`'s lifecycle.
	///
	/// [Job]: crate::job::Job
	/// [Agent]: crate::Agent
	/// [`Born`]: State::Born
	/// [`Idle`]: State::Idle
	/// [`Next`]: State::Next
	/// [`Busy`]: State::Busy
	/// [`Done`]: State::Done
	/// [`Gone`]: State::Gone
	#[derive(Default)]
	pub enum State<C: Config> {
		/// State of a continuation signifying a non-bound future and shared data.
		#[default]
		Born -> {Idle, Gone},
		/// State of a continuation that is waiting for external reactivation.
		Idle -> {Next, Gone},
		/// State of the continuation that is currently active. At most one
		/// continuation may be in this state during a simulation run at any
		/// time.
		Busy -> {Idle, Done},
		/// State of a continuation that is managed by the calendar.
		Next(<C::Plan as Scheduler>::State) -> {Idle, Busy},
		/// State of a completed continuation with a result available for extraction.
		Done(ExitStatus) -> {Gone},
		/// State of a continuation that cannot be scheduled anymore.
		Gone(ExitStatus) -> {}
	}

	/// Meta-state for all non-[`Born`] states.
	///
	/// [`Born`]: State::Born
	pub Init = {Idle, Busy, Next, Done, Gone};

	/// Meta-state for [`Done`] and [`Gone`] states and a subset of the [`Init`]
	/// meta-state.
	///
	/// [`Done`]: State::Done
	/// [`Gone`]: State::Gone
	/// [`Init`]: token::Init
	pub Term: Init = {Done, Gone};
}

impl<C: ?Sized + Config> fmt::Debug for State<C> {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		let mut debug = f.debug_tuple(self.label());

		match self {
			State::Next(plan) => debug.field(plan),
			State::Done(exit) => debug.field(exit),
			State::Gone(exit) => debug.field(exit),
			_ => &mut debug,
		}
		.finish()
	}
}