use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum WorkerState {
Starting,
Ready,
Busy,
Stalled,
Dead,
Stopped,
}
impl WorkerState {
pub fn is_terminal(&self) -> bool {
matches!(self, WorkerState::Dead | WorkerState::Stopped)
}
pub fn is_available(&self) -> bool {
matches!(self, WorkerState::Ready)
}
}
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct WorkerStateMap {
states: HashMap<String, WorkerState>,
}
impl WorkerStateMap {
pub fn get(&self, worker_id: &str) -> WorkerState {
self.states
.get(worker_id)
.copied()
.unwrap_or(WorkerState::Starting)
}
pub fn set(&mut self, worker_id: &str, state: WorkerState) -> Option<WorkerState> {
self.states.insert(worker_id.to_string(), state)
}
pub fn transition(&mut self, worker_id: &str, new_state: WorkerState) -> Option<WorkerState> {
let old = self.get(worker_id);
if old != new_state {
self.set(worker_id, new_state);
Some(old)
} else {
None
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn worker_state_terminal() {
assert!(WorkerState::Dead.is_terminal());
assert!(WorkerState::Stopped.is_terminal());
assert!(!WorkerState::Starting.is_terminal());
assert!(!WorkerState::Ready.is_terminal());
assert!(!WorkerState::Busy.is_terminal());
assert!(!WorkerState::Stalled.is_terminal());
}
#[test]
fn worker_state_available() {
assert!(WorkerState::Ready.is_available());
assert!(!WorkerState::Starting.is_available());
assert!(!WorkerState::Busy.is_available());
assert!(!WorkerState::Stalled.is_available());
assert!(!WorkerState::Dead.is_available());
assert!(!WorkerState::Stopped.is_available());
}
#[test]
fn worker_state_map_transition() {
let mut map = WorkerStateMap::default();
assert_eq!(map.get("w1"), WorkerState::Starting);
let old = map.transition("w1", WorkerState::Ready);
assert_eq!(old, Some(WorkerState::Starting));
assert_eq!(map.get("w1"), WorkerState::Ready);
let old = map.transition("w1", WorkerState::Ready);
assert_eq!(old, None);
}
}