use super::*;
use std::alloc::{Layout, LayoutError};
#[derive(Debug, Clone, Copy)]
pub struct SharedLayout {
pub payload: Layout,
pub allocation: Layout,
}
pub fn shared(payload: Layout) -> Result<SharedLayout, LayoutError> {
Ok(SharedLayout {
payload,
allocation: Layout::new::<[std::sync::atomic::AtomicUsize; 2]>()
.extend(payload)?
.0
.pad_to_align(),
})
}
#[allow(dead_code)]
struct RootState {
publisher: Option<saddle_core::request_context::RequestRootPublisher>,
current: Option<saddle_core::request_context::RequestExecutionView>,
output: Option<EmergencyDiagnosticHandle>,
id: Option<tokio::task::Id>,
panic: Option<saddle_observability::root_diagnostic::RootRequestFailure>,
cancelled: Option<saddle_observability::root_diagnostic::RootRequestFailure>,
cleanup: Option<saddle_observability::root_diagnostic::RootRequestFailure>,
control: Option<saddle_observability::root_diagnostic::RootRequestFailure>,
finished: bool,
}
#[allow(dead_code)]
struct RootOutcome<T> {
result: saddle_observability::root_diagnostic::RootSupervisionReturn<T>,
state: Arc<Mutex<RootState>>,
join_error: Option<tokio::task::JoinError>,
}
#[derive(Debug)]
pub struct TaskLayouts {
pub pair_box: Layout,
pub inner_box: Layout,
pub future: Layout,
pub outcome: Layout,
pub join_result: Layout,
pub legacy_state: SharedLayout,
pub root_state: SharedLayout,
pub root_outcome: Layout,
pub root_join_result: Layout,
}
fn output<A, R>(_: impl FnOnce(A) -> R) -> Layout {
Layout::new::<R>()
}
pub fn pair_layouts<I, M, B, T, E>(_: impl FnOnce(I) -> M) -> Result<TaskLayouts, LayoutError>
where
M: FnOnce(RequestTaskContext) -> B,
B: Future<Output = Result<T, FrameworkRequestFailure<E>>>,
{
Ok(TaskLayouts {
pair_box: output(|(make, context): (M, RequestTaskContext)| {
super::pair_body(make, context)
}),
inner_box: Layout::new::<B>(),
future: Layout::new::<
RequestTaskFuture<std::future::Pending<Result<T, FrameworkRequestFailure<E>>>>,
>(),
outcome: Layout::new::<RequestTaskOutcome<T, E>>(),
join_result: Layout::new::<Result<RequestTaskOutcome<T, E>, tokio::task::JoinError>>(),
legacy_state: shared(Layout::new::<Mutex<State>>())?,
root_state: shared(Layout::new::<Mutex<RootState>>())?,
root_outcome: Layout::new::<RootOutcome<T>>(),
root_join_result: Layout::new::<Result<RootOutcome<T>, tokio::task::JoinError>>(),
})
}
#[derive(Debug, Clone, Copy)]
pub struct TaskBackendLayouts {
pub header: Layout,
pub scheduler: Layout,
pub task_id: Layout,
pub trailer: Layout,
pub cell_alignment: usize,
pub joinset_entry: Layout,
pub joinset_shared: Layout,
}
pub fn task_cell(
backend: TaskBackendLayouts,
future: Layout,
result: Layout,
) -> Result<Layout, LayoutError> {
let union = Layout::from_size_align(
future.size().max(result.size()),
future.align().max(result.align()),
)?
.pad_to_align();
let stage = Layout::new::<u32>().extend(union)?.0.pad_to_align();
let core = backend
.scheduler
.extend(backend.task_id)?
.0
.extend(stage)?
.0
.pad_to_align();
Ok(backend
.header
.extend(core)?
.0
.extend(backend.trailer)?
.0
.align_to(backend.cell_alignment)?
.pad_to_align())
}
pub fn ticket_table(entry: Layout, entries: usize) -> Result<(Layout, usize), LayoutError> {
if entries == 0 {
return Ok((Layout::from_size_align(0, 16)?, 0));
}
let buckets = if entries < 4 {
4
} else if entries < 8 {
8
} else {
entries
.checked_mul(8)
.and_then(|n| n.checked_add(6))
.map(|n| n / 7)
.and_then(usize::checked_next_power_of_two)
.ok_or_else(layout_error)?
};
let bytes = entry
.pad_to_align()
.size()
.checked_mul(buckets)
.ok_or_else(layout_error)?;
let controls = buckets.checked_add(16).ok_or_else(layout_error)?;
let data = Layout::from_size_align(bytes, entry.align().max(16))?;
let result = data.extend(Layout::from_size_align(controls, 1)?)?.0;
Ok((result, buckets))
}
fn layout_error() -> LayoutError {
Layout::from_size_align(usize::MAX, 2).unwrap_err()
}
pub fn bytes(layout: Layout, count: usize) -> Result<usize, LayoutError> {
layout.size().checked_mul(count).ok_or_else(layout_error)
}
pub fn sum(parts: &[usize]) -> Result<usize, LayoutError> {
parts.iter().try_fold(0usize, |total, part| {
total.checked_add(*part).ok_or_else(layout_error)
})
}
pub fn owned_layouts<O, T, E, I, M>(_: impl FnOnce(I) -> M) -> Result<TaskLayouts, LayoutError>
where
O: Send + 'static,
T: 'static,
E: 'static,
M: for<'a> FnOnce(&'a mut O, RequestTaskContext) -> RequestTaskBorrowedFuture<'a, T, E>
+ Send
+ 'static,
{
pair_layouts(|(owner, make): (Arc<tokio::sync::Mutex<O>>, M)| {
move |context| super::owned_body(owner, make, context)
})
}