Merge pull request #534 from RWDai/fix/issue-528-cache-affinity-health

fix(gateway): preserve cache affinity during health updates

apps/aether-gateway/src/orchestration/effects.rs

@@ -243,6 +243,35 @@ fn local_scheduler_affinity_target(plan: &ExecutionPlan) -> Option<SchedulerAffi
     })
 }
 
+async fn local_execution_plan_uses_pool(state: &AppState, plan: &ExecutionPlan) -> bool {
+    let Ok(Some(transport)) = state
+        .read_provider_transport_snapshot(&plan.provider_id, &plan.endpoint_id, &plan.key_id)
+        .await
+    else {
+        return false;
+    };
+
+    admin_provider_pool_config_from_config_value(transport.provider.config.as_ref()).is_some()
+}
+
+async fn local_scheduler_affinity_matches_failed_target(
+    state: &AppState,
+    plan: &ExecutionPlan,
+    cached_target: &SchedulerAffinityTarget,
+    failed_target: &SchedulerAffinityTarget,
+) -> bool {
+    if cached_target == failed_target {
+        return true;
+    }
+    if cached_target.provider_id != failed_target.provider_id
+        || cached_target.endpoint_id != failed_target.endpoint_id
+    {
+        return false;
+    }
+
+    local_execution_plan_uses_pool(state, plan).await
+}
+
 async fn scheduler_cache_affinity_enabled(state: &AppState) -> bool {
     match read_scheduler_ordering_config(state).await {
         Ok(config) => config.scheduling_mode == SchedulerSchedulingMode::CacheAffinity,
@@ -360,7 +389,24 @@ async fn record_attempt_failure_effect(
     }
 
     if let Some(cache_key) = local_scheduler_affinity_cache_key(context.report_context) {
-        let _ = state.remove_scheduler_affinity_cache_entry(&cache_key);
+        let Some(failed_target) = local_scheduler_affinity_target(context.plan) else {
+            return;
+        };
+        let Some(cached_target) =
+            state.read_scheduler_affinity_target(&cache_key, SCHEDULER_AFFINITY_TTL)
+        else {
+            return;
+        };
+        if local_scheduler_affinity_matches_failed_target(
+            state,
+            context.plan,
+            &cached_target,
+            &failed_target,
+        )
+        .await
+        {
+            let _ = state.remove_scheduler_affinity_cache_entry(&cache_key);
+        }
     }
 }
 
@@ -448,7 +494,10 @@ async fn record_adaptive_rate_limit_effect(
     updated_key.status_snapshot = Some(projection.status_snapshot);
     updated_key.updated_at_unix_secs = Some(observed_at_unix_secs);
 
-    if let Err(err) = state.update_provider_catalog_key(&updated_key).await {
+    if let Err(err) = state
+        .update_provider_catalog_key_runtime_state(&updated_key)
+        .await
+    {
         warn!(
             "gateway orchestration effects: failed to persist adaptive rate-limit projection for provider {} endpoint {} key {}: {:?}",
             context.plan.provider_id, context.plan.endpoint_id, context.plan.key_id, err
@@ -496,7 +545,10 @@ async fn record_adaptive_success_effect(
     updated_key.status_snapshot = Some(projection.status_snapshot);
     updated_key.updated_at_unix_secs = Some(observed_at_unix_secs);
 
-    if let Err(err) = state.update_provider_catalog_key(&updated_key).await {
+    if let Err(err) = state
+        .update_provider_catalog_key_runtime_state(&updated_key)
+        .await
+    {
         warn!(
             "gateway orchestration effects: failed to persist adaptive success projection for provider {} endpoint {} key {}: {:?}",
             context.plan.provider_id, context.plan.endpoint_id, context.plan.key_id, err
@@ -1204,6 +1256,20 @@ mod tests {
         .expect("provider should build")
     }
 
+    fn sample_pool_health_provider() -> StoredProviderCatalogProvider {
+        sample_health_provider().with_transport_fields(
+            true,
+            false,
+            false,
+            None,
+            None,
+            None,
+            None,
+            None,
+            Some(json!({"pool_advanced": {}})),
+        )
+    }
+
     fn sample_health_endpoint() -> StoredProviderCatalogEndpoint {
         StoredProviderCatalogEndpoint::new(
             "ep-1".to_string(),
@@ -1266,6 +1332,20 @@ mod tests {
             )
     }
 
+    fn pool_health_state() -> AppState {
+        let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
+            vec![sample_pool_health_provider()],
+            vec![sample_health_endpoint()],
+            vec![sample_health_key()],
+        ));
+        AppState::new()
+            .expect("gateway state should build")
+            .with_data_state_for_tests(
+                GatewayDataState::with_provider_catalog_repository_for_tests(repository)
+                    .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
+            )
+    }
+
     fn health_state_with_key(key: StoredProviderCatalogKey) -> AppState {
         let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
             vec![sample_health_provider()],
@@ -1429,6 +1509,137 @@ mod tests {
             .is_some());
     }
 
+    #[tokio::test]
+    async fn attempt_failure_keeps_scheduler_affinity_for_non_affinity_candidate() {
+        let state = AppState::new().expect("gateway state should build");
+        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");
+        let affinity_target = SchedulerAffinityTarget {
+            provider_id: "prov-2".to_string(),
+            endpoint_id: "ep-2".to_string(),
+            key_id: "key-2".to_string(),
+        };
+
+        state.remember_scheduler_affinity_target(
+            &cache_key,
+            affinity_target.clone(),
+            SCHEDULER_AFFINITY_TTL,
+            16,
+        );
+
+        apply_local_execution_effect(
+            &state,
+            LocalExecutionEffectContext {
+                plan: &plan,
+                report_context: Some(&report_context),
+            },
+            LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
+                status_code: 524,
+                classification: LocalFailoverClassification::RetryUpstreamFailure,
+            }),
+        )
+        .await;
+
+        assert_eq!(
+            state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
+            Some(affinity_target)
+        );
+    }
+
+    #[tokio::test]
+    async fn attempt_failure_keeps_scheduler_affinity_for_non_pool_sibling_key() {
+        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");
+        let affinity_target = SchedulerAffinityTarget {
+            provider_id: "prov-1".to_string(),
+            endpoint_id: "ep-1".to_string(),
+            key_id: "key-2".to_string(),
+        };
+
+        state.remember_scheduler_affinity_target(
+            &cache_key,
+            affinity_target.clone(),
+            SCHEDULER_AFFINITY_TTL,
+            16,
+        );
+
+        apply_local_execution_effect(
+            &state,
+            LocalExecutionEffectContext {
+                plan: &plan,
+                report_context: Some(&report_context),
+            },
+            LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
+                status_code: 524,
+                classification: LocalFailoverClassification::RetryUpstreamFailure,
+            }),
+        )
+        .await;
+
+        assert_eq!(
+            state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
+            Some(affinity_target)
+        );
+    }
+
+    #[tokio::test]
+    async fn attempt_failure_invalidates_scheduler_affinity_for_same_pool_candidate() {
+        let state = pool_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");
+
+        state.remember_scheduler_affinity_target(
+            &cache_key,
+            SchedulerAffinityTarget {
+                provider_id: "prov-1".to_string(),
+                endpoint_id: "ep-1".to_string(),
+                key_id: "key-2".to_string(),
+            },
+            SCHEDULER_AFFINITY_TTL,
+            16,
+        );
+
+        apply_local_execution_effect(
+            &state,
+            LocalExecutionEffectContext {
+                plan: &plan,
+                report_context: Some(&report_context),
+            },
+            LocalExecutionEffect::AttemptFailure(LocalAttemptFailureEffect {
+                status_code: 524,
+                classification: LocalFailoverClassification::RetryUpstreamFailure,
+            }),
+        )
+        .await;
+
+        assert_eq!(
+            state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
+            None
+        );
+    }
+
     #[tokio::test]
     async fn attempt_failure_keeps_scheduler_affinity_for_non_failure_status() {
         let state = AppState::new().expect("gateway state should build");
@@ -1546,6 +1757,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()
@@ -1985,6 +2229,21 @@ mod tests {
     async fn adaptive_rate_limit_effect_updates_adaptive_key_observation() {
         let state = adaptive_state();
         let plan = sample_plan();
+        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");
+        let target = SchedulerAffinityTarget {
+            provider_id: plan.provider_id.clone(),
+            endpoint_id: plan.endpoint_id.clone(),
+            key_id: plan.key_id.clone(),
+        };
+        state.remember_scheduler_affinity_target(
+            &cache_key,
+            target.clone(),
+            SCHEDULER_AFFINITY_TTL,
+            16,
+        );
+        let initial_epoch = state.scheduler_affinity_epoch();
 
         apply_local_execution_effect(
             &state,
@@ -2048,6 +2307,11 @@ mod tests {
                 .and_then(|value| value.get("enforcement_active")),
             Some(&json!(false))
         );
+        assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
+        assert_eq!(
+            state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
+            Some(target)
+        );
     }
 
     #[tokio::test]
@@ -2178,6 +2442,21 @@ mod tests {
             .expect("request candidate should build")],
         );
         let plan = sample_plan();
+        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");
+        let target = SchedulerAffinityTarget {
+            provider_id: plan.provider_id.clone(),
+            endpoint_id: plan.endpoint_id.clone(),
+            key_id: plan.key_id.clone(),
+        };
+        state.remember_scheduler_affinity_target(
+            &cache_key,
+            target.clone(),
+            SCHEDULER_AFFINITY_TTL,
+            16,
+        );
+        let initial_epoch = state.scheduler_affinity_epoch();
 
         apply_local_execution_effect(
             &state,
@@ -2209,5 +2488,10 @@ mod tests {
                 .and_then(Value::as_str),
             Some("high_utilization")
         );
+        assert_eq!(state.scheduler_affinity_epoch(), initial_epoch);
+        assert_eq!(
+            state.read_scheduler_affinity_target(cache_key.as_str(), SCHEDULER_AFFINITY_TTL),
+            Some(target)
+        );
     }
 }

apps/aether-gateway/src/orchestration/report_effects.rs

@@ -424,7 +424,7 @@ async fn sync_grok_quota_from_report_context(
     updated_key.updated_at_unix_secs = Some(now_unix_secs);
 
     Ok(state
-        .update_provider_catalog_key(&updated_key)
+        .update_provider_catalog_key_runtime_state(&updated_key)
         .await?
         .is_some())
 }
@@ -721,7 +721,7 @@ async fn sync_codex_quota_from_response_headers(
     updated_key.updated_at_unix_secs = Some(now_unix_secs);
 
     let updated = state
-        .update_provider_catalog_key(&updated_key)
+        .update_provider_catalog_key_runtime_state(&updated_key)
         .await?
         .is_some();
     if updated {

apps/aether-gateway/src/state/catalog.rs

@@ -615,6 +615,21 @@ impl AppState {
         Ok(updated)
     }
 
+    pub(crate) async fn update_provider_catalog_key_runtime_state(
+        &self,
+        key: &provider_catalog::StoredProviderCatalogKey,
+    ) -> Result<Option<provider_catalog::StoredProviderCatalogKey>, GatewayError> {
+        let updated = self
+            .data
+            .update_provider_catalog_key(key)
+            .await
+            .map_err(|err| GatewayError::Internal(err.to_string()))?;
+        if updated.is_some() {
+            self.invalidate_provider_health_routing_caches();
+        }
+        Ok(updated)
+    }
+
     pub(crate) async fn update_provider_catalog_key_upstream_metadata(
         &self,
         key_id: &str,
@@ -806,7 +821,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 +945,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)
+        );
+    }
 }

apps/aether-gateway/src/state/core.rs

@@ -572,6 +572,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();
     }