1use async_trait::async_trait;
92use parking_lot::RwLock;
93use serde::{Deserialize, Serialize};
94use std::sync::Arc;
95use std::time::{Duration, Instant};
96use tokio::sync::Mutex as AsyncMutex;
97use tokio_util::sync::CancellationToken;
98
99#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
102#[serde(rename_all = "lowercase")]
103pub enum HealthcheckStatus {
104 Healthy,
105 Degraded,
106 Unhealthy,
107}
108
109#[derive(Debug, Clone, Serialize, Deserialize)]
111pub struct HealthcheckResult {
112 pub status: HealthcheckStatus,
113 pub message: Option<String>,
114 #[serde(skip_serializing_if = "Option::is_none")]
120 pub code: Option<String>,
121}
122
123impl HealthcheckResult {
124 #[must_use]
125 pub fn healthy() -> Self {
126 Self {
127 status: HealthcheckStatus::Healthy,
128 message: None,
129 code: None,
130 }
131 }
132
133 #[must_use]
134 pub fn degraded(message: impl Into<String>) -> Self {
135 Self {
136 status: HealthcheckStatus::Degraded,
137 message: Some(message.into()),
138 code: None,
139 }
140 }
141
142 #[must_use]
143 pub fn unhealthy(message: impl Into<String>) -> Self {
144 Self {
145 status: HealthcheckStatus::Unhealthy,
146 message: Some(message.into()),
147 code: None,
148 }
149 }
150
151 #[must_use]
153 pub fn with_code(mut self, code: impl Into<String>) -> Self {
154 self.code = Some(code.into());
155 self
156 }
157}
158
159impl Default for HealthcheckResult {
160 fn default() -> Self {
161 Self::healthy()
162 }
163}
164
165#[async_trait]
172pub trait Healthcheck: Send + Sync + 'static {
173 fn name(&self) -> &'static str;
175
176 async fn check(&self) -> HealthcheckResult {
177 HealthcheckResult::healthy()
178 }
179}
180
181#[derive(Debug, Clone, Serialize, Deserialize)]
183pub struct HealthcheckComponentReport {
184 pub gear: String,
185 pub check: String,
186 pub status: HealthcheckStatus,
187 pub message: Option<String>,
189 #[serde(skip_serializing_if = "Option::is_none")]
191 pub code: Option<String>,
192 pub latency_ms: u64,
193}
194
195#[derive(Debug, Clone, Serialize, Deserialize)]
197pub struct HealthcheckReport {
198 pub status: HealthcheckStatus,
199 pub components: Vec<HealthcheckComponentReport>,
200}
201
202impl HealthcheckReport {
203 #[must_use]
205 pub fn is_ready(&self) -> bool {
206 self.status != HealthcheckStatus::Unhealthy
207 }
208}
209
210struct RegistryEntry {
211 gear: String,
212 check: Arc<dyn Healthcheck>,
213}
214
215const REPORT_CACHE_TTL: Duration = Duration::from_secs(2);
218
219#[derive(Default)]
222pub struct RestHealthcheckRegistry {
223 entries: RwLock<Vec<RegistryEntry>>,
224 cached_report: AsyncMutex<Option<(Instant, HealthcheckReport)>>,
225 cancel: CancellationToken,
227}
228
229impl RestHealthcheckRegistry {
230 #[must_use]
231 pub fn new() -> Self {
232 Self::default()
233 }
234
235 #[must_use]
237 pub fn with_cancellation(cancel: CancellationToken) -> Self {
238 Self {
239 cancel,
240 ..Self::default()
241 }
242 }
243
244 pub fn register(&self, gear: impl Into<String>, check: Arc<dyn Healthcheck>) {
246 self.entries.write().push(RegistryEntry {
247 gear: gear.into(),
248 check,
249 });
250 }
251
252 pub async fn report(&self, timeout_per_check: Duration) -> HealthcheckReport {
256 let mut guard = self.cached_report.lock().await;
257 if let Some((ts, cached)) = guard.as_ref()
258 && ts.elapsed() < REPORT_CACHE_TTL
259 {
260 return cached.clone();
261 }
262
263 let entries: Vec<(_, _)> = {
264 let r = self.entries.read();
265 r.iter()
266 .map(|e| (e.gear.clone(), e.check.clone()))
267 .collect()
268 };
269
270 let report = if entries.is_empty() {
271 HealthcheckReport {
272 status: HealthcheckStatus::Healthy,
273 components: vec![],
274 }
275 } else {
276 let cancel = &self.cancel;
277 let checks = entries
278 .into_iter()
279 .map(|(gear, check)| run_one_check(gear, check, timeout_per_check, cancel));
280 let components = futures_util::future::join_all(checks).await;
281 let aggregate = compute_aggregate(&components);
282 HealthcheckReport {
283 status: aggregate,
284 components,
285 }
286 };
287
288 *guard = Some((Instant::now(), report.clone()));
289 report
290 }
291}
292
293async fn run_one_check(
294 gear: String,
295 check: Arc<dyn Healthcheck>,
296 timeout_per_check: Duration,
297 cancel: &CancellationToken,
298) -> HealthcheckComponentReport {
299 let name = check.name();
300 let start = Instant::now();
301
302 let mut handle = tokio::spawn(async move { check.check().await });
303
304 let (status, message, code) = tokio::select! {
305 result = &mut handle => match result {
306 Ok(r) => (
307 r.status,
308 r.message.as_deref().map(sanitize_message),
309 r.code.as_deref().and_then(sanitize_code),
310 ),
311 Err(join_err) => {
312 tracing::error!(gear = %gear, check = name, error = %join_err, "healthcheck task panicked");
313 (
314 HealthcheckStatus::Unhealthy,
315 Some("health check failed".to_owned()),
316 None,
317 )
318 }
319 },
320 () = tokio::time::sleep(timeout_per_check) => {
321 handle.abort();
322 if let Err(join_err) = handle.await
323 && !join_err.is_cancelled()
324 {
325 tracing::warn!(gear = %gear, check = name, error = %join_err, "healthcheck task join error after timeout abort");
326 }
327 tracing::warn!(gear = %gear, check = name, timeout_ms = timeout_per_check.as_millis(), "healthcheck timed out");
328 (
329 HealthcheckStatus::Unhealthy,
330 Some("health check timed out".to_owned()),
331 None,
332 )
333 }
334 () = cancel.cancelled() => {
335 handle.abort();
336 if let Err(join_err) = handle.await
337 && !join_err.is_cancelled()
338 {
339 tracing::warn!(gear = %gear, check = name, error = %join_err, "healthcheck task join error after shutdown abort");
340 }
341 tracing::warn!(gear = %gear, check = name, "healthcheck cancelled by runtime shutdown");
342 (
343 HealthcheckStatus::Unhealthy,
344 Some("health check cancelled".to_owned()),
345 None,
346 )
347 }
348 };
349
350 HealthcheckComponentReport {
351 gear,
352 check: name.to_owned(),
353 status,
354 message,
355 code,
356 latency_ms: u64::try_from(start.elapsed().as_millis()).unwrap_or(u64::MAX),
357 }
358}
359
360fn compute_aggregate(components: &[HealthcheckComponentReport]) -> HealthcheckStatus {
361 let mut status = HealthcheckStatus::Healthy;
362 for c in components {
363 match c.status {
364 HealthcheckStatus::Unhealthy => return HealthcheckStatus::Unhealthy,
365 HealthcheckStatus::Degraded => {
366 status = HealthcheckStatus::Degraded;
367 }
368 HealthcheckStatus::Healthy => {}
369 }
370 }
371 status
372}
373
374const MAX_MESSAGE_LEN: usize = 256;
375
376static SUSPICIOUS_SUBSTRINGS: &[&str] = &[
383 "password",
384 "passwd",
385 "secret",
386 "token",
387 "bearer",
388 "authorization",
389 "api_key",
390 "apikey",
391 "private_key",
392 "credential",
393 "access_key",
394 "aws_",
395 "client_secret",
396 "postgres://",
397 "postgresql://",
398 "mysql://",
399 "sqlite://",
400 "mongodb://",
401 "redis://",
402 "amqp://",
403 "jdbc:",
404 "panic",
405 "stack backtrace",
406];
407
408fn sanitize_message(msg: &str) -> String {
409 if msg.len() > MAX_MESSAGE_LEN {
410 return "health check failed".to_owned();
411 }
412 let lower = msg.to_lowercase();
413 for sub in SUSPICIOUS_SUBSTRINGS {
414 if lower.contains(sub) {
415 return "health check failed".to_owned();
416 }
417 }
418 msg.to_owned()
419}
420
421const MAX_CODE_LEN: usize = 64;
422
423fn sanitize_code(code: &str) -> Option<String> {
430 let ok = !code.is_empty()
431 && code.len() <= MAX_CODE_LEN
432 && code.bytes().all(|b| {
433 b.is_ascii_lowercase() || b.is_ascii_digit() || matches!(b, b'_' | b'.' | b'-')
434 });
435 if ok {
436 Some(code.to_owned())
437 } else {
438 tracing::debug!(
439 code,
440 "health check code dropped: not a [a-z0-9_.-] identifier"
441 );
442 None
443 }
444}
445
446#[cfg(test)]
447#[cfg_attr(coverage_nightly, coverage(off))]
448mod tests {
449 use super::*;
450 use std::sync::atomic::{AtomicBool, Ordering};
451 use std::time::Duration;
452 use tokio::sync::Notify;
453
454 struct AlwaysDegraded;
455 #[async_trait]
456 impl Healthcheck for AlwaysDegraded {
457 fn name(&self) -> &'static str {
458 "always-degraded"
459 }
460 async fn check(&self) -> HealthcheckResult {
461 HealthcheckResult::degraded("cache warming")
462 }
463 }
464
465 struct AlwaysUnhealthy;
466 #[async_trait]
467 impl Healthcheck for AlwaysUnhealthy {
468 fn name(&self) -> &'static str {
469 "always-unhealthy"
470 }
471 async fn check(&self) -> HealthcheckResult {
472 HealthcheckResult::unhealthy("database unreachable")
473 }
474 }
475
476 struct SlowCheck;
477 #[async_trait]
478 impl Healthcheck for SlowCheck {
479 fn name(&self) -> &'static str {
480 "slow-check"
481 }
482 async fn check(&self) -> HealthcheckResult {
483 tokio::time::sleep(Duration::from_secs(10)).await;
484 HealthcheckResult::healthy()
485 }
486 }
487
488 struct AbortTrackingCheck {
489 entered: Arc<Notify>,
490 dropped: Arc<AtomicBool>,
491 }
492
493 #[async_trait]
494 impl Healthcheck for AbortTrackingCheck {
495 fn name(&self) -> &'static str {
496 "abort-tracking-check"
497 }
498 async fn check(&self) -> HealthcheckResult {
499 struct DropFlag(Arc<AtomicBool>);
500 impl Drop for DropFlag {
501 fn drop(&mut self) {
502 self.0.store(true, Ordering::SeqCst);
503 }
504 }
505
506 let _drop_flag = DropFlag(self.dropped.clone());
507 self.entered.notify_one();
508 std::future::pending::<HealthcheckResult>().await
509 }
510 }
511
512 struct PanickingCheck;
513 #[async_trait]
514 impl Healthcheck for PanickingCheck {
515 fn name(&self) -> &'static str {
516 "panicking-check"
517 }
518 async fn check(&self) -> HealthcheckResult {
519 panic!("intentional panic in healthcheck");
520 }
521 }
522
523 #[test]
524 fn healthcheck_result_default_is_healthy() {
525 let r = HealthcheckResult::default();
526 assert_eq!(r.status, HealthcheckStatus::Healthy);
527 assert!(r.message.is_none());
528 }
529
530 #[tokio::test]
531 async fn empty_registry_report_is_healthy_and_ready() {
532 let registry = RestHealthcheckRegistry::new();
533 let report = registry.report(Duration::from_millis(500)).await;
534 assert_eq!(report.status, HealthcheckStatus::Healthy);
535 assert!(report.is_ready());
536 assert!(report.components.is_empty());
537 }
538
539 #[tokio::test]
540 async fn unhealthy_component_makes_aggregate_unhealthy_and_not_ready() {
541 let registry = RestHealthcheckRegistry::new();
542 registry.register("my-gear", Arc::new(AlwaysUnhealthy));
543 let report = registry.report(Duration::from_millis(500)).await;
544 assert_eq!(report.status, HealthcheckStatus::Unhealthy);
545 assert!(!report.is_ready());
546 }
547
548 #[tokio::test]
549 async fn mixed_degraded_and_unhealthy_aggregate_is_unhealthy() {
550 let registry = RestHealthcheckRegistry::new();
551 registry.register("degraded-gear", Arc::new(AlwaysDegraded));
552 registry.register("unhealthy-gear", Arc::new(AlwaysUnhealthy));
553 let report = registry.report(Duration::from_millis(500)).await;
554 assert_eq!(
555 report.status,
556 HealthcheckStatus::Unhealthy,
557 "unhealthy must take priority over degraded"
558 );
559 assert!(!report.is_ready());
560 }
561
562 #[tokio::test]
563 async fn slow_check_times_out_and_becomes_unhealthy() {
564 let registry = RestHealthcheckRegistry::new();
565 registry.register("slow-gear", Arc::new(SlowCheck));
566 let report = registry.report(Duration::from_millis(100)).await;
567 assert_eq!(report.status, HealthcheckStatus::Unhealthy);
568 let comp = &report.components[0];
569 assert_eq!(comp.status, HealthcheckStatus::Unhealthy);
570 assert_eq!(comp.message.as_deref(), Some("health check timed out"));
571 }
572
573 #[tokio::test]
574 async fn timed_out_check_is_aborted() {
575 let registry = RestHealthcheckRegistry::new();
576 let entered = Arc::new(Notify::new());
577 let dropped = Arc::new(AtomicBool::new(false));
578 registry.register(
579 "slow-gear",
580 Arc::new(AbortTrackingCheck {
581 entered: entered.clone(),
582 dropped: dropped.clone(),
583 }),
584 );
585
586 let report_task =
587 tokio::spawn(async move { registry.report(Duration::from_millis(10)).await });
588 entered.notified().await;
589 let report = report_task.await.expect("report task panicked");
590
591 assert_eq!(report.status, HealthcheckStatus::Unhealthy);
592 assert!(
593 dropped.load(Ordering::SeqCst),
594 "timed-out healthcheck future must be dropped after abort"
595 );
596 }
597
598 #[tokio::test]
599 async fn cancelled_registry_aborts_in_flight_check_and_reports_unhealthy() {
600 let cancel = CancellationToken::new();
601 let registry = RestHealthcheckRegistry::with_cancellation(cancel.clone());
602 let entered = Arc::new(Notify::new());
603 let dropped = Arc::new(AtomicBool::new(false));
604 registry.register(
605 "slow-gear",
606 Arc::new(AbortTrackingCheck {
607 entered: entered.clone(),
608 dropped: dropped.clone(),
609 }),
610 );
611
612 let report_task =
614 tokio::spawn(async move { registry.report(Duration::from_mins(1)).await });
615 entered.notified().await;
616 cancel.cancel();
617 let report = report_task.await.expect("report task panicked");
618
619 assert_eq!(report.status, HealthcheckStatus::Unhealthy);
620 assert_eq!(
621 report.components[0].message.as_deref(),
622 Some("health check cancelled")
623 );
624 assert!(
625 dropped.load(Ordering::SeqCst),
626 "cancelled healthcheck future must be dropped after abort"
627 );
628 }
629
630 #[tokio::test]
631 async fn panicking_check_becomes_unhealthy_without_panicking_caller() {
632 let registry = RestHealthcheckRegistry::new();
633 registry.register("panic-gear", Arc::new(PanickingCheck));
634 let report = registry.report(Duration::from_millis(500)).await;
635 assert_eq!(report.status, HealthcheckStatus::Unhealthy);
636 }
637
638 #[tokio::test]
639 async fn second_call_within_ttl_reuses_cached_report_without_rerunning_checks() {
640 use std::sync::atomic::AtomicUsize;
641
642 struct CountingCheck(Arc<AtomicUsize>);
643 #[async_trait]
644 impl Healthcheck for CountingCheck {
645 fn name(&self) -> &'static str {
646 "counting-check"
647 }
648 async fn check(&self) -> HealthcheckResult {
649 self.0.fetch_add(1, Ordering::SeqCst);
650 HealthcheckResult::healthy()
651 }
652 }
653
654 let calls = Arc::new(AtomicUsize::new(0));
655 let registry = RestHealthcheckRegistry::new();
656 registry.register("counted-gear", Arc::new(CountingCheck(calls.clone())));
657
658 let first = registry.report(Duration::from_millis(500)).await;
659 let second = registry.report(Duration::from_millis(500)).await;
660
661 assert_eq!(first.status, HealthcheckStatus::Healthy);
662 assert_eq!(second.status, HealthcheckStatus::Healthy);
663 assert_eq!(
664 calls.load(Ordering::SeqCst),
665 1,
666 "second report() within REPORT_CACHE_TTL must reuse the cached result"
667 );
668 }
669
670 #[test]
671 fn sensitive_message_is_sanitized() {
672 assert_eq!(
673 sanitize_message("connection to postgres://user:pass@host/db failed"),
674 "health check failed"
675 );
676 assert_eq!(
677 sanitize_message("invalid token in header"),
678 "health check failed"
679 );
680 assert_eq!(
681 sanitize_message("thread panicked at some_file.rs:42"),
682 "health check failed"
683 );
684 assert_eq!(
685 sanitize_message("invalid api_key provided"),
686 "health check failed"
687 );
688 assert_eq!(
689 sanitize_message("connection to mongodb://host/db failed"),
690 "health check failed"
691 );
692 }
693
694 #[test]
695 fn long_message_is_sanitized() {
696 let long = "x".repeat(257);
697 assert_eq!(sanitize_message(&long), "health check failed");
698 }
699
700 #[test]
701 fn clean_message_passes_through() {
702 assert_eq!(
703 sanitize_message("upstream service unavailable"),
704 "upstream service unavailable"
705 );
706 }
707
708 #[test]
709 fn operator_useful_messages_pass_through() {
710 for msg in [
713 "tenant acme-corp migration pending",
714 "SELECT on read replica timed out",
715 "private subnet route unreachable",
716 ] {
717 assert_eq!(sanitize_message(msg), msg);
718 }
719 }
720
721 #[test]
722 fn code_allowlist_accepts_identifiers_and_drops_the_rest() {
723 assert_eq!(
724 sanitize_code("db_unreachable"),
725 Some("db_unreachable".to_owned())
726 );
727 assert_eq!(sanitize_code("v1.2-beta_3"), Some("v1.2-beta_3".to_owned()));
728 assert_eq!(sanitize_code(""), None);
730 assert_eq!(sanitize_code("DB_UNREACHABLE"), None);
731 assert_eq!(sanitize_code("postgres://user:pw@host/db"), None);
732 assert_eq!(sanitize_code("connection failed"), None);
733 assert_eq!(sanitize_code(&"a".repeat(MAX_CODE_LEN + 1)), None);
734 }
735
736 #[tokio::test]
737 async fn component_report_carries_conforming_code() {
738 struct CodedUnhealthy;
739 #[async_trait]
740 impl Healthcheck for CodedUnhealthy {
741 fn name(&self) -> &'static str {
742 "coded"
743 }
744 async fn check(&self) -> HealthcheckResult {
745 HealthcheckResult::unhealthy("database unreachable").with_code("db_unreachable")
746 }
747 }
748
749 let registry = RestHealthcheckRegistry::new();
750 registry.register("gear", Arc::new(CodedUnhealthy));
751 let report = registry.report(Duration::from_millis(500)).await;
752 assert_eq!(report.components[0].code.as_deref(), Some("db_unreachable"));
753 }
754
755 #[tokio::test]
756 async fn component_report_drops_nonconforming_code() {
757 struct BadCode;
758 #[async_trait]
759 impl Healthcheck for BadCode {
760 fn name(&self) -> &'static str {
761 "bad-code"
762 }
763 async fn check(&self) -> HealthcheckResult {
764 HealthcheckResult::unhealthy("down").with_code("postgres://u:p@host/db")
766 }
767 }
768
769 let registry = RestHealthcheckRegistry::new();
770 registry.register("gear", Arc::new(BadCode));
771 let report = registry.report(Duration::from_millis(500)).await;
772 assert!(report.components[0].code.is_none());
773 }
774}