[sgl-router] refactor - main startup logic (#39861)

Co-authored-by: Cursor <cursoragent@cursor.com>

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 sgl_kv_indexer::{GrpcPrefixIndex, PrefixIndex, PrefixIndexConfig};
-use std::sync::Arc;
+use sgl_router::{
-use tokio::signal::unix::{signal, Signal, SignalKind};
+    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.
+const DRAIN_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
-///
+/// Heartbeat escalates INFO -> WARN here: earlier is a routine rollout draining
-/// Idempotent: a second call returns `Ok` without panicking. When
+/// a long response; later the pod risks SIGKILL with work still open.
-/// `try_init` errors, some other code has already installed a subscriber,
+const DRAIN_WARN_AFTER: Duration = Duration::from_secs(30);
-/// so the `tracing::debug!` below is delivered through THAT subscriber —
+
-/// no recursive init.
+// Main components started by this binary:
-///
+// - Engine monitor (`KvEventIndex`): receives load statistics over ZMQ and, without a
-/// `format` selects the output shape: `Json` emits one JSON record per
+//   remote KV indexer, KV events to maintain a local radix tree. Remote prefix lookups
-/// line (target for production / k8s log aggregators), `Text` is the
+//   use `GrpcPrefixIndex` over gRPC.
-/// human-readable default. The `RUST_LOG` environment variable always
+// - Engine discovery (`spawn_discovery`): watches Kubernetes pods or loads static URLs,
-/// wins over `default_level`.
+//   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
+// Provide startup logging before configuration resolution; later installs are no-ops.
-/// 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.
 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]
+fn log_startup(config: &Config) {
-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.
     if let Some(advisory) = sgl_router::config::shutdown_drain_advisory(
-        cfg.server.shutdown_drain_secs,
+        config.server.shutdown_drain_secs,
-        cfg.server.termination_grace_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,
+        config.server.host,
-        cfg.server.port
+        config.server.port
     );
+}
 
-    let tokenizers = Arc::new(
+fn create_external_kv_indexer_client(config: &Config) -> Result<Option<Arc<dyn PrefixIndex>>> {
-        sgl_router::tokenizer::TokenizerRegistry::load_from_config(&cfg)
+    let endpoint = config
-            .context("load tokenizers")?,
+        .model
-    );
+        .cache_aware
-
+        .as_ref()
-    let registry = Arc::new(sgl_router::workers::WorkerRegistry::default());
+        .filter(|cache| {
-    let cache_aware_uses_indexer = cfg.model.policy == PolicyKind::CacheAware
+            config.model.policy == PolicyKind::CacheAware
-        && cfg
+                && cache.prefix_provider == CachePrefixProvider::Indexer
-            .model
+        })
-            .cache_aware
+        .and_then(|cache| cache.kv_indexer_endpoint.as_ref());
-            .as_ref()
+    endpoint
-            .is_some_and(|cache| cache.prefix_provider == CachePrefixProvider::Indexer);
+        .map(|endpoint| {
-    let prefix_index: Option<Arc<dyn sgl_kv_indexer::PrefixIndex>> = cache_aware_uses_indexer
+            GrpcPrefixIndex::new(prefix_index_config(endpoint))
-        .then_some(cfg.model.cache_aware.as_ref())
+                .map(|index| Arc::new(index) as Arc<dyn PrefixIndex>)
-        .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<dyn sgl_kv_indexer::PrefixIndex>)
                 .context("configure KV Indexer client")
         })
-        .transpose()?;
+        .transpose()
+}
 
