Track scheduler affinity epochs and key sorting

This commit is contained in:
fawney19
2026-05-11 01:45:49 +08:00
parent 7b81c77424
commit e3574e1918
38 changed files with 656 additions and 23 deletions

View File

@@ -40,6 +40,26 @@ pub(crate) fn remember_scheduler_affinity_for_candidate(
client_api_format: &str,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
) {
remember_scheduler_affinity_for_candidate_at_epoch(
state,
auth_snapshot,
client_session_affinity,
client_api_format,
requested_model,
candidate,
None,
);
}
pub(crate) fn remember_scheduler_affinity_for_candidate_at_epoch(
state: PlannerAppState<'_>,
auth_snapshot: Option<&GatewayAuthApiKeySnapshot>,
client_session_affinity: Option<&ClientSessionAffinity>,
client_api_format: &str,
requested_model: &str,
candidate: &SchedulerMinimalCandidateSelectionCandidate,
expected_epoch: Option<u64>,
) {
let Some(api_key_id) = auth_snapshot
.map(|snapshot| snapshot.api_key_id.trim())
@@ -56,7 +76,7 @@ pub(crate) fn remember_scheduler_affinity_for_candidate(
return;
};
state.app().remember_scheduler_affinity_target(
let _ = state.app().remember_scheduler_affinity_target_for_epoch(
&cache_key,
SchedulerAffinityTarget {
provider_id: candidate.provider_id.clone(),
@@ -65,5 +85,6 @@ pub(crate) fn remember_scheduler_affinity_for_candidate(
},
SCHEDULER_AFFINITY_TTL,
PLANNER_SCHEDULER_AFFINITY_MAX_ENTRIES,
expected_epoch,
);
}

View File

@@ -15,7 +15,7 @@ use std::sync::Arc;
use tracing::warn;
use uuid::Uuid;
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate;
use crate::ai_serving::planner::candidate_affinity_cache::remember_scheduler_affinity_for_candidate_at_epoch;
use crate::ai_serving::planner::candidate_resolution::{
resolve_and_rank_logical_local_execution_candidates, EligibleLocalExecutionCandidate,
LocalExecutionCandidateKind, SkippedLocalExecutionCandidate,
@@ -810,13 +810,14 @@ pub(crate) fn remember_first_local_candidate_affinity(
.map(str::trim)
.filter(|value| !value.is_empty())
.unwrap_or(first_candidate.candidate.global_model_name.as_str());
remember_scheduler_affinity_for_candidate(
remember_scheduler_affinity_for_candidate_at_epoch(
state,
auth_snapshot,
client_session_affinity,
client_api_format,
affinity_requested_model,
&first_candidate.candidate,
first_candidate.orchestration.scheduler_affinity_epoch,
);
}
@@ -1324,6 +1325,7 @@ mod tests {
candidate_group_id: pool_key_index.map(|_| "pool-group".to_string()),
pool_key_index,
pool_key_lease: None,
scheduler_affinity_epoch: None,
},
ranking: None,
}

View File

@@ -332,6 +332,7 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
Vec<EligibleLocalExecutionCandidate>,
Vec<SkippedLocalExecutionCandidate>,
) {
let scheduler_affinity_epoch = state.app().scheduler_affinity_epoch();
let port = GatewayLocalCandidateResolutionPort {
state,
requested_model,
@@ -350,7 +351,12 @@ async fn resolve_and_rank_local_execution_candidates_with_pool_expansion(
};
match run_ai_candidate_resolution(&port, candidates, request).await {
Ok(outcome) => (outcome.eligible_candidates, outcome.skipped_candidates),
Ok(mut outcome) => {
for candidate in &mut outcome.eligible_candidates {
candidate.orchestration.scheduler_affinity_epoch = Some(scheduler_affinity_epoch);
}
(outcome.eligible_candidates, outcome.skipped_candidates)
}
Err(error) => match error {},
}
}

View File

@@ -119,6 +119,7 @@ pub(crate) async fn maybe_build_local_same_format_provider_decision_payload_for_
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)

View File

@@ -938,10 +938,12 @@ fn apply_ai_pool_orchestration(
mut candidate: EligibleLocalExecutionCandidate,
orchestration: AiPoolCandidateOrchestration,
) -> EligibleLocalExecutionCandidate {
let scheduler_affinity_epoch = candidate.orchestration.scheduler_affinity_epoch;
candidate.orchestration = LocalExecutionCandidateMetadata {
candidate_group_id: orchestration.candidate_group_id,
pool_key_index: orchestration.pool_key_index,
pool_key_lease: None,
scheduler_affinity_epoch,
};
candidate
}
@@ -1127,6 +1129,7 @@ mod tests {
),
pool_key_index: Some(0),
pool_key_lease: None,
scheduler_affinity_epoch: None,
}
);
assert_eq!(reordered[1].orchestration.pool_key_index, Some(1));
@@ -1143,6 +1146,7 @@ mod tests {
),
pool_key_index: None,
pool_key_lease: None,
scheduler_affinity_epoch: None,
}
);
}

View File

