#架构说明

// key/lease.rs

pub struct KeyLease {
    inner: Arc<KeyInner>,
    reserved_tokens: u32,
}

struct KeyInner {
    key: String,
    provider: ProviderId,
    tpm_inflight: AtomicU32,
    tpm_limit: u32,
    // STATE_HEALTHY = 0 | STATE_DEAD = 1
    state: AtomicU8,
    // --- Cooldown (rate-limit driven) ---
    // Unix timestamp millis; 0 = not cooling down.
    // Set by 429 responses. Checked lazily in acquire().
    cool_down_until: AtomicU64,
    // --- Circuit breaker (error driven, independent of cooldown) ---
    // Consecutive non-rate-limit failures (5xx, network, timeout).
    // Reset to 0 on any success. Triggers its own cooldown at threshold.
    // Decoupled from cool_down_until so that a key recovering from a 429
    // cooldown isn't immediately killed by a stale failure count.
    consecutive_failures: AtomicU32,
    failure_cool_down_until: AtomicU64,
}

impl Drop for KeyLease {
    fn drop(&mut self) {
        // Unconditional — runs on success, error, panic, or async cancel
        self.inner
            .tpm_inflight
            .fetch_sub(self.reserved_tokens, Ordering::Release);
    }
}

// key/pool.rs — acquire

const MAX_CAS_ATTEMPTS: usize = 5;

pub fn acquire(&self, estimated_tokens: u32) -> Option<KeyLease> {
    let now_ms = current_millis();
    let start = rand::random::<usize>() % self.keys.len();
    let candidates = self.keys[start..]
        .iter()
        .chain(self.keys[..start].iter())
        .filter(|k| k.state.load(Ordering::Acquire) == STATE_HEALTHY)
        .filter(|k| {
            let rl = k.cool_down_until.load(Ordering::Acquire);
            let cb = k.failure_cool_down_until.load(Ordering::Acquire);
            (rl == 0 || now_ms >= rl) && (cb == 0 || now_ms >= cb)
        });

    let mut attempts = 0;
    for key in candidates {
        if attempts >= MAX_CAS_ATTEMPTS {
            break; // hard cap — avoid long-tail spin
        }
        attempts += 1;

        let cur = key.tpm_inflight.load(Ordering::Relaxed);
        if cur + estimated_tokens > key.tpm_limit {
            continue;
        }

        // Single CAS attempt per key — on failure, skip to next.
        let result = key.tpm_inflight.compare_exchange_weak(
            cur,
            cur + estimated_tokens,
            Ordering::AcqRel,
            Ordering::Relaxed,
        );
        if result.is_ok() {
            return Some(KeyLease {
                inner: Arc::clone(key),
                reserved_tokens: estimated_tokens,
            });
        }
        std::hint::spin_loop();
    }
    None
}

// key/pool.rs — error reporting + circuit breaker

const CIRCUIT_BREAKER_THRESHOLD: u32 = 5;
const CIRCUIT_BREAKER_COOLDOWN_MS: u64 = 30_000; // 30 seconds

pub fn report_error(&self, lease: &KeyLease, err: &ApiError) {
    match err {
        ApiError::Unauthorized => {
            lease.inner.state.store(STATE_DEAD, Ordering::Release);
        }
        ApiError::RateLimited { retry_after } => {
            // Rate-limit cooldown — independent of circuit breaker.
            // A 429 is not an "error"; it's the provider saying "slow down".
            // Does NOT increment consecutive_failures.
            let until = current_millis() + retry_after.as_millis() as u64;
            lease.inner.cool_down_until.fetch_max(until, Ordering::AcqRel);
        }
        ApiError::ServerError(_) | ApiError::NetworkError(_) => {
            // Circuit breaker: writes to failure_cool_down_until, NOT cool_down_until.
            // This decoupling prevents the scenario where a key recovers from a
            // 429 cooldown, gets one 5xx (which would be failure #5 from before
            // the cooldown), and is immediately killed again.
            let prev = lease.inner.consecutive_failures
                .fetch_add(1, Ordering::Relaxed);
            if prev + 1 >= CIRCUIT_BREAKER_THRESHOLD {
                let until = current_millis() + CIRCUIT_BREAKER_COOLDOWN_MS;
                lease.inner.failure_cool_down_until.store(until, Ordering::Release);
                lease.inner.consecutive_failures.store(0, Ordering::Relaxed);
            }
        }
        _ => {}
    }
}

