fix(gateway): preserve cache affinity during health updates

This commit is contained in:
RWDai
2026-05-21 15:51:56 +08:00
parent b84e4a96e2
commit 965a8c79af
3 changed files with 82 additions and 1 deletions

View File

@@ -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()

View File

@@ -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)
);
}
}

View File

@@ -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();
}