Ship production deployment and runtime observability

crates/remora-runtime/Cargo.toml

@@ -0,0 +1,16 @@
+[package]
+name = "remora-runtime"
+edition.workspace = true
+publish = false
+rust-version.workspace = true
+version.workspace = true
+
+[dependencies]
+actix-web = "=4.14.0"
+getrandom = "0.3.4"
+remora-coordination = { path = "../remora-coordination" }
+serde.workspace = true
+sqlx.workspace = true
+tokio.workspace = true
+tracing.workspace = true
+tracing-subscriber.workspace = true

crates/remora-runtime/src/lib.rs

@@ -0,0 +1,358 @@
+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<String>) -> Result<Self, RuntimeError> {
+        Ok(Self {
+            bind: env_parse("REMORA_MANAGEMENT_BIND", "0.0.0.0:9090")?,
+            component: component.into(),
+        })
+    }
+}
+
+#[derive(Clone)]
+pub struct RuntimeStatus {
+    component: Arc<str>,
+    accepting_traffic: Arc<AtomicBool>,
+    background_runs: Arc<AtomicU64>,
+    background_failures: Arc<AtomicU64>,
+    pool: PgPool,
+    coordination: CoordinationStore,
+    readiness_checks: Vec<ReadinessCheck>,
+}
+
+type ReadinessCheck = Arc<dyn Fn() -> Pin<Box<dyn Future<Output = bool> + Send>> + Send + Sync>;
+
+impl RuntimeStatus {
+    #[must_use]
+    pub fn new(
+        component: impl Into<String>,
+        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<F, Fut>(mut self, check: F) -> Self
+    where
+        F: Fn() -> Fut + Send + Sync + 'static,
+        Fut: Future<Output = bool> + 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<String>) -> 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<String, RuntimeError> {
+    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<T>(name: &str, default: &str) -> Result<T, RuntimeError>
+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<Duration, RuntimeError> {
+    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<F, Fut>(
+    name: &'static str,
+    interval: Duration,
+    status: RuntimeStatus,
+    mut operation: F,
+) -> JoinHandle<()>
+where
+    F: FnMut() -> Fut + Send + 'static,
+    Fut: Future<Output = Result<(), String>> + 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<actix_web::dev::ServerHandle> {
+    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<impl MessageBody>,
+) -> Result<ServiceResponse<impl MessageBody>, 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<RuntimeStatus>) -> 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<RuntimeStatus>) -> 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<String> {
+    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<const N: usize>() -> 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());
+    }
+}