perf(gateway): scale request hot paths for 20k streams

Shard and singleflight hot-path caches, batch and prioritize candidate and usage lifecycle persistence, and extend database and pressure-test instrumentation for 20k concurrent streams.
2026-10-10 19:29:50 +08:00 · 2026-07-22 02:11:08 +08:00
parent 7756c0913f
commit fc92c4f431
124 changed files with 36325 additions and 3217 deletions
@@ -1,8 +1,10 @@
 use super::{
    ApiKeyLastUsedDelta, DataLayerError, GatewayDataState, GeminiFileMappingListQuery,
-    GeminiFileMappingStats, ProviderCatalogKeyListQuery, PublicHealthStatusCount,
-    PublicHealthTimelineBucket, StoredGeminiFileMapping, StoredGeminiFileMappingListPage,
-    StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
+    GeminiFileMappingStats, ProviderCatalogKeyAdaptiveStateUpdate,
+    ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery,
+    ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
+    PublicHealthStatusCount, PublicHealthTimelineBucket, StoredGeminiFileMapping,
+    StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
    StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
    StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, StoredRequestCandidate,
    UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
@@ -370,6 +372,34 @@ impl GatewayDataState {
        Ok(updated)
    }

+    pub(crate) async fn update_provider_catalog_key_oauth_runtime_state(
+        &self,
+        key_id: &str,
+        oauth_invalid_at_unix_secs: Option<u64>,
+        oauth_invalid_reason: Option<&str>,
+        encrypted_auth_config_update: Option<&str>,
+        updated_at_unix_secs: Option<u64>,
+    ) -> Result<bool, DataLayerError> {
+        let updated = match &self.provider_catalog_writer {
+            Some(repository) => {
+                repository
+                    .update_key_oauth_runtime_state(
+                        key_id,
+                        oauth_invalid_at_unix_secs,
+                        oauth_invalid_reason,
+                        encrypted_auth_config_update,
+                        updated_at_unix_secs,
+                    )
+                    .await
+            }
+            None => Ok(false),
+        }?;
+        if updated {
+            self.clear_provider_catalog_cache();
+        }
+        Ok(updated)
+    }
+
    pub(crate) async fn create_provider_catalog_key(
        &self,
        key: &StoredProviderCatalogKey,
@@ -681,4 +711,76 @@ impl GatewayDataState {
        }
        Ok(updated)
    }
+
+    pub(crate) async fn reset_provider_catalog_key_error_count(
+        &self,
+        key_id: &str,
+    ) -> Result<bool, DataLayerError> {
+        let updated = match &self.provider_catalog_writer {
+            Some(repository) => repository.reset_key_error_count(key_id).await,
+            None => Ok(false),
+        }?;
+        if updated {
+            self.clear_provider_catalog_cache();
+        }
+        Ok(updated)
+    }
+
+    pub(crate) async fn compare_and_update_provider_catalog_key_adaptive_state(
+        &self,
+        update: &ProviderCatalogKeyAdaptiveStateUpdate,
+    ) -> Result<bool, DataLayerError> {
+        let Some(repository) = &self.provider_catalog_writer else {
+            return Ok(false);
+        };
+        let updated = repository
+            .compare_and_update_key_adaptive_state(update)
+            .await?;
+        // A false CAS result normally means another instance won the write. Drop the
+        // five-second read cache before the caller reloads and retries.
+        self.clear_provider_catalog_cache();
+        Ok(updated)
+    }
+
+    pub(crate) async fn update_provider_catalog_key_runtime_metadata(
+        &self,
+        update: &ProviderCatalogKeyRuntimeMetadataUpdate,
+    ) -> Result<bool, DataLayerError> {
+        let Some(repository) = &self.provider_catalog_writer else {
+            return Ok(false);
+        };
+        let updated = repository.update_key_runtime_metadata(update).await?;
+        // A false result is a namespace CAS conflict.  Drop the read cache so
+        // the caller's retry observes the writer that won the race.
+        self.clear_provider_catalog_cache();
+        Ok(updated)
+    }
+
+    pub(crate) async fn update_provider_catalog_key_status_snapshot(
+        &self,
+        update: &ProviderCatalogKeyStatusSnapshotUpdate,
+    ) -> Result<bool, DataLayerError> {
+        let Some(repository) = &self.provider_catalog_writer else {
+            return Ok(false);
+        };
+        let updated = repository.update_key_status_snapshot(update).await?;
+        if updated {
+            self.clear_provider_catalog_cache();
+        }
+        Ok(updated)
+    }
+
+    pub(crate) async fn compare_and_update_provider_catalog_key_health_state(
+        &self,
+        update: &ProviderCatalogKeyHealthStateUpdate,
+    ) -> Result<bool, DataLayerError> {
+        let Some(repository) = &self.provider_catalog_writer else {
+            return Ok(false);
+        };
+        let updated = repository
+            .compare_and_update_key_health_state(update)
+            .await?;
+        self.clear_provider_catalog_cache();
+        Ok(updated)
+    }
 }