use std::{ env, fmt, future::Future, net::SocketAddr, pin::Pin, str::FromStr, sync::{ Arc, atomic::{AtomicBool, AtomicU64, Ordering}, }, time::Duration, }; use actix_web::{ App, Error, HttpResponse, HttpServer, body::MessageBody, dev::{ServiceRequest, ServiceResponse}, http::header::{HeaderName, HeaderValue}, middleware::{Next, from_fn}, web, }; use remora_coordination::CoordinationStore; use serde::Serialize; use sqlx::PgPool; use tokio::task::JoinHandle; use tracing::Instrument as _; use tracing_subscriber::EnvFilter; #[derive(Clone, Debug)] pub struct ManagementConfig { pub bind: SocketAddr, pub component: String, } impl ManagementConfig { pub fn from_env(component: impl Into) -> Result { Ok(Self { bind: env_parse("REMORA_MANAGEMENT_BIND", "0.0.0.0:9090")?, component: component.into(), }) } } #[derive(Clone)] pub struct RuntimeStatus { component: Arc, accepting_traffic: Arc, background_runs: Arc, background_failures: Arc, pool: PgPool, coordination: CoordinationStore, readiness_checks: Vec, } type ReadinessCheck = Arc Pin + Send>> + Send + Sync>; impl RuntimeStatus { #[must_use] pub fn new( component: impl Into, pool: PgPool, coordination: CoordinationStore, ) -> Self { Self { component: Arc::from(component.into()), accepting_traffic: Arc::new(AtomicBool::new(false)), background_runs: Arc::new(AtomicU64::new(0)), background_failures: Arc::new(AtomicU64::new(0)), pool, coordination, readiness_checks: Vec::new(), } } #[must_use] pub fn with_readiness_check(mut self, check: F) -> Self where F: Fn() -> Fut + Send + Sync + 'static, Fut: Future + Send + 'static, { self.readiness_checks .push(Arc::new(move || Box::pin(check()))); self } pub fn set_accepting_traffic(&self, value: bool) { self.accepting_traffic.store(value, Ordering::Release); } pub fn record_background_run(&self, succeeded: bool) { self.background_runs.fetch_add(1, Ordering::Relaxed); if !succeeded { self.background_failures.fetch_add(1, Ordering::Relaxed); } } async fn ready(&self) -> bool { if !self.accepting_traffic.load(Ordering::Acquire) { return false; } let database = sqlx::query_scalar::<_, i32>("SELECT 1") .fetch_one(&self.pool) .await .is_ok(); let coordination = self.coordination.global_used().await.is_ok(); if !database || !coordination { return false; } for check in &self.readiness_checks { if !check().await { return false; } } true } fn metrics(&self, ready: bool) -> String { let component = prometheus_label(&self.component); format!( "# TYPE remora_build_info gauge\n\ remora_build_info{{component=\"{component}\"}} 1\n\ # TYPE remora_ready gauge\n\ remora_ready{{component=\"{component}\"}} {}\n\ # TYPE remora_background_runs_total counter\n\ remora_background_runs_total{{component=\"{component}\"}} {}\n\ # TYPE remora_background_failures_total counter\n\ remora_background_failures_total{{component=\"{component}\"}} {}\n", u8::from(ready), self.background_runs.load(Ordering::Relaxed), self.background_failures.load(Ordering::Relaxed), ) } } #[derive(Debug)] pub struct RuntimeError(String); impl RuntimeError { #[must_use] pub fn message(message: impl Into) -> Self { Self(message.into()) } } impl fmt::Display for RuntimeError { fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { formatter.write_str(&self.0) } } impl std::error::Error for RuntimeError {} pub fn init_logging() -> Result<(), RuntimeError> { let filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("info")); tracing_subscriber::fmt() .with_env_filter(filter) .json() .with_current_span(true) .with_span_list(true) .try_init() .map_err(|error| RuntimeError::message(format!("initialize logging: {error}"))) } pub fn required_env(name: &str) -> Result { env::var(name) .ok() .filter(|value| !value.is_empty()) .ok_or_else(|| RuntimeError::message(format!("{name} is required"))) } pub fn env_string(name: &str, default: &str) -> String { env::var(name) .ok() .filter(|value| !value.is_empty()) .unwrap_or_else(|| default.to_owned()) } pub fn env_parse(name: &str, default: &str) -> Result where T: FromStr, T::Err: fmt::Display, { env_string(name, default) .parse() .map_err(|error| RuntimeError::message(format!("invalid {name}: {error}"))) } pub fn env_duration(name: &str, default_seconds: u64) -> Result { let seconds = env_parse(name, &default_seconds.to_string())?; if seconds == 0 { return Err(RuntimeError::message(format!( "{name} must be greater than zero" ))); } Ok(Duration::from_secs(seconds)) } pub fn spawn_periodic( name: &'static str, interval: Duration, status: RuntimeStatus, mut operation: F, ) -> JoinHandle<()> where F: FnMut() -> Fut + Send + 'static, Fut: Future> + Send, { tokio::spawn(async move { let mut ticker = tokio::time::interval(interval); ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); loop { ticker.tick().await; let result = operation().await; status.record_background_run(result.is_ok()); match result { Ok(()) => tracing::debug!(background_operation = name, "background pass completed"), Err(error) => { tracing::error!(background_operation = name, error, "background pass failed") } } } }) } pub fn start_management_server( config: ManagementConfig, status: RuntimeStatus, ) -> std::io::Result { let server = HttpServer::new(move || { App::new() .wrap(from_fn(trace_request)) .app_data(web::Data::new(status.clone())) .route("/health/live", web::get().to(live)) .route("/health/ready", web::get().to(ready)) .route("/metrics", web::get().to(metrics)) }) .disable_signals() .bind(config.bind)? .run(); let handle = server.handle(); actix_web::rt::spawn(server); Ok(handle) } pub async fn trace_request( request: ServiceRequest, next: Next, ) -> Result, Error> { let trace_id = request .headers() .get("traceparent") .and_then(|value| value.to_str().ok()) .and_then(parent_trace_id) .unwrap_or_else(random_hex::<16>); let span_id = random_hex::<8>(); let response_trace = format!("00-{trace_id}-{span_id}-01"); let span = tracing::info_span!( "http_request", trace_id, method = %request.method(), path = request.path(), ); let mut response = next.call(request).instrument(span).await?; response.headers_mut().insert( HeaderName::from_static("traceparent"), HeaderValue::from_str(&response_trace).expect("generated traceparent is valid"), ); Ok(response) } async fn live() -> HttpResponse { HttpResponse::Ok().json(Health { status: "live" }) } async fn ready(status: web::Data) -> HttpResponse { if status.ready().await { HttpResponse::Ok().json(Health { status: "ready" }) } else { HttpResponse::ServiceUnavailable().json(Health { status: "not_ready", }) } } async fn metrics(status: web::Data) -> HttpResponse { HttpResponse::Ok() .content_type("text/plain; version=0.0.4") .body(status.metrics(status.ready().await)) } #[derive(Serialize)] struct Health { status: &'static str, } fn prometheus_label(value: &str) -> String { value.replace('\\', "\\\\").replace('"', "\\\"") } fn parent_trace_id(value: &str) -> Option { let mut parts = value.split('-'); let version = parts.next()?; let trace_id = parts.next()?; let parent_id = parts.next()?; let flags = parts.next()?; if parts.next().is_some() || version != "00" || trace_id.len() != 32 || parent_id.len() != 16 || flags.len() != 2 || trace_id.bytes().all(|byte| byte == b'0') || !trace_id.bytes().all(|byte| byte.is_ascii_hexdigit()) || !parent_id.bytes().all(|byte| byte.is_ascii_hexdigit()) || !flags.bytes().all(|byte| byte.is_ascii_hexdigit()) { return None; } Some(trace_id.to_ascii_lowercase()) } fn random_hex() -> String { let mut bytes = [0_u8; N]; getrandom::fill(&mut bytes).expect("operating system randomness is required"); let mut value = String::with_capacity(N * 2); for byte in bytes { use std::fmt::Write as _; write!(value, "{byte:02x}").expect("writing to String cannot fail"); } value } #[cfg(test)] mod tests { use super::*; #[test] fn escapes_prometheus_labels() { assert_eq!(prometheus_label("api\"\\node"), "api\\\"\\\\node"); } #[test] fn runtime_errors_do_not_capture_environment_values() { let error = RuntimeError::message("REMORA_SECRET is required"); assert_eq!(error.to_string(), "REMORA_SECRET is required"); } #[test] fn accepts_only_w3c_trace_identifiers() { assert_eq!( parent_trace_id("00-4bf92f3577b34da6a3ce929d0e0e4736-00f067aa0ba902b7-01"), Some("4bf92f3577b34da6a3ce929d0e0e4736".into()) ); assert!( parent_trace_id("00-00000000000000000000000000000000-00f067aa0ba902b7-01").is_none() ); assert!(parent_trace_id("invalid").is_none()); } }