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

Improve gateway scheduling and runtime admission

Browse Source
This commit is contained in:
elky
2026-06-24 01:53:45 +08:00
parent cf0af8fa1e
commit d336d1a7fa
87 changed files with 9671 additions and 804 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/provider_pool_demand.rs
+142 -49
View File
@@ -1,10 +1,11 @@
use std::sync::{
atomic::{AtomicBool, Ordering},
atomic::{AtomicBool, AtomicUsize, Ordering},
Arc,
};
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use aether_runtime_state::RuntimeState;
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use tokio::task::JoinHandle;
use tracing::debug;
@@ -32,14 +33,24 @@ pub(crate) struct ProviderPoolDemandSnapshot {
}
pub(crate) struct ProviderPoolInFlightGuard {
runtime: Arc<RuntimeState>,
tokens_key: String,
token: String,
stop_renewal: Arc<AtomicBool>,
renew_handle: Option<JoinHandle<()>>,
kind: ProviderPoolInFlightGuardKind,
released: bool,
}
enum ProviderPoolInFlightGuardKind {
Local {
provider_id: String,
counter: Arc<AtomicUsize>,
},
Runtime {
runtime: Arc<RuntimeState>,
tokens_key: String,
token: String,
stop_renewal: Arc<AtomicBool>,
renew_handle: Option<JoinHandle<()>>,
},
}
impl ProviderPoolInFlightGuard {
pub(crate) async fn release(mut self) {
self.release_inner().await;
@@ -50,19 +61,29 @@ impl ProviderPoolInFlightGuard {
return;
}
self.released = true;
self.stop_renewal.store(true, Ordering::Release);
if let Some(handle) = self.renew_handle.take() {
handle.abort();
}
if let Err(err) = self
.runtime
.score_remove(&self.tokens_key, &self.token)
.await
{
debug!(
error = ?err,
"gateway provider pool demand: failed to release in-flight token"
);
match &mut self.kind {
ProviderPoolInFlightGuardKind::Local {
provider_id,
counter,
} => decrement_local_provider_in_flight(provider_id, counter),
ProviderPoolInFlightGuardKind::Runtime {
runtime,
tokens_key,
token,
stop_renewal,
renew_handle,
} => {
stop_renewal.store(true, Ordering::Release);
if let Some(handle) = renew_handle.take() {
handle.abort();
}
if let Err(err) = runtime.score_remove(tokens_key, token).await {
debug!(
error = ?err,
"gateway provider pool demand: failed to release in-flight token"
);
}
}
}
}
}
@@ -73,23 +94,37 @@ impl Drop for ProviderPoolInFlightGuard {
return;
}
self.released = true;
self.stop_renewal.store(true, Ordering::Release);
if let Some(handle) = self.renew_handle.take() {
handle.abort();
}
let runtime = self.runtime.clone();
let tokens_key = self.tokens_key.clone();
let token = self.token.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
if let Err(err) = runtime.score_remove(&tokens_key, &token).await {
debug!(
error = ?err,
"gateway provider pool demand: failed to release dropped in-flight token"
);
match &mut self.kind {
ProviderPoolInFlightGuardKind::Local {
provider_id,
counter,
} => decrement_local_provider_in_flight(provider_id, counter),
ProviderPoolInFlightGuardKind::Runtime {
runtime,
tokens_key,
token,
stop_renewal,
renew_handle,
} => {
stop_renewal.store(true, Ordering::Release);
if let Some(handle) = renew_handle.take() {
handle.abort();
}
});
let runtime = runtime.clone();
let tokens_key = tokens_key.clone();
let token = token.clone();
if let Ok(handle) = tokio::runtime::Handle::try_current() {
handle.spawn(async move {
if let Err(err) = runtime.score_remove(&tokens_key, &token).await {
debug!(
error = ?err,
"gateway provider pool demand: failed to release dropped in-flight token"
);
}
});
}
}
}
}
}
@@ -106,6 +141,46 @@ fn in_flight_tokens_key(provider_id: &str) -> String {
format!("{PROVIDER_POOL_IN_FLIGHT_TOKENS_PREFIX}:{provider_id}")
}
fn local_provider_in_flight_counts() -> &'static DashMap<String, Arc<AtomicUsize>> {
static COUNTS: std::sync::OnceLock<DashMap<String, Arc<AtomicUsize>>> =
std::sync::OnceLock::new();
COUNTS.get_or_init(DashMap::new)
}
fn increment_local_provider_in_flight(provider_id: &str) -> Arc<AtomicUsize> {
let counter_ref = local_provider_in_flight_counts()
.entry(provider_id.to_string())
.or_insert_with(|| Arc::new(AtomicUsize::new(0)));
counter_ref.fetch_add(1, Ordering::AcqRel);
counter_ref.clone()
}
fn decrement_local_provider_in_flight(provider_id: &str, counter: &AtomicUsize) {
let mut current = counter.load(Ordering::Acquire);
while current > 0 {
match counter.compare_exchange_weak(
current,
current - 1,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => break,
Err(next) => current = next,
}
}
if counter.load(Ordering::Acquire) == 0 {
let _ = local_provider_in_flight_counts()
.remove_if(provider_id, |_, stored| stored.load(Ordering::Acquire) == 0);
}
}
fn local_provider_live_in_flight_count(provider_id: &str) -> usize {
local_provider_in_flight_counts()
.get(provider_id)
.map(|counter| counter.load(Ordering::Acquire))
.unwrap_or(0)
}
fn demand_snapshot_key(provider_id: &str) -> String {
format!("{PROVIDER_POOL_DEMAND_SNAPSHOT_PREFIX}:{provider_id}")
}
@@ -185,6 +260,17 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard(
return None;
}
if runtime.is_memory() {
let counter = increment_local_provider_in_flight(provider_id);
return Some(ProviderPoolInFlightGuard {
kind: ProviderPoolInFlightGuardKind::Local {
provider_id: provider_id.to_string(),
counter,
},
released: false,
});
}
let tokens_key = in_flight_tokens_key(provider_id);
let token = build_in_flight_token(request_id, candidate_id, key_id);
match tokio::time::timeout(
@@ -221,11 +307,13 @@ pub(crate) async fn acquire_provider_pool_in_flight_guard(
);
Some(ProviderPoolInFlightGuard {
runtime,
tokens_key,
token,
stop_renewal,
renew_handle: Some(renew_handle),
kind: ProviderPoolInFlightGuardKind::Runtime {
runtime,
tokens_key,
token,
stop_renewal,
renew_handle: Some(renew_handle),
},
released: false,
})
}
@@ -238,6 +326,9 @@ pub(crate) async fn provider_pool_live_in_flight_count(
if provider_id.is_empty() {
return 0;
}
if runtime.is_memory() {
return local_provider_live_in_flight_count(provider_id);
}
let key = in_flight_tokens_key(provider_id);
let now_ms = current_unix_ms() as f64;
if let Err(err) = runtime.score_remove_by_score(&key, now_ms).await {
@@ -389,9 +480,10 @@ mod tests {
#[tokio::test]
async fn in_flight_guard_tracks_and_releases_provider_tokens() {
let runtime = Arc::new(RuntimeState::memory(MemoryRuntimeStateConfig::default()));
let provider_id = "provider-guard-release";
let guard = acquire_provider_pool_in_flight_guard(
runtime.clone(),
"provider-1",
provider_id,
"request-1",
Some("candidate-1"),
"key-1",
@@ -400,14 +492,14 @@ mod tests {
.expect("guard should be acquired");
assert_eq!(
provider_pool_live_in_flight_count(runtime.as_ref(), "provider-1").await,
provider_pool_live_in_flight_count(runtime.as_ref(), provider_id).await,
1
);
guard.release().await;
assert_eq!(
provider_pool_live_in_flight_count(runtime.as_ref(), "provider-1").await,
provider_pool_live_in_flight_count(runtime.as_ref(), provider_id).await,
0
);
}
@@ -415,26 +507,27 @@ mod tests {
#[tokio::test]
async fn demand_snapshot_uses_instant_in_flight_for_fast_rise_and_ema_for_fall() {
let runtime = RuntimeState::memory(MemoryRuntimeStateConfig::default());
let provider_id = "provider-demand-snapshot";
let mut guards = Vec::new();
for idx in 0..10 {
let guard = acquire_provider_pool_in_flight_guard(
Arc::new(runtime.clone()),
"provider-1",
provider_id,
"request-1",
Some(&format!("candidate-{idx}")),
"key-1",
)
.await
.expect("guard");
std::mem::forget(guard);
guards.push(guard);
}
let high = sample_provider_pool_demand(&runtime, "provider-1", 100, 50).await;
let high = sample_provider_pool_demand(&runtime, provider_id, 100, 50).await;
assert_eq!(high.in_flight, 10);
assert_eq!(high.desired_hot, 12);
let key = in_flight_tokens_key("provider-1");
let _ = runtime.score_remove_by_score(&key, f64::INFINITY).await;
let low = sample_provider_pool_demand(&runtime, "provider-1", 100, 50).await;
drop(guards);
let low = sample_provider_pool_demand(&runtime, provider_id, 100, 50).await;
assert_eq!(low.in_flight, 0);
assert!(low.ema_in_flight > 0.0);
assert!(low.desired_hot >= PROVIDER_POOL_DEMAND_FLOOR);