From 802c2a3fb092d8ec6ed07ba7291e2444dcb227a3 Mon Sep 17 00:00:00 2001 From: Tyler Longwell Date: Fri, 13 Mar 2026 23:09:05 -0400 Subject: [PATCH] feat: optional UDS listener + health endpoints for Kubernetes deployment - Add SPROUT_UDS_PATH env var for optional Unix Domain Socket binding - Dual-bind TCP + UDS via tokio::select! (gated with #[cfg(unix)]) - Add /_liveness (always 200) and /_readiness (checks MySQL + Redis) - Concurrent readiness checks with 2s global timeout - Safe UDS cleanup (only removes socket files, not arbitrary paths) - Add Db::ping() method for readiness health checks - Add redis_pool to AppState for readiness probe - Zero behavior change when SPROUT_UDS_PATH is unset --- crates/sprout-db/src/lib.rs | 5 +++ crates/sprout-relay/src/config.rs | 9 ++++++ crates/sprout-relay/src/main.rs | 53 +++++++++++++++++++++++++++++-- crates/sprout-relay/src/router.rs | 53 ++++++++++++++++++++++++++++++- crates/sprout-relay/src/state.rs | 5 +++ 5 files changed, 121 insertions(+), 4 deletions(-) diff --git a/crates/sprout-db/src/lib.rs b/crates/sprout-db/src/lib.rs index b36b52d458..34b92fbaf1 100644 --- a/crates/sprout-db/src/lib.rs +++ b/crates/sprout-db/src/lib.rs @@ -925,6 +925,11 @@ impl Db { partition::ensure_future_partitions(&self.pool, months_ahead).await } + /// Returns `true` if the database is reachable (used by readiness probes). + pub async fn ping(&self) -> bool { + sqlx::query("SELECT 1").execute(&self.pool).await.is_ok() + } + // ── Workflows ───────────────────────────────────────────────────────────── /// Creates a new workflow definition and returns its UUID. diff --git a/crates/sprout-relay/src/config.rs b/crates/sprout-relay/src/config.rs index 1d3f2aad1d..d2452cacf7 100644 --- a/crates/sprout-relay/src/config.rs +++ b/crates/sprout-relay/src/config.rs @@ -45,6 +45,9 @@ pub struct Config { /// Optional hex-encoded private key for the relay's signing keypair. /// If absent, a fresh keypair is generated at startup. pub relay_private_key: Option, + /// Optional Unix Domain Socket path. When set, the relay also listens on this + /// UDS for traffic (e.g. service mesh sidecar). Health probes still use TCP. + pub uds_path: Option, } impl Config { @@ -118,6 +121,11 @@ impl Config { let relay_private_key = std::env::var("SPROUT_RELAY_PRIVATE_KEY").ok(); + let uds_path = std::env::var("SPROUT_UDS_PATH") + .ok() + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()); + Ok(Self { bind_addr, database_url, @@ -132,6 +140,7 @@ impl Config { require_auth_token, cors_origins, relay_private_key, + uds_path, }) } } diff --git a/crates/sprout-relay/src/main.rs b/crates/sprout-relay/src/main.rs index 37e2746504..0a4dd5a43f 100644 --- a/crates/sprout-relay/src/main.rs +++ b/crates/sprout-relay/src/main.rs @@ -61,6 +61,7 @@ async fn main() -> anyhow::Result<()> { cfg.create_pool(Some(deadpool_redis::Runtime::Tokio1)) .map_err(|e| anyhow::anyhow!("Redis pool creation failed: {e}"))? }; + let redis_health_pool = redis_pool.clone(); // cheap Arc clone — shared with readiness handler let pubsub = Arc::new( PubSubManager::new(&config.redis_url, redis_pool) .await @@ -101,6 +102,7 @@ async fn main() -> anyhow::Result<()> { let state = Arc::new(AppState::new( config.clone(), db, + redis_health_pool, audit, pubsub, auth, @@ -184,14 +186,59 @@ async fn main() -> anyhow::Result<()> { let router = build_router(Arc::clone(&state)); - let listener = tokio::net::TcpListener::bind(&config.bind_addr) + serve_dual(router, &config).await +} + +/// Bind TCP (always) and optionally UDS, then run both concurrently. +/// If either listener exits, the function returns and the process exits. +async fn serve_dual(router: axum::Router, config: &Config) -> anyhow::Result<()> { + let tcp_listener = tokio::net::TcpListener::bind(&config.bind_addr) .await .map_err(|e| anyhow::anyhow!("Failed to bind {}: {e}", config.bind_addr))?; + info!(addr = %config.bind_addr, "sprout-relay TCP listening"); + + #[cfg(unix)] + if let Some(ref uds_path) = config.uds_path { + // Only remove stale socket files — refuse to delete non-socket paths. + use std::os::unix::fs::FileTypeExt as _; + match std::fs::symlink_metadata(uds_path) { + Ok(meta) if meta.file_type().is_socket() => { + let _ = std::fs::remove_file(uds_path); + } + Ok(_) => { + return Err(anyhow::anyhow!( + "SPROUT_UDS_PATH {uds_path} exists but is not a socket — refusing to delete" + )); + } + Err(_) => {} // Path doesn't exist — fine, we'll create it + } + let uds_listener = tokio::net::UnixListener::bind(uds_path) + .map_err(|e| anyhow::anyhow!("Failed to bind UDS {uds_path}: {e}"))?; + info!(path = %uds_path, "sprout-relay UDS listening"); + + let router_uds = router.clone(); + tokio::select! { + r = axum::serve( + tcp_listener, + router.into_make_service_with_connect_info::(), + ) => r.map_err(|e| anyhow::anyhow!("TCP server error: {e}")), + r = axum::serve(uds_listener, router_uds.into_make_service()) => + r.map_err(|e| anyhow::anyhow!("UDS server error: {e}")), + }?; + return Ok(()); + } - info!(addr = %config.bind_addr, "sprout-relay listening"); + #[cfg(not(unix))] + if config.uds_path.is_some() { + tracing::warn!( + "SPROUT_UDS_PATH is set but UDS is not supported on this platform — \ + falling back to TCP only" + ); + } + // TCP-only path (non-unix, or unix without uds_path set). axum::serve( - listener, + tcp_listener, router.into_make_service_with_connect_info::(), ) .await diff --git a/crates/sprout-relay/src/router.rs b/crates/sprout-relay/src/router.rs index dc15de495a..77558f2399 100644 --- a/crates/sprout-relay/src/router.rs +++ b/crates/sprout-relay/src/router.rs @@ -9,6 +9,7 @@ use axum::{ routing::{delete, get, post, put}, Router, }; +use serde_json::json; use tower_http::cors::{AllowOrigin, CorsLayer}; use tower_http::limit::RequestBodyLimitLayer; use tower_http::trace::TraceLayer; @@ -26,6 +27,8 @@ pub fn build_router(state: Arc) -> Router { .route("/info", get(relay_info_handler)) .route("/.well-known/nostr.json", get(api::nip05::nostr_nip05)) .route("/health", get(health_handler)) + .route("/_liveness", get(liveness_handler)) + .route("/_readiness", get(readiness_handler)) // Token self-service routes .route( "/api/tokens", @@ -153,12 +156,23 @@ pub fn build_router(state: Arc) -> Router { /// /// Uses `axum::extract::Request` to manually attempt WS upgrade, so non-WS /// requests aren't rejected by the extractor. +/// +/// `ConnectInfo` is read from request extensions rather than as an extractor — +/// UDS connections have no `SocketAddr`, so the extractor would panic. +/// TCP connections populate it via `into_make_service_with_connect_info`; UDS +/// connections fall back to `0.0.0.0:0`. async fn nip11_or_ws_handler( State(state): State>, - ConnectInfo(addr): ConnectInfo, headers: HeaderMap, req: axum::extract::Request, ) -> impl IntoResponse { + // Read peer address from extensions (set by TCP serve; absent for UDS). + let addr = req + .extensions() + .get::>() + .map(|ci| ci.0) + .unwrap_or_else(|| std::net::SocketAddr::from(([0, 0, 0, 0], 0))); + let accept = headers .get("accept") .and_then(|v| v.to_str().ok()) @@ -185,6 +199,43 @@ async fn health_handler() -> impl IntoResponse { (StatusCode::OK, "ok") } +async fn liveness_handler() -> impl IntoResponse { + (StatusCode::OK, "ok") +} + +/// Readiness probe — checks MySQL and Redis connectivity. +/// +/// Returns 200 `{"status":"ready"}` when both are reachable, or +/// 503 `{"status":"not_ready","mysql":bool,"redis":bool}` otherwise. +/// +/// Note: Redis pool exhaustion (all connections in use) also returns 503. +/// This is intentionally conservative — if we can't acquire a connection, +/// the pod shouldn't receive traffic. +async fn readiness_handler(State(state): State>) -> impl IntoResponse { + use std::time::Duration; + + let check = async { + let (mysql_ok, redis_ok) = tokio::join!(state.db.ping(), async { + state.redis_pool.get().await.is_ok() + },); + (mysql_ok, redis_ok) + }; + + let (mysql_ok, redis_ok) = tokio::time::timeout(Duration::from_secs(2), check) + .await + .unwrap_or((false, false)); + + if mysql_ok && redis_ok { + (StatusCode::OK, Json(json!({"status": "ready"}))).into_response() + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({"status": "not_ready", "mysql": mysql_ok, "redis": redis_ok})), + ) + .into_response() + } +} + /// Build a CORS layer from the configured origins list. /// /// If `cors_origins` is empty (dev default), returns a permissive layer. diff --git a/crates/sprout-relay/src/state.rs b/crates/sprout-relay/src/state.rs index 2551e7f24b..ea11682725 100644 --- a/crates/sprout-relay/src/state.rs +++ b/crates/sprout-relay/src/state.rs @@ -9,6 +9,7 @@ use tokio::sync::{mpsc, Semaphore}; use tokio_util::sync::CancellationToken; use uuid::Uuid; +use deadpool_redis; use sprout_audit::AuditService; use sprout_auth::AuthService; use sprout_db::Db; @@ -83,6 +84,8 @@ pub struct AppState { pub config: Arc, /// Database connection pool. pub db: Db, + /// Redis pool for readiness health checks. + pub redis_pool: deadpool_redis::Pool, /// Audit event service. pub audit: Arc, /// Pub/sub manager for broadcasting events to subscribers. @@ -122,6 +125,7 @@ impl AppState { pub fn new( config: Config, db: Db, + redis_pool: deadpool_redis::Pool, audit: AuditService, pubsub: Arc, auth: AuthService, @@ -134,6 +138,7 @@ impl AppState { Self { config: Arc::new(config), db, + redis_pool, audit: Arc::new(audit), pubsub, auth: Arc::new(auth),