main
use std::sync::Arc;
use std::sync::atomic::{AtomicU8, Ordering};
use std::time::Duration;
use actix_web::HttpResponse;
use sea_orm::{ConnectionTrait, DbBackend, Statement};
use api::http::{Router, State, get};
use cache::AppRedis;
use db::AppDatabase;
use crate::AppState;
const UP: u8 = 1;
const DOWN: u8 = 2;
/// 连续失败达到该次数才置 DOWN：单次采样抖动不应把 Pod 摘出负载均衡。
const DOWN_AFTER_FAILURES: u8 = 2;
/// 探针自身超时，须小于 k8s `timeoutSeconds`（3s），否则超时也会被计为失败。
const READY_PROBE_TIMEOUT: Duration = Duration::from_secs(2);
/// 后台采样得到的依赖健康状态；探针只读取这里的值，不发起探测。
/// 原子量初值为 0，表示尚未采样（同样按未就绪处理）。
#[derive(Clone, Default)]
pub struct HealthCache {
inner: Arc<HealthCacheInner>,
}
#[derive(Default)]
struct HealthCacheInner {
rocketmq: AtomicU8,
rocketmq_failures: AtomicU8,
s3: AtomicU8,
s3_failures: AtomicU8,
}
impl HealthCache {
pub fn mark_rocketmq(&self, up: bool) {
mark(&self.inner.rocketmq, &self.inner.rocketmq_failures, up);
}
pub fn mark_s3(&self, up: bool) {
mark(&self.inner.s3, &self.inner.s3_failures, up);
}
fn dependencies(&self) -> [(&'static str, bool); 2] {
[
(
"rocketmq",
self.inner.rocketmq.load(Ordering::Relaxed) == UP,
),
("s3", self.inner.s3.load(Ordering::Relaxed) == UP),
]
}
}
/// 成功即置 UP；失败累计到阈值才置 DOWN，避免单次抖动摘流量。
fn mark(state: &AtomicU8, failures: &AtomicU8, up: bool) {
if up {
failures.store(0, Ordering::Relaxed);
state.store(UP, Ordering::Relaxed);
return;
}
let count = failures.fetch_add(1, Ordering::Relaxed).saturating_add(1);
if count >= DOWN_AFTER_FAILURES {
state.store(DOWN, Ordering::Relaxed);
}
}
#[derive(Clone)]
struct HealthState {
database: Option<Arc<AppDatabase>>,
redis: AppRedis,
health: HealthCache,
}
async fn healthz() -> HttpResponse {
HttpResponse::Ok().body("ok")
}
async fn readyz(State(state): State<HealthState>) -> HttpResponse {
let db_ok = match state.database.as_ref() {
None => true,
Some(db) => actix_web::rt::time::timeout(
READY_PROBE_TIMEOUT,
db.write().query_one_raw(Statement::from_sql_and_values(
DbBackend::Postgres,
"SELECT 1",
[],
)),
)
.await
.is_ok_and(|result| result.is_ok()),
};
let redis_ok = actix_web::rt::time::timeout(READY_PROBE_TIMEOUT, state.redis.ping())
.await
.is_ok_and(|result| result.is_ok());
let mut failed = Vec::new();
if !db_ok {
failed.push("database");
}
if !redis_ok {
failed.push("redis");
}
for (dependency, up) in state.health.dependencies() {
if !up {
failed.push(dependency);
}
}
if failed.is_empty() {
HttpResponse::Ok().body("ready")
} else {
HttpResponse::ServiceUnavailable().body(format!("not ready: {}", failed.join(",")))
}
}
pub fn router_with(
database: Option<Arc<AppDatabase>>,
redis: AppRedis,
health: HealthCache,
) -> Router {
Router::<()>::new()
.route("/healthz", get(healthz))
.route("/readyz", get(readyz))
.with_state(HealthState {
database,
redis,
health,
})
}
pub fn router(state: AppState) -> Router {
router_with(
Some(state.database.clone()),
state.redis.clone(),
state.health.clone(),
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn down_needs_consecutive_failures_and_recovers_immediately() {
let state = AtomicU8::new(0);
let failures = AtomicU8::new(0);
// 尚未采样 + 首次失败：仍保持「未就绪」而不立刻判 DOWN。
mark(&state, &failures, false);
assert_eq!(state.load(Ordering::Relaxed), 0);
mark(&state, &failures, false);
assert_eq!(state.load(Ordering::Relaxed), DOWN);
// 一次成功立刻恢复，并清零失败计数。
mark(&state, &failures, true);
assert_eq!(state.load(Ordering::Relaxed), UP);
assert_eq!(failures.load(Ordering::Relaxed), 0);
// 恢复后再次单次失败仍然是 UP（阈值重新累计）。
mark(&state, &failures, false);
assert_eq!(state.load(Ordering::Relaxed), UP);
}
}