pub fn report_success(&self, lease: &KeyLease) {
    lease.inner.consecutive_failures.store(0, Ordering::Relaxed);
}

// limiter/window.rs

pub struct SlidingWindow {
    window: Duration,
    limit: u32,
    timestamps: Mutex<VecDeque<Instant>>,
}

impl SlidingWindow {
    pub fn try_acquire(&self) -> bool {
        let now = Instant::now();
        let mut ts = self.timestamps.lock().unwrap();

        // Evict entries that have fallen outside the window
        let cutoff = now - self.window;
        while ts.front().map_or(false, |t| *t < cutoff) {
            ts.pop_front();
        }

        if ts.len() < self.limit as usize {
            ts.push_back(now);
            true
        } else {
            false
        }
    }
}

// budget/tracker.rs

/// 1 USD = 1_000_000 micro-dollars
pub type MicroDollar = u64;

pub struct BudgetTracker {
    limit: MicroDollar,
    used: AtomicU64,
}

impl BudgetTracker {
    pub fn try_reserve(&self, estimated: MicroDollar) -> bool {
        loop {
            let cur = self.used.load(Ordering::Acquire);
            if cur + estimated > self.limit {
                return false;
            }
            if self.used
                .compare_exchange_weak(
                    cur,
                    cur + estimated,
                    Ordering::AcqRel,
                    Ordering::Relaxed,
                )
                .is_ok()
            {
                return true;
            }
        }
    }

    /// Called after response: correct estimated → actual delta
    pub fn settle(&self, estimated: MicroDollar, actual: MicroDollar) {
        if actual > estimated {
            self.used.fetch_add(actual - estimated, Ordering::Relaxed);
        } else {
            self.used.fetch_sub(estimated - actual, Ordering::Relaxed);
        }
    }
}

// gateway/mod.rs

pub async fn call(&self, req: LlmRequest, cancel: CancellationToken) -> Result<LlmResponse> {
    let est_tokens = req.estimated_tokens();
    let est_cost   = pricing::estimate(est_tokens, &req.model);

    // 1. Acquire key (select + CAS reserve in one step)
    let lease = self.pool.acquire(est_tokens)
        .ok_or(Error::NoAvailableKey)?;

    // 2. Pre-occupy budget
    if !self.budget.try_reserve(est_cost) {
        return Err(Error::BudgetExceeded);
    }

    // 3. Rate limit check
    if !lease.inner.rpm_window.try_acquire() {
        self.budget.settle(est_cost, 0);
        return Err(Error::RateLimited);
    }

    // 4. Execute with cancellation propagation.
    //    If upstream cancels, we abort the HTTP request — not just drop
    //    the future. This prevents phantom inflight: the provider stops
    //    processing, and our TPM counter reflects reality.
    let result = tokio::select! {
        res = self.dispatcher.call(&lease, &req) => res,
        _ = cancel.cancelled() => Err(ApiError::Cancelled),
    };

    // 5. Accounting — always runs before lease drops
    match &result {
        Ok(resp) => {
            let actual = pricing::actual(&resp.usage, &req.model);
            self.budget.settle(est_cost, actual);
            self.pool.report_success(&lease);
        }
        Err(ApiError::Cancelled) => {
            self.budget.settle(est_cost, 0);
            // Not a key failure — don't touch circuit breaker
        }
        Err(e) => {
            self.budget.settle(est_cost, 0);
            self.pool.report_error(&lease, e);
        }
    }

    // lease drops here → TPM returned ✓
    result
}

// key/registry.rs

pub struct PoolRegistry {
    pools: HashMap<(ProviderId, ModelId), KeyPool>,
}

impl PoolRegistry {
    pub fn acquire(
        &self,
        provider: &ProviderId,
        model: &ModelId,
        estimated_tokens: u32,
    ) -> Option<KeyLease> {
        self.pools
            .get(&(provider.clone(), model.clone()))?
            .acquire(estimated_tokens)
    }
}

let lease = self.registry.acquire(&req.provider, &req.model, est_tokens)
    .ok_or(Error::NoAvailableKey)?;

// Per-key metrics (all atomic, zero-cost read)
struct KeyMetrics {
    tpm_inflight:       u32,     // current TPM occupation
    tpm_limit:          u32,     // configured cap
    rpm_remaining:      u32,     // SlidingWindow::remaining()
    cool_down_until:    u64,     // 0 = healthy
    failure_cool_down:  u64,     // 0 = healthy
    consecutive_fails:  u32,     // circuit breaker counter
    state:              u8,      // HEALTHY / DEAD
}