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

refactor: lazy pool key scheduling

Browse Source
This commit is contained in:
fawney19
2026-05-11 00:12:05 +08:00
parent 1a0f1a7b72
commit bacb14e5f0
26 changed files with 1856 additions and 223 deletions
Download Patch File Download Diff File Expand all files Collapse all files

74
apps/aether-gateway/src/orchestration/attempt.rs
View File

@@ -1,3 +1,4 @@
use aether_runtime_state::RuntimeLockLease;
use aether_scheduler_core::parse_request_candidate_report_context;
use serde_json::Value;
@@ -29,8 +30,14 @@ impl ExecutionAttemptIdentity {
pub(crate) struct LocalExecutionCandidateMetadata {
pub(crate) candidate_group_id: Option<String>,
pub(crate) pool_key_index: Option<u32>,
pub(crate) pool_key_lease: Option<RuntimeLockLease>,
}
pub(crate) const POOL_KEY_LEASE_KEY_REPORT_FIELD: &str = "pool_key_lease_key";
pub(crate) const POOL_KEY_LEASE_OWNER_REPORT_FIELD: &str = "pool_key_lease_owner";
pub(crate) const POOL_KEY_LEASE_TOKEN_REPORT_FIELD: &str = "pool_key_lease_token";
pub(crate) const POOL_KEY_LEASE_TTL_MS_REPORT_FIELD: &str = "pool_key_lease_ttl_ms";
pub(crate) fn attempt_identity_from_report_context(
report_context: Option<&Value>,
) -> Option<ExecutionAttemptIdentity> {
@@ -57,9 +64,65 @@ pub(crate) fn local_execution_candidate_metadata_from_report_context(
.and_then(|value| value.get("pool_key_index"))
.and_then(Value::as_u64)
.and_then(|value| u32::try_from(value).ok()),
pool_key_lease: pool_key_lease_from_report_context(report_context),
}
}
pub(crate) fn insert_pool_key_lease_report_context_fields(
extra_fields: &mut serde_json::Map<String, Value>,
lease: Option<&RuntimeLockLease>,
) {
let Some(lease) = lease else {
return;
};
extra_fields.insert(
POOL_KEY_LEASE_KEY_REPORT_FIELD.to_string(),
Value::String(lease.key.clone()),
);
extra_fields.insert(
POOL_KEY_LEASE_OWNER_REPORT_FIELD.to_string(),
Value::String(lease.owner.clone()),
);
extra_fields.insert(
POOL_KEY_LEASE_TOKEN_REPORT_FIELD.to_string(),
Value::String(lease.token.clone()),
);
extra_fields.insert(
POOL_KEY_LEASE_TTL_MS_REPORT_FIELD.to_string(),
Value::Number(lease.ttl_ms.into()),
);
}
fn pool_key_lease_from_report_context(report_context: Option<&Value>) -> Option<RuntimeLockLease> {
let report_context = report_context?;
let key = report_context
.get(POOL_KEY_LEASE_KEY_REPORT_FIELD)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let owner = report_context
.get(POOL_KEY_LEASE_OWNER_REPORT_FIELD)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let token = report_context
.get(POOL_KEY_LEASE_TOKEN_REPORT_FIELD)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())?;
let ttl_ms = report_context
.get(POOL_KEY_LEASE_TTL_MS_REPORT_FIELD)
.and_then(Value::as_u64)
.filter(|value| *value > 0)?;
Some(RuntimeLockLease {
key: key.to_string(),
owner: owner.to_string(),
token: token.to_string(),
ttl_ms,
})
}
pub(crate) fn build_local_attempt_identities(
candidate_index: u32,
transport: &GatewayProviderTransportSnapshot,
@@ -127,6 +190,7 @@ mod tests {
GatewayProviderTransportEndpoint, GatewayProviderTransportKey,
GatewayProviderTransportProvider, GatewayProviderTransportSnapshot,
};
use aether_runtime_state::RuntimeLockLease;
fn sample_transport(
provider_max_retries: Option<i32>,
@@ -412,6 +476,10 @@ mod tests {
let metadata = local_execution_candidate_metadata_from_report_context(Some(&json!({
"candidate_group_id": "group-1",
"pool_key_index": 3,
"pool_key_lease_key": "ap:provider-1:lease:key-1",
"pool_key_lease_owner": "gateway-1",
"pool_key_lease_token": "gateway-1:token-1",
"pool_key_lease_ttl_ms": 900000,
})));
assert_eq!(
@@ -419,6 +487,12 @@ mod tests {
LocalExecutionCandidateMetadata {
candidate_group_id: Some("group-1".to_string()),
pool_key_index: Some(3),
pool_key_lease: Some(RuntimeLockLease {
key: "ap:provider-1:lease:key-1".to_string(),
owner: "gateway-1".to_string(),
token: "gateway-1:token-1".to_string(),
ttl_ms: 900000,
}),
}
);
}

24
apps/aether-gateway/src/orchestration/effects.rs
View File

@@ -27,8 +27,9 @@ use crate::handlers::shared::provider_pool::admin_provider_pool_config_from_conf
use crate::handlers::shared::provider_pool::{
admin_provider_pool_key_circuit_breaker_reason, record_admin_provider_pool_error,
record_admin_provider_pool_stream_timeout, record_admin_provider_pool_success,
AdminProviderPoolConfig,
release_admin_provider_pool_key_lease, AdminProviderPoolConfig,
};
use crate::orchestration::local_execution_candidate_metadata_from_report_context;
use crate::scheduler::affinity::SCHEDULER_AFFINITY_TTL;
use crate::AppState;
@@ -129,19 +130,40 @@ pub(crate) async fn apply_local_execution_effect(
}
LocalExecutionEffect::PoolSuccessSync { payload } => {
record_sync_pool_success_effect(state, context, payload).await;
release_pool_key_lease_effect(state, context).await;
}
LocalExecutionEffect::PoolSuccessStream { payload } => {
record_stream_pool_success_effect(state, context, payload).await;
release_pool_key_lease_effect(state, context).await;
}
LocalExecutionEffect::PoolError(effect) => {
record_pool_error_effect(state, context, effect).await;
release_pool_key_lease_effect(state, context).await;
}
LocalExecutionEffect::PoolStreamTimeout => {
record_pool_stream_timeout_effect(state, context).await;
release_pool_key_lease_effect(state, context).await;
}
}
}
async fn release_pool_key_lease_effect(state: &AppState, context: LocalExecutionEffectContext<'_>) {
let metadata = local_execution_candidate_metadata_from_report_context(context.report_context);
let Some(lease) = metadata.pool_key_lease else {
return;
};
if let Err(err) =
release_admin_provider_pool_key_lease(state.runtime_state.as_ref(), &lease).await
{
warn!(
error = ?err,
provider_id = %context.plan.provider_id,
key_id = %context.plan.key_id,
"gateway orchestration effects: failed to release pool key lease"
);
}
}
fn report_context_string_field<'a>(
report_context: Option<&'a Value>,
field: &str,

3
apps/aether-gateway/src/orchestration/mod.rs
View File

@@ -17,7 +17,8 @@ pub(crate) use self::adaptive::{
LocalAdaptiveRateLimitProjection, LocalAdaptiveSuccessProjection,
};
pub(crate) use self::attempt::{
attempt_identity_from_report_context, build_local_attempt_identities, local_attempt_slot_count,
attempt_identity_from_report_context, build_local_attempt_identities,
insert_pool_key_lease_report_context_fields, local_attempt_slot_count,
local_execution_candidate_metadata_from_report_context, ExecutionAttemptIdentity,
LocalExecutionCandidateMetadata,
};