From c1f4a5f9ee6f34a53b1a70962c6dbce630db9059 Mon Sep 17 00:00:00 2001 From: Raudur Refur <261796072+refur-nfn@users.noreply.github.com> Date: Wed, 1 Jul 2026 12:12:38 +0300 Subject: [PATCH] feat(toolkit): add REST healthcheck infrastructure - Concurrent healthcheck execution with timeout protection - HealthcheckRegistry for registering and running per-gear probes - Message sanitization to prevent secret leaks in responses - API Gateway integration for /health and /readyz endpoints - Comprehensive test coverage with panic/timeout/degradation scenarios Signed-off-by: Raudur Refur <261796072+refur-nfn@users.noreply.github.com> --- gears/system/api-gateway/src/config.rs | 31 +- gears/system/api-gateway/src/gear.rs | 182 ++++- gears/system/api-gateway/src/web.rs | 50 +- .../api-gateway/tests/access_log_tests.rs | 6 +- .../api-gateway/tests/auth_middleware.rs | 54 +- .../api-gateway/tests/body_limit_tests.rs | 18 +- gears/system/api-gateway/tests/cors_tests.rs | 18 +- .../api-gateway/tests/health_endpoints.rs | 264 +++++++ .../api-gateway/tests/http_metrics_tests.rs | 36 +- .../api-gateway/tests/license_middleware.rs | 30 +- .../api-gateway/tests/middleware_order.rs | 12 +- .../api-gateway/tests/rate_limit_tests.rs | 18 +- libs/toolkit/src/contracts.rs | 30 +- libs/toolkit/src/healthcheck.rs | 642 ++++++++++++++++++ libs/toolkit/src/lib.rs | 8 + libs/toolkit/src/registry.rs | 2 + libs/toolkit/src/runtime/host_runtime.rs | 34 +- libs/toolkit/tests/macro_tests.rs | 2 + 18 files changed, 1344 insertions(+), 93 deletions(-) create mode 100644 gears/system/api-gateway/tests/health_endpoints.rs create mode 100644 libs/toolkit/src/healthcheck.rs diff --git a/gears/system/api-gateway/src/config.rs b/gears/system/api-gateway/src/config.rs index 0ca959f4f..176a45bb3 100644 --- a/gears/system/api-gateway/src/config.rs +++ b/gears/system/api-gateway/src/config.rs @@ -8,8 +8,12 @@ fn default_body_limit_bytes() -> usize { 16 * 1024 * 1024 } +fn default_healthcheck_timeout_ms() -> u64 { + 500 +} + /// API gateway configuration - reused from `api_gateway` gear -#[derive(Debug, Clone, Deserialize, Serialize, Default)] +#[derive(Debug, Clone, Deserialize, Serialize)] #[serde(deny_unknown_fields)] #[allow(clippy::struct_excessive_bools)] pub struct ApiGatewayConfig { @@ -57,6 +61,31 @@ pub struct ApiGatewayConfig { /// HTTP metrics configuration. #[serde(default)] pub metrics: MetricsConfig, + + /// Per-check `/health`/`/readyz` timeout (ms); raise for slow dependencies. + #[serde(default = "default_healthcheck_timeout_ms")] + pub healthcheck_timeout_ms: u64, +} + +impl Default for ApiGatewayConfig { + // Manual (not derived) so defaults match the `#[serde(default = ...)]` fns above; + // a derived Default would zero each field (e.g. timeout 0 = instant probe failure). + fn default() -> Self { + Self { + bind_addr: String::default(), + enable_docs: false, + cors_enabled: false, + cors: None, + openapi: OpenApiConfig::default(), + defaults: Defaults::default(), + auth_disabled: false, + require_auth_by_default: default_require_auth_by_default(), + prefix_path: String::default(), + route_policies: RoutePoliciesConfig::default(), + metrics: MetricsConfig::default(), + healthcheck_timeout_ms: default_healthcheck_timeout_ms(), + } + } } #[derive(Debug, Clone, Deserialize, Serialize)] diff --git a/gears/system/api-gateway/src/gear.rs b/gears/system/api-gateway/src/gear.rs index 402bbe34b..4b760e847 100644 --- a/gears/system/api-gateway/src/gear.rs +++ b/gears/system/api-gateway/src/gear.rs @@ -3,7 +3,7 @@ //! Contains the `ApiGateway` gear struct and its trait implementations. use async_trait::async_trait; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use arc_swap::ArcSwap; use dashmap::DashMap; @@ -12,12 +12,12 @@ use anyhow::Result; use axum::error_handling::HandleErrorLayer; use axum::http::Method; use axum::middleware::from_fn_with_state; -use axum::{Router, extract::DefaultBodyLimit, middleware::from_fn, routing::get}; +use axum::{Extension, Router, extract::DefaultBodyLimit, middleware::from_fn, routing::get}; use parking_lot::Mutex; use std::net::SocketAddr; use std::time::Duration; use tokio_util::sync::CancellationToken; -use toolkit::api::{OpenApiRegistry, OpenApiRegistryImpl}; +use toolkit::api::{OpenApiRegistry, OpenApiRegistryImpl, OperationBuilder}; use toolkit::lifecycle::ReadySignal; use tower::{BoxError, ServiceBuilder}; use tower_http::{ @@ -77,6 +77,13 @@ pub struct ApiGateway { pub(crate) final_router: Mutex>, // AuthN Resolver client (resolved during init, None when auth_disabled) pub(crate) authn_client: Mutex>>, + // Readiness registry, set once from `rest_prepare`; `OnceLock` = lock-free reads. + pub(crate) healthcheck_registry: OnceLock>, + // Built-once health router. `health_routes` registers OperationSpecs into the shared + // OpenAPI registry, so it must run exactly once per instance; caching the built router + // makes repeat `build_router`/`rebuild_and_cache_router` calls idempotent (no duplicate + // handler_id/route registration and its ERROR log). + pub(crate) health_router: OnceLock, // Duplicate detection (per (method, path) and per handler id) pub(crate) registered_routes: DashMap<(Method, String), ()>, @@ -92,24 +99,103 @@ impl Default for ApiGateway { router_cache: RouterCache::new(default_router), final_router: Mutex::new(None), authn_client: Mutex::new(None), + healthcheck_registry: OnceLock::new(), + health_router: OnceLock::new(), registered_routes: DashMap::new(), registered_handlers: DashMap::new(), } } } +// Built-in health-probe paths, shared by route registration and auth policy so they can't drift. +const HEALTH_DETAIL_PATH: &str = "/health"; +const HEALTHZ_PATH: &str = "/healthz"; +const READYZ_PATH: &str = "/readyz"; + impl ApiGateway { - fn apply_prefix_nesting(mut router: Router, prefix: &str) -> Router { - if prefix.is_empty() { - return router; + /// Nest the middleware-wrapped `router` under `prefix`, then merge the + /// pre-built `health_router` at root (always unprefixed). + /// + /// Both routers must have their own middleware already applied by the caller: + /// `router`'s auth matching is keyed on unprefixed `OperationBuilder` paths, so + /// its layers must run before `nest()` strips the prefix; health must stay at + /// root regardless of `prefix`, hence merged after nesting rather than nested. + fn apply_prefix_nesting(router: Router, prefix: &str, health_router: Router) -> Router { + let nested = if prefix.is_empty() { + router + } else { + Router::new().nest(prefix, router) + }; + + health_router.merge(nested) + } + + /// Health router with its own middleware stack, ready to merge via [`apply_prefix_nesting`]. + /// + /// Built once and cached: [`health_routes`](Self::health_routes) registers `OperationSpec`s + /// into the shared `OpenAPI` registry, so re-running it would trip duplicate detection. The + /// health deps (`hc_registry`, `authn_client`, `healthcheck_timeout`) are all fixed after + /// init/`rest_prepare`, so the first build is authoritative; later calls reuse the clone. + fn build_health_router( + &self, + hc_registry: Arc, + authn_client: Option>, + healthcheck_timeout: Duration, + ) -> Result { + if let Some(cached) = self.health_router.get() { + return Ok(cached.clone()); } + let built = self.apply_middleware_stack( + self.health_routes(hc_registry, healthcheck_timeout), + authn_client, + )?; + // First build wins; a concurrent racer's `set` is a no-op (router builds are + // sequential during lifecycle setup, so this is a belt-and-braces guard). + drop(self.health_router.set(built.clone())); + Ok(built) + } - let top = Router::new() - .route("/health", get(web::health_check)) - .route("/healthz", get(|| async { "ok" })); + /// Health-probe routes, defined via [`OperationBuilder`] so each route's auth intent is + /// co-located with its definition (and flows into the route policy via `OperationSpec`, + /// not a hand-maintained set). `/health` is authenticated (per-component detail); `/healthz` + /// and `/readyz` are public (status code only). Deps injected via `Extension`. + fn health_routes( + &self, + hc_registry: Arc, + healthcheck_timeout: Duration, + ) -> Router { + use http::StatusCode; + + let router = OperationBuilder::get(HEALTH_DETAIL_PATH) + .tag("Health Check") + .summary("Detailed readiness report") + .authenticated() + .no_license_required() + .handler(web::health_detail) + .json_response(StatusCode::OK, "Aggregate healthy or degraded") + .json_response(StatusCode::SERVICE_UNAVAILABLE, "Aggregate unhealthy") + .register(Router::new(), self); + + let router = OperationBuilder::get(HEALTHZ_PATH) + .tag("Health Check") + .summary("Liveness probe") + .public() + .handler(|| async { "ok" }) + .text_response(StatusCode::OK, "Service is live", "text/plain") + .register(router, self); - router = Router::new().nest(prefix, router); - top.merge(router) + let router = OperationBuilder::get(READYZ_PATH) + .tag("Health Check") + .summary("Readiness probe") + .public() + .handler(web::readyz_check) + .no_content_response(StatusCode::OK, "Ready (healthy or degraded)") + .no_content_response(StatusCode::SERVICE_UNAVAILABLE, "Not ready (unhealthy)") + .register(router, self); + + router + .layer(Extension(hc_registry)) + .layer(Extension(web::HealthcheckTimeout(healthcheck_timeout))) } /// Create a new `ApiGateway` instance with the given configuration @@ -122,6 +208,8 @@ impl ApiGateway { router_cache: RouterCache::new(default_router), final_router: Mutex::new(None), authn_client: Mutex::new(None), + healthcheck_registry: OnceLock::new(), + health_router: OnceLock::new(), registered_routes: DashMap::new(), registered_handlers: DashMap::new(), } @@ -157,10 +245,9 @@ impl ApiGateway { let mut authenticated_routes = std::collections::HashSet::new(); let mut public_routes = std::collections::HashSet::new(); - // Always mark built-in health check routes as public - public_routes.insert((Method::GET, "/health".to_owned())); - public_routes.insert((Method::GET, "/healthz".to_owned())); - + // Health probes (`/health` authenticated, `/healthz` + `/readyz` public) declare + // their auth via `OperationBuilder` in `health_routes`, so they flow into the sets + // through the `operation_specs` loop below like every other route. public_routes.insert((Method::GET, "/docs".to_owned())); public_routes.insert((Method::GET, "/openapi.json".to_owned())); @@ -256,7 +343,11 @@ impl ApiGateway { // 14) Propagate MatchedPath to response extensions (route_layer — innermost). // This copies MatchedPath from the request (populated by Axum route matching) // into the response so outer layer() middleware (metrics) can read it. - router = router.route_layer(from_fn(middleware::http_metrics::propagate_matched_path)); + // `route_layer` panics on a routeless router — reachable when no REST provider + // has registered anything yet. + if router.has_routes() { + router = router.route_layer(from_fn(middleware::http_metrics::propagate_matched_path)); + } let config = self.get_cached_config(); @@ -484,19 +575,28 @@ impl ApiGateway { } tracing::debug!("Building new router (standalone/fallback mode)"); - // In standalone mode (no REST pipeline), register both health endpoints here. - // In normal operation, rest_prepare() registers these instead. - let mut router = Router::new() - .route("/health", get(web::health_check)) - .route("/healthz", get(|| async { "ok" })); + // Standalone/fallback path; normally rest_prepare()/rest_finalize() wire health. + let hc_registry = self.healthcheck_registry.get().cloned().unwrap_or_else(|| { + tracing::warn!( + "healthcheck_registry not set (rest_prepare was not called); \ + building standalone router with an empty registry" + ); + Arc::new(toolkit::RestHealthcheckRegistry::new()) + }); - // Apply all middleware layers including auth, above the router + // No "main" routes here — the empty router is tolerated (`has_routes` guard) + // and health is merged in unprefixed with its own middleware. let authn_client = self.authn_client.lock().clone(); - router = self.apply_middleware_stack(router, authn_client)?; + let router = self.apply_middleware_stack(Router::new(), authn_client.clone())?; let config = self.get_cached_config(); let prefix = Self::normalize_prefix_path(&config.prefix_path)?; - router = Self::apply_prefix_nesting(router, &prefix); + let health_router = self.build_health_router( + hc_registry, + authn_client, + Duration::from_millis(config.healthcheck_timeout_ms), + )?; + let router = Self::apply_prefix_nesting(router, &prefix, health_router); // Cache the built router for future use self.router_cache.store(router.clone()); @@ -708,16 +808,15 @@ impl toolkit::contracts::ApiGatewayCapability for ApiGateway { &self, _ctx: &toolkit::context::GearCtx, router: axum::Router, + hc_registry: Arc, ) -> anyhow::Result { - // Add health check endpoints: - // - /health: detailed JSON response with status and timestamp - // - /healthz: simple "ok" liveness probe (Kubernetes-style) - let router = router - .route("/health", get(web::health_check)) - .route("/healthz", get(|| async { "ok" })); - - // You may attach global middlewares here (trace, compression, cors), but do not start server. - tracing::debug!("REST host prepared base router with health check endpoints"); + // Store for use when health routes are added in rest_finalize. A second set + // means rest_prepare ran twice — a lifecycle bug; fail fast rather than mask it. + if self.healthcheck_registry.set(hc_registry).is_err() { + anyhow::bail!("healthcheck_registry already set; rest_prepare called more than once"); + } + + tracing::debug!("REST host prepared base router"); Ok(router) } @@ -725,6 +824,7 @@ impl toolkit::contracts::ApiGatewayCapability for ApiGateway { &self, _ctx: &toolkit::context::GearCtx, mut router: axum::Router, + hc_registry: Arc, ) -> anyhow::Result { let config = self.get_cached_config(); @@ -732,18 +832,26 @@ impl toolkit::contracts::ApiGatewayCapability for ApiGateway { router = self.add_openapi_routes(router)?; } - // Apply middleware stack (including auth) to the final router + // Middleware on the main router before nesting (auth matching keyed on + // unprefixed OperationBuilder paths; layers run before nest() strips the prefix). tracing::debug!("Applying middleware stack to finalized router"); let authn_client = self.authn_client.lock().clone(); - router = self.apply_middleware_stack(router, authn_client)?; + router = self.apply_middleware_stack(router, authn_client.clone())?; let prefix = Self::normalize_prefix_path(&config.prefix_path)?; - router = Self::apply_prefix_nesting(router, &prefix); + let health_router = self.build_health_router( + hc_registry, + authn_client, + Duration::from_millis(config.healthcheck_timeout_ms), + )?; + router = Self::apply_prefix_nesting(router, &prefix, health_router); // Keep the finalized router to be used by `serve()` *self.final_router.lock() = Some(router.clone()); - tracing::info!("REST host finalized router with OpenAPI endpoints and auth middleware"); + tracing::info!( + "REST host finalized router with OpenAPI endpoints, health checks, and auth middleware" + ); Ok(router) } diff --git a/gears/system/api-gateway/src/web.rs b/gears/system/api-gateway/src/web.rs index d24f95210..74eab796e 100644 --- a/gears/system/api-gateway/src/web.rs +++ b/gears/system/api-gateway/src/web.rs @@ -1,10 +1,19 @@ use axum::{ + extract::Extension, http::StatusCode, - response::{Html, Json}, + response::{Html, IntoResponse, Json, Response}, routing::{MethodRouter, get}, }; use chrono::{SecondsFormat, Utc}; -use serde_json::{Value, json}; +use serde_json::json; +use std::sync::Arc; +use std::time::Duration; +use toolkit::RestHealthcheckRegistry; + +/// Per-check timeout (from `ApiGatewayConfig::healthcheck_timeout_ms`), injected as an +/// `Extension`. A newtype avoids colliding with any other `Duration` extension. +#[derive(Clone, Copy)] +pub struct HealthcheckTimeout(pub Duration); /// Returns a 501 Not Implemented handler for operations without implementations #[allow(dead_code)] @@ -20,11 +29,38 @@ pub fn placeholder_handler_501() -> MethodRouter { }) } -pub async fn health_check() -> Json { - Json(json!({ - "status": "healthy", - "timestamp": Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true) - })) +/// `GET /health`: runs all checks, returns JSON with per-component detail. Requires +/// auth (detail not for public callers). 200 (Healthy/Degraded), 503 (Unhealthy). +pub async fn health_detail( + Extension(registry): Extension>, + Extension(timeout): Extension, +) -> Response { + let report = registry.report(timeout.0).await; + let status_code = status_for_report(report.is_ready()); + let body = Json(json!({ + "status": report.status, + "timestamp": Utc::now().to_rfc3339_opts(SecondsFormat::Millis, true), + "components": report.components, + })); + (status_code, body).into_response() +} + +/// `GET /readyz`: runs all checks, public, status code only (no detail). +/// 200 (Healthy/Degraded), 503 (Unhealthy). +pub async fn readyz_check( + Extension(registry): Extension>, + Extension(timeout): Extension, +) -> StatusCode { + let report = registry.report(timeout.0).await; + status_for_report(report.is_ready()) +} + +fn status_for_report(report_ready: bool) -> StatusCode { + if report_ready { + StatusCode::OK + } else { + StatusCode::SERVICE_UNAVAILABLE + } } #[cfg(not(feature = "embed_elements"))] diff --git a/gears/system/api-gateway/tests/access_log_tests.rs b/gears/system/api-gateway/tests/access_log_tests.rs index 38fcf852f..b7fc42b83 100644 --- a/gears/system/api-gateway/tests/access_log_tests.rs +++ b/gears/system/api-gateway/tests/access_log_tests.rs @@ -349,7 +349,11 @@ async fn e2e_full_middleware_stack_logs_remote_addr() -> anyhow::Result<()> { .handler(get(e2e_handler)) .register(Router::new(), &api); - let app = api.rest_finalize(&ctx, router)?; + let app = api.rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + )?; // Set up capturing layer for the access log. let layer = CapturingLayer::default(); diff --git a/gears/system/api-gateway/tests/auth_middleware.rs b/gears/system/api-gateway/tests/auth_middleware.rs index 9f09b366a..c3b940048 100644 --- a/gears/system/api-gateway/tests/auth_middleware.rs +++ b/gears/system/api-gateway/tests/auth_middleware.rs @@ -228,7 +228,11 @@ async fn test_auth_disabled_mode() { // Finalize router (applies middleware) let router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); // Test protected route WITHOUT token (should work because auth is disabled) @@ -290,7 +294,11 @@ async fn test_public_routes_accessible() { // First call rest_prepare to add built-in routes let router = Router::new(); let router = api_gateway - .rest_prepare(&api_ctx, router) + .rest_prepare( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to prepare"); // Then register test gear routes @@ -301,7 +309,11 @@ async fn test_public_routes_accessible() { // Finally finalize let router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); // Test built-in health endpoints @@ -364,7 +376,11 @@ async fn test_public_routes_with_prefix_accessible() { // First call rest_prepare to add built-in routes let router = Router::new(); let router = api_gateway - .rest_prepare(&api_ctx, router) + .rest_prepare( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to prepare"); // Then register test gear routes @@ -375,7 +391,11 @@ async fn test_public_routes_with_prefix_accessible() { // Finally finalize let router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); // Test built-in health endpoints @@ -459,7 +479,11 @@ async fn test_middleware_always_inserts_security_ctx() { .expect("Failed to register routes"); let router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); // Make request to protected handler that extracts SecurityContext @@ -570,7 +594,11 @@ async fn test_route_pattern_matching_with_path_params() { .expect("Failed to register routes"); let router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); // Test that /tests/v1/api/users/123 is accessible (matches /tests/v1/api/users/{id}) @@ -638,7 +666,11 @@ async fn test_route_pattern_matching_with_prefix_path_params() { .expect("Failed to register routes"); let router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); // Test that /tests/v1/api/users/123 is accessible (matches /tests/v1/api/users/{id}) @@ -777,7 +809,11 @@ async fn create_router(config: serde_json::Value, mock: MockAuthNResolverClient) .expect("Failed to register routes"); api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize") } diff --git a/gears/system/api-gateway/tests/body_limit_tests.rs b/gears/system/api-gateway/tests/body_limit_tests.rs index e270a8a76..560eee491 100644 --- a/gears/system/api-gateway/tests/body_limit_tests.rs +++ b/gears/system/api-gateway/tests/body_limit_tests.rs @@ -116,7 +116,11 @@ async fn test_body_limit_configured() { .expect("Failed to register routes"); let _final_router = api_gateway - .rest_finalize(&ctx, router) + .rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize router"); // Verify router builds with custom body limit @@ -141,7 +145,11 @@ async fn test_body_limit_with_cors() { .expect("Failed to register routes"); let _final_router = api_gateway - .rest_finalize(&ctx, router) + .rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize router"); // Both CORS and body limit should be active @@ -180,7 +188,11 @@ async fn test_default_body_limit() { .expect("Failed to register routes"); let _final_router = api_gateway - .rest_finalize(&ctx, router) + .rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize router"); // Verify default body limit is applied (16MB) diff --git a/gears/system/api-gateway/tests/cors_tests.rs b/gears/system/api-gateway/tests/cors_tests.rs index f582b88e3..3430e7ef9 100644 --- a/gears/system/api-gateway/tests/cors_tests.rs +++ b/gears/system/api-gateway/tests/cors_tests.rs @@ -145,7 +145,11 @@ async fn test_cors_layer_builds_with_config() { // Build the final router with CORS middleware let _final_router = api_gateway - .rest_finalize(&ctx, router) + .rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize router"); // Verify router builds successfully with CORS enabled @@ -165,7 +169,11 @@ async fn test_cors_permissive_mode() { .expect("Failed to register routes"); let _final_router = api_gateway - .rest_finalize(&ctx, router) + .rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize router"); // Verify permissive CORS builds successfully @@ -199,7 +207,11 @@ async fn test_cors_disabled() { .expect("Failed to register routes"); let _final_router = api_gateway - .rest_finalize(&ctx, router) + .rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize router"); // Verify router builds without CORS layer diff --git a/gears/system/api-gateway/tests/health_endpoints.rs b/gears/system/api-gateway/tests/health_endpoints.rs new file mode 100644 index 000000000..213b038a7 --- /dev/null +++ b/gears/system/api-gateway/tests/health_endpoints.rs @@ -0,0 +1,264 @@ +#![allow(clippy::unwrap_used, clippy::expect_used)] + +//! Integration tests for health endpoints (/health, /healthz, /readyz). + +use async_trait::async_trait; +use axum::{ + Router, + body::Body, + http::{Request, StatusCode}, +}; +use serde_json::json; +use std::sync::Arc; +use toolkit::{ClientHub, Gear, Healthcheck, HealthcheckResult}; +use tower::ServiceExt; +use uuid::Uuid; + +// ---------- helpers ---------- + +/// A `ConfigProvider` backed by a single JSON blob, keyed by gear name. +struct JsonConfigProvider(serde_json::Value); + +impl toolkit::config::ConfigProvider for JsonConfigProvider { + fn get_gear_config(&self, gear: &str) -> Option<&serde_json::Value> { + self.0.get(gear) + } +} + +/// Build a fully finalized router (`rest_prepare` + `rest_finalize`, auth disabled) +/// mounted at `prefix` ("" for no prefix), with a healthcheck registry pre-populated +/// by `setup`. Exercises the same path production wiring uses, so tests catch +/// real routing bugs (e.g. health routes leaking into a prefixed subtree). +async fn build_router_with_prefix( + prefix: &str, + setup: impl FnOnce(&toolkit::RestHealthcheckRegistry), +) -> Router { + use api_gateway::ApiGateway; + use toolkit::contracts::ApiGatewayCapability; + + let registry = Arc::new(toolkit::RestHealthcheckRegistry::new()); + setup(®istry); + + let config = json!({ + "api-gateway": { + "config": { + "bind_addr": "0.0.0.0:0", + "auth_disabled": true, + "prefix_path": prefix, + } + } + }); + let ctx = toolkit::GearCtx::new( + "api-gateway", + Uuid::new_v4(), + Arc::new(JsonConfigProvider(config)), + Arc::new(ClientHub::new()), + tokio_util::sync::CancellationToken::new(), + ); + + let gw = ApiGateway::default(); + gw.init(&ctx).await.expect("init failed"); + let router = gw + .rest_prepare(&ctx, Router::new(), registry.clone()) + .expect("rest_prepare failed"); + gw.rest_finalize(&ctx, router, registry) + .expect("rest_finalize failed") +} + +/// Build a router with no prefix. See [`build_router_with_prefix`]. +async fn build_router(setup: impl FnOnce(&toolkit::RestHealthcheckRegistry)) -> Router { + build_router_with_prefix("", setup).await +} + +async fn get(router: Router, path: &str) -> axum::http::Response { + router + .oneshot(Request::builder().uri(path).body(Body::empty()).unwrap()) + .await + .unwrap() +} + +async fn body_json(resp: axum::http::Response) -> serde_json::Value { + let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + serde_json::from_slice(&bytes).unwrap() +} + +// ---------- check stubs ---------- + +struct HealthyCheck; +#[async_trait] +impl Healthcheck for HealthyCheck { + fn name(&self) -> &'static str { + "always-healthy" + } +} + +struct DegradedCheck; +#[async_trait] +impl Healthcheck for DegradedCheck { + fn name(&self) -> &'static str { + "always-degraded" + } + async fn check(&self) -> HealthcheckResult { + HealthcheckResult::degraded("cache warming up") + } +} + +struct UnhealthyCheck; +#[async_trait] +impl Healthcheck for UnhealthyCheck { + fn name(&self) -> &'static str { + "always-unhealthy" + } + async fn check(&self) -> HealthcheckResult { + HealthcheckResult::unhealthy("database unreachable") + } +} + +// ---------- /healthz ---------- + +#[tokio::test] +async fn healthz_returns_200_with_no_checks() { + let router = build_router(|_| {}).await; + let resp = get(router, "/healthz").await; + assert_eq!(resp.status(), StatusCode::OK); +} + +#[tokio::test] +async fn healthz_returns_200_even_when_unhealthy_check_registered() { + let router = build_router(|reg| { + reg.register("gear", Arc::new(UnhealthyCheck)); + }) + .await; + let resp = get(router, "/healthz").await; + assert_eq!(resp.status(), StatusCode::OK); +} + +// ---------- /readyz ---------- + +#[tokio::test] +async fn readyz_returns_200_with_no_checks() { + let router = build_router(|_| {}).await; + let resp = get(router, "/readyz").await; + assert_eq!(resp.status(), StatusCode::OK); +} + +#[tokio::test] +async fn readyz_returns_200_when_all_healthy() { + let router = build_router(|reg| { + reg.register("gear", Arc::new(HealthyCheck)); + }) + .await; + let resp = get(router, "/readyz").await; + assert_eq!(resp.status(), StatusCode::OK); +} + +#[tokio::test] +async fn readyz_returns_200_when_degraded() { + let router = build_router(|reg| { + reg.register("gear", Arc::new(DegradedCheck)); + }) + .await; + let resp = get(router, "/readyz").await; + assert_eq!(resp.status(), StatusCode::OK); +} + +#[tokio::test] +async fn readyz_returns_503_when_unhealthy() { + let router = build_router(|reg| { + reg.register("gear", Arc::new(UnhealthyCheck)); + }) + .await; + let resp = get(router, "/readyz").await; + assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE); +} + +// ---------- /health ---------- + +#[tokio::test] +async fn health_returns_json_with_components() { + let router = build_router(|reg| { + reg.register("my-gear", Arc::new(HealthyCheck)); + }) + .await; + let resp = get(router, "/health").await; + assert_eq!(resp.status(), StatusCode::OK); + let body = body_json(resp).await; + assert_eq!(body["status"], "healthy"); + assert!(body["timestamp"].is_string()); + let components = body["components"].as_array().unwrap(); + assert_eq!(components.len(), 1); + assert_eq!(components[0]["gear"], "my-gear"); + assert_eq!(components[0]["status"], "healthy"); +} + +#[tokio::test] +async fn health_returns_503_when_unhealthy() { + let router = build_router(|reg| { + reg.register("gear", Arc::new(UnhealthyCheck)); + }) + .await; + let resp = get(router, "/health").await; + assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE); + let body = body_json(resp).await; + assert_eq!(body["status"], "unhealthy"); + assert!( + !body["components"] + .as_array() + .unwrap() + .iter() + .any(|c| c["status"] == "healthy"), + "no component should report healthy when aggregate is unhealthy" + ); +} + +#[tokio::test] +async fn health_returns_degraded_status_when_degraded() { + let router = build_router(|reg| { + reg.register("gear", Arc::new(DegradedCheck)); + }) + .await; + let resp = get(router, "/health").await; + assert_eq!(resp.status(), StatusCode::OK); + let body = body_json(resp).await; + assert_eq!(body["status"], "degraded"); + assert!( + !body["components"] + .as_array() + .unwrap() + .iter() + .any(|c| c["status"] == "unhealthy"), + "no component should report unhealthy when aggregate is only degraded" + ); +} + +// ---------- outside prefix ---------- + +#[tokio::test] +async fn health_endpoints_accessible_outside_prefix() { + let router = build_router_with_prefix("/cf", |_| {}).await; + + // Health endpoints must be reachable at root paths, not nested under /cf. + for path in ["/healthz", "/readyz", "/health"] { + let resp = get(router.clone(), path).await; + assert_eq!(resp.status(), StatusCode::OK, "GET {path} should be 200"); + } +} + +#[tokio::test] +async fn health_endpoints_not_duplicated_under_prefix() { + let router = build_router_with_prefix("/cf", |_| {}).await; + + // Health routes must NOT also be reachable under the prefix — regression test + // for a bug where rest_prepare merged health routes into the router before it + // got nested, making them reachable at both `/health` and `/cf/health`. + for path in ["/cf/healthz", "/cf/readyz", "/cf/health"] { + let resp = get(router.clone(), path).await; + assert_eq!( + resp.status(), + StatusCode::NOT_FOUND, + "GET {path} should be 404 (health routes are root-only)" + ); + } +} diff --git a/gears/system/api-gateway/tests/http_metrics_tests.rs b/gears/system/api-gateway/tests/http_metrics_tests.rs index 48a02d914..b75326b24 100644 --- a/gears/system/api-gateway/tests/http_metrics_tests.rs +++ b/gears/system/api-gateway/tests/http_metrics_tests.rs @@ -158,7 +158,11 @@ async fn metrics_capture_successful_request() -> Result<()> { .json_response(StatusCode::OK, "OK") .handler(axum::routing::get(ok_handler)) .register(Router::new(), &api); - let app = api.rest_finalize(&ctx, router)?; + let app = api.rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + )?; let res = app .oneshot( @@ -212,7 +216,11 @@ async fn metrics_capture_mime_rejection() -> Result<()> { .json_response(StatusCode::OK, "OK") .handler(axum::routing::post(ok_handler)) .register(Router::new(), &api); - let app = api.rest_finalize(&ctx, router)?; + let app = api.rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + )?; let res = app .oneshot( @@ -275,7 +283,11 @@ async fn metrics_capture_rate_limit() -> Result<()> { .json_response(StatusCode::OK, "OK") .handler(axum::routing::get(ok_handler)) .register(Router::new(), &api); - let app = api.rest_finalize(&ctx, router)?; + let app = api.rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + )?; // First request — succeeds and consumes the token let res1 = app @@ -335,7 +347,11 @@ async fn metrics_route_attribute_uses_template() -> Result<()> { .json_response(StatusCode::OK, "OK") .handler(axum::routing::get(ok_handler)) .register(Router::new(), &api); - let app = api.rest_finalize(&ctx, router)?; + let app = api.rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + )?; let res = app .oneshot( @@ -387,7 +403,11 @@ async fn metrics_unmatched_route() -> Result<()> { .json_response(StatusCode::OK, "OK") .handler(axum::routing::get(ok_handler)) .register(Router::new(), &api); - let app = api.rest_finalize(&ctx, router)?; + let app = api.rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + )?; let res = app .oneshot( @@ -453,7 +473,11 @@ async fn metrics_prefix_applied_to_instrument_names() -> Result<()> { .json_response(StatusCode::OK, "OK") .handler(axum::routing::get(ok_handler)) .register(Router::new(), &api); - let app = api.rest_finalize(&ctx, router)?; + let app = api.rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + )?; let res = app .oneshot( diff --git a/gears/system/api-gateway/tests/license_middleware.rs b/gears/system/api-gateway/tests/license_middleware.rs index 045645c49..0eaf552fc 100644 --- a/gears/system/api-gateway/tests/license_middleware.rs +++ b/gears/system/api-gateway/tests/license_middleware.rs @@ -157,7 +157,11 @@ async fn rejects_non_base_feature_requirement() { .expect("Failed to register routes"); router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); let response = router @@ -214,7 +218,11 @@ async fn rejects_non_base_feature_requirement_with_prefix() { .expect("Failed to register routes"); router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); let response = router @@ -269,7 +277,11 @@ async fn allows_base_feature_requirement() { .expect("Failed to register routes"); let router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); let response = router @@ -312,7 +324,11 @@ async fn allows_base_feature_requirement_with_prefix() { .expect("Failed to register routes"); let router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); let response = router @@ -354,7 +370,11 @@ async fn allows_no_license_requirement() { .expect("Failed to register routes"); let router = api_gateway - .rest_finalize(&api_ctx, router) + .rest_finalize( + &api_ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize"); let response = router diff --git a/gears/system/api-gateway/tests/middleware_order.rs b/gears/system/api-gateway/tests/middleware_order.rs index 12b682089..c56a1a3e7 100644 --- a/gears/system/api-gateway/tests/middleware_order.rs +++ b/gears/system/api-gateway/tests/middleware_order.rs @@ -84,7 +84,11 @@ async fn real_middlewares_observe_documented_order() -> Result<()> { .register(router, &api); // Apply the real gateway middleware stack. - let app = api.rest_finalize(&ctx, router)?; + let app = api.rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + )?; // -------------------- // Req1: invalid Content-Type -> should be rejected by MIME validation (BAD_REQUEST / 400), @@ -203,7 +207,11 @@ async fn real_middlewares_observe_documented_order_with_prefix() -> Result<()> { .register(router, &api); // Apply the real gateway middleware stack. - let app = api.rest_finalize(&ctx, router)?; + let app = api.rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + )?; // -------------------- // Req1: invalid Content-Type -> should be rejected by MIME validation (BAD_REQUEST / 400), diff --git a/gears/system/api-gateway/tests/rate_limit_tests.rs b/gears/system/api-gateway/tests/rate_limit_tests.rs index dd4b53a4f..7f8a3ce11 100644 --- a/gears/system/api-gateway/tests/rate_limit_tests.rs +++ b/gears/system/api-gateway/tests/rate_limit_tests.rs @@ -173,7 +173,11 @@ async fn test_rate_limit_enforcement() { // Build the final router with middleware let _final_router = api_gateway - .rest_finalize(&ctx, router) + .rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize router"); // Note: Full HTTP testing would require starting a server and making real requests @@ -291,7 +295,11 @@ async fn test_rate_limit_returns_canonical_problem_with_headers() { .expect("Failed to register routes"); let app = api_gateway - .rest_finalize(&ctx, router) + .rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize router"); // First request consumes the only token. @@ -389,7 +397,11 @@ async fn test_in_flight_limit_returns_canonical_service_unavailable() { .register(Router::new(), &api_gateway); let app = api_gateway - .rest_finalize(&ctx, router) + .rest_finalize( + &ctx, + router, + Arc::new(toolkit::RestHealthcheckRegistry::new()), + ) .expect("Failed to finalize router"); // Start one slow request and immediately fire a second one while the first holds the only permit. diff --git a/libs/toolkit/src/contracts.rs b/libs/toolkit/src/contracts.rs index 127476d14..f0df447c3 100644 --- a/libs/toolkit/src/contracts.rs +++ b/libs/toolkit/src/contracts.rs @@ -60,22 +60,39 @@ pub trait RestApiCapability: Send + Sync { router: Router, openapi: &dyn OpenApiRegistry, ) -> anyhow::Result; + + /// Readiness healthcheck for this gear, run on every `/readyz` and `/health`. + /// `None` (default) opts out. Since `cf-gears-toolkit` v0.6.13. + /// TODO: gRPC healthcheck hook is a follow-up; this PR is REST-only. + fn healthcheck( + &self, + _ctx: &crate::context::GearCtx, + ) -> Option> { + None + } } /// API Gateway capability: handles gateway hosting with prepare/finalize phases. /// Must be sync. Runs during REST phase, but doesn't start the server. #[allow(dead_code)] pub trait ApiGatewayCapability: Send + Sync + 'static { - /// Prepare a base Router (e.g., global middlewares, /healthz) and optionally touch `OpenAPI` meta. - /// Do NOT start the server here. + /// Prepare a base Router (e.g., global middlewares) and optionally touch `OpenAPI` meta. + /// Do NOT start the server here. `hc_registry` is the runtime's shared healthcheck + /// registry, passed explicitly (not via `ClientHub`, which is for inter-gear clients). /// /// # Errors /// Returns an error if router preparation fails. - fn rest_prepare(&self, ctx: &crate::context::GearCtx, router: Router) - -> anyhow::Result; + fn rest_prepare( + &self, + ctx: &crate::context::GearCtx, + router: Router, + hc_registry: std::sync::Arc, + ) -> anyhow::Result; - /// Finalize before start: attach /openapi.json, /docs, persist the Router internally if needed. - /// Do NOT start the server here. + /// Finalize before start: attach /openapi.json, /docs, health endpoints, persist the + /// Router internally if needed. Do NOT start the server here. + /// + /// `hc_registry` is the same registry instance passed to [`rest_prepare`](Self::rest_prepare). /// /// # Errors /// Returns an error if router finalization fails. @@ -83,6 +100,7 @@ pub trait ApiGatewayCapability: Send + Sync + 'static { &self, ctx: &crate::context::GearCtx, router: Router, + hc_registry: std::sync::Arc, ) -> anyhow::Result; // Return OpenAPI registry of the gear, e.g., to register endpoints diff --git a/libs/toolkit/src/healthcheck.rs b/libs/toolkit/src/healthcheck.rs new file mode 100644 index 000000000..0ac4a7dd4 --- /dev/null +++ b/libs/toolkit/src/healthcheck.rs @@ -0,0 +1,642 @@ +//! REST readiness healthcheck infrastructure for probing gear health. +//! +//! This module provides concurrent healthcheck execution with timeout protection. +//! Gears implement [`Healthcheck`] and register themselves via [`RestHealthcheckRegistry`]. +//! The API Gateway calls [`report()`](RestHealthcheckRegistry::report) on `/health` and `/readyz` +//! requests to aggregate per-component readiness status. +//! +//! # Usage +//! +//! ```rust,ignore +//! use async_trait::async_trait; +//! use std::sync::Arc; +//! use toolkit::{Healthcheck, HealthcheckResult, contracts::RestApiCapability}; +//! +//! struct MyHealthcheck; +//! +//! #[async_trait] +//! impl Healthcheck for MyHealthcheck { +//! fn name(&self) -> &'static str { +//! "my-gear-readiness" +//! } +//! +//! async fn check(&self) -> HealthcheckResult { +//! HealthcheckResult::healthy() +//! } +//! } +//! +//! impl RestApiCapability for MyGear { +//! fn healthcheck( +//! &self, +//! _ctx: &toolkit::context::GearCtx, +//! ) -> Option> { +//! Some(Arc::new(MyHealthcheck)) +//! } +//! +//! fn register_rest( +//! &self, +//! ctx: &toolkit::context::GearCtx, +//! router: axum::Router, +//! openapi: &dyn toolkit::contracts::OpenApiRegistry, +//! ) -> anyhow::Result { +//! Ok(router) +//! } +//! } +//! ``` +//! +//! # Kubernetes probe configuration +//! +//! ```yaml +//! livenessProbe: +//! httpGet: +//! path: /healthz +//! port: http +//! periodSeconds: 10 +//! timeoutSeconds: 2 +//! failureThreshold: 3 +//! +//! readinessProbe: +//! httpGet: +//! path: /readyz +//! port: http +//! periodSeconds: 5 +//! timeoutSeconds: 2 +//! failureThreshold: 3 +//! ``` +//! +//! `/healthz` is liveness — always shallow, never runs user checks. +//! `/readyz` is readiness — runs user REST healthchecks and removes the pod +//! from traffic when unhealthy without triggering a restart. +//! `/health` is the detailed diagnostic endpoint — it requires authentication +//! since its JSON body includes per-component check names and messages. + +use async_trait::async_trait; +use parking_lot::RwLock; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::sync::Mutex as AsyncMutex; +use tokio_util::sync::CancellationToken; + +/// Single-check and aggregate readiness status. +/// Serialized lowercase into `/health`/`/readyz`; variants are a stable API contract. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "lowercase")] +pub enum HealthcheckStatus { + Healthy, + Degraded, + Unhealthy, +} + +/// Result of one [`Healthcheck::check`]; fields are part of the stable `/health` JSON contract. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HealthcheckResult { + pub status: HealthcheckStatus, + pub message: Option, +} + +impl HealthcheckResult { + #[must_use] + pub fn healthy() -> Self { + Self { + status: HealthcheckStatus::Healthy, + message: None, + } + } + + #[must_use] + pub fn degraded(message: impl Into) -> Self { + Self { + status: HealthcheckStatus::Degraded, + message: Some(message.into()), + } + } + + #[must_use] + pub fn unhealthy(message: impl Into) -> Self { + Self { + status: HealthcheckStatus::Unhealthy, + message: Some(message.into()), + } + } +} + +impl Default for HealthcheckResult { + fn default() -> Self { + Self::healthy() + } +} + +/// Readiness probe implemented by a gear. +/// +/// `name` must be an explicit human-readable id (no type paths); it is exposed on +/// `/health`. `check` must be cancellation-safe and should not leak secrets in its +/// message (the registry sanitises anyway); panics and per-check timeouts are caught +/// and mapped to [`HealthcheckStatus::Unhealthy`]. Default `check` returns healthy. +#[async_trait] +pub trait Healthcheck: Send + Sync + 'static { + /// Human-readable check name, exposed verbatim in `/health` JSON. + fn name(&self) -> &'static str; + + async fn check(&self) -> HealthcheckResult { + HealthcheckResult::healthy() + } +} + +/// One gear's healthcheck result. All fields are part of the stable `/health` JSON contract. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HealthcheckComponentReport { + pub gear: String, + pub check: String, + pub status: HealthcheckStatus, + /// Sanitized (see [`sanitize_message`]). + pub message: Option, + pub latency_ms: u64, +} + +/// Aggregate report from [`RestHealthcheckRegistry::report`]; stable `/health`/`/readyz` JSON contract. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HealthcheckReport { + pub status: HealthcheckStatus, + pub components: Vec, +} + +impl HealthcheckReport { + /// Ready unless `Unhealthy` (`Healthy` and `Degraded` both keep the pod in rotation). + #[must_use] + pub fn is_ready(&self) -> bool { + self.status != HealthcheckStatus::Unhealthy + } +} + +struct RegistryEntry { + gear: String, + check: Arc, +} + +/// Bursty probes within this window reuse the last report instead of re-running checks. +/// Kept > the default per-check timeout (500 ms) so a timed-out check still buffers. +const REPORT_CACHE_TTL: Duration = Duration::from_secs(2); + +/// Holds the REST healthchecks registered during REST wiring; the gateway calls +/// [`report`](Self::report) on every `/readyz` and `/health` request. +#[derive(Default)] +pub struct RestHealthcheckRegistry { + entries: RwLock>, + cached_report: AsyncMutex>, + /// Runtime shutdown token; aborts in-flight checks. Default never fires. + cancel: CancellationToken, +} + +impl RestHealthcheckRegistry { + #[must_use] + pub fn new() -> Self { + Self::default() + } + + /// Registry whose in-flight checks are aborted when `cancel` fires. + #[must_use] + pub fn with_cancellation(cancel: CancellationToken) -> Self { + Self { + cancel, + ..Self::default() + } + } + + /// Register a healthcheck for the given gear. + pub fn register(&self, gear: impl Into, check: Arc) { + self.entries.write().push(RegistryEntry { + gear: gear.into(), + check, + }); + } + + /// Run all checks concurrently, aggregate, and cache for [`REPORT_CACHE_TTL`]. + /// Panics and per-check timeouts map to [`HealthcheckStatus::Unhealthy`]. The cache + /// lock is held across the compute so concurrent cache-misses coalesce into one fan-out. + pub async fn report(&self, timeout_per_check: Duration) -> HealthcheckReport { + let mut guard = self.cached_report.lock().await; + if let Some((ts, cached)) = guard.as_ref() + && ts.elapsed() < REPORT_CACHE_TTL + { + return cached.clone(); + } + + let entries: Vec<(_, _)> = { + let r = self.entries.read(); + r.iter() + .map(|e| (e.gear.clone(), e.check.clone())) + .collect() + }; + + let report = if entries.is_empty() { + HealthcheckReport { + status: HealthcheckStatus::Healthy, + components: vec![], + } + } else { + let cancel = &self.cancel; + let checks = entries + .into_iter() + .map(|(gear, check)| run_one_check(gear, check, timeout_per_check, cancel)); + let components = futures_util::future::join_all(checks).await; + let aggregate = compute_aggregate(&components); + HealthcheckReport { + status: aggregate, + components, + } + }; + + *guard = Some((Instant::now(), report.clone())); + report + } +} + +async fn run_one_check( + gear: String, + check: Arc, + timeout_per_check: Duration, + cancel: &CancellationToken, +) -> HealthcheckComponentReport { + let name = check.name(); + let start = Instant::now(); + + let mut handle = tokio::spawn(async move { check.check().await }); + + let (status, message) = tokio::select! { + result = &mut handle => match result { + Ok(r) => (r.status, r.message.as_deref().map(sanitize_message)), + Err(join_err) => { + tracing::error!(gear = %gear, check = name, error = %join_err, "healthcheck task panicked"); + ( + HealthcheckStatus::Unhealthy, + Some("health check failed".to_owned()), + ) + } + }, + () = tokio::time::sleep(timeout_per_check) => { + handle.abort(); + if let Err(join_err) = handle.await + && !join_err.is_cancelled() + { + tracing::warn!(gear = %gear, check = name, error = %join_err, "healthcheck task join error after timeout abort"); + } + tracing::warn!(gear = %gear, check = name, timeout_ms = timeout_per_check.as_millis(), "healthcheck timed out"); + ( + HealthcheckStatus::Unhealthy, + Some("health check timed out".to_owned()), + ) + } + () = cancel.cancelled() => { + handle.abort(); + if let Err(join_err) = handle.await + && !join_err.is_cancelled() + { + tracing::warn!(gear = %gear, check = name, error = %join_err, "healthcheck task join error after shutdown abort"); + } + tracing::warn!(gear = %gear, check = name, "healthcheck cancelled by runtime shutdown"); + ( + HealthcheckStatus::Unhealthy, + Some("health check cancelled".to_owned()), + ) + } + }; + + HealthcheckComponentReport { + gear, + check: name.to_owned(), + status, + message, + latency_ms: u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX), + } +} + +fn compute_aggregate(components: &[HealthcheckComponentReport]) -> HealthcheckStatus { + let mut status = HealthcheckStatus::Healthy; + for c in components { + match c.status { + HealthcheckStatus::Unhealthy => return HealthcheckStatus::Unhealthy, + HealthcheckStatus::Degraded => { + status = HealthcheckStatus::Degraded; + } + HealthcheckStatus::Healthy => {} + } + } + status +} + +const MAX_MESSAGE_LEN: usize = 256; + +// Deliberately broad blocklist: false positives (benign messages collapsed to the +// generic string) are accepted to never leak secrets/DSNs on the public endpoint. +static SUSPICIOUS_SUBSTRINGS: &[&str] = &[ + "password", + "passwd", + "secret", + "token", + "bearer", + "authorization", + "api_key", + "apikey", + "private_key", + "private", + "credential", + "access_key", + "aws_", + "client_secret", + "postgres://", + "postgresql://", + "mysql://", + "sqlite://", + "mongodb://", + "redis://", + "amqp://", + "jdbc:", + "tenant", + "select ", + "insert ", + "update ", + "delete ", + "panic", + "stack backtrace", +]; + +fn sanitize_message(msg: &str) -> String { + if msg.len() > MAX_MESSAGE_LEN { + return "health check failed".to_owned(); + } + let lower = msg.to_lowercase(); + for sub in SUSPICIOUS_SUBSTRINGS { + if lower.contains(sub) { + return "health check failed".to_owned(); + } + } + msg.to_owned() +} + +#[cfg(test)] +#[cfg_attr(coverage_nightly, coverage(off))] +mod tests { + use super::*; + use std::sync::atomic::{AtomicBool, Ordering}; + use std::time::Duration; + use tokio::sync::Notify; + + struct AlwaysDegraded; + #[async_trait] + impl Healthcheck for AlwaysDegraded { + fn name(&self) -> &'static str { + "always-degraded" + } + async fn check(&self) -> HealthcheckResult { + HealthcheckResult::degraded("cache warming") + } + } + + struct AlwaysUnhealthy; + #[async_trait] + impl Healthcheck for AlwaysUnhealthy { + fn name(&self) -> &'static str { + "always-unhealthy" + } + async fn check(&self) -> HealthcheckResult { + HealthcheckResult::unhealthy("database unreachable") + } + } + + struct SlowCheck; + #[async_trait] + impl Healthcheck for SlowCheck { + fn name(&self) -> &'static str { + "slow-check" + } + async fn check(&self) -> HealthcheckResult { + tokio::time::sleep(Duration::from_secs(10)).await; + HealthcheckResult::healthy() + } + } + + struct AbortTrackingCheck { + entered: Arc, + dropped: Arc, + } + + #[async_trait] + impl Healthcheck for AbortTrackingCheck { + fn name(&self) -> &'static str { + "abort-tracking-check" + } + async fn check(&self) -> HealthcheckResult { + struct DropFlag(Arc); + impl Drop for DropFlag { + fn drop(&mut self) { + self.0.store(true, Ordering::SeqCst); + } + } + + let _drop_flag = DropFlag(self.dropped.clone()); + self.entered.notify_one(); + std::future::pending::().await + } + } + + struct PanickingCheck; + #[async_trait] + impl Healthcheck for PanickingCheck { + fn name(&self) -> &'static str { + "panicking-check" + } + async fn check(&self) -> HealthcheckResult { + panic!("intentional panic in healthcheck"); + } + } + + #[test] + fn healthcheck_result_default_is_healthy() { + let r = HealthcheckResult::default(); + assert_eq!(r.status, HealthcheckStatus::Healthy); + assert!(r.message.is_none()); + } + + #[tokio::test] + async fn empty_registry_report_is_healthy_and_ready() { + let registry = RestHealthcheckRegistry::new(); + let report = registry.report(Duration::from_millis(500)).await; + assert_eq!(report.status, HealthcheckStatus::Healthy); + assert!(report.is_ready()); + assert!(report.components.is_empty()); + } + + #[tokio::test] + async fn unhealthy_component_makes_aggregate_unhealthy_and_not_ready() { + let registry = RestHealthcheckRegistry::new(); + registry.register("my-gear", Arc::new(AlwaysUnhealthy)); + let report = registry.report(Duration::from_millis(500)).await; + assert_eq!(report.status, HealthcheckStatus::Unhealthy); + assert!(!report.is_ready()); + } + + #[tokio::test] + async fn mixed_degraded_and_unhealthy_aggregate_is_unhealthy() { + let registry = RestHealthcheckRegistry::new(); + registry.register("degraded-gear", Arc::new(AlwaysDegraded)); + registry.register("unhealthy-gear", Arc::new(AlwaysUnhealthy)); + let report = registry.report(Duration::from_millis(500)).await; + assert_eq!( + report.status, + HealthcheckStatus::Unhealthy, + "unhealthy must take priority over degraded" + ); + assert!(!report.is_ready()); + } + + #[tokio::test] + async fn slow_check_times_out_and_becomes_unhealthy() { + let registry = RestHealthcheckRegistry::new(); + registry.register("slow-gear", Arc::new(SlowCheck)); + let report = registry.report(Duration::from_millis(100)).await; + assert_eq!(report.status, HealthcheckStatus::Unhealthy); + let comp = &report.components[0]; + assert_eq!(comp.status, HealthcheckStatus::Unhealthy); + assert_eq!(comp.message.as_deref(), Some("health check timed out")); + } + + #[tokio::test] + async fn timed_out_check_is_aborted() { + let registry = RestHealthcheckRegistry::new(); + let entered = Arc::new(Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + registry.register( + "slow-gear", + Arc::new(AbortTrackingCheck { + entered: entered.clone(), + dropped: dropped.clone(), + }), + ); + + let report_task = + tokio::spawn(async move { registry.report(Duration::from_millis(10)).await }); + entered.notified().await; + let report = report_task.await.expect("report task panicked"); + + assert_eq!(report.status, HealthcheckStatus::Unhealthy); + assert!( + dropped.load(Ordering::SeqCst), + "timed-out healthcheck future must be dropped after abort" + ); + } + + #[tokio::test] + async fn cancelled_registry_aborts_in_flight_check_and_reports_unhealthy() { + let cancel = CancellationToken::new(); + let registry = RestHealthcheckRegistry::with_cancellation(cancel.clone()); + let entered = Arc::new(Notify::new()); + let dropped = Arc::new(AtomicBool::new(false)); + registry.register( + "slow-gear", + Arc::new(AbortTrackingCheck { + entered: entered.clone(), + dropped: dropped.clone(), + }), + ); + + // Long per-check timeout so cancellation, not the timeout, ends the check. + let report_task = + tokio::spawn(async move { registry.report(Duration::from_mins(1)).await }); + entered.notified().await; + cancel.cancel(); + let report = report_task.await.expect("report task panicked"); + + assert_eq!(report.status, HealthcheckStatus::Unhealthy); + assert_eq!( + report.components[0].message.as_deref(), + Some("health check cancelled") + ); + assert!( + dropped.load(Ordering::SeqCst), + "cancelled healthcheck future must be dropped after abort" + ); + } + + #[tokio::test] + async fn panicking_check_becomes_unhealthy_without_panicking_caller() { + let registry = RestHealthcheckRegistry::new(); + registry.register("panic-gear", Arc::new(PanickingCheck)); + let report = registry.report(Duration::from_millis(500)).await; + assert_eq!(report.status, HealthcheckStatus::Unhealthy); + } + + #[tokio::test] + async fn second_call_within_ttl_reuses_cached_report_without_rerunning_checks() { + use std::sync::atomic::AtomicUsize; + + struct CountingCheck(Arc); + #[async_trait] + impl Healthcheck for CountingCheck { + fn name(&self) -> &'static str { + "counting-check" + } + async fn check(&self) -> HealthcheckResult { + self.0.fetch_add(1, Ordering::SeqCst); + HealthcheckResult::healthy() + } + } + + let calls = Arc::new(AtomicUsize::new(0)); + let registry = RestHealthcheckRegistry::new(); + registry.register("counted-gear", Arc::new(CountingCheck(calls.clone()))); + + let first = registry.report(Duration::from_millis(500)).await; + let second = registry.report(Duration::from_millis(500)).await; + + assert_eq!(first.status, HealthcheckStatus::Healthy); + assert_eq!(second.status, HealthcheckStatus::Healthy); + assert_eq!( + calls.load(Ordering::SeqCst), + 1, + "second report() within REPORT_CACHE_TTL must reuse the cached result" + ); + } + + #[test] + fn sensitive_message_is_sanitized() { + assert_eq!( + sanitize_message("connection to postgres://user:pass@host/db failed"), + "health check failed" + ); + assert_eq!( + sanitize_message("invalid token in header"), + "health check failed" + ); + assert_eq!( + sanitize_message("SELECT * FROM users failed"), + "health check failed" + ); + assert_eq!( + sanitize_message("thread panicked at some_file.rs:42"), + "health check failed" + ); + assert_eq!( + sanitize_message("invalid api_key provided"), + "health check failed" + ); + assert_eq!( + sanitize_message("connection to mongodb://host/db failed"), + "health check failed" + ); + } + + #[test] + fn long_message_is_sanitized() { + let long = "x".repeat(257); + assert_eq!(sanitize_message(&long), "health check failed"); + } + + #[test] + fn clean_message_passes_through() { + assert_eq!( + sanitize_message("upstream service unavailable"), + "upstream service unavailable" + ); + } +} diff --git a/libs/toolkit/src/lib.rs b/libs/toolkit/src/lib.rs index f78d0770a..186e9e565 100644 --- a/libs/toolkit/src/lib.rs +++ b/libs/toolkit/src/lib.rs @@ -136,6 +136,14 @@ pub use directory::{ ServiceInstanceInfo, }; +// REST healthcheck infrastructure. Module private so types have one public path +// (flat re-export below), not two. +mod healthcheck; +pub use healthcheck::{ + Healthcheck, HealthcheckComponentReport, HealthcheckReport, HealthcheckResult, + HealthcheckStatus, RestHealthcheckRegistry, +}; + // GTS schema support pub mod gts; diff --git a/libs/toolkit/src/registry.rs b/libs/toolkit/src/registry.rs index c53dbabe8..61ab57a6b 100644 --- a/libs/toolkit/src/registry.rs +++ b/libs/toolkit/src/registry.rs @@ -1071,6 +1071,7 @@ mod tests { &self, _ctx: &crate::context::GearCtx, router: axum::Router, + _hc_registry: std::sync::Arc, ) -> anyhow::Result { Ok(router) } @@ -1078,6 +1079,7 @@ mod tests { &self, _ctx: &crate::context::GearCtx, router: axum::Router, + _hc_registry: std::sync::Arc, ) -> anyhow::Result { Ok(router) } diff --git a/libs/toolkit/src/runtime/host_runtime.rs b/libs/toolkit/src/runtime/host_runtime.rs index 64cc87cee..9515c4404 100644 --- a/libs/toolkit/src/runtime/host_runtime.rs +++ b/libs/toolkit/src/runtime/host_runtime.rs @@ -439,13 +439,22 @@ impl HostRuntime { // use host as the registry let registry: &dyn crate::contracts::OpenApiRegistry = host.as_registry(); + // Healthcheck registry, passed explicitly to the REST host and providers below + // (not via ClientHub). Seeded with the host's shutdown token so in-flight checks + // are aborted on shutdown. + let hc_registry = Arc::new( + crate::healthcheck::RestHealthcheckRegistry::with_cancellation( + host_ctx.cancellation_token().clone(), + ), + ); + // 1) Host prepare: base Router / global middlewares / basic OAS meta - router = - host.rest_prepare(&host_ctx, router) - .map_err(|source| RegistryError::RestPrepare { - gear: host_entry.name, - source, - })?; + router = host + .rest_prepare(&host_ctx, router, hc_registry.clone()) + .map_err(|source| RegistryError::RestPrepare { + gear: host_entry.name, + source, + })?; // 2) Register all REST providers (in the current discovery order) for e in self.registry.gears() { @@ -463,16 +472,21 @@ impl HostRuntime { gear: e.name, source, })?; + + // Register the gear's readiness healthcheck after successful route registration. + if let Some(hc) = rest.healthcheck(&ctx) { + hc_registry.register(e.name, hc); + } } } // 3) Host finalize: attach /openapi.json and /docs, persist Router if needed (no server start) - router = host.rest_finalize(&host_ctx, router).map_err(|source| { - RegistryError::RestFinalize { + router = host + .rest_finalize(&host_ctx, router, hc_registry) + .map_err(|source| RegistryError::RestFinalize { gear: host_entry.name, source, - } - })?; + })?; Ok(router) } diff --git a/libs/toolkit/tests/macro_tests.rs b/libs/toolkit/tests/macro_tests.rs index 163a45d74..a2d552a33 100644 --- a/libs/toolkit/tests/macro_tests.rs +++ b/libs/toolkit/tests/macro_tests.rs @@ -188,6 +188,7 @@ impl ApiGatewayCapability for TestApiGatewayGear { &self, _ctx: &toolkit::context::GearCtx, router: axum::Router, + _hc_registry: std::sync::Arc, ) -> anyhow::Result { Ok(router) } @@ -196,6 +197,7 @@ impl ApiGatewayCapability for TestApiGatewayGear { &self, _ctx: &toolkit::context::GearCtx, router: axum::Router, + _hc_registry: std::sync::Arc, ) -> anyhow::Result { Ok(router) }