i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-08 02:17:46 +08:00
Code Issues Packages Projects Releases Wiki Activity

perf: queue request candidate persistence

Browse Source
This commit is contained in:
elky
2026-06-22 02:49:17 +08:00
parent f75894acbb
commit 5b7805181b
11 changed files with 1143 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/data/state/core.rs
+8
View File
@@ -322,6 +322,14 @@ impl GatewayDataState {
self.request_candidate_writer.is_some()
}
pub(crate) fn request_candidate_writer(
&self,
) -> Option<
Arc<dyn aether_data_contracts::repository::candidates::RequestCandidateWriteRepository>,
> {
self.request_candidate_writer.clone()
}
pub(crate) fn has_routing_group_reader(&self) -> bool {
self.routing_group_reader.is_some()
}
apps/aether-gateway/src/lib.rs
+1
View File
@@ -61,6 +61,7 @@ mod provider_key_auth;
mod provider_pool_demand;
pub(crate) use aether_provider_transport as provider_transport;
mod rate_limit;
mod request_candidate_queue;
mod request_candidate_runtime;
mod roles;
mod router;
apps/aether-gateway/src/orchestration/effects.rs
+163 -3
View File
@@ -38,7 +38,9 @@ use crate::handlers::shared::provider_pool::{
use crate::orchestration::local_execution_candidate_metadata_from_report_context;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::scheduler::config::{read_scheduler_ordering_config, SchedulerSchedulingMode};
use crate::AppState;
use crate::{
provider_transport::snapshot::GatewayProviderTransportProvider, AppState, GatewayError,
};
const POOL_SCORE_FEEDBACK_GATE_MAX_ENTRIES: usize = 50_000;
const POOL_SCORE_SUCCESS_FEEDBACK_MIN_INTERVAL_ENV: &str =
@@ -900,11 +902,88 @@ async fn record_oauth_invalidation_effect(
"last_request_feedback": {
"source": "oauth_invalidation",
"status_code": effect.status_code,
"reason": invalid_reason
"reason": invalid_reason.as_str()
}
}),
)
.await;
match auto_remove_runtime_oauth_invalid_key(
state,
&transport.provider,
&plan.key_id,
invalid_reason.as_str(),
)
.await
{
Ok(true) => {
tracing::info!(
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
provider_type = %transport.provider.provider_type,
event_name = "auto_removed_oauth_runtime_invalid",
"gateway auto-removed runtime invalid oauth key"
);
}
Ok(false) => {}
Err(err) => {
warn!(
"gateway orchestration effects: failed to auto-remove oauth invalid key for provider {} endpoint {} key {}: {:?}",
plan.provider_id, plan.endpoint_id, plan.key_id, err
);
}
}
}
async fn auto_remove_runtime_oauth_invalid_key(
state: &AppState,
provider: &GatewayProviderTransportProvider,
key_id: &str,
invalid_reason: &str,
) -> Result<bool, GatewayError> {
if !admin_provider_quota_pure::provider_auto_remove_banned_keys(provider.config.as_ref()) {
return Ok(false);
}
let key_ids = [key_id.to_string()];
let Some(key) = state
.read_provider_catalog_keys_by_ids(&key_ids)
.await?
.into_iter()
.next()
else {
return Ok(false);
};
if key.provider_id != provider.id {
return Ok(false);
}
if !admin_provider_quota_pure::should_auto_remove_oauth_invalid_key(
&key,
Some(invalid_reason),
true,
current_unix_secs(),
) {
return Ok(false);
}
let deleted_key_id = key.id.clone();
if !state.delete_provider_catalog_key(&deleted_key_id).await? {
return Ok(false);
}
state
.cleanup_deleted_provider_catalog_refs(
&provider.id,
false,
&[],
std::slice::from_ref(&deleted_key_id),
)
.await?;
let _ = state
.invalidate_local_oauth_refresh_entry(&deleted_key_id)
.await;
Ok(true)
}
fn resolve_local_oauth_invalid_reason(
@@ -1309,6 +1388,16 @@ mod tests {
)
}
fn sample_codex_provider_with_auto_remove() -> StoredProviderCatalogProvider {
let mut provider = sample_codex_provider();
provider.config = Some(json!({
"pool_advanced": {
"auto_remove_banned_keys": true
}
}));
provider
}
fn sample_codex_endpoint() -> StoredProviderCatalogEndpoint {
StoredProviderCatalogEndpoint::new(
"endpoint-codex-cli-local-1".to_string(),
@@ -1363,8 +1452,16 @@ mod tests {
}
fn codex_state() -> AppState {
codex_state_with_provider(sample_codex_provider())
}
fn codex_state_with_auto_remove() -> AppState {
codex_state_with_provider(sample_codex_provider_with_auto_remove())
}
fn codex_state_with_provider(provider: StoredProviderCatalogProvider) -> AppState {
let repository = Arc::new(InMemoryProviderCatalogReadRepository::seed(
vec![sample_codex_provider()],
vec![provider],
vec![sample_codex_endpoint()],
vec![sample_codex_key()],
));
@@ -2279,6 +2376,69 @@ mod tests {
);
}
#[tokio::test]
async fn oauth_invalidation_auto_removes_inactive_pat_owner_when_enabled() {
let state = codex_state_with_auto_remove();
let plan = sample_codex_plan();
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: None,
},
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
status_code: 403,
response_text: Some(
r#"{"error":{"code":"biscuit_baker_service_auth_credential_error_status","message":"Personal access token owner is inactive."},"status":403}"#,
),
}),
)
.await;
let keys = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
.await
.expect("provider catalog keys should load");
assert!(
keys.is_empty(),
"hard-invalid PAT owner should be auto removed"
);
}
#[tokio::test]
async fn oauth_invalidation_auto_remove_keeps_recoverable_expired_token() {
let state = codex_state_with_auto_remove();
let plan = sample_codex_plan();
apply_local_execution_effect(
&state,
LocalExecutionEffectContext {
plan: &plan,
report_context: None,
},
LocalExecutionEffect::OauthInvalidation(LocalOAuthInvalidationEffect {
status_code: 401,
response_text: Some(
r#"{"error":{"message":"session expired","type":"invalid_request_error"}}"#,
),
}),
)
.await;
let stored_key = state
.read_provider_catalog_keys_by_ids(std::slice::from_ref(&plan.key_id))
.await
.expect("provider catalog keys should load")
.into_iter()
.next()
.expect("recoverable expired token should be retained");
assert_eq!(
stored_key.oauth_invalid_reason.as_deref(),
Some("[OAUTH_EXPIRED] session expired")
);
}
#[tokio::test]
async fn health_failure_projection_updates_key_health_for_format() {
let state = health_state();
apps/aether-gateway/src/provider_pool_demand.rs
+36 -9
View File
@@ -15,6 +15,9 @@ const PROVIDER_POOL_DEMAND_SNAPSHOT_PREFIX: &str = "ap:provider_pool:demand";
const PROVIDER_POOL_BURST_PENDING_PREFIX: &str = "ap:quota_probe:burst_pending";
const PROVIDER_POOL_IN_FLIGHT_TOKEN_TTL_MS: u64 = 120_000;
const PROVIDER_POOL_IN_FLIGHT_RENEW_MS: u64 = 30_000;
const PROVIDER_POOL_IN_FLIGHT_ACQUIRE_TIMEOUT_ENV: &str =
"AETHER_GATEWAY_PROVIDER_POOL_IN_FLIGHT_ACQUIRE_TIMEOUT_MS";
const DEFAULT_PROVIDER_POOL_IN_FLIGHT_ACQUIRE_TIMEOUT_MS: u64 = 10;
const PROVIDER_POOL_DEMAND_SNAPSHOT_TTL_SECONDS: u64 = 6 * 60 * 60;
const PROVIDER_POOL_DEMAND_ALPHA: f64 = 0.2;
const PROVIDER_POOL_DEMAND_HEADROOM: f64 = 1.2;
@@ -131,6 +134,17 @@ fn token_expiry_score(now_ms: u64) -> f64 {
now_ms.saturating_add(PROVIDER_POOL_IN_FLIGHT_TOKEN_TTL_MS) as f64
}
fn provider_pool_in_flight_acquire_timeout() -> Duration {
static TIMEOUT: std::sync::OnceLock<Duration> = std::sync::OnceLock::new();
*TIMEOUT.get_or_init(|| {
let millis = std::env::var(PROVIDER_POOL_IN_FLIGHT_ACQUIRE_TIMEOUT_ENV)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.unwrap_or(DEFAULT_PROVIDER_POOL_IN_FLIGHT_ACQUIRE_TIMEOUT_MS);
Duration::from_millis(millis)
})
}
fn spawn_in_flight_renewal(
runtime: Arc<RuntimeState>,
tokens_key: String,
@@ -173,16 +187,29 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard(
let tokens_key = in_flight_tokens_key(provider_id);
let token = build_in_flight_token(request_id, candidate_id, key_id);
if let Err(err) = runtime
.score_set(&tokens_key, &token, token_expiry_score(current_unix_ms()))
.await
match tokio::time::timeout(
provider_pool_in_flight_acquire_timeout(),
runtime.score_set(&tokens_key, &token, token_expiry_score(current_unix_ms())),
)
.await
{
debug!(
provider_id,
error = ?err,
"gateway provider pool demand: failed to acquire in-flight token"
);
return None;
Ok(Ok(())) => {}
Ok(Err(err)) => {
debug!(
provider_id,
error = ?err,
"gateway provider pool demand: failed to acquire in-flight token"
);
return None;
}
Err(_) => {
debug!(
provider_id,
timeout_ms = provider_pool_in_flight_acquire_timeout().as_millis() as u64,
"gateway provider pool demand: skipped in-flight token after acquire timeout"
);
return None;
}
}
let stop_renewal = Arc::new(AtomicBool::new(false));
apps/aether-gateway/src/request_candidate_queue.rs
+751
View File
@@ -0,0 +1,751 @@
use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Duration;
use aether_data_contracts::repository::candidates::{
RequestCandidateStatus, RequestCandidateWriteRepository, UpsertRequestCandidateRecord,
};
use aether_runtime::{MetricKind, MetricSample};
use tokio::sync::mpsc;
use tokio::time::{interval, MissedTickBehavior};
use tracing::{debug, warn};
const MODE_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_WRITE_MODE";
const QUEUE_CAPACITY_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_CAPACITY";
const BATCH_SIZE_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_BATCH_SIZE";
const FLUSH_INTERVAL_MS_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_FLUSH_INTERVAL_MS";
const WORKERS_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_WORKERS";
const QUEUE_FULL_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_QUEUE_FULL";
const DEFAULT_QUEUE_CAPACITY: usize = 65_536;
const DEFAULT_BATCH_SIZE: usize = 512;
const DEFAULT_FLUSH_INTERVAL_MS: u64 = 50;
const DEFAULT_WORKERS: usize = 2;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RequestCandidateWriteMode {
Sync,
Async,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum RequestCandidateQueueFullPolicy {
Drop,
Sync,
}
#[derive(Debug, Clone)]
pub(crate) struct RequestCandidateQueueConfig {
pub(crate) mode: RequestCandidateWriteMode,
pub(crate) capacity: usize,
pub(crate) batch_size: usize,
pub(crate) flush_interval: Duration,
pub(crate) workers: usize,
pub(crate) full_policy: RequestCandidateQueueFullPolicy,
}
impl Default for RequestCandidateQueueConfig {
fn default() -> Self {
Self {
mode: RequestCandidateWriteMode::Sync,
capacity: DEFAULT_QUEUE_CAPACITY,
batch_size: DEFAULT_BATCH_SIZE,
flush_interval: Duration::from_millis(DEFAULT_FLUSH_INTERVAL_MS),
workers: DEFAULT_WORKERS,
full_policy: RequestCandidateQueueFullPolicy::Drop,
}
}
}
impl RequestCandidateQueueConfig {
pub(crate) fn from_env() -> Self {
let mut config = Self::default();
config.mode = match env_string(MODE_ENV).as_deref() {
Some("async") | Some("queued") | Some("queue") => RequestCandidateWriteMode::Async,
_ => RequestCandidateWriteMode::Sync,
};
config.capacity = env_usize(QUEUE_CAPACITY_ENV, DEFAULT_QUEUE_CAPACITY).max(1);
config.batch_size = env_usize(BATCH_SIZE_ENV, DEFAULT_BATCH_SIZE).max(1);
config.flush_interval =
Duration::from_millis(env_u64(FLUSH_INTERVAL_MS_ENV, DEFAULT_FLUSH_INTERVAL_MS).max(1));
config.workers = env_usize(WORKERS_ENV, DEFAULT_WORKERS).clamp(1, 32);
config.full_policy = match env_string(QUEUE_FULL_ENV).as_deref() {
Some("sync") | Some("fallback_sync") | Some("fallback-sync") => {
RequestCandidateQueueFullPolicy::Sync
}
_ => RequestCandidateQueueFullPolicy::Drop,
};
config
}
pub(crate) fn async_enabled(&self) -> bool {
matches!(self.mode, RequestCandidateWriteMode::Async)
}
}
#[derive(Debug, Default)]
struct RequestCandidateQueueMetrics {
queued_current: AtomicUsize,
pending_current: AtomicUsize,
enqueued_total: AtomicU64,
dropped_total: AtomicU64,
flushed_total: AtomicU64,
flush_failed_total: AtomicU64,
flush_batches_total: AtomicU64,
flush_sql_ops_total: AtomicU64,
compacted_total: AtomicU64,
sync_fallback_total: AtomicU64,
}
#[derive(Clone)]
pub(crate) struct RequestCandidateQueueRuntime {
senders: Vec<mpsc::Sender<UpsertRequestCandidateRecord>>,
repository: Arc<dyn RequestCandidateWriteRepository>,
config: RequestCandidateQueueConfig,
metrics: Arc<RequestCandidateQueueMetrics>,
}
impl std::fmt::Debug for RequestCandidateQueueRuntime {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RequestCandidateQueueRuntime")
.field("config", &self.config)
.field(
"queued_current",
&self.metrics.queued_current.load(Ordering::Acquire),
)
.finish_non_exhaustive()
}
}
impl RequestCandidateQueueRuntime {
pub(crate) fn spawn(
repository: Arc<dyn RequestCandidateWriteRepository>,
mut config: RequestCandidateQueueConfig,
) -> Arc<Self> {
config.workers = config.workers.min(config.capacity).max(1);
let mut senders = Vec::with_capacity(config.workers);
let mut receivers = Vec::with_capacity(config.workers);
for worker_index in 0..config.workers {
let capacity = worker_queue_capacity(config.capacity, config.workers, worker_index);
let (sender, receiver) = mpsc::channel(capacity);
senders.push(sender);
receivers.push(receiver);
}
let runtime = Arc::new(Self {
senders,
repository,
config,
metrics: Arc::new(RequestCandidateQueueMetrics::default()),
});
runtime.spawn_workers(receivers);
runtime
}
pub(crate) async fn enqueue_or_fallback(
&self,
record: UpsertRequestCandidateRecord,
) -> Result<(), aether_data::DataLayerError> {
let worker_index = self.worker_index_for(&record);
let Some(sender) = self.senders.get(worker_index) else {
self.metrics
.sync_fallback_total
.fetch_add(1, Ordering::AcqRel);
return self.repository.upsert(record).await.map(|_| ());
};
self.metrics.queued_current.fetch_add(1, Ordering::AcqRel);
self.metrics.pending_current.fetch_add(1, Ordering::AcqRel);
match sender.try_send(record) {
Ok(()) => {
self.metrics.enqueued_total.fetch_add(1, Ordering::AcqRel);
Ok(())
}
Err(mpsc::error::TrySendError::Full(record)) => {
decrement_atomic_usize(&self.metrics.queued_current);
decrement_atomic_usize(&self.metrics.pending_current);
warn!(
event_name = "request_candidate_queue_full",
log_type = "event",
full_policy = ?self.config.full_policy,
worker_index,
queued = self.metrics.queued_current.load(Ordering::Acquire),
capacity = self.config.capacity,
"gateway request candidate async queue is full"
);
match self.config.full_policy {
RequestCandidateQueueFullPolicy::Drop => {
self.metrics.dropped_total.fetch_add(1, Ordering::AcqRel);
Ok(())
}
RequestCandidateQueueFullPolicy::Sync => {
self.metrics
.sync_fallback_total
.fetch_add(1, Ordering::AcqRel);
self.repository.upsert(record).await.map(|_| ())
}
}
}
Err(mpsc::error::TrySendError::Closed(record)) => {
decrement_atomic_usize(&self.metrics.queued_current);
decrement_atomic_usize(&self.metrics.pending_current);
self.metrics
.sync_fallback_total
.fetch_add(1, Ordering::AcqRel);
self.repository.upsert(record).await.map(|_| ())
}
}
}
pub(crate) fn metric_samples(&self) -> Vec<MetricSample> {
vec![
MetricSample::new(
"request_candidate_queue_depth",
"Current number of request candidate records waiting in the async persistence queue.",
MetricKind::Gauge,
self.metrics.queued_current.load(Ordering::Acquire) as u64,
),
MetricSample::new(
"request_candidate_queue_pending_depth",
"Current number of request candidate records accepted into the async persistence queue but not yet flushed.",
MetricKind::Gauge,
self.metrics.pending_current.load(Ordering::Acquire) as u64,
),
MetricSample::new(
"request_candidate_queue_capacity",
"Configured request candidate async persistence queue capacity.",
MetricKind::Gauge,
self.config.capacity as u64,
),
MetricSample::new(
"request_candidate_queue_enqueued_total",
"Total request candidate records accepted into the async persistence queue.",
MetricKind::Counter,
self.metrics.enqueued_total.load(Ordering::Acquire),
),
MetricSample::new(
"request_candidate_queue_dropped_total",
"Total request candidate records dropped because the async persistence queue was full.",
MetricKind::Counter,
self.metrics.dropped_total.load(Ordering::Acquire),
),
MetricSample::new(
"request_candidate_queue_flushed_total",
"Total request candidate records flushed by async persistence workers.",
MetricKind::Counter,
self.metrics.flushed_total.load(Ordering::Acquire),
),
MetricSample::new(
"request_candidate_queue_flush_failed_total",
"Total request candidate records that failed during async persistence flush.",
MetricKind::Counter,
self.metrics.flush_failed_total.load(Ordering::Acquire),
),
MetricSample::new(
"request_candidate_queue_flush_batches_total",
"Total async request candidate persistence flush batches.",
MetricKind::Counter,
self.metrics.flush_batches_total.load(Ordering::Acquire),
),
MetricSample::new(
"request_candidate_queue_flush_sql_ops_total",
"Total repository upsert operations issued by async request candidate persistence workers after compaction.",
MetricKind::Counter,
self.metrics.flush_sql_ops_total.load(Ordering::Acquire),
),
MetricSample::new(
"request_candidate_queue_compacted_total",
"Total request candidate records compacted before async persistence because a later queued record covered the same slot and status.",
MetricKind::Counter,
self.metrics.compacted_total.load(Ordering::Acquire),
),
MetricSample::new(
"request_candidate_queue_sync_fallback_total",
"Total request candidate records synchronously persisted after async queue fallback.",
MetricKind::Counter,
self.metrics.sync_fallback_total.load(Ordering::Acquire),
),
]
}
fn spawn_workers(
self: &Arc<Self>,
receivers: Vec<mpsc::Receiver<UpsertRequestCandidateRecord>>,
) {
for (worker_index, receiver) in receivers.into_iter().enumerate() {
let repository = Arc::clone(&self.repository);
let config = self.config.clone();
let metrics = Arc::clone(&self.metrics);
tokio::spawn(async move {
run_worker(repository, config, metrics, worker_index, receiver).await;
});
}
}
fn worker_index_for(&self, record: &UpsertRequestCandidateRecord) -> usize {
let worker_count = self.senders.len();
if worker_count <= 1 {
return 0;
}
(request_candidate_slot_hash(record) % worker_count as u64) as usize
}
}
async fn run_worker(
repository: Arc<dyn RequestCandidateWriteRepository>,
config: RequestCandidateQueueConfig,
metrics: Arc<RequestCandidateQueueMetrics>,
worker_index: usize,
mut receiver: mpsc::Receiver<UpsertRequestCandidateRecord>,
) {
let mut ticker = interval(config.flush_interval);
ticker.set_missed_tick_behavior(MissedTickBehavior::Delay);
let mut batch = Vec::with_capacity(config.batch_size);
loop {
tokio::select! {
_ = ticker.tick() => {
if !batch.is_empty() {
flush_batch(&repository, &metrics, worker_index, &mut batch).await;
}
}
received = receiver.recv() => {
match received {
Some(record) => {
decrement_atomic_usize(&metrics.queued_current);
batch.push(record);
if batch.len() >= config.batch_size {
flush_batch(&repository, &metrics, worker_index, &mut batch).await;
}
}
None => {
if !batch.is_empty() {
flush_batch(&repository, &metrics, worker_index, &mut batch).await;
}
break;
}
}
}
}
}
}
async fn flush_batch(
repository: &Arc<dyn RequestCandidateWriteRepository>,
metrics: &RequestCandidateQueueMetrics,
worker_index: usize,
batch: &mut Vec<UpsertRequestCandidateRecord>,
) {
let records = std::mem::take(batch);
if records.is_empty() {
return;
}
let source_count = records.len();
let records = compact_same_status_records(records);
let compacted = source_count.saturating_sub(records.len());
if compacted > 0 {
metrics
.compacted_total
.fetch_add(compacted as u64, Ordering::AcqRel);
}
metrics.flush_batches_total.fetch_add(1, Ordering::AcqRel);
let mut failed = 0_u64;
for record in records {
metrics.flush_sql_ops_total.fetch_add(1, Ordering::AcqRel);
if let Err(err) = repository.upsert(record.record).await {
failed = failed.saturating_add(record.source_count as u64);
decrement_atomic_usize_by(&metrics.pending_current, record.source_count);
warn!(
event_name = "request_candidate_async_flush_failed",
log_type = "event",
worker_index,
error = ?err,
"gateway failed to asynchronously persist request candidate"
);
} else {
metrics
.flushed_total
.fetch_add(record.source_count as u64, Ordering::AcqRel);
decrement_atomic_usize_by(&metrics.pending_current, record.source_count);
}
}
if failed > 0 {
metrics
.flush_failed_total
.fetch_add(failed, Ordering::AcqRel);
}
debug!(
event_name = "request_candidate_async_flush_completed",
log_type = "event",
worker_index,
failed,
"gateway completed request candidate async flush batch"
);
}
#[derive(Debug)]
struct CompactedRequestCandidateRecord {
record: UpsertRequestCandidateRecord,
source_count: usize,
}
fn compact_same_status_records(
records: Vec<UpsertRequestCandidateRecord>,
) -> Vec<CompactedRequestCandidateRecord> {
let mut latest_slot_status = HashMap::<(String, u32, u32), (u8, usize)>::new();
let mut compacted = Vec::<CompactedRequestCandidateRecord>::with_capacity(records.len());
for record in records {
let slot = (
record.request_id.clone(),
record.candidate_index,
record.retry_index,
);
let status = request_candidate_status_discriminant(record.status);
match latest_slot_status.get(&slot).copied() {
Some((latest_status, index)) if latest_status == status => {
merge_request_candidate_record(&mut compacted[index].record, record);
compacted[index].source_count = compacted[index].source_count.saturating_add(1);
}
_ => {
latest_slot_status.insert(slot, (status, compacted.len()));
compacted.push(CompactedRequestCandidateRecord {
record,
source_count: 1,
});
}
}
}
compacted
}
fn request_candidate_slot_hash(record: &UpsertRequestCandidateRecord) -> u64 {
const FNV_OFFSET: u64 = 14_695_981_039_346_656_037;
const FNV_PRIME: u64 = 1_099_511_628_211;
let mut hash = FNV_OFFSET;
for byte in record.request_id.as_bytes() {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(FNV_PRIME);
}
for byte in record.candidate_index.to_le_bytes() {
hash ^= u64::from(byte);
hash = hash.wrapping_mul(FNV_PRIME);
}
for byte in record.retry_index.to_le_bytes() {
hash ^= u64::from(byte);
hash = hash.wrapping_mul(FNV_PRIME);
}
hash
}
fn worker_queue_capacity(total_capacity: usize, workers: usize, worker_index: usize) -> usize {
let workers = workers.max(1);
let base = total_capacity / workers;
let remainder = total_capacity % workers;
(base + usize::from(worker_index < remainder)).max(1)
}
fn merge_request_candidate_record(
target: &mut UpsertRequestCandidateRecord,
incoming: UpsertRequestCandidateRecord,
) {
if !incoming.id.trim().is_empty() {
target.id = incoming.id;
}
target.request_id = incoming.request_id;
target.candidate_index = incoming.candidate_index;
target.retry_index = incoming.retry_index;
target.status = incoming.status;
take_if_some(&mut target.user_id, incoming.user_id);
take_if_some(&mut target.api_key_id, incoming.api_key_id);
take_if_some(&mut target.username, incoming.username);
take_if_some(&mut target.api_key_name, incoming.api_key_name);
take_if_some(&mut target.provider_id, incoming.provider_id);
take_if_some(&mut target.endpoint_id, incoming.endpoint_id);
take_if_some(&mut target.key_id, incoming.key_id);
take_if_some(&mut target.skip_reason, incoming.skip_reason);
take_if_some(&mut target.is_cached, incoming.is_cached);
take_if_some(&mut target.status_code, incoming.status_code);
take_if_some(&mut target.error_type, incoming.error_type);
take_if_some(&mut target.error_message, incoming.error_message);
take_if_some(&mut target.latency_ms, incoming.latency_ms);
take_if_some(
&mut target.concurrent_requests,
incoming.concurrent_requests,
);
merge_json_value(&mut target.extra_data, incoming.extra_data);
merge_json_value(
&mut target.required_capabilities,
incoming.required_capabilities,
);
if target.created_at_unix_ms.is_none() {
target.created_at_unix_ms = incoming.created_at_unix_ms;
}
take_if_some(&mut target.started_at_unix_ms, incoming.started_at_unix_ms);
take_if_some(
&mut target.finished_at_unix_ms,
incoming.finished_at_unix_ms,
);
}
fn request_candidate_status_discriminant(status: RequestCandidateStatus) -> u8 {
match status {
RequestCandidateStatus::Available => 0,
RequestCandidateStatus::Unused => 1,
RequestCandidateStatus::Pending => 2,
RequestCandidateStatus::Streaming => 3,
RequestCandidateStatus::Success => 4,
RequestCandidateStatus::Failed => 5,
RequestCandidateStatus::Cancelled => 6,
RequestCandidateStatus::Skipped => 7,
}
}
fn take_if_some<T>(target: &mut Option<T>, incoming: Option<T>) {
if incoming.is_some() {
*target = incoming;
}
}
fn merge_json_value(target: &mut Option<serde_json::Value>, incoming: Option<serde_json::Value>) {
match (target.as_mut(), incoming) {
(Some(serde_json::Value::Object(target)), Some(serde_json::Value::Object(incoming))) => {
target.extend(incoming);
}
(_, Some(incoming)) => {
*target = Some(incoming);
}
(_, None) => {}
}
}
fn env_string(key: &str) -> Option<String> {
std::env::var(key)
.ok()
.map(|value| value.trim().to_ascii_lowercase())
.filter(|value| !value.is_empty())
}
fn env_usize(key: &str, default: usize) -> usize {
std::env::var(key)
.ok()
.and_then(|value| value.trim().parse::<usize>().ok())
.unwrap_or(default)
}
fn env_u64(key: &str, default: u64) -> u64 {
std::env::var(key)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.unwrap_or(default)
}
fn decrement_atomic_usize(value: &AtomicUsize) {
decrement_atomic_usize_by(value, 1);
}
fn decrement_atomic_usize_by(value: &AtomicUsize, amount: usize) {
let _ = value.fetch_update(Ordering::AcqRel, Ordering::Acquire, |current| {
Some(current.saturating_sub(amount))
});
}
#[cfg(test)]
mod tests {
use super::{
compact_same_status_records, RequestCandidateQueueConfig, RequestCandidateQueueRuntime,
};
use aether_data::repository::candidates::InMemoryRequestCandidateRepository;
use aether_data::DataLayerError;
use aether_data_contracts::repository::candidates::{
RequestCandidateReadRepository, RequestCandidateStatus, RequestCandidateWriteRepository,
StoredRequestCandidate, UpsertRequestCandidateRecord,
};
use std::sync::atomic::Ordering;
use std::sync::Arc;
use std::time::Duration;
#[derive(Default)]
struct DelayedPendingRequestCandidateRepository {
inner: InMemoryRequestCandidateRepository,
}
#[async_trait::async_trait]
impl RequestCandidateWriteRepository for DelayedPendingRequestCandidateRepository {
async fn upsert(
&self,
candidate: UpsertRequestCandidateRecord,
) -> Result<StoredRequestCandidate, DataLayerError> {
if candidate.status == RequestCandidateStatus::Pending {
tokio::time::sleep(Duration::from_millis(40)).await;
}
self.inner.upsert(candidate).await
}
async fn delete_created_before(
&self,
created_before_unix_secs: u64,
limit: usize,
) -> Result<usize, DataLayerError> {
self.inner
.delete_created_before(created_before_unix_secs, limit)
.await
}
}
fn record(
request_id: &str,
candidate_index: u32,
retry_index: u32,
status: RequestCandidateStatus,
) -> UpsertRequestCandidateRecord {
UpsertRequestCandidateRecord {
id: format!("{request_id}-{candidate_index}-{retry_index}-{status:?}"),
request_id: request_id.to_string(),
user_id: None,
api_key_id: None,
username: None,
api_key_name: None,
candidate_index,
retry_index,
provider_id: None,
endpoint_id: None,
key_id: None,
status,
skip_reason: None,
is_cached: None,
status_code: None,
error_type: None,
error_message: None,
latency_ms: None,
concurrent_requests: None,
extra_data: None,
required_capabilities: None,
created_at_unix_ms: None,
started_at_unix_ms: None,
finished_at_unix_ms: None,
}
}
#[test]
fn compact_merges_same_slot_and_status_without_dropping_state_transitions() {
let mut first_success = record("req", 0, 0, RequestCandidateStatus::Success);
first_success.provider_id = Some("provider-a".to_string());
first_success.extra_data = Some(serde_json::json!({"first": true}));
let mut second_success = record("req", 0, 0, RequestCandidateStatus::Success);
second_success.latency_ms = Some(123);
second_success.extra_data = Some(serde_json::json!({"second": true}));
let compacted = compact_same_status_records(vec![
record("req", 0, 0, RequestCandidateStatus::Pending),
first_success,
record("req", 0, 1, RequestCandidateStatus::Failed),
second_success,
]);
assert_eq!(compacted.len(), 3);
assert_eq!(compacted[0].record.status, RequestCandidateStatus::Pending);
assert_eq!(compacted[0].source_count, 1);
assert_eq!(compacted[1].record.status, RequestCandidateStatus::Success);
assert_eq!(compacted[1].source_count, 2);
assert_eq!(
compacted[1].record.provider_id.as_deref(),
Some("provider-a")
);
assert_eq!(compacted[1].record.latency_ms, Some(123));
assert_eq!(
compacted[1].record.extra_data,
Some(serde_json::json!({"first": true, "second": true}))
);
assert_eq!(compacted[2].record.status, RequestCandidateStatus::Failed);
assert_eq!(compacted[2].source_count, 1);
}
#[test]
fn compact_preserves_same_slot_status_order_across_transitions() {
let compacted = compact_same_status_records(vec![
record("req", 0, 0, RequestCandidateStatus::Success),
record("req", 0, 0, RequestCandidateStatus::Failed),
record("req", 0, 0, RequestCandidateStatus::Success),
]);
assert_eq!(compacted.len(), 3);
assert_eq!(compacted[0].record.status, RequestCandidateStatus::Success);
assert_eq!(compacted[1].record.status, RequestCandidateStatus::Failed);
assert_eq!(compacted[2].record.status, RequestCandidateStatus::Success);
}
#[tokio::test]
async fn async_queue_flushes_enqueued_records() {
let repository = Arc::new(InMemoryRequestCandidateRepository::default());
let runtime = RequestCandidateQueueRuntime::spawn(
repository.clone(),
RequestCandidateQueueConfig {
mode: super::RequestCandidateWriteMode::Async,
capacity: 16,
batch_size: 2,
flush_interval: Duration::from_millis(10),
workers: 1,
full_policy: super::RequestCandidateQueueFullPolicy::Drop,
},
);
runtime
.enqueue_or_fallback(record("req", 0, 0, RequestCandidateStatus::Success))
.await
.unwrap();
for _ in 0..50 {
let rows = repository.list_by_request_id("req").await.unwrap();
if rows.len() == 1 {
assert_eq!(rows[0].status, RequestCandidateStatus::Success);
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
panic!("async request candidate queue did not flush record in time");
}
#[tokio::test]
async fn async_queue_preserves_same_slot_order_with_multiple_workers() {
let repository = Arc::new(DelayedPendingRequestCandidateRepository::default());
let runtime = RequestCandidateQueueRuntime::spawn(
repository.clone(),
RequestCandidateQueueConfig {
mode: super::RequestCandidateWriteMode::Async,
capacity: 16,
batch_size: 1,
flush_interval: Duration::from_millis(100),
workers: 2,
full_policy: super::RequestCandidateQueueFullPolicy::Drop,
},
);
runtime
.enqueue_or_fallback(record("req-order", 0, 0, RequestCandidateStatus::Pending))
.await
.unwrap();
runtime
.enqueue_or_fallback(record("req-order", 0, 0, RequestCandidateStatus::Success))
.await
.unwrap();
for _ in 0..50 {
if runtime.metrics.pending_current.load(Ordering::Acquire) == 0 {
let rows = repository
.inner
.list_by_request_id("req-order")
.await
.unwrap();
assert_eq!(rows.len(), 1);
assert_eq!(rows[0].status, RequestCandidateStatus::Success);
return;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
panic!("async request candidate queue did not finish ordered same-slot writes in time");
}
}
apps/aether-gateway/src/request_candidate_runtime.rs
+43 -5
View File
@@ -15,6 +15,7 @@ use aether_usage_runtime::build_locally_actionable_report_context_from_request_c
use async_trait::async_trait;
use serde_json::Value;
use std::sync::OnceLock;
use std::time::Duration;
use tracing::{debug, info, warn};
use uuid::Uuid;
@@ -23,6 +24,9 @@ use crate::log_ids::short_request_id;
use crate::GatewayError;
const REQUEST_CANDIDATE_PERSISTENCE_ENV: &str = "AETHER_GATEWAY_REQUEST_CANDIDATE_PERSISTENCE";
const REQUEST_CANDIDATE_SEED_WRITE_TIMEOUT_ENV: &str =
"AETHER_GATEWAY_REQUEST_CANDIDATE_SEED_WRITE_TIMEOUT_MS";
const DEFAULT_REQUEST_CANDIDATE_SEED_WRITE_TIMEOUT_MS: u64 = 10;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum RequestCandidatePersistenceMode {
@@ -67,6 +71,17 @@ fn should_persist_request_candidate_status(status: RequestCandidateStatus) -> bo
}
}
fn request_candidate_seed_write_timeout() -> Duration {
static TIMEOUT: OnceLock<Duration> = OnceLock::new();
*TIMEOUT.get_or_init(|| {
let millis = std::env::var(REQUEST_CANDIDATE_SEED_WRITE_TIMEOUT_ENV)
.ok()
.and_then(|value| value.trim().parse::<u64>().ok())
.unwrap_or(DEFAULT_REQUEST_CANDIDATE_SEED_WRITE_TIMEOUT_MS);
Duration::from_millis(millis)
})
}
#[derive(Debug, Clone)]
pub(crate) struct LocalRequestCandidateStatusSnapshot {
candidate_id: String,
@@ -573,8 +588,15 @@ pub(crate) async fn ensure_execution_request_candidate_slot(
return;
}
let candidate_id = match state.upsert_request_candidate(seed.upsert_record).await {
Ok(Some(stored)) => {
let seed_upsert_record = seed.upsert_record;
let generated_candidate_id = generated_candidate_id.clone();
let candidate_id = match tokio::time::timeout(
request_candidate_seed_write_timeout(),
state.upsert_request_candidate(seed_upsert_record),
)
.await
{
Ok(Ok(Some(stored))) => {
info!(
event_name = "request_candidate_slot_seeded",
log_type = "event",
@@ -588,7 +610,7 @@ pub(crate) async fn ensure_execution_request_candidate_slot(
);
stored.id
}
Ok(None) => {
Ok(Ok(None)) => {
warn!(
event_name = "request_candidate_writer_unavailable",
log_type = "event",
@@ -602,7 +624,7 @@ pub(crate) async fn ensure_execution_request_candidate_slot(
);
generated_candidate_id
}
Err(err) => {
Ok(Err(err)) => {
warn!(
event_name = "request_candidate_slot_seed_failed",
log_type = "event",
@@ -610,7 +632,23 @@ pub(crate) async fn ensure_execution_request_candidate_slot(
error = ?err,
"gateway failed to seed execution request candidate slot"
);
return;
generated_candidate_id
}
Err(_) => {
let timeout_ms = request_candidate_seed_write_timeout().as_millis() as u64;
warn!(
event_name = "request_candidate_slot_seed_timed_out",
log_type = "event",
request_id = %request_id,
candidate_id = %generated_candidate_id,
provider_id = %plan.provider_id,
endpoint_id = %plan.endpoint_id,
key_id = %plan.key_id,
source = "seed",
timeout_ms,
"gateway skipped blocking request candidate seed after timeout"
);
generated_candidate_id
}
};
apps/aether-gateway/src/state/app.rs
+2
View File
@@ -18,6 +18,7 @@ use super::super::cache::{
use super::super::data::GatewayDataState;
use super::super::fallback_metrics;
use super::super::rate_limit::FrontdoorUserRpmLimiter;
use super::super::request_candidate_queue::RequestCandidateQueueRuntime;
use super::super::{provider_transport, usage};
use super::{
AdminBillingCollectorRecord, AdminBillingRuleRecord, AdminPaymentCallbackRecord,
@@ -135,6 +136,7 @@ pub struct AppState {
pub(crate) dashboard_response_cache: Arc<DashboardResponseCache>,
pub(crate) system_config_cache: Arc<SystemConfigCache>,
pub(crate) fallback_metrics: Arc<fallback_metrics::GatewayFallbackMetrics>,
pub(crate) request_candidate_queue: Option<Arc<RequestCandidateQueueRuntime>>,
pub(crate) frontdoor_cors: Option<Arc<FrontdoorCorsConfig>>,
pub(crate) frontdoor_user_rpm: Arc<FrontdoorUserRpmLimiter>,
pub(crate) tunnel: crate::tunnel::EmbeddedTunnelState,
apps/aether-gateway/src/state/core.rs
+29
View File
@@ -20,6 +20,7 @@ use aether_runtime_state::{
RuntimeSemaphoreSnapshot, RuntimeState,
};
use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS;
use tracing::warn;
use super::{
AppState, FrontdoorCorsConfig, FrontdoorRuntimeGuardConfig, LocalExecutionRuntimeMissDiagnostic,
@@ -38,6 +39,9 @@ use super::super::fallback_metrics;
use super::super::fallback_metrics::{GatewayFallbackMetricKind, GatewayFallbackReason};
use super::super::model_fetch::spawn_model_fetch_worker;
use super::super::rate_limit::{FrontdoorUserRpmConfig, FrontdoorUserRpmLimiter};
use super::super::request_candidate_queue::{
RequestCandidateQueueConfig, RequestCandidateQueueRuntime,
};
use super::super::router::RequestAdmissionError;
use super::super::{control::GatewayControlDecision, error::GatewayError};
use super::super::{provider_transport, usage};
@@ -180,6 +184,7 @@ impl AppState {
self.runtime_state.clone(),
);
self.data = data;
self.configure_request_candidate_queue_from_env();
}
pub fn force_close_all_tunnel_proxies(&self) -> usize {
@@ -251,6 +256,7 @@ impl AppState {
dashboard_response_cache: Arc::new(DashboardResponseCache::default()),
system_config_cache: Arc::new(SystemConfigCache::default()),
fallback_metrics: Arc::new(fallback_metrics::GatewayFallbackMetrics::default()),
request_candidate_queue: None,
frontdoor_cors: None,
frontdoor_user_rpm: Arc::new(FrontdoorUserRpmLimiter::new(
FrontdoorUserRpmConfig::default(),
@@ -423,6 +429,26 @@ impl AppState {
self
}
fn configure_request_candidate_queue_from_env(&mut self) {
let config = RequestCandidateQueueConfig::from_env();
self.request_candidate_queue = if config.async_enabled() {
if tokio::runtime::Handle::try_current().is_err() {
warn!(
event_name = "request_candidate_async_queue_unavailable",
log_type = "ops",
"request candidate async queue requested outside a Tokio runtime; falling back to sync persistence"
);
None
} else {
self.data
.request_candidate_writer()
.map(|writer| RequestCandidateQueueRuntime::spawn(writer, config))
}
} else {
None
};
}
pub fn with_distributed_request_concurrency_gate(mut self, gate: RuntimeSemaphore) -> Self {
self.distributed_request_gate = Some(Arc::new(gate));
self
@@ -918,6 +944,9 @@ impl AppState {
if let Some(summary) = self.data.database_pool_summary() {
samples.extend(database_pool_metric_samples(&summary));
}
if let Some(queue) = self.request_candidate_queue.as_ref() {
samples.extend(queue.metric_samples());
}
samples.extend(self.tunnel.metric_samples());
samples.extend(self.fallback_metrics.metric_samples());
samples
apps/aether-gateway/src/state/runtime/candidate_queries.rs
+58
View File
@@ -113,9 +113,67 @@ impl AppState {
&self,
candidate: candidates::UpsertRequestCandidateRecord,
) -> Result<Option<candidates::StoredRequestCandidate>, GatewayError> {
if let Some(queue) = self.request_candidate_queue.as_ref() {
let stored = stored_request_candidate_from_upsert(&candidate)?;
queue
.enqueue_or_fallback(candidate)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))?;
return Ok(Some(stored));
}
self.data
.upsert_request_candidate(candidate)
.await
.map_err(|err| GatewayError::Internal(err.to_string()))
}
}
fn stored_request_candidate_from_upsert(
candidate: &candidates::UpsertRequestCandidateRecord,
) -> Result<candidates::StoredRequestCandidate, GatewayError> {
candidate
.validate()
.map_err(|err| GatewayError::Internal(err.to_string()))?;
candidates::StoredRequestCandidate::new(
candidate.id.clone(),
candidate.request_id.clone(),
candidate.user_id.clone(),
candidate.api_key_id.clone(),
candidate.username.clone(),
candidate.api_key_name.clone(),
candidate.candidate_index.try_into().unwrap_or(i32::MAX),
candidate.retry_index.try_into().unwrap_or(i32::MAX),
candidate.provider_id.clone(),
candidate.endpoint_id.clone(),
candidate.key_id.clone(),
candidate.status,
candidate.skip_reason.clone(),
candidate.is_cached.unwrap_or(false),
candidate.status_code.map(i32::from),
candidate.error_type.clone(),
candidate.error_message.clone(),
candidate
.latency_ms
.map(|value| i32::try_from(value).unwrap_or(i32::MAX)),
candidate
.concurrent_requests
.map(|value| i32::try_from(value).unwrap_or(i32::MAX)),
candidate.extra_data.clone(),
candidate.required_capabilities.clone(),
candidate
.created_at_unix_ms
.or(candidate.started_at_unix_ms)
.or(candidate.finished_at_unix_ms)
.unwrap_or_else(crate::clock::current_unix_ms)
.try_into()
.unwrap_or(i64::MAX),
candidate
.started_at_unix_ms
.map(|value| value.try_into().unwrap_or(i64::MAX)),
candidate
.finished_at_unix_ms
.map(|value| value.try_into().unwrap_or(i64::MAX)),
)
.map_err(|err| GatewayError::Internal(err.to_string()))
}