Skip to content

feat(core): add execute_with_resilience as primary resilience API #811

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
CatherineSue wants to merge 2 commits into from
Closed

feat(core): add execute_with_resilience as primary resilience API #811

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
18b2311
feat(core): add execute_with_resilience as the primary resilience API
CatherineSue Mar 18, 2026
c320e01
fix(core): record CB outcome in retry path, remove URL from 503 body
CatherineSue Mar 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 3 additions & 1 deletion model_gateway/src/core/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,9 @@ pub use openai_protocol::{
model_type::{Endpoint, ModelType},
worker::{ProviderType, WorkerGroupKey},
};
pub use resilience::{resolve_resilience, ResolvedResilience, DEFAULT_RETRYABLE_STATUS_CODES};
pub use resilience::{
execute_with_resilience, resolve_resilience, ResolvedResilience, DEFAULT_RETRYABLE_STATUS_CODES,
};
pub use retry::{is_retryable_status, RetryExecutor};
pub use worker::{
AttachedBody, BasicWorker, ConnectionMode, RuntimeType, Worker, WorkerLoadGuard, WorkerType,
Expand Down
182 changes: 181 additions & 1 deletion model_gateway/src/core/resilience.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,9 +5,15 @@

use std::{collections::HashSet, time::Duration};

use axum::{http::StatusCode, response::IntoResponse};
use openai_protocol::worker::ResilienceUpdate;
use tracing::debug;

use crate::{config::types::RetryConfig, core::circuit_breaker::CircuitBreakerConfig};
use crate::{
config::types::RetryConfig,
core::{circuit_breaker::CircuitBreakerConfig, retry::RetryExecutor, worker::Worker},
observability::metrics::Metrics,
};

/// Default retryable HTTP status codes.
pub const DEFAULT_RETRYABLE_STATUS_CODES: &[u16] = &[408, 429, 500, 502, 503, 504];
Expand Down Expand Up @@ -108,9 +114,183 @@ pub fn resolve_resilience(
(resolved, cb_config)
}

/// Execute a request with the worker's retry and circuit breaker config.
///
/// This is the primary resilience API. Routers call this instead of
/// using `RetryExecutor` directly.
///
/// Circuit breaker outcomes are recorded automatically.
/// Retry decisions use the worker's `is_retryable()` predicate.
pub async fn execute_with_resilience<F, Fut>(
worker: &(dyn Worker + '_),
mut operation: F,
Comment on lines +124 to +126

@chatgpt-codex-connector chatgpt-codex-connector Bot Mar 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Preserve per-attempt worker selection for retries

I checked model_gateway/src/routers/http/router.rs:216-221 and :281-289: each retry currently re-enters select_worker_for_model, and that selector filters w.is_available() at :152-155. This helper bakes a single Worker into the retry loop, so any router migration will retry the same backend instead of failing over to another healthy worker. In multi-worker pools, a transient 429/503 from one instance will become a user-visible failure even when capacity is still available elsewhere.

Useful? React with 👍 / 👎.

) -> axum::response::Response
where
F: FnMut(u32) -> Fut + Send,
Fut: std::future::Future<Output = axum::response::Response> + Send,
{
let resilience = worker.resilience();
let worker_url = worker.url();

// Check circuit breaker before first attempt
if resilience.circuit_breaker_enabled && !worker.circuit_breaker().can_execute() {
debug!(
worker_url = worker_url,
"Circuit breaker open, rejecting request"
);
return (
StatusCode::SERVICE_UNAVAILABLE,
"Upstream worker temporarily unavailable",
)
.into_response();
Comment thread
CatherineSue marked this conversation as resolved.
Comment thread
coderabbitai[bot] marked this conversation as resolved.
}

if !resilience.retry_enabled {
// Single attempt, no retries
let response = operation(0).await;
if resilience.circuit_breaker_enabled {
let success = !worker.is_retryable(&response);
worker.circuit_breaker().record_outcome(success);
}
return response;
}

// Retry loop — delegate to RetryExecutor with worker-aware hooks
let response = RetryExecutor::execute_response_with_retry(
&resilience.retry,
operation,
|res, _attempt| worker.is_retryable(res),
Comment thread
CatherineSue marked this conversation as resolved.
Comment on lines +159 to +162

@chatgpt-codex-connector chatgpt-codex-connector Bot Mar 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Re-check circuit-breaker state before retrying

This path calls can_execute() only once before handing control to RetryExecutor. If another request opens the same worker's breaker while this call is sleeping in backoff, the next attempt still invokes operation on that worker because there is no per-attempt gate here. The current HTTP flow avoids that by selecting an available worker on every attempt (model_gateway/src/routers/http/router.rs:152-155, :216-221, :281-289), so migrating to this helper would let concurrent failure bursts keep hammering a worker that the breaker has already quarantined.

Useful? React with 👍 / 👎.

|delay, attempt| {
Metrics::record_worker_retry_backoff(attempt, delay);
debug!(
worker_url = worker_url,
attempt = attempt,
delay_ms = delay.as_millis() as u64,
"Worker retry backoff"
);
},
|| {
debug!(worker_url = worker_url, "Worker retries exhausted");
},
)
.await;

// Record final outcome for circuit breaker after all retries
if resilience.circuit_breaker_enabled {
let success = !worker.is_retryable(&response);
worker.circuit_breaker().record_outcome(success);
}

response
}

#[cfg(test)]
mod tests {
use std::sync::{
atomic::{AtomicU32, Ordering},
Arc,
};

use axum::{http::StatusCode, response::IntoResponse};

use super::*;
use crate::core::worker_builder::BasicWorkerBuilder;

#[tokio::test]
async fn test_execute_with_resilience_success() {
let worker = BasicWorkerBuilder::new("http://test:8080").build();
let response = execute_with_resilience(&worker, |_attempt| async {
(StatusCode::OK, "ok").into_response()
})
.await;
assert_eq!(response.status(), StatusCode::OK);
}

#[tokio::test]
async fn test_execute_with_resilience_circuit_open() {
let worker = BasicWorkerBuilder::new("http://test:8080").build();
worker.circuit_breaker().force_open();
let response = execute_with_resilience(&worker, |_attempt| async {
(StatusCode::OK, "ok").into_response()
})
.await;
assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
}

#[tokio::test]
async fn test_execute_with_resilience_retries_disabled() {
let resolved = ResolvedResilience {
retry_enabled: false,
..Default::default()
};
let worker = BasicWorkerBuilder::new("http://test:8080")
.resilience(resolved)
.build();

let call_count = Arc::new(AtomicU32::new(0));
let cc = call_count.clone();
let response = execute_with_resilience(&worker, move |_attempt| {
cc.fetch_add(1, Ordering::Relaxed);
async { (StatusCode::SERVICE_UNAVAILABLE, "fail").into_response() }
})
.await;

assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE);
assert_eq!(call_count.load(Ordering::Relaxed), 1); // No retries
}

#[tokio::test]
async fn test_execute_with_resilience_retries_on_retryable_status() {
let resolved = ResolvedResilience {
retry: RetryConfig {
max_retries: 3,
initial_backoff_ms: 1,
max_backoff_ms: 2,
backoff_multiplier: 1.0,
jitter_factor: 0.0,
},
..Default::default()
};
let worker = BasicWorkerBuilder::new("http://test:8080")
.resilience(resolved)
.build();

let call_count = Arc::new(AtomicU32::new(0));
let cc = call_count.clone();
let response = execute_with_resilience(&worker, move |_attempt| {
let count = cc.fetch_add(1, Ordering::Relaxed);
async move {
if count < 2 {
(StatusCode::SERVICE_UNAVAILABLE, "fail").into_response()
} else {
(StatusCode::OK, "ok").into_response()
}
}
})
.await;

assert_eq!(response.status(), StatusCode::OK);
assert_eq!(call_count.load(Ordering::Relaxed), 3);
}
Comment on lines +242 to +274

@coderabbitai coderabbitai Bot Mar 18, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick | 🔵 Trivial

Add a regression test that asserts breaker state updates in the retry-enabled path.

Current tests validate status and retry counts, but there is no explicit assertion that breaker outcome/state changes after the retry loop. A focused assertion here would lock in this fix against regression.

🤖 Prompt for AI Agents 
Verify each finding against the current code and only fix it if needed.

In `@model_gateway/src/core/resilience.rs` around lines 242 - 274, Add an
assertion after the existing response checks in the test that verifies the
circuit-breaker state/outcome was updated by the retry-enabled path: after
calling execute_with_resilience(&worker, ...) and asserting StatusCode::OK and
call_count, query the worker's resilience/breaker (the breaker instance
configured via ResolvedResilience/RetryConfig on BasicWorkerBuilder::resilience)
using the actual breaker API in the codebase (e.g., the breaker state or last
outcome accessor) and assert it reflects a successful outcome (closed/success or
equivalent enum variant). Ensure you reference the same worker variable and use
the breaker state accessor implemented on the worker/resilience types so the
test fails if breaker state is not updated.


#[tokio::test]
async fn test_execute_with_resilience_cb_disabled_ignores_open() {
let resolved = ResolvedResilience {
circuit_breaker_enabled: false,
..Default::default()
};
let worker = BasicWorkerBuilder::new("http://test:8080")
.resilience(resolved)
.build();
worker.circuit_breaker().force_open();

// CB is disabled — request should go through even though CB is open
let response = execute_with_resilience(&worker, |_attempt| async {
(StatusCode::OK, "ok").into_response()
})
.await;
assert_eq!(response.status(), StatusCode::OK);
}

#[test]
fn test_default_resilience() {
Expand Down
Loading