mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
fix(gateway): preserve cache affinity during health updates
This commit is contained in:
@@ -1546,6 +1546,39 @@ mod tests {
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn health_success_keeps_scheduler_affinity_after_health_state_update() {
|
||||
let state = health_state();
|
||||
let plan = sample_plan();
|
||||
let report_context = json!({
|
||||
"api_key_id": "api-key-1",
|
||||
"client_api_format": "openai:chat",
|
||||
"model": "gpt-5",
|
||||
});
|
||||
let cache_key =
|
||||
build_scheduler_affinity_cache_key_for_api_key_id("api-key-1", "openai:chat", "gpt-5")
|
||||
.expect("scheduler affinity cache key should build");
|
||||
|
||||
apply_local_execution_effect(
|
||||
&state,
|
||||
LocalExecutionEffectContext {
|
||||
plan: &plan,
|
||||
report_context: Some(&report_context),
|
||||
},
|
||||
LocalExecutionEffect::HealthSuccess(LocalHealthSuccessEffect),
|
||||
)
|
||||
.await;
|
||||
|
||||
assert_eq!(
|
||||
state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
|
||||
Some(SchedulerAffinityTarget {
|
||||
provider_id: "prov-1".to_string(),
|
||||
endpoint_id: "ep-1".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn load_balance_success_does_not_remember_scheduler_affinity_cache() {
|
||||
let state = AppState::new()
|
||||
|
||||
@@ -806,7 +806,7 @@ impl AppState {
|
||||
.await
|
||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
||||
if updated {
|
||||
self.invalidate_provider_routing_caches();
|
||||
self.invalidate_provider_health_routing_caches();
|
||||
}
|
||||
Ok(updated)
|
||||
}
|
||||
@@ -930,4 +930,47 @@ mod tests {
|
||||
.expect("provider transport should exist after update");
|
||||
assert!(snapshot.provider.keep_priority_on_conversion);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn provider_catalog_health_update_keeps_scheduler_affinity_cache() {
|
||||
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
|
||||
vec![sample_provider()],
|
||||
vec![sample_endpoint()],
|
||||
vec![sample_key()],
|
||||
));
|
||||
let state = AppState::new()
|
||||
.expect("app state should build")
|
||||
.with_data_state_for_tests(
|
||||
GatewayDataState::with_provider_catalog_repository_for_tests(repository)
|
||||
.with_encryption_key_for_tests("test-encryption-key"),
|
||||
);
|
||||
|
||||
let cache_key = "scheduler_affinity:api-key-1:openai:chat:gpt-5";
|
||||
let ttl = Duration::from_secs(300);
|
||||
let target = SchedulerAffinityTarget {
|
||||
provider_id: "provider-1".to_string(),
|
||||
endpoint_id: "endpoint-1".to_string(),
|
||||
key_id: "key-1".to_string(),
|
||||
};
|
||||
state.remember_scheduler_affinity_target(cache_key, target.clone(), ttl, 128);
|
||||
let initial_epoch = state.scheduler_affinity_epoch();
|
||||
|
||||
let health_by_format = serde_json::json!({
|
||||
"openai:chat": {
|
||||
"last_success_at_unix_secs": 1,
|
||||
"consecutive_failures": 0
|
||||
}
|
||||
});
|
||||
let updated = state
|
||||
.update_provider_catalog_key_health_state("key-1", true, Some(&health_by_format), None)
|
||||
.await
|
||||
.expect("key health update should succeed");
|
||||
|
||||
assert!(updated);
|
||||
assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
|
||||
assert_eq!(
|
||||
state.read_scheduler_affinity_target(cache_key, ttl),
|
||||
Some(target)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -578,6 +578,11 @@ impl AppState {
|
||||
self.invalidate_scheduler_affinity_cache();
|
||||
}
|
||||
|
||||
pub(crate) fn invalidate_provider_health_routing_caches(&self) {
|
||||
self.data.clear_minimal_candidate_selection_cache();
|
||||
self.clear_provider_transport_snapshot_cache();
|
||||
}
|
||||
|
||||
pub(crate) fn invalidate_auth_context_cache(&self) {
|
||||
self.auth_context_cache.clear();
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user