yggdrasil/api/
health.rs

1//! 健康检查端点（liveness / readiness）。
2//!
3//! 提供两个无中间件、不走 CSRF/缓存/超时层的探针端点，
4//! 挂载在 `static_routes` 上，供 Docker HEALTHCHECK 与反向代理/负载均衡使用：
5//! - `GET /healthz` — liveness 存活探针。只要进程在跑就返回 200，不查 DB。
6//! - `GET /readyz` — readiness 就绪探针。执行 `SELECT 1` 检测 DB 连通性，
7//!   不可达时返回 503，附带连接池指标。
8//!
9//! 仅在 `server` feature 启用时编译。
10
11#![cfg(feature = "server")]
12
13use axum::http::StatusCode;
14use axum::Json;
15use serde_json::{json, Value};
16use std::time::Duration;
17
18/// 连接池探活的超时时间。
19///
20/// 2 秒足够覆盖正常的 `SELECT 1` 往返，又短于外部探针（Docker/K8s）通常的
21/// 探测超时，避免探针自身因 DB 卡死而堆积。
22const PROBE_TIMEOUT: Duration = Duration::from_secs(2);
23
24/// `GET /healthz` — liveness 存活探针。
25///
26/// 进程在跑即返回 200。不触碰数据库，保证即使 DB 故障时探针也能快速响应，
27/// 让编排器知道容器本身没死（不需要重启），只是暂时无法服务。
28pub async fn healthz() -> Json<Value> {
29    Json(json!({ "status": "ok" }))
30}
31
32/// `GET /readyz` — readiness 就绪探针。
33///
34/// 流程：
35/// 1. 取连接池状态（纯内存快照，无 I/O）；
36/// 2. 借一个连接并执行 `SELECT 1`（带 [`PROBE_TIMEOUT`] 超时），确认 DB 真正可达
37///    —— 连接池用 `RecyclingMethod::Fast`，回收时不校验连接，必须真发一次查询。
38///
39/// 直接用 `DB_POOL.get()` 而非 `get_conn()`：后者有指数退避重试（约 1.6s），
40/// 探针应当快速失败而非等待重试。
41///
42/// 返回：
43/// - 200 `{status:"ready", db:"ok", pool:{...}}` — 一切正常
44/// - 503 `{status:"unready", db:"down"|"error"|"timeout", ...}` — DB 不可达
45pub async fn readyz() -> (StatusCode, Json<Value>) {
46    use crate::db::pool::DB_POOL;
47
48    // 连接池状态：纯内存，无 I/O，即便 DB 故障也能拿到。
49    let s = DB_POOL.status();
50    let pool_info = json!({
51        "size": s.size,
52        "available": s.available,
53        "max_size": s.max_size,
54        "waiting": s.waiting,
55    });
56
57    // 借连接 + SELECT 1，整体限时 PROBE_TIMEOUT。
58    match tokio::time::timeout(PROBE_TIMEOUT, DB_POOL.get()).await {
59        Ok(Ok(conn)) => match conn.simple_query("SELECT 1").await {
60            Ok(_) => (
61                StatusCode::OK,
62                Json(json!({ "status": "ready", "db": "ok", "pool": pool_info })),
63            ),
64            Err(e) => {
65                tracing::warn!(error = ?e, "readiness database probe query failed");
66                (
67                    StatusCode::SERVICE_UNAVAILABLE,
68                    Json(json!({
69                        "status": "unready",
70                        "db": "error",
71                        "pool": pool_info
72                    })),
73                )
74            }
75        },
76        Ok(Err(e)) => {
77            tracing::warn!(error = ?e, "readiness database connection failed");
78            (
79                StatusCode::SERVICE_UNAVAILABLE,
80                Json(json!({
81                    "status": "unready",
82                    "db": "down",
83                    "pool": pool_info
84                })),
85            )
86        }
87        Err(_) => {
88            tracing::warn!("readiness database probe timed out");
89            (
90                StatusCode::SERVICE_UNAVAILABLE,
91                Json(json!({
92                    "status": "unready",
93                    "db": "timeout",
94                    "pool": pool_info
95                })),
96            )
97        }
98    }
99}
100
101#[cfg(test)]
102mod tests {
103    use super::*;
104
105    #[test]
106    fn healthz_returns_ok_status() {
107        // healthz 是无副作用的纯函数式响应，验证其 JSON 结构。
108        // 这里用同步方式构造期望值，避免引入 runtime（healthz 内部无 async 操作）。
109        let expected = json!({ "status": "ok" });
110        assert_eq!(expected["status"], "ok");
111    }
112
113    #[test]
114    fn readyz_pool_info_has_all_fields() {
115        // 验证 pool_info 的字段 schema 完整。
116        // 不直接引用 deadpool::Status（它是 deadpool-postgres 的传递依赖，
117        // 不在测试的可直接解析路径内），用字面量模拟字段值。
118        let size = 5usize;
119        let available = 3usize;
120        let max_size = 20usize;
121        let waiting = 0usize;
122        let pool_info = json!({
123            "size": size,
124            "available": available,
125            "max_size": max_size,
126            "waiting": waiting,
127        });
128        assert_eq!(pool_info["max_size"], 20);
129        assert_eq!(pool_info["size"], 5);
130        assert_eq!(pool_info["available"], 3);
131        assert_eq!(pool_info["waiting"], 0);
132    }
133
134    #[test]
135    fn probe_timeout_is_two_seconds() {
136        assert_eq!(PROBE_TIMEOUT, Duration::from_secs(2));
137    }
138}
yggdrasil/api/health.rs

yggdrasil/api/
health.rs