@@ -18,7 +18,10 @@ use crate::ai_serving::{
use crate::client_session_affinity::{
client_session_affinity_report_context_value, CLIENT_SESSION_AFFINITY_REPORT_CONTEXT_FIELD,
};
use crate::orchestration::{insert_pool_key_lease_report_context_fields, ExecutionAttemptIdentity};
use crate::orchestration::{
insert_pool_key_lease_report_context_fields, ExecutionAttemptIdentity,
SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
};
pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) auth_context: &'a ExecutionRuntimeAuthContext,
@@ -52,6 +55,7 @@ pub(crate) struct LocalExecutionReportContextParts<'a> {
pub(crate) original_request_body_json: Option<&'a Value>,
pub(crate) original_request_body_base64: Option<&'a str>,
pub(crate) client_session_affinity: Option<&'a ClientSessionAffinity>,
pub(crate) scheduler_affinity_epoch: Option<u64>,
pub(crate) client_requested_stream: bool,
pub(crate) upstream_is_stream: bool,
pub(crate) has_envelope: bool,
@@ -89,6 +93,12 @@ pub(crate) fn build_local_execution_report_context(
merge_incoming_tls_fingerprint(&mut extra_fields, incoming_tls);
}
insert_pool_key_lease_report_context_fields(&mut extra_fields, parts.pool_key_lease);
if let Some(epoch) = parts.scheduler_affinity_epoch {
extra_fields.insert(
SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD.to_string(),
Value::Number(epoch.into()),
);
}
insert_request_path_fields(
&mut extra_fields,
parts.request_path,
@@ -278,6 +288,7 @@ mod tests {
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,
client_session_affinity: Some(&client_session_affinity),
scheduler_affinity_epoch: None,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,
@@ -359,6 +370,7 @@ mod tests {
})),
original_request_body_base64: None,
client_session_affinity: None,
scheduler_affinity_epoch: None,
client_requested_stream: false,
upstream_is_stream: true,
has_envelope: false,
@@ -424,6 +436,7 @@ mod tests {
original_request_body_json: Some(&json!({"model": "gpt-5"})),
original_request_body_base64: None,
client_session_affinity: None,
scheduler_affinity_epoch: None,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,

View File

@@ -101,6 +101,7 @@ pub(super) async fn maybe_build_local_gemini_files_decision_payload_for_candidat
original_request_body_json: Some(body_json),
original_request_body_base64: resolved.provider_request_body_base64.as_deref(),
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream: spec_metadata.require_streaming,
has_envelope: false,

View File

@@ -121,6 +121,7 @@ pub(super) async fn maybe_build_local_openai_image_decision_payload_for_candidat
original_request_body_json: Some(body_json),
original_request_body_base64: body_base64,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: spec_metadata.require_streaming,
upstream_is_stream,
has_envelope: false,

View File

@@ -82,6 +82,7 @@ pub(super) async fn maybe_build_local_video_create_decision_payload_for_candidat
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: false,
upstream_is_stream: false,
has_envelope: false,

View File

@@ -127,6 +127,7 @@ pub(super) async fn maybe_build_local_standard_decision_payload_for_candidate(
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)

View File

@@ -116,6 +116,7 @@ pub(crate) async fn maybe_build_local_openai_chat_decision_payload_for_candidate
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)

View File

@@ -112,6 +112,7 @@ pub(crate) async fn maybe_build_local_openai_responses_decision_payload_for_cand
original_request_body_json: Some(body_json),
original_request_body_base64: None,
client_session_affinity: input.client_session_affinity.as_ref(),
scheduler_affinity_epoch: eligible.orchestration.scheduler_affinity_epoch,
client_requested_stream: body_json
.get("stream")
.and_then(serde_json::Value::as_bool)

View File

@@ -5,13 +5,20 @@ pub(crate) use aether_scheduler_core::SchedulerAffinityTarget;
#[derive(Debug, Default)]
pub(crate) struct SchedulerAffinityCache {
entries: ExpiringMap<String, SchedulerAffinityTarget>,
entries: ExpiringMap<String, SchedulerAffinityCacheValue>,
}
#[derive(Debug, Clone)]
struct SchedulerAffinityCacheValue {
target: SchedulerAffinityTarget,
epoch: u64,
}
#[derive(Debug, Clone)]
pub(crate) struct SchedulerAffinitySnapshotEntry {
pub(crate) cache_key: String,
pub(crate) target: SchedulerAffinityTarget,
pub(crate) epoch: u64,
pub(crate) age: Duration,
}
@@ -21,7 +28,19 @@ impl SchedulerAffinityCache {
cache_key: &str,
ttl: Duration,
) -> Option<SchedulerAffinityTarget> {
self.entries.get_fresh(&cache_key.to_string(), ttl)
self.entries
.get_fresh(&cache_key.to_string(), ttl)
.map(|value| value.target)
}
pub(crate) fn get_fresh_for_epoch(
&self,
cache_key: &str,
ttl: Duration,
epoch: u64,
) -> Option<SchedulerAffinityTarget> {
let value = self.entries.get_fresh(&cache_key.to_string(), ttl)?;
(value.epoch == epoch).then_some(value.target)
}
#[cfg_attr(not(test), allow(dead_code))]
@@ -32,11 +51,33 @@ impl SchedulerAffinityCache {
ttl: Duration,
max_entries: usize,
) {
self.entries.insert(cache_key, target, ttl, max_entries);
self.insert_for_epoch(cache_key, target, ttl, max_entries, 0);
}
pub(crate) fn insert_for_epoch(
&self,
cache_key: String,
target: SchedulerAffinityTarget,
ttl: Duration,
max_entries: usize,
epoch: u64,
) {
self.entries.insert(
cache_key,
SchedulerAffinityCacheValue { target, epoch },
ttl,
max_entries,
);
}
pub(crate) fn remove(&self, cache_key: &str) -> Option<SchedulerAffinityTarget> {
self.entries.remove(&cache_key.to_string())
self.entries
.remove(&cache_key.to_string())
.map(|value| value.target)
}
pub(crate) fn clear(&self) {
self.entries.clear();
}
pub(crate) fn fresh_entries(&self, ttl: Duration) -> Vec<SchedulerAffinitySnapshotEntry> {
@@ -46,10 +87,22 @@ impl SchedulerAffinityCache {
.map(
|ExpiringMapFreshEntry { key, value, age }| SchedulerAffinitySnapshotEntry {
cache_key: key,
target: value,
target: value.target,
epoch: value.epoch,
age,
},
)
.collect()
}
pub(crate) fn fresh_entries_for_epoch(
&self,
ttl: Duration,
epoch: u64,
) -> Vec<SchedulerAffinitySnapshotEntry> {
self.fresh_entries(ttl)
.into_iter()
.filter(|entry| entry.epoch == epoch)
.collect()
}
}

