use std::sync::{
atomic::{AtomicU64, Ordering},
Arc,
};
use crate::SessionCapabilities;
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
pub struct SessionAdmission {
capabilities: SessionCapabilities,
}
impl SessionAdmission {
pub const fn new(capabilities: SessionCapabilities) -> Self {
Self { capabilities }
}
pub fn validate(self, realized: SessionCapabilities) -> Result<(), SessionAdmissionError> {
if self.capabilities == realized {
Ok(())
} else {
Err(SessionAdmissionError {
admitted: self.capabilities,
realized,
})
}
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq, thiserror::Error)]
#[error("realized session capabilities {realized:?} do not match pre-materialization admission {admitted:?}")]
pub struct SessionAdmissionError {
admitted: SessionCapabilities,
realized: SessionCapabilities,
}
impl SessionAdmissionError {
pub const fn admitted(self) -> SessionCapabilities {
self.admitted
}
pub const fn realized(self) -> SessionCapabilities {
self.realized
}
}
#[derive(Debug)]
pub struct SessionAuthority {
active: Arc<AtomicU64>,
next_ticket: u64,
}
impl Default for SessionAuthority {
fn default() -> Self {
Self::new()
}
}
impl SessionAuthority {
pub fn new() -> Self {
Self {
active: Arc::new(AtomicU64::new(0)),
next_ticket: 1,
}
}
pub fn require_idle(&self) -> Result<(), SessionAuthorityError> {
if self.active.load(Ordering::Acquire) == 0 {
Ok(())
} else {
Err(SessionAuthorityError::Busy)
}
}
pub fn begin_submission(&mut self) -> Result<SubmissionLease, SessionAuthorityError> {
self.require_idle()?;
let ticket = self.next_ticket;
self.next_ticket = ticket
.checked_add(1)
.ok_or(SessionAuthorityError::TicketExhausted)?;
self.active.store(ticket, Ordering::Release);
Ok(SubmissionLease {
owner: Arc::clone(&self.active),
ticket,
})
}
}
#[derive(Debug, Clone, Copy, Eq, PartialEq, thiserror::Error)]
#[non_exhaustive]
pub enum SessionAuthorityError {
#[error("model session already owns an unresolved submission completion")]
Busy,
#[error("model session submission ticket space exhausted")]
TicketExhausted,
}
#[derive(Debug)]
#[must_use = "the submission lease must be retained until native work is safely resolved"]
pub struct SubmissionLease {
owner: Arc<AtomicU64>,
ticket: u64,
}
impl SubmissionLease {
pub fn resolve(&self) -> bool {
self.owner
.compare_exchange(self.ticket, 0, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
}
}
impl Drop for SubmissionLease {
fn drop(&mut self) {
self.resolve();
}
}
#[cfg(test)]
mod tests;