#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CompletionRole {
Startup = 0,
Flush = 1,
Shutdown = 2,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) struct CompletionTicket {
role: CompletionRole,
generation: u64,
}
impl CompletionTicket {
pub(super) const fn initial_startup() -> Self {
Self {
role: CompletionRole::Startup,
generation: 1,
}
}
pub(super) const fn role(self) -> CompletionRole {
self.role
}
pub(super) const fn generation(self) -> u64 {
self.generation
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum CompletionResult<E> {
Pending,
Ready(Result<(), E>),
Stale,
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(super) enum CompletionError {
Busy,
GenerationExhausted,
Stale,
}
#[derive(Clone, Copy)]
enum CellState<E> {
Free,
Pending,
Cancelled,
Ready(Result<(), E>),
}
#[derive(Clone, Copy)]
struct Cell<E> {
generation: u64,
state: CellState<E>,
}
impl<E: Copy> Cell<E> {
const fn free() -> Self {
Self {
generation: 0,
state: CellState::Free,
}
}
}
pub(super) struct CompletionSet<E> {
cells: [Cell<E>; 3],
}
impl<E: Copy> CompletionSet<E> {
pub(super) const fn new() -> Self {
Self {
cells: [Cell::free(); 3],
}
}
pub(super) fn acquire(
&mut self,
role: CompletionRole,
) -> Result<CompletionTicket, CompletionError> {
let cell = &mut self.cells[role as usize];
if !matches!(cell.state, CellState::Free) {
return Err(CompletionError::Busy);
}
cell.generation = cell
.generation
.checked_add(1)
.ok_or(CompletionError::GenerationExhausted)?;
cell.state = CellState::Pending;
Ok(CompletionTicket {
role,
generation: cell.generation,
})
}
pub(super) fn complete(
&mut self,
ticket: CompletionTicket,
result: Result<(), E>,
) -> Result<bool, CompletionError> {
let cell = &mut self.cells[ticket.role as usize];
if cell.generation != ticket.generation {
return Err(CompletionError::Stale);
}
match cell.state {
CellState::Pending => {
cell.state = CellState::Ready(result);
Ok(true)
}
CellState::Cancelled => {
cell.state = CellState::Free;
Ok(false)
}
CellState::Free | CellState::Ready(_) => Err(CompletionError::Stale),
}
}
pub(super) fn result(&self, ticket: CompletionTicket) -> CompletionResult<E> {
let cell = &self.cells[ticket.role as usize];
if cell.generation != ticket.generation {
return CompletionResult::Stale;
}
match cell.state {
CellState::Free => CompletionResult::Stale,
CellState::Pending => CompletionResult::Pending,
CellState::Cancelled => CompletionResult::Stale,
CellState::Ready(result) => CompletionResult::Ready(result),
}
}
pub(super) fn cancel(
&mut self,
ticket: CompletionTicket,
) -> Result<CompletionResult<E>, CompletionError> {
let cell = &mut self.cells[ticket.role as usize];
if cell.generation != ticket.generation {
return Err(CompletionError::Stale);
}
match cell.state {
CellState::Pending => {
cell.state = CellState::Cancelled;
Ok(CompletionResult::Pending)
}
CellState::Ready(result) => Ok(CompletionResult::Ready(result)),
CellState::Free | CellState::Cancelled => Err(CompletionError::Stale),
}
}
pub(super) fn recycle(&mut self, ticket: CompletionTicket) -> Result<(), CompletionError> {
let cell = &mut self.cells[ticket.role as usize];
if cell.generation != ticket.generation || !matches!(cell.state, CellState::Ready(_)) {
return Err(CompletionError::Stale);
}
cell.state = CellState::Free;
Ok(())
}
pub(super) fn fail_pending(&mut self, error: E) -> [Option<CompletionTicket>; 3] {
let mut notify = [None; 3];
for (index, cell) in self.cells.iter_mut().enumerate() {
match cell.state {
CellState::Pending => {
cell.state = CellState::Ready(Err(error));
notify[index] = Some(CompletionTicket {
role: match index {
0 => CompletionRole::Startup,
1 => CompletionRole::Flush,
2 => CompletionRole::Shutdown,
_ => unreachable!(),
},
generation: cell.generation,
});
}
CellState::Cancelled => cell.state = CellState::Free,
CellState::Free | CellState::Ready(_) => {}
}
}
notify
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn fixed_roles_reject_overlap_and_stale_generation() {
let mut set = CompletionSet::<u8>::new();
let first = set.acquire(CompletionRole::Flush).unwrap();
assert_eq!(
set.acquire(CompletionRole::Flush),
Err(CompletionError::Busy)
);
assert_eq!(set.result(first), CompletionResult::Pending);
assert!(set.complete(first, Ok(())).unwrap());
assert_eq!(set.result(first), CompletionResult::Ready(Ok(())));
set.recycle(first).unwrap();
let second = set.acquire(CompletionRole::Flush).unwrap();
assert_ne!(first, second);
assert_eq!(set.result(first), CompletionResult::Stale);
}
#[test]
fn startup_flush_and_shutdown_use_only_three_inline_cells() {
let mut set = CompletionSet::<u8>::new();
for role in [
CompletionRole::Startup,
CompletionRole::Flush,
CompletionRole::Shutdown,
] {
let ticket = set.acquire(role).unwrap();
assert!(set.complete(ticket, Err(7)).unwrap());
assert_eq!(set.result(ticket), CompletionResult::Ready(Err(7)));
}
assert_eq!(
std::mem::size_of_val(&set),
3 * std::mem::size_of::<Cell<u8>>()
);
}
#[test]
fn completion_and_cancellation_have_one_generation_winner() {
let mut set = CompletionSet::<u8>::new();
let cancelled = set.acquire(CompletionRole::Flush).unwrap();
assert_eq!(set.cancel(cancelled).unwrap(), CompletionResult::Pending);
assert!(!set.complete(cancelled, Ok(())).unwrap());
let reused = set.acquire(CompletionRole::Flush).unwrap();
assert_ne!(cancelled, reused);
assert_eq!(set.result(cancelled), CompletionResult::Stale);
assert!(set.complete(reused, Err(9)).unwrap());
assert_eq!(set.cancel(reused).unwrap(), CompletionResult::Ready(Err(9)));
set.recycle(reused).unwrap();
}
}