use super::{ read_decision_trace, read_provider_transport_snapshot, read_request_candidate_trace, AdjustWalletBalanceInput, AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingMutationOutcome, AdminBillingPresetApplyResult, AdminBillingRuleRecord, AdminBillingRuleWriteInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery, AdminRedeemCodeListQuery, AdminWalletLedgerQuery, AdminWalletListQuery, AdminWalletRefundRequestListQuery, AnnouncementListQuery, AuditLogListQuery, BackgroundTaskListQuery, BackgroundTaskSummary, BillingModelContextCacheKey, BillingModelContextCacheState, BillingModelContextInflightState, BillingPlanRecord, BillingPlanWriteInput, CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput, CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord, CreateManualWalletRechargeInput, CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome, CreateWalletRechargeOrderInput, CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput, CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, DataLayerError, DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput, DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput, GatewayDataState, GatewayProviderTransportSnapshot, LocalVideoTaskReadResponse, PaymentGatewayConfigRecord, PaymentGatewayConfigWriteInput, ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome, RedeemWalletCodeInput, RedeemWalletCodeOutcome, RequestAuditBundle, RequestCandidateTrace, StoredAdminAuditLogPage, StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder, StoredAdminPaymentOrderPage, StoredAdminRedeemCodeBatch, StoredAdminRedeemCodeBatchPage, StoredAdminRedeemCodePage, StoredAdminWalletLedgerPage, StoredAdminWalletListPage, StoredAdminWalletRefund, StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction, StoredAdminWalletTransactionPage, StoredAnnouncement, StoredAnnouncementPage, StoredBackgroundTaskEvent, StoredBackgroundTaskRun, StoredBackgroundTaskRunPage, StoredBillingModelContext, StoredProviderQuotaSnapshot, StoredProviderUsageSummary, StoredRequestUsageAudit, StoredSuspiciousActivity, StoredUsageSettlement, StoredUserAuditLogPage, StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, StoredVideoTask, StoredWalletDailyUsageLedger, StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAnnouncementRecord, UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun, UpsertUsageRecord, UpsertVideoTask, UsageSettlementInput, UserDailyQuotaAvailabilityRecord, UserPlanEntitlementRecord, VideoTaskLookupKey, VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskStatusCount, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult, WalletLookupKey, WalletMutationOutcome, }; use aether_data_contracts::repository::usage::{ PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest, StoredProviderApiKeyWindowUsageSummary, StoredUsageDailySummary, UsageAuditListQuery, UsageCleanupExecutionMode, UsageCleanupSummary, UsageCleanupTargets, UsageCleanupWindow, UsageCounterFlushSummary, UsageCounterHealthSnapshot, UsageCounterPendingHealthSnapshot, UsageDailyHeatmapQuery, }; use aether_runtime_state::RuntimeQueueStore; use aether_video_tasks_core::read_data_backed_video_task_response; use std::time::{Duration, Instant}; use tokio::time::timeout; fn normalize_billing_context_cache_part(value: &str) -> String { value.trim().to_string() } fn normalize_optional_billing_context_cache_part(value: Option<&str>) -> Option { value .map(str::trim) .filter(|value| !value.is_empty()) .map(ToOwned::to_owned) } enum BillingModelContextInflightRegistration<'a> { Leader(BillingModelContextInflightGuard<'a>), Follower(std::sync::Arc), Saturated, } struct BillingModelContextInflightGuard<'a> { state: &'a GatewayDataState, key: Option, inflight_state: std::sync::Arc, admission: Option, } impl<'a> BillingModelContextInflightGuard<'a> { fn new( state: &'a GatewayDataState, key: BillingModelContextCacheKey, inflight_state: std::sync::Arc, admission: tokio::sync::OwnedSemaphorePermit, ) -> Self { Self { state, key: Some(key), inflight_state, admission: Some(admission), } } fn epoch(&self) -> u64 { self.inflight_state.epoch } fn finish(&mut self, error: Option) { let removed = self.key.take().and_then(|key| { self.state.finish_billing_model_context_inflight( &key, &self.inflight_state, self.admission.take(), ) }); self.admission.take(); if let Some(removed) = removed { removed.complete(error.map_or(Ok(()), Err)); } } } impl Drop for BillingModelContextInflightGuard<'_> { fn drop(&mut self) { self.finish(None); } } impl BillingModelContextInflightState { fn complete(&self, result: Result<(), DataLayerError>) { if self.completion.set(result).is_ok() { self.notify.notify_waiters(); } } async fn wait(&self) -> Result<(), DataLayerError> { loop { if let Some(result) = self.completion.get() { return result.clone(); } let mut notified = Box::pin(self.notify.notified()); notified.as_mut().enable(); if let Some(result) = self.completion.get() { return result.clone(); } notified.await; } } } impl Default for BillingModelContextCacheState { fn default() -> Self { Self { entries: aether_cache::ExpiringMap::default(), inflight: std::sync::Mutex::new(std::collections::HashMap::new()), epoch: std::sync::atomic::AtomicU64::new(0), mutation: std::sync::Mutex::new(()), admission: std::sync::Arc::new(tokio::sync::Semaphore::new( GatewayDataState::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT, )), } } } impl GatewayDataState { const MAINTENANCE_POOL_IDLE_RESERVE_ENV: &'static str = "AETHER_GATEWAY_MAINTENANCE_POOL_IDLE_RESERVE"; const MAINTENANCE_POOL_PRESSURE_MAX_DEFER: Duration = Duration::from_secs(30); const BILLING_MODEL_CONTEXT_CACHE_TTL: Duration = Duration::from_secs(30); const BILLING_MODEL_CONTEXT_CACHE_MAX_ENTRIES: usize = 4096; const BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT: usize = 4096; #[cfg(not(test))] const BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_secs(10); #[cfg(test)] const BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_millis(100); pub(crate) async fn run_database_maintenance( &self, table_names: &[&str], ) -> Result { match &self.backends { Some(backends) => backends.run_database_maintenance(table_names).await, None => Ok(DatabaseMaintenanceSummary::default()), } } pub(crate) async fn run_database_migrations( &self, ) -> Result { match &self.backends { Some(backends) => backends.run_database_migrations().await, None => Ok(false), } } pub(crate) async fn run_database_backfills(&self) -> Result { match &self.backends { Some(backends) => backends.run_database_backfills().await, None => Ok(false), } } pub(crate) async fn pending_database_migrations( &self, ) -> Result< Option>, sqlx::migrate::MigrateError, > { match &self.backends { Some(backends) => backends.pending_database_migrations().await, None => Ok(None), } } pub(crate) async fn prepare_database_for_startup( &self, ) -> Result< Option>, sqlx::migrate::MigrateError, > { match &self.backends { Some(backends) => backends.prepare_database_for_startup().await, None => Ok(None), } } pub(crate) async fn warm_database_pool(&self) -> Result<(), DataLayerError> { match &self.backends { Some(backends) => backends.warm_database_pool().await, None => Ok(()), } } pub(crate) async fn pending_database_backfills( &self, ) -> Result< Option>, sqlx::migrate::MigrateError, > { match &self.backends { Some(backends) => backends.pending_database_backfills().await, None => Ok(None), } } pub(crate) fn database_pool_summary(&self) -> Option { self.backends .as_ref() .and_then(|backends| backends.database_pool_summary()) } pub(crate) async fn postgres_observability_snapshot( &self, ) -> Result, DataLayerError> { match &self.backends { Some(backends) => backends.postgres_observability_snapshot().await, None => Ok(None), } } pub(crate) async fn postgres_activity_groups( &self, limit: i64, ) -> Result, DataLayerError> { match &self.backends { Some(backends) => backends.postgres_activity_groups(limit).await, None => Ok(Vec::new()), } } pub(crate) fn database_pool_under_maintenance_pressure(&self) -> bool { self.database_pool_summary() .as_ref() .is_some_and(Self::database_pool_summary_under_maintenance_pressure) } pub(crate) fn database_pool_summary_under_maintenance_pressure( summary: &aether_data::DatabasePoolSummary, ) -> bool { summary.checked_out > 0 && Self::database_pool_available_capacity(summary) <= Self::maintenance_pool_idle_reserve(summary) } pub(crate) fn database_pool_summary_under_usage_worker_pressure( summary: &aether_data::DatabasePoolSummary, ) -> bool { summary.checked_out > 0 && Self::database_pool_available_capacity(summary) <= Self::usage_worker_pool_idle_reserve(summary) } fn database_pool_available_capacity(summary: &aether_data::DatabasePoolSummary) -> usize { let unopened = (summary.max_connections as usize).saturating_sub(summary.pool_size); summary.idle.saturating_add(unopened) } pub(crate) fn maintenance_pool_idle_reserve( summary: &aether_data::DatabasePoolSummary, ) -> usize { if let Some(override_value) = std::env::var(Self::MAINTENANCE_POOL_IDLE_RESERVE_ENV) .ok() .and_then(|value| value.trim().parse::().ok()) { return override_value; } let max_connections = summary.max_connections as usize; if max_connections == 0 { return 0; } let ten_percent_ceil = (max_connections + 9) / 10; ten_percent_ceil.clamp(2, 10).min(max_connections) } fn usage_worker_pool_idle_reserve(summary: &aether_data::DatabasePoolSummary) -> usize { if summary.max_connections <= 1 { return 0; } 1 } pub(crate) fn should_defer_maintenance_for_database_pool_pressure( &self, deferred_since: &mut Option, ) -> bool { Self::should_defer_maintenance_for_pool_pressure_state( self.database_pool_under_maintenance_pressure(), deferred_since, ) } pub(crate) fn should_defer_maintenance_for_pool_pressure_state( pool_under_pressure: bool, deferred_since: &mut Option, ) -> bool { if !pool_under_pressure { *deferred_since = None; return false; } let now = Instant::now(); let since = deferred_since.get_or_insert(now); if now.duration_since(*since) >= Self::MAINTENANCE_POOL_PRESSURE_MAX_DEFER { *deferred_since = None; return false; } true } pub(crate) async fn aggregate_wallet_daily_usage( &self, input: &WalletDailyUsageAggregationInput, ) -> Result { match &self.backends { Some(backends) => backends.aggregate_wallet_daily_usage(input).await, None => Ok(WalletDailyUsageAggregationResult::default()), } } pub(crate) async fn aggregate_stats_hourly( &self, input: &aether_data::StatsHourlyAggregationInput, ) -> Result, DataLayerError> { match &self.backends { Some(backends) => backends.aggregate_stats_hourly(input).await, None => Ok(None), } } pub(crate) async fn aggregate_stats_daily( &self, input: &aether_data::StatsDailyAggregationInput, ) -> Result, DataLayerError> { match &self.backends { Some(backends) => backends.aggregate_stats_daily(input).await, None => Ok(None), } } pub(crate) async fn list_announcements( &self, query: &AnnouncementListQuery, ) -> Result { match &self.announcement_reader { Some(repository) => repository.list_announcements(query).await, None => Ok(StoredAnnouncementPage::default()), } } pub(crate) async fn find_announcement_by_id( &self, announcement_id: &str, ) -> Result, DataLayerError> { match &self.announcement_reader { Some(repository) => repository.find_by_id(announcement_id).await, None => Ok(None), } } pub(crate) async fn list_admin_audit_logs( &self, query: &AuditLogListQuery, ) -> Result { let Some(repository) = self .backends .as_ref() .and_then(|backends| backends.read().audit_logs()) else { return Ok(StoredAdminAuditLogPage { items: Vec::new(), total: 0, }); }; repository.list_admin_audit_logs(query).await } pub(crate) async fn list_admin_suspicious_activities( &self, cutoff_unix_secs: u64, ) -> Result, DataLayerError> { let Some(repository) = self .backends .as_ref() .and_then(|backends| backends.read().audit_logs()) else { return Ok(Vec::new()); }; repository .list_admin_suspicious_activities(cutoff_unix_secs) .await } pub(crate) async fn read_admin_user_behavior_event_counts( &self, user_id: &str, cutoff_unix_secs: u64, ) -> Result, DataLayerError> { let Some(repository) = self .backends .as_ref() .and_then(|backends| backends.read().audit_logs()) else { return Ok(std::collections::BTreeMap::new()); }; repository .read_admin_user_behavior_event_counts(user_id, cutoff_unix_secs) .await } pub(crate) async fn list_user_audit_logs( &self, user_id: &str, query: &AuditLogListQuery, ) -> Result { let Some(repository) = self .backends .as_ref() .and_then(|backends| backends.read().audit_logs()) else { return Ok(StoredUserAuditLogPage { items: Vec::new(), total: 0, }); }; repository.list_user_audit_logs(user_id, query).await } pub(crate) async fn delete_audit_logs_before( &self, cutoff_unix_secs: u64, limit: usize, ) -> Result { let Some(repository) = self .backends .as_ref() .and_then(|backends| backends.read().audit_logs()) else { return Ok(0); }; repository .delete_audit_logs_before(cutoff_unix_secs, limit) .await } pub(crate) async fn count_unread_active_announcements( &self, user_id: &str, now_unix_secs: u64, ) -> Result { match &self.announcement_reader { Some(repository) => { repository .count_unread_active_announcements(user_id, now_unix_secs) .await } None => Ok(0), } } pub(crate) async fn list_required_unread_active_announcements( &self, user_id: &str, now_unix_secs: u64, limit: usize, ) -> Result, DataLayerError> { match &self.announcement_reader { Some(repository) => { repository .list_required_unread_active_announcements(user_id, now_unix_secs, limit) .await } None => Ok(Vec::new()), } } pub(crate) async fn create_announcement( &self, record: CreateAnnouncementRecord, ) -> Result, DataLayerError> { match &self.announcement_writer { Some(repository) => repository.create_announcement(record).await.map(Some), None => Ok(None), } } pub(crate) async fn update_announcement( &self, record: UpdateAnnouncementRecord, ) -> Result, DataLayerError> { match &self.announcement_writer { Some(repository) => repository.update_announcement(record).await, None => Ok(None), } } pub(crate) async fn delete_announcement( &self, announcement_id: &str, ) -> Result { match &self.announcement_writer { Some(repository) => repository.delete_announcement(announcement_id).await, None => Ok(false), } } pub(crate) async fn mark_announcement_as_read( &self, user_id: &str, announcement_id: &str, read_at_unix_secs: u64, ) -> Result { match &self.announcement_writer { Some(repository) => { repository .mark_announcement_as_read(user_id, announcement_id, read_at_unix_secs) .await } None => Ok(false), } } pub(crate) async fn find_video_task( &self, key: VideoTaskLookupKey<'_>, ) -> Result, DataLayerError> { match &self.video_task_reader { Some(repository) => repository.find(key).await, None => Ok(None), } } pub(crate) async fn list_video_task_page( &self, filter: &VideoTaskQueryFilter, offset: usize, limit: usize, ) -> Result, DataLayerError> { match &self.video_task_reader { Some(repository) => repository.list_page(filter, offset, limit).await, None => Ok(Vec::new()), } } pub(crate) async fn list_video_task_page_summary( &self, filter: &VideoTaskQueryFilter, offset: usize, limit: usize, ) -> Result, DataLayerError> { match &self.video_task_reader { Some(repository) => repository.list_page_summary(filter, offset, limit).await, None => Ok(Vec::new()), } } pub(crate) async fn count_video_tasks( &self, filter: &VideoTaskQueryFilter, ) -> Result { match &self.video_task_reader { Some(repository) => repository.count(filter).await, None => Ok(0), } } pub(crate) async fn count_video_tasks_by_status( &self, filter: &VideoTaskQueryFilter, ) -> Result, DataLayerError> { match &self.video_task_reader { Some(repository) => repository.count_by_status(filter).await, None => Ok(Vec::new()), } } pub(crate) async fn count_distinct_video_task_users( &self, filter: &VideoTaskQueryFilter, ) -> Result { match &self.video_task_reader { Some(repository) => repository.count_distinct_users(filter).await, None => Ok(0), } } pub(crate) async fn top_video_task_models( &self, filter: &VideoTaskQueryFilter, limit: usize, ) -> Result, DataLayerError> { match &self.video_task_reader { Some(repository) => repository.top_models(filter, limit).await, None => Ok(Vec::new()), } } pub(crate) async fn count_video_tasks_created_since( &self, filter: &VideoTaskQueryFilter, created_since_unix_secs: u64, ) -> Result { match &self.video_task_reader { Some(repository) => { repository .count_created_since(filter, created_since_unix_secs) .await } None => Ok(0), } } pub(crate) async fn upsert_video_task( &self, task: UpsertVideoTask, ) -> Result, DataLayerError> { match &self.video_task_writer { Some(repository) => repository.upsert(task).await.map(Some), None => Ok(None), } } pub(crate) async fn update_active_video_task( &self, task: UpsertVideoTask, ) -> Result, DataLayerError> { match &self.video_task_writer { Some(repository) => repository.update_if_active(task).await, None => Ok(None), } } pub(crate) async fn claim_due_video_tasks( &self, now_unix_secs: u64, claim_until_unix_secs: u64, limit: usize, ) -> Result, DataLayerError> { match &self.video_task_writer { Some(repository) => { repository .claim_due(now_unix_secs, claim_until_unix_secs, limit) .await } None => Ok(Vec::new()), } } pub(crate) async fn find_wallet( &self, key: WalletLookupKey<'_>, ) -> Result, DataLayerError> { match &self.wallet_reader { Some(repository) => repository.find(key).await, None => Ok(None), } } pub(crate) async fn list_wallets_by_api_key_ids( &self, api_key_ids: &[String], ) -> Result, DataLayerError> { match &self.wallet_reader { Some(repository) => repository.list_wallets_by_api_key_ids(api_key_ids).await, None => Ok(Vec::new()), } } pub(crate) async fn list_wallets_by_user_ids( &self, user_ids: &[String], ) -> Result, DataLayerError> { match &self.wallet_reader { Some(repository) => repository.list_wallets_by_user_ids(user_ids).await, None => Ok(Vec::new()), } } pub(crate) async fn list_admin_wallets( &self, query: &AdminWalletListQuery, ) -> Result { match &self.wallet_reader { Some(repository) => repository.list_admin_wallets(query).await, None => Ok(StoredAdminWalletListPage::default()), } } pub(crate) async fn list_admin_wallet_ledger( &self, query: &AdminWalletLedgerQuery, ) -> Result { match &self.wallet_reader { Some(repository) => repository.list_admin_wallet_ledger(query).await, None => Ok(StoredAdminWalletLedgerPage::default()), } } pub(crate) async fn list_admin_wallet_refund_requests( &self, query: &AdminWalletRefundRequestListQuery, ) -> Result { match &self.wallet_reader { Some(repository) => repository.list_admin_wallet_refund_requests(query).await, None => Ok(StoredAdminWalletRefundRequestPage::default()), } } pub(crate) async fn list_admin_wallet_transactions( &self, wallet_id: &str, limit: usize, offset: usize, ) -> Result { match &self.wallet_reader { Some(repository) => { repository .list_admin_wallet_transactions(wallet_id, limit, offset) .await } None => Ok(StoredAdminWalletTransactionPage::default()), } } pub(crate) async fn find_wallet_today_usage( &self, wallet_id: &str, billing_timezone: &str, ) -> Result, DataLayerError> { match &self.wallet_reader { Some(repository) => { repository .find_wallet_today_usage(wallet_id, billing_timezone) .await } None => Ok(None), } } pub(crate) async fn list_wallet_daily_usage_history( &self, wallet_id: &str, billing_timezone: &str, limit: usize, ) -> Result { match &self.wallet_reader { Some(repository) => { repository .list_wallet_daily_usage_history(wallet_id, billing_timezone, limit) .await } None => Ok(StoredWalletDailyUsageLedgerPage::default()), } } pub(crate) async fn list_admin_wallet_refunds( &self, wallet_id: &str, limit: usize, offset: usize, ) -> Result { match &self.wallet_reader { Some(repository) => { repository .list_admin_wallet_refunds(wallet_id, limit, offset) .await } None => Ok(StoredAdminWalletRefundPage::default()), } } pub(crate) async fn list_admin_payment_orders( &self, query: &AdminPaymentOrderListQuery, ) -> Result { match &self.wallet_reader { Some(repository) => repository.list_admin_payment_orders(query).await, None => Ok(StoredAdminPaymentOrderPage::default()), } } pub(crate) async fn list_admin_payment_callbacks( &self, payment_method: Option<&str>, limit: usize, offset: usize, ) -> Result { match &self.wallet_reader { Some(repository) => { repository .list_admin_payment_callbacks(payment_method, limit, offset) .await } None => Ok(StoredAdminPaymentCallbackPage::default()), } } pub(crate) async fn list_admin_redeem_code_batches( &self, query: &AdminRedeemCodeBatchListQuery, ) -> Result { match &self.wallet_reader { Some(repository) => repository.list_admin_redeem_code_batches(query).await, None => Ok(StoredAdminRedeemCodeBatchPage::default()), } } pub(crate) async fn find_admin_redeem_code_batch( &self, batch_id: &str, ) -> Result, DataLayerError> { match &self.wallet_reader { Some(repository) => repository.find_admin_redeem_code_batch(batch_id).await, None => Ok(None), } } pub(crate) async fn list_admin_redeem_codes( &self, query: &AdminRedeemCodeListQuery, ) -> Result { match &self.wallet_reader { Some(repository) => repository.list_admin_redeem_codes(query).await, None => Ok(StoredAdminRedeemCodePage::default()), } } pub(crate) async fn find_admin_payment_order( &self, order_id: &str, ) -> Result, DataLayerError> { match &self.wallet_reader { Some(repository) => repository.find_admin_payment_order(order_id).await, None => Ok(None), } } pub(crate) async fn list_wallet_payment_orders_by_user_id( &self, user_id: &str, limit: usize, offset: usize, ) -> Result { match &self.wallet_reader { Some(repository) => { repository .list_wallet_payment_orders_by_user_id(user_id, limit, offset) .await } None => Ok(StoredAdminPaymentOrderPage::default()), } } pub(crate) async fn find_wallet_payment_order_by_user_id( &self, user_id: &str, order_id: &str, ) -> Result, DataLayerError> { match &self.wallet_reader { Some(repository) => { repository .find_wallet_payment_order_by_user_id(user_id, order_id) .await } None => Ok(None), } } pub(crate) async fn find_pending_plan_purchase_order_by_user_id( &self, user_id: &str, product_id: &str, ) -> Result, DataLayerError> { match &self.wallet_reader { Some(repository) => { repository .find_pending_plan_purchase_order_by_user_id(user_id, product_id) .await } None => Ok(None), } } pub(crate) async fn find_wallet_refund( &self, wallet_id: &str, refund_id: &str, ) -> Result, DataLayerError> { match &self.wallet_reader { Some(repository) => repository.find_wallet_refund(wallet_id, refund_id).await, None => Ok(None), } } pub(crate) async fn create_wallet_recharge_order( &self, input: CreateWalletRechargeOrderInput, ) -> Result, DataLayerError> { match &self.wallet_writer { Some(repository) => repository .create_wallet_recharge_order(input) .await .map(Some), None => Ok(None), } } pub(crate) async fn create_plan_purchase_order( &self, input: CreatePlanPurchaseOrderInput, ) -> Result, DataLayerError> { match &self.wallet_writer { Some(repository) => repository.create_plan_purchase_order(input).await.map(Some), None => Ok(None), } } pub(crate) async fn create_wallet_refund_request( &self, input: CreateWalletRefundRequestInput, ) -> Result, DataLayerError> { match &self.wallet_writer { Some(repository) => repository .create_wallet_refund_request(input) .await .map(Some), None => Ok(None), } } pub(crate) async fn process_payment_callback( &self, input: ProcessPaymentCallbackInput, ) -> Result, DataLayerError> { match &self.wallet_writer { Some(repository) => repository.process_payment_callback(input).await.map(Some), None => Ok(None), } } pub(crate) async fn adjust_wallet_balance( &self, input: AdjustWalletBalanceInput, ) -> Result, DataLayerError> { match &self.wallet_writer { Some(repository) => repository.adjust_wallet_balance(input).await, None => Ok(None), } } pub(crate) async fn create_manual_wallet_recharge( &self, input: CreateManualWalletRechargeInput, ) -> Result, DataLayerError> { match &self.wallet_writer { Some(repository) => repository.create_manual_wallet_recharge(input).await, None => Ok(None), } } pub(crate) async fn process_admin_wallet_refund( &self, input: ProcessAdminWalletRefundInput, ) -> Result< Option< WalletMutationOutcome<( StoredWalletSnapshot, StoredAdminWalletRefund, StoredAdminWalletTransaction, )>, >, DataLayerError, > { match &self.wallet_writer { Some(repository) => repository .process_admin_wallet_refund(input) .await .map(Some), None => Ok(None), } } pub(crate) async fn complete_admin_wallet_refund( &self, input: CompleteAdminWalletRefundInput, ) -> Result>, DataLayerError> { match &self.wallet_writer { Some(repository) => repository .complete_admin_wallet_refund(input) .await .map(Some), None => Ok(None), } } pub(crate) async fn fail_admin_wallet_refund( &self, input: FailAdminWalletRefundInput, ) -> Result< Option< WalletMutationOutcome<( StoredWalletSnapshot, StoredAdminWalletRefund, Option, )>, >, DataLayerError, > { match &self.wallet_writer { Some(repository) => repository.fail_admin_wallet_refund(input).await.map(Some), None => Ok(None), } } pub(crate) async fn expire_admin_payment_order( &self, order_id: &str, ) -> Result>, DataLayerError> { match &self.wallet_writer { Some(repository) => repository .expire_admin_payment_order(order_id) .await .map(Some), None => Ok(None), } } pub(crate) async fn fail_admin_payment_order( &self, order_id: &str, ) -> Result>, DataLayerError> { match &self.wallet_writer { Some(repository) => repository .fail_admin_payment_order(order_id) .await .map(Some), None => Ok(None), } } pub(crate) async fn credit_admin_payment_order( &self, input: CreditAdminPaymentOrderInput, ) -> Result>, DataLayerError> { match &self.wallet_writer { Some(repository) => repository.credit_admin_payment_order(input).await.map(Some), None => Ok(None), } } pub(crate) async fn create_admin_redeem_code_batch( &self, input: CreateAdminRedeemCodeBatchInput, ) -> Result, DataLayerError> { match &self.wallet_writer { Some(repository) => repository .create_admin_redeem_code_batch(input) .await .map(Some), None => Ok(None), } } pub(crate) async fn disable_admin_redeem_code_batch( &self, input: DisableAdminRedeemCodeBatchInput, ) -> Result>, DataLayerError> { match &self.wallet_writer { Some(repository) => repository .disable_admin_redeem_code_batch(input) .await .map(Some), None => Ok(None), } } pub(crate) async fn delete_admin_redeem_code_batch( &self, input: DeleteAdminRedeemCodeBatchInput, ) -> Result>, DataLayerError> { match &self.wallet_writer { Some(repository) => repository .delete_admin_redeem_code_batch(input) .await .map(Some), None => Ok(None), } } pub(crate) async fn disable_admin_redeem_code( &self, input: DisableAdminRedeemCodeInput, ) -> Result< Option>, DataLayerError, > { match &self.wallet_writer { Some(repository) => repository.disable_admin_redeem_code(input).await.map(Some), None => Ok(None), } } pub(crate) async fn redeem_wallet_code( &self, input: RedeemWalletCodeInput, ) -> Result, DataLayerError> { match &self.wallet_writer { Some(repository) => repository.redeem_wallet_code(input).await.map(Some), None => Ok(None), } } pub(crate) async fn settle_usage( &self, input: UsageSettlementInput, ) -> Result, DataLayerError> { match &self.settlement_writer { Some(repository) => repository.settle_usage(input).await, None => Ok(None), } } pub(crate) async fn reset_due_provider_quotas( &self, now_unix_secs: u64, ) -> Result { match &self.provider_quota_writer { Some(repository) => repository.reset_due(now_unix_secs).await, None => Ok(0), } } pub(crate) async fn find_provider_quota_by_provider_id( &self, provider_id: &str, ) -> Result, DataLayerError> { match &self.provider_quota_reader { Some(repository) => repository.find_by_provider_id(provider_id).await, None => Ok(None), } } pub(crate) async fn find_provider_quotas_by_provider_ids( &self, provider_ids: &[String], ) -> Result, DataLayerError> { match &self.provider_quota_reader { Some(repository) => repository.find_by_provider_ids(provider_ids).await, None => Ok(Vec::new()), } } #[allow(dead_code)] pub(crate) async fn upsert_usage( &self, usage: UpsertUsageRecord, ) -> Result, DataLayerError> { crate::request_diagnostics::observe_db_operation( "usage_upsert", self.database_pool_summary(), async { match &self.usage_writer { Some(repository) => repository.upsert(usage).await.map(Some), None => Ok(None), } }, ) .await } pub(crate) async fn upsert_first_byte_usage( &self, usage: UpsertUsageRecord, ) -> Result<(), DataLayerError> { crate::request_diagnostics::observe_db_operation( "usage_first_byte_upsert", self.database_pool_summary(), async { match &self.usage_writer { Some(repository) => repository.upsert_first_byte(usage).await, None => Ok(()), } }, ) .await } pub(crate) async fn upsert_first_byte_usage_many( &self, usages: Vec, ) -> Result<(), DataLayerError> { if usages.is_empty() { return Ok(()); } crate::request_diagnostics::observe_db_operation( "usage_first_byte_upsert_batch", self.database_pool_summary(), async { match &self.usage_writer { Some(repository) => repository.upsert_first_byte_many(usages).await, None => Ok(()), } }, ) .await } pub(crate) async fn upsert_pending_usage_many( &self, usages: Vec, ) -> Result<(), DataLayerError> { if usages.is_empty() { return Ok(()); } crate::request_diagnostics::observe_db_operation( "usage_pending_upsert_batch", self.database_pool_summary(), async { match &self.usage_writer { Some(repository) => repository.upsert_pending_many(usages).await, None => Ok(()), } }, ) .await } #[allow(dead_code)] pub(crate) async fn rebuild_api_key_usage_stats(&self) -> Result { match &self.usage_writer { Some(repository) => repository.rebuild_api_key_usage_stats().await, None => Ok(0), } } #[allow(dead_code)] pub(crate) async fn rebuild_provider_api_key_usage_stats(&self) -> Result { match &self.usage_writer { Some(repository) => repository.rebuild_provider_api_key_usage_stats().await, None => Ok(0), } } pub(crate) async fn flush_usage_counter_deltas( &self, batch_size: usize, ) -> Result { match &self.usage_writer { Some(repository) => repository.flush_usage_counter_deltas(batch_size).await, None => Ok(UsageCounterFlushSummary::default()), } } pub(crate) async fn cleanup_processed_usage_counter_deltas( &self, cutoff_unix_secs: u64, batch_size: usize, ) -> Result { match &self.usage_writer { Some(repository) => { repository .cleanup_processed_usage_counter_deltas(cutoff_unix_secs, batch_size) .await } None => Ok(0), } } pub(crate) async fn cleanup_stale_pending_requests( &self, cutoff_unix_secs: u64, now_unix_secs: u64, timeout_minutes: u64, batch_size: usize, ) -> Result { match &self.usage_writer { Some(repository) => { repository .cleanup_stale_pending_requests( cutoff_unix_secs, now_unix_secs, timeout_minutes, batch_size, ) .await } None => Ok(PendingUsageCleanupSummary::default()), } } pub(crate) async fn cleanup_usage( &self, window: &UsageCleanupWindow, batch_size: usize, auto_delete_expired_keys: bool, targets: UsageCleanupTargets, mode: UsageCleanupExecutionMode, ) -> Result { match &self.usage_writer { Some(repository) => { repository .cleanup_usage(window, batch_size, auto_delete_expired_keys, targets, mode) .await } None => Ok(UsageCleanupSummary::default()), } } pub(crate) async fn preview_usage_cleanup( &self, window: &UsageCleanupWindow, targets: UsageCleanupTargets, mode: UsageCleanupExecutionMode, ) -> Result { match &self.usage_writer { Some(repository) => { repository .preview_usage_cleanup(window, targets, mode) .await } None => { Ok(aether_data_contracts::repository::usage::UsageCleanupPreviewCounts::default()) } } } pub(crate) async fn find_request_usage_by_request_id( &self, request_id: &str, ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.find_by_request_id(request_id).await, None => Ok(None), } } pub(crate) async fn find_request_usage_by_request_id_shallow( &self, request_id: &str, ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.find_by_request_id_shallow(request_id).await, None => Ok(None), } } pub(crate) async fn find_request_usage_by_id( &self, usage_id: &str, ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.find_by_id(usage_id).await, None => Ok(None), } } pub(crate) async fn list_request_usage_by_ids( &self, usage_ids: &[String], ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.list_by_ids(usage_ids).await, None => Ok(Vec::new()), } } pub(crate) async fn resolve_request_usage_body_ref( &self, body_ref: &str, ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.resolve_body_ref(body_ref).await, None => Ok(None), } } pub(crate) async fn list_usage_audits( &self, query: &UsageAuditListQuery, ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.list_usage_audits(query).await, None => Ok(Vec::new()), } } pub(crate) async fn count_usage_audits( &self, query: &UsageAuditListQuery, ) -> Result { match &self.usage_reader { Some(repository) => repository.count_usage_audits(query).await, None => Ok(0), } } pub(crate) async fn list_usage_audits_by_keyword_search( &self, query: &aether_data_contracts::repository::usage::UsageAuditKeywordSearchQuery, ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.list_usage_audits_by_keyword_search(query).await, None => Ok(Vec::new()), } } pub(crate) async fn count_usage_audits_by_keyword_search( &self, query: &aether_data_contracts::repository::usage::UsageAuditKeywordSearchQuery, ) -> Result { match &self.usage_reader { Some(repository) => repository.count_usage_audits_by_keyword_search(query).await, None => Ok(0), } } pub(crate) async fn aggregate_usage_audits( &self, query: &aether_data_contracts::repository::usage::UsageAuditAggregationQuery, ) -> Result< Vec, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.aggregate_usage_audits(query).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_usage_audits( &self, query: &aether_data_contracts::repository::usage::UsageAuditSummaryQuery, ) -> Result { match &self.usage_reader { Some(repository) => repository.summarize_usage_audits(query).await, None => { Ok(aether_data_contracts::repository::usage::StoredUsageAuditSummary::default()) } } } pub(crate) async fn read_usage_counter_health( &self, ) -> Result { match &self.usage_reader { Some(repository) => repository.read_usage_counter_health().await, None => Ok(UsageCounterHealthSnapshot::default()), } } pub(crate) async fn read_usage_counter_pending_health( &self, ) -> Result { match &self.usage_reader { Some(repository) => repository.read_usage_counter_pending_health().await, None => Ok(UsageCounterPendingHealthSnapshot::default()), } } pub(crate) async fn summarize_usage_totals_by_user_ids( &self, user_ids: &[String], ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => { repository .summarize_usage_totals_by_user_ids(user_ids) .await } None => Ok(Vec::new()), } } pub(crate) async fn summarize_usage_cache_hit_summary( &self, query: &aether_data_contracts::repository::usage::UsageCacheHitSummaryQuery, ) -> Result { match &self.usage_reader { Some(repository) => repository.summarize_usage_cache_hit_summary(query).await, None => { Ok(aether_data_contracts::repository::usage::StoredUsageCacheHitSummary::default()) } } } pub(crate) async fn summarize_usage_settled_cost( &self, query: &aether_data_contracts::repository::usage::UsageSettledCostSummaryQuery, ) -> Result< aether_data_contracts::repository::usage::StoredUsageSettledCostSummary, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.summarize_usage_settled_cost(query).await, None => Ok( aether_data_contracts::repository::usage::StoredUsageSettledCostSummary::default(), ), } } pub(crate) async fn summarize_usage_cache_affinity_hit_summary( &self, query: &aether_data_contracts::repository::usage::UsageCacheAffinityHitSummaryQuery, ) -> Result< aether_data_contracts::repository::usage::StoredUsageCacheAffinityHitSummary, DataLayerError, > { match &self.usage_reader { Some(repository) => repository .summarize_usage_cache_affinity_hit_summary(query) .await, None => Ok( aether_data_contracts::repository::usage::StoredUsageCacheAffinityHitSummary::default(), ), } } pub(crate) async fn list_usage_cache_affinity_intervals( &self, query: &aether_data_contracts::repository::usage::UsageCacheAffinityIntervalQuery, ) -> Result< Vec, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.list_usage_cache_affinity_intervals(query).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_dashboard_usage( &self, query: &aether_data_contracts::repository::usage::UsageDashboardSummaryQuery, ) -> Result { match &self.usage_reader { Some(repository) => repository.summarize_dashboard_usage(query).await, None => Ok( aether_data_contracts::repository::usage::StoredUsageDashboardSummary::default(), ), } } pub(crate) async fn summarize_dashboard_stats( &self, query: &aether_data_contracts::repository::usage::UsageDashboardSummaryQuery, ) -> Result< aether_data_contracts::repository::usage::StoredUsageDashboardStatsSummary, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.summarize_dashboard_stats(query).await, None => Ok( aether_data_contracts::repository::usage::StoredUsageDashboardStatsSummary::default( ), ), } } pub(crate) async fn list_dashboard_daily_breakdown( &self, query: &aether_data_contracts::repository::usage::UsageDashboardDailyBreakdownQuery, ) -> Result< Vec, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.list_dashboard_daily_breakdown(query).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_dashboard_provider_counts( &self, query: &aether_data_contracts::repository::usage::UsageDashboardProviderCountsQuery, ) -> Result< Vec, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.summarize_dashboard_provider_counts(query).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_usage_breakdown( &self, query: &aether_data_contracts::repository::usage::UsageBreakdownSummaryQuery, ) -> Result< Vec, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.summarize_usage_breakdown(query).await, None => Ok(Vec::new()), } } pub(crate) async fn count_monitoring_usage_errors( &self, query: &aether_data_contracts::repository::usage::UsageMonitoringErrorCountQuery, ) -> Result { match &self.usage_reader { Some(repository) => repository.count_monitoring_usage_errors(query).await, None => Ok(0), } } pub(crate) async fn list_monitoring_usage_errors( &self, query: &aether_data_contracts::repository::usage::UsageMonitoringErrorListQuery, ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.list_monitoring_usage_errors(query).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_usage_error_distribution( &self, query: &aether_data_contracts::repository::usage::UsageErrorDistributionQuery, ) -> Result< Vec, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.summarize_usage_error_distribution(query).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_usage_performance_percentiles( &self, query: &aether_data_contracts::repository::usage::UsagePerformancePercentilesQuery, ) -> Result< Vec, DataLayerError, > { match &self.usage_reader { Some(repository) => { repository .summarize_usage_performance_percentiles(query) .await } None => Ok(Vec::new()), } } pub(crate) async fn summarize_usage_provider_performance( &self, query: &aether_data_contracts::repository::usage::UsageProviderPerformanceQuery, ) -> Result< aether_data_contracts::repository::usage::StoredUsageProviderPerformance, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.summarize_usage_provider_performance(query).await, None => Ok( aether_data_contracts::repository::usage::StoredUsageProviderPerformance::default(), ), } } pub(crate) async fn summarize_usage_cost_savings( &self, query: &aether_data_contracts::repository::usage::UsageCostSavingsSummaryQuery, ) -> Result< aether_data_contracts::repository::usage::StoredUsageCostSavingsSummary, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.summarize_usage_cost_savings(query).await, None => Ok( aether_data_contracts::repository::usage::StoredUsageCostSavingsSummary::default(), ), } } pub(crate) async fn summarize_usage_time_series( &self, query: &aether_data_contracts::repository::usage::UsageTimeSeriesQuery, ) -> Result< Vec, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.summarize_usage_time_series(query).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_usage_leaderboard( &self, query: &aether_data_contracts::repository::usage::UsageLeaderboardQuery, ) -> Result< Vec, DataLayerError, > { match &self.usage_reader { Some(repository) => repository.summarize_usage_leaderboard(query).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_usage_daily_heatmap( &self, query: &UsageDailyHeatmapQuery, ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.summarize_usage_daily_heatmap(query).await, None => Ok(Vec::new()), } } pub(crate) async fn list_recent_usage_audits( &self, user_id: Option<&str>, limit: usize, ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => repository.list_recent_usage_audits(user_id, limit).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_usage_total_tokens_by_api_key_ids( &self, api_key_ids: &[String], ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => { repository .summarize_total_tokens_by_api_key_ids(api_key_ids) .await } None => Ok(std::collections::BTreeMap::new()), } } pub(crate) async fn summarize_usage_by_provider_api_key_ids( &self, provider_api_key_ids: &[String], ) -> Result< std::collections::BTreeMap< String, aether_data_contracts::repository::usage::StoredProviderApiKeyUsageSummary, >, DataLayerError, > { match &self.usage_reader { Some(repository) => { repository .summarize_usage_by_provider_api_key_ids(provider_api_key_ids) .await } None => Ok(std::collections::BTreeMap::new()), } } pub(crate) async fn summarize_usage_by_provider_api_key_windows( &self, requests: &[ProviderApiKeyWindowUsageRequest], ) -> Result, DataLayerError> { match &self.usage_reader { Some(repository) => { repository .summarize_usage_by_provider_api_key_windows(requests) .await } None => Ok(Vec::new()), } } pub(crate) async fn list_users_by_ids( &self, user_ids: &[String], ) -> Result, DataLayerError> { match &self.user_reader { Some(repository) => repository.list_users_by_ids(user_ids).await, None => Ok(Vec::new()), } } pub(crate) async fn list_users_by_username_search( &self, username_search: &str, ) -> Result, DataLayerError> { match &self.user_reader { Some(repository) => { repository .list_users_by_username_search(username_search) .await } None => Ok(Vec::new()), } } pub(crate) async fn list_export_users( &self, ) -> Result, DataLayerError> { match &self.user_reader { Some(repository) => repository.list_export_users().await, None => Ok(Vec::new()), } } pub(crate) async fn list_export_users_page( &self, query: &aether_data::repository::users::UserExportListQuery, ) -> Result, DataLayerError> { match &self.user_reader { Some(repository) => repository.list_export_users_page(query).await, None => Ok(Vec::new()), } } pub(crate) async fn count_export_users( &self, query: &aether_data::repository::users::UserExportListQuery, ) -> Result { match &self.user_reader { Some(repository) => repository.count_export_users(query).await, None => Ok(0), } } pub(crate) async fn summarize_export_users( &self, ) -> Result { match &self.user_reader { Some(repository) => repository.summarize_export_users().await, None => Ok(aether_data::repository::users::UserExportSummary::default()), } } pub(crate) async fn find_export_user_by_id( &self, user_id: &str, ) -> Result, DataLayerError> { match &self.user_reader { Some(repository) => repository.find_export_user_by_id(user_id).await, None => Ok(None), } } pub(crate) async fn read_user_feature_settings( &self, user_id: &str, ) -> Result, DataLayerError> { let user_id = user_id.trim(); if user_id.is_empty() { return Ok(None); } if let Some(user) = self.find_export_user_by_id(user_id).await? { return Ok(user.feature_settings); } Ok(None) } pub(crate) async fn list_non_admin_export_users( &self, ) -> Result, DataLayerError> { match &self.user_reader { Some(repository) => repository.list_non_admin_export_users().await, None => Ok(Vec::new()), } } pub(crate) async fn list_user_auth_by_ids( &self, user_ids: &[String], ) -> Result, DataLayerError> { match &self.user_reader { Some(repository) => repository.list_user_auth_by_ids(user_ids).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_provider_usage_since( &self, provider_id: &str, since_unix_secs: u64, ) -> Result { match &self.usage_reader { Some(repository) => { repository .summarize_provider_usage_since(provider_id, since_unix_secs) .await } None => Ok(StoredProviderUsageSummary::default()), } } pub(crate) fn usage_worker_queue(&self) -> Option> { self.usage_worker_queue.clone() } pub(crate) async fn find_billing_model_context( &self, provider_id: &str, provider_api_key_id: Option<&str>, global_model_name: &str, ) -> Result, DataLayerError> { let key = BillingModelContextCacheKey::ByGlobalModelName { provider_id: normalize_billing_context_cache_part(provider_id), provider_api_key_id: normalize_optional_billing_context_cache_part(provider_api_key_id), global_model_name: normalize_billing_context_cache_part(global_model_name), }; if let Some(value) = self.cached_billing_model_context(&key) { return Ok(value); } loop { match self.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Saturated => { return Err(DataLayerError::TimedOut(format!( "billing model context cache admission saturated for {key:?}" ))); } BillingModelContextInflightRegistration::Follower(inflight_state) => { match timeout( Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT, inflight_state.wait(), ) .await { Ok(Ok(())) => {} Ok(Err(error)) => return Err(error), Err(_) => self.expire_billing_model_context_inflight(&key, &inflight_state), } if let Some(value) = self.cached_billing_model_context(&key) { return Ok(value); } continue; } BillingModelContextInflightRegistration::Leader(mut guard) => { if let Some(value) = self.cached_billing_model_context(&key) { return Ok(value); } let load_epoch = guard.epoch(); let result = match timeout( Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT, self.load_billing_model_context_by_name( key, provider_id, provider_api_key_id, global_model_name, load_epoch, &guard.inflight_state, ), ) .await { Ok(result) => result, Err(_) => Err(DataLayerError::TimedOut( "billing model context load timed out".to_string(), )), }; guard.finish(result.as_ref().err().cloned()); return result; } } } } pub(crate) async fn find_billing_model_context_by_model_id( &self, provider_id: &str, provider_api_key_id: Option<&str>, model_id: &str, ) -> Result, DataLayerError> { let key = BillingModelContextCacheKey::ByModelId { provider_id: normalize_billing_context_cache_part(provider_id), provider_api_key_id: normalize_optional_billing_context_cache_part(provider_api_key_id), model_id: normalize_billing_context_cache_part(model_id), }; if let Some(value) = self.cached_billing_model_context(&key) { return Ok(value); } loop { match self.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Saturated => { return Err(DataLayerError::TimedOut(format!( "billing model context cache admission saturated for {key:?}" ))); } BillingModelContextInflightRegistration::Follower(inflight_state) => { match timeout( Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT, inflight_state.wait(), ) .await { Ok(Ok(())) => {} Ok(Err(error)) => return Err(error), Err(_) => self.expire_billing_model_context_inflight(&key, &inflight_state), } if let Some(value) = self.cached_billing_model_context(&key) { return Ok(value); } continue; } BillingModelContextInflightRegistration::Leader(mut guard) => { if let Some(value) = self.cached_billing_model_context(&key) { return Ok(value); } let load_epoch = guard.epoch(); let result = match timeout( Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT, self.load_billing_model_context_by_model_id( key, provider_id, provider_api_key_id, model_id, load_epoch, &guard.inflight_state, ), ) .await { Ok(result) => result, Err(_) => Err(DataLayerError::TimedOut( "billing model context load timed out".to_string(), )), }; guard.finish(result.as_ref().err().cloned()); return result; } } } } async fn load_billing_model_context_by_name( &self, key: BillingModelContextCacheKey, provider_id: &str, provider_api_key_id: Option<&str>, global_model_name: &str, load_epoch: u64, load_flight: &std::sync::Arc, ) -> Result, DataLayerError> { crate::request_diagnostics::observe_db_operation( "billing_model_context", self.database_pool_summary(), async { match &self.billing_reader { Some(repository) => { let value = repository .find_model_context(provider_id, provider_api_key_id, global_model_name) .await?; self.remember_billing_model_context( key, value.clone(), load_epoch, load_flight, ); Ok(value) } None => { self.remember_billing_model_context(key, None, load_epoch, load_flight); Ok(None) } } }, ) .await } async fn load_billing_model_context_by_model_id( &self, key: BillingModelContextCacheKey, provider_id: &str, provider_api_key_id: Option<&str>, model_id: &str, load_epoch: u64, load_flight: &std::sync::Arc, ) -> Result, DataLayerError> { crate::request_diagnostics::observe_db_operation( "billing_model_context", self.database_pool_summary(), async { match &self.billing_reader { Some(repository) => { let value = repository .find_model_context_by_model_id( provider_id, provider_api_key_id, model_id, ) .await?; self.remember_billing_model_context( key, value.clone(), load_epoch, load_flight, ); Ok(value) } None => { self.remember_billing_model_context(key, None, load_epoch, load_flight); Ok(None) } } }, ) .await } fn register_billing_model_context_inflight( &self, key: &BillingModelContextCacheKey, ) -> BillingModelContextInflightRegistration<'_> { let mut inflight = self .billing_model_context_cache .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(inflight_state) = inflight.get(key) { return BillingModelContextInflightRegistration::Follower(std::sync::Arc::clone( inflight_state, )); } if inflight.len() >= Self::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT { return BillingModelContextInflightRegistration::Saturated; } let Ok(admission) = std::sync::Arc::clone(&self.billing_model_context_cache.admission).try_acquire_owned() else { return BillingModelContextInflightRegistration::Saturated; }; let inflight_state = std::sync::Arc::new(BillingModelContextInflightState { epoch: self .billing_model_context_cache .epoch .load(std::sync::atomic::Ordering::Acquire), completion: std::sync::OnceLock::new(), notify: tokio::sync::Notify::new(), }); inflight.insert(key.clone(), std::sync::Arc::clone(&inflight_state)); BillingModelContextInflightRegistration::Leader(BillingModelContextInflightGuard::new( self, key.clone(), inflight_state, admission, )) } fn finish_billing_model_context_inflight( &self, key: &BillingModelContextCacheKey, inflight_state: &std::sync::Arc, admission: Option, ) -> Option> { let mut inflight = self .billing_model_context_cache .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); drop(admission); if inflight .get(key) .is_some_and(|current| std::sync::Arc::ptr_eq(current, inflight_state)) { inflight.remove(key) } else { None } } fn expire_billing_model_context_inflight( &self, key: &BillingModelContextCacheKey, inflight_state: &std::sync::Arc, ) { let _mutation = self .billing_model_context_cache .mutation .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); let removed = { let mut inflight = self .billing_model_context_cache .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if inflight .get(key) .is_some_and(|current| std::sync::Arc::ptr_eq(current, inflight_state)) { inflight.remove(key) } else { None } }; drop(_mutation); if let Some(removed) = removed { tracing::warn!( event_name = "billing_model_context_cache_inflight_expired", log_type = "ops", cache_key = ?key, wait_timeout_ms = Self::BILLING_MODEL_CONTEXT_CACHE_INFLIGHT_WAIT_TIMEOUT.as_millis() as u64, "gateway billing model context cache expired stale inflight load" ); removed.complete(Ok(())); } } fn cached_billing_model_context( &self, key: &BillingModelContextCacheKey, ) -> Option> { self.billing_model_context_cache .entries .get_fresh(key, Self::BILLING_MODEL_CONTEXT_CACHE_TTL) } fn remember_billing_model_context( &self, key: BillingModelContextCacheKey, value: Option, load_epoch: u64, load_flight: &std::sync::Arc, ) { let _mutation = self .billing_model_context_cache .mutation .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if load_epoch != self .billing_model_context_cache .epoch .load(std::sync::atomic::Ordering::Acquire) { return; } let inflight = self .billing_model_context_cache .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if !inflight .get(&key) .is_some_and(|current| std::sync::Arc::ptr_eq(current, load_flight)) { return; } self.billing_model_context_cache.entries.insert( key, value, Self::BILLING_MODEL_CONTEXT_CACHE_TTL, Self::BILLING_MODEL_CONTEXT_CACHE_MAX_ENTRIES, ); } pub(super) fn clear_billing_model_context_cache(&self) { let _mutation = self .billing_model_context_cache .mutation .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); self.billing_model_context_cache .epoch .fetch_add(1, std::sync::atomic::Ordering::AcqRel); self.billing_model_context_cache.entries.clear(); let inflight_states = self .billing_model_context_cache .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .drain() .map(|(_, state)| state) .collect::>(); drop(_mutation); if !inflight_states.is_empty() { tracing::warn!( event_name = "billing_model_context_cache_inflight_cleared", log_type = "ops", "gateway billing model context cache cleared in-flight loads" ); for inflight_state in inflight_states { inflight_state.complete(Ok(())); } } } pub(crate) async fn admin_billing_enabled_default_value_exists( &self, api_format: &str, task_type: &str, dimension_name: &str, existing_id: Option<&str>, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => { repository .admin_billing_enabled_default_value_exists( api_format, task_type, dimension_name, existing_id, ) .await } None => Ok(None), } } pub(crate) async fn create_admin_billing_rule( &self, input: &AdminBillingRuleWriteInput, ) -> Result, DataLayerError> { let result = match &self.billing_reader { Some(repository) => repository.create_admin_billing_rule(input).await, None => Ok(AdminBillingMutationOutcome::Unavailable), }; if result.is_ok() { self.clear_billing_model_context_cache(); } result } pub(crate) async fn list_admin_billing_rules( &self, task_type: Option<&str>, is_enabled: Option, page: u32, page_size: u32, ) -> Result, u64)>, DataLayerError> { match &self.billing_reader { Some(repository) => { repository .list_admin_billing_rules(task_type, is_enabled, page, page_size) .await } None => Ok(None), } } pub(crate) async fn find_admin_billing_rule( &self, rule_id: &str, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.find_admin_billing_rule(rule_id).await, None => Ok(None), } } pub(crate) async fn update_admin_billing_rule( &self, rule_id: &str, input: &AdminBillingRuleWriteInput, ) -> Result, DataLayerError> { let result = match &self.billing_reader { Some(repository) => repository.update_admin_billing_rule(rule_id, input).await, None => Ok(AdminBillingMutationOutcome::Unavailable), }; if result.is_ok() { self.clear_billing_model_context_cache(); } result } pub(crate) async fn create_admin_billing_collector( &self, input: &AdminBillingCollectorWriteInput, ) -> Result, DataLayerError> { let result = match &self.billing_reader { Some(repository) => repository.create_admin_billing_collector(input).await, None => Ok(AdminBillingMutationOutcome::Unavailable), }; if result.is_ok() { self.clear_billing_model_context_cache(); } result } pub(crate) async fn list_admin_billing_collectors( &self, api_format: Option<&str>, task_type: Option<&str>, dimension_name: Option<&str>, is_enabled: Option, page: u32, page_size: u32, ) -> Result, u64)>, DataLayerError> { match &self.billing_reader { Some(repository) => { repository .list_admin_billing_collectors( api_format, task_type, dimension_name, is_enabled, page, page_size, ) .await } None => Ok(None), } } pub(crate) async fn find_admin_billing_collector( &self, collector_id: &str, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.find_admin_billing_collector(collector_id).await, None => Ok(None), } } pub(crate) async fn update_admin_billing_collector( &self, collector_id: &str, input: &AdminBillingCollectorWriteInput, ) -> Result, DataLayerError> { let result = match &self.billing_reader { Some(repository) => { repository .update_admin_billing_collector(collector_id, input) .await } None => Ok(AdminBillingMutationOutcome::Unavailable), }; if result.is_ok() { self.clear_billing_model_context_cache(); } result } pub(crate) async fn apply_admin_billing_preset( &self, preset: &str, mode: &str, collectors: &[AdminBillingCollectorWriteInput], ) -> Result, DataLayerError> { let result = match &self.billing_reader { Some(repository) => { repository .apply_admin_billing_preset(preset, mode, collectors) .await } None => Ok(AdminBillingMutationOutcome::Unavailable), }; if result.is_ok() { self.clear_billing_model_context_cache(); } result } pub(crate) async fn find_payment_gateway_config( &self, provider: &str, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.find_payment_gateway_config(provider).await, None => Ok(None), } } pub(crate) async fn upsert_payment_gateway_config( &self, input: &PaymentGatewayConfigWriteInput, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.upsert_payment_gateway_config(input).await, None => Ok(AdminBillingMutationOutcome::Unavailable), } } pub(crate) async fn list_billing_plans( &self, include_disabled: bool, ) -> Result>, DataLayerError> { match &self.billing_reader { Some(repository) => repository.list_billing_plans(include_disabled).await, None => Ok(None), } } pub(crate) async fn find_billing_plan( &self, plan_id: &str, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.find_billing_plan(plan_id).await, None => Ok(None), } } pub(crate) async fn create_billing_plan( &self, input: &BillingPlanWriteInput, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.create_billing_plan(input).await, None => Ok(AdminBillingMutationOutcome::Unavailable), } } pub(crate) async fn update_billing_plan( &self, plan_id: &str, input: &BillingPlanWriteInput, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.update_billing_plan(plan_id, input).await, None => Ok(AdminBillingMutationOutcome::Unavailable), } } pub(crate) async fn set_billing_plan_enabled( &self, plan_id: &str, enabled: bool, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.set_billing_plan_enabled(plan_id, enabled).await, None => Ok(AdminBillingMutationOutcome::Unavailable), } } pub(crate) async fn delete_billing_plan( &self, plan_id: &str, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.delete_billing_plan(plan_id).await, None => Ok(AdminBillingMutationOutcome::Unavailable), } } pub(crate) async fn list_user_plan_entitlements( &self, user_id: &str, ) -> Result>, DataLayerError> { match &self.billing_reader { Some(repository) => repository.list_user_plan_entitlements(user_id).await, None => Ok(None), } } pub(crate) async fn revoke_user_plan_entitlement( &self, user_id: &str, entitlement_id: &str, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => { repository .revoke_user_plan_entitlement(user_id, entitlement_id) .await } None => Ok(AdminBillingMutationOutcome::Unavailable), } } pub(crate) async fn find_user_daily_quota_availability( &self, user_id: &str, ) -> Result, DataLayerError> { match &self.billing_reader { Some(repository) => repository.find_user_daily_quota_availability(user_id).await, None => Ok(None), } } pub(crate) async fn read_request_candidate_trace( &self, request_id: &str, attempted_only: bool, ) -> Result, DataLayerError> { read_request_candidate_trace(self, request_id, attempted_only).await } pub(crate) async fn read_decision_trace( &self, request_id: &str, attempted_only: bool, ) -> Result, DataLayerError> { read_decision_trace(self, request_id, attempted_only).await } pub(crate) async fn read_request_usage_audit( &self, request_id: &str, ) -> Result, DataLayerError> { self.find_request_usage_by_request_id(request_id).await } pub(crate) async fn read_request_usage_audit_shallow( &self, request_id: &str, ) -> Result, DataLayerError> { self.find_request_usage_by_request_id_shallow(request_id) .await } pub(crate) async fn read_request_audit_bundle( &self, request_id: &str, attempted_only: bool, now_unix_secs: u64, ) -> Result, DataLayerError> { aether_data::repository::audit::read_request_audit_bundle( self, request_id, attempted_only, now_unix_secs, ) .await } #[allow(dead_code)] pub(crate) async fn read_provider_transport_snapshot( &self, provider_id: &str, endpoint_id: &str, key_id: &str, ) -> Result, DataLayerError> { read_provider_transport_snapshot(self, provider_id, endpoint_id, key_id).await } pub(crate) async fn read_video_task_response( &self, route_family: Option<&str>, request_path: &str, ) -> Result, DataLayerError> { read_data_backed_video_task_response(self, route_family, request_path).await } pub(crate) async fn find_background_task_run( &self, run_id: &str, ) -> Result, DataLayerError> { match &self.background_task_reader { Some(repository) => repository.find_run(run_id).await, None => Ok(None), } } pub(crate) async fn list_background_task_runs( &self, query: &BackgroundTaskListQuery, ) -> Result { match &self.background_task_reader { Some(repository) => repository.list_runs(query).await, None => Ok(StoredBackgroundTaskRunPage::default()), } } pub(crate) async fn list_background_task_events( &self, run_id: &str, offset: usize, limit: usize, ) -> Result, DataLayerError> { match &self.background_task_reader { Some(repository) => repository.list_events(run_id, offset, limit).await, None => Ok(Vec::new()), } } pub(crate) async fn summarize_background_task_runs( &self, ) -> Result { match &self.background_task_reader { Some(repository) => repository.summarize_runs().await, None => Ok(BackgroundTaskSummary::default()), } } pub(crate) async fn upsert_background_task_run( &self, run: UpsertBackgroundTaskRun, ) -> Result, DataLayerError> { match &self.background_task_writer { Some(repository) => repository.upsert_run(run).await.map(Some), None => Ok(None), } } pub(crate) async fn request_cancel_background_task_run( &self, run_id: &str, updated_at_unix_secs: u64, ) -> Result { match &self.background_task_writer { Some(repository) => { repository .request_cancel(run_id, updated_at_unix_secs) .await } None => Ok(false), } } pub(crate) async fn upsert_background_task_event( &self, event: UpsertBackgroundTaskEvent, ) -> Result, DataLayerError> { match &self.background_task_writer { Some(repository) => repository.upsert_event(event).await.map(Some), None => Ok(None), } } } #[cfg(test)] mod tests { use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::{Arc, Mutex}; use std::time::Duration; use aether_data::repository::global_models::InMemoryGlobalModelReadRepository; use aether_data::repository::users::{InMemoryUserReadRepository, StoredUserExportRow}; use aether_data_contracts::repository::billing::{ BillingReadRepository, StoredBillingModelContext, }; use aether_data_contracts::repository::global_models::{ StoredAdminGlobalModel, StoredPublicGlobalModel, UpdateAdminGlobalModelRecord, }; use aether_data_contracts::DataLayerError; use async_trait::async_trait; use serde_json::json; use tokio::sync::Barrier; use super::{ BillingModelContextCacheKey, BillingModelContextInflightRegistration, GatewayDataState, }; struct SlowBillingContextRepository { calls: AtomicUsize, context: StoredBillingModelContext, } struct BlockedBillingContextRepository { calls: AtomicUsize, context: Mutex, first_read: Barrier, release_first_read: Barrier, } struct ConcurrentBillingContextRepository { calls: AtomicUsize, context: StoredBillingModelContext, entered: Barrier, release: Barrier, } #[async_trait] impl BillingReadRepository for SlowBillingContextRepository { async fn find_model_context( &self, _provider_id: &str, _provider_api_key_id: Option<&str>, _global_model_name: &str, ) -> Result, aether_data_contracts::DataLayerError> { self.calls.fetch_add(1, Ordering::AcqRel); tokio::time::sleep(Duration::from_millis(25)).await; Ok(Some(self.context.clone())) } } #[async_trait] impl BillingReadRepository for BlockedBillingContextRepository { async fn find_model_context( &self, _provider_id: &str, _provider_api_key_id: Option<&str>, _global_model_name: &str, ) -> Result, aether_data_contracts::DataLayerError> { let call_index = self.calls.fetch_add(1, Ordering::AcqRel); let context = self .context .lock() .expect("mutable billing context lock") .clone(); if call_index == 0 { self.first_read.wait().await; self.release_first_read.wait().await; } Ok(Some(context)) } } #[async_trait] impl BillingReadRepository for ConcurrentBillingContextRepository { async fn find_model_context( &self, _provider_id: &str, _provider_api_key_id: Option<&str>, _global_model_name: &str, ) -> Result, DataLayerError> { self.calls.fetch_add(1, Ordering::AcqRel); self.entered.wait().await; self.release.wait().await; Ok(Some(self.context.clone())) } } fn billing_context() -> StoredBillingModelContext { StoredBillingModelContext::new( "provider-1".to_string(), Some("pay_as_you_go".to_string()), Some("key-1".to_string()), None, None, "global-model-1".to_string(), "gpt-5".to_string(), None, Some(0.02), Some(json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]})), Some("model-1".to_string()), Some("gpt-5-upstream".to_string()), None, None, None, ) .expect("billing context should build") } fn billing_cache_key(global_model_name: impl Into) -> BillingModelContextCacheKey { BillingModelContextCacheKey::ByGlobalModelName { provider_id: "provider-1".to_string(), provider_api_key_id: Some("key-1".to_string()), global_model_name: global_model_name.into(), } } #[tokio::test] async fn billing_model_context_cancelled_leader_cannot_lose_follower_wakeup() { let state = GatewayDataState::default(); let key = billing_cache_key("lost-wakeup"); let leader = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Leader(guard) => guard, _ => panic!("first registration should lead"), }; let follower = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Follower(inflight_state) => inflight_state, _ => panic!("second registration should follow"), }; // Complete before wait() is constructed or polled. A bare global // notify_waiters() broadcast loses this ordering. drop(leader); tokio::time::timeout(Duration::from_millis(100), follower.wait()) .await .expect("cancelled flight must release an unpolled follower") .expect("leader cancellation should allow a retry"); assert!(matches!( state.register_billing_model_context_inflight(&key), BillingModelContextInflightRegistration::Leader(_) )); } #[tokio::test] async fn billing_model_context_failed_flight_fans_out_error() { let state = GatewayDataState::default(); let key = billing_cache_key("failed-flight"); let mut leader = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Leader(guard) => guard, _ => panic!("first registration should lead"), }; let followers = (0..2) .map( |_| match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Follower(inflight_state) => { inflight_state } _ => panic!("same-key registration should follow"), }, ) .collect::>(); leader.finish(Some(DataLayerError::Sql( "forced billing context load failure".to_string(), ))); for follower in followers { let error = tokio::time::timeout(Duration::from_millis(100), follower.wait()) .await .expect("failed flight should release every follower") .expect_err("follower should receive the leader failure"); assert_eq!( error.to_string(), "sql error: forced billing context load failure" ); } assert!(state .billing_model_context_cache .inflight .lock() .unwrap() .is_empty()); } #[tokio::test] async fn billing_model_context_clear_wakes_old_follower_without_removing_replacement() { let state = GatewayDataState::default(); let key = billing_cache_key("clear-replacement"); let old_leader = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Leader(guard) => guard, _ => panic!("first registration should lead"), }; let old_follower = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Follower(inflight_state) => inflight_state, _ => panic!("second registration should follow"), }; let old_epoch = old_leader.epoch(); state.clear_billing_model_context_cache(); let replacement_leader = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Leader(guard) => guard, _ => panic!("clear should allow a replacement leader"), }; let replacement_follower = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Follower(inflight_state) => inflight_state, _ => panic!("registration behind replacement should follow"), }; assert_ne!(replacement_leader.epoch(), old_epoch); drop(old_leader); assert!(state .billing_model_context_cache .inflight .lock() .unwrap() .get(&key) .is_some_and(|current| std::sync::Arc::ptr_eq( current, &replacement_leader.inflight_state ))); tokio::time::timeout(Duration::from_millis(100), old_follower.wait()) .await .expect("clear should wake the invalidated flight") .expect("clear should allow an immediate retry"); drop(replacement_leader); tokio::time::timeout(Duration::from_millis(100), replacement_follower.wait()) .await .expect("old guard must not strand the replacement follower") .expect("replacement completion should succeed"); } #[tokio::test] async fn billing_model_context_timeout_expiration_allows_replacement() { let state = GatewayDataState::default(); let key = billing_cache_key("timeout-replacement"); let old_leader = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Leader(guard) => guard, _ => panic!("first registration should lead"), }; let old_follower = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Follower(inflight_state) => inflight_state, _ => panic!("second registration should follow"), }; state.expire_billing_model_context_inflight(&key, &old_follower); old_follower .wait() .await .expect("expired flight should permit a retry"); let replacement_leader = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Leader(guard) => guard, _ => panic!("timeout should allow a replacement leader"), }; drop(old_leader); assert!(state .billing_model_context_cache .inflight .lock() .unwrap() .get(&key) .is_some_and(|current| std::sync::Arc::ptr_eq( current, &replacement_leader.inflight_state ))); } #[test] fn billing_model_context_expired_leader_cannot_publish_over_replacement() { let state = GatewayDataState::default(); let key = billing_cache_key("timeout-publication-replacement"); let old_leader = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Leader(guard) => guard, _ => panic!("first registration should lead"), }; let old_flight = std::sync::Arc::clone(&old_leader.inflight_state); let load_epoch = old_leader.epoch(); state.expire_billing_model_context_inflight(&key, &old_flight); let replacement = match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Leader(guard) => guard, _ => panic!("expiration should allow a replacement leader"), }; assert_eq!(replacement.epoch(), load_epoch); let mut fresh = billing_context(); fresh.default_price_per_request = Some(2.0); state.remember_billing_model_context( key.clone(), Some(fresh), replacement.epoch(), &replacement.inflight_state, ); let mut stale = billing_context(); stale.default_price_per_request = Some(1.0); state.remember_billing_model_context(key.clone(), Some(stale), load_epoch, &old_flight); let cached = state .cached_billing_model_context(&key) .expect("replacement should publish") .expect("billing context should exist"); assert_eq!(cached.default_price_per_request, Some(2.0)); } #[test] fn billing_model_context_inflight_limit_rejects_only_new_keys() { let state = GatewayDataState::default(); let mut leaders = Vec::with_capacity(GatewayDataState::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT); for index in 0..GatewayDataState::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT { let key = billing_cache_key(format!("model-{index}")); match state.register_billing_model_context_inflight(&key) { BillingModelContextInflightRegistration::Leader(guard) => leaders.push(guard), _ => panic!("unique key below the hard limit should lead"), } } assert!(matches!( state.register_billing_model_context_inflight(&billing_cache_key("overflow")), BillingModelContextInflightRegistration::Saturated )); assert!(matches!( state.register_billing_model_context_inflight(&billing_cache_key("model-0")), BillingModelContextInflightRegistration::Follower(_) )); assert_eq!( state .billing_model_context_cache .inflight .lock() .unwrap() .len(), GatewayDataState::BILLING_MODEL_CONTEXT_CACHE_MAX_INFLIGHT ); drop(leaders); assert!(state .billing_model_context_cache .inflight .lock() .unwrap() .is_empty()); } #[tokio::test] async fn billing_model_context_different_keys_load_concurrently() { let repository = Arc::new(ConcurrentBillingContextRepository { calls: AtomicUsize::new(0), context: billing_context(), entered: Barrier::new(3), release: Barrier::new(3), }); let state = Arc::new(GatewayDataState::with_billing_reader_for_tests( repository.clone(), )); let task_a = { let state = Arc::clone(&state); tokio::spawn(async move { state .find_billing_model_context("provider-1", Some("key-1"), "model-a") .await }) }; let task_b = { let state = Arc::clone(&state); tokio::spawn(async move { state .find_billing_model_context("provider-1", Some("key-1"), "model-b") .await }) }; tokio::time::timeout(Duration::from_secs(1), repository.entered.wait()) .await .expect("different cache keys should enter the repository concurrently"); assert_eq!(repository.calls.load(Ordering::Acquire), 2); repository.release.wait().await; task_a .await .expect("first lookup should join") .expect("first lookup should succeed"); task_b .await .expect("second lookup should join") .expect("second lookup should succeed"); } #[tokio::test] async fn billing_model_context_cache_coalesces_concurrent_loads() { let repository = Arc::new(SlowBillingContextRepository { calls: AtomicUsize::new(0), context: billing_context(), }); let state = Arc::new(GatewayDataState::with_billing_reader_for_tests( repository.clone(), )); let mut tasks = Vec::new(); for _ in 0..16 { let state = Arc::clone(&state); tasks.push(tokio::spawn(async move { state .find_billing_model_context("provider-1", Some("key-1"), "gpt-5") .await .expect("billing context lookup should succeed") .expect("billing context should exist"); })); } for task in tasks { task.await.expect("lookup task should complete"); } assert_eq!(repository.calls.load(Ordering::Acquire), 1); } #[tokio::test] async fn global_model_price_update_invalidates_inflight_billing_context_cache() { let mut initial_context = billing_context(); initial_context.default_price_per_request = None; initial_context.default_tiered_pricing = None; let repository = Arc::new(BlockedBillingContextRepository { calls: AtomicUsize::new(0), context: Mutex::new(initial_context), first_read: Barrier::new(2), release_first_read: Barrier::new(2), }); let mut state = GatewayDataState::with_billing_reader_for_tests(repository.clone()); let stored_global_model = StoredAdminGlobalModel::new( "global-model-1".to_string(), "gpt-5".to_string(), "GPT-5".to_string(), true, None, None, None, None, 1, 1, 0, Some(1_711_000_000), Some(1_711_000_000), ) .expect("stored global model should build"); state.global_model_writer = Some(Arc::new( InMemoryGlobalModelReadRepository::seed(Vec::::new()) .with_admin_global_models([stored_global_model]), )); let state = Arc::new(state); let lookup_state = Arc::clone(&state); let stale_lookup = tokio::spawn(async move { lookup_state .find_billing_model_context("provider-1", Some("key-1"), "gpt-5") .await }); repository.first_read.wait().await; let updated_pricing = json!({"tiers":[{"up_to":null,"input_price_per_1m":3.0,"output_price_per_1m":15.0}]}); let update = UpdateAdminGlobalModelRecord::new( "global-model-1".to_string(), "GPT-5".to_string(), true, None, Some(updated_pricing.clone()), None, None, ) .expect("global model update should build"); repository .context .lock() .expect("mutable billing context lock") .default_tiered_pricing = Some(updated_pricing.clone()); state .update_admin_global_model(&update) .await .expect("global model price update should succeed"); repository.release_first_read.wait().await; let before = stale_lookup .await .expect("initial billing lookup task should complete") .expect("initial billing lookup should succeed") .expect("initial billing context should exist"); assert_eq!(before.default_tiered_pricing, None); assert_eq!(repository.calls.load(Ordering::Acquire), 1); let after = state .find_billing_model_context("provider-1", Some("key-1"), "gpt-5") .await .expect("updated billing lookup should succeed") .expect("updated billing context should exist"); assert_eq!(after.default_tiered_pricing, Some(updated_pricing)); assert_eq!(repository.calls.load(Ordering::Acquire), 2); } #[tokio::test] async fn lists_non_admin_export_users_from_user_reader() { let repository = Arc::new(InMemoryUserReadRepository::seed_export_users(vec![ StoredUserExportRow::new( "user-1".to_string(), Some("alice@example.com".to_string()), true, "alice".to_string(), Some("hash".to_string()), "user".to_string(), "local".to_string(), Some(serde_json::json!(["openai"])), Some(serde_json::json!(["openai:chat"])), Some(serde_json::json!(["gpt-4.1"])), Some(60), Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}})), true, ) .expect("user export row should build"), ])); let state = GatewayDataState::with_user_reader_for_tests(repository); let rows = state .list_non_admin_export_users() .await .expect("export users should succeed"); assert_eq!(rows.len(), 1); assert_eq!(rows[0].username, "alice"); assert!(rows[0].email_verified); assert_eq!(rows[0].password_hash.as_deref(), Some("hash")); assert_eq!(rows[0].allowed_models, Some(vec!["gpt-4.1".to_string()])); assert_eq!( rows[0].model_capability_settings, Some(serde_json::json!({"gpt-4.1": {"cache_1h": true}})) ); } }