af-workflow 0.4.0

//! Provider-local backpressure and circuit breaking.

use std::collections::HashMap;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

use tokio::sync::{OwnedSemaphorePermit, Semaphore};

/// Dispatch priority; higher priorities keep reserved capacity.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum WorkPriority {
    /// Lowest.
    Notification,
    /// Ordinary workflow work.
    Normal,
    /// Funds and control actions.
    Control,
    /// Emergency stop and revocation.
    Emergency,
}

/// Bulkhead, circuit-breaker and backoff limits of a provider.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ProviderLimits {
    /// Total concurrent calls.
    pub concurrency: usize,
    /// Slots kept for control-priority work.
    pub reserved_control_concurrency: usize,
    /// Concurrent calls per tenant.
    pub tenant_concurrency: usize,
    /// Consecutive failures that open the circuit.
    pub failure_threshold: u32,
    /// Open duration before a probe.
    pub recovery_after: Duration,
    /// Per-call deadline.
    pub max_elapsed: Duration,
    /// First backoff.
    pub initial_backoff: Duration,
    /// Backoff cap.
    pub max_backoff: Duration,
}

impl ProviderLimits {
    /// Reject zero or inconsistent limits.
    pub fn validate(&self) -> Result<(), String> {
        if self.concurrency == 0
            || self.tenant_concurrency == 0
            || self.failure_threshold == 0
            || self.recovery_after.is_zero()
            || self.max_elapsed.is_zero()
        {
            return Err("provider limits must be positive".into());
        }
        if self.tenant_concurrency > self.concurrency {
            return Err("tenant concurrency cannot exceed provider concurrency".into());
        }
        if self.reserved_control_concurrency >= self.concurrency {
            return Err("reserved control concurrency must be below provider concurrency".into());
        }
        if self.initial_backoff.is_zero() || self.initial_backoff > self.max_backoff {
            return Err("provider backoff range is invalid".into());
        }
        Ok(())
    }

    /// Exponential backoff with deterministic jitter for `attempt`.
    pub fn backoff(&self, attempt: u32, jitter_seed: u64) -> Duration {
        let factor = 1_u32.checked_shl(attempt.min(20)).unwrap_or(u32::MAX);
        let base = self
            .initial_backoff
            .saturating_mul(factor)
            .min(self.max_backoff);
        let jitter_ceiling = (base.as_millis() / 4).max(1) as u64;
        base.saturating_add(Duration::from_millis(jitter_seed % jitter_ceiling))
            .min(self.max_backoff)
    }
}

/// Why a permit was refused.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ProviderBlock {
    /// The circuit is open.
    CircuitOpen {
        /// Time until a probe is allowed.
        retry_after: Duration,
    },
    /// Every allowed slot is busy.
    Saturated,
}

struct Circuit {
    consecutive_failures: u32,
    opened_at: Option<Instant>,
    probe_in_flight: bool,
}

/// ponytail: one semaphore per tenant is kept in a map; idle entries are swept
/// once the map passes this size so a long tail of one-off tenants cannot grow
/// it without bound. Upgrade path is an LRU if sweeping ever shows up in profiles.
const TENANT_SWEEP_THRESHOLD: usize = 1_024;

/// Bulkhead plus circuit breaker guarding one provider.
pub struct ProviderGate {
    limits: ProviderLimits,
    global: Arc<Semaphore>,
    normal: Arc<Semaphore>,
    tenants: Mutex<HashMap<String, Arc<Semaphore>>>,
    circuit: Arc<Mutex<Circuit>>,
}

/// Holding a permit keeps provider capacity reserved. Dropping it without a
/// `record_*` call releases capacity and clears a recovery probe without
/// touching the failure counter, so a store error or a panic can never leave
/// the circuit stuck open.
pub struct ProviderPermit {
    _global: OwnedSemaphorePermit,
    _normal: Option<OwnedSemaphorePermit>,
    _tenant: Option<OwnedSemaphorePermit>,
    is_probe: bool,
    circuit: Arc<Mutex<Circuit>>,
}

impl Drop for ProviderPermit {
    fn drop(&mut self) {
        if self.is_probe {
            self.circuit
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .probe_in_flight = false;
        }
    }
}

impl ProviderGate {
    /// Gate for validated `limits`.
    pub fn new(limits: ProviderLimits) -> Result<Self, String> {
        limits.validate()?;
        Ok(Self {
            global: Arc::new(Semaphore::new(limits.concurrency)),
            normal: Arc::new(Semaphore::new(
                limits.concurrency - limits.reserved_control_concurrency,
            )),
            limits,
            tenants: Mutex::new(HashMap::new()),
            circuit: Arc::new(Mutex::new(Circuit {
                consecutive_failures: 0,
                opened_at: None,
                probe_in_flight: false,
            })),
        })
    }

