Fix provider pool exhaustion scheduling

This commit is contained in:
fawney19
2026-05-20 20:13:07 +08:00
parent d0981c2fd5
commit c972bbd397
9 changed files with 424 additions and 48 deletions

View File

@@ -709,31 +709,39 @@ impl<'a> PoolKeyCursor<'a> {
} }
async fn refill_queued_candidates(&mut self) -> bool { async fn refill_queued_candidates(&mut self) -> bool {
let mut candidates = Vec::new();
let refill_target = self.window_size.max(1) as usize; let refill_target = self.window_size.max(1) as usize;
// Keep pool expansion bounded; the cursor freezes one small window at a time.
while candidates.len() < refill_target {
let Some(mut page_candidates) = self.next_page_candidates().await else {
break;
};
candidates.append(&mut page_candidates);
}
if candidates.is_empty() { loop {
return false; let mut candidates = Vec::new();
} // Keep pool expansion bounded; the cursor freezes one small window at a time.
while candidates.len() < refill_target {
let Some(mut page_candidates) = self.next_page_candidates().await else {
break;
};
candidates.append(&mut page_candidates);
}
let (mut scheduled, mut skipped) = schedule_pool_page_candidates( if candidates.is_empty() {
self.state, return false;
candidates, }
self.sticky_session_token.as_deref(),
) let (mut scheduled, mut skipped) = schedule_pool_page_candidates(
.await; self.state,
scheduled.truncate(refill_target); candidates,
self.record_skipped_candidates(&skipped); self.sticky_session_token.as_deref(),
self.queued_candidates.extend(scheduled.drain(..)); )
self.skipped_candidates.append(&mut skipped); .await;
!self.queued_candidates.is_empty() self.record_skipped_candidates(&skipped);
self.skipped_candidates.append(&mut skipped);
if scheduled.is_empty() {
continue;
}
scheduled.truncate(refill_target);
self.queued_candidates.extend(scheduled.drain(..));
return true;
}
} }
async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> { async fn next_queued_candidate(&mut self) -> Option<EligibleLocalExecutionCandidate> {
@@ -2798,6 +2806,78 @@ mod tests {
assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&1)); assert_eq!(cursor.skip_reason_counts.get("pool_cooldown"), Some(&1));
} }
#[tokio::test]
async fn pool_key_cursor_continues_after_exhausted_window() {
let provider_config = Some(json!({
"pool_advanced": {
"skip_exhausted_accounts": true
}
}));
let (provider, endpoint, mut keys, rows) = large_pool_fixture(3, provider_config.clone());
for key in keys.iter_mut().take(2) {
key.status_snapshot = Some(json!({
"quota": {
"provider_type": "openai",
"exhausted": true,
"usage_ratio": 1.0,
"windows": [
{
"code": "daily",
"used_ratio": 1.0,
"remaining_ratio": 0.0
}
]
}
}));
}
let data_state =
GatewayDataState::with_provider_catalog_and_minimal_candidate_selection_for_tests(
Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![provider],
vec![endpoint],
keys,
)),
Arc::new(InMemoryMinimalCandidateSelectionReadRepository::seed(rows)),
)
.with_encryption_key_for_tests(aether_crypto::DEVELOPMENT_ENCRYPTION_KEY);
let app = AppState::new()
.expect("state should build")
.with_data_state_for_tests(data_state);
let group = sample_eligible_candidate(
"provider-pool",
"endpoint-1",
"pool-group",
10,
provider_config,
);
let mut cursor = PoolKeyCursor::new(PlannerAppState::new(&app), group, None, None, None);
cursor.window_size = 2;
cursor.page_size = 2;
cursor.max_scanned_keys = 4;
let candidate = cursor
.next_key()
.await
.expect("cursor should scan past an exhausted window");
assert_eq!(candidate.candidate.key_id, "key-00002");
assert_eq!(candidate.orchestration.pool_key_index, Some(0));
assert!(candidate.orchestration.pool_key_lease.is_none());
assert_eq!(
cursor
.skip_reason_counts
.get(aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON),
Some(&2)
);
let skipped = cursor.take_skipped_candidates();
assert_eq!(skipped.len(), 2);
assert!(skipped.iter().all(|candidate| {
candidate.skip_reason == aether_pool_core::POOL_ACCOUNT_EXHAUSTED_SKIP_REASON
}));
}
#[tokio::test] #[tokio::test]
async fn pool_key_cursor_simulates_large_lru_pool_with_lazy_pages_and_dynamic_skips() { async fn pool_key_cursor_simulates_large_lru_pool_with_lazy_pages_and_dynamic_skips() {
const KEY_COUNT: usize = 2048; const KEY_COUNT: usize = 2048;

View File

@@ -2,7 +2,9 @@ use super::state::{
decode_jwt_claims, enrich_admin_provider_oauth_auth_config, json_non_empty_string, decode_jwt_claims, enrich_admin_provider_oauth_auth_config, json_non_empty_string,
json_u64_value, json_u64_value,
}; };
use crate::handlers::admin::admin_provider_pool_config;
use crate::handlers::admin::request::AdminAppState; use crate::handlers::admin::request::AdminAppState;
use crate::maintenance::ensure_provider_key_pool_scores_for_keys;
use crate::provider_key_auth::provider_active_api_formats; use crate::provider_key_auth::provider_active_api_formats;
use crate::GatewayError; use crate::GatewayError;
use aether_data_contracts::repository::provider_catalog::{ use aether_data_contracts::repository::provider_catalog::{
@@ -168,6 +170,7 @@ pub(crate) async fn create_provider_oauth_catalog_key(
.app() .app()
.invalidate_local_oauth_refresh_entry(&key.id) .invalidate_local_oauth_refresh_entry(&key.id)
.await; .await;
seed_provider_oauth_pool_score(state, provider_id, key, now_unix_secs).await;
} }
Ok(created) Ok(created)
} }
@@ -221,10 +224,75 @@ pub(crate) async fn update_existing_provider_oauth_catalog_key(
.app() .app()
.invalidate_local_oauth_refresh_entry(&key.id) .invalidate_local_oauth_refresh_entry(&key.id)
.await; .await;
seed_provider_oauth_pool_score(state, &existing_key.provider_id, key, now_unix_secs).await;
} }
Ok(persisted) Ok(persisted)
} }
async fn seed_provider_oauth_pool_score(
state: &AdminAppState<'_>,
provider_id: &str,
key: &StoredProviderCatalogKey,
now_unix_secs: u64,
) {
let provider_id = provider_id.to_string();
let provider = match state
.read_provider_catalog_providers_by_ids(std::slice::from_ref(&provider_id))
.await
{
Ok(mut providers) => providers.pop(),
Err(err) => {
tracing::debug!(
provider_id = %provider_id,
key_id = %key.id,
error = ?err,
"gateway provider oauth provisioning: failed to read provider for pool score seed"
);
return;
}
};
let Some(provider) = provider else {
return;
};
let Some(pool_config) = admin_provider_pool_config(&provider) else {
return;
};
let endpoints = match state
.list_provider_catalog_endpoints_by_provider_ids(std::slice::from_ref(&provider_id))
.await
{
Ok(endpoints) => endpoints,
Err(err) => {
tracing::debug!(
provider_id = %provider_id,
key_id = %key.id,
error = ?err,
"gateway provider oauth provisioning: failed to read endpoints for pool score seed"
);
return;
}
};
let score_ensure_budget = (pool_config.score_fallback_scan_limit as usize).clamp(1, 50_000);
if let Err(err) = ensure_provider_key_pool_scores_for_keys(
state.as_ref(),
&provider,
&pool_config,
&endpoints,
std::slice::from_ref(key),
now_unix_secs,
score_ensure_budget,
)
.await
{
tracing::debug!(
provider_id = %provider_id,
key_id = %key.id,
error = ?err,
"gateway provider oauth provisioning: failed to seed pool score row"
);
}
}
fn provider_oauth_catalog_key_api_formats( fn provider_oauth_catalog_key_api_formats(
provider_type: &str, provider_type: &str,
api_formats: &[String], api_formats: &[String],

View File

@@ -12,8 +12,12 @@ use aether_data::repository::management_tokens::{
use aether_data::repository::oauth_providers::{ use aether_data::repository::oauth_providers::{
InMemoryOAuthProviderRepository, OAuthProviderReadRepository, InMemoryOAuthProviderRepository, OAuthProviderReadRepository,
}; };
use aether_data::repository::pool_scores::InMemoryPoolMemberScoreRepository;
use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository; use aether_data::repository::provider_catalog::InMemoryProviderCatalogReadRepository;
use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository; use aether_data::repository::proxy_nodes::InMemoryProxyNodeRepository;
use aether_data_contracts::repository::pool_scores::{
GetPoolMemberScoresByIdsQuery, PoolMemberHardState, PoolMemberIdentity, PoolScoreReadRepository,
};
use aether_data_contracts::repository::provider_catalog::{ use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint, ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
}; };
@@ -32,6 +36,7 @@ use super::super::{
use crate::admin_api::{ use crate::admin_api::{
maybe_build_local_admin_provider_oauth_response, AdminAppState, AdminRequestContext, maybe_build_local_admin_provider_oauth_response, AdminAppState, AdminRequestContext,
}; };
use crate::ai_serving::{provider_key_pool_score_id, provider_key_pool_score_scope};
use crate::audit::AdminAuditEvent; use crate::audit::AdminAuditEvent;
use crate::constants::{ use crate::constants::{
GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER, GATEWAY_HEADER, TRUSTED_ADMIN_MANAGEMENT_TOKEN_ID_HEADER, TRUSTED_ADMIN_SESSION_ID_HEADER,
@@ -1930,6 +1935,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
let mut provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 10); let mut provider = sample_provider("provider-chatgpt-web", "chatgpt_web", 10);
provider.provider_type = "chatgpt_web".to_string(); provider.provider_type = "chatgpt_web".to_string();
provider.config = Some(json!({"pool_advanced": {}}));
let endpoint = sample_endpoint( let endpoint = sample_endpoint(
"endpoint-chatgpt-web-image", "endpoint-chatgpt-web-image",
"provider-chatgpt-web", "provider-chatgpt-web",
@@ -1941,6 +1947,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
vec![endpoint], vec![endpoint],
vec![], vec![],
)); ));
let pool_score_repository = Arc::new(InMemoryPoolMemberScoreRepository::default());
let (token_url, token_handle) = start_server(token_server).await; let (token_url, token_handle) = start_server(token_server).await;
let gateway = build_router_with_state( let gateway = build_router_with_state(
@@ -1950,6 +1957,7 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
GatewayDataState::with_provider_catalog_repository_for_tests( GatewayDataState::with_provider_catalog_repository_for_tests(
provider_catalog_repository.clone(), provider_catalog_repository.clone(),
) )
.with_pool_score_repository_for_tests(Arc::clone(&pool_score_repository))
.with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY), .with_encryption_key_for_tests(DEVELOPMENT_ENCRYPTION_KEY),
) )
.with_provider_oauth_token_url_for_tests( .with_provider_oauth_token_url_for_tests(
@@ -2030,6 +2038,20 @@ async fn gateway_batch_imports_chatgpt_web_access_tokens_with_pool_hints() {
assert_eq!(auth_config["plan_type"], "plus"); assert_eq!(auth_config["plan_type"], "plus");
assert_eq!(auth_config["user_id"], "user-pool-image"); assert_eq!(auth_config["user_id"], "user-pool-image");
let score_scope = provider_key_pool_score_scope();
let score_identity =
PoolMemberIdentity::provider_api_key("provider-chatgpt-web", persisted.id.clone());
let scores = pool_score_repository
.get_pool_member_scores_by_ids(&GetPoolMemberScoresByIdsQuery {
ids: vec![provider_key_pool_score_id(&score_identity, &score_scope)],
})
.await
.expect("pool score should load");
assert_eq!(scores.len(), 1);
assert_eq!(scores[0].member_id, persisted.id);
assert_eq!(scores[0].hard_state, PoolMemberHardState::Unknown);
assert!(scores[0].score > 0.0);
gateway_handle.abort(); gateway_handle.abort();
token_handle.abort(); token_handle.abort();
} }

View File

@@ -380,6 +380,81 @@ mod tests {
)); ));
} }
#[test]
fn provider_quota_exhaustion_snapshot_expires_after_reset_at() {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time should be after unix epoch")
.as_secs();
let mut expired = sample_key(None);
expired.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "exhausted",
"exhausted": true,
"updated_at": now.saturating_sub(600),
"windows": [{
"code": "5h",
"used_ratio": 1.0,
"reset_at": now.saturating_sub(60),
"is_exhausted": true
}]
}
}));
assert!(!provider_pool_key_account_quota_exhausted(
&expired, "codex"
));
let mut active = sample_key(None);
active.status_snapshot = Some(json!({
"quota": {
"version": 2,
"provider_type": "codex",
"code": "exhausted",
"exhausted": true,
"updated_at": now,
"windows": [{
"code": "5h",
"used_ratio": 1.0,
"reset_at": now.saturating_add(3600),
"is_exhausted": true
}]
}
}));
assert!(provider_pool_key_account_quota_exhausted(&active, "codex"));
}
#[test]
fn provider_quota_exhaustion_metadata_expires_after_reset_at() {
let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.expect("system time should be after unix epoch")
.as_secs();
assert!(!provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
"updated_at": now.saturating_sub(600),
"primary_used_percent": 100.0,
"primary_reset_at": now.saturating_sub(60)
}
}))),
"codex",
));
assert!(provider_pool_key_account_quota_exhausted(
&sample_key(Some(json!({
"codex": {
"updated_at": now,
"primary_used_percent": 100.0,
"primary_reset_at": now.saturating_add(3600)
}
}))),
"codex",
));
}
#[test] #[test]
fn grok_quota_tier_boundaries_match_pool_modes() { fn grok_quota_tier_boundaries_match_pool_modes() {
assert_eq!( assert_eq!(

View File

@@ -10,8 +10,9 @@ use crate::provider::{
ProviderPoolMemberInput, ProviderPoolMemberInput,
}; };
use crate::quota::{ use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket, provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
}; };
use crate::quota_refresh::ProviderPoolQuotaRequestSpec; use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -255,6 +256,19 @@ fn infer_chatgpt_web_image_quota_limit(
} }
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool { pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if provider_pool_current_unix_secs().is_some_and(|now| {
let mut image_quota = Map::new();
if let Some(value) = bucket.get("image_quota_reset_at") {
image_quota.insert("reset_at".to_string(), value.clone());
}
provider_pool_reset_deadline_elapsed(
&image_quota,
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
now,
)
}) {
return false;
}
if provider_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) { if provider_pool_json_bool(bucket.get("image_quota_blocked")) == Some(true) {
return true; return true;
} }

View File

@@ -10,8 +10,9 @@ use crate::provider::{
ProviderPoolMemberInput, ProviderPoolMemberInput,
}; };
use crate::quota::{ use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket, provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
}; };
use crate::quota_refresh::ProviderPoolQuotaRequestSpec; use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -121,6 +122,37 @@ pub fn build_codex_pool_quota_request(
}) })
} }
fn codex_window_reset_elapsed(bucket: &Map<String, Value>, prefix: &str) -> bool {
let Some(now_unix_secs) = provider_pool_current_unix_secs() else {
return false;
};
let mut window = Map::new();
for (target, source) in [
("reset_at", format!("{prefix}_reset_at")),
("next_reset_at", format!("{prefix}_next_reset_at")),
("reset_seconds", format!("{prefix}_reset_seconds")),
(
"reset_after_seconds",
format!("{prefix}_reset_after_seconds"),
),
] {
if let Some(value) = bucket.get(source.as_str()) {
window.insert(target.to_string(), value.clone());
}
}
provider_pool_reset_deadline_elapsed(
&window,
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
now_unix_secs,
)
}
fn codex_window_used_percent_exhausted(bucket: &Map<String, Value>, prefix: &str) -> bool {
let used_percent_key = format!("{prefix}_used_percent");
provider_pool_json_f64(bucket.get(used_percent_key.as_str()))
.is_some_and(|value| value >= 100.0 && !codex_window_reset_elapsed(bucket, prefix))
}
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool { pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) { if provider_pool_json_bool(bucket.get("credits_unlimited")) == Some(true) {
return false; return false;
@@ -130,7 +162,6 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if !has_window_data && provider_pool_json_bool(bucket.get("has_credits")) == Some(false) { if !has_window_data && provider_pool_json_bool(bucket.get("has_credits")) == Some(false) {
return true; return true;
} }
provider_pool_json_f64(bucket.get("primary_used_percent")).is_some_and(|value| value >= 100.0) codex_window_used_percent_exhausted(bucket, "primary")
|| provider_pool_json_f64(bucket.get("secondary_used_percent")) || codex_window_used_percent_exhausted(bucket, "secondary")
.is_some_and(|value| value >= 100.0)
} }

