use aether_data::{DataBackends, DataLayerError, DatabaseDriver}; use aether_data_contracts::repository::candidate_selection::MinimalCandidateSelectionReadRepository; use aether_data_contracts::repository::candidates::RequestCandidateReadRepository; use aether_data_contracts::repository::provider_catalog::ProviderCatalogReadRepository; use aether_runtime_state::RuntimeQueueStore; use std::sync::Arc; use std::time::Duration; use super::{ GatewayDataConfig, GatewayDataState, StoredSystemConfigEntry, SystemConfigValueCacheState, SystemConfigValueInflightCompletion, SystemConfigValueInflightState, }; const SYSTEM_CONFIG_VALUE_CACHE_TTL: Duration = Duration::from_secs(30); const SYSTEM_CONFIG_VALUE_CACHE_MAX_ENTRIES: usize = 512; const SYSTEM_CONFIG_VALUE_CACHE_MAX_INFLIGHT: usize = 512; enum SystemConfigValueLoadRegistration<'a> { Leader(SystemConfigValueLoadGuard<'a>), Follower(Arc), Saturated, } struct SystemConfigValueLoadGuard<'a> { cache: &'a SystemConfigValueCacheState, key: Option, state: Arc, admission: Option, } impl SystemConfigValueInflightState { async fn wait(&self) -> SystemConfigValueInflightCompletion { loop { if let Some(completion) = self.completion.get() { return completion.clone(); } let mut notified = Box::pin(self.notify.notified()); notified.as_mut().enable(); if let Some(completion) = self.completion.get() { return completion.clone(); } notified.await; } } } impl SystemConfigValueCacheState { fn get(&self, key: &str) -> Option> { self.entries.get_fresh(key, SYSTEM_CONFIG_VALUE_CACHE_TTL) } fn register(&self, key: &str) -> SystemConfigValueLoadRegistration<'_> { { let inflight = self .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(state) = inflight.get(key) { return SystemConfigValueLoadRegistration::Follower(Arc::clone(state)); } } let _mutation = self .mutation .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); let mut inflight = self .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); if let Some(state) = inflight.get(key) { return SystemConfigValueLoadRegistration::Follower(Arc::clone(state)); } if inflight.len() >= SYSTEM_CONFIG_VALUE_CACHE_MAX_INFLIGHT { return SystemConfigValueLoadRegistration::Saturated; } let Ok(admission) = Arc::clone(&self.admission).try_acquire_owned() else { return SystemConfigValueLoadRegistration::Saturated; }; let state = Arc::new(SystemConfigValueInflightState { notify: Arc::new(tokio::sync::Notify::new()), completion: std::sync::OnceLock::new(), }); inflight.insert(key.to_string(), Arc::clone(&state)); SystemConfigValueLoadRegistration::Leader(SystemConfigValueLoadGuard { cache: self, key: Some(key.to_string()), state, admission: Some(admission), }) } fn finish_loaded( &self, key: &str, state: &Arc, admission: tokio::sync::OwnedSemaphorePermit, value: Option, ) { self.finish_current( key, state, admission, SystemConfigValueInflightCompletion::Loaded, Some(value), ); } fn finish_current( &self, key: &str, state: &Arc, admission: tokio::sync::OwnedSemaphorePermit, completion: SystemConfigValueInflightCompletion, cache_value: Option>, ) { let _mutation = self .mutation .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); let mut inflight = self .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); drop(admission); debug_assert!(self.admission.available_permits() > 0); if !inflight .get(key) .is_some_and(|current| Arc::ptr_eq(current, state)) { return; } if let Some(value) = cache_value { self.entries.insert( key.to_string(), value, SYSTEM_CONFIG_VALUE_CACHE_TTL, SYSTEM_CONFIG_VALUE_CACHE_MAX_ENTRIES, ); } let completed = state.completion.set(completion).is_ok(); inflight.remove(key); drop(inflight); drop(_mutation); if completed { state.notify.notify_waiters(); } } fn invalidate(&self, key: &str) { let _mutation = self .mutation .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); self.entries.remove(key); let state = self .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .remove(key); let completed = state.as_ref().is_some_and(|state| { state .completion .set(SystemConfigValueInflightCompletion::Invalidated) .is_ok() }); drop(_mutation); if completed { state .expect("completed state should exist") .notify .notify_waiters(); } } fn clear(&self) { let _mutation = self .mutation .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); self.entries.clear(); let states = self .inflight .lock() .unwrap_or_else(std::sync::PoisonError::into_inner) .drain() .map(|(_, state)| state) .collect::>(); let completed = states .iter() .filter(|state| { state .completion .set(SystemConfigValueInflightCompletion::Invalidated) .is_ok() }) .cloned() .collect::>(); drop(_mutation); for state in completed { state.notify.notify_waiters(); } } } impl Default for SystemConfigValueCacheState { fn default() -> Self { Self { entries: aether_cache::ExpiringMap::default(), inflight: std::sync::Mutex::new(std::collections::HashMap::new()), mutation: std::sync::Mutex::new(()), admission: Arc::new(tokio::sync::Semaphore::new( SYSTEM_CONFIG_VALUE_CACHE_MAX_INFLIGHT, )), } } } impl SystemConfigValueLoadGuard<'_> { fn finish_loaded(&mut self, value: Option) { if let Some(key) = self.key.take() { let admission = self .admission .take() .expect("active system config leader must own admission"); self.cache .finish_loaded(&key, &self.state, admission, value); } } fn finish(&mut self, completion: SystemConfigValueInflightCompletion) { if let Some(key) = self.key.take() { let admission = self .admission .take() .expect("active system config leader must own admission"); self.cache .finish_current(&key, &self.state, admission, completion, None); } } } impl Drop for SystemConfigValueLoadGuard<'_> { fn drop(&mut self) { self.finish(SystemConfigValueInflightCompletion::Cancelled); } } fn current_system_config_updated_at_unix_secs() -> u64 { std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap_or_default() .as_secs() } #[cfg(test)] mod system_config_value_cache_tests { use super::*; fn leader<'a>( cache: &'a SystemConfigValueCacheState, key: &str, ) -> SystemConfigValueLoadGuard<'a> { match cache.register(key) { SystemConfigValueLoadRegistration::Leader(guard) => guard, _ => panic!("first registration should lead"), } } #[tokio::test] async fn completion_before_first_poll_releases_system_config_follower() { let cache = SystemConfigValueCacheState::default(); let mut leader = leader(&cache, "key-a"); let follower = match cache.register("key-a") { SystemConfigValueLoadRegistration::Follower(state) => state, _ => panic!("second registration should follow"), }; leader.finish_loaded(Some(serde_json::json!({"version": 1}))); assert!(matches!( tokio::time::timeout(Duration::from_millis(100), follower.wait()) .await .expect("completed follower must not miss the notification"), SystemConfigValueInflightCompletion::Loaded )); assert_eq!( cache.get("key-a"), Some(Some(serde_json::json!({"version": 1}))) ); } #[tokio::test] async fn system_config_follower_receives_leader_failure() { let cache = SystemConfigValueCacheState::default(); let mut leader = leader(&cache, "key-a"); let follower = match cache.register("key-a") { SystemConfigValueLoadRegistration::Follower(state) => state, _ => panic!("second registration should follow"), }; leader.finish(SystemConfigValueInflightCompletion::Failed( DataLayerError::Sql("forced system config failure".to_string()), )); let completion = tokio::time::timeout(Duration::from_millis(100), follower.wait()) .await .expect("failed load should release its follower"); let SystemConfigValueInflightCompletion::Failed(error) = completion else { panic!("follower should observe the leader failure"); }; assert_eq!(error.to_string(), "sql error: forced system config failure"); } #[tokio::test] async fn invalidation_wins_and_old_guard_preserves_replacement() { let cache = SystemConfigValueCacheState::default(); let mut old_leader = leader(&cache, "key-a"); let old_follower = match cache.register("key-a") { SystemConfigValueLoadRegistration::Follower(state) => state, _ => panic!("second registration should follow"), }; cache.invalidate("key-a"); let replacement = leader(&cache, "key-a"); old_leader.finish_loaded(Some(serde_json::json!({"stale": true}))); assert!(matches!( old_follower.wait().await, SystemConfigValueInflightCompletion::Invalidated )); assert_eq!(cache.get("key-a"), None); assert!(matches!( cache.register("key-a"), SystemConfigValueLoadRegistration::Follower(_) )); drop(replacement); } #[test] fn system_config_states_are_independent_and_inflight_is_bounded() { let first = SystemConfigValueCacheState::default(); let second = SystemConfigValueCacheState::default(); let first_guard = leader(&first, "shared-key"); let second_guard = leader(&second, "shared-key"); drop(first_guard); drop(second_guard); let mut guards = Vec::with_capacity(SYSTEM_CONFIG_VALUE_CACHE_MAX_INFLIGHT); for index in 0..SYSTEM_CONFIG_VALUE_CACHE_MAX_INFLIGHT { let key = format!("bounded-{index}"); guards.push(leader(&first, &key)); } assert!(matches!( first.register("over-capacity"), SystemConfigValueLoadRegistration::Saturated )); drop(guards); } #[test] fn capacity_full_cancelled_system_config_follower_can_retry() { let cache = SystemConfigValueCacheState::default(); let mut active = Vec::with_capacity(SYSTEM_CONFIG_VALUE_CACHE_MAX_INFLIGHT - 1); for index in 0..SYSTEM_CONFIG_VALUE_CACHE_MAX_INFLIGHT - 1 { active.push(leader(&cache, &format!("active-{index}"))); } let current = leader(&cache, "retry-key"); let follower = match cache.register("retry-key") { SystemConfigValueLoadRegistration::Follower(state) => state, _ => panic!("same-key request should follow at full capacity"), }; assert_eq!(cache.admission.available_permits(), 0); assert!(matches!( cache.register("over-capacity"), SystemConfigValueLoadRegistration::Saturated )); drop(current); assert!(matches!( follower.completion.get(), Some(SystemConfigValueInflightCompletion::Cancelled) )); let mut replacement = leader(&cache, "retry-key"); replacement.finish(SystemConfigValueInflightCompletion::Cancelled); assert_eq!(cache.admission.available_permits(), 1); drop(active); } } impl GatewayDataState { pub(crate) fn disabled() -> Self { Self::default() } pub(crate) fn from_config(config: GatewayDataConfig) -> Result { if !config.is_enabled() { return Ok(Self { config, backends: None, auth_api_key_reader: None, auth_api_key_writer: None, auth_module_reader: None, auth_module_writer: None, announcement_reader: None, announcement_writer: None, management_token_reader: None, management_token_writer: None, oauth_provider_reader: None, oauth_provider_writer: None, proxy_node_reader: None, proxy_node_writer: None, billing_reader: None, background_task_reader: None, background_task_writer: None, gemini_file_mapping_reader: None, gemini_file_mapping_writer: None, global_model_reader: None, global_model_writer: None, minimal_candidate_selection_reader: None, request_candidate_reader: None, request_candidate_writer: None, provider_catalog_reader: None, provider_catalog_writer: None, pool_score_reader: None, pool_score_writer: None, provider_quota_reader: None, provider_quota_writer: None, routing_group_reader: None, routing_group_writer: None, usage_reader: None, usage_writer: None, user_reader: None, user_preferences: None, usage_worker_queue: None, video_task_reader: None, video_task_writer: None, wallet_reader: None, wallet_writer: None, settlement_writer: None, system_config_values: None, system_config_value_cache: Default::default(), billing_model_context_cache: Default::default(), }); } let backends = DataBackends::from_config(config.to_data_layer_config())?; let auth_api_key_reader = backends.read().auth_api_keys().map(|repository| { Arc::new(super::auth_api_key_cache::CachedAuthApiKeyReadRepository::new(repository)) as Arc }); let auth_api_key_writer = backends.write().auth_api_keys(); let auth_module_reader = backends.read().auth_modules(); let auth_module_writer = backends.write().auth_modules(); let announcement_reader = backends.read().announcements(); let announcement_writer = backends.write().announcements(); let management_token_reader = backends.read().management_tokens(); let management_token_writer = backends.write().management_tokens(); let oauth_provider_reader = backends.read().oauth_providers(); let oauth_provider_writer = backends.write().oauth_providers(); let proxy_node_reader = backends.read().proxy_nodes(); let proxy_node_writer = backends.write().proxy_nodes(); let billing_reader = backends.read().billing(); let background_task_reader = backends.read().background_tasks(); let background_task_writer = backends.write().background_tasks(); let gemini_file_mapping_reader = backends.read().gemini_file_mappings(); let global_model_reader = backends.read().global_models(); let global_model_writer = backends.write().global_models(); let minimal_candidate_selection_reader = backends .read() .minimal_candidate_selection() .map(|repository| { Arc::new( super::candidate_cache::CachedMinimalCandidateSelectionReadRepository::new( repository, ), ) as Arc }); let request_candidate_reader = backends.read().request_candidates().map(|repository| { Arc::new( super::request_candidate_cache::CachedRequestCandidateReadRepository::new( repository, ), ) as Arc }); let request_candidate_writer = backends.write().request_candidates(); let gemini_file_mapping_writer = backends.write().gemini_file_mappings(); let provider_catalog_reader = backends.read().provider_catalog().map(|repository| { Arc::new( super::provider_catalog_cache::CachedProviderCatalogReadRepository::new(repository), ) as Arc }); let provider_catalog_writer = backends.write().provider_catalog(); let pool_score_reader = backends.read().pool_scores(); let pool_score_writer = backends.write().pool_scores(); let provider_quota_reader = backends.read().provider_quotas(); let provider_quota_writer = backends.write().provider_quotas(); let routing_group_reader = backends.read().routing_groups().map(|repository| { Arc::new(super::routing_group_cache::CachedRoutingGroupReadRepository::new( repository, )) as Arc }); let routing_group_writer = backends.write().routing_groups(); let usage_reader = backends.read().usage(); let usage_writer = backends.write().usage(); let user_reader = backends.read().users(); let usage_worker_queue = None; let video_task_reader = backends.read().video_tasks(); let video_task_writer = backends.write().video_tasks(); let wallet_reader = backends.read().wallets(); let wallet_writer = backends.write().wallets(); let settlement_writer = backends.write().settlement(); Ok(Self { config, backends: Some(backends), auth_api_key_reader, auth_api_key_writer, auth_module_reader, auth_module_writer, announcement_reader, announcement_writer, management_token_reader, management_token_writer, oauth_provider_reader, oauth_provider_writer, proxy_node_reader, proxy_node_writer, billing_reader, background_task_reader, background_task_writer, gemini_file_mapping_reader, gemini_file_mapping_writer, global_model_reader, global_model_writer, minimal_candidate_selection_reader, request_candidate_reader, request_candidate_writer, provider_catalog_reader, provider_catalog_writer, pool_score_reader, pool_score_writer, provider_quota_reader, provider_quota_writer, routing_group_reader, routing_group_writer, usage_reader, usage_writer, user_reader, user_preferences: None, usage_worker_queue, video_task_reader, video_task_writer, wallet_reader, wallet_writer, settlement_writer, system_config_values: None, system_config_value_cache: Default::default(), billing_model_context_cache: Default::default(), }) } pub(crate) fn has_backends(&self) -> bool { self.backends.is_some() } pub(crate) fn with_usage_worker_queue( mut self, queue: Option>, ) -> Self { self.usage_worker_queue = queue; self } pub(crate) fn has_database_maintenance_backend(&self) -> bool { self.backends .as_ref() .is_some_and(|backends| backends.has_database_maintenance_backend()) } pub(crate) fn has_database_pool_summary(&self) -> bool { self.backends .as_ref() .is_some_and(|backends| backends.has_database_pool_summary()) } pub(crate) fn has_wallet_daily_usage_aggregation_backend(&self) -> bool { self.backends .as_ref() .is_some_and(|backends| backends.has_wallet_daily_usage_aggregation_backend()) } pub(crate) fn has_stats_hourly_aggregation_backend(&self) -> bool { self.backends .as_ref() .is_some_and(|backends| backends.has_stats_hourly_aggregation_backend()) } pub(crate) fn has_stats_daily_aggregation_backend(&self) -> bool { self.backends .as_ref() .is_some_and(|backends| backends.has_stats_daily_aggregation_backend()) } pub(crate) fn has_auth_api_key_reader(&self) -> bool { self.auth_api_key_reader.is_some() } pub(crate) fn has_auth_api_key_writer(&self) -> bool { self.auth_api_key_writer.is_some() } pub(crate) fn has_auth_module_writer(&self) -> bool { self.auth_module_writer.is_some() } pub(crate) fn has_announcement_reader(&self) -> bool { self.announcement_reader.is_some() } pub(crate) fn has_announcement_writer(&self) -> bool { self.announcement_writer.is_some() } pub(crate) fn has_background_task_reader(&self) -> bool { self.background_task_reader.is_some() } pub(crate) fn has_background_task_writer(&self) -> bool { self.background_task_writer.is_some() } pub(crate) fn has_audit_log_reader(&self) -> bool { self.backends .as_ref() .and_then(|backends| backends.read().audit_logs()) .is_some() } pub(crate) fn has_management_token_reader(&self) -> bool { self.management_token_reader.is_some() } pub(crate) fn has_management_token_writer(&self) -> bool { self.management_token_writer.is_some() } pub(crate) fn has_gemini_file_mapping_reader(&self) -> bool { self.gemini_file_mapping_reader.is_some() } pub(crate) fn has_gemini_file_mapping_writer(&self) -> bool { self.gemini_file_mapping_writer.is_some() } pub(crate) fn has_global_model_reader(&self) -> bool { self.global_model_reader.is_some() } pub(crate) fn has_global_model_writer(&self) -> bool { self.global_model_writer.is_some() } #[allow(dead_code)] pub(crate) fn has_minimal_candidate_selection_reader(&self) -> bool { self.minimal_candidate_selection_reader.is_some() } pub(crate) fn clear_minimal_candidate_selection_cache(&self) { if let Some(repository) = &self.minimal_candidate_selection_reader { repository.clear_local_cache(); } } pub(crate) fn clear_routing_group_cache(&self) { if let Some(repository) = &self.routing_group_reader { repository.clear_local_cache(); } } pub(crate) fn clear_provider_catalog_cache(&self) { if let Some(repository) = &self.provider_catalog_reader { repository.clear_local_cache(); } } pub(crate) fn has_request_candidate_reader(&self) -> bool { self.request_candidate_reader.is_some() } pub(crate) fn has_request_candidate_writer(&self) -> bool { self.request_candidate_writer.is_some() } pub(crate) fn request_candidate_writer( &self, ) -> Option< Arc, > { self.request_candidate_writer.clone() } pub(crate) fn has_routing_group_reader(&self) -> bool { self.routing_group_reader.is_some() } pub(crate) fn has_routing_group_writer(&self) -> bool { self.routing_group_writer.is_some() } pub(crate) fn has_provider_catalog_reader(&self) -> bool { self.provider_catalog_reader.is_some() } pub(crate) fn has_provider_catalog_writer(&self) -> bool { self.provider_catalog_writer.is_some() } pub(crate) fn has_pool_score_reader(&self) -> bool { self.pool_score_reader.is_some() } pub(crate) fn has_pool_score_writer(&self) -> bool { self.pool_score_writer.is_some() } pub(crate) fn has_proxy_node_reader(&self) -> bool { self.proxy_node_reader.is_some() } pub(crate) fn has_proxy_node_writer(&self) -> bool { self.proxy_node_writer.is_some() } pub(crate) fn has_system_config_store(&self) -> bool { self.system_config_values.is_some() || self .backends .as_ref() .is_some_and(|backends| backends.has_system_config_backend()) } pub(crate) fn database_driver(&self) -> Option { self.backends .as_ref() .and_then(|backends| backends.database_driver()) } pub(crate) fn has_provider_quota_writer(&self) -> bool { self.provider_quota_writer.is_some() } pub(crate) fn has_usage_reader(&self) -> bool { self.usage_reader.is_some() } pub(crate) fn has_user_reader(&self) -> bool { self.user_reader.is_some() } pub(crate) fn has_usage_writer(&self) -> bool { self.usage_writer.is_some() } pub(crate) fn has_usage_counter_flush_backend(&self) -> bool { self.has_usage_writer() && database_driver_supports_usage_counter_flush(self.database_driver()) } pub(crate) fn has_usage_worker_queue(&self) -> bool { self.usage_worker_queue.is_some() } pub(crate) fn has_video_task_reader(&self) -> bool { self.video_task_reader.is_some() } pub(crate) fn has_video_task_writer(&self) -> bool { self.video_task_writer.is_some() } pub(crate) fn has_wallet_reader(&self) -> bool { self.wallet_reader.is_some() } #[cfg(test)] pub(crate) fn without_wallet_reader_for_tests(mut self) -> Self { self.wallet_reader = None; self } pub(crate) fn has_wallet_writer(&self) -> bool { self.wallet_writer.is_some() } pub(crate) fn has_settlement_writer(&self) -> bool { self.settlement_writer.is_some() } #[allow(dead_code)] pub(crate) fn encryption_key(&self) -> Option<&str> { self.config.encryption_key() } pub(crate) async fn find_system_config_value( &self, key: &str, ) -> Result, DataLayerError> { if let Some(values) = &self.system_config_values { return Ok(values .read() .expect("system config values lock") .get(key) .map(|entry| entry.value.clone())); } if let Some(value) = self.system_config_value_cache.get(key) { return Ok(value); } loop { match self.system_config_value_cache.register(key) { SystemConfigValueLoadRegistration::Saturated => { return Err(DataLayerError::TimedOut(format!( "system config cache admission saturated for key '{key}'" ))); } SystemConfigValueLoadRegistration::Follower(state) => match state.wait().await { SystemConfigValueInflightCompletion::Failed(error) => { return Err(error); } SystemConfigValueInflightCompletion::Loaded => { if let Some(value) = self.system_config_value_cache.get(key) { return Ok(value); } } SystemConfigValueInflightCompletion::Cancelled | SystemConfigValueInflightCompletion::Invalidated => {} }, SystemConfigValueLoadRegistration::Leader(mut guard) => { if let Some(value) = self.system_config_value_cache.get(key) { guard.finish(SystemConfigValueInflightCompletion::Loaded); return Ok(value); } match self.load_system_config_value_uncached(key).await { Ok(value) => { guard.finish_loaded(value.clone()); return Ok(value); } Err(error) => { guard .finish(SystemConfigValueInflightCompletion::Failed(error.clone())); return Err(error); } } } } } } async fn load_system_config_value_uncached( &self, key: &str, ) -> Result, DataLayerError> { let Some(backends) = self.backends.as_ref() else { return Ok(None); }; crate::request_diagnostics::observe_db_operation( "system_config_value", self.database_pool_summary(), backends.find_system_config_value(key), ) .await } pub(crate) async fn find_system_config_value_strong( &self, key: &str, ) -> Result, DataLayerError> { if let Some(values) = &self.system_config_values { return Ok(values .read() .expect("system config values lock") .get(key) .map(|entry| entry.value.clone())); } let Some(backends) = self.backends.as_ref() else { return Ok(None); }; crate::request_diagnostics::observe_db_operation( "system_config_value_strong", self.database_pool_summary(), backends.find_system_config_value(key), ) .await } pub(crate) async fn compare_and_set_system_config_string_value( &self, key: &str, expected: &str, replacement: &str, ) -> Result { if let Some(values) = &self.system_config_values { let updated = { let mut values = values.write().expect("system config values lock"); match values.get_mut(key) { Some(entry) if entry.value.as_str() == Some(expected) => { entry.value = serde_json::Value::String(replacement.to_string()); entry.updated_at_unix_secs = Some(current_system_config_updated_at_unix_secs()); true } _ => false, } }; self.clear_cached_system_config_value(key); return Ok(updated); } let result = match self.backends.as_ref() { Some(backends) => { crate::request_diagnostics::observe_db_operation( "system_config_compare_and_set", self.database_pool_summary(), backends.compare_and_set_system_config_string_value(key, expected, replacement), ) .await } None => Ok(false), }; self.clear_cached_system_config_value(key); result } pub(crate) async fn upsert_system_config_value( &self, key: &str, value: &serde_json::Value, description: Option<&str>, ) -> Result { Ok(self .upsert_system_config_entry(key, value, description) .await? .value) } pub(crate) async fn list_system_config_entries( &self, ) -> Result, DataLayerError> { if let Some(values) = &self.system_config_values { return Ok(values .read() .expect("system config values lock") .values() .cloned() .collect()); } let Some(backends) = self.backends.as_ref() else { return Ok(Vec::new()); }; backends.list_system_config_entries().await } pub(crate) async fn upsert_system_config_entry( &self, key: &str, value: &serde_json::Value, description: Option<&str>, ) -> Result { if let Some(values) = &self.system_config_values { let mut values = values.write().expect("system config values lock"); let description = description .map(ToOwned::to_owned) .or_else(|| values.get(key).and_then(|entry| entry.description.clone())); let entry = StoredSystemConfigEntry { key: key.to_string(), value: value.clone(), description, updated_at_unix_secs: Some(current_system_config_updated_at_unix_secs()), }; values.insert(key.to_string(), entry.clone()); self.clear_cached_system_config_value(key); return Ok(entry); } if let Some(backends) = self.backends.as_ref() { if let Some(entry) = backends .upsert_system_config_entry(key, value, description) .await? { self.clear_cached_system_config_value(key); return Ok(entry); } } self.clear_cached_system_config_value(key); Ok(StoredSystemConfigEntry { key: key.to_string(), value: value.clone(), description: description.map(ToOwned::to_owned), updated_at_unix_secs: Some(current_system_config_updated_at_unix_secs()), }) } pub(crate) async fn delete_system_config_value( &self, key: &str, ) -> Result { if let Some(values) = &self.system_config_values { let deleted = values .write() .expect("system config values lock") .remove(key) .is_some(); self.clear_cached_system_config_value(key); return Ok(deleted); } let Some(backends) = self.backends.as_ref() else { return Ok(false); }; let deleted = backends.delete_system_config_value(key).await?; self.clear_cached_system_config_value(key); Ok(deleted) } fn clear_cached_system_config_value(&self, key: &str) { self.system_config_value_cache.invalidate(key); } pub(crate) async fn read_admin_system_stats( &self, ) -> Result { match self.backends.as_ref() { Some(backends) => backends.read_admin_system_stats().await, None => Ok(super::AdminSystemStats::default()), } } pub(crate) async fn purge_admin_system_data( &self, target: aether_data::repository::system::AdminSystemPurgeTarget, ) -> Result { let purges_config = matches!( &target, aether_data::repository::system::AdminSystemPurgeTarget::Config ); if purges_config { if let Some(values) = &self.system_config_values { let mut values = values.write().expect("system config values lock"); let deleted = values.len() as u64; values.clear(); self.system_config_value_cache.clear(); let mut summary = aether_data::repository::system::AdminSystemPurgeSummary::default(); summary.add("system_configs", deleted); return Ok(summary); } } let result = match self.backends.as_ref() { Some(backends) => backends.purge_admin_system_data(target).await, None => Ok(aether_data::repository::system::AdminSystemPurgeSummary::default()), }; if purges_config && result.is_ok() { self.system_config_value_cache.clear(); } result } pub(crate) async fn export_admin_system_usage_aggregates( &self, ) -> Result { match self.backends.as_ref() { Some(backends) => backends.export_admin_system_usage_aggregates().await, None => { Ok(aether_data::repository::system::AdminSystemUsageAggregateSnapshot::default()) } } } pub(crate) async fn import_admin_system_usage_aggregates( &self, snapshot: &aether_data::repository::system::AdminSystemUsageAggregateSnapshot, user_id_map: &std::collections::BTreeMap, api_key_id_map: &std::collections::BTreeMap, mode: aether_data::repository::system::AdminSystemUsageAggregateImportMode, ) -> Result< aether_data::repository::system::AdminSystemUsageAggregateImportSummary, DataLayerError, > { match self.backends.as_ref() { Some(backends) => { backends .import_admin_system_usage_aggregates( snapshot, user_id_map, api_key_id_map, mode, ) .await } None => Ok( aether_data::repository::system::AdminSystemUsageAggregateImportSummary::default(), ), } } pub(crate) async fn purge_admin_request_bodies_batch( &self, batch_size: usize, ) -> Result { match self.backends.as_ref() { Some(backends) => backends.purge_admin_request_bodies_batch(batch_size).await, None => Ok(aether_data::repository::system::AdminSystemPurgeSummary::default()), } } } fn database_driver_supports_usage_counter_flush(driver: Option) -> bool { matches!( driver, Some(DatabaseDriver::Postgres | DatabaseDriver::Mysql | DatabaseDriver::Sqlite) ) } #[cfg(test)] mod usage_counter_flush_backend_tests { use super::database_driver_supports_usage_counter_flush; use aether_data::DatabaseDriver; #[test] fn every_sql_driver_supports_usage_counter_flush() { assert!(database_driver_supports_usage_counter_flush(Some( DatabaseDriver::Postgres ))); assert!(database_driver_supports_usage_counter_flush(Some( DatabaseDriver::Mysql ))); assert!(database_driver_supports_usage_counter_flush(Some( DatabaseDriver::Sqlite ))); assert!(!database_driver_supports_usage_counter_flush(None)); } }