use chrono::{DateTime, Duration, Utc};
use dashmap::DashMap;
use crate::core::provisioning_stage::ProvisioningStage;
const TERMINAL_TTL_MINUTES: i64 = 10;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ProvisioningLifecycle {
Provisioning,
Ready,
Failed,
}
impl ProvisioningLifecycle {
pub fn wire(&self) -> &'static str {
match self {
ProvisioningLifecycle::Provisioning => "provisioning",
ProvisioningLifecycle::Ready => "ready",
ProvisioningLifecycle::Failed => "failed",
}
}
pub fn is_terminal(&self) -> bool {
!matches!(self, ProvisioningLifecycle::Provisioning)
}
}
#[derive(Debug, Clone)]
pub struct ProvisioningProgress {
pub lifecycle: ProvisioningLifecycle,
pub stage: Option<ProvisioningStage>,
pub detail: Option<String>,
pub session_id: Option<String>,
pub name: Option<String>,
pub error: Option<String>,
pub started_at: DateTime<Utc>,
pub finished_at: Option<DateTime<Utc>>,
}
impl ProvisioningProgress {
fn new_in_flight() -> Self {
Self {
lifecycle: ProvisioningLifecycle::Provisioning,
stage: None,
detail: None,
session_id: None,
name: None,
error: None,
started_at: Utc::now(),
finished_at: None,
}
}
}
#[derive(Debug, Default)]
pub struct ProvisioningRegistry {
entries: DashMap<String, ProvisioningProgress>,
}
impl ProvisioningRegistry {
pub fn begin(&self, job_id: &str) {
self.prune_stale();
self.entries
.insert(job_id.to_string(), ProvisioningProgress::new_in_flight());
}
pub fn update_stage(&self, job_id: &str, stage: ProvisioningStage, detail: Option<String>) {
if let Some(mut e) = self.entries.get_mut(job_id)
&& !e.lifecycle.is_terminal()
{
e.stage = Some(stage);
e.detail = detail;
}
}
pub fn finish_ready(&self, job_id: &str, session_id: String, name: String) {
if let Some(mut e) = self.entries.get_mut(job_id) {
e.lifecycle = ProvisioningLifecycle::Ready;
e.stage = Some(ProvisioningStage::Complete);
e.detail = None;
e.session_id = Some(session_id);
e.name = Some(name);
e.finished_at = Some(Utc::now());
}
}
pub fn finish_failed(&self, job_id: &str, error: String) {
if let Some(mut e) = self.entries.get_mut(job_id) {
e.lifecycle = ProvisioningLifecycle::Failed;
e.error = Some(error);
e.finished_at = Some(Utc::now());
}
}
pub fn get(&self, job_id: &str) -> Option<ProvisioningProgress> {
self.entries.get(job_id).map(|e| e.clone())
}
pub fn prune_stale(&self) {
let cutoff = Utc::now() - Duration::minutes(TERMINAL_TTL_MINUTES);
self.entries.retain(|_, e| match e.finished_at {
Some(finished) => finished >= cutoff,
None => true,
});
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn registry_lifecycle_provisioning_to_ready() {
let reg = ProvisioningRegistry::default();
reg.begin("job-1");
let snap = reg.get("job-1").expect("entry after begin");
assert_eq!(snap.lifecycle, ProvisioningLifecycle::Provisioning);
assert!(snap.stage.is_none());
reg.update_stage(
"job-1",
ProvisioningStage::CloningRepo,
Some("Receiving objects: 42%".into()),
);
let snap = reg.get("job-1").expect("entry after stage");
assert_eq!(snap.stage, Some(ProvisioningStage::CloningRepo));
assert_eq!(snap.detail.as_deref(), Some("Receiving objects: 42%"));
reg.finish_ready("job-1", "sess-9".into(), "tm-repo-01".into());
let snap = reg.get("job-1").expect("entry after ready");
assert_eq!(snap.lifecycle, ProvisioningLifecycle::Ready);
assert_eq!(snap.stage, Some(ProvisioningStage::Complete));
assert_eq!(snap.session_id.as_deref(), Some("sess-9"));
assert_eq!(snap.name.as_deref(), Some("tm-repo-01"));
assert!(snap.detail.is_none());
assert!(snap.finished_at.is_some());
}
#[test]
fn registry_lifecycle_provisioning_to_failed() {
let reg = ProvisioningRegistry::default();
reg.begin("job-2");
reg.finish_failed("job-2", "workspace provisioning failed: boom".into());
let snap = reg.get("job-2").expect("entry after failed");
assert_eq!(snap.lifecycle, ProvisioningLifecycle::Failed);
assert_eq!(
snap.error.as_deref(),
Some("workspace provisioning failed: boom")
);
assert!(snap.finished_at.is_some());
}
#[test]
fn update_stage_ignored_after_terminal() {
let reg = ProvisioningRegistry::default();
reg.begin("job-3");
reg.finish_ready("job-3", "sess".into(), "name".into());
reg.update_stage("job-3", ProvisioningStage::CloningRepo, None);
let snap = reg.get("job-3").expect("entry");
assert_eq!(snap.lifecycle, ProvisioningLifecycle::Ready);
assert_eq!(snap.stage, Some(ProvisioningStage::Complete));
}
#[test]
fn prune_stale_removes_only_old_terminal_entries() {
let reg = ProvisioningRegistry::default();
reg.begin("in-flight");
reg.begin("fresh-done");
reg.finish_ready("fresh-done", "s".into(), "n".into());
reg.begin("old-done");
reg.finish_failed("old-done", "err".into());
if let Some(mut e) = reg.entries.get_mut("old-done") {
e.finished_at = Some(Utc::now() - Duration::minutes(TERMINAL_TTL_MINUTES + 1));
}
reg.prune_stale();
assert!(reg.get("in-flight").is_some(), "in-flight must survive");
assert!(
reg.get("fresh-done").is_some(),
"fresh terminal must survive"
);
assert!(
reg.get("old-done").is_none(),
"stale terminal must be pruned"
);
}
}