View File

@@ -7,8 +7,9 @@ use crate::provider::{
ProviderPoolMemberInput, ProviderPoolMemberInput,
}; };
use crate::quota::{ use crate::quota::{
provider_pool_json_bool, provider_pool_json_f64, provider_pool_metadata_bucket, provider_pool_current_unix_secs, provider_pool_json_bool, provider_pool_json_f64,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_metadata_bucket, provider_pool_quota_snapshot_exhausted_decision,
provider_pool_reset_deadline_elapsed, provider_pool_timestamp_unix_secs,
}; };
pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")]; pub const GROK_QUOTA_WINDOWS_BASIC: &[(&str, &str)] = &[("quota_fast", "fast")];
@@ -191,6 +192,8 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
let mut model_count = 0usize; let mut model_count = 0usize;
let mut exhausted_count = 0usize; let mut exhausted_count = 0usize;
let now_unix_secs = provider_pool_current_unix_secs();
let bucket_updated_at = provider_pool_timestamp_unix_secs(bucket.get("updated_at"));
for (model_key, item) in models.iter() { for (model_key, item) in models.iter() {
if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) { if !supported_mode_keys.is_empty() && !supported_mode_keys.contains(&model_key.as_str()) {
continue; continue;
@@ -207,12 +210,14 @@ pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
continue; continue;
} }
model_count += 1; model_count += 1;
if provider_pool_json_bool(item.get("is_exhausted")) == Some(true) let quota_exhausted = provider_pool_json_bool(item.get("is_exhausted")) == Some(true)
|| provider_pool_json_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0) || provider_pool_json_f64(item.get("used_percent")).is_some_and(|value| value >= 100.0)
|| provider_pool_json_f64(item.get("remaining")).is_some_and(|value| value <= 0.0) || provider_pool_json_f64(item.get("remaining")).is_some_and(|value| value <= 0.0)
|| provider_pool_json_f64(item.get("remaining_fraction")) || provider_pool_json_f64(item.get("remaining_fraction"))
.is_some_and(|value| value <= 0.0) .is_some_and(|value| value <= 0.0);
{ let reset_elapsed = now_unix_secs
.is_some_and(|now| provider_pool_reset_deadline_elapsed(item, bucket_updated_at, now));
if quota_exhausted && !reset_elapsed {
exhausted_count += 1; exhausted_count += 1;
} }
} }

