fix: restore provider key circuit breaker backoff

2026-10-07 18:07:47 +08:00 · 2026-05-19 13:57:36 +08:00
parent 6a1da6a5ff
commit 7e9ca88e00
14 changed files with 525 additions and 42 deletions
@@ -19,7 +19,8 @@ use tracing::warn;
 use super::{
    local_failover_error_message, project_local_adaptive_rate_limit,
    project_local_adaptive_success, project_local_failure_health, project_local_key_circuit_closed,
-    project_local_key_circuit_open, project_local_success_health, LocalFailoverClassification,
+    project_local_key_circuit_failure, project_local_key_circuit_open,
+    project_local_success_health, LocalFailoverClassification,
 };
 use crate::ai_serving::extract_pool_sticky_session_token;
 use crate::client_session_affinity::{
@@ -521,21 +522,38 @@ async fn record_health_failure_effect(
    else {
        return;
    };
+    let observed_at_unix_secs = current_unix_secs();
    let Some(health_by_format) = project_local_failure_health(
        current_key.health_by_format.as_ref(),
        api_format,
        effect.classification,
        effect.status_code,
-        current_unix_secs(),
+        observed_at_unix_secs,
    ) else {
        return;
    };
+    let consecutive_failures = health_by_format
+        .get(api_format)
+        .and_then(|value| value.get("consecutive_failures"))
+        .and_then(Value::as_u64)
+        .unwrap_or(0);
+    let circuit_breaker_by_format = project_local_key_circuit_failure(
+        current_key.circuit_breaker_by_format.as_ref(),
+        api_format,
+        observed_at_unix_secs,
+        consecutive_failures,
+        current_key.max_probe_interval_minutes,
+    );
+    let circuit_breaker_update = circuit_breaker_by_format
+        .as_ref()
+        .or(current_key.circuit_breaker_by_format.as_ref());

    if let Err(err) = state
-        .update_provider_catalog_key_format_health(
+        .update_provider_catalog_key_health_state(
            &context.plan.key_id,
-            api_format,
-            &health_by_format,
+            current_key.is_active,
+            Some(&health_by_format),
+            circuit_breaker_update,
        )
        .await
    {
@@ -705,6 +723,7 @@ async fn open_pool_key_circuit_breaker(
        api_format,
        reason,
        current_unix_secs(),
+        current_key.max_probe_interval_minutes,
    ) else {
        return;
    };
@@ -1831,6 +1850,51 @@ mod tests {
        );
    }

+    #[tokio::test]
+    async fn health_failure_opens_circuit_after_eight_consecutive_failures() {
+        let state = health_state();
+        let plan = sample_plan();
+
+        for _ in 0..8 {
+            apply_local_execution_effect(
+                &state,
+                LocalExecutionEffectContext {
+                    plan: &plan,
+                    report_context: None,
+                },
+                LocalExecutionEffect::HealthFailure(LocalHealthFailureEffect {
+                    status_code: 503,
+                    classification: LocalFailoverClassification::RetryUpstreamFailure,
+                }),
+            )
+            .await;
+        }
+
+        let stored_key = state
+            .read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
+            .await
+            .expect("provider catalog keys should load")
+            .into_iter()
+            .next()
+            .expect("stored key should exist");
+        let circuit = stored_key
+            .circuit_breaker_by_format
+            .as_ref()
+            .and_then(|value| value.get("openai:chat"))
+            .expect("format circuit should be stored");
+        assert_eq!(circuit["open"], json!(true));
+        assert_eq!(circuit["reason"], json!("consecutive_failures_8"));
+        assert_eq!(circuit["probe_interval_minutes"], json!(1));
+        assert!(circuit["next_probe_at_unix_secs"].as_u64().is_some());
+        assert_eq!(
+            circuit["request_results_window"]
+                .as_array()
+                .map(Vec::len)
+                .unwrap_or_default(),
+            8
+        );
+    }
+
    #[tokio::test]
    async fn health_success_projection_resets_key_health_for_format() {
        let state = health_state();
@@ -4,7 +4,8 @@ use super::LocalFailoverClassification;
 use crate::handlers::shared::unix_secs_to_rfc3339;

 const LOCAL_HEALTH_SCORE_FLOOR: f64 = 0.2;
-const LOCAL_KEY_CIRCUIT_PROBE_DELAY_SECS: u64 = 32 * 60;
+pub(crate) const LOCAL_KEY_CIRCUIT_FAILURE_THRESHOLD: u64 = 8;
+const LOCAL_KEY_CIRCUIT_MAX_PROBE_INTERVAL_MINUTES: u64 = 32;

 pub(crate) fn project_local_failure_health(
    current_health_by_format: Option<&Value>,
@@ -79,6 +80,7 @@ pub(crate) fn project_local_key_circuit_open(
    api_format: &str,
    reason: &str,
    observed_at_unix_secs: u64,
+    max_probe_interval_minutes: i32,
 ) -> Option<Value> {
    let api_format = api_format.trim();
    let reason = reason.trim();
@@ -86,23 +88,143 @@ pub(crate) fn project_local_key_circuit_open(
        return None;
    }

-    let next_probe_at_unix_secs =
-        observed_at_unix_secs.saturating_add(LOCAL_KEY_CIRCUIT_PROBE_DELAY_SECS);
    let mut circuit_by_format = current_circuit_by_format
        .and_then(Value::as_object)
        .cloned()
        .unwrap_or_default();
+    let current = circuit_by_format
+        .get(api_format)
+        .and_then(Value::as_object)
+        .cloned()
+        .unwrap_or_default();
+    let max_probe_interval_minutes =
+        normalize_max_probe_interval_minutes(max_probe_interval_minutes);
+    let probe_interval_minutes =
+        next_circuit_probe_interval_minutes(&current, max_probe_interval_minutes);
+    let next_probe_at_unix_secs =
+        next_probe_at_unix_secs(observed_at_unix_secs, probe_interval_minutes);
+    let open_at = current
+        .get("open_at")
+        .filter(|_| current_bool(&current, "open"))
+        .cloned()
+        .unwrap_or_else(|| json!(unix_secs_to_rfc3339(observed_at_unix_secs)));
+    let half_open_failures = current
+        .get("half_open_failures")
+        .and_then(Value::as_u64)
+        .unwrap_or(0)
+        .saturating_add(u64::from(current_bool(&current, "open")));
+    let request_results_window =
+        append_request_result_window(&current, observed_at_unix_secs, false);
    circuit_by_format.insert(
        api_format.to_string(),
        json!({
            "open": true,
-            "open_at": unix_secs_to_rfc3339(observed_at_unix_secs),
+            "open_at": open_at,
            "reason": reason,
            "next_probe_at": unix_secs_to_rfc3339(next_probe_at_unix_secs),
            "next_probe_at_unix_secs": next_probe_at_unix_secs,
+            "probe_interval_minutes": probe_interval_minutes,
+            "max_probe_interval_minutes": max_probe_interval_minutes,
+            "last_failure_at": unix_secs_to_rfc3339(observed_at_unix_secs),
+            "last_probe_failure_at": if half_open_failures > 0 {
+                json!(unix_secs_to_rfc3339(observed_at_unix_secs))
+            } else {
+                Value::Null
+            },
            "half_open_until": Value::Null,
            "half_open_successes": 0,
-            "half_open_failures": 0,
+            "half_open_failures": half_open_failures,
+            "request_results_window": request_results_window,
+        }),
+    );
+    Some(Value::Object(circuit_by_format))
+}
+
+pub(crate) fn project_local_key_circuit_failure(
+    current_circuit_by_format: Option<&Value>,
+    api_format: &str,
+    observed_at_unix_secs: u64,
+    consecutive_failures: u64,
+    max_probe_interval_minutes: i32,
+) -> Option<Value> {
+    let api_format = api_format.trim();
+    if api_format.is_empty() {
+        return None;
+    }
+
+    let mut circuit_by_format = current_circuit_by_format
+        .and_then(Value::as_object)
+        .cloned()
+        .unwrap_or_default();
+    let current = circuit_by_format
+        .get(api_format)
+        .and_then(Value::as_object)
+        .cloned()
+        .unwrap_or_default();
+    let request_results_window =
+        append_request_result_window(&current, observed_at_unix_secs, false);
+    let already_open = current_bool(&current, "open");
+    if !already_open && consecutive_failures < LOCAL_KEY_CIRCUIT_FAILURE_THRESHOLD {
+        circuit_by_format.insert(
+            api_format.to_string(),
+            json!({
+                "open": false,
+                "open_at": Value::Null,
+                "reason": Value::Null,
+                "next_probe_at": Value::Null,
+                "next_probe_at_unix_secs": Value::Null,
+                "probe_interval_minutes": 0,
+                "max_probe_interval_minutes": normalize_max_probe_interval_minutes(max_probe_interval_minutes),
+                "failure_count": consecutive_failures,
+                "last_failure_at": unix_secs_to_rfc3339(observed_at_unix_secs),
+                "last_probe_failure_at": Value::Null,
+                "half_open_until": Value::Null,
+                "half_open_successes": 0,
+                "half_open_failures": 0,
+                "request_results_window": request_results_window,
+            }),
+        );
+        return Some(Value::Object(circuit_by_format));
+    }
+
+    let max_probe_interval_minutes =
+        normalize_max_probe_interval_minutes(max_probe_interval_minutes);
+    let probe_interval_minutes =
+        next_circuit_probe_interval_minutes(&current, max_probe_interval_minutes);
+    let next_probe_at_unix_secs =
+        next_probe_at_unix_secs(observed_at_unix_secs, probe_interval_minutes);
+    let open_at = current
+        .get("open_at")
+        .filter(|_| already_open)
+        .cloned()
+        .unwrap_or_else(|| json!(unix_secs_to_rfc3339(observed_at_unix_secs)));
+    let half_open_failures = current
+        .get("half_open_failures")
+        .and_then(Value::as_u64)
+        .unwrap_or(0)
+        .saturating_add(u64::from(already_open));
+
+    circuit_by_format.insert(
+        api_format.to_string(),
+        json!({
+            "open": true,
+            "open_at": open_at,
+            "reason": format!("consecutive_failures_{LOCAL_KEY_CIRCUIT_FAILURE_THRESHOLD}"),
+            "next_probe_at": unix_secs_to_rfc3339(next_probe_at_unix_secs),
+            "next_probe_at_unix_secs": next_probe_at_unix_secs,
+            "probe_interval_minutes": probe_interval_minutes,
+            "max_probe_interval_minutes": max_probe_interval_minutes,
+            "failure_count": consecutive_failures,
+            "last_failure_at": unix_secs_to_rfc3339(observed_at_unix_secs),
+            "last_probe_failure_at": if already_open {
+                json!(unix_secs_to_rfc3339(observed_at_unix_secs))
+            } else {
+                Value::Null
+            },
+            "half_open_until": Value::Null,
+            "half_open_successes": 0,
+            "half_open_failures": half_open_failures,
+            "request_results_window": request_results_window,
        }),
    );
    Some(Value::Object(circuit_by_format))
@@ -137,6 +259,60 @@ pub(crate) fn project_local_key_circuit_closed(
    Some(Value::Object(circuit_by_format))
 }

+fn current_bool(current: &serde_json::Map<String, Value>, field: &str) -> bool {
+    current.get(field).and_then(Value::as_bool).unwrap_or(false)
+}
+
+fn normalize_max_probe_interval_minutes(value: i32) -> u64 {
+    value.clamp(0, LOCAL_KEY_CIRCUIT_MAX_PROBE_INTERVAL_MINUTES as i32) as u64
+}
+
+fn next_circuit_probe_interval_minutes(
+    current: &serde_json::Map<String, Value>,
+    max_probe_interval_minutes: u64,
+) -> u64 {
+    if max_probe_interval_minutes == 0 {
+        return 0;
+    }
+    if !current_bool(current, "open") {
+        return 1.min(max_probe_interval_minutes);
+    }
+    current
+        .get("probe_interval_minutes")
+        .and_then(Value::as_u64)
+        .unwrap_or(1)
+        .max(1)
+        .saturating_mul(2)
+        .min(max_probe_interval_minutes)
+}
+
+fn next_probe_at_unix_secs(observed_at_unix_secs: u64, interval_minutes: u64) -> u64 {
+    observed_at_unix_secs.saturating_add(interval_minutes.saturating_mul(60))
+}
+
+fn append_request_result_window(
+    current: &serde_json::Map<String, Value>,
+    observed_at_unix_secs: u64,
+    ok: bool,
+) -> Value {
+    let mut window = current
+        .get("request_results_window")
+        .and_then(Value::as_array)
+        .cloned()
+        .unwrap_or_default();
+    window.push(json!({
+        "ts": observed_at_unix_secs,
+        "ok": ok,
+    }));
+    let keep = usize::try_from(LOCAL_KEY_CIRCUIT_FAILURE_THRESHOLD)
+        .unwrap_or(8)
+        .max(1);
+    if window.len() > keep {
+        window = window.split_off(window.len() - keep);
+    }
+    Value::Array(window)
+}
+
 fn local_candidate_failure_should_project_health(
    classification: LocalFailoverClassification,
    status_code: u16,
@@ -182,7 +358,8 @@ mod tests {

    use super::{
        project_local_failure_health, project_local_key_circuit_closed,
-        project_local_key_circuit_open, project_local_success_health,
+        project_local_key_circuit_failure, project_local_key_circuit_open,
+        project_local_success_health,
    };
    use crate::orchestration::LocalFailoverClassification;

@@ -269,6 +446,7 @@ mod tests {
            "openai:chat",
            "account_deactivated_401",
            1_760_000_000,
+            32,
        )
        .expect("projection should exist");

@@ -279,7 +457,47 @@ mod tests {
        );
        assert_eq!(
            projected["openai:chat"]["next_probe_at_unix_secs"],
-            json!(1_760_001_920u64)
+            json!(1_760_000_060u64)
+        );
+        assert_eq!(projected["openai:chat"]["probe_interval_minutes"], json!(1));
+    }
+
+    #[test]
+    fn consecutive_failure_circuit_opens_after_threshold_and_backs_off() {
+        let before_threshold =
+            project_local_key_circuit_failure(None, "openai:chat", 1_760_000_000, 7, 32)
+                .expect("projection should exist");
+        assert_eq!(before_threshold["openai:chat"]["open"], json!(false));
+
+        let opened = project_local_key_circuit_failure(
+            Some(&before_threshold),
+            "openai:chat",
+            1_760_000_060,
+            8,
+            32,
+        )
+        .expect("projection should exist");
+        assert_eq!(opened["openai:chat"]["open"], json!(true));
+        assert_eq!(
+            opened["openai:chat"]["reason"],
+            json!("consecutive_failures_8")
+        );
+        assert_eq!(opened["openai:chat"]["probe_interval_minutes"], json!(1));
+        assert_eq!(
+            opened["openai:chat"]["next_probe_at_unix_secs"],
+            json!(1_760_000_120u64)
+        );
+
+        let backed_off =
+            project_local_key_circuit_failure(Some(&opened), "openai:chat", 1_760_000_120, 9, 32)
+                .expect("projection should exist");
+        assert_eq!(
+            backed_off["openai:chat"]["probe_interval_minutes"],
+            json!(2)
+        );
+        assert_eq!(
+            backed_off["openai:chat"]["next_probe_at_unix_secs"],
+            json!(1_760_000_240u64)
        );
    }

@@ -34,7 +34,8 @@ pub(crate) use self::effects::{
    LocalPoolErrorEffect,
 };
 pub(crate) use self::health::{
-    project_local_failure_health, project_local_key_circuit_closed, project_local_key_circuit_open,
+    project_local_failure_health, project_local_key_circuit_closed,
+    project_local_key_circuit_failure, project_local_key_circuit_open,
    project_local_success_health,
 };
 pub(crate) use self::policy::{