mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-12 14:10:19 +08:00
Improve gateway scheduling and runtime admission
This commit is contained in:
elky
@@ -1,6 +1,6 @@
|
||||
use std::collections::BTreeMap;
|
||||
use std::sync::RwLock;
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
use std::time::{Duration, SystemTime, UNIX_EPOCH};
|
||||
|
||||
use async_trait::async_trait;
|
||||
|
||||
@@ -26,11 +26,14 @@ struct MemoryAuthApiKeyIndex {
|
||||
export_by_api_key_id: BTreeMap<String, StoredAuthApiKeyExportRecord>,
|
||||
by_key_hash: BTreeMap<String, String>,
|
||||
touch_counts: BTreeMap<String, usize>,
|
||||
snapshot_lookup_counts: BTreeMap<String, usize>,
|
||||
key_hash_lookup_counts: BTreeMap<String, usize>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Default)]
|
||||
pub struct InMemoryAuthApiKeySnapshotRepository {
|
||||
index: RwLock<MemoryAuthApiKeyIndex>,
|
||||
lookup_delay: Option<Duration>,
|
||||
}
|
||||
|
||||
impl InMemoryAuthApiKeySnapshotRepository {
|
||||
@@ -99,10 +102,18 @@ impl InMemoryAuthApiKeySnapshotRepository {
|
||||
export_by_api_key_id,
|
||||
by_key_hash,
|
||||
touch_counts: BTreeMap::new(),
|
||||
snapshot_lookup_counts: BTreeMap::new(),
|
||||
key_hash_lookup_counts: BTreeMap::new(),
|
||||
}),
|
||||
lookup_delay: None,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn with_lookup_delay_for_tests(mut self, delay: Duration) -> Self {
|
||||
self.lookup_delay = Some(delay);
|
||||
self
|
||||
}
|
||||
|
||||
pub fn with_export_records<I>(mut self, items: I) -> Self
|
||||
where
|
||||
I: IntoIterator<Item = StoredAuthApiKeyExportRecord>,
|
||||
@@ -129,6 +140,26 @@ impl InMemoryAuthApiKeySnapshotRepository {
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub fn snapshot_lookup_count(&self, api_key_id: &str) -> usize {
|
||||
self.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock")
|
||||
.snapshot_lookup_counts
|
||||
.get(api_key_id)
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub fn key_hash_lookup_count(&self, key_hash: &str) -> usize {
|
||||
self.index
|
||||
.read()
|
||||
.expect("auth api key snapshot repository lock")
|
||||
.key_hash_lookup_counts
|
||||
.get(key_hash)
|
||||
.copied()
|
||||
.unwrap_or(0)
|
||||
}
|
||||
|
||||
pub(crate) fn apply_usage_stats_delta(
|
||||
&self,
|
||||
api_key_id: &str,
|
||||
@@ -200,27 +231,46 @@ impl AuthApiKeyReadRepository for InMemoryAuthApiKeySnapshotRepository {
|
||||
&self,
|
||||
key: AuthApiKeyLookupKey<'_>,
|
||||
) -> Result<Option<StoredAuthApiKeySnapshot>, DataLayerError> {
|
||||
let index = self
|
||||
if let Some(delay) = self.lookup_delay {
|
||||
tokio::time::sleep(delay).await;
|
||||
}
|
||||
let mut index = self
|
||||
.index
|
||||
.read()
|
||||
.write()
|
||||
.expect("auth api key snapshot repository lock");
|
||||
Ok(match key {
|
||||
AuthApiKeyLookupKey::KeyHash(key_hash) => index
|
||||
.by_key_hash
|
||||
.get(key_hash)
|
||||
.and_then(|api_key_id| index.by_api_key_id.get(api_key_id))
|
||||
.cloned(),
|
||||
AuthApiKeyLookupKey::KeyHash(key_hash) => {
|
||||
*index
|
||||
.key_hash_lookup_counts
|
||||
.entry(key_hash.to_string())
|
||||
.or_insert(0) += 1;
|
||||
index
|
||||
.by_key_hash
|
||||
.get(key_hash)
|
||||
.and_then(|api_key_id| index.by_api_key_id.get(api_key_id))
|
||||
.cloned()
|
||||
}
|
||||
AuthApiKeyLookupKey::ApiKeyId(api_key_id) => {
|
||||
*index
|
||||
.snapshot_lookup_counts
|
||||
.entry(api_key_id.to_string())
|
||||
.or_insert(0) += 1;
|
||||
index.by_api_key_id.get(api_key_id).cloned()
|
||||
}
|
||||
AuthApiKeyLookupKey::UserApiKeyIds {
|
||||
user_id,
|
||||
api_key_id,
|
||||
} => index
|
||||
.by_api_key_id
|
||||
.get(api_key_id)
|
||||
.filter(|snapshot| snapshot.user_id == user_id)
|
||||
.cloned(),
|
||||
} => {
|
||||
*index
|
||||
.snapshot_lookup_counts
|
||||
.entry(api_key_id.to_string())
|
||||
.or_insert(0) += 1;
|
||||
index
|
||||
.by_api_key_id
|
||||
.get(api_key_id)
|
||||
.filter(|snapshot| snapshot.user_id == user_id)
|
||||
.cloned()
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -8,6 +8,7 @@ use super::{
|
||||
RequestCandidateStatus, RequestCandidateWriteRepository, StoredRequestCandidate,
|
||||
UpsertRequestCandidateRecord,
|
||||
};
|
||||
use crate::driver::postgres::PostgresTransaction;
|
||||
use crate::driver::postgres::PostgresTransactionRunner;
|
||||
use crate::{error::SqlxResultExt, DataLayerError};
|
||||
use aether_data_query::{push_eq, push_in, push_limit, WhereClause};
|
||||
@@ -182,6 +183,109 @@ RETURNING
|
||||
CAST(EXTRACT(EPOCH FROM finished_at) * 1000 AS BIGINT) AS finished_at_unix_ms
|
||||
"#;
|
||||
|
||||
const UPSERT_CONFLICT_SQL: &str = r#"
|
||||
ON CONFLICT (request_id, candidate_index, retry_index)
|
||||
DO UPDATE SET
|
||||
user_id = COALESCE(EXCLUDED.user_id, request_candidates.user_id),
|
||||
api_key_id = COALESCE(EXCLUDED.api_key_id, request_candidates.api_key_id),
|
||||
username = COALESCE(EXCLUDED.username, request_candidates.username),
|
||||
api_key_name = COALESCE(EXCLUDED.api_key_name, request_candidates.api_key_name),
|
||||
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
|
||||
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
|
||||
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
|
||||
status = EXCLUDED.status,
|
||||
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
|
||||
is_cached = COALESCE(EXCLUDED.is_cached, request_candidates.is_cached),
|
||||
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
|
||||
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
|
||||
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
|
||||
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
|
||||
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
|
||||
extra_data = CASE
|
||||
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
|
||||
WHEN EXCLUDED.extra_data IS NULL THEN request_candidates.extra_data
|
||||
WHEN json_typeof(request_candidates.extra_data) = 'object'
|
||||
AND json_typeof(EXCLUDED.extra_data) = 'object'
|
||||
THEN (request_candidates.extra_data::jsonb || EXCLUDED.extra_data::jsonb)::json
|
||||
ELSE EXCLUDED.extra_data
|
||||
END,
|
||||
required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities),
|
||||
created_at = CASE
|
||||
WHEN request_candidates.created_at <= TO_TIMESTAMP(1)
|
||||
THEN EXCLUDED.created_at
|
||||
ELSE request_candidates.created_at
|
||||
END,
|
||||
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
|
||||
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
|
||||
"#;
|
||||
|
||||
const UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL: &str = r#"
|
||||
ON CONFLICT (request_id, candidate_index, retry_index)
|
||||
DO UPDATE SET
|
||||
user_id = COALESCE(EXCLUDED.user_id, request_candidates.user_id),
|
||||
api_key_id = COALESCE(EXCLUDED.api_key_id, request_candidates.api_key_id),
|
||||
username = COALESCE(EXCLUDED.username, request_candidates.username),
|
||||
api_key_name = COALESCE(EXCLUDED.api_key_name, request_candidates.api_key_name),
|
||||
provider_id = COALESCE(EXCLUDED.provider_id, request_candidates.provider_id),
|
||||
endpoint_id = COALESCE(EXCLUDED.endpoint_id, request_candidates.endpoint_id),
|
||||
key_id = COALESCE(EXCLUDED.key_id, request_candidates.key_id),
|
||||
status = EXCLUDED.status,
|
||||
skip_reason = COALESCE(EXCLUDED.skip_reason, request_candidates.skip_reason),
|
||||
is_cached = request_candidates.is_cached,
|
||||
status_code = COALESCE(EXCLUDED.status_code, request_candidates.status_code),
|
||||
error_type = COALESCE(EXCLUDED.error_type, request_candidates.error_type),
|
||||
error_message = COALESCE(EXCLUDED.error_message, request_candidates.error_message),
|
||||
latency_ms = COALESCE(EXCLUDED.latency_ms, request_candidates.latency_ms),
|
||||
concurrent_requests = COALESCE(EXCLUDED.concurrent_requests, request_candidates.concurrent_requests),
|
||||
extra_data = CASE
|
||||
WHEN request_candidates.extra_data IS NULL THEN EXCLUDED.extra_data
|
||||
WHEN EXCLUDED.extra_data IS NULL THEN request_candidates.extra_data
|
||||
WHEN json_typeof(request_candidates.extra_data) = 'object'
|
||||
AND json_typeof(EXCLUDED.extra_data) = 'object'
|
||||
THEN (request_candidates.extra_data::jsonb || EXCLUDED.extra_data::jsonb)::json
|
||||
ELSE EXCLUDED.extra_data
|
||||
END,
|
||||
required_capabilities = COALESCE(EXCLUDED.required_capabilities, request_candidates.required_capabilities),
|
||||
created_at = CASE
|
||||
WHEN request_candidates.created_at <= TO_TIMESTAMP(1)
|
||||
THEN EXCLUDED.created_at
|
||||
ELSE request_candidates.created_at
|
||||
END,
|
||||
started_at = COALESCE(EXCLUDED.started_at, request_candidates.started_at),
|
||||
finished_at = COALESCE(EXCLUDED.finished_at, request_candidates.finished_at)
|
||||
"#;
|
||||
|
||||
const UPSERT_MANY_PREFIX_SQL: &str = r#"
|
||||
INSERT INTO request_candidates (
|
||||
id,
|
||||
request_id,
|
||||
user_id,
|
||||
api_key_id,
|
||||
username,
|
||||
api_key_name,
|
||||
candidate_index,
|
||||
retry_index,
|
||||
provider_id,
|
||||
endpoint_id,
|
||||
key_id,
|
||||
status,
|
||||
skip_reason,
|
||||
is_cached,
|
||||
status_code,
|
||||
error_type,
|
||||
error_message,
|
||||
latency_ms,
|
||||
concurrent_requests,
|
||||
extra_data,
|
||||
required_capabilities,
|
||||
created_at,
|
||||
started_at,
|
||||
finished_at
|
||||
)
|
||||
"#;
|
||||
|
||||
const MAX_POSTGRES_REQUEST_CANDIDATE_UPSERT_ROWS: usize = 1_000;
|
||||
|
||||
const DELETE_CREATED_BEFORE_SQL: &str = r#"
|
||||
DELETE FROM request_candidates
|
||||
WHERE id IN (
|
||||
@@ -501,6 +605,44 @@ impl SqlxRequestCandidateReadRepository {
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn upsert_many(
|
||||
&self,
|
||||
candidates: Vec<UpsertRequestCandidateRecord>,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
if candidates.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
let rows = candidates
|
||||
.into_iter()
|
||||
.map(BatchUpsertRequestCandidateRow::try_from)
|
||||
.collect::<Result<Vec<_>, _>>()?;
|
||||
|
||||
self.tx_runner
|
||||
.run_read_write(|tx| {
|
||||
Box::pin(async move {
|
||||
let mut persisted = 0usize;
|
||||
for ordered_batch in split_request_candidate_upsert_batches(rows) {
|
||||
let (explicit_is_cached, inherited_is_cached): (Vec<_>, Vec<_>) =
|
||||
ordered_batch
|
||||
.into_iter()
|
||||
.partition(|row| row.is_cached.is_some());
|
||||
|
||||
persisted = persisted.saturating_add(
|
||||
execute_partitioned_upsert_many_batch(tx, &explicit_is_cached, true)
|
||||
.await?,
|
||||
);
|
||||
persisted = persisted.saturating_add(
|
||||
execute_partitioned_upsert_many_batch(tx, &inherited_is_cached, false)
|
||||
.await?,
|
||||
);
|
||||
}
|
||||
|
||||
Ok(persisted)
|
||||
}) as BoxFuture<'_, Result<usize, DataLayerError>>
|
||||
})
|
||||
.await
|
||||
}
|
||||
|
||||
pub async fn delete_created_before(
|
||||
&self,
|
||||
created_before_unix_secs: u64,
|
||||
@@ -524,6 +666,178 @@ impl SqlxRequestCandidateReadRepository {
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_partitioned_upsert_many_batch(
|
||||
tx: &mut PostgresTransaction,
|
||||
rows: &[BatchUpsertRequestCandidateRow],
|
||||
overwrite_is_cached: bool,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
let mut persisted = 0usize;
|
||||
for chunk in rows.chunks(MAX_POSTGRES_REQUEST_CANDIDATE_UPSERT_ROWS) {
|
||||
persisted = persisted
|
||||
.saturating_add(execute_upsert_many_batch(tx, chunk, overwrite_is_cached).await?);
|
||||
}
|
||||
Ok(persisted)
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
struct BatchUpsertRequestCandidateRow {
|
||||
id: String,
|
||||
request_id: String,
|
||||
user_id: Option<String>,
|
||||
api_key_id: Option<String>,
|
||||
username: Option<String>,
|
||||
api_key_name: Option<String>,
|
||||
candidate_index: i32,
|
||||
retry_index: i32,
|
||||
provider_id: Option<String>,
|
||||
endpoint_id: Option<String>,
|
||||
key_id: Option<String>,
|
||||
status: &'static str,
|
||||
skip_reason: Option<String>,
|
||||
is_cached: Option<bool>,
|
||||
status_code: Option<i32>,
|
||||
error_type: Option<String>,
|
||||
error_message: Option<String>,
|
||||
latency_ms: Option<i32>,
|
||||
concurrent_requests: Option<i32>,
|
||||
extra_data: Option<serde_json::Value>,
|
||||
required_capabilities: Option<serde_json::Value>,
|
||||
created_at_unix_ms: Option<f64>,
|
||||
started_at_unix_ms: Option<f64>,
|
||||
finished_at_unix_ms: Option<f64>,
|
||||
}
|
||||
|
||||
impl TryFrom<UpsertRequestCandidateRecord> for BatchUpsertRequestCandidateRow {
|
||||
type Error = DataLayerError;
|
||||
|
||||
fn try_from(candidate: UpsertRequestCandidateRecord) -> Result<Self, Self::Error> {
|
||||
candidate.validate()?;
|
||||
Ok(Self {
|
||||
id: if candidate.id.trim().is_empty() {
|
||||
Uuid::new_v4().to_string()
|
||||
} else {
|
||||
candidate.id
|
||||
},
|
||||
request_id: candidate.request_id,
|
||||
user_id: candidate.user_id,
|
||||
api_key_id: candidate.api_key_id,
|
||||
username: candidate.username,
|
||||
api_key_name: candidate.api_key_name,
|
||||
candidate_index: to_i32(candidate.candidate_index)?,
|
||||
retry_index: to_i32(candidate.retry_index)?,
|
||||
provider_id: candidate.provider_id,
|
||||
endpoint_id: candidate.endpoint_id,
|
||||
key_id: candidate.key_id,
|
||||
status: status_to_database(candidate.status),
|
||||
skip_reason: candidate.skip_reason,
|
||||
is_cached: candidate.is_cached,
|
||||
status_code: candidate.status_code.map(i32::from),
|
||||
error_type: candidate.error_type,
|
||||
error_message: candidate.error_message,
|
||||
latency_ms: candidate.latency_ms.map(to_i32_u64).transpose()?,
|
||||
concurrent_requests: candidate.concurrent_requests.map(to_i32).transpose()?,
|
||||
extra_data: candidate.extra_data,
|
||||
required_capabilities: candidate.required_capabilities,
|
||||
created_at_unix_ms: candidate.created_at_unix_ms.map(|value| value as f64),
|
||||
started_at_unix_ms: candidate.started_at_unix_ms.map(|value| value as f64),
|
||||
finished_at_unix_ms: candidate.finished_at_unix_ms.map(|value| value as f64),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
async fn execute_upsert_many_batch(
|
||||
tx: &mut PostgresTransaction,
|
||||
rows: &[BatchUpsertRequestCandidateRow],
|
||||
overwrite_is_cached: bool,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
if rows.is_empty() {
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
let mut builder = QueryBuilder::<Postgres>::new(UPSERT_MANY_PREFIX_SQL);
|
||||
builder.push_values(rows, |mut values, row| {
|
||||
values
|
||||
.push_bind(row.id.clone())
|
||||
.push_bind(row.request_id.clone())
|
||||
.push_bind(row.user_id.clone())
|
||||
.push_bind(row.api_key_id.clone())
|
||||
.push_bind(row.username.clone())
|
||||
.push_bind(row.api_key_name.clone())
|
||||
.push_bind(row.candidate_index)
|
||||
.push_bind(row.retry_index)
|
||||
.push_bind(row.provider_id.clone())
|
||||
.push_bind(row.endpoint_id.clone())
|
||||
.push_bind(row.key_id.clone())
|
||||
.push_bind(row.status)
|
||||
.push_bind(row.skip_reason.clone())
|
||||
.push_bind(row.is_cached.unwrap_or(false))
|
||||
.push_bind(row.status_code)
|
||||
.push_bind(row.error_type.clone())
|
||||
.push_bind(row.error_message.clone())
|
||||
.push_bind(row.latency_ms)
|
||||
.push_bind(row.concurrent_requests)
|
||||
.push_bind(row.extra_data.clone())
|
||||
.push_bind(row.required_capabilities.clone())
|
||||
.push("COALESCE(CASE WHEN ")
|
||||
.push_bind_unseparated(row.created_at_unix_ms)
|
||||
.push_unseparated(" IS NOT NULL AND ")
|
||||
.push_bind_unseparated(row.created_at_unix_ms)
|
||||
.push_unseparated(" > 1000.0 THEN TO_TIMESTAMP(")
|
||||
.push_bind_unseparated(row.created_at_unix_ms)
|
||||
.push_unseparated(" / 1000.0) END, TO_TIMESTAMP(")
|
||||
.push_bind_unseparated(row.started_at_unix_ms)
|
||||
.push_unseparated(" / 1000.0), TO_TIMESTAMP(")
|
||||
.push_bind_unseparated(row.finished_at_unix_ms)
|
||||
.push_unseparated(" / 1000.0), NOW())")
|
||||
.push("TO_TIMESTAMP(")
|
||||
.push_bind_unseparated(row.started_at_unix_ms)
|
||||
.push_unseparated(" / 1000.0)")
|
||||
.push("TO_TIMESTAMP(")
|
||||
.push_bind_unseparated(row.finished_at_unix_ms)
|
||||
.push_unseparated(" / 1000.0)");
|
||||
});
|
||||
builder.push(upsert_many_conflict_sql(overwrite_is_cached));
|
||||
let result = builder
|
||||
.build()
|
||||
.execute(&mut **tx)
|
||||
.await
|
||||
.map_postgres_err()?;
|
||||
Ok(usize::try_from(result.rows_affected()).unwrap_or(rows.len()))
|
||||
}
|
||||
|
||||
fn split_request_candidate_upsert_batches(
|
||||
rows: Vec<BatchUpsertRequestCandidateRow>,
|
||||
) -> Vec<Vec<BatchUpsertRequestCandidateRow>> {
|
||||
let mut batches = Vec::new();
|
||||
let mut current = Vec::new();
|
||||
let mut seen = std::collections::HashSet::<(String, i32, i32)>::new();
|
||||
|
||||
for row in rows {
|
||||
let key = (row.request_id.clone(), row.candidate_index, row.retry_index);
|
||||
if seen.contains(&key) && !current.is_empty() {
|
||||
batches.push(current);
|
||||
current = Vec::new();
|
||||
seen.clear();
|
||||
}
|
||||
seen.insert(key);
|
||||
current.push(row);
|
||||
}
|
||||
|
||||
if !current.is_empty() {
|
||||
batches.push(current);
|
||||
}
|
||||
|
||||
batches
|
||||
}
|
||||
|
||||
fn upsert_many_conflict_sql(overwrite_is_cached: bool) -> &'static str {
|
||||
if overwrite_is_cached {
|
||||
UPSERT_CONFLICT_SQL
|
||||
} else {
|
||||
UPSERT_CONFLICT_INHERIT_IS_CACHED_SQL
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl RequestCandidateReadRepository for SqlxRequestCandidateReadRepository {
|
||||
async fn list_by_request_id(
|
||||
@@ -600,6 +914,13 @@ impl RequestCandidateWriteRepository for SqlxRequestCandidateReadRepository {
|
||||
Self::upsert(self, candidate).await
|
||||
}
|
||||
|
||||
async fn upsert_many(
|
||||
&self,
|
||||
candidates: Vec<UpsertRequestCandidateRecord>,
|
||||
) -> Result<usize, DataLayerError> {
|
||||
Self::upsert_many(self, candidates).await
|
||||
}
|
||||
|
||||
async fn delete_created_before(
|
||||
&self,
|
||||
created_before_unix_secs: u64,
|
||||
|
||||
Reference in New Issue
Block a user