    /// Acquire a permit for `tenant` at `priority`.
    pub fn try_acquire(
        &self,
        tenant_id: &str,
        priority: WorkPriority,
        now: Instant,
    ) -> Result<ProviderPermit, ProviderBlock> {
        let is_probe = {
            let mut circuit = self
                .circuit
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            match circuit.opened_at {
                None => false,
                Some(opened_at) => {
                    let elapsed = now.saturating_duration_since(opened_at);
                    if elapsed < self.limits.recovery_after || circuit.probe_in_flight {
                        return Err(ProviderBlock::CircuitOpen {
                            retry_after: self.limits.recovery_after.saturating_sub(elapsed),
                        });
                    }
                    circuit.probe_in_flight = true;
                    true
                }
            }
        };
        let normal = if matches!(priority, WorkPriority::Notification | WorkPriority::Normal) {
            match self.normal.clone().try_acquire_owned() {
                Ok(permit) => Some(permit),
                Err(_) => {
                    self.release_probe(is_probe);
                    return Err(ProviderBlock::Saturated);
                }
            }
        } else {
            None
        };
        let global = match self.global.clone().try_acquire_owned() {
            Ok(permit) => permit,
            Err(_) => {
                self.release_probe(is_probe);
                return Err(ProviderBlock::Saturated);
            }
        };
        let tenant = if priority == WorkPriority::Emergency {
            None
        } else {
            let tenant = {
                let mut tenants = self
                    .tenants
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner);
                if tenants.len() >= TENANT_SWEEP_THRESHOLD {
                    let full = self.limits.tenant_concurrency;
                    tenants.retain(|_, semaphore| semaphore.available_permits() < full);
                }
                tenants
                    .entry(tenant_id.to_owned())
                    .or_insert_with(|| Arc::new(Semaphore::new(self.limits.tenant_concurrency)))
                    .clone()
            };
            match tenant.try_acquire_owned() {
                Ok(permit) => Some(permit),
                Err(_) => {
                    self.release_probe(is_probe);
                    return Err(ProviderBlock::Saturated);
                }
            }
        };
        Ok(ProviderPermit {
            _global: global,
            _normal: normal,
            _tenant: tenant,
            is_probe,
            circuit: self.circuit.clone(),
        })
    }

    /// Return a permit after success; closes a half-open circuit.
    pub fn record_success(&self, permit: ProviderPermit) {
        let mut circuit = self
            .circuit
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        circuit.consecutive_failures = 0;
        circuit.opened_at = None;
        drop(circuit);
        drop(permit);
    }

    /// Return a permit after failure; may open the circuit.
    pub fn record_failure(&self, permit: ProviderPermit, now: Instant) {
        let mut circuit = self
            .circuit
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        circuit.consecutive_failures = circuit.consecutive_failures.saturating_add(1);
        if permit.is_probe || circuit.consecutive_failures >= self.limits.failure_threshold {
            circuit.opened_at = Some(now);
        }
        drop(circuit);
        drop(permit);
    }

    /// Give capacity back without judging the provider. Use this when the
    /// caller failed before or after the provider ran (store conflict, stale
    /// claim) so infrastructure errors do not trip the provider's circuit.
    pub fn release(&self, permit: ProviderPermit) {
        drop(permit);
    }

    fn release_probe(&self, is_probe: bool) {
        if is_probe {
            self.circuit
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .probe_in_flight = false;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn limits() -> ProviderLimits {
        ProviderLimits {
            concurrency: 2,
            reserved_control_concurrency: 1,
            tenant_concurrency: 1,
            failure_threshold: 2,
            recovery_after: Duration::from_secs(10),
            max_elapsed: Duration::from_secs(30),
            initial_backoff: Duration::from_millis(100),
            max_backoff: Duration::from_secs(1),
        }
    }

    #[test]
    fn tenant_bulkhead_and_circuit_recovery_are_bounded() {
        let gate = ProviderGate::new(limits()).unwrap();
        let now = Instant::now();
        let first = gate.try_acquire("one", WorkPriority::Normal, now).unwrap();
        let emergency = gate
            .try_acquire("one", WorkPriority::Emergency, now)
            .expect("reserved control capacity must remain available");
        gate.record_success(emergency);
        assert_eq!(
            gate.try_acquire("two", WorkPriority::Normal, now).err(),
            Some(ProviderBlock::Saturated)
        );
        gate.record_failure(first, now);
        let second = gate.try_acquire("two", WorkPriority::Normal, now).unwrap();
        gate.record_failure(second, now);
        assert!(matches!(
            gate.try_acquire("one", WorkPriority::Normal, now),
            Err(ProviderBlock::CircuitOpen { .. })
        ));
        let probe = gate
            .try_acquire("one", WorkPriority::Normal, now + Duration::from_secs(10))
            .unwrap();
        gate.record_success(probe);
        assert!(gate
            .try_acquire("one", WorkPriority::Normal, now + Duration::from_secs(10))
            .is_ok());
    }

    #[test]
    fn dropped_probe_permit_does_not_wedge_the_circuit() {
        let gate = ProviderGate::new(limits()).unwrap();
        let now = Instant::now();
        for tenant in ["one", "two"] {
            let permit = gate.try_acquire(tenant, WorkPriority::Normal, now).unwrap();
            gate.record_failure(permit, now);
        }
        let later = now + Duration::from_secs(10);
        let probe = gate
            .try_acquire("one", WorkPriority::Normal, later)
            .unwrap();
        assert!(matches!(
            gate.try_acquire("two", WorkPriority::Normal, later),
            Err(ProviderBlock::CircuitOpen { .. })
        ));
        gate.release(probe);
        let probe = gate
            .try_acquire("two", WorkPriority::Normal, later)
            .expect("released probe must allow the next probe");
        gate.record_success(probe);
        assert!(gate.try_acquire("one", WorkPriority::Normal, later).is_ok());
    }

    #[test]
    fn tenant_map_is_swept_once_it_grows() {
        let gate = ProviderGate::new(ProviderLimits {
            concurrency: 4,
            tenant_concurrency: 1,
            ..limits()
        })
        .unwrap();
        let now = Instant::now();
        for index in 0..TENANT_SWEEP_THRESHOLD + 5 {
            let permit = gate
                .try_acquire(&format!("tenant-{index}"), WorkPriority::Normal, now)
                .unwrap();
            gate.record_success(permit);
        }
        assert!(gate.tenants.lock().unwrap().len() <= TENANT_SWEEP_THRESHOLD + 1);
    }

    #[test]
    fn backoff_is_capped() {
        let limits = limits();
        assert_eq!(limits.backoff(30, 99), Duration::from_secs(1));
    }
}