tandem-server 0.4.18

HTTP server for Tandem engine APIs
Documentation
use crate::util::time::now_ms;
use std::collections::HashMap;

#[derive(Debug, Default, Clone)]
pub struct ProviderRateLimitStatus {
    pub active_requests: usize,
    pub is_throttled: bool,
    pub throttled_until_ms: Option<u64>,
}

#[derive(Debug, Default)]
pub struct RateLimitManager {
    providers: HashMap<String, ProviderRateLimitStatus>,
}

impl RateLimitManager {
    pub fn new() -> Self {
        Self {
            providers: HashMap::new(),
        }
    }

    pub fn record_throttle(&mut self, provider: &str, retry_after_ms: u64) {
        let status = self.providers.entry(provider.to_string()).or_default();
        status.is_throttled = true;
        status.throttled_until_ms = Some(now_ms() + retry_after_ms);
    }

    pub fn is_provider_throttled(&self, provider: &str) -> bool {
        if let Some(status) = self.providers.get(provider) {
            if status.is_throttled {
                if let Some(until) = status.throttled_until_ms {
                    return now_ms() < until;
                }
                return true;
            }
        }
        false
    }

    pub fn clear_throttle(&mut self, provider: &str) {
        if let Some(status) = self.providers.get_mut(provider) {
            status.is_throttled = false;
            status.throttled_until_ms = None;
        }
    }
}