use std::sync::atomic::{AtomicU8, AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
const STATE_WARMING: u8 = 0;
const STATE_READY: u8 = 1;
const STATE_DEGRADED: u8 = 2;
const STATE_FAILED: u8 = 3;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum SubsystemState {
Warming,
Ready,
Degraded,
Failed,
}
impl SubsystemState {
fn to_u8(self) -> u8 {
match self {
Self::Warming => STATE_WARMING,
Self::Ready => STATE_READY,
Self::Degraded => STATE_DEGRADED,
Self::Failed => STATE_FAILED,
}
}
fn from_u8(v: u8) -> Self {
match v {
STATE_READY => Self::Ready,
STATE_DEGRADED => Self::Degraded,
STATE_FAILED => Self::Failed,
_ => Self::Warming,
}
}
}
use serde::{Deserialize, Serialize};
pub struct ReadinessGate {
state_store: AtomicU8,
engine: AtomicU8,
deadline_secs: AtomicU64,
}
impl ReadinessGate {
pub fn new(deadline_secs: u64) -> Self {
Self {
state_store: AtomicU8::new(STATE_WARMING),
engine: AtomicU8::new(STATE_WARMING),
deadline_secs: AtomicU64::new(deadline_secs),
}
}
pub fn set_deadline_secs(&self, secs: u64) {
self.deadline_secs.store(secs, Ordering::SeqCst);
}
pub fn deadline_secs(&self) -> u64 {
self.deadline_secs.load(Ordering::SeqCst)
}
pub fn set_state_store(&self, s: SubsystemState) {
self.state_store.store(s.to_u8(), Ordering::SeqCst);
}
pub fn set_engine(&self, s: SubsystemState) {
self.engine.store(s.to_u8(), Ordering::SeqCst);
}
pub fn state_store_state(&self) -> SubsystemState {
SubsystemState::from_u8(self.state_store.load(Ordering::SeqCst))
}
pub fn engine_state(&self) -> SubsystemState {
SubsystemState::from_u8(self.engine.load(Ordering::SeqCst))
}
pub fn is_ready(&self) -> bool {
let s = self.state_store_state();
let e = self.engine_state();
let s_ok = s == SubsystemState::Ready || s == SubsystemState::Degraded;
let e_ok = e == SubsystemState::Ready || e == SubsystemState::Degraded;
s_ok && e_ok
}
pub fn enforce_deadline(&self) {
let deadline = self.deadline_secs.load(Ordering::SeqCst);
if deadline == 0 {
return;
}
if self.is_ready() {
return;
}
let now = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs();
if now < deadline {
return;
}
if self.state_store_state() == SubsystemState::Warming {
self.set_state_store(SubsystemState::Degraded);
}
if self.engine_state() == SubsystemState::Warming {
self.set_engine(SubsystemState::Degraded);
}
}
}
impl std::fmt::Debug for ReadinessGate {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ReadinessGate")
.field("state_store", &self.state_store_state())
.field("engine", &self.engine_state())
.field("is_ready", &self.is_ready())
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn new_starts_warming_and_not_ready() {
let g = ReadinessGate::new(0);
assert!(!g.is_ready());
assert_eq!(g.state_store_state(), SubsystemState::Warming);
assert_eq!(g.engine_state(), SubsystemState::Warming);
}
#[test]
fn both_ready_means_ready() {
let g = ReadinessGate::new(0);
g.set_state_store(SubsystemState::Ready);
g.set_engine(SubsystemState::Ready);
assert!(g.is_ready());
}
#[test]
fn engine_degraded_still_counts_as_ready() {
let g = ReadinessGate::new(0);
g.set_state_store(SubsystemState::Ready);
g.set_engine(SubsystemState::Degraded);
assert!(g.is_ready());
}
#[test]
fn engine_failed_keeps_gate_closed() {
let g = ReadinessGate::new(0);
g.set_state_store(SubsystemState::Ready);
g.set_engine(SubsystemState::Failed);
assert!(!g.is_ready());
}
#[test]
fn state_store_not_ready_keeps_gate_closed() {
let g = ReadinessGate::new(0);
g.set_engine(SubsystemState::Ready);
assert!(!g.is_ready());
}
#[test]
fn deadline_elapsed_promotes_warming_to_degraded() {
let g = ReadinessGate::new(1);
std::thread::sleep(std::time::Duration::from_millis(1100));
g.enforce_deadline();
assert_eq!(g.state_store_state(), SubsystemState::Degraded);
assert_eq!(g.engine_state(), SubsystemState::Degraded);
assert!(g.is_ready());
}
#[test]
fn deadline_not_yet_elapsed_keeps_warming() {
let deadline = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_secs()
+ 60;
let g = ReadinessGate::new(deadline);
g.enforce_deadline();
assert_eq!(g.state_store_state(), SubsystemState::Warming);
assert!(!g.is_ready());
}
#[test]
fn deadline_zero_disables_enforcement() {
let g = ReadinessGate::new(0);
g.enforce_deadline();
assert_eq!(g.state_store_state(), SubsystemState::Warming);
}
}