remora/crates/remora-runtime/src/lib.rs

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());
    }
}