View File

@@ -262,6 +262,14 @@ fn admin_monitoring_json_request_count(
})
}
fn admin_monitoring_json_scheduler_affinity_epoch(
object: &serde_json::Map<String, serde_json::Value>,
) -> Option<u64> {
object
.get("scheduler_affinity_epoch")
.and_then(serde_json::Value::as_u64)
}
pub(super) fn admin_monitoring_cache_affinity_record(
raw_key: &str,
raw_value: &str,
@@ -300,12 +308,14 @@ pub(super) fn admin_monitoring_cache_affinity_record(
expire_at: object.get("expire_at").cloned(),
request_count: request_count.unwrap_or(0),
request_count_known: request_count.is_some(),
scheduler_affinity_epoch: None,
})
}
pub(super) fn admin_monitoring_scheduler_affinity_record(
cache_key: &str,
target: &SchedulerAffinityTarget,
epoch: u64,
age: Duration,
ttl: Duration,
now_unix_secs: u64,
@@ -329,6 +339,7 @@ pub(super) fn admin_monitoring_scheduler_affinity_record(
expire_at: Some(serde_json::json!(expire_at)),
request_count: 0,
request_count_known: false,
scheduler_affinity_epoch: Some(epoch),
})
}
@@ -371,6 +382,7 @@ pub(super) fn admin_monitoring_scheduler_affinity_record_from_raw(
expire_at: object.get("expire_at").cloned(),
request_count: request_count.unwrap_or(0),
request_count_known: request_count.is_some(),
scheduler_affinity_epoch: admin_monitoring_json_scheduler_affinity_epoch(object),
})
}

View File

@@ -16,6 +16,7 @@ pub(in super::super) async fn build_admin_monitoring_cache_flush_response(
.iter()
.map(|item| item.raw_key.clone())
.collect::<Vec<_>>();
state.as_ref().invalidate_scheduler_affinity_cache();
let deleted = delete_admin_monitoring_cache_affinity_raw_keys(state, &raw_keys).await?;
clear_admin_monitoring_scheduler_affinity_entries(state, &raw_affinities);

View File

@@ -157,6 +157,7 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
records.push(record);
}
};
let current_scheduler_affinity_epoch = state.as_ref().scheduler_affinity_epoch();
{
let patterns = affinity_keys
@@ -210,6 +211,13 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
let Some(record) = record else {
continue;
};
if key.contains("scheduler_affinity:")
&& record
.scheduler_affinity_epoch
.is_some_and(|epoch| epoch != current_scheduler_affinity_epoch)
{
continue;
}
if affinity_keys.is_some_and(|keys| !keys.contains(&record.affinity_key)) {
continue;
}
@@ -237,6 +245,7 @@ async fn list_admin_monitoring_cache_affinity_records_matching(
let Some(record) = admin_monitoring_scheduler_affinity_record(
&entry.cache_key,
&entry.target,
entry.epoch,
entry.age,
SCHEDULER_AFFINITY_TTL,
now_unix_secs,

View File

@@ -13,6 +13,7 @@ pub(super) struct AdminMonitoringCacheAffinityRecord {
pub(super) expire_at: Option<serde_json::Value>,
pub(super) request_count: u64,
pub(super) request_count_known: bool,
pub(super) scheduler_affinity_epoch: Option<u64>,
}
pub(super) struct AdminMonitoringCacheSnapshot {

View File

@@ -191,6 +191,24 @@ fn admin_pool_sort_keys_for_request(keys: &mut [StoredProviderCatalogKey], sort:
}
}
fn admin_pool_repository_key_order(sort: AdminPoolKeySort) -> ProviderCatalogKeyListOrder {
match (sort.field, sort.direction) {
(AdminPoolKeySortField::Default, _) => ProviderCatalogKeyListOrder::Name,
(AdminPoolKeySortField::ImportedAt, AdminPoolKeySortDirection::Asc) => {
ProviderCatalogKeyListOrder::CreatedAtAsc
}
(AdminPoolKeySortField::ImportedAt, AdminPoolKeySortDirection::Desc) => {
ProviderCatalogKeyListOrder::CreatedAtDesc
}
(AdminPoolKeySortField::LastUsedAt, AdminPoolKeySortDirection::Asc) => {
ProviderCatalogKeyListOrder::LastUsedAtAsc
}
(AdminPoolKeySortField::LastUsedAt, AdminPoolKeySortDirection::Desc) => {
ProviderCatalogKeyListOrder::LastUsedAtDesc
}
}
}
pub(super) async fn build_admin_pool_list_keys_response(
state: &AdminAppState<'_>,
request_context: &AdminRequestContext<'_>,
@@ -305,7 +323,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
.take(page_size)
.collect::<Vec<_>>();
(keys, total)
} else if !quick_selectors.is_empty() || sort.field != AdminPoolKeySortField::Default {
} else if !quick_selectors.is_empty() {
let mut keys = state
.list_provider_catalog_keys_by_provider_ids(std::slice::from_ref(&provider.id))
.await?
@@ -354,7 +372,7 @@ pub(super) async fn build_admin_pool_list_keys_response(
},
offset: page_offset,
limit: page_size,
order: ProviderCatalogKeyListOrder::Name,
order: admin_pool_repository_key_order(sort),
})
.await?;
(key_page.items, key_page.total)

View File

@@ -31,8 +31,10 @@ pub(crate) struct LocalExecutionCandidateMetadata {
pub(crate) candidate_group_id: Option<String>,
pub(crate) pool_key_index: Option<u32>,
pub(crate) pool_key_lease: Option<RuntimeLockLease>,
pub(crate) scheduler_affinity_epoch: Option<u64>,
}
pub(crate) const SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD: &str = "scheduler_affinity_epoch";
pub(crate) const POOL_KEY_LEASE_KEY_REPORT_FIELD: &str = "pool_key_lease_key";
pub(crate) const POOL_KEY_LEASE_OWNER_REPORT_FIELD: &str = "pool_key_lease_owner";
pub(crate) const POOL_KEY_LEASE_TOKEN_REPORT_FIELD: &str = "pool_key_lease_token";
@@ -65,6 +67,9 @@ pub(crate) fn local_execution_candidate_metadata_from_report_context(
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok()),
pool_key_lease: pool_key_lease_from_report_context(report_context),
scheduler_affinity_epoch: report_context
.and_then(|value| value.get(SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD))
.and_then(Value::as_u64),
}
}
@@ -493,6 +498,7 @@ mod tests {
token: "gateway-1:token-1".to_string(),
ttl_ms: 900000,
}),
scheduler_affinity_epoch: None,
}
);
}

