litellm-rs 0.9.0

//! Single-key Lua deployment circuit state (cluster-safe: `KEYS[1]` only).

use super::pool::{RedisLiveConnection, RedisPool};
use crate::utils::error::gateway_error::{GatewayError, Result};
use std::collections::HashMap;
use std::sync::{Arc, OnceLock};
use std::time::Duration;

const CIRCUIT_OPERATION_TIMEOUT: Duration = Duration::from_secs(5);

const CIRCUIT_SCRIPT: &str = r#"
local op = ARGV[1]
local now = tonumber(ARGV[2]) or 0
local epoch = tonumber(ARGV[3]) or 0
local token = ARGV[4]
local allowed = tonumber(ARGV[5]) or 3
local min_req = tonumber(ARGV[6]) or 10
local cooldown = tonumber(ARGV[7]) or 5
local success_th = tonumber(ARGV[8]) or 3
local reason = tonumber(ARGV[9]) or 0

local function load_state()
  local f = tonumber(redis.call('HGET', KEYS[1], 'f') or '0') or 0
  local r = tonumber(redis.call('HGET', KEYS[1], 'r') or '0') or 0
  local tot = tonumber(redis.call('HGET', KEYS[1], 'tot') or '0') or 0
  local fail = tonumber(redis.call('HGET', KEYS[1], 'fail') or '0') or 0
  local opened = tonumber(redis.call('HGET', KEYS[1], 'opened') or '0') or 0
  local consec = tonumber(redis.call('HGET', KEYS[1], 'consec') or '0') or 0
  local e = tonumber(redis.call('HGET', KEYS[1], 'e') or '0') or 0
  local h = tonumber(redis.call('HGET', KEYS[1], 'h') or '1') or 1
  local owner = redis.call('HGET', KEYS[1], 'owner')
  if not owner then owner = '' end
  local own_until = tonumber(redis.call('HGET', KEYS[1], 'own_until') or '0') or 0
  return f, r, tot, fail, opened, consec, e, h, owner, own_until
end

local function save(f, r, tot, fail, opened, consec, e, h, owner, own_until)
  redis.call('HSET', KEYS[1], 'f', f, 'r', r, 'tot', tot, 'fail', fail,
    'opened', opened, 'consec', consec, 'e', e, 'h', h, 'owner', owner,
    'own_until', own_until)
end

local f, r, tot, fail, opened, consec, e, h, owner, own_until = load_state()
local function claim_owner()
  if owner == '' or own_until <= now then
    consec = 0
    owner = token
    own_until = now + cooldown
    if own_until < now + 1 then own_until = now + 1 end
    h = 2
  end
end
local function clear_owner()
  owner = ''
  own_until = 0
end
if e ~= epoch then
  f = 0
  r = 0
  e = epoch
end

if op == 'fail' then
  tot = tot + 1
  fail = fail + 1
  f = f + 1
  consec = 0
  if opened > now then
    h = 4
  elseif opened > 0 then
    opened = now + cooldown
    clear_owner()
    h = 4
  else
    local trip = 0
    if reason == 1 then
      trip = 1
    elseif reason == 2 then
      if tot >= min_req and (fail * 100 / tot) > 50 then trip = 1 end
    elseif f >= allowed and (r + f) >= min_req then
      trip = 1
    end
    if trip == 1 then
      opened = now + cooldown
      clear_owner()
      h = 4
    elseif h == 1 or h == 0 then
      h = 2
    end
  end
elseif op == 'ok' then
  tot = tot + 1
  r = r + 1
  if opened == 0 or (owner == token and own_until > now) then
    consec = consec + 1
  end
  if opened > now then
    h = 4
  elseif opened > 0 then
    if owner == token and own_until > now then
      if consec >= success_th then
        opened = 0
        clear_owner()
        h = 1
      else
        claim_owner()
      end
    end
  elseif h == 2 and consec >= success_th then
    h = 1
  end
end

local status = 0
local owned = 1
if opened > now then
  status = 1
  owned = 0
