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

fix: harden frontdoor and usage ingestion

Browse Source
This commit is contained in:
fawney19
2026-05-22 23:57:38 +08:00
parent 6ef6cbade2
commit d18b13a91a
25 changed files with 1363 additions and 252 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/data/state/candidate_cache.rs
+327 -28
View File
@@ -1,4 +1,4 @@
use std::collections::HashSet;
use std::collections::HashMap;
use std::future::Future;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};
@@ -13,15 +13,26 @@ use aether_data_contracts::repository::candidate_selection::{
};
use async_trait::async_trait;
use tokio::sync::Notify;
use tokio::time::timeout;
use tracing::warn;
const CANDIDATE_SELECTION_CACHE_TTL: Duration = Duration::from_secs(5);
const CANDIDATE_SELECTION_CACHE_MAX_ENTRIES: usize = 4096;
#[cfg(not(test))]
const CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT: Duration = Duration::from_secs(10);
#[cfg(test)]
const CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT: Duration = Duration::from_millis(50);
#[cfg(not(test))]
const CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_secs(10);
#[cfg(test)]
const CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT: Duration = Duration::from_millis(50);
pub(super) struct CachedMinimalCandidateSelectionReadRepository {
inner: Arc<dyn MinimalCandidateSelectionReadRepository>,
entries: ExpiringMap<CandidateSelectionCacheKey, Vec<StoredMinimalCandidateSelectionRow>>,
inflight: Mutex<HashSet<CandidateSelectionCacheKey>>,
inflight: Mutex<HashMap<CandidateSelectionCacheKey, u64>>,
inflight_notify: Notify,
next_inflight_token: AtomicU64,
epoch: AtomicU64,
}
@@ -30,8 +41,9 @@ impl CachedMinimalCandidateSelectionReadRepository {
Self {
inner,
entries: ExpiringMap::new(),
inflight: Mutex::new(HashSet::new()),
inflight: Mutex::new(HashMap::new()),
inflight_notify: Notify::new(),
next_inflight_token: AtomicU64::new(1),
epoch: AtomicU64::new(0),
}
}
@@ -52,67 +64,169 @@ impl CachedMinimalCandidateSelectionReadRepository {
loop {
let notified = self.inflight_notify.notified();
match self.register_inflight(&key) {
InflightRegistration::Bypass => return load().await,
InflightRegistration::Bypass => {
return load_candidate_selection_rows_with_timeout(&key, load()).await;
}
InflightRegistration::Follower => {
notified.await;
if timeout(CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT, notified)
.await
.is_err()
{
self.expire_inflight(&key);
}
if let Some(rows) = self.entries.get_fresh(&key, CANDIDATE_SELECTION_CACHE_TTL)
{
return Ok(rows);
}
continue;
}
InflightRegistration::Leader => {}
}
let load_epoch = self.epoch.load(Ordering::Acquire);
let result = load().await;
if let Ok(rows) = &result {
if load_epoch == self.epoch.load(Ordering::Acquire) {
self.entries.insert(
key.clone(),
rows.clone(),
CANDIDATE_SELECTION_CACHE_TTL,
CANDIDATE_SELECTION_CACHE_MAX_ENTRIES,
);
InflightRegistration::Leader(token) => {
let mut guard = InflightGuard::new(self, key.clone(), token);
let load_epoch = self.epoch.load(Ordering::Acquire);
let result = load_candidate_selection_rows_with_timeout(&key, load()).await;
if let Ok(rows) = &result {
if load_epoch == self.epoch.load(Ordering::Acquire) {
self.entries.insert(
key.clone(),
rows.clone(),
CANDIDATE_SELECTION_CACHE_TTL,
CANDIDATE_SELECTION_CACHE_MAX_ENTRIES,
);
}
}
guard.finish();
return result;
}
}
self.finish_inflight(&key);
return result;
}
}
fn register_inflight(&self, key: &CandidateSelectionCacheKey) -> InflightRegistration {
match self.inflight.lock() {
Ok(mut inflight) => {
if inflight.insert(key.clone()) {
InflightRegistration::Leader
} else {
InflightRegistration::Follower
if inflight.contains_key(key) {
return InflightRegistration::Follower;
}
let token = self.next_inflight_token.fetch_add(1, Ordering::AcqRel);
inflight.insert(key.clone(), token);
InflightRegistration::Leader(token)
}
Err(_) => InflightRegistration::Bypass,
}
}
fn finish_inflight(&self, key: &CandidateSelectionCacheKey) {
fn finish_inflight(&self, key: &CandidateSelectionCacheKey, token: u64) {
let mut removed = false;
if let Ok(mut inflight) = self.inflight.lock() {
inflight.remove(key);
if inflight.get(key).copied() == Some(token) {
inflight.remove(key);
removed = true;
}
}
if removed {
self.inflight_notify.notify_waiters();
}
}
fn expire_inflight(&self, key: &CandidateSelectionCacheKey) {
let mut removed = false;
if let Ok(mut inflight) = self.inflight.lock() {
removed = inflight.remove(key).is_some();
}
if removed {
warn!(
event_name = "candidate_selection_cache_inflight_expired",
log_type = "ops",
cache_key = ?key,
wait_timeout_ms = CANDIDATE_SELECTION_CACHE_INFLIGHT_WAIT_TIMEOUT.as_millis() as u64,
"gateway candidate selection cache expired stale inflight load"
);
self.inflight_notify.notify_waiters();
}
self.inflight_notify.notify_waiters();
}
fn clear(&self) {
self.epoch.fetch_add(1, Ordering::AcqRel);
self.entries.clear();
let mut cleared_inflight = false;
if let Ok(mut inflight) = self.inflight.lock() {
cleared_inflight = !inflight.is_empty();
inflight.clear();
}
if cleared_inflight {
warn!(
event_name = "candidate_selection_cache_inflight_cleared",
log_type = "ops",
"gateway candidate selection cache cleared in-flight loads"
);
self.inflight_notify.notify_waiters();
}
}
}
enum InflightRegistration {
Leader,
Leader(u64),
Follower,
Bypass,
}
struct InflightGuard<'a> {
cache: &'a CachedMinimalCandidateSelectionReadRepository,
key: Option<CandidateSelectionCacheKey>,
token: u64,
}
impl<'a> InflightGuard<'a> {
fn new(
cache: &'a CachedMinimalCandidateSelectionReadRepository,
key: CandidateSelectionCacheKey,
token: u64,
) -> Self {
Self {
cache,
key: Some(key),
token,
}
}
fn finish(&mut self) {
if let Some(key) = self.key.take() {
self.cache.finish_inflight(&key, self.token);
}
}
}
impl Drop for InflightGuard<'_> {
fn drop(&mut self) {
self.finish();
}
}
async fn load_candidate_selection_rows_with_timeout<Fut>(
key: &CandidateSelectionCacheKey,
load: Fut,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>
where
Fut: Future<Output = Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError>>,
{
match timeout(CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT, load).await {
Ok(result) => result,
Err(_) => {
warn!(
event_name = "candidate_selection_cache_load_timeout",
log_type = "ops",
cache_key = ?key,
timeout_ms = CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT.as_millis() as u64,
"gateway candidate selection cache load timed out"
);
Err(DataLayerError::TimedOut(format!(
"candidate selection cache load exceeded {}ms for {key:?}",
CANDIDATE_SELECTION_CACHE_LOAD_TIMEOUT.as_millis()
)))
}
}
}
#[async_trait]
impl MinimalCandidateSelectionReadRepository for CachedMinimalCandidateSelectionReadRepository {
fn clear_local_cache(&self) {
@@ -286,6 +400,7 @@ fn normalize_api_format_key(api_format: &str) -> String {
#[cfg(test)]
mod tests {
use super::*;
use std::future::pending;
use std::sync::atomic::AtomicUsize;
struct StubCandidateSelectionRepository {
@@ -361,6 +476,77 @@ mod tests {
}
}
struct FirstLoadPendingThenFastRepository {
calls: AtomicUsize,
}
impl FirstLoadPendingThenFastRepository {
fn new() -> Self {
Self {
calls: AtomicUsize::new(0),
}
}
fn calls(&self) -> usize {
self.calls.load(Ordering::SeqCst)
}
async fn load(&self) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
let call = self.calls.fetch_add(1, Ordering::SeqCst);
if call == 0 {
pending::<()>().await;
}
Ok(Vec::new())
}
}
#[async_trait]
impl MinimalCandidateSelectionReadRepository for FirstLoadPendingThenFastRepository {
async fn list_for_exact_api_format(
&self,
_api_format: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_for_exact_api_format_and_global_model(
&self,
_api_format: &str,
_global_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_for_exact_api_format_and_requested_model(
&self,
_api_format: &str,
_requested_model_name: &str,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_for_exact_api_format_and_requested_model_page(
&self,
_query: &StoredRequestedModelCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_pool_key_rows_for_group(
&self,
_query: &StoredPoolKeyCandidateRowsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
async fn list_pool_key_rows_for_group_key_ids(
&self,
_query: &StoredPoolKeyCandidateRowsByKeyIdsQuery,
) -> Result<Vec<StoredMinimalCandidateSelectionRow>, DataLayerError> {
self.load().await
}
}
#[tokio::test]
async fn candidate_selection_cache_coalesces_concurrent_loads() {
let inner = Arc::new(StubCandidateSelectionRepository::new(
@@ -398,4 +584,117 @@ mod tests {
cache.list_for_exact_api_format("openai").await.unwrap();
assert_eq!(inner.calls(), 2);
}
#[tokio::test]
async fn candidate_selection_cache_releases_inflight_when_leader_is_cancelled() {
let inner = Arc::new(FirstLoadPendingThenFastRepository::new());
let cache = Arc::new(CachedMinimalCandidateSelectionReadRepository::new(
inner.clone(),
));
let leader_cache = cache.clone();
let leader = tokio::spawn(async move {
leader_cache
.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
)
.await
});
tokio::time::sleep(Duration::from_millis(10)).await;
leader.abort();
let _ = leader.await;
tokio::time::timeout(
Duration::from_millis(200),
cache.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
),
)
.await
.expect("cancelled leader must not leave a permanent inflight wait")
.unwrap();
assert_eq!(inner.calls(), 2);
}
#[tokio::test]
async fn candidate_selection_cache_times_out_and_clears_stuck_load() {
let inner = Arc::new(FirstLoadPendingThenFastRepository::new());
let cache = CachedMinimalCandidateSelectionReadRepository::new(inner.clone());
let err = cache
.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
)
.await
.unwrap_err();
assert!(matches!(err, DataLayerError::TimedOut(_)));
cache
.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
)
.await
.unwrap();
assert_eq!(inner.calls(), 2);
}
#[tokio::test]
async fn candidate_selection_cache_clear_releases_inflight_waiters() {
let inner = Arc::new(FirstLoadPendingThenFastRepository::new());
let cache = Arc::new(CachedMinimalCandidateSelectionReadRepository::new(
inner.clone(),
));
let leader_cache = cache.clone();
let leader = tokio::spawn(async move {
leader_cache
.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
)
.await
});
tokio::time::sleep(Duration::from_millis(10)).await;
cache.clear_local_cache();
tokio::time::timeout(
Duration::from_millis(200),
cache.list_for_exact_api_format_and_requested_model_page(
&StoredRequestedModelCandidateRowsQuery {
api_format: "openai:chat".to_string(),
requested_model_name: "gpt-5.5".to_string(),
offset: 0,
limit: 64,
},
),
)
.await
.expect("cache clear must release stale inflight waiters")
.unwrap();
assert_eq!(inner.calls(), 2);
leader.abort();
let _ = leader.await;
}
}