390 lines
11 KiB
Rust
390 lines
11 KiB
Rust
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::*;
|
|
use actix_web::{http::StatusCode, test as actix_test};
|
|
|
|
#[actix_web::test]
|
|
async fn trace_header() {
|
|
let app = actix_test::init_service(
|
|
App::new()
|
|
.wrap(from_fn(trace_request))
|
|
.route("/", web::get().to(live)),
|
|
)
|
|
.await;
|
|
let trace_id = "4bf92f3577b34da6a3ce929d0e0e4736";
|
|
let request = actix_test::TestRequest::get()
|
|
.uri("/")
|
|
.insert_header(("traceparent", format!("00-{trace_id}-00f067aa0ba902b7-01")))
|
|
.to_request();
|
|
let response = actix_test::call_service(&app, request).await;
|
|
|
|
assert_eq!(response.status(), StatusCode::OK);
|
|
let trace = response
|
|
.headers()
|
|
.get("traceparent")
|
|
.unwrap()
|
|
.to_str()
|
|
.unwrap();
|
|
assert!(trace.starts_with(&format!("00-{trace_id}-")));
|
|
assert!(trace.ends_with("-01"));
|
|
assert_eq!(trace.len(), 55);
|
|
assert_eq!(
|
|
actix_test::read_body(response).await,
|
|
r#"{"status":"live"}"#
|
|
);
|
|
}
|
|
|
|
#[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());
|
|
}
|
|
}
|