use std::sync::Arc;
use std::time::Duration;
use async_trait::async_trait;
#[cfg(doc)]
use super::registry::StatusRefresher;
use super::registry::{ComponentProbe, MIN_REFRESH_INTERVAL, StatusSettings};
use super::{Component, ComponentObservation, StatusReason};
use crate::backends::BackendFailure;
use crate::backends::catalog::RefusalReason;
use crate::backends::catalog_runtime::CatalogStatus;
use crate::backends::control_plane::postgres::ControlPlaneSettings;
use crate::backends::control_plane::{ControlPlaneStore, StatusProbeAdmission};
use crate::backends::health::BackendHealth;
pub struct CatalogProbe {
status: Arc<CatalogStatus>,
}
impl CatalogProbe {
pub fn new(status: Arc<CatalogStatus>) -> Self {
Self { status }
}
}
const fn catalogue_reason(reason: RefusalReason) -> StatusReason {
match reason {
RefusalReason::Unreachable => StatusReason::Unreachable,
RefusalReason::Denied => StatusReason::PermissionDenied,
RefusalReason::NotJson => StatusReason::PayloadCorrupt,
RefusalReason::Schema => StatusReason::SchemaIncompatible,
RefusalReason::UnsupportedEndpoint
| RefusalReason::Oversized
| RefusalReason::IdMismatch
| RefusalReason::Identifier
| RefusalReason::UnknownStatus
| RefusalReason::UnknownModality
| RefusalReason::Price
| RefusalReason::UnknownTierType
| RefusalReason::DuplicateTier
| RefusalReason::NeutralPrice
| RefusalReason::UncanonicalizableText
| RefusalReason::AmbiguousModelKey
| RefusalReason::Content
| RefusalReason::NotRetained
| RefusalReason::UnsolicitedUnchanged
| RefusalReason::Unknown => StatusReason::ValidationRejected,
}
}
#[async_trait]
impl ComponentProbe for CatalogProbe {
fn component(&self) -> Component {
Component::Catalogue
}
async fn observe(&self) -> ComponentObservation {
match self.status.report() {
None => ComponentObservation::unavailable(
Component::Catalogue,
StatusReason::Unknown,
"catalogue has no completed import".to_owned(),
),
Some(report) => match report.last_refusal {
None => ComponentObservation::ok(Component::Catalogue),
Some(reason) => {
let observation = if report.active.is_some() {
ComponentObservation::degraded
} else {
ComponentObservation::unavailable
};
observation(
Component::Catalogue,
catalogue_reason(reason),
format!("catalogue import refused: {}", reason.as_str()),
)
}
},
}
}
}
pub struct ControlPlaneProbe {
store: Arc<dyn ControlPlaneStore>,
}
impl ControlPlaneProbe {
pub fn new(store: Arc<dyn ControlPlaneStore>) -> Self {
Self { store }
}
pub fn pacing(settings: &ControlPlaneSettings) -> StatusSettings {
let bounds = settings.status_probe_timeout(1);
let spacing = settings.connect_timeout.clamp(SPACING, MAX_SPACING);
let pacing = derived(
Component::ControlPlane,
bounds,
spacing,
MIN_REFRESH_INTERVAL,
);
if pacing.probe_timeout < bounds {
tracing::warn!(
component = "control_plane",
store_bound_ms = bounds.as_millis() as u64,
probe_timeout_ms = pacing.probe_timeout.as_millis() as u64,
refresh_interval_ms = pacing.refresh_interval.as_millis() as u64,
"control plane timeouts exceed the observable cadence; the probe will report \
timeout for calls the store is still entitled to complete"
);
}
pacing
}
async fn observe_with_status_probe(
&self,
admission: Option<StatusProbeAdmission>,
) -> ComponentObservation {
match self.store.health_with_status_probe(admission).await {
Ok(()) => ComponentObservation::ok(Component::ControlPlane),
Err(error) => {
let reason = StatusReason::from_failure(error.category());
let detail = format!("{}: {error}", self.store.name());
if reason == StatusReason::Unreachable {
ComponentObservation::unavailable(Component::ControlPlane, reason, detail)
} else {
ComponentObservation::degraded(Component::ControlPlane, reason, detail)
}
}
}
}
}
const SPACING: Duration = Duration::from_secs(1);
const MAX_SPACING: Duration = Duration::from_secs(30);
pub const MAX_REFRESH_INTERVAL: Duration = Duration::from_secs(2 * 60);
pub const MAX_PROBE_TIMEOUT: Duration = Duration::from_secs(4 * 60);
pub const MAX_STALENESS_BUDGET: Duration = Duration::from_secs(5 * 60);
fn derived(
component: Component,
bound: Duration,
spacing: Duration,
floor: Duration,
) -> StatusSettings {
let refresh_interval = bound
.saturating_add(spacing)
.max(floor)
.min(MAX_REFRESH_INTERVAL);
let probe_timeout = bound.min(refresh_interval.saturating_sub(SPACING));
StatusSettings {
probe_timeout,
refresh_interval,
staleness_budget: refresh_interval
.saturating_mul(3)
.max(
refresh_interval
.saturating_add(probe_timeout)
.saturating_add(spacing),
)
.min(MAX_STALENESS_BUDGET),
enabled: vec![component],
}
}
pub const BACKEND_REFRESH_FLOOR: Duration = super::registry::EXPORT_INTERVAL;
pub struct BackendProbe {
component: Component,
health: Arc<dyn BackendHealth>,
}
impl BackendProbe {
pub fn new(component: Component, health: Arc<dyn BackendHealth>) -> Self {
Self { component, health }
}
pub fn pacing(component: Component, health: &Arc<dyn BackendHealth>) -> StatusSettings {
derived(
component,
health.bound().min(MAX_PROBE_TIMEOUT),
SPACING,
BACKEND_REFRESH_FLOOR,
)
}
}
#[async_trait]
impl ComponentProbe for BackendProbe {
fn component(&self) -> Component {
self.component
}
fn begin<'a>(
&'a self,
_fallback: Duration,
) -> (
Duration,
std::pin::Pin<Box<dyn std::future::Future<Output = ComponentObservation> + Send + 'a>>,
) {
let timeout = self.health.bound().min(MAX_PROBE_TIMEOUT);
(timeout, Box::pin(self.observe()))
}
async fn observe(&self) -> ComponentObservation {
match self.health.check().await {
Ok(()) => ComponentObservation::ok(self.component),
Err(failure) => {
let reason = StatusReason::from_failure(failure.category());
let detail = format!("{}: {}", self.health.backend(), failure.detail());
if reason == StatusReason::Unreachable {
ComponentObservation::unavailable(self.component, reason, detail)
} else {
ComponentObservation::degraded(self.component, reason, detail)
}
}
}
}
}
#[async_trait]
impl ComponentProbe for ControlPlaneProbe {
fn component(&self) -> Component {
Component::ControlPlane
}
fn begin<'a>(
&'a self,
fallback: Duration,
) -> (
Duration,
std::pin::Pin<Box<dyn std::future::Future<Output = ComponentObservation> + Send + 'a>>,
) {
let admission = self.store.status_probe_admission();
let timeout = admission
.as_ref()
.map(StatusProbeAdmission::timeout)
.unwrap_or(fallback)
.min(MAX_PROBE_TIMEOUT);
(timeout, Box::pin(self.observe_with_status_probe(admission)))
}
async fn observe(&self) -> ComponentObservation {
self.observe_with_status_probe(None).await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::backends::control_plane::ControlPlaneError;
use crate::backends::health::HealthFailure;
use crate::backends::{Capabilities, FailureCategory};
use crate::desired_state::oracle::InMemoryControlPlane;
use crate::desired_state::{
AccessDenial, AuditEvent, DenialPage, LoadedRevision, RevisionCandidate, RevisionId,
RevisionManifest,
};
use crate::status::ComponentState;
use crate::status::registry::{CachedStatusRegistry, StatusRefresher};
use std::sync::Mutex;
use std::sync::atomic::{AtomicUsize, Ordering};
use tracing_subscriber::layer::SubscriberExt as _;
type Health = Box<dyn Fn() -> Result<(), ControlPlaneError> + Send + Sync>;
#[derive(Clone, Default)]
struct CapturedLogs(Arc<Mutex<Vec<u8>>>);
impl CapturedLogs {
fn rendered(&self) -> String {
String::from_utf8(self.0.lock().expect("not poisoned").clone()).expect("utf-8 logs")
}
}
impl std::io::Write for CapturedLogs {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.0.lock().expect("not poisoned").extend_from_slice(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
impl<'writer> tracing_subscriber::fmt::MakeWriter<'writer> for CapturedLogs {
type Writer = Self;
fn make_writer(&'writer self) -> Self::Writer {
self.clone()
}
}
struct Answering {
inner: Arc<InMemoryControlPlane>,
health: Health,
health_delay: Option<Duration>,
probe_timeout: Option<Duration>,
}
#[async_trait]
impl ControlPlaneStore for Answering {
fn name(&self) -> &'static str {
self.inner.name()
}
fn capabilities(&self) -> Capabilities {
self.inner.capabilities()
}
fn status_probe_admission(&self) -> Option<StatusProbeAdmission> {
self.probe_timeout.map(StatusProbeAdmission::standalone)
}
async fn health(&self) -> Result<(), ControlPlaneError> {
if let Some(delay) = self.health_delay {
tokio::time::sleep(delay).await;
}
(self.health)()
}
async fn desired_revision(&self) -> Result<Option<RevisionId>, ControlPlaneError> {
self.inner.desired_revision().await
}
async fn load_manifest(
&self,
id: RevisionId,
) -> Result<RevisionManifest, ControlPlaneError> {
self.inner.load_manifest(id).await
}
async fn load_revision(&self, id: RevisionId) -> Result<LoadedRevision, ControlPlaneError> {
self.inner.load_revision(id).await
}
async fn publish_revision(
&self,
candidate: RevisionCandidate,
) -> Result<RevisionManifest, ControlPlaneError> {
self.inner.publish_revision(candidate).await
}
async fn audit_trail(&self, id: RevisionId) -> Result<Vec<AuditEvent>, ControlPlaneError> {
self.inner.audit_trail(id).await
}
async fn record_denial(&self, denial: &AccessDenial) -> Result<(), ControlPlaneError> {
self.inner.record_denial(denial).await
}
async fn denials(
&self,
page: &DenialPage,
limit: usize,
) -> Result<Vec<AccessDenial>, ControlPlaneError> {
self.inner.denials(page, limit).await
}
}
fn probing(health: Health) -> ControlPlaneProbe {
ControlPlaneProbe::new(Arc::new(Answering {
inner: Arc::new(InMemoryControlPlane::new()),
health,
health_delay: None,
probe_timeout: None,
}))
}
fn healthy() -> Health {
Box::new(|| Ok(()))
}
fn failing(error: fn() -> ControlPlaneError) -> Health {
Box::new(move || Err(error()))
}
#[test]
fn the_probe_outlives_every_bound_the_store_is_allowed_to_take() {
let settings = ControlPlaneSettings {
connect_timeout: Duration::from_secs(5),
operation_timeout: Duration::from_secs(30),
..ControlPlaneSettings::default()
};
let pacing = ControlPlaneProbe::pacing(&settings);
let queued_behind_an_operation = settings.operation_timeout;
let reconnect = settings.connect_timeout;
let own_call = settings.operation_timeout;
assert!(
pacing.probe_timeout >= queued_behind_an_operation + reconnect + own_call,
"{:?} cuts a call the store would have completed",
pacing.probe_timeout
);
}
#[test]
fn the_probe_timeout_expands_for_every_operation_already_in_the_queue() {
let settings = ControlPlaneSettings {
connect_timeout: Duration::from_secs(5),
operation_timeout: Duration::from_secs(30),
..ControlPlaneSettings::default()
};
let queued = 3;
let expected = settings.status_probe_timeout(queued);
let probe = ControlPlaneProbe::new(Arc::new(Answering {
inner: Arc::new(InMemoryControlPlane::new()),
health: healthy(),
health_delay: None,
probe_timeout: Some(expected),
}));
let (timeout, observation) = probe.begin(Duration::from_secs(1));
drop(observation);
assert_eq!(
timeout, expected,
"the health probe must budget for all queued operations, not one fixed slot"
);
assert_eq!(expected, Duration::from_secs(125));
}
#[test]
fn a_deep_queue_cannot_extend_a_probe_past_metric_expiration() {
let settings = ControlPlaneSettings {
connect_timeout: Duration::from_secs(5),
operation_timeout: Duration::from_secs(30),
..ControlPlaneSettings::default()
};
let deep_queue = 100;
let uncapped = settings.status_probe_timeout(deep_queue);
let probe = ControlPlaneProbe::new(Arc::new(Answering {
inner: Arc::new(InMemoryControlPlane::new()),
health: healthy(),
health_delay: None,
probe_timeout: Some(uncapped),
}));
let (timeout, observation) = probe.begin(Duration::from_secs(1));
drop(observation);
assert!(uncapped > MAX_PROBE_TIMEOUT);
assert_eq!(timeout, MAX_PROBE_TIMEOUT);
assert!(MAX_PROBE_TIMEOUT < Duration::from_secs(5 * 60));
}
#[tokio::test(start_paused = true)]
async fn a_queued_healthy_probe_is_not_recorded_as_a_timeout() {
let settings = ControlPlaneSettings {
connect_timeout: Duration::from_secs(5),
operation_timeout: Duration::from_secs(30),
..ControlPlaneSettings::default()
};
let queued = 3;
let probe_timeout = settings.status_probe_timeout(queued);
let health_delay = probe_timeout - Duration::from_secs(1);
let probe = ControlPlaneProbe::new(Arc::new(Answering {
inner: Arc::new(InMemoryControlPlane::new()),
health: healthy(),
health_delay: Some(health_delay),
probe_timeout: Some(probe_timeout),
}));
let registry = Arc::new(CachedStatusRegistry::new(
StatusSettings {
refresh_interval: probe_timeout + Duration::from_secs(1),
probe_timeout: Duration::from_secs(1),
staleness_budget: Duration::from_secs(300),
enabled: vec![Component::ControlPlane],
},
Arc::new(crate::convergence::SystemClock),
));
let refresher = StatusRefresher::new(Arc::clone(®istry), vec![Arc::new(probe)]);
let round = tokio::spawn(async move { refresher.refresh_once().await });
tokio::task::yield_now().await;
tokio::time::advance(health_delay).await;
round.await.expect("the refresher does not panic");
let observed = registry
.view()
.components
.into_iter()
.find(|observed| observed.component == Component::ControlPlane)
.expect("control plane is reported");
assert_eq!(observed.state, ComponentState::Ok);
assert_eq!(observed.reason, None);
}
#[test]
fn the_derived_pacing_is_valid_for_every_configurable_bound() {
for (connect_ms, operation_ms) in [
(1_u64, 1_u64),
(100, 250),
(5_000, 30_000),
(60_000, 600_000),
] {
let settings = ControlPlaneSettings {
connect_timeout: Duration::from_millis(connect_ms),
operation_timeout: Duration::from_millis(operation_ms),
..ControlPlaneSettings::default()
};
let pacing = ControlPlaneProbe::pacing(&settings);
assert_eq!(
pacing.validate(),
Ok(()),
"connect {connect_ms}ms, operation {operation_ms}ms produced {pacing:?}"
);
assert_eq!(pacing.enabled, vec![Component::ControlPlane]);
assert!(
pacing.refresh_interval <= MAX_REFRESH_INTERVAL,
"connect {connect_ms}ms, operation {operation_ms}ms outruns the pipeline: \
{pacing:?}"
);
assert!(
pacing.staleness_budget > pacing.refresh_interval + pacing.probe_timeout,
"connect {connect_ms}ms, operation {operation_ms}ms would report stale between two \
healthy rounds: {pacing:?}"
);
assert!(
pacing.staleness_budget <= MAX_STALENESS_BUDGET,
"connect {connect_ms}ms, operation {operation_ms}ms would page for an observation \
the replica still calls fresh: {pacing:?}"
);
}
}
#[test]
fn a_store_slower_than_the_pipeline_is_capped_and_the_capping_is_announced() {
let settings = ControlPlaneSettings {
connect_timeout: Duration::from_secs(60),
operation_timeout: Duration::from_secs(600),
..ControlPlaneSettings::default()
};
let bounds = settings.operation_timeout * 2 + settings.connect_timeout;
let logs = CapturedLogs::default();
let dispatch = tracing::Dispatch::new(
tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.with_ansi(false)
.with_writer(logs.clone()),
),
);
let pacing = {
let _default = tracing::dispatcher::set_default(&dispatch);
ControlPlaneProbe::pacing(&settings)
};
assert_eq!(pacing.refresh_interval, MAX_REFRESH_INTERVAL);
assert_eq!(pacing.probe_timeout, MAX_REFRESH_INTERVAL - SPACING);
assert!(pacing.probe_timeout < bounds);
assert_eq!(pacing.staleness_budget, MAX_STALENESS_BUDGET);
assert!(pacing.staleness_budget > pacing.refresh_interval + pacing.probe_timeout);
assert_eq!(pacing.validate(), Ok(()));
let rendered = logs.rendered();
assert!(
rendered.contains("exceed the observable cadence") && rendered.contains("WARN"),
"capping the probe below the store's bounds is announced: {rendered}"
);
let quiet = CapturedLogs::default();
let dispatch = tracing::Dispatch::new(
tracing_subscriber::registry().with(
tracing_subscriber::fmt::layer()
.with_ansi(false)
.with_writer(quiet.clone()),
),
);
{
let _default = tracing::dispatcher::set_default(&dispatch);
ControlPlaneProbe::pacing(&ControlPlaneSettings::default());
}
assert_eq!(quiet.rendered(), "");
}
#[tokio::test]
async fn a_reachable_control_plane_is_ok_and_says_nothing_else() {
let observation = probing(healthy()).observe().await;
assert_eq!(observation.state, ComponentState::Ok);
assert_eq!(observation.reason, None);
assert_eq!(observation.detail, None);
}
#[tokio::test]
async fn unreachable_and_refusing_are_different_observations() {
let unreachable = probing(failing(|| ControlPlaneError::Unavailable {
backend: "postgres",
message: "connection refused".to_owned(),
}))
.observe()
.await;
assert_eq!(unreachable.state, ComponentState::Unavailable);
assert_eq!(unreachable.reason, Some(StatusReason::Unreachable));
let refusing = probing(failing(|| ControlPlaneError::Denied {
backend: "postgres",
message: "permission denied for relation revisions".to_owned(),
}))
.observe()
.await;
assert_eq!(refusing.state, ComponentState::Degraded);
assert_eq!(refusing.reason, Some(StatusReason::PermissionDenied));
}
#[tokio::test]
async fn the_backend_message_stays_on_the_detail() {
let observation = probing(failing(|| ControlPlaneError::Unavailable {
backend: "postgres",
message: "host=db.internal port=5432: connection refused".to_owned(),
}))
.observe()
.await;
let detail = observation.detail.expect("a failure carries a detail");
assert!(detail.contains("connection refused"), "{detail}");
}
struct Answer {
result: Box<dyn Fn() -> Result<(), HealthFailure> + Send + Sync>,
bound: Duration,
delay: Option<Duration>,
checks: Arc<AtomicUsize>,
}
impl Answer {
fn new(result: impl Fn() -> Result<(), HealthFailure> + Send + Sync + 'static) -> Self {
Self {
result: Box::new(result),
bound: Duration::from_secs(5),
delay: None,
checks: Arc::new(AtomicUsize::new(0)),
}
}
fn reachable() -> Self {
Self::new(|| Ok(()))
}
}
#[async_trait]
impl BackendHealth for Answer {
fn backend(&self) -> &'static str {
"redis"
}
fn bound(&self) -> Duration {
self.bound
}
async fn check(&self) -> Result<(), HealthFailure> {
self.checks.fetch_add(1, Ordering::Relaxed);
if let Some(delay) = self.delay {
tokio::time::sleep(delay).await;
}
(self.result)()
}
}
fn observing(component: Component, health: Answer) -> BackendProbe {
BackendProbe::new(component, Arc::new(health))
}
#[tokio::test]
async fn a_reachable_store_is_ok_and_says_nothing_else() {
let observation = observing(Component::RateLimitStore, Answer::reachable())
.observe()
.await;
assert_eq!(observation.component, Component::RateLimitStore);
assert_eq!(observation.state, ComponentState::Ok);
assert_eq!(observation.reason, None);
assert_eq!(observation.detail, None);
}
#[tokio::test]
async fn a_store_that_answers_and_refuses_is_degraded_not_unavailable() {
let unreachable = observing(
Component::BudgetStore,
Answer::new(|| Err(HealthFailure::unavailable("connection refused"))),
)
.observe()
.await;
assert_eq!(unreachable.state, ComponentState::Unavailable);
assert_eq!(unreachable.reason, Some(StatusReason::Unreachable));
let refusing = observing(
Component::BudgetStore,
Answer::new(|| {
Err(HealthFailure::new(
FailureCategory::Denied,
"WRONGPASS invalid username-password pair",
))
}),
)
.observe()
.await;
assert_eq!(refusing.state, ComponentState::Degraded);
assert_eq!(refusing.reason, Some(StatusReason::PermissionDenied));
}
#[tokio::test]
async fn the_stores_message_stays_on_the_detail() {
let observation = observing(
Component::RevocationStore,
Answer::new(|| Err(HealthFailure::unavailable("io error: connection reset"))),
)
.observe()
.await;
let detail = observation.detail.expect("a failure carries a detail");
assert!(detail.starts_with("redis: "), "{detail}");
assert!(detail.contains("connection reset"), "{detail}");
}
#[test]
fn a_store_is_probed_under_its_own_bound_not_the_registrys() {
let mut health = Answer::reachable();
health.bound = Duration::from_secs(3);
let probe = observing(Component::RateLimitStore, health);
let (timeout, _) = probe.begin(Duration::from_secs(90));
assert_eq!(timeout, Duration::from_secs(3));
}
#[test]
fn a_stores_pacing_is_valid_and_bounded_however_it_is_configured() {
for bound in [
Duration::from_millis(1),
Duration::from_secs(5),
Duration::from_secs(60 * 60),
] {
let mut answer = Answer::reachable();
answer.bound = bound;
let health: Arc<dyn BackendHealth> = Arc::new(answer);
let pacing = BackendProbe::pacing(Component::BudgetStore, &health);
assert_eq!(pacing.validate(), Ok(()), "{bound:?}");
assert!(
pacing.refresh_interval >= BACKEND_REFRESH_FLOOR,
"{bound:?}"
);
assert!(pacing.refresh_interval <= MAX_REFRESH_INTERVAL, "{bound:?}");
assert!(pacing.probe_timeout <= MAX_PROBE_TIMEOUT, "{bound:?}");
assert!(pacing.staleness_budget <= MAX_STALENESS_BUDGET, "{bound:?}");
assert_eq!(pacing.enabled, vec![Component::BudgetStore]);
}
}
#[tokio::test(start_paused = true)]
async fn a_store_is_asked_once_per_round_and_a_stuck_check_is_abandoned() {
let mut answer = Answer::reachable();
answer.bound = Duration::from_millis(50);
answer.delay = Some(Duration::from_secs(30));
let checks = Arc::clone(&answer.checks);
let health: Arc<dyn BackendHealth> = Arc::new(answer);
let pacing = BackendProbe::pacing(Component::RateLimitStore, &health);
let registry = Arc::new(CachedStatusRegistry::new(
pacing,
Arc::new(crate::convergence::SystemClock),
));
let refresher = StatusRefresher::new(
Arc::clone(®istry),
vec![Arc::new(BackendProbe::new(
Component::RateLimitStore,
health,
))],
);
refresher.refresh_once().await;
assert_eq!(checks.load(Ordering::Relaxed), 1);
let component = registry
.view()
.components
.into_iter()
.find(|component| component.component == Component::RateLimitStore)
.expect("the enabled component is reported");
assert_eq!(component.state, ComponentState::Unavailable);
assert_eq!(component.reason, Some(StatusReason::Timeout));
let _ = registry.view();
assert_eq!(checks.load(Ordering::Relaxed), 1);
}
}