elseif opened > 0 then
  status = 2
  claim_owner()
  if owner == token then owned = 1 else owned = 0 end
end

if op ~= 'observe' or opened > 0 then
  save(f, r, tot, fail, opened, consec, e, h, owner, own_until)
end
-- Refresh active history without changing cooldown or probe-owner deadlines.
redis.call('EXPIRE', KEYS[1], math.max(600, opened - now + 1, own_until - now + 1))
return {status, opened, f, consec, h, owned}
"#;

const STATUS_OPEN: i64 = 1;
const STATUS_HALF: i64 = 2;

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct CircuitState {
    pub status: i64,
    pub opened_until: i64,
    pub fails: i64,
    pub consecutive_successes: i64,
    pub health: i64,
    pub owned: bool,
}

impl CircuitState {
    pub(crate) fn blocks_selection(self) -> bool {
        self.status == STATUS_OPEN || (self.status == STATUS_HALF && !self.owned)
    }
}

pub(crate) struct CircuitArgs<'a> {
    pub op: &'a str,
    pub now_secs: i64,
    pub window_epoch: i64,
    pub token: &'a str,
    pub allowed_fails: i64,
    pub min_requests: i64,
    pub cooldown_secs: i64,
    pub success_threshold: i64,
    pub reason: i64,
}

fn circuit_script() -> &'static redis::Script {
    static SCRIPT: OnceLock<redis::Script> = OnceLock::new();
    SCRIPT.get_or_init(|| redis::Script::new(CIRCUIT_SCRIPT))
}

fn circuit_runtime_connections()
-> &'static parking_lot::Mutex<HashMap<String, Arc<RedisLiveConnection>>> {
    static CACHE: OnceLock<parking_lot::Mutex<HashMap<String, Arc<RedisLiveConnection>>>> =
        OnceLock::new();
    CACHE.get_or_init(|| parking_lot::Mutex::new(HashMap::new()))
}

async fn connection_on_current_runtime(pool: &RedisPool) -> Result<Arc<RedisLiveConnection>> {
    let cache_key = format!("{}|{}", pool.config.url, pool.config.cluster);
    {
        let cache = circuit_runtime_connections().lock();
        if let Some(conn) = cache.get(&cache_key) {
            return Ok(conn.clone());
        }
    }
    // Reuse the bounded, non-replaying transport for uncertain routing writes.
    let conn = Arc::new(
        pool.open_budget_connection(CIRCUIT_OPERATION_TIMEOUT)
            .await?,
    );
    let mut cache = circuit_runtime_connections().lock();
    Ok(cache.entry(cache_key).or_insert(conn).clone())
}

fn parse_circuit_state(values: Vec<i64>) -> Result<CircuitState> {
    if values.len() != 6 {
        return Err(GatewayError::Storage(format!(
            "Unexpected Redis circuit result length: {}",
            values.len()
        )));
    }
    Ok(CircuitState {
        status: values[0],
        opened_until: values[1].max(0),
        fails: values[2].max(0),
        consecutive_successes: values[3].max(0),
        health: values[4],
        owned: values[5] == 1,
    })
}

impl RedisPool {
    pub(crate) fn circuit_key(deployment_id: &str) -> String {
        format!("litellm-rs:circuit:v1:{deployment_id}")
    }

