mirror of
https://github.com/fawney19/Aether.git
synced 2026-09-02 09:20:22 +08:00
fix: harden frontdoor and usage ingestion
This commit is contained in:
@@ -26,6 +26,8 @@ pub(crate) const EXECUTION_PATH_EXECUTION_RUNTIME_STREAM: &str = "execution_runt
|
|||||||
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_SYNC: &str = "control_execute_sync";
|
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_SYNC: &str = "control_execute_sync";
|
||||||
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_STREAM: &str = "control_execute_stream";
|
pub(crate) const EXECUTION_PATH_CONTROL_EXECUTE_STREAM: &str = "control_execute_stream";
|
||||||
pub(crate) const EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS: &str = "local_execution_runtime_miss";
|
pub(crate) const EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS: &str = "local_execution_runtime_miss";
|
||||||
|
pub(crate) const EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT: &str =
|
||||||
|
"local_execution_planning_timeout";
|
||||||
pub(crate) const EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED: &str =
|
pub(crate) const EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED: &str =
|
||||||
"local_api_key_concurrency_limited";
|
"local_api_key_concurrency_limited";
|
||||||
pub(crate) const API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS: u64 = 150;
|
pub(crate) const API_KEY_CONCURRENCY_WAIT_TIMEOUT_MS: u64 = 150;
|
||||||
|
|||||||
@@ -1,4 +1,4 @@
|
|||||||
use std::collections::HashSet;
|
use std::collections::HashMap;
|
||||||
use std::future::Future;
|
use std::future::Future;
|
||||||
use std::sync::atomic::{AtomicU64, Ordering};
|
use std::sync::atomic::{AtomicU64, Ordering};
|
||||||
use std::sync::{Arc, Mutex};
|
use std::sync::{Arc, Mutex};
|
||||||
@@ -13,15 +13,26 @@ use aether_data_contracts::repository::candidate_selection::{
|
|||||||
};
|
};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
use tokio::sync::Notify;
|
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_TTL: Duration = Duration::from_secs(5);
|
||||||
const CANDIDATE_SELECTION_CACHE_MAX_ENTRIES: usize = 4096;
|
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 {
|
pub(super) struct CachedMinimalCandidateSelectionReadRepository {
|
||||||
inner: Arc<dyn MinimalCandidateSelectionReadRepository>,
|
inner: Arc<dyn MinimalCandidateSelectionReadRepository>,
|
||||||
entries: ExpiringMap<CandidateSelectionCacheKey, Vec<StoredMinimalCandidateSelectionRow>>,
|
entries: ExpiringMap<CandidateSelectionCacheKey, Vec<StoredMinimalCandidateSelectionRow>>,
|
||||||
inflight: Mutex<HashSet<CandidateSelectionCacheKey>>,
|
inflight: Mutex<HashMap<CandidateSelectionCacheKey, u64>>,
|
||||||
inflight_notify: Notify,
|
inflight_notify: Notify,
|
||||||
|
next_inflight_token: AtomicU64,
|
||||||
epoch: AtomicU64,
|
epoch: AtomicU64,
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -30,8 +41,9 @@ impl CachedMinimalCandidateSelectionReadRepository {
|
|||||||
Self {
|
Self {
|
||||||
inner,
|
inner,
|
||||||
entries: ExpiringMap::new(),
|
entries: ExpiringMap::new(),
|
||||||
inflight: Mutex::new(HashSet::new()),
|
inflight: Mutex::new(HashMap::new()),
|
||||||
inflight_notify: Notify::new(),
|
inflight_notify: Notify::new(),
|
||||||
|
next_inflight_token: AtomicU64::new(1),
|
||||||
epoch: AtomicU64::new(0),
|
epoch: AtomicU64::new(0),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -52,67 +64,169 @@ impl CachedMinimalCandidateSelectionReadRepository {
|
|||||||
loop {
|
loop {
|
||||||
let notified = self.inflight_notify.notified();
|
let notified = self.inflight_notify.notified();
|
||||||
match self.register_inflight(&key) {
|
match self.register_inflight(&key) {
|
||||||
InflightRegistration::Bypass => return load().await,
|
InflightRegistration::Bypass => {
|
||||||
|
return load_candidate_selection_rows_with_timeout(&key, load()).await;
|
||||||
|
}
|
||||||
InflightRegistration::Follower => {
|
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)
|
if let Some(rows) = self.entries.get_fresh(&key, CANDIDATE_SELECTION_CACHE_TTL)
|
||||||
{
|
{
|
||||||
return Ok(rows);
|
return Ok(rows);
|
||||||
}
|
}
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
InflightRegistration::Leader => {}
|
InflightRegistration::Leader(token) => {
|
||||||
}
|
let mut guard = InflightGuard::new(self, key.clone(), token);
|
||||||
|
let load_epoch = self.epoch.load(Ordering::Acquire);
|
||||||
let load_epoch = self.epoch.load(Ordering::Acquire);
|
let result = load_candidate_selection_rows_with_timeout(&key, load()).await;
|
||||||
let result = load().await;
|
if let Ok(rows) = &result {
|
||||||
if let Ok(rows) = &result {
|
if load_epoch == self.epoch.load(Ordering::Acquire) {
|
||||||
if load_epoch == self.epoch.load(Ordering::Acquire) {
|
self.entries.insert(
|
||||||
self.entries.insert(
|
key.clone(),
|
||||||
key.clone(),
|
rows.clone(),
|
||||||
rows.clone(),
|
CANDIDATE_SELECTION_CACHE_TTL,
|
||||||
CANDIDATE_SELECTION_CACHE_TTL,
|
CANDIDATE_SELECTION_CACHE_MAX_ENTRIES,
|
||||||
CANDIDATE_SELECTION_CACHE_MAX_ENTRIES,
|
);
|
||||||
);
|
}
|
||||||
|
}
|
||||||
|
guard.finish();
|
||||||
|
return result;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
self.finish_inflight(&key);
|
|
||||||
return result;
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn register_inflight(&self, key: &CandidateSelectionCacheKey) -> InflightRegistration {
|
fn register_inflight(&self, key: &CandidateSelectionCacheKey) -> InflightRegistration {
|
||||||
match self.inflight.lock() {
|
match self.inflight.lock() {
|
||||||
Ok(mut inflight) => {
|
Ok(mut inflight) => {
|
||||||
if inflight.insert(key.clone()) {
|
if inflight.contains_key(key) {
|
||||||
InflightRegistration::Leader
|
return InflightRegistration::Follower;
|
||||||
} else {
|
|
||||||
InflightRegistration::Follower
|
|
||||||
}
|
}
|
||||||
|
let token = self.next_inflight_token.fetch_add(1, Ordering::AcqRel);
|
||||||
|
inflight.insert(key.clone(), token);
|
||||||
|
InflightRegistration::Leader(token)
|
||||||
}
|
}
|
||||||
Err(_) => InflightRegistration::Bypass,
|
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() {
|
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) {
|
fn clear(&self) {
|
||||||
self.epoch.fetch_add(1, Ordering::AcqRel);
|
self.epoch.fetch_add(1, Ordering::AcqRel);
|
||||||
self.entries.clear();
|
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 {
|
enum InflightRegistration {
|
||||||
Leader,
|
Leader(u64),
|
||||||
Follower,
|
Follower,
|
||||||
Bypass,
|
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]
|
#[async_trait]
|
||||||
impl MinimalCandidateSelectionReadRepository for CachedMinimalCandidateSelectionReadRepository {
|
impl MinimalCandidateSelectionReadRepository for CachedMinimalCandidateSelectionReadRepository {
|
||||||
fn clear_local_cache(&self) {
|
fn clear_local_cache(&self) {
|
||||||
@@ -286,6 +400,7 @@ fn normalize_api_format_key(api_format: &str) -> String {
|
|||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
|
use std::future::pending;
|
||||||
use std::sync::atomic::AtomicUsize;
|
use std::sync::atomic::AtomicUsize;
|
||||||
|
|
||||||
struct StubCandidateSelectionRepository {
|
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]
|
#[tokio::test]
|
||||||
async fn candidate_selection_cache_coalesces_concurrent_loads() {
|
async fn candidate_selection_cache_coalesces_concurrent_loads() {
|
||||||
let inner = Arc::new(StubCandidateSelectionRepository::new(
|
let inner = Arc::new(StubCandidateSelectionRepository::new(
|
||||||
@@ -398,4 +584,117 @@ mod tests {
|
|||||||
cache.list_for_exact_api_format("openai").await.unwrap();
|
cache.list_for_exact_api_format("openai").await.unwrap();
|
||||||
assert_eq!(inner.calls(), 2);
|
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;
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -11,12 +11,42 @@ use crate::insert_header_if_missing;
|
|||||||
|
|
||||||
#[derive(Debug)]
|
#[derive(Debug)]
|
||||||
pub(crate) enum GatewayError {
|
pub(crate) enum GatewayError {
|
||||||
UpstreamUnavailable { trace_id: String, message: String },
|
UpstreamUnavailable {
|
||||||
ControlUnavailable { trace_id: String, message: String },
|
trace_id: String,
|
||||||
Client { status: StatusCode, message: String },
|
message: String,
|
||||||
|
},
|
||||||
|
ControlUnavailable {
|
||||||
|
trace_id: String,
|
||||||
|
message: String,
|
||||||
|
},
|
||||||
|
LocalExecutionPlanningTimeout {
|
||||||
|
trace_id: String,
|
||||||
|
phase: &'static str,
|
||||||
|
timeout_ms: u64,
|
||||||
|
},
|
||||||
|
Client {
|
||||||
|
status: StatusCode,
|
||||||
|
message: String,
|
||||||
|
},
|
||||||
Internal(String),
|
Internal(String),
|
||||||
}
|
}
|
||||||
|
|
||||||
|
impl GatewayError {
|
||||||
|
pub(crate) fn into_message(self) -> String {
|
||||||
|
match self {
|
||||||
|
Self::UpstreamUnavailable { message, .. }
|
||||||
|
| Self::ControlUnavailable { message, .. }
|
||||||
|
| Self::Client { message, .. }
|
||||||
|
| Self::Internal(message) => message,
|
||||||
|
Self::LocalExecutionPlanningTimeout {
|
||||||
|
phase, timeout_ms, ..
|
||||||
|
} => {
|
||||||
|
format!("local execution planning timed out in {phase} after {timeout_ms}ms")
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
impl IntoResponse for GatewayError {
|
impl IntoResponse for GatewayError {
|
||||||
fn into_response(self) -> Response<Body> {
|
fn into_response(self) -> Response<Body> {
|
||||||
match self {
|
match self {
|
||||||
@@ -56,6 +86,33 @@ impl IntoResponse for GatewayError {
|
|||||||
);
|
);
|
||||||
response
|
response
|
||||||
}
|
}
|
||||||
|
Self::LocalExecutionPlanningTimeout {
|
||||||
|
trace_id,
|
||||||
|
phase,
|
||||||
|
timeout_ms,
|
||||||
|
} => {
|
||||||
|
warn!(
|
||||||
|
trace_id = %trace_id,
|
||||||
|
phase,
|
||||||
|
timeout_ms,
|
||||||
|
"gateway local execution planning timed out"
|
||||||
|
);
|
||||||
|
let body = Json(json!({
|
||||||
|
"error": {
|
||||||
|
"message": "gateway local execution planning timed out",
|
||||||
|
"trace_id": trace_id,
|
||||||
|
}
|
||||||
|
}));
|
||||||
|
let mut response = (StatusCode::GATEWAY_TIMEOUT, body).into_response();
|
||||||
|
let _ =
|
||||||
|
insert_header_if_missing(response.headers_mut(), TRACE_ID_HEADER, &trace_id);
|
||||||
|
let _ = insert_header_if_missing(
|
||||||
|
response.headers_mut(),
|
||||||
|
GATEWAY_HEADER,
|
||||||
|
"rust-phase3b",
|
||||||
|
);
|
||||||
|
response
|
||||||
|
}
|
||||||
Self::Client { status, message } => (
|
Self::Client { status, message } => (
|
||||||
status,
|
status,
|
||||||
Json(json!({
|
Json(json!({
|
||||||
|
|||||||
@@ -125,7 +125,16 @@ where
|
|||||||
decision,
|
decision,
|
||||||
plan_kind,
|
plan_kind,
|
||||||
};
|
};
|
||||||
run_dynamic_attempt_loop(&port, &mut source).await
|
run_dynamic_attempt_loop(
|
||||||
|
&port,
|
||||||
|
&mut source,
|
||||||
|
trace_id,
|
||||||
|
plan_kind,
|
||||||
|
state
|
||||||
|
.frontdoor_runtime_guards
|
||||||
|
.local_execution_planning_timeout,
|
||||||
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
.instrument(span)
|
.instrument(span)
|
||||||
.await
|
.await
|
||||||
@@ -281,7 +290,16 @@ where
|
|||||||
decision,
|
decision,
|
||||||
plan_kind,
|
plan_kind,
|
||||||
};
|
};
|
||||||
run_dynamic_attempt_loop(&port, &mut source).await
|
run_dynamic_attempt_loop(
|
||||||
|
&port,
|
||||||
|
&mut source,
|
||||||
|
trace_id,
|
||||||
|
plan_kind,
|
||||||
|
state
|
||||||
|
.frontdoor_runtime_guards
|
||||||
|
.local_execution_planning_timeout,
|
||||||
|
)
|
||||||
|
.await
|
||||||
}
|
}
|
||||||
.instrument(span)
|
.instrument(span)
|
||||||
.await
|
.await
|
||||||
@@ -290,6 +308,9 @@ where
|
|||||||
async fn run_dynamic_attempt_loop<Port, Source, Attempt>(
|
async fn run_dynamic_attempt_loop<Port, Source, Attempt>(
|
||||||
port: &Port,
|
port: &Port,
|
||||||
source: &mut Source,
|
source: &mut Source,
|
||||||
|
trace_id: &str,
|
||||||
|
plan_kind: &str,
|
||||||
|
planning_timeout: Duration,
|
||||||
) -> Result<LocalExecutionRequestOutcome, GatewayError>
|
) -> Result<LocalExecutionRequestOutcome, GatewayError>
|
||||||
where
|
where
|
||||||
Port: AiAttemptLoopPort<
|
Port: AiAttemptLoopPort<
|
||||||
@@ -303,7 +324,9 @@ where
|
|||||||
{
|
{
|
||||||
let mut last_attempted = None;
|
let mut last_attempted = None;
|
||||||
|
|
||||||
while let Some(attempt) = source.next_execution_attempt().await? {
|
while let Some(attempt) =
|
||||||
|
next_execution_attempt_with_timeout(source, trace_id, plan_kind, planning_timeout).await?
|
||||||
|
{
|
||||||
last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context()));
|
last_attempted = Some((attempt.execution_plan().clone(), attempt.report_context()));
|
||||||
if let Some(response) = port.execute_attempt(&attempt).await? {
|
if let Some(response) = port.execute_attempt(&attempt).await? {
|
||||||
let remaining = source.drain_execution_attempts().await?;
|
let remaining = source.drain_execution_attempts().await?;
|
||||||
@@ -322,6 +345,37 @@ where
|
|||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn next_execution_attempt_with_timeout<Source, Attempt>(
|
||||||
|
source: &mut Source,
|
||||||
|
trace_id: &str,
|
||||||
|
plan_kind: &str,
|
||||||
|
planning_timeout: Duration,
|
||||||
|
) -> Result<Option<Attempt>, GatewayError>
|
||||||
|
where
|
||||||
|
Source: LocalExecutionAttemptSource<Attempt>,
|
||||||
|
{
|
||||||
|
match timeout(planning_timeout, source.next_execution_attempt()).await {
|
||||||
|
Ok(result) => result,
|
||||||
|
Err(_) => {
|
||||||
|
let timeout_ms = planning_timeout.as_millis() as u64;
|
||||||
|
warn!(
|
||||||
|
event_name = "local_execution_candidate_planning_timeout",
|
||||||
|
log_type = "ops",
|
||||||
|
trace_id,
|
||||||
|
plan_kind,
|
||||||
|
timeout_ms,
|
||||||
|
phase = "next_execution_attempt",
|
||||||
|
"gateway timed out while planning the next local execution candidate"
|
||||||
|
);
|
||||||
|
Err(GatewayError::LocalExecutionPlanningTimeout {
|
||||||
|
trace_id: trace_id.to_string(),
|
||||||
|
phase: "next_execution_attempt",
|
||||||
|
timeout_ms,
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
struct StreamAttemptLoopPort<'a> {
|
struct StreamAttemptLoopPort<'a> {
|
||||||
state: &'a AppState,
|
state: &'a AppState,
|
||||||
trace_id: &'a str,
|
trace_id: &'a str,
|
||||||
@@ -632,6 +686,20 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
struct PendingAttemptSource;
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl LocalExecutionAttemptSource<()> for PendingAttemptSource {
|
||||||
|
async fn next_execution_attempt(&mut self) -> Result<Option<()>, GatewayError> {
|
||||||
|
std::future::pending::<()>().await;
|
||||||
|
Ok(None)
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn drain_execution_attempts(&mut self) -> Result<Vec<()>, GatewayError> {
|
||||||
|
Ok(Vec::new())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn test_plan(timeouts: Option<ExecutionTimeouts>) -> ExecutionPlan {
|
fn test_plan(timeouts: Option<ExecutionTimeouts>) -> ExecutionPlan {
|
||||||
ExecutionPlan {
|
ExecutionPlan {
|
||||||
request_id: "req_watchdog".to_string(),
|
request_id: "req_watchdog".to_string(),
|
||||||
@@ -656,6 +724,33 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn next_execution_attempt_times_out_instead_of_waiting_forever() {
|
||||||
|
let mut source = PendingAttemptSource;
|
||||||
|
|
||||||
|
let err = next_execution_attempt_with_timeout(
|
||||||
|
&mut source,
|
||||||
|
"trace-planning-timeout",
|
||||||
|
"openai_responses_sync",
|
||||||
|
Duration::from_millis(5),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect_err("pending candidate planning should time out");
|
||||||
|
|
||||||
|
match err {
|
||||||
|
GatewayError::LocalExecutionPlanningTimeout {
|
||||||
|
trace_id,
|
||||||
|
phase,
|
||||||
|
timeout_ms,
|
||||||
|
} => {
|
||||||
|
assert_eq!(trace_id, "trace-planning-timeout");
|
||||||
|
assert_eq!(phase, "next_execution_attempt");
|
||||||
|
assert_eq!(timeout_ms, 5);
|
||||||
|
}
|
||||||
|
other => panic!("unexpected error: {other:?}"),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn test_report_context() -> serde_json::Value {
|
fn test_report_context() -> serde_json::Value {
|
||||||
json!({
|
json!({
|
||||||
"request_id": "req_watchdog",
|
"request_id": "req_watchdog",
|
||||||
|
|||||||
@@ -317,12 +317,7 @@ pub(super) async fn execute_provider_quota_plan(
|
|||||||
match state.execute_execution_runtime_sync_plan(None, &plan).await {
|
match state.execute_execution_runtime_sync_plan(None, &plan).await {
|
||||||
Ok(result) => Ok(ProviderQuotaExecutionOutcome::Response(result)),
|
Ok(result) => Ok(ProviderQuotaExecutionOutcome::Response(result)),
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
let error = match err {
|
let error = err.into_message();
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
};
|
|
||||||
let proxy_node_id = plan
|
let proxy_node_id = plan
|
||||||
.proxy
|
.proxy
|
||||||
.as_ref()
|
.as_ref()
|
||||||
|
|||||||
@@ -301,12 +301,7 @@ fn admin_provider_ops_decode_response_bytes(
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn admin_provider_ops_gateway_error_message(error: GatewayError) -> String {
|
fn admin_provider_ops_gateway_error_message(error: GatewayError) -> String {
|
||||||
match error {
|
error.into_message()
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn admin_provider_ops_verify_execution_error_message(error: &str) -> String {
|
pub(super) fn admin_provider_ops_verify_execution_error_message(error: &str) -> String {
|
||||||
|
|||||||
@@ -16,7 +16,12 @@ use aether_runtime_state::{DataLayerError, RuntimeState};
|
|||||||
use futures_util::future::join_all;
|
use futures_util::future::join_all;
|
||||||
use std::collections::{BTreeMap, BTreeSet};
|
use std::collections::{BTreeMap, BTreeSet};
|
||||||
use std::time::{SystemTime, UNIX_EPOCH};
|
use std::time::{SystemTime, UNIX_EPOCH};
|
||||||
use tracing::warn;
|
use tracing::{info, warn};
|
||||||
|
|
||||||
|
const DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 512;
|
||||||
|
const MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT: usize = 10_000;
|
||||||
|
const POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV: &str =
|
||||||
|
"AETHER_GATEWAY_ADMIN_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT";
|
||||||
|
|
||||||
fn current_unix_secs() -> u64 {
|
fn current_unix_secs() -> u64 {
|
||||||
SystemTime::now()
|
SystemTime::now()
|
||||||
@@ -29,6 +34,20 @@ fn should_load_active_probe_members(pool_config: &AdminProviderPoolConfig) -> bo
|
|||||||
pool_config.probing_enabled
|
pool_config.probing_enabled
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn pool_runtime_window_metric_key_limit() -> usize {
|
||||||
|
std::env::var(POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT_ENV)
|
||||||
|
.ok()
|
||||||
|
.and_then(|value| value.trim().parse::<usize>().ok())
|
||||||
|
.filter(|value| *value > 0)
|
||||||
|
.unwrap_or(DEFAULT_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
|
||||||
|
.clamp(1, MAX_POOL_RUNTIME_WINDOW_METRIC_KEY_LIMIT)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn bounded_runtime_window_metric_key_ids(key_ids: &[String], limit: usize) -> &[String] {
|
||||||
|
let end = key_ids.len().min(limit.max(1));
|
||||||
|
&key_ids[..end]
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) async fn read_admin_provider_pool_cooldown_counts(
|
pub(crate) async fn read_admin_provider_pool_cooldown_counts(
|
||||||
runtime: &RuntimeState,
|
runtime: &RuntimeState,
|
||||||
provider_ids: &[String],
|
provider_ids: &[String],
|
||||||
@@ -54,8 +73,21 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
|
|||||||
) -> AdminProviderPoolRuntimeState {
|
) -> AdminProviderPoolRuntimeState {
|
||||||
let mut state = AdminProviderPoolRuntimeState::default();
|
let mut state = AdminProviderPoolRuntimeState::default();
|
||||||
let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
|
let cooldown_keys = pool_cooldown_keys(provider_id, key_ids);
|
||||||
let cost_keys = pool_cost_keys(provider_id, key_ids);
|
let metric_key_limit = pool_runtime_window_metric_key_limit();
|
||||||
let latency_keys = pool_latency_keys(provider_id, key_ids);
|
let metric_key_ids = bounded_runtime_window_metric_key_ids(key_ids, metric_key_limit);
|
||||||
|
if metric_key_ids.len() < key_ids.len() {
|
||||||
|
info!(
|
||||||
|
event_name = "admin_pool_runtime_window_metrics_truncated",
|
||||||
|
log_type = "event",
|
||||||
|
provider_id,
|
||||||
|
total_key_count = key_ids.len(),
|
||||||
|
scanned_key_count = metric_key_ids.len(),
|
||||||
|
metric_key_limit,
|
||||||
|
"gateway limited admin pool runtime cost/latency window reads"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
let cost_keys = pool_cost_keys(provider_id, metric_key_ids);
|
||||||
|
let latency_keys = pool_latency_keys(provider_id, metric_key_ids);
|
||||||
let sticky_sessions_enabled = pool_config.sticky_session_ttl_seconds > 0
|
let sticky_sessions_enabled = pool_config.sticky_session_ttl_seconds > 0
|
||||||
&& admin_provider_pool_cache_affinity_enabled(pool_config);
|
&& admin_provider_pool_cache_affinity_enabled(pool_config);
|
||||||
|
|
||||||
@@ -179,7 +211,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
|
|||||||
.map(|cost_key| runtime.score_range_by_min(cost_key, cost_window_start)),
|
.map(|cost_key| runtime.score_range_by_min(cost_key, cost_window_start)),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
for (key_id, members) in key_ids.iter().zip(cost_results) {
|
for (key_id, members) in metric_key_ids.iter().zip(cost_results) {
|
||||||
let total = members
|
let total = members
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
.iter()
|
.iter()
|
||||||
@@ -197,7 +229,7 @@ pub(crate) async fn read_admin_provider_pool_runtime_state(
|
|||||||
.map(|latency_key| runtime.score_range_by_min(latency_key, latency_window_start)),
|
.map(|latency_key| runtime.score_range_by_min(latency_key, latency_window_start)),
|
||||||
)
|
)
|
||||||
.await;
|
.await;
|
||||||
for (key_id, members) in key_ids.iter().zip(latency_results) {
|
for (key_id, members) in metric_key_ids.iter().zip(latency_results) {
|
||||||
let samples = members
|
let samples = members
|
||||||
.unwrap_or_default()
|
.unwrap_or_default()
|
||||||
.iter()
|
.iter()
|
||||||
@@ -265,3 +297,30 @@ pub(crate) async fn read_admin_provider_pool_key_cooldown_reason(
|
|||||||
.kv_get(&pool_cooldown_key(provider_id, key_id))
|
.kv_get(&pool_cooldown_key(provider_id, key_id))
|
||||||
.await
|
.await
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::bounded_runtime_window_metric_key_ids;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn runtime_window_metric_key_ids_are_bounded() {
|
||||||
|
let key_ids = vec![
|
||||||
|
"key-1".to_string(),
|
||||||
|
"key-2".to_string(),
|
||||||
|
"key-3".to_string(),
|
||||||
|
];
|
||||||
|
|
||||||
|
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 2);
|
||||||
|
|
||||||
|
assert_eq!(bounded, &key_ids[..2]);
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn runtime_window_metric_key_ids_keep_at_least_one_key() {
|
||||||
|
let key_ids = vec!["key-1".to_string(), "key-2".to_string()];
|
||||||
|
|
||||||
|
let bounded = bounded_runtime_window_metric_key_ids(&key_ids, 0);
|
||||||
|
|
||||||
|
assert_eq!(bounded, &key_ids[..1]);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|||||||
@@ -662,10 +662,5 @@ fn admin_provider_oauth_decode_response_bytes(
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn admin_provider_oauth_gateway_error_message(error: GatewayError) -> String {
|
fn admin_provider_oauth_gateway_error_message(error: GatewayError) -> String {
|
||||||
match error {
|
error.into_message()
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -354,12 +354,7 @@ pub(crate) async fn maybe_build_internal_finalize_video_response(
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn gateway_error_message(error: GatewayError) -> String {
|
pub(crate) fn gateway_error_message(error: GatewayError) -> String {
|
||||||
match error {
|
error.into_message()
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) fn build_internal_tunnel_heartbeat_ack(
|
pub(crate) fn build_internal_tunnel_heartbeat_ack(
|
||||||
|
|||||||
@@ -21,13 +21,13 @@ use crate::constants::{
|
|||||||
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
EXECUTION_PATH_EXECUTION_RUNTIME_STREAM, EXECUTION_PATH_EXECUTION_RUNTIME_SYNC,
|
||||||
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED,
|
EXECUTION_PATH_LOCAL_AI_PUBLIC, EXECUTION_PATH_LOCAL_API_KEY_CONCURRENCY_LIMITED,
|
||||||
EXECUTION_PATH_LOCAL_AUTH_DENIED, EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED,
|
EXECUTION_PATH_LOCAL_AUTH_DENIED, EXECUTION_PATH_LOCAL_EXECUTION_LOOP_DETECTED,
|
||||||
EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS, EXECUTION_PATH_LOCAL_INVALID_REQUEST,
|
EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT, EXECUTION_PATH_LOCAL_EXECUTION_RUNTIME_MISS,
|
||||||
EXECUTION_PATH_LOCAL_OVERLOADED, EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED,
|
EXECUTION_PATH_LOCAL_INVALID_REQUEST, EXECUTION_PATH_LOCAL_OVERLOADED,
|
||||||
EXECUTION_PATH_LOCAL_RATE_LIMITED, EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND,
|
EXECUTION_PATH_LOCAL_PROXY_PASSTHROUGH_REMOVED, EXECUTION_PATH_LOCAL_RATE_LIMITED,
|
||||||
EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH, EXECUTION_RUNTIME_LOOP_GUARD_HEADER,
|
EXECUTION_PATH_LOCAL_ROUTE_NOT_FOUND, EXECUTION_PATH_PUBLIC_PROXY_PASSTHROUGH,
|
||||||
FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER, FORWARDED_PROTO_HEADER, GATEWAY_HEADER,
|
EXECUTION_RUNTIME_LOOP_GUARD_HEADER, FORWARDED_FOR_HEADER, FORWARDED_HOST_HEADER,
|
||||||
LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER, TRACE_ID_HEADER,
|
FORWARDED_PROTO_HEADER, GATEWAY_HEADER, LOCAL_EXECUTION_RUNTIME_MISS_REASON_HEADER,
|
||||||
TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
|
TRACE_ID_HEADER, TRUSTED_AUTH_ACCESS_ALLOWED_HEADER, TRUSTED_AUTH_API_KEY_ID_HEADER,
|
||||||
TRUSTED_AUTH_BALANCE_HEADER, TRUSTED_AUTH_USER_ID_HEADER, TUNNEL_AFFINITY_FORWARDED_BY_HEADER,
|
TRUSTED_AUTH_BALANCE_HEADER, TRUSTED_AUTH_USER_ID_HEADER, TUNNEL_AFFINITY_FORWARDED_BY_HEADER,
|
||||||
TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER,
|
TUNNEL_AFFINITY_OWNER_INSTANCE_HEADER,
|
||||||
};
|
};
|
||||||
@@ -65,7 +65,11 @@ use axum::extract::{ConnectInfo, Request, State};
|
|||||||
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
|
use axum::http::{self, header::HeaderName, header::HeaderValue, Response};
|
||||||
use futures_util::StreamExt;
|
use futures_util::StreamExt;
|
||||||
use sha2::{Digest, Sha256};
|
use sha2::{Digest, Sha256};
|
||||||
use std::{collections::BTreeMap, time::Instant};
|
use std::{
|
||||||
|
collections::BTreeMap,
|
||||||
|
error::Error as StdError,
|
||||||
|
time::{Duration, Instant},
|
||||||
|
};
|
||||||
use tracing::{debug, info, warn};
|
use tracing::{debug, info, warn};
|
||||||
|
|
||||||
const OPENAI_CHAT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
|
const OPENAI_CHAT_LOCAL_EXECUTION_RUNTIME_MISS_DETAIL: &str =
|
||||||
@@ -89,23 +93,192 @@ const LOCAL_EXECUTION_LOOP_DETECTED_DETAIL: &str =
|
|||||||
"Gateway detected an execution runtime request loop back into the local frontdoor";
|
"Gateway detected an execution runtime request loop back into the local frontdoor";
|
||||||
const AUTH_API_KEY_CONCURRENCY_LIMIT_REACHED_DETAIL: &str =
|
const AUTH_API_KEY_CONCURRENCY_LIMIT_REACHED_DETAIL: &str =
|
||||||
"当前 API Key 并发请求数已达上限,请稍后重试";
|
"当前 API Key 并发请求数已达上限,请稍后重试";
|
||||||
|
const REQUEST_BODY_READ_TIMEOUT_DETAIL: &str =
|
||||||
|
"Request body read timed out before the gateway could route the request";
|
||||||
|
const REQUEST_BODY_READ_FAILED_DETAIL: &str = "Failed to read request body";
|
||||||
|
const LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL: &str =
|
||||||
|
"当前 AI 请求在本地执行规划阶段超时,请稍后重试";
|
||||||
const EXECUTION_PATH_TUNNEL_AFFINITY_FORWARD: &str = "tunnel_affinity_forward";
|
const EXECUTION_PATH_TUNNEL_AFFINITY_FORWARD: &str = "tunnel_affinity_forward";
|
||||||
const MANAGEMENT_TOKEN_PREFIX: &str = "ae-";
|
const MANAGEMENT_TOKEN_PREFIX: &str = "ae-";
|
||||||
const LEGACY_MANAGEMENT_TOKEN_PREFIX: &str = "ae_";
|
const LEGACY_MANAGEMENT_TOKEN_PREFIX: &str = "ae_";
|
||||||
|
|
||||||
fn build_request_body_normalization_error_response(
|
#[derive(Debug, Clone, Copy)]
|
||||||
|
struct RequestBodyBufferPolicy {
|
||||||
|
max_bytes: u64,
|
||||||
|
read_timeout: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RequestBodyBufferPolicy {
|
||||||
|
fn from_state(state: &AppState) -> Self {
|
||||||
|
Self {
|
||||||
|
max_bytes: crate::headers::max_request_body_bytes(),
|
||||||
|
read_timeout: state.frontdoor_runtime_guards.request_body_read_timeout,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
fn for_tests(max_bytes: u64, read_timeout: Duration) -> Self {
|
||||||
|
Self {
|
||||||
|
max_bytes,
|
||||||
|
read_timeout,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
enum RequestBodyBufferError {
|
||||||
|
Normalization(RequestBodyNormalizationError),
|
||||||
|
TooLarge { limit_bytes: u64 },
|
||||||
|
Timeout { timeout_ms: u64 },
|
||||||
|
ReadFailed { message: String },
|
||||||
|
}
|
||||||
|
|
||||||
|
impl RequestBodyBufferError {
|
||||||
|
fn http_status(&self) -> http::StatusCode {
|
||||||
|
match self {
|
||||||
|
Self::Normalization(error) => error.http_status(),
|
||||||
|
Self::TooLarge { .. } => http::StatusCode::PAYLOAD_TOO_LARGE,
|
||||||
|
Self::Timeout { .. } => http::StatusCode::REQUEST_TIMEOUT,
|
||||||
|
Self::ReadFailed { .. } => http::StatusCode::BAD_REQUEST,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn client_message(&self) -> String {
|
||||||
|
match self {
|
||||||
|
Self::Normalization(error) => error.client_message(),
|
||||||
|
Self::TooLarge { limit_bytes } => format!("Request body exceeds {limit_bytes} bytes"),
|
||||||
|
Self::Timeout { .. } => REQUEST_BODY_READ_TIMEOUT_DETAIL.to_string(),
|
||||||
|
Self::ReadFailed { .. } => REQUEST_BODY_READ_FAILED_DETAIL.to_string(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn reason(&self) -> &'static str {
|
||||||
|
match self {
|
||||||
|
Self::Normalization(error) => match error {
|
||||||
|
RequestBodyNormalizationError::UnsupportedContentEncoding(_) => {
|
||||||
|
"unsupported_content_encoding"
|
||||||
|
}
|
||||||
|
RequestBodyNormalizationError::DecodeFailed { .. } => "decode_failed",
|
||||||
|
RequestBodyNormalizationError::DecompressedBodyTooLarge { .. } => {
|
||||||
|
"decompressed_body_too_large"
|
||||||
|
}
|
||||||
|
RequestBodyNormalizationError::RequestBodyTooLarge { .. } => {
|
||||||
|
"request_body_too_large"
|
||||||
|
}
|
||||||
|
},
|
||||||
|
Self::TooLarge { .. } => "request_body_too_large",
|
||||||
|
Self::Timeout { .. } => "request_body_read_timeout",
|
||||||
|
Self::ReadFailed { .. } => "request_body_read_failed",
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn buffer_and_normalize_request_body(
|
||||||
|
request_body: &mut Option<Body>,
|
||||||
|
headers: &mut http::HeaderMap,
|
||||||
|
body_owner_expectation: &'static str,
|
||||||
|
trace_id: &str,
|
||||||
|
method: &http::Method,
|
||||||
|
path_and_query: &str,
|
||||||
|
phase: &'static str,
|
||||||
|
policy: RequestBodyBufferPolicy,
|
||||||
|
) -> Result<Bytes, RequestBodyBufferError> {
|
||||||
|
if let Err(err) =
|
||||||
|
crate::headers::check_request_content_length_with_limit(headers, policy.max_bytes)
|
||||||
|
{
|
||||||
|
return Err(RequestBodyBufferError::Normalization(err));
|
||||||
|
}
|
||||||
|
|
||||||
|
let read_started_at = Instant::now();
|
||||||
|
let timeout_ms = policy.read_timeout.as_millis() as u64;
|
||||||
|
info!(
|
||||||
|
event_name = "frontdoor_request_body_buffer_started",
|
||||||
|
log_type = "event",
|
||||||
|
trace_id,
|
||||||
|
method = %method,
|
||||||
|
path = %path_and_query,
|
||||||
|
phase,
|
||||||
|
max_body_bytes = policy.max_bytes,
|
||||||
|
timeout_ms,
|
||||||
|
"gateway started buffering request body"
|
||||||
|
);
|
||||||
|
|
||||||
|
let body_limit = usize::try_from(policy.max_bytes).unwrap_or(usize::MAX);
|
||||||
|
let body = match tokio::time::timeout(
|
||||||
|
policy.read_timeout,
|
||||||
|
to_bytes(
|
||||||
|
request_body.take().expect(body_owner_expectation),
|
||||||
|
body_limit,
|
||||||
|
),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
{
|
||||||
|
Ok(Ok(body)) => body,
|
||||||
|
Ok(Err(err)) if request_body_collection_exceeded_limit(&err) => {
|
||||||
|
return Err(RequestBodyBufferError::TooLarge {
|
||||||
|
limit_bytes: policy.max_bytes,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Ok(Err(err)) => {
|
||||||
|
return Err(RequestBodyBufferError::ReadFailed {
|
||||||
|
message: err.to_string(),
|
||||||
|
});
|
||||||
|
}
|
||||||
|
Err(_) => {
|
||||||
|
return Err(RequestBodyBufferError::Timeout { timeout_ms });
|
||||||
|
}
|
||||||
|
};
|
||||||
|
|
||||||
|
let normalized = crate::headers::normalize_request_body_headers_and_bytes_with_limit(
|
||||||
|
headers,
|
||||||
|
body,
|
||||||
|
policy.max_bytes,
|
||||||
|
)
|
||||||
|
.map_err(RequestBodyBufferError::Normalization)?;
|
||||||
|
info!(
|
||||||
|
event_name = "frontdoor_request_body_buffer_completed",
|
||||||
|
log_type = "event",
|
||||||
|
trace_id,
|
||||||
|
method = %method,
|
||||||
|
path = %path_and_query,
|
||||||
|
phase,
|
||||||
|
body_bytes = normalized.len(),
|
||||||
|
elapsed_ms = read_started_at.elapsed().as_millis() as u64,
|
||||||
|
"gateway completed request body buffering"
|
||||||
|
);
|
||||||
|
Ok(normalized)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn request_body_collection_exceeded_limit(error: &(dyn StdError + 'static)) -> bool {
|
||||||
|
let mut current = Some(error);
|
||||||
|
while let Some(error) = current {
|
||||||
|
if error.to_string().contains("length limit exceeded") {
|
||||||
|
return true;
|
||||||
|
}
|
||||||
|
current = error.source();
|
||||||
|
}
|
||||||
|
false
|
||||||
|
}
|
||||||
|
|
||||||
|
fn build_request_body_buffer_error_response(
|
||||||
trace_id: &str,
|
trace_id: &str,
|
||||||
request_context: &GatewayPublicRequestContext,
|
request_context: &GatewayPublicRequestContext,
|
||||||
error: &RequestBodyNormalizationError,
|
error: &RequestBodyBufferError,
|
||||||
) -> Result<Response<Body>, GatewayError> {
|
) -> Result<Response<Body>, GatewayError> {
|
||||||
warn!(
|
warn!(
|
||||||
event_name = "frontdoor_request_body_normalization_failed",
|
event_name = "frontdoor_request_body_buffer_failed",
|
||||||
log_type = "ops",
|
log_type = "ops",
|
||||||
trace_id,
|
trace_id,
|
||||||
method = %request_context.request_method,
|
method = %request_context.request_method,
|
||||||
path = %request_context.request_path_and_query(),
|
path = %request_context.request_path_and_query(),
|
||||||
error = %error,
|
status_code = error.http_status().as_u16(),
|
||||||
"gateway rejected request with invalid encoded body"
|
reason = error.reason(),
|
||||||
|
detail = %error.client_message(),
|
||||||
|
read_error = match error {
|
||||||
|
RequestBodyBufferError::ReadFailed { message } => message.as_str(),
|
||||||
|
_ => "",
|
||||||
|
},
|
||||||
|
"gateway rejected request body before local execution planning"
|
||||||
);
|
);
|
||||||
build_local_http_error_response(
|
build_local_http_error_response(
|
||||||
trace_id,
|
trace_id,
|
||||||
@@ -115,36 +288,16 @@ fn build_request_body_normalization_error_response(
|
|||||||
)
|
)
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn buffer_and_normalize_request_body(
|
fn finalize_request_body_buffer_rejection(
|
||||||
request_body: &mut Option<Body>,
|
|
||||||
headers: &mut http::HeaderMap,
|
|
||||||
body_owner_expectation: &'static str,
|
|
||||||
) -> Result<Result<Bytes, RequestBodyNormalizationError>, GatewayError> {
|
|
||||||
if let Err(err) = crate::headers::check_request_content_length(headers) {
|
|
||||||
return Ok(Err(err));
|
|
||||||
}
|
|
||||||
let body = to_bytes(
|
|
||||||
request_body.take().expect(body_owner_expectation),
|
|
||||||
usize::MAX,
|
|
||||||
)
|
|
||||||
.await
|
|
||||||
.map_err(|err| GatewayError::Internal(err.to_string()))?;
|
|
||||||
Ok(crate::headers::normalize_request_body_headers_and_bytes(
|
|
||||||
headers, body,
|
|
||||||
))
|
|
||||||
}
|
|
||||||
|
|
||||||
fn finalize_request_body_normalization_rejection(
|
|
||||||
state: &AppState,
|
state: &AppState,
|
||||||
request_context: &GatewayPublicRequestContext,
|
request_context: &GatewayPublicRequestContext,
|
||||||
remote_addr: &std::net::SocketAddr,
|
remote_addr: &std::net::SocketAddr,
|
||||||
started_at: &std::time::Instant,
|
started_at: &std::time::Instant,
|
||||||
trace_id: &str,
|
trace_id: &str,
|
||||||
request_permit: Option<aether_runtime::AdmissionPermit>,
|
request_permit: Option<aether_runtime::AdmissionPermit>,
|
||||||
error: &RequestBodyNormalizationError,
|
error: &RequestBodyBufferError,
|
||||||
) -> Result<Response<Body>, GatewayError> {
|
) -> Result<Response<Body>, GatewayError> {
|
||||||
let response =
|
let response = build_request_body_buffer_error_response(trace_id, request_context, error)?;
|
||||||
build_request_body_normalization_error_response(trace_id, request_context, error)?;
|
|
||||||
Ok(finalize_gateway_response_with_context(
|
Ok(finalize_gateway_response_with_context(
|
||||||
state,
|
state,
|
||||||
response,
|
response,
|
||||||
@@ -156,6 +309,59 @@ fn finalize_request_body_normalization_rejection(
|
|||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn local_execution_planning_timeout_parts(error: &GatewayError) -> Option<(&'static str, u64)> {
|
||||||
|
match error {
|
||||||
|
GatewayError::LocalExecutionPlanningTimeout {
|
||||||
|
phase, timeout_ms, ..
|
||||||
|
} => Some((*phase, *timeout_ms)),
|
||||||
|
_ => None,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn finalize_local_execution_planning_timeout(
|
||||||
|
state: &AppState,
|
||||||
|
request_context: &GatewayPublicRequestContext,
|
||||||
|
remote_addr: &std::net::SocketAddr,
|
||||||
|
started_at: &std::time::Instant,
|
||||||
|
trace_id: &str,
|
||||||
|
request_permit: Option<aether_runtime::AdmissionPermit>,
|
||||||
|
control_decision: Option<&GatewayControlDecision>,
|
||||||
|
phase: &'static str,
|
||||||
|
timeout_ms: u64,
|
||||||
|
) -> Result<Response<Body>, GatewayError> {
|
||||||
|
warn!(
|
||||||
|
event_name = "frontdoor_local_execution_planning_timeout",
|
||||||
|
log_type = "ops",
|
||||||
|
trace_id,
|
||||||
|
method = %request_context.request_method,
|
||||||
|
path = %request_context.request_path_and_query(),
|
||||||
|
route_family = control_decision
|
||||||
|
.and_then(|decision| decision.route_family.as_deref())
|
||||||
|
.unwrap_or("-"),
|
||||||
|
route_kind = control_decision
|
||||||
|
.and_then(|decision| decision.route_kind.as_deref())
|
||||||
|
.unwrap_or("-"),
|
||||||
|
phase,
|
||||||
|
timeout_ms,
|
||||||
|
"gateway failed local execution before a candidate could be selected"
|
||||||
|
);
|
||||||
|
let response = build_local_http_error_response(
|
||||||
|
trace_id,
|
||||||
|
control_decision,
|
||||||
|
http::StatusCode::GATEWAY_TIMEOUT,
|
||||||
|
LOCAL_EXECUTION_PLANNING_TIMEOUT_DETAIL,
|
||||||
|
)?;
|
||||||
|
Ok(finalize_gateway_response_with_context(
|
||||||
|
state,
|
||||||
|
response,
|
||||||
|
remote_addr,
|
||||||
|
request_context,
|
||||||
|
EXECUTION_PATH_LOCAL_EXECUTION_PLANNING_TIMEOUT,
|
||||||
|
started_at,
|
||||||
|
request_permit,
|
||||||
|
))
|
||||||
|
}
|
||||||
|
|
||||||
fn local_execution_outcome_label(outcome: &LocalExecutionRequestOutcome) -> &'static str {
|
fn local_execution_outcome_label(outcome: &LocalExecutionRequestOutcome) -> &'static str {
|
||||||
match outcome {
|
match outcome {
|
||||||
LocalExecutionRequestOutcome::Responded(_) => "responded",
|
LocalExecutionRequestOutcome::Responded(_) => "responded",
|
||||||
@@ -991,16 +1197,22 @@ pub(crate) async fn proxy_request(
|
|||||||
}
|
}
|
||||||
let mut request_body = Some(body);
|
let mut request_body = Some(body);
|
||||||
let local_proxy_body = if local_proxy_route_requires_buffered_body(&request_context) {
|
let local_proxy_body = if local_proxy_route_requires_buffered_body(&request_context) {
|
||||||
|
let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state);
|
||||||
let body = buffer_and_normalize_request_body(
|
let body = buffer_and_normalize_request_body(
|
||||||
&mut request_body,
|
&mut request_body,
|
||||||
&mut parts.headers,
|
&mut parts.headers,
|
||||||
"local proxy body buffering should own request body",
|
"local proxy body buffering should own request body",
|
||||||
|
&trace_id,
|
||||||
|
&parts.method,
|
||||||
|
&request_context.request_path_and_query(),
|
||||||
|
"local_proxy",
|
||||||
|
body_buffer_policy,
|
||||||
)
|
)
|
||||||
.await?;
|
.await;
|
||||||
match body {
|
match body {
|
||||||
Ok(body) => Some(body),
|
Ok(body) => Some(body),
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
return finalize_request_body_normalization_rejection(
|
return finalize_request_body_buffer_rejection(
|
||||||
&state,
|
&state,
|
||||||
&request_context,
|
&request_context,
|
||||||
&remote_addr,
|
&remote_addr,
|
||||||
@@ -1171,16 +1383,22 @@ pub(crate) async fn proxy_request(
|
|||||||
&& request_enables_control_execute(&parts.headers);
|
&& request_enables_control_execute(&parts.headers);
|
||||||
|
|
||||||
let buffered_body = if should_buffer_body {
|
let buffered_body = if should_buffer_body {
|
||||||
|
let body_buffer_policy = RequestBodyBufferPolicy::from_state(&state);
|
||||||
let body = buffer_and_normalize_request_body(
|
let body = buffer_and_normalize_request_body(
|
||||||
&mut request_body,
|
&mut request_body,
|
||||||
&mut parts.headers,
|
&mut parts.headers,
|
||||||
"buffered auth/execution runtime path should own request body",
|
"buffered auth/execution runtime path should own request body",
|
||||||
|
&trace_id,
|
||||||
|
&parts.method,
|
||||||
|
&request_context.request_path_and_query(),
|
||||||
|
"auth_execution",
|
||||||
|
body_buffer_policy,
|
||||||
)
|
)
|
||||||
.await?;
|
.await;
|
||||||
match body {
|
match body {
|
||||||
Ok(body) => Some(body),
|
Ok(body) => Some(body),
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
return finalize_request_body_normalization_rejection(
|
return finalize_request_body_buffer_rejection(
|
||||||
&state,
|
&state,
|
||||||
&request_context,
|
&request_context,
|
||||||
&remote_addr,
|
&remote_addr,
|
||||||
@@ -1335,14 +1553,34 @@ pub(crate) async fn proxy_request(
|
|||||||
let stream_request = request_wants_stream(&request_context, &parts.headers, buffered_body);
|
let stream_request = request_wants_stream(&request_context, &parts.headers, buffered_body);
|
||||||
let mut local_execution_exhaustion = None;
|
let mut local_execution_exhaustion = None;
|
||||||
if stream_request {
|
if stream_request {
|
||||||
let stream_outcome = maybe_execute_stream_request(
|
let stream_outcome = match maybe_execute_stream_request(
|
||||||
&state,
|
&state,
|
||||||
&parts,
|
&parts,
|
||||||
buffered_body,
|
buffered_body,
|
||||||
&trace_id,
|
&trace_id,
|
||||||
control_decision,
|
control_decision,
|
||||||
)
|
)
|
||||||
.await?;
|
.await
|
||||||
|
{
|
||||||
|
Ok(outcome) => outcome,
|
||||||
|
Err(err) => {
|
||||||
|
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err)
|
||||||
|
{
|
||||||
|
return finalize_local_execution_planning_timeout(
|
||||||
|
&state,
|
||||||
|
&request_context,
|
||||||
|
&remote_addr,
|
||||||
|
&started_at,
|
||||||
|
&trace_id,
|
||||||
|
request_permit.take(),
|
||||||
|
control_decision,
|
||||||
|
phase,
|
||||||
|
timeout_ms,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
};
|
||||||
debug!(
|
debug!(
|
||||||
event_name = "proxy_stream_local_execute_outcome",
|
event_name = "proxy_stream_local_execute_outcome",
|
||||||
log_type = "debug",
|
log_type = "debug",
|
||||||
@@ -1380,9 +1618,34 @@ pub(crate) async fn proxy_request(
|
|||||||
LocalExecutionRequestOutcome::NoPath => {}
|
LocalExecutionRequestOutcome::NoPath => {}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
match maybe_execute_sync_request(&state, &parts, buffered_body, &trace_id, control_decision)
|
let sync_outcome = match maybe_execute_sync_request(
|
||||||
.await?
|
&state,
|
||||||
|
&parts,
|
||||||
|
buffered_body,
|
||||||
|
&trace_id,
|
||||||
|
control_decision,
|
||||||
|
)
|
||||||
|
.await
|
||||||
{
|
{
|
||||||
|
Ok(outcome) => outcome,
|
||||||
|
Err(err) => {
|
||||||
|
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err) {
|
||||||
|
return finalize_local_execution_planning_timeout(
|
||||||
|
&state,
|
||||||
|
&request_context,
|
||||||
|
&remote_addr,
|
||||||
|
&started_at,
|
||||||
|
&trace_id,
|
||||||
|
request_permit.take(),
|
||||||
|
control_decision,
|
||||||
|
phase,
|
||||||
|
timeout_ms,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
match sync_outcome {
|
||||||
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
|
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
|
||||||
let execution_runtime_response = restore_redacted_sync_execution_response(
|
let execution_runtime_response = restore_redacted_sync_execution_response(
|
||||||
execution_runtime_response,
|
execution_runtime_response,
|
||||||
@@ -1406,15 +1669,35 @@ pub(crate) async fn proxy_request(
|
|||||||
LocalExecutionRequestOutcome::NoPath => {}
|
LocalExecutionRequestOutcome::NoPath => {}
|
||||||
}
|
}
|
||||||
if parts.method != http::Method::POST {
|
if parts.method != http::Method::POST {
|
||||||
match maybe_execute_stream_request(
|
let stream_outcome = match maybe_execute_stream_request(
|
||||||
&state,
|
&state,
|
||||||
&parts,
|
&parts,
|
||||||
buffered_body,
|
buffered_body,
|
||||||
&trace_id,
|
&trace_id,
|
||||||
control_decision,
|
control_decision,
|
||||||
)
|
)
|
||||||
.await?
|
.await
|
||||||
{
|
{
|
||||||
|
Ok(outcome) => outcome,
|
||||||
|
Err(err) => {
|
||||||
|
if let Some((phase, timeout_ms)) = local_execution_planning_timeout_parts(&err)
|
||||||
|
{
|
||||||
|
return finalize_local_execution_planning_timeout(
|
||||||
|
&state,
|
||||||
|
&request_context,
|
||||||
|
&remote_addr,
|
||||||
|
&started_at,
|
||||||
|
&trace_id,
|
||||||
|
request_permit.take(),
|
||||||
|
control_decision,
|
||||||
|
phase,
|
||||||
|
timeout_ms,
|
||||||
|
);
|
||||||
|
}
|
||||||
|
return Err(err);
|
||||||
|
}
|
||||||
|
};
|
||||||
|
match stream_outcome {
|
||||||
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
|
LocalExecutionRequestOutcome::Responded(execution_runtime_response) => {
|
||||||
let execution_runtime_response = restore_redacted_stream_execution_response(
|
let execution_runtime_response = restore_redacted_stream_execution_response(
|
||||||
execution_runtime_response,
|
execution_runtime_response,
|
||||||
@@ -1977,14 +2260,17 @@ fn local_execution_runtime_miss_route_detail(
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
use super::{
|
use super::{
|
||||||
api_key_remote_ip_allowed, diagnostic_is_auth_api_key_concurrency_limited,
|
api_key_remote_ip_allowed, buffer_and_normalize_request_body,
|
||||||
local_execution_runtime_miss_detail, restore_redacted_stream_execution_response,
|
diagnostic_is_auth_api_key_concurrency_limited, local_execution_runtime_miss_detail,
|
||||||
restore_redacted_sync_execution_response, GatewayControlDecision,
|
restore_redacted_stream_execution_response, restore_redacted_sync_execution_response,
|
||||||
LocalExecutionRuntimeMissDiagnostic,
|
GatewayControlDecision, LocalExecutionRuntimeMissDiagnostic, RequestBodyBufferError,
|
||||||
|
RequestBodyBufferPolicy,
|
||||||
};
|
};
|
||||||
use axum::body::{to_bytes, Body};
|
use axum::body::{to_bytes, Body, Bytes};
|
||||||
use axum::http::{header, Response};
|
use axum::http::{header, HeaderMap, Method, Response};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
@@ -2112,6 +2398,58 @@ mod tests {
|
|||||||
assert!(!message.contains(&sentinel));
|
assert!(!message.contains(&sentinel));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn request_body_buffer_rejects_chunked_body_when_limit_is_exceeded() {
|
||||||
|
let mut body = Some(Body::from(Bytes::from_static(b"abcdef")));
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
|
||||||
|
let err = buffer_and_normalize_request_body(
|
||||||
|
&mut body,
|
||||||
|
&mut headers,
|
||||||
|
"test owns body",
|
||||||
|
"trace-body-large",
|
||||||
|
&Method::POST,
|
||||||
|
"/v1/responses",
|
||||||
|
"test",
|
||||||
|
RequestBodyBufferPolicy::for_tests(5, Duration::from_secs(1)),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect_err("body exceeding the ingress limit should fail");
|
||||||
|
|
||||||
|
assert!(matches!(
|
||||||
|
err,
|
||||||
|
RequestBodyBufferError::TooLarge { limit_bytes: 5 }
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn request_body_buffer_times_out_instead_of_waiting_forever() {
|
||||||
|
let stream = async_stream::stream! {
|
||||||
|
yield Ok::<Bytes, std::io::Error>(Bytes::from_static(b"{"));
|
||||||
|
std::future::pending::<()>().await;
|
||||||
|
};
|
||||||
|
let mut body = Some(Body::from_stream(stream));
|
||||||
|
let mut headers = HeaderMap::new();
|
||||||
|
|
||||||
|
let err = buffer_and_normalize_request_body(
|
||||||
|
&mut body,
|
||||||
|
&mut headers,
|
||||||
|
"test owns body",
|
||||||
|
"trace-body-timeout",
|
||||||
|
&Method::POST,
|
||||||
|
"/v1/responses",
|
||||||
|
"test",
|
||||||
|
RequestBodyBufferPolicy::for_tests(1024, Duration::from_millis(5)),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect_err("body buffering should time out");
|
||||||
|
|
||||||
|
assert!(matches!(
|
||||||
|
err,
|
||||||
|
RequestBodyBufferError::Timeout { timeout_ms: 5 }
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn runtime_miss_detail_returns_model_specific_stream_message_when_candidates_are_unavailable() {
|
fn runtime_miss_detail_returns_model_specific_stream_message_when_candidates_are_unavailable() {
|
||||||
let decision = GatewayControlDecision::synthetic(
|
let decision = GatewayControlDecision::synthetic(
|
||||||
|
|||||||
@@ -136,12 +136,7 @@ pub(super) fn announcements_internal_error_response(detail: impl Into<String>) -
|
|||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn announcements_internal_detail(err: GatewayError) -> String {
|
pub(super) fn announcements_internal_detail(err: GatewayError) -> String {
|
||||||
match err {
|
err.into_message()
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(super) fn parse_optional_rfc3339_unix_secs(
|
pub(super) fn parse_optional_rfc3339_unix_secs(
|
||||||
|
|||||||
@@ -21,6 +21,10 @@ static MAX_REQUEST_BODY_BYTES: LazyLock<u64> = LazyLock::new(|| {
|
|||||||
.saturating_mul(1024 * 1024)
|
.saturating_mul(1024 * 1024)
|
||||||
});
|
});
|
||||||
|
|
||||||
|
pub(crate) fn max_request_body_bytes() -> u64 {
|
||||||
|
*MAX_REQUEST_BODY_BYTES
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn extract_or_generate_trace_id(headers: &http::HeaderMap) -> String {
|
pub(crate) fn extract_or_generate_trace_id(headers: &http::HeaderMap) -> String {
|
||||||
header_value_str(headers, TRACE_ID_HEADER).unwrap_or_else(|| Uuid::new_v4().to_string())
|
header_value_str(headers, TRACE_ID_HEADER).unwrap_or_else(|| Uuid::new_v4().to_string())
|
||||||
}
|
}
|
||||||
@@ -246,9 +250,21 @@ impl std::error::Error for RequestBodyNormalizationError {}
|
|||||||
pub(crate) fn normalize_request_body_headers_and_bytes(
|
pub(crate) fn normalize_request_body_headers_and_bytes(
|
||||||
headers: &mut http::HeaderMap,
|
headers: &mut http::HeaderMap,
|
||||||
body_bytes: Bytes,
|
body_bytes: Bytes,
|
||||||
|
) -> Result<Bytes, RequestBodyNormalizationError> {
|
||||||
|
normalize_request_body_headers_and_bytes_with_limit(
|
||||||
|
headers,
|
||||||
|
body_bytes,
|
||||||
|
max_request_body_bytes(),
|
||||||
|
)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn normalize_request_body_headers_and_bytes_with_limit(
|
||||||
|
headers: &mut http::HeaderMap,
|
||||||
|
body_bytes: Bytes,
|
||||||
|
limit_bytes: u64,
|
||||||
) -> Result<Bytes, RequestBodyNormalizationError> {
|
) -> Result<Bytes, RequestBodyNormalizationError> {
|
||||||
let body_was_encoded = !request_content_encodings(headers).is_empty();
|
let body_was_encoded = !request_content_encodings(headers).is_empty();
|
||||||
let decoded = decoded_request_body_bytes(headers, body_bytes.as_ref())?;
|
let decoded = decoded_request_body_bytes_with_limit(headers, body_bytes.as_ref(), limit_bytes)?;
|
||||||
if !body_was_encoded {
|
if !body_was_encoded {
|
||||||
return Ok(body_bytes);
|
return Ok(body_bytes);
|
||||||
}
|
}
|
||||||
@@ -264,7 +280,13 @@ pub(crate) fn normalize_request_body_headers_and_bytes(
|
|||||||
pub(crate) fn check_request_content_length(
|
pub(crate) fn check_request_content_length(
|
||||||
headers: &http::HeaderMap,
|
headers: &http::HeaderMap,
|
||||||
) -> Result<(), RequestBodyNormalizationError> {
|
) -> Result<(), RequestBodyNormalizationError> {
|
||||||
let limit = *MAX_REQUEST_BODY_BYTES;
|
check_request_content_length_with_limit(headers, max_request_body_bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn check_request_content_length_with_limit(
|
||||||
|
headers: &http::HeaderMap,
|
||||||
|
limit: u64,
|
||||||
|
) -> Result<(), RequestBodyNormalizationError> {
|
||||||
let declared = header_value_str(headers, http::header::CONTENT_LENGTH.as_str())
|
let declared = header_value_str(headers, http::header::CONTENT_LENGTH.as_str())
|
||||||
.and_then(|value| value.trim().parse::<u64>().ok());
|
.and_then(|value| value.trim().parse::<u64>().ok());
|
||||||
if declared.is_some_and(|value| value > limit) {
|
if declared.is_some_and(|value| value > limit) {
|
||||||
@@ -276,10 +298,17 @@ pub(crate) fn check_request_content_length(
|
|||||||
pub(crate) fn decoded_request_body_bytes<'a>(
|
pub(crate) fn decoded_request_body_bytes<'a>(
|
||||||
headers: &http::HeaderMap,
|
headers: &http::HeaderMap,
|
||||||
body_bytes: &'a [u8],
|
body_bytes: &'a [u8],
|
||||||
|
) -> Result<Cow<'a, [u8]>, RequestBodyNormalizationError> {
|
||||||
|
decoded_request_body_bytes_with_limit(headers, body_bytes, max_request_body_bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
pub(crate) fn decoded_request_body_bytes_with_limit<'a>(
|
||||||
|
headers: &http::HeaderMap,
|
||||||
|
body_bytes: &'a [u8],
|
||||||
|
limit: u64,
|
||||||
) -> Result<Cow<'a, [u8]>, RequestBodyNormalizationError> {
|
) -> Result<Cow<'a, [u8]>, RequestBodyNormalizationError> {
|
||||||
let encodings = request_content_encodings(headers);
|
let encodings = request_content_encodings(headers);
|
||||||
if encodings.is_empty() {
|
if encodings.is_empty() {
|
||||||
let limit = *MAX_REQUEST_BODY_BYTES;
|
|
||||||
if body_bytes.len() as u64 > limit {
|
if body_bytes.len() as u64 > limit {
|
||||||
return Err(RequestBodyNormalizationError::RequestBodyTooLarge { limit_bytes: limit });
|
return Err(RequestBodyNormalizationError::RequestBodyTooLarge { limit_bytes: limit });
|
||||||
}
|
}
|
||||||
@@ -288,7 +317,7 @@ pub(crate) fn decoded_request_body_bytes<'a>(
|
|||||||
|
|
||||||
let mut decoded = body_bytes.to_vec();
|
let mut decoded = body_bytes.to_vec();
|
||||||
for encoding in encodings.iter().rev() {
|
for encoding in encodings.iter().rev() {
|
||||||
decoded = decode_single_request_body(encoding, decoded.as_slice())?;
|
decoded = decode_single_request_body_with_limit(encoding, decoded.as_slice(), limit)?;
|
||||||
}
|
}
|
||||||
Ok(Cow::Owned(decoded))
|
Ok(Cow::Owned(decoded))
|
||||||
}
|
}
|
||||||
@@ -310,11 +339,19 @@ fn request_content_encodings(headers: &http::HeaderMap) -> Vec<String> {
|
|||||||
fn decode_single_request_body(
|
fn decode_single_request_body(
|
||||||
encoding: &str,
|
encoding: &str,
|
||||||
body_bytes: &[u8],
|
body_bytes: &[u8],
|
||||||
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
|
decode_single_request_body_with_limit(encoding, body_bytes, max_request_body_bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decode_single_request_body_with_limit(
|
||||||
|
encoding: &str,
|
||||||
|
body_bytes: &[u8],
|
||||||
|
limit: u64,
|
||||||
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
match encoding {
|
match encoding {
|
||||||
"gzip" | "x-gzip" => decode_gzip_body(encoding, body_bytes),
|
"gzip" | "x-gzip" => decode_gzip_body_with_limit(encoding, body_bytes, limit),
|
||||||
"deflate" => decode_deflate_body(encoding, body_bytes),
|
"deflate" => decode_deflate_body_with_limit(encoding, body_bytes, limit),
|
||||||
"zstd" => decode_zstd_body(encoding, body_bytes),
|
"zstd" => decode_zstd_body_with_limit(encoding, body_bytes, limit),
|
||||||
_ => Err(RequestBodyNormalizationError::UnsupportedContentEncoding(
|
_ => Err(RequestBodyNormalizationError::UnsupportedContentEncoding(
|
||||||
encoding.to_string(),
|
encoding.to_string(),
|
||||||
)),
|
)),
|
||||||
@@ -324,27 +361,43 @@ fn decode_single_request_body(
|
|||||||
fn decode_gzip_body(
|
fn decode_gzip_body(
|
||||||
encoding: &str,
|
encoding: &str,
|
||||||
body_bytes: &[u8],
|
body_bytes: &[u8],
|
||||||
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
|
decode_gzip_body_with_limit(encoding, body_bytes, max_request_body_bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decode_gzip_body_with_limit(
|
||||||
|
encoding: &str,
|
||||||
|
body_bytes: &[u8],
|
||||||
|
limit: u64,
|
||||||
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
let mut decoder = GzDecoder::new(body_bytes);
|
let mut decoder = GzDecoder::new(body_bytes);
|
||||||
read_request_decoder_to_end(encoding, &mut decoder)
|
read_request_decoder_to_end_with_limit(encoding, &mut decoder, limit)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn decode_deflate_body(
|
fn decode_deflate_body(
|
||||||
encoding: &str,
|
encoding: &str,
|
||||||
body_bytes: &[u8],
|
body_bytes: &[u8],
|
||||||
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
|
decode_deflate_body_with_limit(encoding, body_bytes, max_request_body_bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decode_deflate_body_with_limit(
|
||||||
|
encoding: &str,
|
||||||
|
body_bytes: &[u8],
|
||||||
|
limit: u64,
|
||||||
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
let mut zlib_decoder = ZlibDecoder::new(body_bytes);
|
let mut zlib_decoder = ZlibDecoder::new(body_bytes);
|
||||||
match read_request_decoder_to_end(encoding, &mut zlib_decoder) {
|
match read_request_decoder_to_end_with_limit(encoding, &mut zlib_decoder, limit) {
|
||||||
Ok(decoded) => Ok(decoded),
|
Ok(decoded) => Ok(decoded),
|
||||||
Err(err @ RequestBodyNormalizationError::DecompressedBodyTooLarge { .. }) => Err(err),
|
Err(err @ RequestBodyNormalizationError::DecompressedBodyTooLarge { .. }) => Err(err),
|
||||||
Err(zlib_error) => {
|
Err(zlib_error) => {
|
||||||
let mut raw_decoder = DeflateDecoder::new(body_bytes);
|
let mut raw_decoder = DeflateDecoder::new(body_bytes);
|
||||||
read_request_decoder_to_end(encoding, &mut raw_decoder).map_err(|raw_error| {
|
read_request_decoder_to_end_with_limit(encoding, &mut raw_decoder, limit).map_err(
|
||||||
RequestBodyNormalizationError::DecodeFailed {
|
|raw_error| RequestBodyNormalizationError::DecodeFailed {
|
||||||
encoding: encoding.to_string(),
|
encoding: encoding.to_string(),
|
||||||
reason: format!("{zlib_error}; raw deflate fallback failed: {raw_error}"),
|
reason: format!("{zlib_error}; raw deflate fallback failed: {raw_error}"),
|
||||||
}
|
},
|
||||||
})
|
)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -352,6 +405,14 @@ fn decode_deflate_body(
|
|||||||
fn decode_zstd_body(
|
fn decode_zstd_body(
|
||||||
encoding: &str,
|
encoding: &str,
|
||||||
body_bytes: &[u8],
|
body_bytes: &[u8],
|
||||||
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
|
decode_zstd_body_with_limit(encoding, body_bytes, max_request_body_bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn decode_zstd_body_with_limit(
|
||||||
|
encoding: &str,
|
||||||
|
body_bytes: &[u8],
|
||||||
|
limit: u64,
|
||||||
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
let mut decoder = zstd::stream::read::Decoder::new(body_bytes).map_err(|err| {
|
let mut decoder = zstd::stream::read::Decoder::new(body_bytes).map_err(|err| {
|
||||||
RequestBodyNormalizationError::DecodeFailed {
|
RequestBodyNormalizationError::DecodeFailed {
|
||||||
@@ -359,14 +420,21 @@ fn decode_zstd_body(
|
|||||||
reason: err.to_string(),
|
reason: err.to_string(),
|
||||||
}
|
}
|
||||||
})?;
|
})?;
|
||||||
read_request_decoder_to_end(encoding, &mut decoder)
|
read_request_decoder_to_end_with_limit(encoding, &mut decoder, limit)
|
||||||
}
|
}
|
||||||
|
|
||||||
fn read_request_decoder_to_end(
|
fn read_request_decoder_to_end(
|
||||||
encoding: &str,
|
encoding: &str,
|
||||||
decoder: &mut impl Read,
|
decoder: &mut impl Read,
|
||||||
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
let limit = *MAX_REQUEST_BODY_BYTES;
|
read_request_decoder_to_end_with_limit(encoding, decoder, max_request_body_bytes())
|
||||||
|
}
|
||||||
|
|
||||||
|
fn read_request_decoder_to_end_with_limit(
|
||||||
|
encoding: &str,
|
||||||
|
decoder: &mut impl Read,
|
||||||
|
limit: u64,
|
||||||
|
) -> Result<Vec<u8>, RequestBodyNormalizationError> {
|
||||||
let mut limited = decoder.take(limit.saturating_add(1));
|
let mut limited = decoder.take(limit.saturating_add(1));
|
||||||
let mut out = Vec::new();
|
let mut out = Vec::new();
|
||||||
limited
|
limited
|
||||||
|
|||||||
@@ -551,12 +551,7 @@ fn endpoint_for_self_check(
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn gateway_error_message(err: GatewayError) -> String {
|
fn gateway_error_message(err: GatewayError) -> String {
|
||||||
match err {
|
err.into_message()
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
||||||
fn update_summary_from_outcome(
|
fn update_summary_from_outcome(
|
||||||
|
|||||||
@@ -157,10 +157,5 @@ fn decode_response_bytes(bytes: &[u8], content_encoding: Option<&str>) -> Option
|
|||||||
}
|
}
|
||||||
|
|
||||||
fn gateway_error_to_oauth_error(error: GatewayError) -> OAuthError {
|
fn gateway_error_to_oauth_error(error: GatewayError) -> OAuthError {
|
||||||
match error {
|
OAuthError::Transport(error.into_message())
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => OAuthError::Transport(message),
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -23,6 +23,17 @@ use super::{
|
|||||||
LocalProviderDeleteTaskState, ProviderTransportSnapshotCacheKey,
|
LocalProviderDeleteTaskState, ProviderTransportSnapshotCacheKey,
|
||||||
};
|
};
|
||||||
|
|
||||||
|
const DEFAULT_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 120_000;
|
||||||
|
const MIN_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 1_000;
|
||||||
|
const MAX_REQUEST_BODY_READ_TIMEOUT_MS: u64 = 600_000;
|
||||||
|
const REQUEST_BODY_READ_TIMEOUT_MS_ENV: &str = "AETHER_GATEWAY_REQUEST_BODY_READ_TIMEOUT_MS";
|
||||||
|
|
||||||
|
const DEFAULT_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 30_000;
|
||||||
|
const MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 500;
|
||||||
|
const MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS: u64 = 120_000;
|
||||||
|
const LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV: &str =
|
||||||
|
"AETHER_GATEWAY_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS";
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
type TestExecutionRuntimeSyncOverrideFn = dyn Fn(
|
type TestExecutionRuntimeSyncOverrideFn = dyn Fn(
|
||||||
&aether_contracts::ExecutionPlan,
|
&aether_contracts::ExecutionPlan,
|
||||||
@@ -43,6 +54,52 @@ impl std::fmt::Debug for TestExecutionRuntimeSyncOverride {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub(crate) struct FrontdoorRuntimeGuardConfig {
|
||||||
|
pub(crate) request_body_read_timeout: Duration,
|
||||||
|
pub(crate) local_execution_planning_timeout: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl FrontdoorRuntimeGuardConfig {
|
||||||
|
pub(crate) fn from_env() -> Self {
|
||||||
|
Self {
|
||||||
|
request_body_read_timeout: env_duration_ms(
|
||||||
|
REQUEST_BODY_READ_TIMEOUT_MS_ENV,
|
||||||
|
DEFAULT_REQUEST_BODY_READ_TIMEOUT_MS,
|
||||||
|
MIN_REQUEST_BODY_READ_TIMEOUT_MS,
|
||||||
|
MAX_REQUEST_BODY_READ_TIMEOUT_MS,
|
||||||
|
),
|
||||||
|
local_execution_planning_timeout: env_duration_ms(
|
||||||
|
LOCAL_EXECUTION_PLANNING_TIMEOUT_MS_ENV,
|
||||||
|
DEFAULT_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
|
||||||
|
MIN_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
|
||||||
|
MAX_LOCAL_EXECUTION_PLANNING_TIMEOUT_MS,
|
||||||
|
),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
pub(crate) fn for_tests(
|
||||||
|
request_body_read_timeout: Duration,
|
||||||
|
local_execution_planning_timeout: Duration,
|
||||||
|
) -> Self {
|
||||||
|
Self {
|
||||||
|
request_body_read_timeout,
|
||||||
|
local_execution_planning_timeout,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn env_duration_ms(key: &str, default_ms: u64, min_ms: u64, max_ms: u64) -> Duration {
|
||||||
|
let ms = std::env::var(key)
|
||||||
|
.ok()
|
||||||
|
.and_then(|value| value.trim().parse::<u64>().ok())
|
||||||
|
.filter(|value| *value > 0)
|
||||||
|
.unwrap_or(default_ms)
|
||||||
|
.clamp(min_ms, max_ms);
|
||||||
|
Duration::from_millis(ms)
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone)]
|
#[derive(Debug, Clone)]
|
||||||
pub struct AppState {
|
pub struct AppState {
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -54,6 +111,7 @@ pub struct AppState {
|
|||||||
pub(crate) usage_runtime: Arc<usage::UsageRuntime>,
|
pub(crate) usage_runtime: Arc<usage::UsageRuntime>,
|
||||||
pub(crate) video_tasks: Arc<VideoTaskService>,
|
pub(crate) video_tasks: Arc<VideoTaskService>,
|
||||||
pub(crate) video_task_poller: Option<VideoTaskPollerConfig>,
|
pub(crate) video_task_poller: Option<VideoTaskPollerConfig>,
|
||||||
|
pub(crate) frontdoor_runtime_guards: Arc<FrontdoorRuntimeGuardConfig>,
|
||||||
pub(crate) request_gate: Option<Arc<ConcurrencyGate>>,
|
pub(crate) request_gate: Option<Arc<ConcurrencyGate>>,
|
||||||
pub(crate) distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
|
pub(crate) distributed_request_gate: Option<Arc<RuntimeSemaphore>>,
|
||||||
pub(crate) client: reqwest::Client,
|
pub(crate) client: reqwest::Client,
|
||||||
|
|||||||
@@ -21,7 +21,9 @@ use aether_runtime_state::{
|
|||||||
};
|
};
|
||||||
use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS;
|
use aether_scheduler_core::PROVIDER_KEY_RPM_WINDOW_SECS;
|
||||||
|
|
||||||
use super::{AppState, FrontdoorCorsConfig, LocalExecutionRuntimeMissDiagnostic};
|
use super::{
|
||||||
|
AppState, FrontdoorCorsConfig, FrontdoorRuntimeGuardConfig, LocalExecutionRuntimeMissDiagnostic,
|
||||||
|
};
|
||||||
|
|
||||||
use super::super::async_task::{
|
use super::super::async_task::{
|
||||||
spawn_video_task_poller, VideoTaskPollerConfig, VideoTaskService, VideoTaskTruthSourceMode,
|
spawn_video_task_poller, VideoTaskPollerConfig, VideoTaskService, VideoTaskTruthSourceMode,
|
||||||
@@ -227,6 +229,7 @@ impl AppState {
|
|||||||
VideoTaskTruthSourceMode::PythonSyncReport,
|
VideoTaskTruthSourceMode::PythonSyncReport,
|
||||||
)),
|
)),
|
||||||
video_task_poller: None,
|
video_task_poller: None,
|
||||||
|
frontdoor_runtime_guards: Arc::new(FrontdoorRuntimeGuardConfig::from_env()),
|
||||||
request_gate: None,
|
request_gate: None,
|
||||||
distributed_request_gate: None,
|
distributed_request_gate: None,
|
||||||
client,
|
client,
|
||||||
|
|||||||
@@ -60,12 +60,7 @@ impl provider_transport::VideoTaskTransportSnapshotLookup for AppState {
|
|||||||
) -> Result<Option<GatewayProviderTransportSnapshot>, String> {
|
) -> Result<Option<GatewayProviderTransportSnapshot>, String> {
|
||||||
self.read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
|
self.read_provider_transport_snapshot(provider_id, endpoint_id, key_id)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| match err {
|
.map_err(GatewayError::into_message)
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -77,12 +72,7 @@ impl ModelFetchTransportRuntime for AppState {
|
|||||||
) -> Result<Option<LocalResolvedOAuthRequestAuth>, String> {
|
) -> Result<Option<LocalResolvedOAuthRequestAuth>, String> {
|
||||||
AppState::resolve_local_oauth_request_auth(self, transport)
|
AppState::resolve_local_oauth_request_auth(self, transport)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| match err {
|
.map_err(GatewayError::into_message)
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn resolve_model_fetch_proxy(
|
async fn resolve_model_fetch_proxy(
|
||||||
@@ -99,12 +89,7 @@ impl ModelFetchTransportRuntime for AppState {
|
|||||||
) -> Result<ExecutionResult, String> {
|
) -> Result<ExecutionResult, String> {
|
||||||
execution_runtime::execute_execution_runtime_sync_plan(self, None, plan)
|
execution_runtime::execute_execution_runtime_sync_plan(self, None, plan)
|
||||||
.await
|
.await
|
||||||
.map_err(|err| match err {
|
.map_err(GatewayError::into_message)
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
})
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -27,6 +27,7 @@ pub(crate) use self::admin_types::{
|
|||||||
UserDailyQuotaAvailabilityRecord, UserPlanEntitlementRecord,
|
UserDailyQuotaAvailabilityRecord, UserPlanEntitlementRecord,
|
||||||
};
|
};
|
||||||
pub use self::app::AppState;
|
pub use self::app::AppState;
|
||||||
|
pub(crate) use self::app::FrontdoorRuntimeGuardConfig;
|
||||||
pub(crate) use self::cache::{
|
pub(crate) use self::cache::{
|
||||||
CachedProviderTransportSnapshot, AUTH_API_KEY_LAST_USED_MAX_ENTRIES,
|
CachedProviderTransportSnapshot, AUTH_API_KEY_LAST_USED_MAX_ENTRIES,
|
||||||
AUTH_API_KEY_LAST_USED_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES,
|
AUTH_API_KEY_LAST_USED_TTL, PROVIDER_TRANSPORT_SNAPSHOT_CACHE_MAX_ENTRIES,
|
||||||
|
|||||||
@@ -1448,12 +1448,7 @@ impl AppState {
|
|||||||
.map_err(
|
.map_err(
|
||||||
|err| provider_transport::LocalOAuthRefreshError::InvalidResponse {
|
|err| provider_transport::LocalOAuthRefreshError::InvalidResponse {
|
||||||
provider_type,
|
provider_type,
|
||||||
message: match err {
|
message: err.into_message(),
|
||||||
GatewayError::UpstreamUnavailable { message, .. }
|
|
||||||
| GatewayError::ControlUnavailable { message, .. }
|
|
||||||
| GatewayError::Client { message, .. }
|
|
||||||
| GatewayError::Internal(message) => message,
|
|
||||||
},
|
|
||||||
},
|
},
|
||||||
)?;
|
)?;
|
||||||
let response_body_text = local_oauth_execution_body_text(&result);
|
let response_body_text = local_oauth_execution_body_text(&result);
|
||||||
|
|||||||
@@ -11,7 +11,7 @@ use aether_data_contracts::repository::video_tasks::{
|
|||||||
};
|
};
|
||||||
use serde_json::json;
|
use serde_json::json;
|
||||||
|
|
||||||
use super::{AppState, GatewayDataState};
|
use super::{AppState, FrontdoorRuntimeGuardConfig, GatewayDataState};
|
||||||
use crate::{provider_transport, usage};
|
use crate::{provider_transport, usage};
|
||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
@@ -31,6 +31,14 @@ impl AppState {
|
|||||||
self
|
self
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub(crate) fn with_frontdoor_runtime_guard_config_for_tests(
|
||||||
|
mut self,
|
||||||
|
config: FrontdoorRuntimeGuardConfig,
|
||||||
|
) -> Self {
|
||||||
|
self.frontdoor_runtime_guards = Arc::new(config);
|
||||||
|
self
|
||||||
|
}
|
||||||
|
|
||||||
pub(crate) fn with_tunnel_identity_for_tests(
|
pub(crate) fn with_tunnel_identity_for_tests(
|
||||||
mut self,
|
mut self,
|
||||||
instance_id: &str,
|
instance_id: &str,
|
||||||
|
|||||||
@@ -0,0 +1,21 @@
|
|||||||
|
-- Usage is a historical fact table. Terminal usage events can arrive after
|
||||||
|
-- mutable catalog/auth rows have been disabled or deleted, so these snapshot
|
||||||
|
-- identity columns must not make ingestion depend on current dimension rows.
|
||||||
|
|
||||||
|
ALTER TABLE ONLY public.usage
|
||||||
|
DROP CONSTRAINT IF EXISTS usage_provider_id_fkey;
|
||||||
|
|
||||||
|
ALTER TABLE ONLY public.usage
|
||||||
|
DROP CONSTRAINT IF EXISTS usage_provider_endpoint_id_fkey;
|
||||||
|
|
||||||
|
ALTER TABLE ONLY public.usage
|
||||||
|
DROP CONSTRAINT IF EXISTS usage_provider_api_key_id_fkey;
|
||||||
|
|
||||||
|
ALTER TABLE ONLY public.usage
|
||||||
|
DROP CONSTRAINT IF EXISTS usage_api_key_id_fkey;
|
||||||
|
|
||||||
|
ALTER TABLE ONLY public.usage
|
||||||
|
DROP CONSTRAINT IF EXISTS usage_user_id_fkey;
|
||||||
|
|
||||||
|
ALTER TABLE ONLY public.usage
|
||||||
|
DROP CONSTRAINT IF EXISTS usage_wallet_id_fkey;
|
||||||
@@ -583,80 +583,12 @@ END $mig$;
|
|||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: usage usage_api_key_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
-- Usage is a historical fact table. Snapshot identity columns such as
|
||||||
|
-- user_id/api_key_id/provider_endpoint_id/provider_api_key_id/wallet_id
|
||||||
|
-- intentionally do not carry foreign keys because terminal usage events may
|
||||||
|
-- arrive after those mutable dimension rows have been disabled or deleted.
|
||||||
--
|
--
|
||||||
|
|
||||||
DO $mig$ BEGIN
|
|
||||||
ALTER TABLE ONLY public.usage
|
|
||||||
ADD CONSTRAINT usage_api_key_id_fkey FOREIGN KEY (api_key_id) REFERENCES public.api_keys(id) ON DELETE SET NULL;
|
|
||||||
EXCEPTION
|
|
||||||
WHEN duplicate_object THEN NULL;
|
|
||||||
WHEN duplicate_table THEN NULL;
|
|
||||||
WHEN invalid_table_definition THEN NULL;
|
|
||||||
END $mig$;
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
|
||||||
-- Name: usage usage_provider_api_key_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
|
||||||
--
|
|
||||||
|
|
||||||
DO $mig$ BEGIN
|
|
||||||
ALTER TABLE ONLY public.usage
|
|
||||||
ADD CONSTRAINT usage_provider_api_key_id_fkey FOREIGN KEY (provider_api_key_id) REFERENCES public.provider_api_keys(id) ON DELETE SET NULL;
|
|
||||||
EXCEPTION
|
|
||||||
WHEN duplicate_object THEN NULL;
|
|
||||||
WHEN duplicate_table THEN NULL;
|
|
||||||
WHEN invalid_table_definition THEN NULL;
|
|
||||||
END $mig$;
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
|
||||||
-- Name: usage usage_provider_endpoint_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
|
||||||
--
|
|
||||||
|
|
||||||
DO $mig$ BEGIN
|
|
||||||
ALTER TABLE ONLY public.usage
|
|
||||||
ADD CONSTRAINT usage_provider_endpoint_id_fkey FOREIGN KEY (provider_endpoint_id) REFERENCES public.provider_endpoints(id) ON DELETE SET NULL;
|
|
||||||
EXCEPTION
|
|
||||||
WHEN duplicate_object THEN NULL;
|
|
||||||
WHEN duplicate_table THEN NULL;
|
|
||||||
WHEN invalid_table_definition THEN NULL;
|
|
||||||
END $mig$;
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
|
||||||
-- Name: usage usage_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
|
||||||
--
|
|
||||||
|
|
||||||
DO $mig$ BEGIN
|
|
||||||
ALTER TABLE ONLY public.usage
|
|
||||||
ADD CONSTRAINT usage_user_id_fkey FOREIGN KEY (user_id) REFERENCES public.users(id) ON DELETE SET NULL;
|
|
||||||
EXCEPTION
|
|
||||||
WHEN duplicate_object THEN NULL;
|
|
||||||
WHEN duplicate_table THEN NULL;
|
|
||||||
WHEN invalid_table_definition THEN NULL;
|
|
||||||
END $mig$;
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
|
||||||
-- Name: usage usage_wallet_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
|
||||||
--
|
|
||||||
|
|
||||||
DO $mig$ BEGIN
|
|
||||||
ALTER TABLE ONLY public.usage
|
|
||||||
ADD CONSTRAINT usage_wallet_id_fkey FOREIGN KEY (wallet_id) REFERENCES public.wallets(id) ON DELETE SET NULL;
|
|
||||||
EXCEPTION
|
|
||||||
WHEN duplicate_object THEN NULL;
|
|
||||||
WHEN duplicate_table THEN NULL;
|
|
||||||
WHEN invalid_table_definition THEN NULL;
|
|
||||||
END $mig$;
|
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
--
|
--
|
||||||
-- Name: user_model_usage_counts user_model_usage_counts_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
-- Name: user_model_usage_counts user_model_usage_counts_user_id_fkey; Type: FK CONSTRAINT; Schema: public; Owner: -
|
||||||
--
|
--
|
||||||
|
|||||||
@@ -7,7 +7,7 @@ use tracing::info;
|
|||||||
// Generated by build.rs from schema/bootstrap/postgres.
|
// Generated by build.rs from schema/bootstrap/postgres.
|
||||||
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
|
pub(crate) static EMPTY_DATABASE_SNAPSHOT_SQL: &str =
|
||||||
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
|
include_str!(concat!(env!("OUT_DIR"), "/empty_database_snapshot.sql"));
|
||||||
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260520010000;
|
pub(crate) const EMPTY_DATABASE_SNAPSHOT_CUTOFF_VERSION: i64 = 20260522000000;
|
||||||
|
|
||||||
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
|
const PUBLIC_BASE_TABLE_COUNT_SQL: &str = r#"
|
||||||
SELECT COUNT(*)::BIGINT
|
SELECT COUNT(*)::BIGINT
|
||||||
|
|||||||
@@ -312,6 +312,7 @@ fn empty_database_snapshot_covers_current_cutoff_versions() {
|
|||||||
20260519130000,
|
20260519130000,
|
||||||
20260520000000,
|
20260520000000,
|
||||||
20260520010000,
|
20260520010000,
|
||||||
|
20260522000000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
@@ -393,6 +394,34 @@ fn empty_database_snapshot_sql_includes_usage_body_blobs_and_audit_admin_role()
|
|||||||
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("usage_count bigint DEFAULT 0 NOT NULL"));
|
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("usage_count bigint DEFAULT 0 NOT NULL"));
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn usage_identity_foreign_keys_are_decoupled_for_historical_ingestion() {
|
||||||
|
let migration = POSTGRES_MIGRATOR
|
||||||
|
.iter()
|
||||||
|
.find(|migration| migration.version == 20260522000000)
|
||||||
|
.expect("usage identity foreign key decoupling migration should be embedded");
|
||||||
|
|
||||||
|
for constraint in [
|
||||||
|
"usage_provider_id_fkey",
|
||||||
|
"usage_provider_endpoint_id_fkey",
|
||||||
|
"usage_provider_api_key_id_fkey",
|
||||||
|
"usage_api_key_id_fkey",
|
||||||
|
"usage_user_id_fkey",
|
||||||
|
"usage_wallet_id_fkey",
|
||||||
|
] {
|
||||||
|
assert!(
|
||||||
|
migration
|
||||||
|
.sql
|
||||||
|
.contains(format!("DROP CONSTRAINT IF EXISTS {constraint}").as_str()),
|
||||||
|
"migration should drop {constraint}"
|
||||||
|
);
|
||||||
|
assert!(
|
||||||
|
!EMPTY_DATABASE_SNAPSHOT_SQL.contains(format!("ADD CONSTRAINT {constraint}").as_str()),
|
||||||
|
"fresh bootstrap snapshot should not recreate {constraint}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn empty_database_snapshot_sql_includes_payment_gateway_and_plans() {
|
fn empty_database_snapshot_sql_includes_payment_gateway_and_plans() {
|
||||||
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("payment_provider character varying(64)"));
|
assert!(EMPTY_DATABASE_SNAPSHOT_SQL.contains("payment_provider character varying(64)"));
|
||||||
@@ -1179,6 +1208,7 @@ fn pending_migrations_from_applied_skips_versions_already_applied() {
|
|||||||
20260519130000,
|
20260519130000,
|
||||||
20260520000000,
|
20260520000000,
|
||||||
20260520010000,
|
20260520010000,
|
||||||
|
20260522000000,
|
||||||
]
|
]
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -193,16 +193,83 @@ impl UsageQueueWorker {
|
|||||||
let event = match UsageEvent::from_stream_fields(&entry.fields) {
|
let event = match UsageEvent::from_stream_fields(&entry.fields) {
|
||||||
Ok(event) => event,
|
Ok(event) => event,
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
|
warn!(
|
||||||
|
event_name = "usage_worker_entry_decode_dead_lettered",
|
||||||
|
log_type = "ops",
|
||||||
|
worker_consumer = %self.consumer,
|
||||||
|
worker_group = %self.config.consumer_group,
|
||||||
|
entry_id = %entry.id,
|
||||||
|
error = %err,
|
||||||
|
"usage worker moved malformed queue entry to dead letter"
|
||||||
|
);
|
||||||
self.queue.push_dead_letter(entry, &err.to_string()).await?;
|
self.queue.push_dead_letter(entry, &err.to_string()).await?;
|
||||||
return Ok(true);
|
return Ok(true);
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
self.recorder.record_usage_event(&event).await?;
|
match self.recorder.record_usage_event(&event).await {
|
||||||
Ok(true)
|
Ok(()) => Ok(true),
|
||||||
|
Err(err) if usage_event_record_error_is_permanent(&err) => {
|
||||||
|
warn!(
|
||||||
|
event_name = "usage_worker_entry_record_dead_lettered",
|
||||||
|
log_type = "ops",
|
||||||
|
worker_consumer = %self.consumer,
|
||||||
|
worker_group = %self.config.consumer_group,
|
||||||
|
entry_id = %entry.id,
|
||||||
|
request_id = %event.request_id,
|
||||||
|
event_type = ?event.event_type,
|
||||||
|
provider_name = %event.data.provider_name,
|
||||||
|
model = %event.data.model,
|
||||||
|
api_format = event.data.api_format.as_deref().unwrap_or(""),
|
||||||
|
provider_id = event.data.provider_id.as_deref().unwrap_or(""),
|
||||||
|
provider_endpoint_id = event.data.provider_endpoint_id.as_deref().unwrap_or(""),
|
||||||
|
provider_api_key_id = event.data.provider_api_key_id.as_deref().unwrap_or(""),
|
||||||
|
error = %err,
|
||||||
|
"usage worker moved non-retryable usage event to dead letter"
|
||||||
|
);
|
||||||
|
self.queue.push_dead_letter(entry, &err.to_string()).await?;
|
||||||
|
Ok(true)
|
||||||
|
}
|
||||||
|
Err(err) => {
|
||||||
|
warn!(
|
||||||
|
event_name = "usage_worker_entry_record_retryable_failed",
|
||||||
|
log_type = "ops",
|
||||||
|
worker_consumer = %self.consumer,
|
||||||
|
worker_group = %self.config.consumer_group,
|
||||||
|
entry_id = %entry.id,
|
||||||
|
request_id = %event.request_id,
|
||||||
|
event_type = ?event.event_type,
|
||||||
|
provider_name = %event.data.provider_name,
|
||||||
|
model = %event.data.model,
|
||||||
|
api_format = event.data.api_format.as_deref().unwrap_or(""),
|
||||||
|
provider_id = event.data.provider_id.as_deref().unwrap_or(""),
|
||||||
|
provider_endpoint_id = event.data.provider_endpoint_id.as_deref().unwrap_or(""),
|
||||||
|
provider_api_key_id = event.data.provider_api_key_id.as_deref().unwrap_or(""),
|
||||||
|
error = %err,
|
||||||
|
"usage worker will retry usage event after record failure"
|
||||||
|
);
|
||||||
|
Err(err)
|
||||||
|
}
|
||||||
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
fn usage_event_record_error_is_permanent(err: &DataLayerError) -> bool {
|
||||||
|
match err {
|
||||||
|
DataLayerError::InvalidConfiguration(_)
|
||||||
|
| DataLayerError::InvalidInput(_)
|
||||||
|
| DataLayerError::UnexpectedValue(_) => true,
|
||||||
|
DataLayerError::Postgres(message) | DataLayerError::Sql(message) => {
|
||||||
|
database_error_is_known_permanent(message)
|
||||||
|
}
|
||||||
|
DataLayerError::Redis(_) | DataLayerError::TimedOut(_) => false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
fn database_error_is_known_permanent(message: &str) -> bool {
|
||||||
|
message.contains("SQLSTATE 23503") || message.contains("violates foreign key constraint")
|
||||||
|
}
|
||||||
|
|
||||||
pub fn build_usage_queue_worker<T>(
|
pub fn build_usage_queue_worker<T>(
|
||||||
runner: Arc<dyn RuntimeQueueStore>,
|
runner: Arc<dyn RuntimeQueueStore>,
|
||||||
data: Arc<T>,
|
data: Arc<T>,
|
||||||
@@ -287,16 +354,23 @@ fn consumer_name() -> String {
|
|||||||
|
|
||||||
#[cfg(test)]
|
#[cfg(test)]
|
||||||
mod tests {
|
mod tests {
|
||||||
use std::sync::Mutex;
|
use std::sync::{Arc, Mutex};
|
||||||
|
|
||||||
use aether_data_contracts::repository::settlement::{
|
use aether_data_contracts::repository::settlement::{
|
||||||
StoredUsageSettlement, UsageSettlementInput,
|
StoredUsageSettlement, UsageSettlementInput,
|
||||||
};
|
};
|
||||||
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
|
use aether_data_contracts::repository::usage::{StoredRequestUsageAudit, UpsertUsageRecord};
|
||||||
|
use aether_data_contracts::DataLayerError;
|
||||||
|
use aether_runtime_state::{MemoryRuntimeStateConfig, RuntimeQueueStore, RuntimeState};
|
||||||
use async_trait::async_trait;
|
use async_trait::async_trait;
|
||||||
|
|
||||||
use super::{write_event_record, ManualProxyNodeCounter, UsageRecordWriter};
|
use super::{
|
||||||
use crate::{UsageEvent, UsageEventData, UsageEventType, UsageSettlementWriter};
|
usage_event_record_error_is_permanent, write_event_record, ManualProxyNodeCounter,
|
||||||
|
UsageEventRecorder, UsageQueueWorker, UsageRecordWriter,
|
||||||
|
};
|
||||||
|
use crate::{
|
||||||
|
UsageEvent, UsageEventData, UsageEventType, UsageRuntimeConfig, UsageSettlementWriter,
|
||||||
|
};
|
||||||
|
|
||||||
#[derive(Default)]
|
#[derive(Default)]
|
||||||
struct TestUsageStore {
|
struct TestUsageStore {
|
||||||
@@ -304,6 +378,11 @@ mod tests {
|
|||||||
settlements: Mutex<Vec<UsageSettlementInput>>,
|
settlements: Mutex<Vec<UsageSettlementInput>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[derive(Default)]
|
||||||
|
struct SelectiveFailingRecorder {
|
||||||
|
calls: Mutex<Vec<String>>,
|
||||||
|
}
|
||||||
|
|
||||||
#[async_trait]
|
#[async_trait]
|
||||||
impl UsageRecordWriter for TestUsageStore {
|
impl UsageRecordWriter for TestUsageStore {
|
||||||
async fn upsert_usage_record(
|
async fn upsert_usage_record(
|
||||||
@@ -392,6 +471,22 @@ mod tests {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[async_trait]
|
||||||
|
impl UsageEventRecorder for SelectiveFailingRecorder {
|
||||||
|
async fn record_usage_event(&self, event: &UsageEvent) -> Result<(), DataLayerError> {
|
||||||
|
self.calls
|
||||||
|
.lock()
|
||||||
|
.expect("calls lock")
|
||||||
|
.push(event.request_id.clone());
|
||||||
|
if event.request_id == "req-worker-poison" {
|
||||||
|
return Err(DataLayerError::UnexpectedValue(
|
||||||
|
"permanent test error".to_string(),
|
||||||
|
));
|
||||||
|
}
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn sample_event() -> UsageEvent {
|
fn sample_event() -> UsageEvent {
|
||||||
UsageEvent::new(
|
UsageEvent::new(
|
||||||
UsageEventType::Completed,
|
UsageEventType::Completed,
|
||||||
@@ -436,4 +531,104 @@ mod tests {
|
|||||||
assert_eq!(settlements.len(), 1);
|
assert_eq!(settlements.len(), 1);
|
||||||
assert_eq!(settlements[0].request_id, "req-worker-123");
|
assert_eq!(settlements[0].request_id, "req-worker-123");
|
||||||
}
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn usage_event_record_error_classifies_permanent_failures() {
|
||||||
|
assert!(usage_event_record_error_is_permanent(
|
||||||
|
&DataLayerError::UnexpectedValue("bad payload".to_string())
|
||||||
|
));
|
||||||
|
assert!(usage_event_record_error_is_permanent(
|
||||||
|
&DataLayerError::Postgres(
|
||||||
|
"error returned from database: violates foreign key constraint (SQLSTATE 23503)"
|
||||||
|
.to_string()
|
||||||
|
)
|
||||||
|
));
|
||||||
|
assert!(!usage_event_record_error_is_permanent(
|
||||||
|
&DataLayerError::Redis("connection refused".to_string())
|
||||||
|
));
|
||||||
|
assert!(!usage_event_record_error_is_permanent(
|
||||||
|
&DataLayerError::TimedOut("postgres acquire".to_string())
|
||||||
|
));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn process_entries_dead_letters_permanent_record_error_and_continues() {
|
||||||
|
let runner = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
|
||||||
|
let queue_runner: Arc<dyn RuntimeQueueStore> = runner.clone();
|
||||||
|
let recorder = Arc::new(SelectiveFailingRecorder::default());
|
||||||
|
let config = UsageRuntimeConfig {
|
||||||
|
enabled: true,
|
||||||
|
stream_key: "usage:test:worker:events".to_string(),
|
||||||
|
consumer_group: "usage:test:worker:group".to_string(),
|
||||||
|
dlq_stream_key: "usage:test:worker:dlq".to_string(),
|
||||||
|
consumer_batch_size: 10,
|
||||||
|
consumer_block_ms: 1,
|
||||||
|
..UsageRuntimeConfig::default()
|
||||||
|
};
|
||||||
|
let worker = UsageQueueWorker::new(queue_runner, recorder.clone(), config)
|
||||||
|
.expect("worker should build");
|
||||||
|
worker
|
||||||
|
.queue
|
||||||
|
.ensure_consumer_group()
|
||||||
|
.await
|
||||||
|
.expect("group should initialize");
|
||||||
|
|
||||||
|
let mut poison = sample_event();
|
||||||
|
poison.request_id = "req-worker-poison".to_string();
|
||||||
|
let mut ok = sample_event();
|
||||||
|
ok.request_id = "req-worker-ok".to_string();
|
||||||
|
worker
|
||||||
|
.queue
|
||||||
|
.enqueue(&poison)
|
||||||
|
.await
|
||||||
|
.expect("poison event should enqueue");
|
||||||
|
worker
|
||||||
|
.queue
|
||||||
|
.enqueue(&ok)
|
||||||
|
.await
|
||||||
|
.expect("ok event should enqueue");
|
||||||
|
|
||||||
|
let entries = worker
|
||||||
|
.queue
|
||||||
|
.read_group(&worker.consumer)
|
||||||
|
.await
|
||||||
|
.expect("events should read");
|
||||||
|
assert_eq!(entries.len(), 2);
|
||||||
|
|
||||||
|
worker
|
||||||
|
.process_entries(entries)
|
||||||
|
.await
|
||||||
|
.expect("permanent failure should not block batch");
|
||||||
|
|
||||||
|
assert_eq!(
|
||||||
|
recorder.calls.lock().expect("calls lock").as_slice(),
|
||||||
|
["req-worker-poison", "req-worker-ok"]
|
||||||
|
);
|
||||||
|
|
||||||
|
runner
|
||||||
|
.ensure_consumer_group(
|
||||||
|
"usage:test:worker:dlq",
|
||||||
|
"usage:test:worker:dlq-group",
|
||||||
|
"0-0",
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("dlq group should initialize");
|
||||||
|
let dlq_entries = runner
|
||||||
|
.read_group(
|
||||||
|
"usage:test:worker:dlq",
|
||||||
|
"usage:test:worker:dlq-group",
|
||||||
|
"usage-test-dlq-consumer",
|
||||||
|
10,
|
||||||
|
Some(1),
|
||||||
|
)
|
||||||
|
.await
|
||||||
|
.expect("dlq should read");
|
||||||
|
assert_eq!(dlq_entries.len(), 1);
|
||||||
|
let payload = dlq_entries[0]
|
||||||
|
.fields
|
||||||
|
.get("payload")
|
||||||
|
.expect("dlq payload should exist");
|
||||||
|
assert!(payload.contains("req-worker-poison"));
|
||||||
|
assert!(payload.contains("permanent test error"));
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user