-    // Build the local prefix index and block metadata used by the Radix Tree
+fn prefix_index_config(indexer: &KvIndexerEndpointConfig) -> PrefixIndexConfig {
-    // provider. An external Indexer only needs hash metadata, so it does not
+    PrefixIndexConfig {
-    // subscribe to the local KV-event stream.
+        endpoint: indexer.url.clone(),
-    let block_size_oracle = sgl_router::policies::kv_events::BlockSizeOracle::new();
+        query_deadline: Duration::from_millis(indexer.query_timeout_ms),
-    let kv_event_http = reqwest::Client::builder()
+        max_inflight: indexer.query_max_inflight,
-        .timeout(std::time::Duration::from_secs(2))
+    }
+}
+
+fn start_engine_state_monitor(use_external_indexer: bool) -> Arc<KvEventIndex> {
+    let http = reqwest::Client::builder()
+        .timeout(Duration::from_secs(2))
         .build()
         .expect("default http client builds");
-    let kv_index = if prefix_index.is_some() {
+    if use_external_indexer {
-        sgl_router::policies::kv_events::KvEventIndex::new_metadata_only_with_http_and_oracle(
+        // External indexing still needs worker hash metadata and engine load, but no local KV tree.
-            kv_event_http,
+        KvEventIndex::new_metadata_only_with_http_and_oracle(http, BlockSizeOracle::new())
-            Arc::clone(&block_size_oracle),
-        )
     } else {
-        sgl_router::policies::kv_events::KvEventIndex::new_with_http_and_oracle(
+        KvEventIndex::new_with_http(http)
-            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")?,
-    );
 
-    // Shared ActiveLoadRegistry + janitor task. The janitor reaps
+fn start_local_inflight_tracker(config: &Config) -> (Arc<ActiveLoadRegistry>, JanitorHandle) {
-    // request entries whose lifetime exceeded `stale_request_timeout`,
+    let timeout_secs = config.active_load.stale_request_timeout_secs;
-    // so a leaked guard (proxy task panic, etc.) does not inflate a
+    let local_inflight_requests =
-    // worker's load forever. The registry is built BEFORE the manager
+        ActiveLoadRegistry::new(Arc::new(SystemTimeClock), Duration::from_secs(timeout_secs));
-    // is spawned so the manager can call `forget_worker` on
+    // Reap stale requests at one tenth of their timeout, bounded to 1–60 seconds.
-    // `DiscoveryEvent::Removed`.
+    let sweep_interval = Duration::from_secs((timeout_secs / 10).clamp(1, 60));
-    let stale_timeout = std::time::Duration::from_secs(cfg.active_load.stale_request_timeout_secs);
+    let inflight_cleanup = spawn_janitor(Arc::clone(&local_inflight_requests), sweep_interval);
-    let active_load = sgl_router::policies::active_load::ActiveLoadRegistry::new(
+    (local_inflight_requests, inflight_cleanup)
-        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);
 
-    // Spawn discovery + manager tasks.
+async fn start_worker_discovery_and_manager(
-    // The manager resolves each worker's wire protocol from its `/server_info`
+    config: &Config,
-    // and stamps it onto the registered worker. The proxy holds one client per
+    worker_registry: &Arc<WorkerRegistry>,
-    // protocol and selects by the worker's protocol per request, so the manager
+    engine_state: &Arc<KvEventIndex>,
-    // needs no proxy handle.
+    local_inflight_requests: &Arc<ActiveLoadRegistry>,
-    let (event_rx, discovery_handle) = sgl_router::discovery::spawn_discovery(&cfg)
+) -> Result<(JoinHandle<()>, JoinHandle<()>)> {
-        .await
+    let (worker_events, discovery_handle) =
-        .context("spawn discovery")?;
+        spawn_discovery(config).await.context("spawn discovery")?;
-    let kv_index_opt: Option<Arc<sgl_router::policies::kv_events::KvEventIndex>> =
+    // Keep engine subscriptions and local request counters in sync with worker membership.
-        Some(Arc::clone(&kv_index));
+    let worker_manager_handle = tokio::spawn(manager::run_with_config(
-    let manager_handle = tokio::spawn(sgl_router::workers::manager::run_with_config(
+        worker_events,
-        event_rx,
+        Arc::clone(worker_registry),
-        registry.clone(),
+        Some(Arc::new(config.clone())),
-        Some(Arc::new(cfg.clone())),
+        Some(Arc::clone(engine_state)),
-        kv_index_opt,
+        Some(Arc::clone(local_inflight_requests)),
-        Some(Arc::clone(&active_load)),
     ));
+    Ok((discovery_handle, worker_manager_handle))
+}
 
+fn build_app_context(
+    config: &Config,
+    tokenizers: Arc<TokenizerRegistry>,
+    worker_registry: Arc<WorkerRegistry>,
+    routing_policies: Arc<PolicyRegistry>,
+    local_inflight_requests: Arc<ActiveLoadRegistry>,
+    engine_state: &KvEventIndex,
+    external_kv_indexer_client: Option<Arc<dyn PrefixIndex>>,
+) -> Result<Arc<AppContext>> {
+    let block_size_oracle = engine_state.block_size_oracle();
     let proxy = Arc::new(
-        sgl_router::proxy::Proxy::new(std::time::Duration::from_secs(
+        Proxy::new(Duration::from_secs(config.proxy.request_timeout_secs))
-            cfg.proxy.request_timeout_secs,
+            .context("build proxy client")?,
-        ))
-        .context("build proxy client")?,
     );
 
-    let mut app_ctx = sgl_router::server::app_context::AppContext::with_active_load(
+    let mut app_context = AppContext::with_active_load(
-        cfg.clone(),
+        config.clone(),
         tokenizers,
         proxy,
-        registry,
+        worker_registry,
-        policies,
+        routing_policies,
-        active_load,
+        local_inflight_requests,
     );
-    app_ctx.prefix_index = prefix_index;
+    app_context.prefix_index = external_kv_indexer_client;
-    app_ctx.radix_tree_prefix_provider = (cfg.model.policy == PolicyKind::CacheAware
+    app_context.radix_tree_prefix_provider = (config.model.policy == PolicyKind::CacheAware
-        && cfg
+        && config
             .model
             .cache_aware
             .as_ref()
             .is_some_and(|cache| cache.prefix_provider == CachePrefixProvider::RadixTree))
-    .then(|| {
+    .then(|| RadixTreePrefixProvider::new(engine_state.tree(), Arc::clone(&block_size_oracle)));
-        sgl_router::policies::prefix_provider::RadixTreePrefixProvider::new(
+    app_context.block_size_oracle = block_size_oracle;
-            kv_index.tree(),
+    app_context.engine_load = engine_state.engine_load();
-            Arc::clone(&block_size_oracle),
+    app_context.kv_metrics = engine_state.metrics_source();
-        )
+    Ok(Arc::new(app_context))
-    });
+}
-    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();
 
-    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<u64>,
+}
 
-    let bind = format!("{}:{}", cfg.server.host, cfg.server.port);
+async fn serve(
-    let listener = tokio::net::TcpListener::bind(&bind)
+    listener: TcpListener,
+    app_context: Arc<AppContext>,
+    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}"))?;
+        .context("axum serve");
-    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::<std::time::Instant>);
-    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");
     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.
+async fn report_drain_progress(
-    // The janitor handle's drop signals cancellation; we additionally
+    app_context: Arc<AppContext>,
-    // await `shutdown` so the task joins cleanly before the process
+    mut drain_rx: watch::Receiver<Option<Instant>>,
-    // 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
+    // Start reporting only after the readiness pause, when axum begins draining requests.
-    // charged to terminationGracePeriodSeconds.
+    let Ok(started) = drain_rx
-    discovery_handle.abort();
+        .wait_for(Option::is_some)
-    manager_handle.abort();
+        .await
-    janitor_handle.shutdown().await;
+        .map(|at| at.expect("wait_for only resolves once the instant is published"))
-    // The ERROR arms exist because otherwise the log says "shutdown complete"
+    else {
-    // at INFO and the error leaves the process through `Termination`, never
+        return;
-    // through `tracing` — so a severity-based alert sees nothing wrong with a
+    };
-    // crashed router. `None` means the server stopped without ever reaching the
+    let mut ticker = tokio::time::interval(DRAIN_HEARTBEAT_INTERVAL);
-    // drain, which is not the same as draining instantly, so it gets its own
+    ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay);
-    // message rather than `inflight_drain_secs = 0`.
+    ticker.tick().await; // the first tick completes immediately
-    match (&server_result, inflight_drain_secs) {
+
+    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<u64>) {
+    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.
+// SIGTERM pauses for readiness propagation; SIGINT goes straight to the in-flight drain.
-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.
 async fn shutdown_signal(
     mut sigterm: Signal,
     mut sigint: Signal,
-    ctx: Arc<sgl_router::server::app_context::AppContext>,
+    app_context: Arc<AppContext>,
-    drain: std::time::Duration,
+    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::<()>();
+    let (expedite_tx, expedite_rx) = 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 mut expedite_tx = sigterm_first.then_some(expedite_tx);
-    // Hand both streams to a task that outlives this future, on EITHER branch.
+    // Keep consuming signals during the in-flight drain; tokio never restores default handlers.
-    // 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.
     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<tokio::sync::oneshot::Sender<()>>,
+    expedite_tx: &mut Option<oneshot::Sender<()>>,
     sigterm_first: bool,
 ) -> FurtherSignal {
-    // The first further signal cuts the readiness pause short, so an operator
+    // A failed `send` means the pause already elapsed; it must fall through to
-    // watching a stuck rollout is not held for a window that has stopped being
+    // the notice below. Discarding the `Err` once swallowed the first
-    // useful.
+    // 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());
     }
 }