Improve gateway scheduling and runtime admission

2026-10-10 11:19:50 +08:00 · 2026-06-24 01:53:45 +08:00
parent cf0af8fa1e
commit d336d1a7fa
87 changed files with 9671 additions and 804 deletions
@@ -0,0 +1,269 @@
+use std::sync::Arc;
+use std::time::{Duration, Instant};
+
+use aether_contracts::ExecutionPlan;
+use aether_runtime::{ConcurrencyGate, ConcurrencyPermit, MetricKind, MetricLabel, MetricSample};
+use dashmap::DashMap;
+use tokio::time::timeout;
+use url::Url;
+
+use crate::stage_metrics::observe_gateway_stage_ms;
+use crate::GatewayError;
+
+const GATE_NAME: &str = "gateway_upstream_target";
+const DEFAULT_METRIC_TARGET_LIMIT: usize = 32;
+const METRIC_TARGET_LIMIT_ENV: &str = "AETHER_GATEWAY_UPSTREAM_TARGET_GATE_METRIC_LIMIT";
+
+#[derive(Debug)]
+pub(crate) struct UpstreamTargetAdmission {
+    limit: Option<usize>,
+    queue_budget: Duration,
+    gates: DashMap<String, Arc<ConcurrencyGate>>,
+}
+
+#[derive(Debug)]
+pub(crate) struct UpstreamTargetAdmissionPermit {
+    _permit: ConcurrencyPermit,
+}
+
+impl UpstreamTargetAdmission {
+    pub(crate) fn new(limit: Option<usize>, queue_budget: Duration) -> Self {
+        Self {
+            limit,
+            queue_budget,
+            gates: DashMap::new(),
+        }
+    }
+
+    pub(crate) async fn acquire(
+        &self,
+        plan: &ExecutionPlan,
+        trace_id: &str,
+    ) -> Result<Option<UpstreamTargetAdmissionPermit>, GatewayError> {
+        let Some(limit) = self.limit else {
+            return Ok(None);
+        };
+        let key = upstream_target_key(plan);
+        let gate = self
+            .gates
+            .entry(key.clone())
+            .or_insert_with(|| Arc::new(ConcurrencyGate::new(GATE_NAME, limit)))
+            .clone();
+        let started_at = Instant::now();
+        let permit = match timeout(self.queue_budget, gate.acquire()).await {
+            Ok(Ok(permit)) => permit,
+            Ok(Err(err)) => return Err(GatewayError::Internal(err.to_string())),
+            Err(_) => {
+                tracing::debug!(
+                    event_name = "gateway_upstream_target_admission_timeout",
+                    log_type = "ops",
+                    trace_id,
+                    target = key.as_str(),
+                    limit,
+                    queue_budget_ms = self.queue_budget.as_millis() as u64,
+                    "gateway upstream target admission gate timed out"
+                );
+                return Err(GatewayError::AdmissionTimeout {
+                    trace_id: trace_id.to_string(),
+                    gate: GATE_NAME,
+                    queue_budget_ms: self.queue_budget.as_millis() as u64,
+                });
+            }
+        };
+        observe_gateway_stage_ms(
+            "stream_upstream_target_admission",
+            started_at.elapsed().as_millis() as u64,
+        );
+        Ok(Some(UpstreamTargetAdmissionPermit { _permit: permit }))
+    }
+
+    pub(crate) fn metric_samples(&self) -> Vec<MetricSample> {
+        let mut samples = vec![MetricSample::new(
+            "upstream_target_gate_active_targets",
+            "Number of upstream targets currently tracked by the gateway upstream target admission gates.",
+            MetricKind::Gauge,
+            self.gates.len() as u64,
+        )];
+
+        let Some(limit) = self.limit else {
+            return samples;
+        };
+
+        samples.push(MetricSample::new(
+            "upstream_target_gate_limit",
+            "Configured per-upstream-target admission gate limit.",
+            MetricKind::Gauge,
+            limit as u64,
+        ));
+
+        let mut snapshots = self
+            .gates
+            .iter()
+            .map(|entry| {
+                let snapshot = entry.value().snapshot();
+                (
+                    entry.key().clone(),
+                    snapshot.in_flight,
+                    snapshot.available_permits,
+                    snapshot.high_watermark,
+                    snapshot.rejected,
+                )
+            })
+            .collect::<Vec<_>>();
+        snapshots.sort_by(|left, right| {
+            right
+                .1
+                .cmp(&left.1)
+                .then_with(|| right.3.cmp(&left.3))
+                .then_with(|| right.4.cmp(&left.4))
+        });
+
+        let metric_target_limit = upstream_target_metric_limit();
+        for (target, in_flight, available, high_watermark, rejected) in
+            snapshots.into_iter().take(metric_target_limit)
+        {
+            let labels = vec![MetricLabel::new("target", target)];
+            samples.push(
+                MetricSample::new(
+                    "upstream_target_gate_in_flight",
+                    "Current number of in-flight operations for an upstream target admission gate.",
+                    MetricKind::Gauge,
+                    in_flight as u64,
+                )
+                .with_labels(labels.clone()),
+            );
+            samples.push(
+                MetricSample::new(
+                    "upstream_target_gate_available_permits",
+                    "Currently available permits for an upstream target admission gate.",
+                    MetricKind::Gauge,
+                    available as u64,
+                )
+                .with_labels(labels.clone()),
+            );
+            samples.push(
+                MetricSample::new(
+                    "upstream_target_gate_high_watermark",
+                    "Highest observed in-flight count for an upstream target admission gate.",
+                    MetricKind::Gauge,
+                    high_watermark as u64,
+                )
+                .with_labels(labels.clone()),
+            );
+            samples.push(
+                MetricSample::new(
+                    "upstream_target_gate_rejected_total",
+                    "Number of operations rejected by an upstream target admission gate.",
+                    MetricKind::Counter,
+                    rejected,
+                )
+                .with_labels(labels),
+            );
+        }
+
+        samples
+    }
+}
+
+pub(crate) fn upstream_target_key(plan: &ExecutionPlan) -> String {
+    let parsed = Url::parse(plan.url.as_str()).ok();
+    let Some(url) = parsed else {
+        return fallback_target_key(plan);
+    };
+    let scheme = url.scheme().to_ascii_lowercase();
+    let Some(host) = url.host_str().map(|host| host.to_ascii_lowercase()) else {
+        return fallback_target_key(plan);
+    };
+    let port = url
+        .port_or_known_default()
+        .map(|port| port.to_string())
+        .unwrap_or_else(|| "-".to_string());
+    let proxy = plan
+        .proxy
+        .as_ref()
+        .and_then(|proxy| proxy.url.as_deref())
+        .map(str::trim)
+        .filter(|value| !value.is_empty())
+        .unwrap_or("-");
+    format!("{scheme}://{host}:{port}|proxy={proxy}")
+}
+
+fn fallback_target_key(plan: &ExecutionPlan) -> String {
+    format!(
+        "unparsed|provider={}|endpoint={}|url={}",
+        plan.provider_id, plan.endpoint_id, plan.url
+    )
+}
+
+fn upstream_target_metric_limit() -> usize {
+    std::env::var(METRIC_TARGET_LIMIT_ENV)
+        .ok()
+        .and_then(|value| value.trim().parse::<usize>().ok())
+        .unwrap_or(DEFAULT_METRIC_TARGET_LIMIT)
+}
+
+#[cfg(test)]
+mod tests {
+    use std::time::Duration;
+
+    use aether_contracts::{ExecutionPlan, RequestBody};
+    use serde_json::json;
+
+    use super::*;
+
+    fn test_plan(url: &str) -> ExecutionPlan {
+        ExecutionPlan {
+            request_id: "req-upstream-target".to_string(),
+            candidate_id: Some("cand-upstream-target".to_string()),
+            provider_name: Some("provider".to_string()),
+            provider_id: "provider_id".to_string(),
+            endpoint_id: "endpoint_id".to_string(),
+            key_id: "key_id".to_string(),
+            method: "POST".to_string(),
+            url: url.to_string(),
+            headers: Default::default(),
+            content_type: Some("application/json".to_string()),
+            content_encoding: None,
+            body: RequestBody::from_json(json!({"stream": true})),
+            stream: true,
+            client_api_format: "openai".to_string(),
+            provider_api_format: "openai".to_string(),
+            model_name: Some("model".to_string()),
+            proxy: None,
+            transport_profile: None,
+            timeouts: None,
+        }
+    }
+
+    #[test]
+    fn upstream_target_key_ignores_path_and_query() {
+        let left = test_plan("http://127.0.0.1:18181/v1/chat/completions?x=1");
+        let right = test_plan("http://127.0.0.1:18181/v1/responses");
+
+        assert_eq!(upstream_target_key(&left), upstream_target_key(&right));
+    }
+
+    #[tokio::test]
+    async fn acquire_times_out_when_target_gate_is_saturated() {
+        let admission = UpstreamTargetAdmission::new(Some(1), Duration::from_millis(1));
+        let plan = test_plan("http://127.0.0.1:18181/v1/chat/completions");
+        let _first = admission
+            .acquire(&plan, "trace-upstream-target")
+            .await
+            .expect("first acquire should succeed")
+            .expect("gate enabled");
+
+        let err = admission
+            .acquire(&plan, "trace-upstream-target")
+            .await
+            .expect_err("second acquire should time out");
+
+        assert!(matches!(
+            err,
+            GatewayError::AdmissionTimeout {
+                gate: "gateway_upstream_target",
+                ..
+            }
+        ));
+    }
+}