fix(stargate): settle reservations after backend RTT and update

docs/observability/metrics/llm-request-router/metrics.md

@@ -38,3 +38,5 @@ unbounded request fields as metric labels.
 | `stargate_proxy_duration_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model`, `inference_server_id` | Time to first byte from upstream in seconds. |
 | `stargate_routing_duration_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model` | Load-balancer decision time in seconds. |
 | `stargate_active_inference_servers` | Gauge | `llm-request-router:9090/metrics` | `routing_key`, `model` | Currently routable inference servers for a routing target. |
+| `stargate_routing_reservations_active` | Gauge | `llm-request-router:9090/metrics` | `routing_key`, `model` | Local routing reservations not yet released or pruned. A reservation counts until its backend sends a registration update received at or after the soft TTL; time alone never ends it. Release after a queue estimate mismatch and removing or deregistering the backend also end it, so a backend that stops reporting keeps this gauge raised until it reports or deregisters. |
+| `stargate_routing_reservation_ttl_seconds` | Histogram | `llm-request-router:9090/metrics` | `routing_key`, `model` | RTT-based soft reservation TTL in seconds after applying the configured bounds. After it, the next registration update from the backend ends the reservation. |

src/libraries/rust/stargate/crates/stargate/src/control_plane/registration.rs

@@ -15,7 +15,7 @@
 
 use std::collections::BTreeSet;
 use std::sync::Arc;
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use futures::{Stream, StreamExt};
 use tokio_util::sync::CancellationToken;
