use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
use axum::extract::State;
use axum::http::{Method, StatusCode};
use axum::response::{IntoResponse, Response};
use axum::routing::get;
use axum::Router;
use gregg_protocol::v2::{HealthResponseV2, StatusPayloadV2};
use gregg_protocol::{HealthResponse, ReadinessState, StatusSnapshot};
use tokio::net::TcpListener;
use tokio::sync::{broadcast, RwLock};
use tracing::info;
use crate::server::error::{ServerConfigError, ServerError};
pub mod error;
const V1_UNAVAILABLE_MESSAGE: &str = "schema v1 status is unavailable on this platform";
#[allow(clippy::cast_possible_truncation)]
fn now_unix_ms() -> Option<u64> {
match SystemTime::now().duration_since(UNIX_EPOCH) {
Ok(duration) => Some(duration.as_millis() as u64),
Err(error) => {
tracing::debug!(
%error,
"system clock precedes the Unix epoch; cached snapshots remain age-uncheckable until corrected"
);
None
}
}
}
#[derive(Debug, Clone)]
pub struct Config {
pub host: IpAddr,
pub port: u16,
pub sample_interval_ms: u64,
pub max_consecutive_failures: u32,
pub max_snapshot_age: Duration,
}
impl Default for Config {
fn default() -> Self {
Self {
host: IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED),
port: 11310,
sample_interval_ms: 1000,
max_consecutive_failures: 0,
max_snapshot_age: Duration::ZERO,
}
}
}
impl Config {
pub fn validate(&self) -> Result<(), ServerConfigError> {
if self.port == 0 {
return Err(ServerConfigError::InvalidPort(self.port));
}
if self.sample_interval_ms < 250 || self.sample_interval_ms > 60000 {
return Err(ServerConfigError::InvalidSampleInterval(
self.sample_interval_ms,
));
}
Ok(())
}
#[must_use]
pub fn socket_addr(&self) -> SocketAddr {
SocketAddr::new(self.host, self.port)
}
}
#[derive(Debug, Clone)]
pub struct ServerState {
published: Arc<RwLock<PublishedState>>,
max_consecutive_failures: u32,
max_snapshot_age: Duration,
}
#[derive(Debug)]
struct PublishedState {
snapshot: Option<Arc<StatusSnapshot>>,
snapshot_v2: Option<Arc<StatusPayloadV2>>,
last_observed_at_unix_ms: Option<u64>,
health: HealthResponse,
health_v2: HealthResponseV2,
consecutive_failures: u32,
}
impl Default for ServerState {
fn default() -> Self {
Self::new()
}
}
impl ServerState {
#[must_use]
pub fn new() -> Self {
Self::with_stale_policy(0, Duration::ZERO)
}
#[must_use]
pub fn with_stale_policy(max_consecutive_failures: u32, max_snapshot_age: Duration) -> Self {
Self {
published: Arc::new(RwLock::new(PublishedState {
snapshot: None,
snapshot_v2: None,
last_observed_at_unix_ms: None,
health: HealthResponse::warming(),
health_v2: HealthResponseV2::warming(),
consecutive_failures: 0,
})),
max_consecutive_failures,
max_snapshot_age,
}
}
pub async fn update_snapshot(&self, snap: StatusSnapshot, payload_v2: StatusPayloadV2) {
let observed_at_unix_ms = snap
.observed_at_unix_ms
.max(payload_v2.snapshot.observed_at_unix_ms);
let health = HealthResponse::ready(snap.clone());
let health_v2 = HealthResponseV2::ready(payload_v2.snapshot.clone());
let arc_snap = Arc::new(snap);
let arc_snap_v2 = Arc::new(payload_v2);
let mut state = self.published.write().await;
state.snapshot = Some(arc_snap);
state.snapshot_v2 = Some(arc_snap_v2);
state.health = health;
state.health_v2 = health_v2;
state.last_observed_at_unix_ms = Some(observed_at_unix_ms);
state.consecutive_failures = 0;
}
pub async fn update_snapshot_v2_only(&self, payload_v2: StatusPayloadV2) {
let observed_at_unix_ms = payload_v2.snapshot.observed_at_unix_ms;
let health_v2 = HealthResponseV2::ready(payload_v2.snapshot.clone());
let mut state = self.published.write().await;
state.snapshot = None;
state.snapshot_v2 = Some(Arc::new(payload_v2));
state.health = HealthResponse::failed(
gregg_protocol::HealthCategory::NotServing,
V1_UNAVAILABLE_MESSAGE,
);
state.health_v2 = health_v2;
state.last_observed_at_unix_ms = Some(observed_at_unix_ms);
state.consecutive_failures = 0;
}
pub async fn update_snapshot_v1_only(&self, snap: StatusSnapshot) {
let observed_at_unix_ms = snap.observed_at_unix_ms;
let health = HealthResponse::ready(snap.clone());
let mut state = self.published.write().await;
state.snapshot = Some(Arc::new(snap));
state.snapshot_v2 = None;
state.health = health;
state.health_v2 = HealthResponseV2::failed(
gregg_protocol::HealthCategory::NotServing,
"schema v2 status is unavailable from this sampler",
);
state.last_observed_at_unix_ms = Some(observed_at_unix_ms);
state.consecutive_failures = 0;
}
pub async fn set_warming(&self) {
let mut state = self.published.write().await;
state.snapshot = None;
state.snapshot_v2 = None;
state.last_observed_at_unix_ms = None;
state.health = HealthResponse::warming();
state.health_v2 = HealthResponseV2::warming();
state.consecutive_failures = 0;
}
pub async fn set_failed(&self, msg: &str) {
let mut state = self.published.write().await;
state.consecutive_failures = state.consecutive_failures.saturating_add(1);
let prev = state.consecutive_failures;
if state.health.category != Some(gregg_protocol::HealthCategory::NotServing) {
state.health =
HealthResponse::failed(gregg_protocol::HealthCategory::CollectorFailure, msg);
}
if state.health_v2.category != Some(gregg_protocol::HealthCategory::NotServing) {
state.health_v2 =
HealthResponseV2::failed(gregg_protocol::HealthCategory::CollectorFailure, msg);
}
tracing::debug!(
consecutive_failures = prev,
max = self.max_consecutive_failures,
"server failure recorded"
);
}
#[must_use]
pub async fn consecutive_failures(&self) -> u32 {
self.published.read().await.consecutive_failures
}
async fn v1_status_data(
&self,
now_unix_ms: Option<u64>,
) -> (Option<Arc<StatusSnapshot>>, HealthResponse, bool) {
let state = self.published.read().await;
let snapshot_is_stale = self.is_stale(&state, now_unix_ms);
let mut health = state.health.clone();
if snapshot_is_stale && health.state == ReadinessState::Ready {
health = HealthResponse::failed(
gregg_protocol::HealthCategory::CollectorFailure,
"cached snapshot is stale",
);
}
(state.snapshot.clone(), health, snapshot_is_stale)
}
async fn v2_status_data(
&self,
now_unix_ms: Option<u64>,
) -> (Option<Arc<StatusPayloadV2>>, HealthResponseV2, bool) {
let state = self.published.read().await;
let snapshot_is_stale = self.is_stale(&state, now_unix_ms);
let mut health_v2 = state.health_v2.clone();
if snapshot_is_stale && health_v2.state == ReadinessState::Ready {
health_v2 = HealthResponseV2::failed(
gregg_protocol::HealthCategory::CollectorFailure,
"cached snapshot is stale",
);
}
(state.snapshot_v2.clone(), health_v2, snapshot_is_stale)
}
fn is_stale(&self, state: &PublishedState, now_unix_ms: Option<u64>) -> bool {
if self.max_consecutive_failures > 0 {
let failures = state.consecutive_failures;
if failures >= self.max_consecutive_failures {
return true;
}
}
if !self.max_snapshot_age.is_zero() {
let Some(now_unix_ms) = now_unix_ms else {
return true;
};
if let Some(observed_at_unix_ms) = state.last_observed_at_unix_ms {
let age_ms = now_unix_ms.checked_sub(observed_at_unix_ms);
if age_ms.is_none_or(|age| u128::from(age) >= self.max_snapshot_age.as_millis()) {
return true;
}
}
}
false
}
pub async fn snapshot(&self) -> Option<Arc<StatusSnapshot>> {
self.published.read().await.snapshot.clone()
}
pub async fn snapshot_v2(&self) -> Option<Arc<StatusPayloadV2>> {
self.published.read().await.snapshot_v2.clone()
}
pub async fn health(&self) -> HealthResponse {
self.published.read().await.health.clone()
}
pub async fn health_v2(&self) -> HealthResponseV2 {
self.published.read().await.health_v2.clone()
}
}
pub async fn serve(
listener: TcpListener,
state: ServerState,
mut shutdown: broadcast::Receiver<()>,
) -> Result<(), ServerError> {
let addr = listener.local_addr().map_err(ServerError::Runtime)?;
let app = Router::new()
.route("/", get(status_handler))
.route("/v1/status", get(status_handler))
.route("/v2/status", get(status_handler_v2))
.route("/healthz", get(health_handler))
.route("/v2/healthz", get(health_handler_v2))
.fallback(fallback_handler)
.with_state(state);
info!("greggd listening on {addr}");
axum::serve(listener, app)
.with_graceful_shutdown(async move {
let _ = shutdown.recv().await;
info!("shutdown signal received, stopping HTTP server");
})
.await
.map_err(ServerError::Runtime)
}
async fn status_handler(State(state): State<ServerState>) -> Response {
let now = now_unix_ms();
let (snap, health_state, snapshot_is_stale) = state.v1_status_data(now).await;
if let Some(snap) = snap {
if snapshot_is_stale {
return health_response(&health_state, StatusCode::SERVICE_UNAVAILABLE);
}
let body = match serde_json::to_vec(&*snap) {
Ok(body) => body,
Err(e) => {
let error_body = serde_json::to_vec(&serde_json::json!({"error": e.to_string()}))
.unwrap_or_else(|_| b"{\"error\":\"serialization failed\"}".to_vec());
return (
StatusCode::INTERNAL_SERVER_ERROR,
[("content-type", "application/json")],
error_body,
)
.into_response();
}
};
return (StatusCode::OK, [("content-type", "application/json")], body).into_response();
}
health_response(&health_state, StatusCode::SERVICE_UNAVAILABLE)
}
async fn health_handler(State(state): State<ServerState>) -> Response {
let now = now_unix_ms();
let (_, health_state, snapshot_is_stale) = state.v1_status_data(now).await;
let status =
if health_state.state == gregg_protocol::ReadinessState::Ready && !snapshot_is_stale {
StatusCode::OK
} else {
StatusCode::SERVICE_UNAVAILABLE
};
health_response(&health_state, status)
}
async fn fallback_handler(method: Method, uri: axum::http::Uri) -> (StatusCode, String) {
(StatusCode::NOT_FOUND, format!("{method} {uri} not found"))
}
fn health_response(health: &HealthResponse, status: StatusCode) -> Response {
let body = match serde_json::to_vec(&health) {
Ok(body) => body,
Err(e) => {
let error_body = serde_json::to_vec(&serde_json::json!({"error": e.to_string()}))
.unwrap_or_else(|_| b"{\"error\":\"serialization failed\"}".to_vec());
return (
StatusCode::INTERNAL_SERVER_ERROR,
[("content-type", "application/json")],
error_body,
)
.into_response();
}
};
(status, [("content-type", "application/json")], body).into_response()
}
async fn status_handler_v2(State(state): State<ServerState>) -> Response {
let now = now_unix_ms();
let (snap, health_state, snapshot_is_stale) = state.v2_status_data(now).await;
if let Some(snap) = snap {
if snapshot_is_stale {
return health_response_v2(&health_state, StatusCode::SERVICE_UNAVAILABLE);
}
let body = match serde_json::to_vec(&*snap) {
Ok(body) => body,
Err(e) => {
let error_body = serde_json::to_vec(&serde_json::json!({"error": e.to_string()}))
.unwrap_or_else(|_| b"{\"error\":\"serialization failed\"}".to_vec());
return (
StatusCode::INTERNAL_SERVER_ERROR,
[("content-type", "application/json")],
error_body,
)
.into_response();
}
};
return (StatusCode::OK, [("content-type", "application/json")], body).into_response();
}
health_response_v2(&health_state, StatusCode::SERVICE_UNAVAILABLE)
}
async fn health_handler_v2(State(state): State<ServerState>) -> Response {
let now = now_unix_ms();
let (_, health_state, snapshot_is_stale) = state.v2_status_data(now).await;
let status =
if health_state.state == gregg_protocol::ReadinessState::Ready && !snapshot_is_stale {
StatusCode::OK
} else {
StatusCode::SERVICE_UNAVAILABLE
};
health_response_v2(&health_state, status)
}
fn health_response_v2(health: &HealthResponseV2, status: StatusCode) -> Response {
let body = match serde_json::to_vec(&health) {
Ok(body) => body,
Err(e) => {
let error_body = serde_json::to_vec(&serde_json::json!({"error": e.to_string()}))
.unwrap_or_else(|_| b"{\"error\":\"serialization failed\"}".to_vec());
return (
StatusCode::INTERNAL_SERVER_ERROR,
[("content-type", "application/json")],
error_body,
)
.into_response();
}
};
(status, [("content-type", "application/json")], body).into_response()
}
#[cfg(test)]
mod tests;