View File

@@ -248,12 +248,16 @@ fn remember_successful_local_scheduler_affinity(
let Some(target) = local_scheduler_affinity_target(context.plan) else {
return;
};
let expected_epoch =
local_execution_candidate_metadata_from_report_context(context.report_context)
.scheduler_affinity_epoch;
state.remember_scheduler_affinity_target(
let _ = state.remember_scheduler_affinity_target_for_epoch(
&cache_key,
target,
SCHEDULER_AFFINITY_TTL,
LOCAL_EXECUTION_SCHEDULER_AFFINITY_MAX_ENTRIES,
expected_epoch,
);
}

View File

@@ -20,7 +20,7 @@ pub(crate) use self::attempt::{
attempt_identity_from_report_context, build_local_attempt_identities,
insert_pool_key_lease_report_context_fields, local_attempt_slot_count,
local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
LocalExecutionCandidateMetadata,
LocalExecutionCandidateMetadata, SCHEDULER_AFFINITY_EPOCH_REPORT_FIELD,
};
pub(crate) use self::classifier::{
classify_local_failover, local_failover_error_message, LocalFailoverClassification,

View File

@@ -58,6 +58,22 @@ fn report_context_key_id(report_context: Option<&Value>) -> Option<String> {
.map(ToOwned::to_owned)
}
fn report_context_provider_response_headers(
report_context: Option<&Value>,
) -> Option<BTreeMap<String, String>> {
let headers = report_context
.and_then(|context| context.get("provider_response_headers"))
.and_then(Value::as_object)?;
let mut out = BTreeMap::new();
for (key, value) in headers {
let Some(value) = value.as_str() else {
continue;
};
out.insert(key.clone(), value.to_string());
}
(!out.is_empty()).then_some(out)
}
fn is_volatile_compare_field(key: &str) -> bool {
key == "updated_at" || key.ends_with("_reset_seconds") || key.ends_with("_reset_after_seconds")
}
@@ -333,7 +349,12 @@ async fn sync_codex_quota_from_response_headers(
};
let now_unix_secs = current_unix_secs();
let Some(parsed) = admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs)
let provider_headers = report_context_provider_response_headers(report_context);
let parsed_from_provider_headers = provider_headers.as_ref().and_then(|headers| {
admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs)
});
let Some(parsed) = parsed_from_provider_headers
.or_else(|| admin_provider_quota_pure::parse_codex_usage_headers(headers, now_unix_secs))
else {
return Ok(false);
};

View File

