From 671630abf1b783cafcfc9277f7afdc3f56abc338 Mon Sep 17 00:00:00 2001 From: Kan Wu Date: Sun, 20 Sep 2026 01:55:34 -0700 Subject: [PATCH] [sgl-router] refactor - main startup logic (#39861) Co-authored-by: Cursor --- experimental/sgl-router/src/main.rs | 630 +++++++++++++--------------- 1 file changed, 296 insertions(+), 334 deletions(-) diff --git a/experimental/sgl-router/src/main.rs b/experimental/sgl-router/src/main.rs index 91ce23a73..71e6ecdab 100644 --- a/experimental/sgl-router/src/main.rs +++ b/experimental/sgl-router/src/main.rs @@ -3,21 +3,125 @@ use anyhow::{Context, Result}; use clap::Parser; -use sgl_router::config::{CachePrefixProvider, Cli, LogFormat, PolicyKind}; -use std::sync::Arc; -use tokio::signal::unix::{signal, Signal, SignalKind}; +use sgl_kv_indexer::{GrpcPrefixIndex, PrefixIndex, PrefixIndexConfig}; +use sgl_router::{ + config::{CachePrefixProvider, Cli, Config, KvIndexerEndpointConfig, LogFormat, PolicyKind}, + discovery::spawn_discovery, + policies::{ + active_load::{spawn_janitor, ActiveLoadRegistry, JanitorHandle, SystemTimeClock}, + factory::build_registry as build_policy_registry, + kv_events::{BlockSizeOracle, KvEventIndex}, + prefix_provider::RadixTreePrefixProvider, + PolicyRegistry, + }, + proxy::Proxy, + server::{app::build_router, app_context::AppContext, shutdown::drain_for_termination}, + tokenizer::TokenizerRegistry, + workers::{manager, WorkerRegistry}, +}; +use std::{ + sync::Arc, + time::{Duration, Instant}, +}; +use tokio::{ + net::TcpListener, + signal::unix::{signal, Signal, SignalKind}, + sync::{oneshot, watch}, + task::JoinHandle, +}; -/// Install the global tracing subscriber. -/// -/// Idempotent: a second call returns `Ok` without panicking. When -/// `try_init` errors, some other code has already installed a subscriber, -/// so the `tracing::debug!` below is delivered through THAT subscriber — -/// no recursive init. -/// -/// `format` selects the output shape: `Json` emits one JSON record per -/// line (target for production / k8s log aggregators), `Text` is the -/// human-readable default. The `RUST_LOG` environment variable always -/// wins over `default_level`. +const DRAIN_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5); +/// Heartbeat escalates INFO -> WARN here: earlier is a routine rollout draining +/// a long response; later the pod risks SIGKILL with work still open. +const DRAIN_WARN_AFTER: Duration = Duration::from_secs(30); + +// Main components started by this binary: +// - Engine monitor (`KvEventIndex`): receives load statistics over ZMQ and, without a +// remote KV indexer, KV events to maintain a local radix tree. Remote prefix lookups +// use `GrpcPrefixIndex` over gRPC. +// - Engine discovery (`spawn_discovery`): watches Kubernetes pods or loads static URLs, +// sending `DiscoveryEvent`s to `manager::run_with_config` to update `WorkerRegistry`. +// - HTTP server (`axum::serve`): serves OpenAI APIs, health/readiness, and metrics +// through routes built by `build_router`, sharing state via `AppContext`. +#[tokio::main] +async fn main() -> Result<()> { + // Resolve CLI configuration and set up startup logging. + let cli = Cli::parse(); + install_bootstrap_subscriber(); + let config = cli + .into_config() + .context("resolve configuration from CLI flags")?; + init_tracing( + &config.observability.log_level, + config.observability.log_format, + )?; + + // Buffer termination signals before tokenizer loading or discovery can block startup. + let (sigterm, sigint) = install_signal_handlers()?; + log_startup(&config); + + // Load tokenizers used to prepare requests for routing. + let tokenizers = + Arc::new(TokenizerRegistry::load_from_config(&config).context("load tokenizers")?); + + // (Optional) Create a gRPC client only when routing uses an external KV indexer. + let external_kv_indexer_client = create_external_kv_indexer_client(&config)?; + + // Monitor engine-reported KV-cache events and load statistics for routing. + let engine_state = start_engine_state_monitor(external_kv_indexer_client.is_some()); + + // Build the policies that choose which workers receive each request. + let routing_policies = Arc::new( + build_policy_registry( + &config, + engine_state.tree(), + engine_state.block_size_oracle(), + ) + .context("build policy registry")?, + ); + + // Track this router's local view of in-flight requests. + let (local_inflight_requests, inflight_cleanup) = start_local_inflight_tracker(&config); + + // Discovery feeds worker changes to the manager, which maintains this routing catalog. + let worker_registry = Arc::new(WorkerRegistry::default()); + let (discovery_handle, worker_manager_handle) = start_worker_discovery_and_manager( + &config, + &worker_registry, + &engine_state, + &local_inflight_requests, + ) + .await?; + + // Share routing dependencies with HTTP handlers and mark startup complete. + let app_context = build_app_context( + &config, + tokenizers, + worker_registry, + routing_policies, + local_inflight_requests, + &engine_state, + external_kv_indexer_client, + )?; + app_context.mark_ready(); + + // Serve HTTP requests until shutdown, allowing in-flight requests to finish. + let listen_addr = format!("{}:{}", config.server.host, config.server.port); + let listener = TcpListener::bind(&listen_addr) + .await + .with_context(|| format!("bind {listen_addr}"))?; + tracing::info!("listening on {listen_addr}"); + let outcome = serve(listener, app_context, sigterm, sigint).await; + + // Stop background tasks once the HTTP server has finished draining. + discovery_handle.abort(); + worker_manager_handle.abort(); + inflight_cleanup.shutdown().await; + log_shutdown(&outcome.result, outcome.inflight_drain_secs); + outcome.result +} + +// Respect RUST_LOG and tolerate an already-installed subscriber. fn init_tracing(default_level: &str, format: LogFormat) -> Result<()> { let filter = tracing_subscriber::EnvFilter::try_from_default_env() .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new(default_level)); @@ -33,9 +137,6 @@ fn init_tracing(default_level: &str, format: LogFormat) -> Result<()> { .try_init(), }; if let Err(e) = install_result { - // A second install attempt; the existing subscriber is fine. - // Surface the attempted default level so an operator can see - // what we tried. tracing::debug!( default_level = %default_level, ?format, @@ -46,13 +147,7 @@ fn init_tracing(default_level: &str, format: LogFormat) -> Result<()> { Ok(()) } -/// Install a minimal text-format subscriber BEFORE config resolution so a -/// config-resolution error has somewhere to surface. The real subscriber -/// (driven by `Config.observability`) is installed after; the second -/// `try_init` is a no-op because a subscriber is already present. -/// The bootstrap subscriber respects `RUST_LOG` so an operator can -/// debug startup with `RUST_LOG=debug` even when configuration resolution -/// fails. +// Provide startup logging before configuration resolution; later installs are no-ops. fn install_bootstrap_subscriber() { let filter = tracing_subscriber::EnvFilter::try_from_default_env() .unwrap_or_else(|_| tracing_subscriber::EnvFilter::new("info")); @@ -62,59 +157,16 @@ fn install_bootstrap_subscriber() { .try_init(); } -/// How often to report progress while axum drains in-flight requests. That -/// phase is unbounded, so without a heartbeat a pod SIGKILLed at -/// `terminationGracePeriodSeconds` leaves no evidence of what it was waiting on. -const DRAIN_HEARTBEAT_INTERVAL: std::time::Duration = std::time::Duration::from_secs(5); - -/// How long the in-flight drain may run before the heartbeat escalates from -/// INFO to WARN. Under this, a pod finishing a long streaming completion is -/// doing exactly what the drain is for, and logging it at WARN would fire on -/// every routine rollout — training operators to filter router WARNs, which -/// are also where `further termination signal ignored` and the drain advisory -/// land. Past it the pod is at real risk of being SIGKILLed with work open. -const DRAIN_WARN_AFTER: std::time::Duration = std::time::Duration::from_secs(30); - -/// Install SIGTERM and SIGINT handlers up front so a failure here surfaces -/// before `axum::serve` starts. If installation fails (rare: container -/// without signal capability, seccomp policy), we return an error and the -/// process exits cleanly rather than running deaf to k8s termination. fn install_signal_handlers() -> Result<(Signal, Signal)> { let sigterm = signal(SignalKind::terminate()).context("install SIGTERM handler")?; let sigint = signal(SignalKind::interrupt()).context("install SIGINT handler")?; Ok((sigterm, sigint)) } -#[tokio::main] -async fn main() -> Result<()> { - let cli = Cli::parse(); - // Bootstrap subscriber so a config-resolution error has structured - // output. The configured-format subscriber installs after this and - // becomes a no-op via try_init's idempotency. - install_bootstrap_subscriber(); - let cfg = cli - .into_config() - .context("resolve configuration from CLI flags")?; - - init_tracing(&cfg.observability.log_level, cfg.observability.log_format)?; - - // Before anything slow — tokenizer download, discovery, bind. kubelet can - // SIGTERM a pod mid-rollout while it is still starting, and until the - // handlers exist that signal takes the default disposition: instant death, - // no readiness flip, no drain, no log. tokio's `Signal` buffers a - // notification until first polled, so installing here loses nothing and - // costs one deferred shutdown instead of a silent kill. - let (sigterm, sigint) = install_signal_handlers()?; - - // Emitted here rather than from `Config::validate`: this is startup advice - // about the deployment, not a validation failure, and keeping it out of - // `validate` leaves that function free of side effects. It also runs after - // the configured subscriber is installed. Static message, values as - // structured fields: a message that varies with the configured seconds - // cannot be grouped or deduped by a log aggregator. +fn log_startup(config: &Config) { if let Some(advisory) = sgl_router::config::shutdown_drain_advisory( - cfg.server.shutdown_drain_secs, - cfg.server.termination_grace_secs, + config.server.shutdown_drain_secs, + config.server.termination_grace_secs, ) { tracing::warn!( shutdown_drain_secs = advisory.shutdown_drain_secs, @@ -128,224 +180,195 @@ async fn main() -> Result<()> { } tracing::info!( - configured_decode_policy = ?cfg.model.decode_policy, + configured_decode_policy = ?config.model.decode_policy, "sgl-router {} starting on {}:{}", env!("CARGO_PKG_VERSION"), - cfg.server.host, - cfg.server.port + config.server.host, + config.server.port ); +} - let tokenizers = Arc::new( - sgl_router::tokenizer::TokenizerRegistry::load_from_config(&cfg) - .context("load tokenizers")?, - ); - - let registry = Arc::new(sgl_router::workers::WorkerRegistry::default()); - let cache_aware_uses_indexer = cfg.model.policy == PolicyKind::CacheAware - && cfg - .model - .cache_aware - .as_ref() - .is_some_and(|cache| cache.prefix_provider == CachePrefixProvider::Indexer); - let prefix_index: Option> = cache_aware_uses_indexer - .then_some(cfg.model.cache_aware.as_ref()) - .flatten() - .and_then(|cache| cache.kv_indexer_endpoint.as_ref()) - .map(|indexer| { - let config = prefix_index_config(indexer); - sgl_kv_indexer::GrpcPrefixIndex::new(config) - .map(|index| Arc::new(index) as Arc) +fn create_external_kv_indexer_client(config: &Config) -> Result>> { + let endpoint = config + .model + .cache_aware + .as_ref() + .filter(|cache| { + config.model.policy == PolicyKind::CacheAware + && cache.prefix_provider == CachePrefixProvider::Indexer + }) + .and_then(|cache| cache.kv_indexer_endpoint.as_ref()); + endpoint + .map(|endpoint| { + GrpcPrefixIndex::new(prefix_index_config(endpoint)) + .map(|index| Arc::new(index) as Arc) .context("configure KV Indexer client") }) - .transpose()?; + .transpose() +} - // Build the local prefix index and block metadata used by the Radix Tree - // provider. An external Indexer only needs hash metadata, so it does not - // subscribe to the local KV-event stream. - let block_size_oracle = sgl_router::policies::kv_events::BlockSizeOracle::new(); - let kv_event_http = reqwest::Client::builder() - .timeout(std::time::Duration::from_secs(2)) +fn prefix_index_config(indexer: &KvIndexerEndpointConfig) -> PrefixIndexConfig { + PrefixIndexConfig { + endpoint: indexer.url.clone(), + query_deadline: Duration::from_millis(indexer.query_timeout_ms), + max_inflight: indexer.query_max_inflight, + } +} + +fn start_engine_state_monitor(use_external_indexer: bool) -> Arc { + let http = reqwest::Client::builder() + .timeout(Duration::from_secs(2)) .build() .expect("default http client builds"); - let kv_index = if prefix_index.is_some() { - sgl_router::policies::kv_events::KvEventIndex::new_metadata_only_with_http_and_oracle( - kv_event_http, - Arc::clone(&block_size_oracle), - ) + if use_external_indexer { + // External indexing still needs worker hash metadata and engine load, but no local KV tree. + KvEventIndex::new_metadata_only_with_http_and_oracle(http, BlockSizeOracle::new()) } else { - sgl_router::policies::kv_events::KvEventIndex::new_with_http_and_oracle( - kv_event_http, - Arc::clone(&block_size_oracle), - ) - }; - let policies = Arc::new( - sgl_router::policies::factory::build_registry( - &cfg, - kv_index.tree(), - Arc::clone(&block_size_oracle), - ) - .context("build policy registry")?, - ); + KvEventIndex::new_with_http(http) + } +} - // Shared ActiveLoadRegistry + janitor task. The janitor reaps - // request entries whose lifetime exceeded `stale_request_timeout`, - // so a leaked guard (proxy task panic, etc.) does not inflate a - // worker's load forever. The registry is built BEFORE the manager - // is spawned so the manager can call `forget_worker` on - // `DiscoveryEvent::Removed`. - let stale_timeout = std::time::Duration::from_secs(cfg.active_load.stale_request_timeout_secs); - let active_load = sgl_router::policies::active_load::ActiveLoadRegistry::new( - Arc::new(sgl_router::policies::active_load::SystemTimeClock), - stale_timeout, - ); - // Sweep cadence is 1/10 of the configured timeout, clamped to - // [1 s, 60 s]. A short timeout (test setting) needs frequent - // sweeps to fire within the test's window; a long timeout - // (production) doesn't need sub-minute checks. - let sweep_interval = std::time::Duration::from_secs( - (cfg.active_load.stale_request_timeout_secs / 10).clamp(1, 60), - ); - let janitor_handle = - sgl_router::policies::active_load::spawn_janitor(Arc::clone(&active_load), sweep_interval); +fn start_local_inflight_tracker(config: &Config) -> (Arc, JanitorHandle) { + let timeout_secs = config.active_load.stale_request_timeout_secs; + let local_inflight_requests = + ActiveLoadRegistry::new(Arc::new(SystemTimeClock), Duration::from_secs(timeout_secs)); + // Reap stale requests at one tenth of their timeout, bounded to 1–60 seconds. + let sweep_interval = Duration::from_secs((timeout_secs / 10).clamp(1, 60)); + let inflight_cleanup = spawn_janitor(Arc::clone(&local_inflight_requests), sweep_interval); + (local_inflight_requests, inflight_cleanup) +} - // Spawn discovery + manager tasks. - // The manager resolves each worker's wire protocol from its `/server_info` - // and stamps it onto the registered worker. The proxy holds one client per - // protocol and selects by the worker's protocol per request, so the manager - // needs no proxy handle. - let (event_rx, discovery_handle) = sgl_router::discovery::spawn_discovery(&cfg) - .await - .context("spawn discovery")?; - let kv_index_opt: Option> = - Some(Arc::clone(&kv_index)); - let manager_handle = tokio::spawn(sgl_router::workers::manager::run_with_config( - event_rx, - registry.clone(), - Some(Arc::new(cfg.clone())), - kv_index_opt, - Some(Arc::clone(&active_load)), +async fn start_worker_discovery_and_manager( + config: &Config, + worker_registry: &Arc, + engine_state: &Arc, + local_inflight_requests: &Arc, +) -> Result<(JoinHandle<()>, JoinHandle<()>)> { + let (worker_events, discovery_handle) = + spawn_discovery(config).await.context("spawn discovery")?; + // Keep engine subscriptions and local request counters in sync with worker membership. + let worker_manager_handle = tokio::spawn(manager::run_with_config( + worker_events, + Arc::clone(worker_registry), + Some(Arc::new(config.clone())), + Some(Arc::clone(engine_state)), + Some(Arc::clone(local_inflight_requests)), )); + Ok((discovery_handle, worker_manager_handle)) +} +fn build_app_context( + config: &Config, + tokenizers: Arc, + worker_registry: Arc, + routing_policies: Arc, + local_inflight_requests: Arc, + engine_state: &KvEventIndex, + external_kv_indexer_client: Option>, +) -> Result> { + let block_size_oracle = engine_state.block_size_oracle(); let proxy = Arc::new( - sgl_router::proxy::Proxy::new(std::time::Duration::from_secs( - cfg.proxy.request_timeout_secs, - )) - .context("build proxy client")?, + Proxy::new(Duration::from_secs(config.proxy.request_timeout_secs)) + .context("build proxy client")?, ); - let mut app_ctx = sgl_router::server::app_context::AppContext::with_active_load( - cfg.clone(), + let mut app_context = AppContext::with_active_load( + config.clone(), tokenizers, proxy, - registry, - policies, - active_load, + worker_registry, + routing_policies, + local_inflight_requests, ); - app_ctx.prefix_index = prefix_index; - app_ctx.radix_tree_prefix_provider = (cfg.model.policy == PolicyKind::CacheAware - && cfg + app_context.prefix_index = external_kv_indexer_client; + app_context.radix_tree_prefix_provider = (config.model.policy == PolicyKind::CacheAware + && config .model .cache_aware .as_ref() .is_some_and(|cache| cache.prefix_provider == CachePrefixProvider::RadixTree)) - .then(|| { - sgl_router::policies::prefix_provider::RadixTreePrefixProvider::new( - kv_index.tree(), - Arc::clone(&block_size_oracle), - ) - }); - app_ctx.block_size_oracle = block_size_oracle; - app_ctx.engine_load = kv_index.engine_load(); - app_ctx.kv_metrics = kv_index.metrics_source(); - let ctx = Arc::new(app_ctx); - ctx.mark_ready(); + .then(|| RadixTreePrefixProvider::new(engine_state.tree(), Arc::clone(&block_size_oracle))); + app_context.block_size_oracle = block_size_oracle; + app_context.engine_load = engine_state.engine_load(); + app_context.kv_metrics = engine_state.metrics_source(); + Ok(Arc::new(app_context)) +} - let app = sgl_router::server::app::build_router(ctx.clone()); +/// How serving ended; `inflight_drain_secs` is `None` when the server stopped +/// without ever reaching the in-flight drain. +struct ServeOutcome { + result: Result<()>, + inflight_drain_secs: Option, +} - let bind = format!("{}:{}", cfg.server.host, cfg.server.port); - let listener = tokio::net::TcpListener::bind(&bind) +async fn serve( + listener: TcpListener, + app_context: Arc, + sigterm: Signal, + sigint: Signal, +) -> ServeOutcome { + let app = build_router(Arc::clone(&app_context)); + let drain = app_context.config.server.shutdown_drain(); + let (drain_tx, drain_rx) = watch::channel(None); + let heartbeat = tokio::spawn(report_drain_progress( + Arc::clone(&app_context), + drain_rx.clone(), + )); + let result = axum::serve(listener, app) + .with_graceful_shutdown(async move { + shutdown_signal(sigterm, sigint, app_context, drain).await; + let _ = drain_tx.send(Some(Instant::now())); + }) .await - .with_context(|| format!("bind {bind}"))?; - tracing::info!("listening on {bind}"); - - // Published the moment the readiness drain finishes, i.e. when axum starts - // its in-flight drain. That phase — not the pause, and not the uptime - // before it — is what the heartbeat below reports on. - let (inflight_drain_tx, inflight_drain_rx) = - tokio::sync::watch::channel(None::); - let shutdown_ctx = ctx.clone(); - let drain = cfg.server.shutdown_drain(); - let serve = axum::serve(listener, app).with_graceful_shutdown(async move { - shutdown_signal(sigterm, sigint, shutdown_ctx, drain).await; - let _ = inflight_drain_tx.send(Some(std::time::Instant::now())); - }); - - let heartbeat_ctx = ctx.clone(); - let mut heartbeat_rx = inflight_drain_rx.clone(); - let heartbeat = tokio::spawn(async move { - // Stay silent until the in-flight drain actually begins; an `Err` here - // means the sender went away without one, so there is nothing to report. - let Ok(started) = heartbeat_rx - .wait_for(Option::is_some) - .await - .map(|at| at.expect("wait_for only resolves once the instant is published")) - else { - return; - }; - let mut ticker = tokio::time::interval(DRAIN_HEARTBEAT_INTERVAL); - // Delay, not the default Burst: the runtime stalling is exactly the - // condition this heartbeat exists to report, and Burst would answer it - // with a clump of back-dated ticks instead of one line per interval. - ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); - ticker.tick().await; // the first tick completes immediately - - // One message, two severities: duplicating the text across an if/else - // is how the two arms drift apart. - macro_rules! heartbeat { - ($level:ident, $elapsed:expr) => { - tracing::$level!( - elapsed_secs = $elapsed, - // What axum is actually waiting on: every open exchange, - // on every route, until its response body finishes. - inflight_http = heartbeat_ctx.inflight_http.count(), - // The proxied subset, to separate "waiting on a worker" - // from "waiting on a client that stopped reading". - inflight_proxied = heartbeat_ctx.active_load.inflight_count(), - "still draining in-flight requests; this phase is unbounded and ends at \ - SIGKILL when terminationGracePeriodSeconds expires", - ) - }; - } - loop { - ticker.tick().await; - let elapsed = started.elapsed(); - if elapsed < DRAIN_WARN_AFTER { - heartbeat!(info, elapsed.as_secs()); - } else { - heartbeat!(warn, elapsed.as_secs()); - } - } - }); - let server_result = serve.await.context("axum serve"); + .context("axum serve"); heartbeat.abort(); - let inflight_drain_secs = inflight_drain_rx.borrow().map(|at| at.elapsed().as_secs()); + let inflight_drain_secs = drain_rx.borrow().map(|at| at.elapsed().as_secs()); + ServeOutcome { + result, + inflight_drain_secs, + } +} - // Best-effort: cancel discovery + manager + janitor on shutdown. - // The janitor handle's drop signals cancellation; we additionally - // await `shutdown` so the task joins cleanly before the process - // exits — useful for tracing tail logs. `JanitorHandle::shutdown` caps its - // own join at 2 s, so it cannot hang the exit — though those 2 s are still - // charged to terminationGracePeriodSeconds. - discovery_handle.abort(); - manager_handle.abort(); - janitor_handle.shutdown().await; - // The ERROR arms exist because otherwise the log says "shutdown complete" - // at INFO and the error leaves the process through `Termination`, never - // through `tracing` — so a severity-based alert sees nothing wrong with a - // crashed router. `None` means the server stopped without ever reaching the - // drain, which is not the same as draining instantly, so it gets its own - // message rather than `inflight_drain_secs = 0`. - match (&server_result, inflight_drain_secs) { +async fn report_drain_progress( + app_context: Arc, + mut drain_rx: watch::Receiver>, +) { + // Start reporting only after the readiness pause, when axum begins draining requests. + let Ok(started) = drain_rx + .wait_for(Option::is_some) + .await + .map(|at| at.expect("wait_for only resolves once the instant is published")) + else { + return; + }; + let mut ticker = tokio::time::interval(DRAIN_HEARTBEAT_INTERVAL); + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + ticker.tick().await; // the first tick completes immediately + + macro_rules! heartbeat { + ($level:ident, $elapsed:expr) => { + tracing::$level!( + elapsed_secs = $elapsed, + inflight_http = app_context.inflight_http.count(), + inflight_proxied = app_context.active_load.inflight_count(), + "still draining in-flight requests; this phase is unbounded and ends at \ + SIGKILL when terminationGracePeriodSeconds expires", + ) + }; + } + loop { + ticker.tick().await; + let elapsed = started.elapsed(); + if elapsed < DRAIN_WARN_AFTER { + heartbeat!(info, elapsed.as_secs()); + } else { + heartbeat!(warn, elapsed.as_secs()); + } + } +} + +fn log_shutdown(result: &Result<()>, inflight_drain_secs: Option) { + match (result, inflight_drain_secs) { (Ok(()), Some(inflight_drain_secs)) => { tracing::info!(inflight_drain_secs, "shutdown complete") } @@ -362,36 +385,14 @@ async fn main() -> Result<()> { "the server exited with an error before any termination signal", ), } - server_result } -/// Build the external Indexer client with the Router's bounded query settings. -fn prefix_index_config( - indexer: &sgl_router::config::KvIndexerEndpointConfig, -) -> sgl_kv_indexer::PrefixIndexConfig { - sgl_kv_indexer::PrefixIndexConfig { - endpoint: indexer.url.clone(), - query_deadline: std::time::Duration::from_millis(indexer.query_timeout_ms), - max_inflight: indexer.query_max_inflight, - } -} - -/// Resolve when a termination signal arrives, then hand control to axum's -/// graceful drain. On SIGTERM (k8s pod termination) first run the readiness -/// drain — flip `/readyz` to 503 and keep serving for `drain` so the endpoint -/// removal reaches kube-proxy before we stop accepting, closing the -/// rolling-update race. SIGINT (local Ctrl-C) skips the readiness drain -/// entirely — no 503 flip, no pause — so dev iteration does not pay it. -/// -/// Either way axum's own in-flight drain runs afterwards and is unbounded: a -/// long streaming completion still holds the process until it finishes or -/// `terminationGracePeriodSeconds` expires. A further termination signal cuts -/// the pause short but cannot reach that phase; it is logged instead. +// SIGTERM pauses for readiness propagation; SIGINT goes straight to the in-flight drain. async fn shutdown_signal( mut sigterm: Signal, mut sigint: Signal, - ctx: Arc, - drain: std::time::Duration, + app_context: Arc, + drain: Duration, ) { let sigterm_first = tokio::select! { _ = sigterm.recv() => { @@ -404,15 +405,9 @@ async fn shutdown_signal( } }; - let (expedite_tx, expedite_rx) = tokio::sync::oneshot::channel::<()>(); - // Only the SIGTERM path runs a pause, so only it has something to cut - // short; on the SIGINT path the first further signal goes straight to the - // warning below. + let (expedite_tx, expedite_rx) = oneshot::channel::<()>(); let mut expedite_tx = sigterm_first.then_some(expedite_tx); - // Hand both streams to a task that outlives this future, on EITHER branch. - // Dropping them here would make every later signal vanish: tokio never - // restores the default disposition, so the process would neither expedite - // nor die, and nothing would be logged. + // Keep consuming signals during the in-flight drain; tokio never restores default handlers. tokio::spawn(async move { loop { let delivered = tokio::select! { @@ -420,8 +415,6 @@ async fn shutdown_signal( delivered = sigint.recv() => delivered, }; if delivered.is_none() { - // The signal driver is gone (runtime shutting down). Looping - // would spin without ever receiving again. return; } handle_further_signal(&mut expedite_tx, sigterm_first); @@ -432,45 +425,28 @@ async fn shutdown_signal( let expedite = async move { let _ = expedite_rx.await; }; - sgl_router::server::shutdown::drain_for_termination(&ctx, drain, expedite).await; + drain_for_termination(&app_context, drain, expedite).await; } } -/// What a termination signal past the first one achieved. #[derive(Debug, PartialEq, Eq)] enum FurtherSignal { - /// Cut the readiness pause short. Expedited, - /// Arrived with no pause left to cut short, and was reported as such. Ignored, } -/// Handle one termination signal past the first, and say what it did. -/// -/// A failed `send` is not a lost race to shrug at: it means the pause already -/// ended on its own and dropped the receiver with it, which is the same -/// "nothing left to cut short" state as a spent `expedite_tx` and must reach -/// the same notice. Discarding that `Err` is what made the *first* signal after -/// the pause disappear, so only a second one was ever reported — the opposite -/// of this function's contract. -/// -/// Split out of the watcher task because that task owns real `Signal` streams -/// and cannot be driven from a test; this is where the decision lives, so this -/// is what a test can pin. fn handle_further_signal( - expedite_tx: &mut Option>, + expedite_tx: &mut Option>, sigterm_first: bool, ) -> FurtherSignal { - // The first further signal cuts the readiness pause short, so an operator - // watching a stuck rollout is not held for a window that has stopped being - // useful. + // A failed `send` means the pause already elapsed; it must fall through to + // the notice below. Discarding the `Err` once swallowed the first + // post-pause signal (see the regression test). if let Some(tx) = expedite_tx.take() { if tx.send(()).is_ok() { return FurtherSignal::Expedited; } } - // Two messages, because the two states call for different conclusions: one - // ran a pause that has since passed, the other never ran one at all. if sigterm_first { tracing::warn!( "further termination signal ignored: the readiness pause is over and the \ @@ -491,27 +467,24 @@ mod tests { #[test] fn prefix_index_config_preserves_router_limits() { - let config = prefix_index_config(&sgl_router::config::KvIndexerEndpointConfig { + let config = prefix_index_config(&KvIndexerEndpointConfig { url: "http://127.0.0.1:50051".to_string(), query_timeout_ms: 25, query_max_inflight: 17, }); assert_eq!(config.endpoint, "http://127.0.0.1:50051"); - assert_eq!(config.query_deadline, std::time::Duration::from_millis(25)); + assert_eq!(config.query_deadline, Duration::from_millis(25)); assert_eq!(config.max_inflight, 17); } #[tokio::test] async fn install_signal_handlers_returns_both() { - // Pins the contract that handler installation works on a standard - // tokio runtime. If this fails on a sandboxed runner, the real - // service would also fail to install — which is the point. assert!(install_signal_handlers().is_ok()); } #[test] fn a_further_signal_expedites_a_running_pause() { - let (tx, mut rx) = tokio::sync::oneshot::channel::<()>(); + let (tx, mut rx) = oneshot::channel::<()>(); let mut expedite_tx = Some(tx); assert_eq!( handle_further_signal(&mut expedite_tx, true), @@ -527,13 +500,9 @@ mod tests { ); } - /// The regression this function exists for. When the pause elapses on its - /// own, the receiver goes with it while the watcher still holds the sender - /// — so the very next signal hits a `send` that fails. Discarding that - /// `Err` swallowed it, and only the signal AFTER it was ever reported. #[test] fn the_first_signal_after_the_pause_ends_is_reported_not_swallowed() { - let (tx, rx) = tokio::sync::oneshot::channel::<()>(); + let (tx, rx) = oneshot::channel::<()>(); drop(rx); // the pause elapsed and dropped its receiver let mut expedite_tx = Some(tx); assert_eq!( @@ -541,17 +510,12 @@ mod tests { FurtherSignal::Ignored, "a send into a dropped receiver expedites nothing and must say so", ); - // ...and every later signal behaves identically, rather than the second - // one being the first to report anything. assert_eq!( handle_further_signal(&mut expedite_tx, true), FurtherSignal::Ignored, ); } - /// SIGINT never runs a readiness pause, so `expedite_tx` is `None` from the - /// start and every further Ctrl-C is reported rather than expediting a - /// pause that does not exist. #[test] fn a_further_signal_on_the_sigint_path_is_always_ignored() { let mut expedite_tx = None; @@ -569,8 +533,6 @@ mod tests { #[test] fn init_tracing_accepts_json_format() { - // Doesn't matter whether we win or lose the race against another - // subscriber install — the function must return Ok either way. assert!(init_tracing("info", LogFormat::Json).is_ok()); } }