@@ -105,7 +105,7 @@ pub(super) async fn process_registration_stream(
     stop: CancellationToken,
 ) {
     let first_update = match next_registration_update(&mut stream, idle_timeout, &stop).await {
-        Ok(Some(update)) => update,
+        Ok(Some(received)) => received,
         Ok(None) => return,
         Err(error) => {
             warn!(
@@ -117,7 +117,7 @@ pub(super) async fn process_registration_stream(
         }
     };
     let session = match RegistrationSession::start(
-        &first_update,
+        &first_update.update,
         state,
         connection,
         auth_result.routing_key.as_deref(),
@@ -138,7 +138,7 @@ async fn next_registration_update(
     stream: &mut (impl Stream<Item = Result<InferenceServerRegistration, Status>> + Unpin),
     idle_timeout: Option<Duration>,
     stop: &CancellationToken,
-) -> Result<Option<InferenceServerRegistration>, RegistrationError> {
+) -> Result<Option<ReceivedRegistrationUpdate>, RegistrationError> {
     let next = if let Some(idle_timeout) = idle_timeout {
         // Once shutdown begins, cleanup must win over another ready stream update.
         tokio::select! {
@@ -162,13 +162,25 @@ async fn next_registration_update(
         Some(Err(status)) => return Err(RegistrationError::StreamRead(status)),
         Some(Ok(update)) => update,
     };
+    // Stamp receipt before any connection or RTT await: routing reservations
+    // compare it with their soft expiry to decide whether this update can
+    // include their request.
+    let received_at = Instant::now();
     debug!(
         inference_server_id = %update.inference_server_id,
         cluster_id = %update.cluster_id,
         model_ids = ?sorted_model_ids(&update),
         "received inference servers update"
     );
-    Ok(Some(update))
+    Ok(Some(ReceivedRegistrationUpdate {
+        update,
+        received_at,
+    }))
+}
+
+struct ReceivedRegistrationUpdate {
+    update: InferenceServerRegistration,
+    received_at: Instant,
 }
 
 fn log_deregistration(
@@ -288,14 +300,14 @@ impl RegistrationSession {
 
     async fn run(
         mut self,
-        mut update: InferenceServerRegistration,
+        mut received: ReceivedRegistrationUpdate,
         stream: &mut (impl Stream<Item = Result<InferenceServerRegistration, Status>> + Unpin),
         responses: &flume::Sender<Result<InferenceServerAck, Status>>,
         idle_timeout: Option<Duration>,
         stop: &CancellationToken,
     ) {
         let outcome = loop {
-            match self.apply_update(&update, stop).await {
+            match self.apply_update(&received, stop).await {
                 Ok(ApplyUpdateOutcome::Ack(ack)) => {
                     if !send_registration_response(responses, stop, Ok(ack)).await {
                         break if stop.is_cancelled() {
@@ -313,7 +325,7 @@ impl RegistrationSession {
                 }
             }
             match next_registration_update(stream, idle_timeout, stop).await {
-                Ok(Some(next_update)) => update = next_update,
+                Ok(Some(next_update)) => received = next_update,
                 Ok(None) => break Ok(()),
                 Err(error) => break Err(error),
             }
@@ -323,9 +335,13 @@ impl RegistrationSession {
 
     async fn apply_update(
         &mut self,
-        update: &InferenceServerRegistration,
+        received: &ReceivedRegistrationUpdate,
         stop: &CancellationToken,
     ) -> Result<ApplyUpdateOutcome, Status> {
+        let ReceivedRegistrationUpdate {
+            update,
+            received_at,
+        } = received;
         if let Some(status) = validate_running_update(self.registration.identity(), update) {
             return Err(status);
         }
@@ -341,7 +357,13 @@ impl RegistrationSession {
                 // Keep the no-ack retry behavior while clearing any stale route
                 // that depended on the now-unavailable direct connection.
                 self.state
-                    .apply_registration_update(&self.registration, update, false, None)
+                    .apply_received_registration_update(
+                        &self.registration,
+                        update,
+                        false,
+                        None,
+                        *received_at,
+                    )
                     .await;
                 return Ok(ApplyUpdateOutcome::Skip);
             }
@@ -356,7 +378,13 @@ impl RegistrationSession {
         }
 
         self.state
-            .apply_registration_update(&self.registration, update, reverse_connected, rtt)
+            .apply_received_registration_update(
+                &self.registration,
+                update,
+                reverse_connected,
+                rtt,
+                *received_at,
+            )
             .await;
 
         Ok(ApplyUpdateOutcome::Ack(build_registration_ack(
@@ -639,7 +667,13 @@ mod tests {
         assert_eq!(state.candidates_for_target(&target).await.len(), 1);
 
         let outcome = session
-            .apply_update(&update, &CancellationToken::new())
+            .apply_update(
+                &ReceivedRegistrationUpdate {
+                    update,
+                    received_at: Instant::now(),
+                },
+                &CancellationToken::new(),
+            )
             .await
             .expect("unavailable direct update should be handled");
 

src/libraries/rust/stargate/crates/stargate/src/http_proxy.rs

@@ -33,6 +33,8 @@ use crate::tunnel::{QuicHttpProxy, QuicTunnelConfig};
 mod attempt;
 mod diagnostics;
 mod request;
+#[cfg(test)]
+mod reservation_tests;
 mod retry;
 mod routing;
 mod run;
@@ -120,6 +122,8 @@ pub struct ProxyAppState {
     pub readiness: ReadinessState,
     pub lb_router: Arc<LoadBalancerRouter>,
     pub metrics: Arc<StargateMetrics>,
+    pub routing_reservation_ttl_min: std::time::Duration,
+    pub routing_reservation_ttl_max: std::time::Duration,
     pub retry: ProxyRetryConfig,
     pub debug_config: DebugConfig,
 }
@@ -247,7 +251,9 @@ mod test_support {
     use crate::routing_state::StargateState;
     use crate::tunnel::{QuicHttpProxy, QuicTunnelConfig};
 
-    use super::{DebugConfig, ProxyAppState, ProxyRetryConfig, ProxyTrafficState, ReadinessState, readyz};
+    use super::{
+        DebugConfig, ProxyAppState, ProxyRetryConfig, ProxyTrafficState, ReadinessState, readyz,
+    };
 
     pub(super) fn test_proxy_app_state() -> ProxyAppState {
         test_proxy_app_state_with_lb_config(LoadBalancerConfig::default())
@@ -286,6 +292,8 @@ mod test_support {
                     .expect("load balancer should initialize"),
             ),
             metrics,
+            routing_reservation_ttl_min: std::time::Duration::from_millis(1),
+            routing_reservation_ttl_max: std::time::Duration::from_millis(1000),
             retry: ProxyRetryConfig::default(),
             debug_config: DebugConfig::default(),
         }
@@ -308,7 +316,10 @@ mod test_support {
         let ready_token = CancellationToken::new();
         app.readiness = ReadinessState::warming_up(ready_token.clone());
 
-        assert_eq!(readyz(State(app.clone())).await, StatusCode::SERVICE_UNAVAILABLE);
+        assert_eq!(
+            readyz(State(app.clone())).await,
+            StatusCode::SERVICE_UNAVAILABLE
+        );
 
         ready_token.cancel();