use std::collections::BTreeMap;
use std::future::Future;
use std::pin::Pin;
use std::sync::Arc;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use super::{
BusinessCalendar, MemoryObligationStore, Obligation, ObligationRecord, ObligationState,
ObligationStatus, ObligationStore, SlaError,
};
use crate::job::{JobInfo, JobUniqueness, JobUniquenessWindow};
use crate::{AppState, AutumnResult};
pub const CHECK_JOB: &str = "autumn_sla_check";
pub const ESCALATE_JOB: &str = "autumn_sla_escalate";
const ESCALATE_UNIQUE_MS: u64 = 24 * 60 * 60 * 1_000;
const MAX_ATTEMPTS: u32 = 5;
const BACKOFF_MS: u64 = 1_000;
const RELEASE_ATTEMPTS: u32 = 5;
const RELEASE_BACKOFF: std::time::Duration = std::time::Duration::from_millis(100);
type BreachFuture = Pin<Box<dyn Future<Output = AutumnResult<()>> + Send>>;
type BreachHandler = Arc<dyn Fn(AppState, SlaBreach) -> BreachFuture + Send + Sync>;
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct SlaBreach {
pub key: String,
pub obligation: String,
pub subject: String,
pub calendar: String,
pub zone: String,
pub generation: uuid::Uuid,
pub started_at: DateTime<Utc>,
pub due_at: DateTime<Utc>,
pub escalated_at: DateTime<Utc>,
pub token: uuid::Uuid,
}
#[derive(Debug, Serialize, Deserialize)]
struct CheckArgs {
key: String,
#[serde(default)]
generation: Option<uuid::Uuid>,
due_at: DateTime<Utc>,
#[serde(default)]
requeue: u32,
#[serde(default)]
revision: u64,
}
struct SlaEngine {
calendars: BTreeMap<String, BusinessCalendar>,
store: Arc<dyn ObligationStore>,
handlers: BTreeMap<String, BreachHandler>,
fallback: Option<BreachHandler>,
}
pub struct SlaPlugin {
calendars: BTreeMap<String, BusinessCalendar>,
store: Arc<dyn ObligationStore>,
handlers: BTreeMap<String, BreachHandler>,
fallback: Option<BreachHandler>,
}
impl std::fmt::Debug for SlaPlugin {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("SlaPlugin")
.field("calendars", &self.calendars.keys().collect::<Vec<_>>())
.field("handlers", &self.handlers.keys().collect::<Vec<_>>())
.field("fallback", &self.fallback.is_some())
.finish_non_exhaustive()
}
}
impl Default for SlaPlugin {
fn default() -> Self {
Self::new()
}
}
impl SlaPlugin {
#[must_use]
pub fn new() -> Self {
Self {
calendars: BTreeMap::new(),
store: Arc::new(MemoryObligationStore::new()),
handlers: BTreeMap::new(),
fallback: None,
}
}
#[must_use]
pub fn calendar(mut self, name: impl Into<String>, calendar: BusinessCalendar) -> Self {
self.calendars.insert(name.into(), calendar);
self
}
#[must_use]
pub fn store(mut self, store: impl ObligationStore) -> Self {
self.store = Arc::new(store);
self
}
#[must_use]
pub fn on_breach<F, Fut>(mut self, obligation: impl Into<String>, handler: F) -> Self
where
F: Fn(AppState, SlaBreach) -> Fut + Send + Sync + 'static,
Fut: Future<Output = AutumnResult<()>> + Send + 'static,
{
self.handlers.insert(obligation.into(), boxed(handler));
self
}
#[must_use]
pub fn on_any_breach<F, Fut>(mut self, handler: F) -> Self
where
F: Fn(AppState, SlaBreach) -> Fut + Send + Sync + 'static,
Fut: Future<Output = AutumnResult<()>> + Send + 'static,
{
self.fallback = Some(boxed(handler));
self
}
}
fn boxed<F, Fut>(handler: F) -> BreachHandler
where
F: Fn(AppState, SlaBreach) -> Fut + Send + Sync + 'static,
Fut: Future<Output = AutumnResult<()>> + Send + 'static,
{
Arc::new(move |state, breach| Box::pin(handler(state, breach)))
}
impl crate::plugin::Plugin for SlaPlugin {
fn build(self, app: crate::app::AppBuilder) -> crate::app::AppBuilder {
let engine = SlaEngine {
calendars: self.calendars,
store: self.store,
handlers: self.handlers,
fallback: self.fallback,
};
app.state_initializer(move |state| state.insert_extension(engine))
.jobs(vec![
sla_job(
CHECK_JOB,
check_job,
&["key", "generation", "due_at", "requeue", "revision"],
JobUniquenessWindow::Running,
),
sla_job(
ESCALATE_JOB,
escalate_job,
&["key", "generation", "due_at"],
JobUniquenessWindow::TtlMs(ESCALATE_UNIQUE_MS),
),
])
}
}
fn sla_job(
name: &str,
handler: crate::job::JobHandler,
unique_by: &[&str],
window: JobUniquenessWindow,
) -> JobInfo {
let mut info = JobInfo::new(name, MAX_ATTEMPTS, BACKOFF_MS, handler);
info.uniqueness = Some(JobUniqueness {
by: unique_by.iter().map(|field| (*field).to_owned()).collect(),
window,
});
info
}
#[derive(Clone)]
pub struct Sla {
state: AppState,
engine: Arc<SlaEngine>,
}
impl std::fmt::Debug for Sla {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Sla")
.field(
"calendars",
&self.engine.calendars.keys().collect::<Vec<_>>(),
)
.finish_non_exhaustive()
}
}
impl Sla {
pub fn from_state(state: &AppState) -> Result<Self, SlaError> {
let engine = state
.extension::<SlaEngine>()
.ok_or(SlaError::NotInstalled)?;
Ok(Self {
state: state.clone(),
engine,
})
}
#[must_use]
pub fn now(&self) -> DateTime<Utc> {
self.state.clock().now()
}
#[must_use]
pub fn calendar(&self, name: &str) -> Option<&BusinessCalendar> {
self.engine.calendars.get(name)
}
pub fn status(&self, obligation: &Obligation) -> Result<ObligationStatus, SlaError> {
let calendar = self.calendar_of(obligation)?;
let zone = self.zone_of(obligation, calendar);
Ok(obligation.status_with(calendar, zone, self.now()))
}
pub async fn track(&self, obligation: &Obligation) -> Result<ObligationStatus, SlaError> {
let now = self.now();
let calendar = self.calendar_of(obligation)?;
if obligation.budget().resolve(calendar).is_zero() {
return Err(SlaError::InvalidDuration(format!(
"{}: the budget is zero working time; set it with Obligation::within",
obligation.key()
)));
}
let zone = self.zone_of(obligation, calendar);
let resolved = obligation
.clone()
.zone(zone)
.starting_at(obligation.started_at().unwrap_or(now));
let key = resolved.key();
let Some(due_at) =
super::store::real_deadline(resolved.status_with(calendar, zone, now).due_at)
else {
return Err(SlaError::NoDeadline(key));
};
let store = &self.engine.store;
let generation = self.state.entropy().uuid_v4();
let (mut record, created) = store
.insert(ObligationRecord::new(resolved, generation).with_due_at(due_at))
.await?;
if !created && let Some(met) = obligation.met() {
store.mark_met(&key, record.generation, met).await?;
if let Some(fresh) = store.get(&key).await?
&& fresh.generation == record.generation
{
record = fresh;
}
}
self.schedule(&key, &record, now).await
}
async fn schedule(
&self,
key: &str,
record: &ObligationRecord,
now: DateTime<Utc>,
) -> Result<ObligationStatus, SlaError> {
let status = self.status_of(record, now)?;
let Some(due) = status.due_at else {
return Err(SlaError::NoDeadline(key.to_owned()));
};
if status.escalated_at.is_none() && status.state != ObligationState::Met {
self.schedule_check(key, record.generation, due, 0).await?;
}
Ok(status)
}
pub async fn meet(&self, key: &str) -> Result<Option<ObligationStatus>, SlaError> {
let store = &self.engine.store;
if let Some(record) = store.get(key).await? {
store.mark_met(key, record.generation, self.now()).await?;
}
self.get(key).await
}
pub async fn get(&self, key: &str) -> Result<Option<ObligationStatus>, SlaError> {
let now = self.now();
self.engine
.store
.get(key)
.await?
.map(|record| self.status_of(&record, now))
.transpose()
}
pub async fn statuses(&self) -> Result<Vec<ObligationStatus>, SlaError> {
let now = self.now();
self.engine
.store
.list()
.await?
.iter()
.map(|record| self.status_of(record, now))
.collect()
}
pub async fn reconcile(&self) -> Result<usize, SlaError> {
let now = self.now();
let mut scheduled = 0_usize;
let store = &self.engine.store;
for record in store.list().await? {
let status = self.calendar_status(&record, now)?;
if status.escalated_at.is_some() {
continue;
}
let due = super::store::real_deadline(status.due_at);
let stored = due.unwrap_or_else(super::store::no_deadline);
if record.due_at != Some(stored)
&& !store
.set_due(&status.key, record.generation, stored)
.await?
{
continue;
}
let Some(due) = due else { continue };
if status.state == ObligationState::Met {
continue;
}
let revision = self.state.entropy().next_u64().max(1);
self.schedule_check(&status.key, record.generation, due, revision)
.await?;
scheduled = scheduled.saturating_add(1);
}
Ok(scheduled)
}
pub async fn forget(&self, key: &str) -> Result<bool, SlaError> {
self.engine.store.remove(key).await
}
fn calendar_of(&self, obligation: &Obligation) -> Result<&BusinessCalendar, SlaError> {
let name = obligation.calendar_name();
self.calendar(name)
.ok_or_else(|| SlaError::UnknownCalendar(name.to_owned()))
}
fn zone_of(&self, obligation: &Obligation, calendar: &BusinessCalendar) -> chrono_tz::Tz {
obligation
.time_zone()
.or_else(|| calendar.home_zone())
.unwrap_or_else(|| self.state.config_arc().time_zone.default_tz())
}
fn calendar_status(
&self,
record: &ObligationRecord,
now: DateTime<Utc>,
) -> Result<ObligationStatus, SlaError> {
let calendar = self.calendar_of(&record.obligation)?;
let zone = self.zone_of(&record.obligation, calendar);
let mut status = record.obligation.status_with(calendar, zone, now);
status.escalated_at = record.escalated_at;
Ok(status)
}
fn status_of(
&self,
record: &ObligationRecord,
now: DateTime<Utc>,
) -> Result<ObligationStatus, SlaError> {
let Some(due) = record.due_at else {
return self.calendar_status(record, now);
};
let calendar = self.calendar_of(&record.obligation)?;
let zone = self.zone_of(&record.obligation, calendar);
let mut status = if due == super::store::no_deadline() {
record.obligation.status_without_due(calendar, zone, now)
} else {
record.obligation.status_with_due(calendar, zone, now, due)
};
status.escalated_at = record.escalated_at;
Ok(status)
}
async fn schedule_check(
&self,
key: &str,
generation: uuid::Uuid,
due_at: DateTime<Utc>,
revision: u64,
) -> Result<(), SlaError> {
let args = CheckArgs {
key: key.to_owned(),
generation: Some(generation),
due_at,
requeue: 0,
revision,
};
self.enqueue(CHECK_JOB, &args, Some(due_at)).await
}
async fn enqueue<T: Serialize + Sync>(
&self,
name: &str,
args: &T,
due_at: Option<DateTime<Utc>>,
) -> Result<(), SlaError> {
let client = self
.state
.extension::<crate::job::JobClient>()
.ok_or(SlaError::NoJobRuntime)?;
let payload = serde_json::to_value(args).map_err(|err| SlaError::Job(err.to_string()))?;
client
.enqueue_due(name, payload, due_at)
.await
.map_err(|err| SlaError::Job(err.to_string()))
}
async fn requeue_if_early(&self, args: &CheckArgs) -> Result<bool, SlaError> {
let Some(early) = args
.due_at
.signed_duration_since(self.now())
.to_std()
.ok()
.filter(|early| !early.is_zero())
else {
return Ok(false);
};
if args.requeue > 0 {
tracing::warn!(
key = %args.key,
?early,
requeue = args.requeue,
"SLA check ran before its deadline again; the job queue clock leads the app clock"
);
}
let again = CheckArgs {
key: args.key.clone(),
generation: args.generation,
due_at: args.due_at,
requeue: args.requeue.saturating_add(1),
revision: args.revision,
};
let client = self
.state
.extension::<crate::job::JobClient>()
.ok_or(SlaError::NoJobRuntime)?;
let payload = serde_json::to_value(&again).map_err(|err| SlaError::Job(err.to_string()))?;
client
.enqueue_relative(CHECK_JOB, payload, early)
.await
.map_err(|err| SlaError::Job(err.to_string()))?;
Ok(true)
}
async fn check(
&self,
key: &str,
generation: Option<uuid::Uuid>,
due_hint: DateTime<Utc>,
) -> Result<(), SlaError> {
let now = self.now();
let store = &self.engine.store;
let Some(record) = store.get(key).await? else {
tracing::warn!(
%key,
"SLA check found no tracked obligation; replicas need a shared ObligationStore"
);
return Ok(());
};
if generation.is_some_and(|generation| generation != record.generation) {
return Ok(());
}
if record.escalated_at.is_some() {
return Ok(());
}
let status = self.status_of(&record, now)?;
let due_at = match (status.state, status.due_at) {
(ObligationState::Breached, Some(due_at)) => due_at,
(ObligationState::Running | ObligationState::Paused, Some(due_at))
if due_at != due_hint =>
{
return self.schedule_check(key, record.generation, due_at, 0).await;
}
_ => return Ok(()),
};
if !store
.claim_escalation(key, record.generation, due_at, now)
.await?
{
return Ok(());
}
let obligation = &record.obligation;
let breach = SlaBreach {
key: key.to_owned(),
obligation: obligation.name().to_owned(),
subject: obligation.subject().to_owned(),
calendar: obligation.calendar_name().to_owned(),
zone: status.zone.name().to_owned(),
generation: record.generation,
started_at: status.started_at,
due_at,
escalated_at: now,
token: self.state.entropy().uuid_v4(),
};
if let Err(err) = self.enqueue(ESCALATE_JOB, &breach, None).await {
if let Err(release) = release_claim(store.as_ref(), key, record.generation, now).await {
tracing::error!(
%key,
enqueue = %err,
release = %release,
"SLA escalation is claimed but not enqueued; this escalation is lost"
);
}
return Err(err);
}
Ok(())
}
async fn met_on_time(&self, breach: &SlaBreach) -> Result<bool, SlaError> {
let store = &self.engine.store;
let Some(record) = store.get(&breach.key).await? else {
return Ok(false);
};
let on_time = record.generation == breach.generation
&& record.dispatch_token.is_none()
&& record
.obligation
.met()
.is_some_and(|met| met <= breach.due_at);
if on_time && let Some(claimed_at) = record.escalated_at {
store
.release_escalation(&breach.key, breach.generation, claimed_at)
.await?;
}
Ok(on_time)
}
fn handler_for(&self, obligation: &str) -> Option<BreachHandler> {
self.engine
.handlers
.get(obligation)
.or(self.engine.fallback.as_ref())
.cloned()
}
}
fn check_job(state: AppState, payload: Value) -> BreachFuture {
Box::pin(async move {
let args: CheckArgs = serde_json::from_value(payload)?;
let sla = Sla::from_state(&state)?;
if !sla.requeue_if_early(&args).await? {
sla.check(&args.key, args.generation, args.due_at).await?;
}
Ok(())
})
}
fn escalate_job(state: AppState, payload: Value) -> BreachFuture {
Box::pin(async move {
let breach: SlaBreach = serde_json::from_value(payload)?;
let sla = Sla::from_state(&state)?;
if sla.met_on_time(&breach).await? {
return Ok(());
}
if !sla
.engine
.store
.begin_dispatch(&breach.key, breach.generation, breach.due_at, breach.token)
.await?
{
tracing::info!(key = %breach.key, "SLA escalation already dispatched");
return Ok(());
}
let Some(handler) = sla.handler_for(&breach.obligation) else {
tracing::warn!(key = %breach.key, "SLA breach has no handler");
return Ok(());
};
handler(state, breach).await
})
}
impl axum::extract::FromRequestParts<AppState> for Sla {
type Rejection = crate::AutumnError;
async fn from_request_parts(
_parts: &mut axum::http::request::Parts,
state: &AppState,
) -> Result<Self, Self::Rejection> {
Ok(Self::from_state(state)?)
}
}
pub(super) async fn release_claim(
store: &dyn ObligationStore,
key: &str,
generation: uuid::Uuid,
claimed_at: DateTime<Utc>,
) -> Result<(), SlaError> {
let mut delay = RELEASE_BACKOFF;
for attempt in 1..RELEASE_ATTEMPTS {
match store.release_escalation(key, generation, claimed_at).await {
Ok(()) => return Ok(()),
Err(err) => {
tracing::warn!(%key, attempt, error = %err, "SLA claim release failed; trying again");
tokio::time::sleep(delay).await;
delay = delay.saturating_mul(2);
}
}
}
store.release_escalation(key, generation, claimed_at).await
}