[model-gateway] extract circuit breaker state struct (#15098)
This commit is contained in:
@@ -53,6 +53,24 @@ impl std::fmt::Display for CircuitState {
|
||||
}
|
||||
}
|
||||
|
||||
impl CircuitState {
|
||||
pub fn as_str(&self) -> &'static str {
|
||||
match self {
|
||||
CircuitState::Closed => "closed",
|
||||
CircuitState::Open => "open",
|
||||
CircuitState::HalfOpen => "half_open",
|
||||
}
|
||||
}
|
||||
|
||||
pub fn to_int(&self) -> u8 {
|
||||
match self {
|
||||
CircuitState::Closed => 0u8,
|
||||
CircuitState::Open => 1u8,
|
||||
CircuitState::HalfOpen => 2u8,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Circuit breaker implementation
|
||||
#[derive(Debug)]
|
||||
pub struct CircuitBreaker {
|
||||
@@ -196,16 +214,8 @@ impl CircuitBreaker {
|
||||
}
|
||||
}
|
||||
|
||||
let from = match old_state {
|
||||
CircuitState::Closed => "closed",
|
||||
CircuitState::Open => "open",
|
||||
CircuitState::HalfOpen => "half_open",
|
||||
};
|
||||
let to = match new_state {
|
||||
CircuitState::Closed => "closed",
|
||||
CircuitState::Open => "open",
|
||||
CircuitState::HalfOpen => "half_open",
|
||||
};
|
||||
let from = old_state.as_str();
|
||||
let to = new_state.as_str();
|
||||
info!("Circuit breaker state transition: {} -> {}", from, to);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -16,7 +16,7 @@ use super::{
|
||||
CircuitBreaker, Endpoint, ModelCard, ModelType, ProviderType, WorkerError, WorkerResult,
|
||||
};
|
||||
use crate::{
|
||||
core::{BasicWorkerBuilder, CircuitState, DPAwareWorkerBuilder},
|
||||
core::{BasicWorkerBuilder, DPAwareWorkerBuilder},
|
||||
observability::metrics::RouterMetrics,
|
||||
protocols::worker_spec::WorkerInfo,
|
||||
routers::grpc::client::GrpcClient,
|
||||
@@ -137,24 +137,12 @@ pub trait Worker: Send + Sync + fmt::Debug {
|
||||
let after = self.circuit_breaker().state();
|
||||
|
||||
if before != after {
|
||||
let from = match before {
|
||||
CircuitState::Closed => "closed",
|
||||
CircuitState::Open => "open",
|
||||
CircuitState::HalfOpen => "half_open",
|
||||
};
|
||||
let to = match after {
|
||||
CircuitState::Closed => "closed",
|
||||
CircuitState::Open => "open",
|
||||
CircuitState::HalfOpen => "half_open",
|
||||
};
|
||||
let from = before.as_str();
|
||||
let to = after.as_str();
|
||||
RouterMetrics::record_cb_state_transition(self.url(), from, to);
|
||||
}
|
||||
|
||||
let state_code = match self.circuit_breaker().state() {
|
||||
CircuitState::Closed => 0u8,
|
||||
CircuitState::Open => 1u8,
|
||||
CircuitState::HalfOpen => 2u8,
|
||||
};
|
||||
let state_code = self.circuit_breaker().state().to_int();
|
||||
RouterMetrics::set_cb_state(self.url(), state_code);
|
||||
|
||||
// Update consecutive failures/successes gauges
|
||||
@@ -1180,7 +1168,7 @@ mod tests {
|
||||
use std::{thread, time::Duration};
|
||||
|
||||
use super::*;
|
||||
use crate::core::CircuitBreakerConfig;
|
||||
use crate::core::{CircuitBreakerConfig, CircuitState};
|
||||
|
||||
#[test]
|
||||
fn test_worker_type_display() {
|
||||
|
||||
Reference in New Issue
Block a user