mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 01:10:23 +08:00
Track scheduler affinity epochs and key sorting
This commit is contained in:
@@ -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,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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 {},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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),
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -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);
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
);
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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);
|
||||
};
|
||||
|
||||
@@ -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,
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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>;
|
||||
}
|
||||
|
||||
@@ -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>,
|
||||
|
||||
@@ -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)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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> {
|
||||
|
||||
@@ -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()
|
||||
{
|
||||
|
||||
@@ -516,6 +516,10 @@ pub enum ProviderCatalogKeyListOrder {
|
||||
#[default]
|
||||
Name,
|
||||
CreatedAt,
|
||||
CreatedAtAsc,
|
||||
CreatedAtDesc,
|
||||
LastUsedAtAsc,
|
||||
LastUsedAtDesc,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, PartialEq, Eq, Default)]
|
||||
|
||||
@@ -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: -
|
||||
--
|
||||
|
||||
@@ -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);
|
||||
@@ -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: -
|
||||
--
|
||||
|
||||
@@ -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: -
|
||||
--
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
@@ -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");
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user