    pub(crate) async fn circuit_invoke(
        &self,
        key: &str,
        args: CircuitArgs<'_>,
    ) -> Result<CircuitState> {
        if self.noop_mode {
            return Err(GatewayError::Storage(
                "circuit redis backend is unavailable".to_string(),
            ));
        }

        let mut connection = None;
        let result =
            tokio::time::timeout(CIRCUIT_OPERATION_TIMEOUT, async {
                let _permit =
                    self.semaphore.clone().acquire_owned().await.map_err(|_| {
                        GatewayError::Internal("Redis semaphore closed".to_string())
                    })?;
                let cached = connection_on_current_runtime(self).await?;
                connection = Some(Arc::clone(&cached));
                let mut conn = cached.as_ref().clone();

                let values: Vec<i64> = circuit_script()
                    .key(key)
                    .arg(args.op)
                    .arg(args.now_secs)
                    .arg(args.window_epoch)
                    .arg(args.token)
                    .arg(args.allowed_fails)
                    .arg(args.min_requests)
                    .arg(args.cooldown_secs)
                    .arg(args.success_threshold)
                    .arg(args.reason)
                    .invoke_async(&mut conn)
                    .await
                    .map_err(GatewayError::from)?;
                parse_circuit_state(values)
            })
            .await
            .unwrap_or_else(|_| {
                Err(GatewayError::Unavailable(
                    "Redis circuit operation deadline exceeded".into(),
                ))
            });
        if result.is_err()
            && let Some(connection) = connection
        {
            let key = format!("{}|{}", self.config.url, self.config.cluster);
            let mut cache = circuit_runtime_connections().lock();
            // A late failure from an old connection must preserve a replacement.
            if cache
                .get(&key)
                .is_some_and(|current| Arc::ptr_eq(current, &connection))
            {
                cache.remove(&key);
            }
        }
        result
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::config::models::storage::RedisConfig;

    #[tokio::test]
    async fn expired_probe_owner_cannot_supply_recovery_successes() {
        let Ok(url) = std::env::var("REDIS_URL") else {
            assert!(std::env::var("CI").is_err(), "REDIS_URL is required in CI");
            return;
        };
        let pool = RedisPool::new(&RedisConfig {
            url,
            enabled: true,
            allow_degraded: false,
            ..RedisConfig::default()
        })
        .await
        .unwrap();
        let key = RedisPool::circuit_key(&uuid::Uuid::new_v4().to_string());
        async fn invoke(
            pool: &RedisPool,
            key: &str,
            op: &str,
            now: i64,
            token: &str,
            cooldown: i64,
        ) -> CircuitState {
            // This test owns a short-lived runtime; do not populate the
            // production bridge's process-lifetime connection cache.
            let mut conn = pool.open_live_connection().await.unwrap();
            let values = circuit_script()
                .key(key)
                .arg(op)
                .arg(now)
                .arg(0)
                .arg(token)
                .arg(1)
                .arg(1)
                .arg(cooldown)
                .arg(2)
                .arg(1)
                .invoke_async(&mut conn)
                .await
                .unwrap();
            parse_circuit_state(values).unwrap()
        }
        assert!(
            invoke(&pool, &key, "fail", 100, "a", 10)
                .await
                .blocks_selection()
        );
        assert!(
            !invoke(&pool, &key, "observe", 110, "a", 10)
                .await
                .blocks_selection()
        );
        assert!(
            invoke(&pool, &key, "observe", 119, "b", 10)
                .await
                .blocks_selection()
        );
        assert!(
            !invoke(&pool, &key, "observe", 120, "b", 10)
                .await
                .blocks_selection()
        );
        invoke(&pool, &key, "ok", 121, "a", 10).await;
        let first = invoke(&pool, &key, "ok", 122, "b", 10).await;
        assert_eq!(
            first.status, STATUS_HALF,
            "expired owner's late success must not close the circuit"
        );
        assert_eq!(first.consecutive_successes, 1);
        assert!(
            invoke(&pool, &key, "observe", 122, "a", 10)
                .await
                .blocks_selection()
        );
        assert_eq!(invoke(&pool, &key, "ok", 123, "b", 10).await.status, 0);
        let mut conn = pool.open_live_connection().await.unwrap();
        let history: (i64, i64, i64) = redis::cmd("HMGET")
            .arg(&key)
            .arg(&["tot", "fail", "opened"])
            .query_async(&mut conn)
            .await
            .unwrap();
        assert_eq!(history, (4, 1, 0));
        redis::cmd("PEXPIRE")
            .arg(&key)
            .arg(60_000)
            .query_async::<i64>(&mut conn)
            .await
            .unwrap();
        invoke(&pool, &key, "observe", 123, "a", 10).await;
        let ttl: i64 = redis::cmd("TTL")
            .arg(&key)
            .query_async(&mut conn)
            .await
            .unwrap();
        assert!(
            (570..=600).contains(&ttl),
            "closed observe must refresh idle lifetime: {ttl}"
        );
        let retained: (i64, i64, i64) = redis::cmd("HMGET")
            .arg(&key)
            .arg(&["tot", "fail", "opened"])
            .query_async(&mut conn)
            .await
            .unwrap();
        assert_eq!(
            retained, history,
            "active observations must retain cumulative history"
        );
        pool.delete(&key).await.unwrap();

        let key = RedisPool::circuit_key(&uuid::Uuid::new_v4().to_string());
        let opened = invoke(&pool, &key, "fail", 100, "a", 900).await;
        assert_eq!(opened.opened_until, 1_000);
        let ttl: i64 = redis::cmd("TTL")
            .arg(&key)
            .query_async(&mut conn)
            .await
            .unwrap();
        assert!(
            (871..=901).contains(&ttl),
            "long cooldown must outlive idle TTL: {ttl}"
        );
        let observed = invoke(&pool, &key, "observe", 200, "b", 900).await;
        assert_eq!(
            observed.opened_until, 1_000,
            "observing must not extend the open deadline"
        );
        let ttl: i64 = redis::cmd("TTL")
            .arg(&key)
            .query_async(&mut conn)
            .await
            .unwrap();
        assert!((771..=801).contains(&ttl));
        let probe = invoke(&pool, &key, "observe", 1_000, "a", 900).await;
        assert_eq!(probe.status, STATUS_HALF);
        assert!(probe.owned);
        let ttl: i64 = redis::cmd("TTL")
            .arg(&key)
            .query_async(&mut conn)
            .await
            .unwrap();
        assert!(
            (871..=901).contains(&ttl),
            "live probe owner must outlive idle TTL: {ttl}"
        );
        let foreign = invoke(&pool, &key, "observe", 1_001, "b", 900).await;
        assert!(foreign.blocks_selection());
        let deadlines: (i64, i64) = redis::cmd("HMGET")
            .arg(&key)
            .arg(&["opened", "own_until"])
            .query_async(&mut conn)
            .await
            .unwrap();
        assert_eq!(deadlines, (1_000, 1_900));
        pool.delete(&key).await.unwrap();

        // A closed observation of an unused namespace must not create a hash.
        invoke(&pool, &key, "observe", 2_000, "a", 900).await;
        let exists: bool = redis::cmd("EXISTS")
            .arg(&key)
            .query_async(&mut conn)
            .await
            .unwrap();
        assert!(!exists);
    }

    #[test]
    fn parses_circuit_state() {
        let open = parse_circuit_state(vec![1, 10, 3, 0, 4, 0]).unwrap();
        assert!(open.blocks_selection());
        assert_eq!(open.opened_until, 10);

        let half_owned = parse_circuit_state(vec![2, 5, 3, 1, 2, 1]).unwrap();
        assert!(!half_owned.blocks_selection());
        assert!(half_owned.owned);

        let half_foreign = parse_circuit_state(vec![2, 5, 3, 0, 2, 0]).unwrap();
        assert!(half_foreign.blocks_selection());
        assert!(parse_circuit_state(vec![1, 0]).is_err());
    }

    #[tokio::test]
    async fn noop_redis_pool_fails_closed_on_circuit_observe() {
        let pool = RedisPool::new(&RedisConfig {
            enabled: false,
            ..RedisConfig::default()
        })
        .await
        .expect("disabled Redis should create a no-op pool");

        let err = pool
            .circuit_invoke(
                "litellm-rs:circuit:v1:noop-test",
                CircuitArgs {
                    op: "observe",
                    now_secs: 1,
                    window_epoch: 0,
                    token: "tok",
                    allowed_fails: 3,
                    min_requests: 1,
                    cooldown_secs: 5,
                    success_threshold: 1,
                    reason: 0,
                },
            )
            .await
            .expect_err("no-op Redis must not fail-open circuit observe");
        assert!(err.to_string().contains("unavailable"));
    }
}