use std::sync::atomic::{AtomicU8, Ordering};
#[repr(u8)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd)]
pub(crate) enum ModelState {
Uninit = 0,
Loading = 1,
Ready = 2,
Failed = 3,
}
impl ModelState {
pub(crate) const fn from_u8(v: u8) -> Self {
match v {
1 => Self::Loading,
2 => Self::Ready,
3 => Self::Failed,
_ => Self::Uninit,
}
}
}
pub(crate) struct AtomicModelState(AtomicU8);
impl AtomicModelState {
pub(crate) const fn new(state: ModelState) -> Self {
Self(AtomicU8::new(state as u8))
}
pub(crate) fn load(&self, order: Ordering) -> ModelState {
ModelState::from_u8(self.0.load(order))
}
pub(crate) fn store(&self, state: ModelState, order: Ordering) {
self.0.store(state as u8, order);
}
#[must_use]
pub(crate) fn is_ready(&self) -> bool {
self.load(Ordering::Acquire) == ModelState::Ready
}
pub(crate) fn compare_exchange(
&self,
expected: ModelState,
new: ModelState,
success: Ordering,
failure: Ordering,
) -> Result<ModelState, ModelState> {
self.0
.compare_exchange(expected as u8, new as u8, success, failure)
.map(ModelState::from_u8)
.map_err(ModelState::from_u8)
}
#[must_use]
pub(crate) fn transition(&self, from: ModelState, to: ModelState) -> bool {
self.compare_exchange(from, to, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
}
}
pub(crate) struct ModelLoadGuard<'a>(&'a AtomicModelState);
impl<'a> ModelLoadGuard<'a> {
pub(crate) const fn new(state: &'a AtomicModelState) -> Self {
Self(state)
}
}
impl Drop for ModelLoadGuard<'_> {
fn drop(&mut self) {
self.0
.compare_exchange(
ModelState::Loading,
ModelState::Failed,
Ordering::AcqRel,
Ordering::Acquire,
)
.ok();
}
}
#[cfg(test)]
mod tests {
use super::{AtomicModelState, ModelLoadGuard, ModelState};
use std::sync::atomic::Ordering;
#[test]
fn guard_transitions_loading_to_failed_on_drop() {
let state = AtomicModelState::new(ModelState::Loading);
let guard = ModelLoadGuard::new(&state);
drop(guard);
assert_eq!(state.load(Ordering::Acquire), ModelState::Failed);
}
#[test]
fn guard_drop_is_noop_on_terminal_states() {
let ready = AtomicModelState::new(ModelState::Ready);
drop(ModelLoadGuard::new(&ready));
assert_eq!(ready.load(Ordering::Acquire), ModelState::Ready);
let failed = AtomicModelState::new(ModelState::Failed);
drop(ModelLoadGuard::new(&failed));
assert_eq!(failed.load(Ordering::Acquire), ModelState::Failed);
}
}