View File

@@ -11,8 +11,9 @@ use crate::provider::{
ProviderPoolMemberInput, ProviderPoolMemberInput,
}; };
use crate::quota::{ use crate::quota::{
provider_pool_json_f64, provider_pool_metadata_bucket, provider_pool_current_unix_secs, provider_pool_json_f64, provider_pool_metadata_bucket,
provider_pool_quota_snapshot_exhausted_decision, provider_pool_quota_snapshot_exhausted_decision, provider_pool_reset_deadline_elapsed,
provider_pool_timestamp_unix_secs,
}; };
use crate::quota_refresh::ProviderPoolQuotaRequestSpec; use crate::quota_refresh::ProviderPoolQuotaRequestSpec;
@@ -151,6 +152,15 @@ fn normalize_kiro_version(value: &str) -> &str {
} }
pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool { pub(crate) fn quota_exhausted_from_bucket(bucket: &Map<String, Value>) -> bool {
if provider_pool_current_unix_secs().is_some_and(|now| {
provider_pool_reset_deadline_elapsed(
bucket,
provider_pool_timestamp_unix_secs(bucket.get("updated_at")),
now,
)
}) {
return false;
}
if provider_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) { if provider_pool_json_f64(bucket.get("remaining")).is_some_and(|value| value <= 0.0) {
return true; return true;
} }

View File

@@ -1,3 +1,5 @@
use std::time::{SystemTime, UNIX_EPOCH};
use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey; use aether_data_contracts::repository::provider_catalog::StoredProviderCatalogKey;
use serde_json::{json, Map, Value}; use serde_json::{json, Map, Value};
@@ -135,7 +137,7 @@ pub(crate) fn provider_pool_json_f64(value: Option<&Value>) -> Option<f64> {
.filter(|value| value.is_finite()) .filter(|value| value.is_finite())
} }
fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> { pub(crate) fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
let mut timestamp = provider_pool_json_f64(value)?; let mut timestamp = provider_pool_json_f64(value)?;
if timestamp <= 0.0 { if timestamp <= 0.0 {
return None; return None;
@@ -146,6 +148,49 @@ fn provider_pool_timestamp_unix_secs(value: Option<&Value>) -> Option<u64> {
Some(timestamp as u64) Some(timestamp as u64)
} }
pub(crate) fn provider_pool_current_unix_secs() -> Option<u64> {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.ok()
.map(|duration| duration.as_secs())
}
fn provider_pool_reset_deadline_unix_secs(
item: &Map<String, Value>,
fallback_observed_at: Option<u64>,
) -> Option<u64> {
provider_pool_timestamp_unix_secs(item.get("reset_at"))
.or_else(|| provider_pool_timestamp_unix_secs(item.get("next_reset_at")))
.or_else(|| {
let reset_seconds = provider_pool_json_f64(item.get("reset_seconds"))
.or_else(|| provider_pool_json_f64(item.get("reset_after_seconds")))?;
if reset_seconds < 0.0 {
return None;
}
let base = provider_pool_timestamp_unix_secs(item.get("observed_at"))
.or_else(|| provider_pool_timestamp_unix_secs(item.get("updated_at")))
.or(fallback_observed_at)?;
Some(base.saturating_add(reset_seconds.ceil() as u64))
})
}
pub(crate) fn provider_pool_reset_deadline_elapsed(
item: &Map<String, Value>,
fallback_observed_at: Option<u64>,
now_unix_secs: u64,
) -> bool {
provider_pool_reset_deadline_unix_secs(item, fallback_observed_at)
.is_some_and(|reset_at| reset_at <= now_unix_secs)
}
fn provider_pool_quota_window_is_exhausted(window: &Map<String, Value>) -> bool {
provider_pool_json_bool(window.get("is_exhausted"))
.or_else(|| {
provider_pool_json_f64(window.get("used_ratio")).map(|value| value >= 1.0 - 1e-6)
})
.unwrap_or(false)
}
fn provider_pool_quota_snapshot_matches_provider( fn provider_pool_quota_snapshot_matches_provider(
quota_snapshot: &Map<String, Value>, quota_snapshot: &Map<String, Value>,
provider_type: &str, provider_type: &str,
@@ -203,19 +248,45 @@ pub(crate) fn provider_pool_quota_snapshot_exhausted_decision(
} }
let exhausted = provider_pool_json_bool(quota_snapshot.get("exhausted"))?; let exhausted = provider_pool_json_bool(quota_snapshot.get("exhausted"))?;
if exhausted { if exhausted {
let windows_max_ratio = quota_snapshot let now_unix_secs = provider_pool_current_unix_secs();
let snapshot_observed_at =
provider_pool_timestamp_unix_secs(quota_snapshot.get("observed_at"))
.or_else(|| provider_pool_timestamp_unix_secs(quota_snapshot.get("updated_at")));
if let Some(windows) = quota_snapshot
.get("windows") .get("windows")
.and_then(Value::as_array) .and_then(Value::as_array)
.filter(|w| !w.is_empty()) .filter(|windows| !windows.is_empty())
.and_then(|windows| { {
windows let mut saw_exhausted_window = false;
.iter() let mut saw_active_exhausted_window = false;
.filter_map(Value::as_object) let mut windows_max_ratio = None::<f64>;
.filter_map(|w| w.get("used_ratio"))
.filter_map(Value::as_f64) for window in windows.iter().filter_map(Value::as_object) {
.max_by(f64::total_cmp) if let Some(ratio) = provider_pool_json_f64(window.get("used_ratio")) {
}); windows_max_ratio =
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) { Some(windows_max_ratio.map_or(ratio, |current| current.max(ratio)));
}
if provider_pool_quota_window_is_exhausted(window) {
saw_exhausted_window = true;
let reset_elapsed = now_unix_secs.is_some_and(|now| {
provider_pool_reset_deadline_elapsed(window, snapshot_observed_at, now)
});
if !reset_elapsed {
saw_active_exhausted_window = true;
}
}
}
if saw_exhausted_window {
return Some(saw_active_exhausted_window);
}
if windows_max_ratio.is_some_and(|ratio| ratio < 1.0 - 1e-6) {
return Some(false);
}
} else if now_unix_secs.is_some_and(|now| {
provider_pool_reset_deadline_elapsed(quota_snapshot, snapshot_observed_at, now)
}) {
return Some(false); return Some(false);
} }
} }