use alloc::boxed::Box;
use core::any::Any;
use core::future::Future;
use proxima_core::ProximaError;
use crate::{BackgroundHandle, CoreId, Runtime, SpawnError};
pub trait RuntimeExt: Runtime {
fn spawn<F>(&self, future: F)
where
F: Future<Output = ()> + 'static,
{
self.spawn_on_current_core(Box::pin(future));
}
fn core(&self, placement: impl Into<Placement>) -> OnCore<'_, Self> {
OnCore {
runtime: self,
placement: placement.into(),
}
}
fn blocking<R, W>(&self, work: W) -> Blocking<'_, Self, W, R>
where
R: Send + 'static,
W: FnOnce() -> Result<R, ProximaError> + Send + 'static,
{
Blocking {
runtime: self,
work,
_result: core::marker::PhantomData,
}
}
}
impl<R: Runtime + ?Sized> RuntimeExt for R {}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Placement {
Only(CoreId),
AnyOther,
Any,
AnyOf(alloc::vec::Vec<CoreId>),
}
impl From<CoreId> for Placement {
fn from(core_id: CoreId) -> Self {
Placement::Only(core_id)
}
}
impl From<usize> for Placement {
fn from(core_id: usize) -> Self {
Placement::Only(CoreId(core_id))
}
}
impl<T: Into<CoreId>> From<alloc::vec::Vec<T>> for Placement {
fn from(cores: alloc::vec::Vec<T>) -> Self {
Placement::AnyOf(cores.into_iter().map(Into::into).collect())
}
}
impl<T: Into<CoreId>, const N: usize> From<[T; N]> for Placement {
fn from(cores: [T; N]) -> Self {
Placement::AnyOf(cores.into_iter().map(Into::into).collect())
}
}
impl Placement {
fn candidates(&self, current: CoreId, cores: usize) -> alloc::vec::Vec<CoreId> {
match self {
Placement::Only(core_id) => alloc::vec![*core_id],
Placement::AnyOf(list) => list.clone(),
Placement::Any => (0..cores)
.map(|offset| CoreId((current.0 + offset) % cores))
.collect(),
Placement::AnyOther => (1..cores)
.map(|offset| CoreId((current.0 + offset) % cores))
.collect(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Exhausted {
pub attempts: u32,
pub last: Option<SpawnError>,
}
pub struct OnCore<'rt, R: ?Sized> {
runtime: &'rt R,
placement: Placement,
}
impl<'rt, R: Runtime + ?Sized> OnCore<'rt, R> {
pub fn spawn<F>(self, future: F) -> Result<(), SpawnError>
where
F: Future<Output = ()> + Send + 'static,
{
let candidates = self
.placement
.candidates(self.runtime.current_core(), self.runtime.num_cores());
let Some(core_id) = candidates.first().copied() else {
return Err(SpawnError::Disconnected);
};
self.runtime.spawn_on_core(core_id, Box::pin(future))
}
pub fn spawn_with<Factory, F, Decide>(
self,
factory: Factory,
mut decide: Decide,
) -> Result<CoreId, Exhausted>
where
Factory: Fn() -> F,
F: Future<Output = ()> + Send + 'static,
Decide: FnMut(u32, &SpawnError) -> bool,
{
let candidates = self
.placement
.candidates(self.runtime.current_core(), self.runtime.num_cores());
let mut attempts = 0u32;
let mut last = None;
for core_id in candidates {
match self.runtime.spawn_on_core(core_id, Box::pin(factory())) {
Ok(()) => return Ok(core_id),
Err(err) => {
let keep_going = decide(attempts, &err);
attempts = attempts.saturating_add(1);
last = Some(err);
if !keep_going {
break;
}
}
}
}
Err(Exhausted { attempts, last })
}
}
pub struct Blocking<'rt, R: ?Sized, W, Out> {
runtime: &'rt R,
work: W,
_result: core::marker::PhantomData<fn() -> Out>,
}
impl<'rt, R, W, Out> Blocking<'rt, R, W, Out>
where
R: Runtime + ?Sized,
Out: Send + 'static,
W: FnOnce() -> Result<Out, ProximaError> + Send + 'static,
{
pub fn spawn(self) -> impl Future<Output = Result<Out, ProximaError>> + Send {
let handle: BackgroundHandle<Box<dyn Any + Send>> =
self.runtime.spawn_background_blocking(Box::new(move || {
(self.work)().map(|value| Box::new(value) as Box<dyn Any + Send>)
}));
async move {
let boxed = handle.await?;
boxed
.downcast::<Out>()
.map(|value| *value)
.map_err(|_| ProximaError::Config("background payload type changed".into()))
}
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use alloc::sync::Arc;
use core::sync::atomic::{AtomicUsize, Ordering};
use super::*;
fn block_on<F: Future>(future: F) -> F::Output {
let mut future = core::pin::pin!(future);
let waker = core::task::Waker::noop();
let mut cx = core::task::Context::from_waker(waker);
loop {
if let core::task::Poll::Ready(value) = future.as_mut().poll(&mut cx) {
return value;
}
}
}
#[derive(Default)]
struct StubRuntime {
on_core: AtomicUsize,
background: AtomicUsize,
full_for_first: AtomicUsize,
}
impl Runtime for StubRuntime {
fn spawn_on_current_core(
&self,
_f: core::pin::Pin<Box<dyn Future<Output = ()> + 'static>>,
) {
}
fn spawn_on_core(
&self,
_core_id: CoreId,
_future: core::pin::Pin<Box<dyn Future<Output = ()> + Send + 'static>>,
) -> Result<(), SpawnError> {
let attempt = self.on_core.fetch_add(1, Ordering::Relaxed);
if attempt < self.full_for_first.load(Ordering::Relaxed) {
return Err(SpawnError::InboxFull);
}
Ok(())
}
fn spawn_factory_on_core(
&self,
_core_id: CoreId,
_factory: Box<
dyn FnOnce() -> core::pin::Pin<Box<dyn Future<Output = ()> + 'static>> + Send,
>,
) -> Result<(), SpawnError> {
Ok(())
}
fn spawn_background_blocking(
&self,
work: Box<dyn FnOnce() -> Result<Box<dyn Any + Send>, ProximaError> + Send>,
) -> BackgroundHandle<Box<dyn Any + Send>> {
self.background.fetch_add(1, Ordering::Relaxed);
let result = work();
Box::pin(async move { result })
}
fn timer_at(
&self,
_deadline: std::time::Instant,
) -> core::pin::Pin<Box<dyn Future<Output = ()>>> {
Box::pin(async {})
}
fn num_cores(&self) -> usize {
1
}
fn current_core(&self) -> CoreId {
CoreId(0)
}
}
#[test]
fn builder_is_reachable_through_a_trait_object() {
let stub = Arc::new(StubRuntime::default());
let erased: Arc<dyn Runtime> = stub.clone();
erased.core(CoreId(0)).spawn(async {}).expect("chosen core");
assert_eq!(stub.on_core.load(Ordering::Relaxed), 1);
}
#[test]
fn any_other_excludes_the_current_core() {
let candidates = Placement::AnyOther.candidates(CoreId(2), 4);
assert_eq!(candidates, alloc::vec![CoreId(3), CoreId(0), CoreId(1)]);
assert!(!candidates.contains(&CoreId(2)));
}
#[test]
fn any_includes_the_current_core_last_in_wrap_order() {
assert_eq!(
Placement::Any.candidates(CoreId(2), 4),
alloc::vec![CoreId(2), CoreId(3), CoreId(0), CoreId(1)]
);
}
#[test]
fn any_other_on_one_core_yields_no_candidate() {
assert!(Placement::AnyOther.candidates(CoreId(0), 1).is_empty());
let stub = Arc::new(StubRuntime::default());
let erased: Arc<dyn Runtime> = stub.clone();
let outcome = erased.core(Placement::AnyOther).spawn(async {});
assert_eq!(outcome, Err(SpawnError::Disconnected));
assert_eq!(stub.on_core.load(Ordering::Relaxed), 0);
}
#[test]
fn a_bare_index_is_a_specific_core() {
assert_eq!(Placement::from(3usize), Placement::Only(CoreId(3)));
}
#[test]
fn a_bare_literal_array_is_a_candidate_list() {
let stub = Arc::new(StubRuntime::default());
let erased: Arc<dyn Runtime> = stub.clone();
erased
.core([1, 2, 3])
.spawn(async {})
.expect("first candidate");
assert_eq!(stub.on_core.load(Ordering::Relaxed), 1);
}
#[test]
fn arrays_and_vecs_of_either_element_type_become_candidate_lists() {
let expected = Placement::AnyOf(alloc::vec![CoreId(1), CoreId(2)]);
assert_eq!(Placement::from([1usize, 2]), expected);
assert_eq!(Placement::from([CoreId(1), CoreId(2)]), expected);
assert_eq!(Placement::from(alloc::vec![1usize, 2]), expected);
assert_eq!(Placement::from(alloc::vec![CoreId(1), CoreId(2)]), expected);
}
#[test]
fn a_full_core_falls_through_to_the_next_candidate() {
let stub = Arc::new(StubRuntime::default());
stub.full_for_first.store(2, Ordering::Relaxed);
let erased: Arc<dyn Runtime> = stub.clone();
let outcome = erased
.core([0, 1, 2])
.spawn_with(|| async {}, |_, err| *err == SpawnError::InboxFull);
assert_eq!(outcome, Ok(CoreId(2)), "reports which core took it");
assert_eq!(
stub.on_core.load(Ordering::Relaxed),
3,
"two full, third accepted"
);
}
#[test]
fn a_decline_stops_before_the_remaining_candidates() {
let stub = Arc::new(StubRuntime::default());
stub.full_for_first.store(3, Ordering::Relaxed);
let erased: Arc<dyn Runtime> = stub.clone();
let outcome = erased.core([0, 1, 2]).spawn_with(|| async {}, |_, _| false);
assert_eq!(
outcome,
Err(Exhausted {
attempts: 1,
last: Some(SpawnError::InboxFull)
})
);
assert_eq!(
stub.on_core.load(Ordering::Relaxed),
1,
"stopped after the first"
);
}
#[test]
fn no_candidates_reports_zero_attempts_and_no_cause() {
let stub = Arc::new(StubRuntime::default());
let erased: Arc<dyn Runtime> = stub.clone();
let outcome = erased
.core(Placement::AnyOther)
.spawn_with(|| async {}, |_, _| true);
assert_eq!(
outcome,
Err(Exhausted {
attempts: 0,
last: None
})
);
assert_eq!(stub.on_core.load(Ordering::Relaxed), 0);
}
#[test]
fn every_candidate_full_reports_exhaustion_not_the_last_error() {
let stub = Arc::new(StubRuntime::default());
stub.full_for_first.store(9, Ordering::Relaxed);
let erased: Arc<dyn Runtime> = stub.clone();
let outcome = erased.core([0, 1, 2]).spawn_with(|| async {}, |_, _| true);
assert_eq!(
outcome,
Err(Exhausted {
attempts: 3,
last: Some(SpawnError::InboxFull)
})
);
}
#[test]
fn blocking_returns_the_concrete_type_not_an_any_box() {
let stub = Arc::new(StubRuntime::default());
let erased: Arc<dyn Runtime> = stub.clone();
let fut = erased.blocking(|| Ok(7u32)).spawn();
let value = block_on(fut).expect("background work");
assert_eq!(value, 7u32);
assert_eq!(stub.background.load(Ordering::Relaxed), 1);
}
}