1use std::sync::Arc;
9#[cfg(feature = "db")]
10use std::sync::RwLock;
11use std::sync::atomic::{AtomicBool, Ordering};
12
13use crate::extract::State;
14use axum::Json;
15use axum::http::StatusCode;
16use axum::response::IntoResponse;
17use serde::Serialize;
18
19pub trait ProvideProbeState {
25 fn probes(&self) -> &ProbeState;
28
29 fn health_detailed(&self) -> bool;
32
33 fn profile(&self) -> &str;
35
36 fn uptime_display(&self) -> String;
39
40 #[cfg(feature = "db")]
43 fn pool(
44 &self,
45 ) -> Option<&diesel_async::pooled_connection::deadpool::Pool<crate::db::RuntimeConnection>>;
46
47 #[cfg(feature = "db")]
49 fn replica_pool(
50 &self,
51 ) -> Option<&diesel_async::pooled_connection::deadpool::Pool<crate::db::RuntimeConnection>>
52 {
53 None
54 }
55
56 fn health_indicator_registry(&self) -> Option<&crate::actuator::HealthIndicatorRegistry> {
61 None
62 }
63
64 fn mark_startup_complete(&self) {
89 self.probes().mark_startup_complete();
90 }
91}
92
93#[derive(Clone, Debug, Default)]
99pub struct ProbeState {
100 startup_complete: Arc<AtomicBool>,
101 shutdown: Arc<Shutdown>,
102 #[cfg(feature = "db")]
103 replica_dependency: Arc<RwLock<ReplicaDependency>>,
104 #[cfg(feature = "db")]
105 db_checks: Arc<DbChecks>,
106}
107
108#[derive(Debug, Default)]
114struct Shutdown {
115 draining: AtomicBool,
117 signal: tokio_util::sync::CancellationToken,
121}
122
123#[cfg(feature = "db")]
126#[derive(Debug)]
127struct DbChecks {
128 primary: DbPingCheck,
129 replica: DbPingCheck,
130 primary_gates_readiness: AtomicBool,
132}
133
134#[cfg(feature = "db")]
135impl Default for DbChecks {
136 fn default() -> Self {
137 Self {
138 primary: DbPingCheck::new("primary"),
139 replica: DbPingCheck::new("replica"),
140 primary_gates_readiness: AtomicBool::new(true),
141 }
142 }
143}
144
145#[cfg(feature = "db")]
146#[derive(Clone, Debug, PartialEq, Eq)]
147pub(crate) struct ReplicaMigrationCheck {
148 pub(crate) primary_url: String,
149 pub(crate) replica_url: String,
150}
151
152#[cfg(feature = "db")]
153#[derive(Clone, Debug, PartialEq, Eq)]
154struct ReplicaDependency {
155 configured: bool,
156 fallback: crate::config::ReplicaFallback,
157 connection_ready: bool,
158 migrations_ready: bool,
159 migration_check: Option<ReplicaMigrationCheck>,
160 detail: Option<String>,
161 max_lag: Option<std::time::Duration>,
163 lag: Option<std::time::Duration>,
165 lag_at: Option<tokio::time::Instant>,
168 lag_started: Option<tokio::time::Instant>,
172 lag_detail: Option<String>,
174}
175
176#[cfg(feature = "db")]
177impl ReplicaDependency {
178 const fn base_ready(&self) -> bool {
180 !self.configured || (self.connection_ready && self.migrations_ready)
181 }
182
183 fn lag_age(&self) -> std::time::Duration {
185 self.lag_at
186 .map_or(std::time::Duration::ZERO, |at| at.elapsed())
187 }
188
189 fn lag_ok(&self) -> bool {
191 self.max_lag.is_none_or(|max| {
192 self.lag
193 .is_some_and(|lag| lag <= max && self.lag_age() <= sample_max_age(max))
194 })
195 }
196
197 fn lag_problem(&self) -> Option<String> {
198 if self.lag_ok() {
199 return None;
200 }
201 let max = self.max_lag.map_or(0, duration_ms);
202 Some(match (self.lag, &self.lag_detail) {
203 (Some(lag), _) if lag <= self.max_lag.unwrap_or_default() => format!(
204 "replica lag sample is {}ms old; the limit is {}ms",
205 duration_ms(self.lag_age()),
206 duration_ms(sample_max_age(self.max_lag.unwrap_or_default()))
207 ),
208 (Some(lag), _) => format!(
209 "replica lag {}ms exceeds database.replica_max_lag_ms {max}ms",
210 duration_ms(lag)
211 ),
212 (None, Some(detail)) => detail.clone(),
213 (None, None) => "replica lag is not measured yet".to_owned(),
214 })
215 }
216}
217
218#[cfg(feature = "db")]
223fn sample_max_age(max_lag: std::time::Duration) -> std::time::Duration {
224 max_lag
225 .saturating_mul(2)
226 .max(std::time::Duration::from_secs(1))
227}
228
229#[cfg(feature = "db")]
230fn duration_ms(duration: std::time::Duration) -> u64 {
231 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
232}
233
234#[cfg(feature = "db")]
235impl Default for ReplicaDependency {
236 fn default() -> Self {
237 Self {
238 configured: false,
239 fallback: crate::config::ReplicaFallback::default(),
240 connection_ready: true,
241 migrations_ready: true,
242 migration_check: None,
243 detail: None,
244 max_lag: None,
245 lag: None,
246 lag_at: None,
247 lag_started: None,
248 lag_detail: None,
249 }
250 }
251}
252
253#[cfg(feature = "db")]
254use crate::db_ping::{DbPingCheck, DbPingStatus, DbPool};
255
256impl ProbeState {
257 #[must_use]
259 pub fn pending_startup() -> Self {
260 Self::default()
261 }
262
263 #[must_use]
265 pub fn starting() -> Self {
266 Self::pending_startup()
267 }
268
269 #[must_use]
271 pub fn ready_for_test() -> Self {
272 let state = Self::pending_startup();
273 state.mark_startup_complete();
274 state
275 }
276
277 pub fn mark_startup_complete(&self) {
279 self.startup_complete.store(true, Ordering::Relaxed);
280 }
281
282 pub fn set_startup_complete(&self, complete: bool) {
284 self.startup_complete.store(complete, Ordering::Relaxed);
285 }
286
287 pub fn begin_shutdown(&self) {
289 self.shutdown.draining.store(true, Ordering::Relaxed);
290 }
291
292 pub(crate) fn shutdown_signal(&self) -> &tokio_util::sync::CancellationToken {
295 &self.shutdown.signal
296 }
297
298 pub fn begin_draining(&self) {
300 self.begin_shutdown();
301 }
302
303 pub fn set_draining(&self, draining: bool) {
305 self.shutdown.draining.store(draining, Ordering::Relaxed);
306 }
307
308 #[cfg(feature = "db")]
314 pub fn configure_replica_dependency(&self, fallback: crate::config::ReplicaFallback) {
315 let mut dependency = self
316 .replica_dependency
317 .write()
318 .expect("replica dependency lock poisoned");
319 *dependency = ReplicaDependency {
320 configured: true,
321 fallback,
322 connection_ready: false,
323 migrations_ready: true,
324 migration_check: None,
325 detail: Some("replica has not passed a readiness check".to_owned()),
326 max_lag: dependency.max_lag,
327 lag: None,
328 lag_at: None,
329 lag_started: None,
330 lag_detail: None,
331 };
332 }
333
334 #[cfg(feature = "db")]
340 pub fn configure_db_check(&self, cache_ttl: std::time::Duration, timeout: std::time::Duration) {
341 self.db_checks.primary.configure(cache_ttl, timeout);
342 self.db_checks.replica.configure(cache_ttl, timeout);
343 }
344
345 #[cfg(feature = "db")]
349 pub fn set_db_readiness(&self, gate: bool) {
350 self.db_checks
351 .primary_gates_readiness
352 .store(gate, Ordering::Relaxed);
353 }
354
355 #[cfg(feature = "db")]
356 fn primary_gates_readiness(&self) -> bool {
357 self.db_checks
358 .primary_gates_readiness
359 .load(Ordering::Relaxed)
360 }
361
362 #[cfg(all(test, feature = "db"))]
364 pub(crate) fn with_primary_ping(ping: Arc<dyn crate::db_ping::DbPing>) -> Self {
365 Self {
366 db_checks: Arc::new(DbChecks {
367 primary: DbPingCheck::with_ping("primary", ping),
368 ..DbChecks::default()
369 }),
370 ..Self::default()
371 }
372 }
373
374 #[cfg(all(test, feature = "db"))]
376 pub(crate) fn with_db_pings(
377 primary: Arc<dyn crate::db_ping::DbPing>,
378 replica: Arc<dyn crate::db_ping::DbPing>,
379 ) -> Self {
380 Self {
381 db_checks: Arc::new(DbChecks {
382 primary: DbPingCheck::with_ping("primary", primary),
383 replica: DbPingCheck::with_ping("replica", replica),
384 ..DbChecks::default()
385 }),
386 ..Self::default()
387 }
388 }
389
390 #[cfg(feature = "db")]
392 pub(crate) async fn check_primary_db(&self, pool: &DbPool) -> DbPingStatus {
393 self.db_checks.primary.check(pool).await
394 }
395
396 #[cfg(feature = "db")]
399 pub(crate) async fn check_primary_db_uncached(pool: &DbPool) -> DbPingStatus {
400 let check = DbPingCheck::new("primary");
401 check.configure(
402 std::time::Duration::ZERO,
403 crate::health_cache::DEFAULT_PING_TIMEOUT,
404 );
405 check.check(pool).await
406 }
407
408 #[cfg(feature = "db")]
410 pub(crate) fn configure_replica_migration_check(
411 &self,
412 primary_url: impl Into<String>,
413 replica_url: impl Into<String>,
414 ) {
415 let mut dependency = self
416 .replica_dependency
417 .write()
418 .expect("replica dependency lock poisoned");
419 dependency.migration_check = Some(ReplicaMigrationCheck {
420 primary_url: primary_url.into(),
421 replica_url: replica_url.into(),
422 });
423 }
424
425 #[cfg(feature = "db")]
431 pub fn mark_replica_connection_ready(&self) {
432 let mut dependency = self
433 .replica_dependency
434 .write()
435 .expect("replica dependency lock poisoned");
436 dependency.connection_ready = true;
437 if dependency.migrations_ready {
438 dependency.detail = None;
439 }
440 }
441
442 #[cfg(feature = "db")]
448 pub fn mark_replica_connection_unready(&self, detail: impl Into<String>) {
449 let mut dependency = self
450 .replica_dependency
451 .write()
452 .expect("replica dependency lock poisoned");
453 dependency.connection_ready = false;
454 dependency.detail = Some(detail.into());
455 }
456
457 #[cfg(feature = "db")]
463 pub fn mark_replica_migrations_ready(&self) {
464 let mut dependency = self
465 .replica_dependency
466 .write()
467 .expect("replica dependency lock poisoned");
468 dependency.migrations_ready = true;
469 if dependency.connection_ready {
470 dependency.detail = None;
471 }
472 }
473
474 #[cfg(feature = "db")]
476 pub(crate) fn mark_replica_migrations_unready(&self, detail: impl Into<String>) {
477 let mut dependency = self
478 .replica_dependency
479 .write()
480 .expect("replica dependency lock poisoned");
481 dependency.migrations_ready = false;
482 dependency.detail = Some(detail.into());
483 }
484
485 #[cfg(feature = "db")]
491 pub fn mark_replica_ready(&self) {
492 let mut dependency = self
493 .replica_dependency
494 .write()
495 .expect("replica dependency lock poisoned");
496 dependency.connection_ready = true;
497 dependency.migrations_ready = true;
498 dependency.detail = None;
499 }
500
501 #[cfg(feature = "db")]
507 pub fn mark_replica_unready(&self, detail: impl Into<String>) {
508 let mut dependency = self
509 .replica_dependency
510 .write()
511 .expect("replica dependency lock poisoned");
512 dependency.connection_ready = false;
513 dependency.migrations_ready = false;
514 dependency.detail = Some(detail.into());
515 }
516
517 #[must_use]
519 pub fn is_startup_complete(&self) -> bool {
520 self.startup_complete.load(Ordering::Relaxed)
521 }
522
523 #[must_use]
525 pub fn is_shutting_down(&self) -> bool {
526 self.shutdown.draining.load(Ordering::Relaxed)
527 }
528
529 #[must_use]
531 pub fn draining(&self) -> bool {
532 self.is_shutting_down()
533 }
534
535 #[cfg(feature = "db")]
536 pub(crate) fn replica_allows_readiness(&self) -> bool {
537 let dependency = self
538 .replica_dependency
539 .read()
540 .expect("replica dependency lock poisoned");
541 dependency.base_ready()
544 || matches!(dependency.fallback, crate::config::ReplicaFallback::Primary)
545 }
546
547 #[cfg(feature = "db")]
548 pub(crate) fn should_route_reads_to_replica(&self) -> bool {
549 let dependency = self
550 .replica_dependency
551 .read()
552 .expect("replica dependency lock poisoned");
553 dependency.base_ready() && dependency.lag_ok()
554 }
555
556 #[cfg(feature = "db")]
557 pub(crate) fn should_fallback_reads_to_primary(&self) -> bool {
558 let (base_ready, lag_ok, fallback) = {
559 let dependency = self
560 .replica_dependency
561 .read()
562 .expect("replica dependency lock poisoned");
563 (
564 dependency.base_ready(),
565 dependency.lag_ok(),
566 dependency.fallback,
567 )
568 };
569 let stale_only = base_ready && !lag_ok;
572 let down_and_allowed =
573 !base_ready && matches!(fallback, crate::config::ReplicaFallback::Primary);
574 stale_only || down_and_allowed
575 }
576
577 #[cfg(feature = "db")]
583 pub fn configure_replica_max_lag(&self, max_lag: Option<std::time::Duration>) {
584 self.replica_dependency
585 .write()
586 .expect("replica dependency lock poisoned")
587 .max_lag = max_lag;
588 }
589
590 #[cfg(feature = "db")]
596 pub fn record_replica_lag(&self, lag: std::time::Duration) {
597 let mut dependency = self
598 .replica_dependency
599 .write()
600 .expect("replica dependency lock poisoned");
601 let now = tokio::time::Instant::now();
602 dependency.lag = Some(lag);
603 dependency.lag_at = Some(now);
604 dependency.lag_started = Some(now);
605 dependency.lag_detail = None;
606 }
607
608 #[cfg(feature = "db")]
615 pub fn mark_replica_lag_unknown(&self, detail: impl Into<String>) {
616 let mut dependency = self
617 .replica_dependency
618 .write()
619 .expect("replica dependency lock poisoned");
620 dependency.lag = None;
621 dependency.lag_at = None;
622 dependency.lag_started = Some(tokio::time::Instant::now());
623 dependency.lag_detail = Some(detail.into());
624 }
625
626 #[cfg(feature = "db")]
630 fn record_replica_lag_sample(
631 &self,
632 started: tokio::time::Instant,
633 result: Result<std::time::Duration, String>,
634 ) -> bool {
635 let mut dependency = self
636 .replica_dependency
637 .write()
638 .expect("replica dependency lock poisoned");
639 if dependency.lag_started.is_some_and(|newer| started < newer) {
640 return false;
641 }
642 dependency.lag_started = Some(started);
643 match result {
644 Ok(lag) => {
645 dependency.lag = Some(lag);
646 dependency.lag_at = Some(started);
648 dependency.lag_detail = None;
649 }
650 Err(error) => {
651 dependency.lag = None;
652 dependency.lag_at = None;
653 dependency.lag_detail = Some(format!("replica lag check failed: {error}"));
654 }
655 }
656 true
657 }
658
659 #[cfg(feature = "db")]
665 #[must_use]
666 pub fn replica_lag(&self) -> Option<std::time::Duration> {
667 self.replica_dependency
668 .read()
669 .expect("replica dependency lock poisoned")
670 .lag
671 }
672
673 #[cfg(feature = "db")]
675 pub(crate) fn replica_lag_ok(&self) -> bool {
676 self.replica_dependency
677 .read()
678 .expect("replica dependency lock poisoned")
679 .lag_ok()
680 }
681
682 #[cfg(feature = "db")]
684 pub(crate) fn replica_max_lag(&self) -> Option<std::time::Duration> {
685 self.replica_dependency
686 .read()
687 .expect("replica dependency lock poisoned")
688 .max_lag
689 }
690
691 #[cfg(feature = "db")]
697 #[must_use]
698 pub fn replica_status(&self) -> Option<ReplicaStatus> {
699 let dependency = self
700 .replica_dependency
701 .read()
702 .expect("replica dependency lock poisoned");
703 if !dependency.configured && dependency.max_lag.is_none() {
704 return None;
705 }
706 let base_ready = dependency.base_ready();
707 Some(ReplicaStatus {
708 ready: base_ready && dependency.lag_ok(),
709 lag_ms: dependency.lag.map(duration_ms),
710 max_lag_ms: dependency.max_lag.map(duration_ms),
711 detail: if base_ready {
713 dependency.lag_problem()
714 } else {
715 dependency.detail.clone()
716 }
717 .map(|detail| crate::db_url::redact_targets_in_message(&detail)),
718 })
719 }
720
721 #[cfg(feature = "db")]
722 pub(crate) fn replica_configured(&self) -> bool {
723 self.replica_dependency
724 .read()
725 .expect("replica dependency lock poisoned")
726 .configured
727 }
728
729 #[cfg(feature = "db")]
730 pub(crate) fn replica_migration_check(&self) -> Option<ReplicaMigrationCheck> {
731 self.replica_dependency
732 .read()
733 .expect("replica dependency lock poisoned")
734 .migration_check
735 .clone()
736 }
737}
738
739#[derive(Clone, Copy)]
740enum ProbeKind {
741 Live,
742 Ready,
743 Startup,
744}
745
746#[derive(Serialize)]
747pub(crate) struct ProbeResponse {
748 status: &'static str,
749 #[serde(skip_serializing_if = "Option::is_none")]
750 version: Option<&'static str>,
751 #[serde(skip_serializing_if = "Option::is_none")]
752 profile: Option<String>,
753 #[serde(skip_serializing_if = "Option::is_none")]
754 uptime: Option<String>,
755 #[serde(skip_serializing_if = "Option::is_none")]
756 pool: Option<PoolStatus>,
757 #[cfg(feature = "db")]
758 #[serde(skip_serializing_if = "Option::is_none")]
759 replica: Option<ReplicaStatus>,
760 #[serde(skip_serializing_if = "Option::is_none")]
761 database: Option<DatabaseStatus>,
762}
763
764#[cfg(feature = "db")]
766#[derive(Clone, Debug, PartialEq, Eq, Serialize)]
767pub struct ReplicaStatus {
768 pub ready: bool,
770 pub lag_ms: Option<u64>,
772 pub max_lag_ms: Option<u64>,
774 #[serde(skip_serializing_if = "Option::is_none")]
776 pub detail: Option<String>,
777}
778
779#[derive(Serialize)]
781pub(crate) struct DatabaseStatus {
782 status: &'static str,
783 #[serde(skip_serializing_if = "Option::is_none")]
784 error: Option<String>,
785}
786
787#[derive(Serialize)]
788pub(crate) struct PoolStatus {
789 size: u64,
790 available: u64,
791 waiting: u64,
792}
793
794#[allow(clippy::missing_const_for_fn, unused_variables)]
799fn dependency_readiness<S: ProvideProbeState>(state: &S) -> (bool, Option<PoolStatus>) {
800 #[cfg(feature = "db")]
801 {
802 let replica_ready_for_policy = state.probes().replica_allows_readiness();
803 let pool_status = state.pool().map(|pool| {
804 let status = pool.status();
805 PoolStatus {
806 size: status.max_size as u64,
807 available: status.available as u64,
808 waiting: status.waiting as u64,
809 }
810 });
811
812 (replica_ready_for_policy, pool_status)
813 }
814
815 #[cfg(not(feature = "db"))]
816 {
817 (true, None)
818 }
819}
820
821#[cfg(feature = "db")]
822async fn refresh_replica_readiness<S: ProvideProbeState + Sync>(state: &S) {
823 if !state.probes().replica_configured() {
824 return;
825 }
826
827 let Some(replica_pool) = state.replica_pool() else {
828 state
829 .probes()
830 .mark_replica_connection_unready("replica pool is not available");
831 return;
832 };
833
834 let status = state.probes().db_checks.replica.check(replica_pool).await;
837 if status.up {
838 state.probes().mark_replica_connection_ready();
839 refresh_replica_migration_readiness(state).await;
840 } else {
841 state.probes().mark_replica_connection_unready(
842 status
843 .error
844 .unwrap_or_else(|| "replica ping failed".to_owned()),
845 );
846 }
847}
848
849#[cfg(feature = "db")]
850async fn refresh_replica_migration_readiness<S: ProvideProbeState + Sync>(state: &S) {
851 refresh_replica_migration_readiness_with(state, |check| {
852 crate::migrate::check_replica_migration_readiness_blocking(
853 check.primary_url,
854 check.replica_url,
855 )
856 })
857 .await;
858}
859
860#[cfg(feature = "db")]
861async fn refresh_replica_migration_readiness_with<S, F, Fut>(state: &S, check_readiness: F)
862where
863 S: ProvideProbeState + Sync,
864 F: FnOnce(ReplicaMigrationCheck) -> Fut,
865 Fut: std::future::Future<Output = crate::migrate::ReplicaMigrationReadiness>,
866{
867 let Some(check) = state.probes().replica_migration_check() else {
868 return;
869 };
870
871 let readiness = check_readiness(check).await;
872
873 if readiness.is_ready() {
874 state.probes().mark_replica_migrations_ready();
875 } else if let Some(detail) = readiness.detail() {
876 state.probes().mark_replica_migrations_unready(detail);
877 }
878}
879
880#[cfg(feature = "db")]
882pub(crate) fn replica_lag_query_budget(max_lag: std::time::Duration) -> std::time::Duration {
883 max_lag.max(std::time::Duration::from_secs(1))
884}
885
886#[cfg(feature = "db")]
888struct DiscardUnlessFinished<M: deadpool::managed::Manager>(Option<deadpool::managed::Object<M>>);
889
890#[cfg(feature = "db")]
891impl<M: deadpool::managed::Manager> Drop for DiscardUnlessFinished<M> {
892 fn drop(&mut self) {
893 if let Some(conn) = self.0.take() {
894 drop(deadpool::managed::Object::take(conn));
896 }
897 }
898}
899
900#[cfg(feature = "db")]
907pub(crate) async fn refresh_replica_lag_bounded<M, F>(
908 probes: &ProbeState,
909 conn: deadpool::managed::Object<M>,
910 budget: std::time::Duration,
911 measure: F,
912) where
913 M: deadpool::managed::Manager,
914 F: for<'c> FnOnce(
915 &'c mut deadpool::managed::Object<M>,
916 ) -> std::pin::Pin<
917 Box<dyn std::future::Future<Output = Result<std::time::Duration, String>> + Send + 'c>,
918 >,
919{
920 let started = tokio::time::Instant::now();
921 let mut guard = DiscardUnlessFinished(Some(conn));
922 let Some(conn) = guard.0.as_mut() else {
923 return;
924 };
925 let measured = tokio::time::timeout(budget, measure(conn))
926 .await
927 .map_or_else(
928 |_| Err(format!("query took over {}ms", budget.as_millis())),
929 |result| {
930 drop(guard.0.take());
932 result
933 },
934 );
935 apply_replica_lag_sample(probes, started, measured);
936}
937
938#[cfg(feature = "db")]
942pub(crate) fn apply_replica_lag_sample(
943 probes: &ProbeState,
944 started: tokio::time::Instant,
945 measured: Result<std::time::Duration, String>,
946) {
947 let was_fresh = probes.replica_lag_ok();
948 if !probes.record_replica_lag_sample(started, measured) {
949 return;
950 }
951 let fresh = probes.replica_lag_ok();
952 if was_fresh && !fresh {
953 let detail = probes
954 .replica_status()
955 .and_then(|status| status.detail)
956 .unwrap_or_default();
957 tracing::warn!(target: "autumn::db", %detail, "replica is stale: reads use the primary");
958 } else if !was_fresh && fresh {
959 tracing::info!(target: "autumn::db", "replica is fresh again: reads use the replica");
960 }
961}
962
963fn probe_response<S: ProvideProbeState>(
964 state: &S,
965 kind: ProbeKind,
966 checks_ready: bool,
967) -> (StatusCode, Json<ProbeResponse>) {
968 let startup_complete = state.probes().is_startup_complete();
969 let shutting_down = state.probes().is_shutting_down();
970 let (dependencies_ready, pool_status) = dependency_readiness(state);
971
972 let (status_code, status) = match kind {
973 ProbeKind::Live => (StatusCode::OK, "ok"),
974 ProbeKind::Startup if startup_complete => (StatusCode::OK, "ok"),
975 ProbeKind::Startup => (StatusCode::SERVICE_UNAVAILABLE, "starting"),
976 ProbeKind::Ready
977 if startup_complete && !shutting_down && dependencies_ready && checks_ready =>
978 {
979 (StatusCode::OK, "ok")
980 }
981 ProbeKind::Ready => (StatusCode::SERVICE_UNAVAILABLE, "degraded"),
982 };
983
984 let detailed = state.health_detailed();
985 let body = ProbeResponse {
986 status,
987 version: if detailed {
988 Some(env!("CARGO_PKG_VERSION"))
989 } else {
990 None
991 },
992 profile: if detailed {
993 Some(state.profile().to_owned())
994 } else {
995 None
996 },
997 uptime: if detailed {
998 Some(state.uptime_display())
999 } else {
1000 None
1001 },
1002 pool: if detailed { pool_status } else { None },
1003 #[cfg(feature = "db")]
1004 replica: if detailed && matches!(kind, ProbeKind::Ready) {
1005 state.probes().replica_status()
1006 } else {
1007 None
1008 },
1009 database: None,
1010 };
1011
1012 (status_code, Json(body))
1013}
1014
1015async fn check_readiness_indicators<S: ProvideProbeState + Sync>(state: &S) -> bool {
1020 let Some(registry) = state.health_indicator_registry() else {
1021 return true;
1022 };
1023 let results = registry.run_readiness().await;
1024 let statuses: Vec<crate::actuator::HealthStatus> =
1025 results.iter().map(|r| r.output.status).collect();
1026 crate::actuator::HealthIndicatorRegistry::aggregate_status(&statuses).is_healthy()
1027}
1028
1029#[cfg(feature = "db")]
1031async fn check_primary_db<S: ProvideProbeState + Sync>(state: &S) -> Option<DbPingStatus> {
1032 let pool = state.pool()?;
1033 Some(state.probes().check_primary_db(pool).await)
1034}
1035
1036async fn readiness<S: ProvideProbeState + Sync>(state: &S) -> (StatusCode, Json<ProbeResponse>) {
1038 let probes = state.probes();
1041 if !probes.is_startup_complete() || probes.is_shutting_down() {
1042 #[cfg(feature = "db")]
1043 refresh_replica_readiness(state).await;
1044 return probe_response(state, ProbeKind::Ready, true);
1045 }
1046
1047 #[cfg(feature = "db")]
1048 {
1049 let primary_gates = state.probes().primary_gates_readiness();
1055 let primary_check = async {
1056 if primary_gates {
1057 check_primary_db(state).await
1058 } else {
1059 None
1060 }
1061 };
1062 let ((), primary, indicators_ready) = tokio::join!(
1063 refresh_replica_readiness(state),
1064 primary_check,
1065 check_readiness_indicators(state)
1066 );
1067 let primary_ready = primary.as_ref().is_none_or(|status| status.up);
1068 let (code, Json(mut body)) =
1069 probe_response(state, ProbeKind::Ready, primary_ready && indicators_ready);
1070 if state.health_detailed() {
1071 body.database = primary.map(|status| DatabaseStatus {
1072 status: if status.up { "ok" } else { "down" },
1073 error: status.error,
1074 });
1075 }
1076 (code, Json(body))
1077 }
1078
1079 #[cfg(not(feature = "db"))]
1080 {
1081 let indicators_ready = check_readiness_indicators(state).await;
1082 probe_response(state, ProbeKind::Ready, indicators_ready)
1083 }
1084}
1085
1086pub async fn live_handler<S: ProvideProbeState + Send + Sync + 'static>(
1088 State(state): State<S>,
1089) -> impl IntoResponse {
1090 probe_response(&state, ProbeKind::Live, true)
1091}
1092
1093pub async fn ready_handler<S: ProvideProbeState + Send + Sync + 'static>(
1095 State(state): State<S>,
1096) -> impl IntoResponse {
1097 readiness(&state).await
1098}
1099
1100pub async fn startup_handler<S: ProvideProbeState + Send + Sync + 'static>(
1102 State(state): State<S>,
1103) -> impl IntoResponse {
1104 probe_response(&state, ProbeKind::Startup, true)
1105}
1106
1107pub(crate) async fn readiness_response<S: ProvideProbeState + Sync>(
1109 state: &S,
1110) -> (StatusCode, Json<ProbeResponse>) {
1111 readiness(state).await
1112}
1113
1114#[cfg(test)]
1115mod tests {
1116 use super::*;
1117
1118 struct TestProbeState {
1119 probes: ProbeState,
1120 health_detailed: bool,
1121 profile: String,
1122 }
1123
1124 impl ProvideProbeState for TestProbeState {
1125 fn probes(&self) -> &ProbeState {
1126 &self.probes
1127 }
1128
1129 fn health_detailed(&self) -> bool {
1130 self.health_detailed
1131 }
1132
1133 fn profile(&self) -> &str {
1134 &self.profile
1135 }
1136
1137 fn uptime_display(&self) -> String {
1138 "test uptime".to_string()
1139 }
1140
1141 #[cfg(feature = "db")]
1142 fn pool(
1143 &self,
1144 ) -> Option<&diesel_async::pooled_connection::deadpool::Pool<crate::db::RuntimeConnection>>
1145 {
1146 None
1147 }
1148 }
1149
1150 impl TestProbeState {
1151 fn new() -> Self {
1152 Self {
1153 probes: ProbeState::pending_startup(),
1154 health_detailed: true,
1155 profile: "test".to_string(),
1156 }
1157 }
1158 }
1159
1160 #[test]
1161 fn test_live_handler_returns_ok() {
1162 let state = TestProbeState::new();
1163 let (status, Json(response)) = probe_response(&state, ProbeKind::Live, true);
1164 assert_eq!(status, StatusCode::OK);
1165 assert_eq!(response.status, "ok");
1166 }
1167
1168 #[tokio::test]
1169 async fn test_startup_handler_pending() {
1170 let state = TestProbeState::new();
1171 let (status, Json(response)) = probe_response(&state, ProbeKind::Startup, true);
1172 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1173 assert_eq!(response.status, "starting");
1174 }
1175
1176 #[tokio::test]
1177 async fn test_startup_handler_complete() {
1178 let state = TestProbeState::new();
1179 state.mark_startup_complete();
1180 let (status, Json(response)) = probe_response(&state, ProbeKind::Startup, true);
1181 assert_eq!(status, StatusCode::OK);
1182 assert_eq!(response.status, "ok");
1183 }
1184
1185 #[tokio::test]
1186 async fn test_ready_handler_pending_startup() {
1187 let state = TestProbeState::new();
1188 let (status, Json(response)) = probe_response(&state, ProbeKind::Ready, true);
1189 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1190 assert_eq!(response.status, "degraded");
1191 }
1192
1193 #[tokio::test]
1194 async fn test_ready_handler_complete_startup() {
1195 let state = TestProbeState::new();
1196 state.mark_startup_complete();
1197 let (status, Json(response)) = probe_response(&state, ProbeKind::Ready, true);
1198 assert_eq!(status, StatusCode::OK);
1199 assert_eq!(response.status, "ok");
1200 }
1201
1202 #[tokio::test]
1203 async fn test_ready_handler_shutting_down() {
1204 let state = TestProbeState::new();
1205 state.mark_startup_complete();
1206 state.probes().begin_shutdown();
1207 let (status, Json(response)) = probe_response(&state, ProbeKind::Ready, true);
1208 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1209 assert_eq!(response.status, "degraded");
1210 }
1211
1212 #[cfg(feature = "db")]
1213 #[tokio::test]
1214 async fn replica_lag_alone_keeps_the_pod_ready_and_is_reported() {
1215 let state = TestProbeState::new();
1216 state.mark_startup_complete();
1217 let probes = state.probes();
1218 probes.configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1219 probes.mark_replica_ready();
1220 probes.configure_replica_max_lag(Some(std::time::Duration::from_secs(1)));
1221 probes.record_replica_lag(std::time::Duration::from_secs(30));
1222
1223 let (status, Json(response)) = probe_response(&state, ProbeKind::Ready, true);
1224
1225 assert_eq!(
1226 status,
1227 StatusCode::OK,
1228 "reads use the primary, so stay ready"
1229 );
1230 let replica = response
1231 .replica
1232 .expect("detailed probe reports the replica");
1233 assert!(!replica.ready);
1234 assert_eq!(replica.lag_ms, Some(30_000));
1235 assert_eq!(replica.max_lag_ms, Some(1_000));
1236 let body = serde_json::to_value(&replica).unwrap();
1237 assert_eq!(body["lag_ms"], 30_000);
1238 }
1239
1240 #[cfg(feature = "db")]
1241 #[tokio::test]
1242 async fn replica_lag_refresh_records_the_measured_lag() {
1243 let state = TestProbeState::new();
1244 let probes = state.probes();
1245 probes.configure_replica_dependency(crate::config::ReplicaFallback::Primary);
1246 probes.mark_replica_ready();
1247 probes.configure_replica_max_lag(Some(std::time::Duration::from_millis(500)));
1248
1249 apply_replica_lag_sample(
1250 probes,
1251 tokio::time::Instant::now(),
1252 Ok(std::time::Duration::from_millis(900)),
1253 );
1254 assert_eq!(
1255 probes.replica_lag(),
1256 Some(std::time::Duration::from_millis(900))
1257 );
1258 assert!(!probes.should_route_reads_to_replica());
1259 assert!(probes.should_fallback_reads_to_primary());
1260
1261 apply_replica_lag_sample(
1262 probes,
1263 tokio::time::Instant::now(),
1264 Err("no route to host".to_owned()),
1265 );
1266 assert_eq!(probes.replica_lag(), None);
1267 assert!(!probes.should_route_reads_to_replica());
1268
1269 apply_replica_lag_sample(
1270 probes,
1271 tokio::time::Instant::now(),
1272 Ok(std::time::Duration::ZERO),
1273 );
1274 assert!(probes.should_route_reads_to_replica());
1275 }
1276
1277 #[cfg(feature = "db")]
1278 #[tokio::test(start_paused = true)]
1279 async fn an_old_lag_sample_ages_out() {
1280 let probes = ProbeState::default();
1281 probes.configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1282 probes.mark_replica_ready();
1283 probes.configure_replica_max_lag(Some(std::time::Duration::from_secs(1)));
1284 probes.record_replica_lag(std::time::Duration::from_millis(100));
1285 assert!(probes.should_route_reads_to_replica());
1286
1287 tokio::time::advance(std::time::Duration::from_millis(1500)).await;
1289 assert!(probes.should_route_reads_to_replica());
1290
1291 tokio::time::advance(std::time::Duration::from_secs(1)).await;
1292 assert!(
1293 !probes.should_route_reads_to_replica(),
1294 "no new sample for 2.5s, so freshness is unknown"
1295 );
1296 assert!(probes.should_fallback_reads_to_primary());
1297 let detail = probes.replica_status().unwrap().detail.unwrap();
1298 assert!(detail.contains("old"), "{detail}");
1299 }
1300
1301 #[cfg(feature = "db")]
1302 #[tokio::test(start_paused = true)]
1303 async fn a_small_lag_limit_does_not_flap_between_samples() {
1304 let probes = ProbeState::default();
1305 probes.configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1306 probes.mark_replica_ready();
1307 probes.configure_replica_max_lag(Some(std::time::Duration::from_millis(100)));
1308 probes.record_replica_lag(std::time::Duration::ZERO);
1309 tokio::time::advance(std::time::Duration::from_millis(300)).await;
1311 assert!(probes.should_route_reads_to_replica());
1312 }
1313
1314 #[cfg(feature = "db")]
1315 #[tokio::test(start_paused = true)]
1316 async fn an_older_lag_measurement_does_not_overwrite_a_newer_one() {
1317 let probes = ProbeState::default();
1318 probes.configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1319 probes.mark_replica_ready();
1320 probes.configure_replica_max_lag(Some(std::time::Duration::from_secs(1)));
1321 let earlier = tokio::time::Instant::now();
1323 tokio::time::advance(std::time::Duration::from_millis(10)).await;
1324 let later = tokio::time::Instant::now();
1325 apply_replica_lag_sample(&probes, later, Ok(std::time::Duration::from_secs(30)));
1326 assert!(
1327 !probes.should_route_reads_to_replica(),
1328 "the replica is stale"
1329 );
1330 apply_replica_lag_sample(&probes, earlier, Ok(std::time::Duration::ZERO));
1331 assert!(
1332 !probes.should_route_reads_to_replica(),
1333 "an older measurement must not route reads back to a stale replica"
1334 );
1335 assert_eq!(
1336 probes.replica_lag(),
1337 Some(std::time::Duration::from_secs(30))
1338 );
1339 tokio::time::advance(std::time::Duration::from_millis(10)).await;
1341 apply_replica_lag_sample(
1342 &probes,
1343 tokio::time::Instant::now(),
1344 Ok(std::time::Duration::ZERO),
1345 );
1346 assert!(probes.should_route_reads_to_replica());
1347 }
1348
1349 #[cfg(feature = "db")]
1350 #[test]
1351 fn replica_status_is_absent_without_a_replica() {
1352 assert!(ProbeState::default().replica_status().is_none());
1353 }
1354
1355 #[cfg(feature = "db")]
1356 #[tokio::test]
1357 async fn ready_fails_when_replica_is_unready_and_policy_is_fail_readiness() {
1358 let state = TestProbeState::new();
1359 state.mark_startup_complete();
1360 state
1361 .probes()
1362 .configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1363 state
1364 .probes()
1365 .mark_replica_unready("replica migrations lag primary");
1366
1367 let (status, Json(response)) = probe_response(&state, ProbeKind::Ready, true);
1368
1369 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1370 assert_eq!(response.status, "degraded");
1371 }
1372
1373 #[cfg(feature = "db")]
1374 #[tokio::test]
1375 async fn ready_fails_when_replica_is_configured_but_not_checked() {
1376 let state = TestProbeState::new();
1377 state.mark_startup_complete();
1378 state
1379 .probes()
1380 .configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1381
1382 let (status, Json(response)) = readiness_response(&state).await;
1383
1384 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1385 assert_eq!(response.status, "degraded");
1386 }
1387
1388 #[cfg(feature = "db")]
1389 #[test]
1390 fn replica_migration_lag_can_recover_without_resetting_connection_readiness() {
1391 let probes = ProbeState::ready_for_test();
1392 probes.configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1393 probes.mark_replica_connection_ready();
1394 probes.mark_replica_migrations_unready("replica migrations lag primary");
1395
1396 assert!(!probes.replica_allows_readiness());
1397 assert!(!probes.should_route_reads_to_replica());
1398
1399 probes.mark_replica_connection_ready();
1400 assert!(!probes.replica_allows_readiness());
1401 assert!(!probes.should_route_reads_to_replica());
1402
1403 probes.mark_replica_migrations_ready();
1404 assert!(probes.replica_allows_readiness());
1405 assert!(probes.should_route_reads_to_replica());
1406 }
1407
1408 #[cfg(feature = "db")]
1409 #[test]
1410 fn replica_migration_retry_urls_are_stored_for_readiness_rechecks() {
1411 let probes = ProbeState::ready_for_test();
1412 probes.configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1413 probes.configure_replica_migration_check(
1414 "postgres://localhost/primary",
1415 "postgres://localhost/replica",
1416 );
1417
1418 let check = probes
1419 .replica_migration_check()
1420 .expect("migration check should be configured");
1421
1422 assert_eq!(check.primary_url, "postgres://localhost/primary");
1423 assert_eq!(check.replica_url, "postgres://localhost/replica");
1424 }
1425
1426 #[cfg(feature = "db")]
1427 #[tokio::test]
1428 async fn replica_migration_readiness_rechecks_after_initial_ready_state() {
1429 let state = TestProbeState::new();
1430 state.mark_startup_complete();
1431 state
1432 .probes()
1433 .configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1434 state.probes().configure_replica_migration_check(
1435 "postgres://localhost/primary",
1436 "postgres://localhost/replica",
1437 );
1438 state.probes().mark_replica_connection_ready();
1439 state.probes().mark_replica_migrations_ready();
1440
1441 let checks = std::sync::atomic::AtomicUsize::new(0);
1442 refresh_replica_migration_readiness_with(&state, |check| {
1443 checks.fetch_add(1, Ordering::Relaxed);
1444 assert_eq!(check.primary_url, "postgres://localhost/primary");
1445 assert_eq!(check.replica_url, "postgres://localhost/replica");
1446 std::future::ready(crate::migrate::ReplicaMigrationReadiness::Stale {
1447 primary_latest: Some("20260511000000".to_owned()),
1448 replica_latest: Some("20260510000000".to_owned()),
1449 })
1450 })
1451 .await;
1452
1453 assert_eq!(checks.load(Ordering::Relaxed), 1);
1454 assert!(!state.probes().replica_allows_readiness());
1455 assert!(!state.probes().should_route_reads_to_replica());
1456 }
1457
1458 #[cfg(feature = "db")]
1459 #[tokio::test]
1460 async fn ready_allows_primary_fallback_when_replica_is_unready() {
1461 let state = TestProbeState::new();
1462 state.mark_startup_complete();
1463 state
1464 .probes()
1465 .configure_replica_dependency(crate::config::ReplicaFallback::Primary);
1466 state
1467 .probes()
1468 .mark_replica_unready("replica migrations lag primary");
1469
1470 let (status, Json(response)) = probe_response(&state, ProbeKind::Ready, true);
1471
1472 assert_eq!(status, StatusCode::OK);
1473 assert_eq!(response.status, "ok");
1474 }
1475
1476 #[tokio::test]
1477 async fn test_probe_state_set_draining() {
1478 let state = ProbeState::starting();
1479 assert!(!state.draining());
1480 state.set_draining(true);
1481 assert!(state.draining());
1482 }
1483
1484 #[tokio::test]
1485 async fn test_probe_state_set_startup_complete() {
1486 let state = ProbeState::starting();
1487 assert!(!state.is_startup_complete());
1488 state.set_startup_complete(true);
1489 assert!(state.is_startup_complete());
1490 }
1491
1492 #[tokio::test]
1493 async fn test_ready_for_test() {
1494 let state = ProbeState::ready_for_test();
1495 assert!(state.is_startup_complete());
1496 }
1497
1498 #[tokio::test]
1499 async fn test_health_detailed_false() {
1500 let mut state = TestProbeState::new();
1501 state.health_detailed = false;
1502
1503 let (_, Json(response)) = probe_response(&state, ProbeKind::Live, true);
1504 assert!(response.version.is_none());
1505 assert!(response.profile.is_none());
1506 assert!(response.uptime.is_none());
1507 assert!(response.pool.is_none());
1508 }
1509
1510 #[tokio::test]
1511 async fn test_begin_draining() {
1512 let state = ProbeState::ready_for_test();
1513 assert!(!state.draining());
1514 state.begin_draining();
1515 assert!(state.draining());
1516 }
1517
1518 struct TestProbeStateWithIndicators {
1521 probes: ProbeState,
1522 health_detailed: bool,
1523 profile: String,
1524 registry: crate::actuator::HealthIndicatorRegistry,
1525 }
1526
1527 impl ProvideProbeState for TestProbeStateWithIndicators {
1528 fn probes(&self) -> &ProbeState {
1529 &self.probes
1530 }
1531
1532 fn health_detailed(&self) -> bool {
1533 self.health_detailed
1534 }
1535
1536 fn profile(&self) -> &str {
1537 &self.profile
1538 }
1539
1540 fn uptime_display(&self) -> String {
1541 "test uptime".to_string()
1542 }
1543
1544 #[cfg(feature = "db")]
1545 fn pool(
1546 &self,
1547 ) -> Option<&diesel_async::pooled_connection::deadpool::Pool<crate::db::RuntimeConnection>>
1548 {
1549 None
1550 }
1551
1552 fn health_indicator_registry(&self) -> Option<&crate::actuator::HealthIndicatorRegistry> {
1553 Some(&self.registry)
1554 }
1555 }
1556
1557 impl TestProbeStateWithIndicators {
1558 fn new(registry: crate::actuator::HealthIndicatorRegistry) -> Self {
1559 let probes = ProbeState::pending_startup();
1560 probes.mark_startup_complete();
1561 Self {
1562 probes,
1563 health_detailed: true,
1564 profile: "test".to_string(),
1565 registry,
1566 }
1567 }
1568 }
1569
1570 struct AlwaysDown;
1571 impl crate::actuator::HealthIndicator for AlwaysDown {
1572 fn check(&self) -> futures::future::BoxFuture<'_, crate::actuator::HealthCheckOutput> {
1573 Box::pin(async { crate::actuator::HealthCheckOutput::down() })
1574 }
1575 }
1576
1577 struct AlwaysUp;
1578 impl crate::actuator::HealthIndicator for AlwaysUp {
1579 fn check(&self) -> futures::future::BoxFuture<'_, crate::actuator::HealthCheckOutput> {
1580 Box::pin(async { crate::actuator::HealthCheckOutput::up() })
1581 }
1582 }
1583
1584 #[tokio::test]
1585 async fn ready_is_degraded_when_readiness_indicator_is_down() {
1586 let registry = crate::actuator::HealthIndicatorRegistry::new();
1587 registry
1588 .register(
1589 "svc",
1590 crate::actuator::IndicatorGroup::Readiness,
1591 std::sync::Arc::new(AlwaysDown),
1592 )
1593 .unwrap();
1594 let state = TestProbeStateWithIndicators::new(registry);
1595 let (status, Json(response)) = readiness_response(&state).await;
1596 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1597 assert_eq!(response.status, "degraded");
1598 }
1599
1600 #[tokio::test]
1601 async fn ready_is_ok_when_readiness_indicator_is_up() {
1602 let registry = crate::actuator::HealthIndicatorRegistry::new();
1603 registry
1604 .register(
1605 "svc",
1606 crate::actuator::IndicatorGroup::Readiness,
1607 std::sync::Arc::new(AlwaysUp),
1608 )
1609 .unwrap();
1610 let state = TestProbeStateWithIndicators::new(registry);
1611 let (status, Json(response)) = readiness_response(&state).await;
1612 assert_eq!(status, StatusCode::OK);
1613 assert_eq!(response.status, "ok");
1614 }
1615
1616 #[tokio::test]
1617 async fn ready_is_ok_when_health_only_indicator_is_down() {
1618 let registry = crate::actuator::HealthIndicatorRegistry::new();
1619 registry
1620 .register(
1621 "svc",
1622 crate::actuator::IndicatorGroup::HealthOnly,
1623 std::sync::Arc::new(AlwaysDown),
1624 )
1625 .unwrap();
1626 let state = TestProbeStateWithIndicators::new(registry);
1627 let (status, Json(response)) = readiness_response(&state).await;
1628 assert_eq!(status, StatusCode::OK);
1629 assert_eq!(response.status, "ok");
1630 }
1631
1632 #[cfg(feature = "db")]
1633 mod bounded_lag {
1634 use super::*;
1635 use std::time::Duration;
1636
1637 struct Unit;
1639
1640 impl deadpool::managed::Manager for Unit {
1641 type Type = ();
1642 type Error = std::convert::Infallible;
1643
1644 async fn create(&self) -> Result<(), Self::Error> {
1645 Ok(())
1646 }
1647
1648 async fn recycle(
1649 &self,
1650 _conn: &mut (),
1651 _metrics: &deadpool::managed::Metrics,
1652 ) -> deadpool::managed::RecycleResult<Self::Error> {
1653 Ok(())
1654 }
1655 }
1656
1657 fn pool() -> deadpool::managed::Pool<Unit> {
1658 deadpool::managed::Pool::builder(Unit)
1659 .max_size(1)
1660 .build()
1661 .unwrap()
1662 }
1663
1664 fn probes() -> ProbeState {
1665 let probes = ProbeState::default();
1666 probes.configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1667 probes.mark_replica_ready();
1668 probes.configure_replica_max_lag(Some(Duration::from_secs(1)));
1669 probes
1670 }
1671
1672 #[tokio::test(start_paused = true)]
1673 async fn a_hung_lag_query_is_bounded_and_its_connection_discarded() {
1674 let (pool, probes) = (pool(), probes());
1675 let conn = pool.get().await.unwrap();
1676 refresh_replica_lag_bounded(&probes, conn, Duration::from_secs(1), |_| {
1677 Box::pin(std::future::pending())
1678 })
1679 .await;
1680 assert!(!probes.replica_lag_ok(), "the lag is unknown");
1681 assert_eq!(pool.status().size, 0, "the connection is not reused");
1682 }
1683
1684 #[tokio::test(start_paused = true)]
1685 async fn a_cancelled_lag_query_discards_its_connection() {
1686 let (pool, probes) = (pool(), probes());
1687 let conn = pool.get().await.unwrap();
1688 let probe = refresh_replica_lag_bounded(&probes, conn, Duration::from_secs(10), |_| {
1689 Box::pin(std::future::pending())
1690 });
1691 assert!(
1693 tokio::time::timeout(Duration::from_millis(10), probe)
1694 .await
1695 .is_err()
1696 );
1697 assert_eq!(pool.status().size, 0, "the connection is not reused");
1698 }
1699
1700 #[tokio::test(start_paused = true)]
1701 async fn a_finished_lag_query_returns_its_connection() {
1702 let (pool, probes) = (pool(), probes());
1703 let conn = pool.get().await.unwrap();
1704 refresh_replica_lag_bounded(&probes, conn, Duration::from_secs(1), |_| {
1705 Box::pin(async { Ok(Duration::from_millis(5)) })
1706 })
1707 .await;
1708 assert!(probes.replica_lag_ok());
1709 assert_eq!(pool.status().size, 1);
1710 assert_eq!(pool.status().available, 1, "back in the pool");
1711 }
1712 }
1713
1714 #[cfg(feature = "db")]
1717 mod primary_db {
1718 use super::*;
1719 use std::sync::atomic::AtomicUsize;
1720 use std::time::Duration;
1721
1722 struct FakePing {
1724 calls: AtomicUsize,
1725 delay: Duration,
1726 result: Result<(), String>,
1727 }
1728
1729 impl FakePing {
1730 fn new(delay: Duration, result: Result<(), String>) -> Arc<Self> {
1731 Arc::new(Self {
1732 calls: AtomicUsize::new(0),
1733 delay,
1734 result,
1735 })
1736 }
1737
1738 fn calls(&self) -> usize {
1739 self.calls.load(Ordering::SeqCst)
1740 }
1741 }
1742
1743 impl crate::db_ping::DbPing for FakePing {
1744 fn ping(
1745 self: Arc<Self>,
1746 _pool: DbPool,
1747 _budget: Duration,
1748 ) -> futures::future::BoxFuture<'static, Result<(), String>> {
1749 Box::pin(async move {
1750 self.calls.fetch_add(1, Ordering::SeqCst);
1751 tokio::time::sleep(self.delay).await;
1752 self.result.clone()
1753 })
1754 }
1755 }
1756
1757 struct DbProbeState {
1758 probes: ProbeState,
1759 pool: DbPool,
1760 replica: Option<DbPool>,
1761 }
1762
1763 impl ProvideProbeState for DbProbeState {
1764 fn probes(&self) -> &ProbeState {
1765 &self.probes
1766 }
1767
1768 fn health_detailed(&self) -> bool {
1769 true
1770 }
1771
1772 fn profile(&self) -> &'static str {
1773 "test"
1774 }
1775
1776 fn uptime_display(&self) -> String {
1777 String::new()
1778 }
1779
1780 fn pool(&self) -> Option<&DbPool> {
1781 Some(&self.pool)
1782 }
1783
1784 fn replica_pool(&self) -> Option<&DbPool> {
1785 self.replica.as_ref()
1786 }
1787 }
1788
1789 fn idle_pool(url: &str) -> DbPool {
1791 #[cfg(feature = "sqlite")]
1792 let url = {
1793 let _ = url;
1794 "sqlite::memory:"
1795 };
1796 let config = crate::config::DatabaseConfig {
1797 url: Some(url.to_owned()),
1798 pool_size: 1,
1799 connect_timeout_secs: 30,
1801 ..Default::default()
1802 };
1803 crate::db::create_pool(&config)
1804 .expect("pool config is valid")
1805 .expect("pool url is set")
1806 }
1807
1808 fn state_with(ping: Arc<dyn crate::db_ping::DbPing>) -> DbProbeState {
1809 let probes = ProbeState::with_primary_ping(ping);
1810 probes.mark_startup_complete();
1811 DbProbeState {
1812 probes,
1813 pool: idle_pool("postgres://autumn@127.0.0.1:1/unused"),
1814 replica: None,
1815 }
1816 }
1817
1818 #[tokio::test(start_paused = true)]
1819 async fn ready_is_ok_when_primary_answers() {
1820 let ping = FakePing::new(Duration::ZERO, Ok(()));
1821 let state = state_with(ping.clone());
1822
1823 let (status, Json(body)) = readiness_response(&state).await;
1824
1825 assert_eq!(status, StatusCode::OK);
1826 assert_eq!(body.status, "ok");
1827 assert_eq!(ping.calls(), 1);
1828 }
1829
1830 #[tokio::test(start_paused = true)]
1831 async fn ready_is_degraded_when_primary_ping_fails() {
1832 let ping = FakePing::new(Duration::ZERO, Err("connection refused".to_owned()));
1833 let state = state_with(ping);
1834
1835 let (status, Json(body)) = readiness_response(&state).await;
1836
1837 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1838 let database = body.database.expect("detailed body reports the database");
1839 assert_eq!(database.status, "down");
1840 assert!(
1841 database
1842 .error
1843 .as_deref()
1844 .is_some_and(|e| e.contains("connection refused")),
1845 "error: {:?}",
1846 database.error
1847 );
1848 }
1849
1850 #[tokio::test(start_paused = true)]
1851 async fn ready_is_degraded_within_timeout_when_primary_hangs() {
1852 let ping = FakePing::new(Duration::from_secs(3_600), Ok(()));
1853 let state = state_with(ping);
1854 state
1855 .probes
1856 .configure_db_check(Duration::from_secs(1), Duration::from_millis(250));
1857
1858 let started = tokio::time::Instant::now();
1859 let (status, Json(body)) = readiness_response(&state).await;
1860
1861 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1862 assert_eq!(started.elapsed(), Duration::from_millis(250));
1863 let error = body.database.and_then(|d| d.error).unwrap_or_default();
1864 assert!(error.contains("timed out"), "error: {error}");
1865 }
1866
1867 #[tokio::test(start_paused = true)]
1868 async fn concurrent_probes_ping_primary_at_most_once_per_window() {
1869 let ping = FakePing::new(Duration::from_millis(20), Ok(()));
1870 let state = Arc::new(state_with(ping.clone()));
1871 state
1872 .probes
1873 .configure_db_check(Duration::from_secs(1), Duration::from_secs(2));
1874
1875 let burst = || {
1876 let tasks: Vec<_> = (0..64)
1877 .map(|_| {
1878 let state = Arc::clone(&state);
1879 tokio::spawn(async move { readiness_response(&*state).await.0 })
1880 })
1881 .collect();
1882 async move {
1883 for task in tasks {
1884 assert_eq!(task.await.unwrap(), StatusCode::OK);
1885 }
1886 }
1887 };
1888
1889 burst().await;
1890 assert_eq!(ping.calls(), 1, "one ping for the first window");
1891
1892 tokio::time::advance(Duration::from_millis(500)).await;
1893 burst().await;
1894 assert_eq!(ping.calls(), 1, "still inside the first window");
1895
1896 tokio::time::advance(Duration::from_millis(600)).await;
1897 burst().await;
1898 assert_eq!(ping.calls(), 2, "one ping for the second window");
1899 }
1900
1901 #[tokio::test(start_paused = true)]
1902 async fn zero_ttl_pings_primary_on_every_probe() {
1903 let ping = FakePing::new(Duration::ZERO, Ok(()));
1904 let state = state_with(ping.clone());
1905 state
1906 .probes
1907 .configure_db_check(Duration::ZERO, Duration::from_secs(2));
1908
1909 for _ in 0..3 {
1910 let _ = readiness_response(&state).await;
1911 }
1912
1913 assert_eq!(ping.calls(), 3);
1914 }
1915
1916 #[tokio::test(start_paused = true)]
1917 async fn primary_and_replica_pings_run_at_the_same_time() {
1918 let primary = FakePing::new(Duration::from_millis(400), Ok(()));
1919 let replica = FakePing::new(Duration::from_millis(400), Ok(()));
1920 let probes = ProbeState::with_db_pings(primary.clone(), replica.clone());
1921 probes.mark_startup_complete();
1922 probes.configure_replica_dependency(crate::config::ReplicaFallback::FailReadiness);
1923 let state = DbProbeState {
1924 probes,
1925 pool: idle_pool("postgres://autumn@127.0.0.1:1/unused"),
1926 replica: Some(idle_pool("postgres://autumn@127.0.0.1:1/replica")),
1927 };
1928
1929 let started = tokio::time::Instant::now();
1930 let (status, _) = readiness_response(&state).await;
1931
1932 assert_eq!(status, StatusCode::OK);
1933 assert_eq!(primary.calls(), 1);
1934 assert_eq!(replica.calls(), 1);
1935 assert_eq!(started.elapsed(), Duration::from_millis(400));
1937 }
1938
1939 #[tokio::test(start_paused = true)]
1940 async fn ready_does_not_ping_primary_while_draining() {
1941 let ping = FakePing::new(Duration::ZERO, Ok(()));
1942 let state = state_with(ping.clone());
1943 state.probes.begin_draining();
1944
1945 let (status, _) = readiness_response(&state).await;
1946
1947 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
1948 assert_eq!(ping.calls(), 0);
1949 }
1950
1951 #[tokio::test(start_paused = true)]
1952 async fn db_readiness_off_does_not_ping_or_wait_for_the_primary() {
1953 let ping = FakePing::new(Duration::from_secs(3_600), Ok(()));
1956 let state = state_with(ping.clone());
1957 state.probes.set_db_readiness(false);
1958
1959 let started = tokio::time::Instant::now();
1960 let (status, Json(body)) = readiness_response(&state).await;
1961
1962 assert_eq!(status, StatusCode::OK);
1963 assert_eq!(started.elapsed(), Duration::ZERO);
1964 assert!(body.database.is_none());
1965 assert_eq!(ping.calls(), 0);
1966 }
1967
1968 #[tokio::test(start_paused = true)]
1969 async fn clones_share_one_ping_cache() {
1970 let ping = FakePing::new(Duration::ZERO, Ok(()));
1971 let state = state_with(ping.clone());
1972 let clone = DbProbeState {
1973 probes: state.probes.clone(),
1974 pool: state.pool.clone(),
1975 replica: None,
1976 };
1977
1978 let _ = readiness_response(&state).await;
1979 let _ = readiness_response(&clone).await;
1980
1981 assert_eq!(ping.calls(), 1);
1982 }
1983
1984 #[cfg(feature = "sqlite")]
1985 #[tokio::test]
1986 async fn sqlite_dedicated_ping_is_up_and_does_not_use_the_pool() {
1987 let probes = ProbeState::ready_for_test();
1988 probes.configure_db_check(Duration::ZERO, Duration::from_secs(5));
1989 let state = DbProbeState {
1990 probes,
1991 pool: idle_pool("sqlite::memory:"),
1992 replica: None,
1993 };
1994
1995 let (status, _) = readiness_response(&state).await;
1996
1997 assert_eq!(status, StatusCode::OK);
1998 assert_eq!(state.pool.status().size, 0, "the ping is not in the pool");
1999 }
2000
2001 #[cfg(not(feature = "sqlite"))]
2004 async fn black_hole() -> (std::net::SocketAddr, tokio::task::JoinHandle<()>) {
2005 let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
2006 .await
2007 .expect("bind black hole");
2008 let addr = listener.local_addr().expect("local addr");
2009 let task = tokio::spawn(async move {
2010 while let Ok((socket, _)) = listener.accept().await {
2011 tokio::spawn(async move {
2013 let _socket = socket;
2014 std::future::pending::<()>().await;
2015 });
2016 }
2017 });
2018 (addr, task)
2019 }
2020
2021 #[cfg(not(feature = "sqlite"))]
2022 #[tokio::test]
2023 async fn ready_is_ok_when_pool_is_saturated_and_primary_answers() {
2024 let (addr, server) = black_hole().await;
2028 let ping = FakePing::new(Duration::ZERO, Ok(()));
2029 let probes = ProbeState::with_primary_ping(ping);
2030 probes.mark_startup_complete();
2031 let state = DbProbeState {
2032 probes,
2033 pool: idle_pool(&format!("postgres://autumn@{addr}/autumn")),
2034 replica: None,
2035 };
2036 let holders: Vec<_> = (0..2)
2037 .map(|_| {
2038 let pool = state.pool.clone();
2039 tokio::spawn(async move { drop(pool.get().await) })
2040 })
2041 .collect();
2042 tokio::time::timeout(Duration::from_secs(5), async {
2043 while state.pool.status().waiting == 0 {
2044 tokio::time::sleep(Duration::from_millis(10)).await;
2045 }
2046 })
2047 .await
2048 .expect("a request waits for a connection");
2049 assert_eq!(state.pool.status().available, 0);
2050
2051 let (status, _) = readiness_response(&state).await;
2052
2053 assert_eq!(status, StatusCode::OK);
2054 for holder in holders {
2055 holder.abort();
2056 }
2057 server.abort();
2058 }
2059
2060 #[cfg(not(feature = "sqlite"))]
2061 #[tokio::test]
2062 async fn dedicated_ping_is_down_within_timeout_when_primary_hangs() {
2063 let (addr, server) = black_hole().await;
2064 let probes = ProbeState::ready_for_test();
2065 probes.configure_db_check(Duration::ZERO, Duration::from_millis(300));
2066 let state = DbProbeState {
2067 probes,
2068 pool: idle_pool(&format!("postgres://autumn@{addr}/autumn")),
2069 replica: None,
2070 };
2071
2072 let started = std::time::Instant::now();
2073 let (status, _) = readiness_response(&state).await;
2074 let elapsed = started.elapsed();
2075
2076 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
2077 assert!(
2078 elapsed < Duration::from_millis(300) + Duration::from_secs(3),
2079 "took {elapsed:?}"
2080 );
2081 server.abort();
2082 }
2083
2084 #[cfg(not(feature = "sqlite"))]
2085 #[tokio::test]
2086 async fn dedicated_ping_is_down_when_primary_refuses() {
2087 let addr = {
2088 let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
2089 listener.local_addr().expect("local addr")
2090 };
2091 let probes = ProbeState::ready_for_test();
2092 probes.configure_db_check(Duration::ZERO, Duration::from_secs(2));
2093 let state = DbProbeState {
2094 probes,
2095 pool: idle_pool(&format!("postgres://autumn@{addr}/autumn")),
2096 replica: None,
2097 };
2098
2099 let (status, _) = readiness_response(&state).await;
2100
2101 assert_eq!(status, StatusCode::SERVICE_UNAVAILABLE);
2102 }
2103 }
2104}