diff --git a/crates/aisix-server/src/heartbeat.rs b/crates/aisix-server/src/heartbeat.rs new file mode 100644 index 00000000..5a8e9c9d --- /dev/null +++ b/crates/aisix-server/src/heartbeat.rs @@ -0,0 +1,233 @@ +//! Periodic `POST /dp/heartbeat` so cp-api knows the DP is alive. +//! +//! Protocol: prd-09 §9.3.5 and §9.10. The DP sends `{ dp_id, +//! uptime_seconds, version }` over HTTPS with `Authorization: +//! Bearer ` — Phase 1 auth (Phase 2 upgrades to mTLS client +//! cert once cp-api terminates mTLS). +//! +//! Shape: +//! - spawned once from `main` after registration/cert load is +//! complete +//! - ticks at the interval returned by the register response +//! (default 15s) +//! - individual heartbeats fail fast on network errors; the +//! ticker keeps running so a transient outage doesn't stop the +//! DP from being seen when the CP comes back +//! - cancelled via the shared `watch::Receiver` so graceful +//! shutdown doesn't leave an in-flight request dangling + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use anyhow::{anyhow, Context}; +use serde::Serialize; +use tokio::sync::watch; + +/// Configuration captured at register time. All three fields come +/// from the register response (see prd-09 §9.3.5) and are immutable +/// for the life of the heartbeat worker. +#[derive(Debug, Clone)] +pub struct HeartbeatConfig { + pub url: String, + pub dp_id: String, + pub interval: Duration, +} + +impl HeartbeatConfig { + /// Clamp the server-suggested interval into a safe band. Defence + /// against a buggy CP config that returns 0 or a week. + pub fn sanitised(url: String, dp_id: String, interval: Duration) -> Self { + const MIN: Duration = Duration::from_secs(5); + const MAX: Duration = Duration::from_secs(300); + let interval = interval.clamp(MIN, MAX); + Self { + url, + dp_id, + interval, + } + } +} + +/// Spawn the heartbeat worker. Returns the JoinHandle so `main` can +/// await it at shutdown. Errors during individual heartbeats are +/// logged, not propagated — a heartbeat that can't reach the CP is +/// noisy, not fatal. +pub fn spawn( + cfg: HeartbeatConfig, + mut cancel: watch::Receiver, +) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + run(cfg, &mut cancel).await; + }) +} + +async fn run(cfg: HeartbeatConfig, cancel: &mut watch::Receiver) { + let client = match build_client() { + Ok(c) => Arc::new(c), + Err(e) => { + tracing::error!(error = %e, "heartbeat: build reqwest client failed; disabled"); + return; + } + }; + let started = Instant::now(); + let mut ticker = tokio::time::interval(cfg.interval); + // Skip the catch-up fire — we want the first beat to happen + // immediately at spawn but subsequent ones to follow the tick + // schedule without bursting if we fall behind. + ticker.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + + tracing::info!( + url = %cfg.url, + dp_id = %cfg.dp_id, + interval_secs = cfg.interval.as_secs(), + "heartbeat started", + ); + + loop { + tokio::select! { + _ = ticker.tick() => { + let uptime = started.elapsed().as_secs() as i64; + match send(&client, &cfg, uptime).await { + Ok(()) => tracing::debug!("heartbeat ok"), + Err(e) => tracing::warn!(error = %e, "heartbeat failed"), + } + } + _ = cancel.changed() => { + if *cancel.borrow() { + tracing::info!("heartbeat shutting down"); + return; + } + } + } + } +} + +#[derive(Debug, Serialize)] +struct HeartbeatBody<'a> { + dp_id: &'a str, + uptime_seconds: i64, + version: &'a str, +} + +async fn send(client: &reqwest::Client, cfg: &HeartbeatConfig, uptime: i64) -> anyhow::Result<()> { + let resp = client + .post(&cfg.url) + .bearer_auth(&cfg.dp_id) + .json(&HeartbeatBody { + dp_id: &cfg.dp_id, + uptime_seconds: uptime, + version: env!("CARGO_PKG_VERSION"), + }) + .send() + .await + .with_context(|| format!("POST {}", cfg.url))?; + + let status = resp.status(); + if !status.is_success() { + let body = resp.text().await.unwrap_or_default(); + return Err(anyhow!( + "heartbeat {} returned {} — {}", + cfg.url, + status, + body.trim().chars().take(200).collect::() + )); + } + Ok(()) +} + +fn build_client() -> anyhow::Result { + reqwest::Client::builder() + .timeout(Duration::from_secs(10)) + .user_agent(format!("aisix-dp/{}", env!("CARGO_PKG_VERSION"))) + .build() + .context("build reqwest client") +} + +#[cfg(test)] +mod tests { + use super::*; + use wiremock::matchers::{body_string_contains, header, method, path}; + use wiremock::{Mock, MockServer, ResponseTemplate}; + + fn cfg(url: String) -> HeartbeatConfig { + HeartbeatConfig::sanitised(url, "dp_test_node_42".into(), Duration::from_millis(50)) + } + + #[tokio::test] + async fn send_posts_dp_id_and_bearer() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/dp/heartbeat")) + .and(header("Authorization", "Bearer dp_test_node_42")) + .and(body_string_contains("\"dp_id\":\"dp_test_node_42\"")) + .and(body_string_contains("\"uptime_seconds\":")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({ + "ok": true + }))) + .mount(&server) + .await; + + let c = build_client().unwrap(); + send(&c, &cfg(format!("{}/dp/heartbeat", server.uri())), 7) + .await + .unwrap(); + } + + #[tokio::test] + async fn send_propagates_non_success_body() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/dp/heartbeat")) + .respond_with(ResponseTemplate::new(404).set_body_json(serde_json::json!({ + "error": {"code": "DP_NOT_FOUND", "message": "no registered DP matches this id"} + }))) + .mount(&server) + .await; + + let c = build_client().unwrap(); + let err = send(&c, &cfg(format!("{}/dp/heartbeat", server.uri())), 7) + .await + .unwrap_err(); + let s = format!("{err:#}"); + assert!(s.contains("404"), "expected status in error: {s}"); + assert!(s.contains("DP_NOT_FOUND"), "expected body in error: {s}"); + } + + #[tokio::test] + async fn run_stops_on_cancel() { + // Start a server that 200s fast enough that the first tick + // completes, then cancel and make sure the task exits. + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/dp/heartbeat")) + .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({"ok": true}))) + .mount(&server) + .await; + + let (tx, rx) = watch::channel(false); + let handle = spawn(cfg(format!("{}/dp/heartbeat", server.uri())), rx); + + tokio::time::sleep(Duration::from_millis(150)).await; + tx.send(true).unwrap(); + + // Runs to completion within a small grace window. + tokio::time::timeout(Duration::from_secs(2), handle) + .await + .expect("heartbeat did not stop after cancel") + .unwrap(); + } + + #[test] + fn sanitised_interval_clamps_extremes() { + let a = + HeartbeatConfig::sanitised("http://x".into(), "id".into(), Duration::from_millis(10)); + assert_eq!(a.interval, Duration::from_secs(5)); + + let b = + HeartbeatConfig::sanitised("http://x".into(), "id".into(), Duration::from_secs(86_400)); + assert_eq!(b.interval, Duration::from_secs(300)); + + let c = HeartbeatConfig::sanitised("http://x".into(), "id".into(), Duration::from_secs(30)); + assert_eq!(c.interval, Duration::from_secs(30)); + } +} diff --git a/crates/aisix-server/src/main.rs b/crates/aisix-server/src/main.rs index 46a6e9a8..0ff254cd 100644 --- a/crates/aisix-server/src/main.rs +++ b/crates/aisix-server/src/main.rs @@ -15,6 +15,7 @@ use std::path::PathBuf; use std::sync::Arc; +mod heartbeat; mod register; use aisix_admin::{AdminState, ConfigStore, EtcdConfigStore}; @@ -61,38 +62,57 @@ async fn main() -> anyhow::Result<()> { /// Factored out of `main` so the integration tests can drive the full /// startup with a real config struct and still use `#[tokio::test]`. async fn run(mut cfg: Config) -> anyhow::Result<()> { - // If this is a managed tenant and the mTLS bundle isn't on disk - // yet, perform the one-shot `POST /dp/register` exchange against - // the aisix.cloud control plane. The response fills in - // `cfg.etcd.endpoints` + `cfg.etcd.tls` so the regular etcd - // connect path below is oblivious to whether certs came from an - // out-of-band install or just-now registration. - if cfg.managed.is_managed() - && !register::bundle_exists(&cfg.managed.mtls_dir) - && cfg.managed.registration_enabled() - { - tracing::info!("managed mode: registering with aisix.cloud CP"); - let r = register::register_and_persist(&cfg.managed) - .await - .map_err(|e| anyhow::anyhow!("DP registration failed: {e:#}"))?; - tracing::info!( - dp_id = %r.dp_id, - gateway_id = %r.gateway_id, - etcd = %r.etcd_endpoint, - "registered with control plane", - ); - // Override the static config with what the CP handed back. - // Endpoints get the https:// scheme re-attached (the CP sends - // a bare host:port, but tonic-based etcd-client expects a - // full URL for TLS endpoints). - cfg.etcd.endpoints = vec![format!("https://{}", r.etcd_endpoint)]; - cfg.etcd.tls = Some(EtcdTlsConfig { - ca_cert_file: r.ca_cert_path.to_string_lossy().into_owned(), - client_cert_file: r.client_cert_path.to_string_lossy().into_owned(), - client_key_file: r.client_key_path.to_string_lossy().into_owned(), - domain_name: None, // derive from endpoint host - }); - } + // Managed-mode bootstrap. If we have to register (first boot), + // we also capture the heartbeat config the CP sent back so the + // worker can be spawned a few lines below. If the bundle is + // already on disk (subsequent boot), we synthesise the same + // values from config + dp_id_file. + let heartbeat_cfg: Option = if cfg.managed.is_managed() { + if !register::bundle_exists(&cfg.managed.mtls_dir) && cfg.managed.registration_enabled() { + tracing::info!("managed mode: registering with aisix.cloud CP"); + let r = register::register_and_persist(&cfg.managed) + .await + .map_err(|e| anyhow::anyhow!("DP registration failed: {e:#}"))?; + tracing::info!( + dp_id = %r.dp_id, + gateway_id = %r.gateway_id, + etcd = %r.etcd_endpoint, + "registered with control plane", + ); + // Override the static config with what the CP handed back. + // Endpoints get the https:// scheme re-attached (the CP + // sends a bare host:port, but tonic-based etcd-client + // expects a full URL for TLS endpoints). + cfg.etcd.endpoints = vec![format!("https://{}", r.etcd_endpoint)]; + cfg.etcd.tls = Some(EtcdTlsConfig { + ca_cert_file: r.ca_cert_path.to_string_lossy().into_owned(), + client_cert_file: r.client_cert_path.to_string_lossy().into_owned(), + client_key_file: r.client_key_path.to_string_lossy().into_owned(), + domain_name: None, // derive from endpoint host + }); + Some(heartbeat::HeartbeatConfig::sanitised( + r.heartbeat_url, + r.dp_id, + r.heartbeat_interval, + )) + } else if register::bundle_exists(&cfg.managed.mtls_dir) { + // Bundle persisted from a previous boot; load the dp_id + // and synthesise heartbeat config from the configured + // cp_base_url. Registration doesn't re-run. + match load_heartbeat_config_from_disk(&cfg.managed) { + Ok(h) => Some(h), + Err(e) => { + tracing::warn!(error = %e, + "managed mode: heartbeat worker disabled (dp_id unreadable)"); + None + } + } + } else { + None + } + } else { + None + }; // Steps 4-6: etcd + supervisor. let connect_options = build_etcd_connect_options(&cfg.etcd)?; @@ -124,6 +144,10 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { let (cancel_tx, cancel_rx) = watch::channel(false); let watch_task = tokio::spawn(supervisor.clone().run(cancel_rx.clone())); + // Spawn heartbeat worker if we have a config for it. The + // JoinHandle is awaited after graceful shutdown below so the + // final in-flight beat drains cleanly. + let heartbeat_task = heartbeat_cfg.map(|h| heartbeat::spawn(h, cancel_rx.clone())); // Steps 7-8: build Hub, shared components, then routers. let hub = Arc::new(build_hub()); @@ -243,6 +267,9 @@ async fn run(mut cfg: Config) -> anyhow::Result<()> { let _ = cancel_tx.send(true); let _ = signal_task.await; let _ = watch_task.await; + if let Some(task) = heartbeat_task { + let _ = task.await; + } tracing::info!("aisix shut down cleanly"); Ok(()) } @@ -323,6 +350,39 @@ fn default_domain_from_endpoint(endpoint: &str) -> anyhow::Result { Ok(host.to_string()) } +/// Synthesise a HeartbeatConfig when the mTLS bundle is already on +/// disk from a previous boot. Reads `managed.dp_id_file` and +/// combines with `managed.cp_base_url` — the register response is +/// not available on this code path. +/// +/// Returns an error (not None) when the user has configured managed +/// mode AND the bundle exists BUT the dp_id is unreadable — that's +/// an inconsistent on-disk state an operator should know about. +fn load_heartbeat_config_from_disk( + managed: &aisix_core::ManagedConfig, +) -> anyhow::Result { + let base = managed + .cp_base_url + .as_deref() + .filter(|s| !s.is_empty()) + .ok_or_else(|| { + anyhow::anyhow!("managed.cp_base_url must be set for heartbeat on subsequent boots") + })?; + let dp_id = std::fs::read_to_string(&managed.dp_id_file) + .map_err(|e| anyhow::anyhow!("read dp_id from {}: {e}", managed.dp_id_file))? + .trim() + .to_string(); + if dp_id.is_empty() { + anyhow::bail!("dp_id file {} is empty", managed.dp_id_file); + } + let url = format!("{}/dp/heartbeat", base.trim_end_matches('/')); + Ok(heartbeat::HeartbeatConfig::sanitised( + url, + dp_id, + std::time::Duration::from_secs(15), + )) +} + /// Register all four provider bridges on a fresh Hub. The Hub is /// created once at startup; future dynamic reload lands behind the /// same `register()` call.