@@ -47,13 +47,14 @@ pub(super) fn remember_scheduler_affinity(
affinity_cache_key: Option<&str>,
state: &(impl SchedulerRuntimeState + ?Sized),
candidate: &SchedulerMinimalCandidateSelectionCandidate,
expected_epoch: Option<u64>,
) {
let Some(cache_key) = affinity_cache_key else {
return;
};
let (provider_id, endpoint_id, key_id) = candidate_key(candidate);
state.remember_scheduler_affinity_target(
let _ = state.remember_scheduler_affinity_target_for_epoch(
cache_key,
SchedulerAffinityTarget {
provider_id,
@@ -62,5 +63,6 @@ pub(super) fn remember_scheduler_affinity(
},
SCHEDULER_AFFINITY_TTL,
SCHEDULER_AFFINITY_MAX_ENTRIES,
expected_epoch,
);
}

View File

@@ -46,6 +46,7 @@ pub(super) async fn select_minimal_candidate(
now_unix_secs: u64,
enable_model_directives: bool,
) -> Result<Option<SchedulerMinimalCandidateSelectionCandidate>, GatewayError> {
let affinity_epoch = runtime_state.scheduler_affinity_epoch();
let affinity_cache_key = build_scheduler_affinity_cache_key(
auth_snapshot,
api_format,
@@ -68,7 +69,12 @@ pub(super) async fn select_minimal_candidate(
.into_iter()
.next();
if let Some(candidate) = selected.as_ref() {
remember_scheduler_affinity(affinity_cache_key.as_deref(), runtime_state, candidate);
remember_scheduler_affinity(
affinity_cache_key.as_deref(),
runtime_state,
candidate,
Some(affinity_epoch),
);
}
Ok(selected)
}

View File

@@ -42,6 +42,8 @@ pub(crate) trait SchedulerRuntimeState {
ttl: Duration,
) -> Option<SchedulerAffinityTarget>;
fn scheduler_affinity_epoch(&self) -> u64;
fn remember_scheduler_affinity_target(
&self,
cache_key: &str,
@@ -50,6 +52,15 @@ pub(crate) trait SchedulerRuntimeState {
max_entries: usize,
);
fn remember_scheduler_affinity_target_for_epoch(
&self,
cache_key: &str,
target: SchedulerAffinityTarget,
ttl: Duration,
max_entries: usize,
expected_epoch: Option<u64>,
) -> bool;
async fn read_scheduler_ordering_config(&self)
-> Result<SchedulerOrderingConfig, GatewayError>;
}

View File

@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::sync::atomic::AtomicU64;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
@@ -60,6 +61,7 @@ pub struct AppState {
pub(crate) oauth_refresh: Arc<provider_transport::LocalOAuthRefreshCoordinator>,
pub(crate) direct_plan_bypass_cache: Arc<DirectPlanBypassCache>,
pub(crate) scheduler_affinity_cache: Arc<SchedulerAffinityCache>,
pub(crate) scheduler_affinity_epoch: Arc<AtomicU64>,
pub(crate) dashboard_response_cache: Arc<DashboardResponseCache>,
pub(crate) system_config_cache: Arc<SystemConfigCache>,
pub(crate) fallback_metrics: Arc<fallback_metrics::GatewayFallbackMetrics>,

View File

@@ -1,4 +1,5 @@
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::time::Duration;
@@ -75,6 +76,7 @@ impl AppState {
cache_key: &str,
target: &SchedulerAffinityTarget,
ttl: Duration,
epoch: u64,
) {
if self.runtime_state.is_memory() {
return;
@@ -85,6 +87,7 @@ impl AppState {
let cache_key = cache_key.to_string();
let runtime_state = self.runtime_state.clone();
let scheduler_affinity_epoch = self.scheduler_affinity_epoch.clone();
let provider_id = target.provider_id.clone();
let endpoint_id = target.endpoint_id.clone();
let key_id = target.key_id.clone();
@@ -93,6 +96,9 @@ impl AppState {
let expire_at = now_unix_secs.saturating_add(ttl_seconds);
handle.spawn(async move {
if scheduler_affinity_epoch.load(Ordering::Acquire) != epoch {
return;
}
let existing = runtime_state
.kv_get(&cache_key)
.await
@@ -117,8 +123,12 @@ impl AppState {
"created_at": created_at,
"expire_at": expire_at,
"request_count": request_count,
"scheduler_affinity_epoch": epoch,
});
if let Ok(serialized) = serde_json::to_string(&payload) {
if scheduler_affinity_epoch.load(Ordering::Acquire) != epoch {
return;
}
let _ = runtime_state
.kv_set(
&cache_key,
@@ -202,6 +212,7 @@ impl AppState {
oauth_refresh: Arc::new(provider_transport::LocalOAuthRefreshCoordinator::new()),
direct_plan_bypass_cache: Arc::new(DirectPlanBypassCache::default()),
scheduler_affinity_cache: Arc::new(SchedulerAffinityCache::default()),
scheduler_affinity_epoch: Arc::new(AtomicU64::new(0)),
dashboard_response_cache: Arc::new(DashboardResponseCache::default()),
system_config_cache: Arc::new(SystemConfigCache::default()),
fallback_metrics: Arc::new(fallback_metrics::GatewayFallbackMetrics::default()),
@@ -958,12 +969,29 @@ impl AppState {
self.scheduler_affinity_cache.remove(cache_key).is_some()
}
pub(crate) fn scheduler_affinity_epoch(&self) -> u64 {
self.scheduler_affinity_epoch.load(Ordering::Acquire)
}
pub(crate) fn invalidate_scheduler_affinity_cache(&self) -> u64 {
let next_epoch = self
.scheduler_affinity_epoch
.fetch_add(1, Ordering::AcqRel)
.saturating_add(1);
self.scheduler_affinity_cache.clear();
next_epoch
}
pub(crate) fn read_scheduler_affinity_target(
&self,
cache_key: &str,
ttl: Duration,
) -> Option<SchedulerAffinityTarget> {
self.scheduler_affinity_cache.get_fresh(cache_key, ttl)
self.scheduler_affinity_cache.get_fresh_for_epoch(
cache_key,
ttl,
self.scheduler_affinity_epoch(),
)
}
pub(crate) fn remember_scheduler_affinity_target(
@@ -973,16 +1001,45 @@ impl AppState {
ttl: Duration,
max_entries: usize,
) {
self.spawn_scheduler_affinity_redis_write(cache_key, &target, ttl);
self.scheduler_affinity_cache
.insert(cache_key.to_string(), target, ttl, max_entries);
let epoch = self.scheduler_affinity_epoch();
self.remember_scheduler_affinity_target_for_epoch(
cache_key,
target,
ttl,
max_entries,
Some(epoch),
);
}
pub(crate) fn remember_scheduler_affinity_target_for_epoch(
&self,
cache_key: &str,
target: SchedulerAffinityTarget,
ttl: Duration,
max_entries: usize,
expected_epoch: Option<u64>,
) -> bool {
let epoch = expected_epoch.unwrap_or_else(|| self.scheduler_affinity_epoch());
if self.scheduler_affinity_epoch() != epoch {
return false;
}
self.spawn_scheduler_affinity_redis_write(cache_key, &target, ttl, epoch);
self.scheduler_affinity_cache.insert_for_epoch(
cache_key.to_string(),
target,
ttl,
max_entries,
epoch,
);
true
}
pub(crate) fn list_scheduler_affinity_entries(
&self,
ttl: Duration,
) -> Vec<SchedulerAffinitySnapshotEntry> {
self.scheduler_affinity_cache.fresh_entries(ttl)
self.scheduler_affinity_cache
.fresh_entries_for_epoch(ttl, self.scheduler_affinity_epoch())
}
pub fn with_video_task_store_path(
@@ -1205,4 +1262,57 @@ mod tests {
Some(json!("new"))
);
}
#[test]
fn scheduler_affinity_epoch_blocks_stale_rewarm_after_invalidation() {
let state = AppState::new().expect("app state should build");
let cache_key = "scheduler_affinity:api-key-1:openai:chat:gpt-5";
let ttl = std::time::Duration::from_secs(300);
let first_target = crate::cache::SchedulerAffinityTarget {
provider_id: "provider-old".to_string(),
endpoint_id: "endpoint-old".to_string(),
key_id: "key-old".to_string(),
};
let next_target = crate::cache::SchedulerAffinityTarget {
provider_id: "provider-new".to_string(),
endpoint_id: "endpoint-new".to_string(),
key_id: "key-new".to_string(),
};
let initial_epoch = state.scheduler_affinity_epoch();
assert!(state.remember_scheduler_affinity_target_for_epoch(
cache_key,
first_target.clone(),
ttl,
16,
Some(initial_epoch),
));
assert_eq!(
state.read_scheduler_affinity_target(cache_key, ttl),
Some(first_target)
);
let next_epoch = state.invalidate_scheduler_affinity_cache();
assert!(!state.remember_scheduler_affinity_target_for_epoch(
cache_key,
next_target.clone(),
ttl,
16,
Some(initial_epoch),
));
assert_eq!(state.read_scheduler_affinity_target(cache_key, ttl), None);
assert!(state.remember_scheduler_affinity_target_for_epoch(
cache_key,
next_target.clone(),
ttl,
16,
Some(next_epoch),
));
assert_eq!(
state.read_scheduler_affinity_target(cache_key, ttl),
Some(next_target)
);
}
}

View File

@@ -332,6 +332,10 @@ impl SchedulerRuntimeState for AppState {
AppState::read_scheduler_affinity_target(self, cache_key, ttl)
}
fn scheduler_affinity_epoch(&self) -> u64 {
AppState::scheduler_affinity_epoch(self)
}
fn remember_scheduler_affinity_target(
&self,
cache_key: &str,
@@ -342,6 +346,24 @@ impl SchedulerRuntimeState for AppState {
AppState::remember_scheduler_affinity_target(self, cache_key, target, ttl, max_entries);
}
fn remember_scheduler_affinity_target_for_epoch(
&self,
cache_key: &str,
target: SchedulerAffinityTarget,
ttl: Duration,
max_entries: usize,
expected_epoch: Option<u64>,
) -> bool {
AppState::remember_scheduler_affinity_target_for_epoch(
self,
cache_key,
target,
ttl,
max_entries,
expected_epoch,
)
}
async fn read_scheduler_ordering_config(
&self,
) -> Result<crate::scheduler::config::SchedulerOrderingConfig, GatewayError> {

View File

@@ -801,6 +801,59 @@ mod tests {
assert_eq!(quota.get("updated_at"), quota.get("observed_at"));
}
#[tokio::test]
async fn submit_sync_report_updates_codex_quota_from_provider_response_headers() {
crate::orchestration::clear_local_report_effect_caches_for_tests();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider(
"provider-codex-sync-provider-headers",
"codex",
)],
Vec::new(),
vec![sample_provider_catalog_key(
"key-codex-sync-provider-headers",
"provider-codex-sync-provider-headers",
)],
));
let state = build_provider_catalog_test_state(Arc::clone(&provider_catalog_repository));
submit_sync_report(
&state,
GatewaySyncReportRequest {
trace_id: "trace-codex-reporting-sync-provider-headers".to_string(),
report_kind: "openai_responses_sync_success".to_string(),
report_context: Some(json!({
"request_id": "req-codex-reporting-sync-provider-headers",
"key_id": "key-codex-sync-provider-headers",
"provider_response_headers": sample_codex_paid_headers()
})),
status_code: 200,
headers: BTreeMap::new(),
body_json: None,
client_body_json: None,
body_base64: None,
telemetry: None,
},
)
.await
.expect("sync report should stay local");
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-sync-provider-headers".to_string()])
.await
.expect("keys should list");
let codex = reloaded[0]
.upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(serde_json::Value::as_object)
.expect("codex metadata should exist");
assert_eq!(codex.get("primary_used_percent"), Some(&json!(31.0)));
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
}
#[tokio::test]
async fn submit_stream_report_updates_codex_quota_from_response_headers() {
crate::orchestration::clear_local_report_effect_caches_for_tests();
@@ -866,6 +919,61 @@ mod tests {
assert_eq!(quota.get("updated_at"), quota.get("observed_at"));
}
#[tokio::test]
async fn submit_stream_report_updates_codex_quota_from_provider_response_headers() {
crate::orchestration::clear_local_report_effect_caches_for_tests();
let provider_catalog_repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider_catalog_provider(
"provider-codex-stream-provider-headers",
"codex",
)],
Vec::new(),
vec![sample_provider_catalog_key(
"key-codex-stream-provider-headers",
"provider-codex-stream-provider-headers",
)],
));
let state = build_provider_catalog_test_state(Arc::clone(&provider_catalog_repository));
submit_stream_report(
&state,
GatewayStreamReportRequest {
trace_id: "trace-codex-reporting-stream-provider-headers".to_string(),
report_kind: "openai_responses_stream_success".to_string(),
report_context: Some(json!({
"request_id": "req-codex-reporting-stream-provider-headers",
"key_id": "key-codex-stream-provider-headers",
"provider_response_headers": sample_codex_paid_headers()
})),
status_code: 200,
headers: BTreeMap::new(),
provider_body_base64: None,
provider_body_state: None,
client_body_base64: None,
client_body_state: None,
terminal_summary: None,
telemetry: None,
},
)
.await
.expect("stream report should stay local");
let reloaded = provider_catalog_repository
.list_keys_by_ids(&["key-codex-stream-provider-headers".to_string()])
.await
.expect("keys should list");
let codex = reloaded[0]
.upstream_metadata
.as_ref()
.and_then(serde_json::Value::as_object)
.and_then(|metadata| metadata.get("codex"))
.and_then(serde_json::Value::as_object)
.expect("codex metadata should exist");
assert_eq!(codex.get("primary_used_percent"), Some(&json!(31.0)));
assert_eq!(codex.get("secondary_used_percent"), Some(&json!(100.0)));
}
#[tokio::test]
async fn submit_sync_report_stores_gemini_file_mapping_locally_when_payload_contains_file_json()
{

View File

@@ -516,6 +516,10 @@ pub enum ProviderCatalogKeyListOrder {
#[default]
Name,
CreatedAt,
CreatedAtAsc,
CreatedAtDesc,
LastUsedAtAsc,
LastUsedAtDesc,
}
#[derive(Debug, Clone, PartialEq, Eq, Default)]

View File

@@ -2657,6 +2657,22 @@ CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active ON public.provi
--
-- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_created_at_desc ON public.provider_api_keys USING btree (provider_id, created_at DESC NULLS LAST, name, id);
--
-- Name: idx_provider_api_keys_provider_last_used_at_desc; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_last_used_at_desc ON public.provider_api_keys USING btree (provider_id, last_used_at DESC NULLS LAST, name, id);
--
-- Name: idx_provider_api_keys_provider_id; Type: INDEX; Schema: public; Owner: -
--

View File

@@ -0,0 +1,5 @@
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_created_at_desc
ON public.provider_api_keys USING btree (provider_id, created_at DESC NULLS LAST, name, id);
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_last_used_at_desc
ON public.provider_api_keys USING btree (provider_id, last_used_at DESC NULLS LAST, name, id);

View File

@@ -133,6 +133,22 @@ CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active ON public.provi
--
-- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_created_at_desc ON public.provider_api_keys USING btree (provider_id, created_at DESC NULLS LAST, name, id);
--
-- Name: idx_provider_api_keys_provider_last_used_at_desc; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_last_used_at_desc ON public.provider_api_keys USING btree (provider_id, last_used_at DESC NULLS LAST, name, id);
--
-- Name: idx_provider_api_keys_provider_id; Type: INDEX; Schema: public; Owner: -
--

View File

@@ -125,6 +125,22 @@ CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_active ON public.provi
--
-- Name: idx_provider_api_keys_provider_created_at_desc; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_created_at_desc ON public.provider_api_keys USING btree (provider_id, created_at DESC NULLS LAST, name, id);
--
-- Name: idx_provider_api_keys_provider_last_used_at_desc; Type: INDEX; Schema: public; Owner: -
--
CREATE INDEX IF NOT EXISTS idx_provider_api_keys_provider_last_used_at_desc ON public.provider_api_keys USING btree (provider_id, last_used_at DESC NULLS LAST, name, id);
--
-- Name: idx_provider_api_keys_provider_id; Type: INDEX; Schema: public; Owner: -
--

View File

@@ -7,7 +7,7 @@ use tracing::info;
// Generated by build.rs from schema/bootstrap/postgres.
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260510120000;
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260511000000;
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
SELECT COUNT(*)::BIGINT

View File

@@ -298,6 +298,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
20260509120000,
20260510000000,
20260510120000,
20260511000000,
]
);
}
@@ -1082,6 +1083,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
20260509120000,
20260510000000,
20260510120000,
20260511000000,
]
);
}

View File

@@ -1,3 +1,4 @@
use std::cmp::Ordering;
use std::collections::BTreeMap;
use std::sync::RwLock;
@@ -454,6 +455,20 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
})
.cloned()
.collect::<Vec<_>>();
fn compare_optional_u64_null_last(
left: Option<u64>,
right: Option<u64>,
descending: bool,
) -> Ordering {
match (left, right) {
(Some(left), Some(right)) if descending => right.cmp(&left),
(Some(left), Some(right)) => left.cmp(&right),
(Some(_), None) => Ordering::Less,
(None, Some(_)) => Ordering::Greater,
(None, None) => Ordering::Equal,
}
}
match query.order {
ProviderCatalogKeyListOrder::Name => {
keys.sort_by(|left, right| {
@@ -475,6 +490,50 @@ impl ProviderCatalogReadRepository for InMemoryProviderCatalogReadRepository {
.then(left.id.cmp(&right.id))
});
}
ProviderCatalogKeyListOrder::CreatedAtAsc => {
keys.sort_by(|left, right| {
compare_optional_u64_null_last(
left.created_at_unix_ms,
right.created_at_unix_ms,
false,
)
.then(left.name.cmp(&right.name))
.then(left.id.cmp(&right.id))
});
}
ProviderCatalogKeyListOrder::CreatedAtDesc => {
keys.sort_by(|left, right| {
compare_optional_u64_null_last(
left.created_at_unix_ms,
right.created_at_unix_ms,
true,
)
.then(left.name.cmp(&right.name))
.then(left.id.cmp(&right.id))
});
}
ProviderCatalogKeyListOrder::LastUsedAtAsc => {
keys.sort_by(|left, right| {
compare_optional_u64_null_last(
left.last_used_at_unix_secs,
right.last_used_at_unix_secs,
false,
)
.then(left.name.cmp(&right.name))
.then(left.id.cmp(&right.id))
});
}
ProviderCatalogKeyListOrder::LastUsedAtDesc => {
keys.sort_by(|left, right| {
compare_optional_u64_null_last(
left.last_used_at_unix_secs,
right.last_used_at_unix_secs,
true,
)
.then(left.name.cmp(&right.name))
.then(left.id.cmp(&right.id))
});
}
}
let total = keys.len();
let items = keys
@@ -1028,6 +1087,70 @@ mod tests {
);
}
#[tokio::test]
async fn paginates_provider_keys_by_pool_sort_fields() {
let mut old = sample_key("key-1", "provider-1");
old.name = "old".to_string();
old.created_at_unix_ms = Some(10);
old.last_used_at_unix_secs = Some(30);
let mut fresh = sample_key("key-2", "provider-1");
fresh.name = "fresh".to_string();
fresh.created_at_unix_ms = Some(20);
fresh.last_used_at_unix_secs = Some(10);
let mut unused = sample_key("key-3", "provider-1");
unused.name = "unused".to_string();
unused.created_at_unix_ms = None;
unused.last_used_at_unix_secs = None;
let repository = InMemoryProviderCatalogReadRepository::seed(
vec![sample_provider("provider-1")],
vec![],
vec![old, fresh, unused],
);
let imported = repository
.list_keys_page(&ProviderCatalogKeyListQuery {
provider_id: "provider-1".to_string(),
search: None,
is_active: None,
offset: 0,
limit: 2,
order: ProviderCatalogKeyListOrder::CreatedAtDesc,
})
.await
.expect("keys should page");
assert_eq!(imported.total, 3);
assert_eq!(
imported
.items
.iter()
.map(|item| item.id.as_str())
.collect::<Vec<_>>(),
vec!["key-2", "key-1"]
);
let last_used = repository
.list_keys_page(&ProviderCatalogKeyListQuery {
provider_id: "provider-1".to_string(),
search: None,
is_active: None,
offset: 0,
limit: 3,
order: ProviderCatalogKeyListOrder::LastUsedAtDesc,
})
.await
.expect("keys should page");
assert_eq!(
last_used
.items
.iter()
.map(|item| item.id.as_str())
.collect::<Vec<_>>(),
vec!["key-1", "key-2", "key-3"]
);
}
#[tokio::test]
async fn summarizes_provider_key_stats() {
let mut inactive = sample_key("key-2", "provider-1");

View File

@@ -571,6 +571,18 @@ ORDER BY provider_priority ASC, name ASC
ProviderCatalogKeyListOrder::CreatedAt => {
"internal_priority ASC, COALESCE(created_at, TO_TIMESTAMP(0)) ASC, id ASC"
}
ProviderCatalogKeyListOrder::CreatedAtAsc => {
"created_at ASC NULLS LAST, name ASC, id ASC"
}
ProviderCatalogKeyListOrder::CreatedAtDesc => {
"created_at DESC NULLS LAST, name ASC, id ASC"
}
ProviderCatalogKeyListOrder::LastUsedAtAsc => {
"last_used_at ASC NULLS LAST, name ASC, id ASC"
}
ProviderCatalogKeyListOrder::LastUsedAtDesc => {
"last_used_at DESC NULLS LAST, name ASC, id ASC"
}
};
let count_row = sqlx::query(