Skip to content
Merged
Show file tree
Hide file tree
Changes from 7 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
60 changes: 40 additions & 20 deletions crates/aisix-admin/src/lib.rs
Original file line number Diff line number Diff line change
@@ -1,7 +1,12 @@
//! aisix-admin — Admin API + Playground (:3001).
//!
//! Mounts the admin surface behind admin-key bearer auth:
//! - `GET /health`
//! Public admin-listener endpoints:
//! - `GET /livez`
//! - `GET /metrics`
//! - `GET /admin/openapi.json`
//! - `GET /admin/openapi-scalar`
//!
//! Admin-key protected routes:
//! - `GET|POST /admin/v1/models`
//! - `GET|PUT|DELETE /admin/v1/models/:id`
//! - `GET|POST /admin/v1/apikeys`
Expand Down Expand Up @@ -48,12 +53,11 @@ pub use state::AdminState;
pub use store::{ConfigStore, InMemoryStore, StoreError};

use axum::routing::{get, post};
use axum::{http::StatusCode, Json, Router};
use serde_json::json;
use axum::{http::StatusCode, response::Response, Router};

pub fn build_router(state: AdminState) -> Router {
Router::new()
.route("/health", get(health))
.route("/livez", get(livez))
.route("/metrics", get(metrics_handler))
// OpenAPI scalar UI is unauthenticated like /metrics — admin
// listener is private in production.
Expand Down Expand Up @@ -139,18 +143,11 @@ pub fn build_router(state: AdminState) -> Router {
.with_state(state)
}

async fn health(
async fn livez(
axum::extract::State(state): axum::extract::State<AdminState>,
) -> (StatusCode, Json<serde_json::Value>) {
let snap = state.snapshot.load();
(
StatusCode::OK,
Json(json!({
"status": "ok",
"models": snap.models.len(),
"apikeys": snap.apikeys.len(),
})),
)
axum::extract::Query(params): axum::extract::Query<std::collections::HashMap<String, String>>,
) -> Response {
aisix_proxy::health::livez_response(&state.livez_state, params.contains_key("verbose"))
}

/// Prometheus `/metrics` endpoint. Unauthenticated by design — the admin
Expand Down Expand Up @@ -328,16 +325,39 @@ mod tests {
}

#[tokio::test]
async fn health_reports_snapshot_counts() {
async fn livez_reports_plain_ok_by_default() {
let app = build_router(build_state());
let req = Request::builder()
.uri("/health")
.uri("/livez")
.body(Body::empty())
.unwrap();
let resp = run(app, req).await;
assert_eq!(resp.status(), StatusCode::OK);
let v = body_json(resp).await;
assert_eq!(v["status"], "ok");
let bytes = to_bytes(resp.into_body(), 65536).await.unwrap();
assert_eq!(std::str::from_utf8(&bytes).unwrap(), "ok");
}

#[tokio::test]
async fn livez_rejects_non_get_requests() {
let app = build_router(build_state());
let req = Request::builder()
.method("POST")
.uri("/livez")
.body(Body::empty())
.unwrap();
let resp = run(app, req).await;
assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED);
}

#[tokio::test]
async fn health_route_is_not_found() {
let app = build_router(build_state());
let req = Request::builder()
.uri("/health")
.body(Body::empty())
.unwrap();
let resp = run(app, req).await;
assert_eq!(resp.status(), StatusCode::NOT_FOUND);
}

#[tokio::test]
Expand Down
35 changes: 29 additions & 6 deletions crates/aisix-admin/src/openapi.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,11 +24,23 @@ const OPENAPI_JSON: &str = r##"{
"description": "Admin surface for the standalone aisix gateway. All `/admin/v1/*` routes require Bearer admin-key auth (configured via `admin.admin_keys`). Errors use {\"error_msg\": \"...\"}.\n\nIn managed mode (aisix.cloud tenant) the admin listener is not bound — the dashboard owns CRUD via the AISIX-Cloud control plane."
},
"paths": {
"/health": {
"/livez": {
"get": {
"summary": "process health + snapshot counts",
"summary": "minimal public liveness probe",
"security": [],
"responses": {"200": {"description": "OK"}}
"responses": {
"200": {
"description": "OK",
"content": {
"text/plain": {
"schema": {
"type": "string",
"enum": ["ok"]
}
}
}
Comment on lines +40 to +59
Comment on lines +40 to +59
}
}
}
},
"/metrics": {
Expand Down Expand Up @@ -399,7 +411,7 @@ mod tests {
serde_json::from_str(OPENAPI_JSON).expect("OPENAPI_JSON must parse");
// Every route mounted in build_router should be documented.
for path in [
"/health",
"/livez",
"/metrics",
"/admin/openapi.json",
"/admin/openapi-scalar",
Expand Down Expand Up @@ -448,15 +460,15 @@ mod tests {

#[tokio::test]
async fn openapi_unauthenticated_routes_carry_empty_security() {
// /health, /metrics, and the openapi self-references are public
// /livez, /metrics, and the openapi self-references are public
// (mirrors the `unauthenticated like /metrics` design note in
// build_router). The spec must mark them with security: [] so
// Scalar's "Try it" doesn't prompt for an admin key on those
// routes.
let parsed: serde_json::Value =
serde_json::from_str(OPENAPI_JSON).expect("OPENAPI_JSON must parse");
for path in [
"/health",
"/livez",
"/metrics",
"/admin/openapi.json",
"/admin/openapi-scalar",
Expand All @@ -469,6 +481,17 @@ mod tests {
}
}

#[tokio::test]
async fn openapi_livez_documents_plain_ok() {
let parsed: serde_json::Value =
serde_json::from_str(OPENAPI_JSON).expect("OPENAPI_JSON must parse");
let schema = &parsed["paths"]["/livez"]["get"]["responses"]["200"]["content"]["text/plain"]
["schema"];

assert_eq!(schema["type"], "string");
assert_eq!(schema["enum"], serde_json::json!(["ok"]));
}

#[tokio::test]
async fn scalar_html_loads_the_spec_url() {
let html = openapi_scalar().await;
Expand Down
14 changes: 12 additions & 2 deletions crates/aisix-admin/src/state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
//! Holds:
//! - the bootstrap-config-provided `admin_keys` (auth)
//! - the `ConfigStore` trait object (CRUD backend)
//! - a `SnapshotHandle` for the /health endpoint (snapshot counts)
//! - a `SnapshotHandle` for authenticated operator-health handlers
//! - an optional `Metrics` handle — when present, `/metrics` renders
//! the same Prometheus exposition that the proxy's middleware writes to
//!
Expand All @@ -15,7 +15,7 @@ use aisix_core::snapshot::SnapshotHandle;
use aisix_core::{AdminConfig, AisixSnapshot};
use aisix_etcd::WatchStatus;
use aisix_obs::Metrics;
use aisix_proxy::HealthTracker;
use aisix_proxy::{HealthTracker, LivezState};
use axum::Router;
use std::sync::Arc;

Expand All @@ -36,6 +36,8 @@ pub struct AdminState {
/// stream that would otherwise let the gateway serve a stale
/// snapshot indefinitely. See issue #114.
pub watch_status: Option<WatchStatus>,
/// Shared liveness state for the public `GET /livez` endpoint.
pub livez_state: Arc<LivezState>,
/// Proxy router shared for the `/playground/chat/completions` endpoint.
/// The playground handler calls `router.oneshot(req)` so the request
/// goes through the full proxy middleware stack (auth, rate-limit, bridge)
Expand All @@ -56,6 +58,7 @@ impl AdminState {
metrics: None,
health_tracker: None,
watch_status: None,
livez_state: Arc::new(LivezState::new()),
proxy_router: None,
}
}
Expand Down Expand Up @@ -83,6 +86,13 @@ impl AdminState {
self
}

/// Share the public liveness state with the proxy so both listeners
/// report the same shutdown signal.
pub fn with_livez_state(mut self, livez_state: Arc<LivezState>) -> Self {
self.livez_state = livez_state;
self
}

/// Wire the proxy router so the playground endpoint can forward
/// requests to it in-process via `tower::ServiceExt::oneshot`.
pub fn with_proxy_router(mut self, router: Router) -> Self {
Expand Down
10 changes: 5 additions & 5 deletions crates/aisix-etcd/src/supervisor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -297,7 +297,7 @@ impl<P: ConfigProvider> Supervisor<P> {
/// stamp the cache with the etcd `load_all` revision even when the
/// resulting entry set is empty (so the file still reflects when
/// the DP last successfully reached the CP). Also stamps
/// `WatchStatus.last_apply_at` so `/health` reflects the
/// `WatchStatus.last_apply_at` so `/admin/v1/health` reflects the
/// successful round-trip with etcd even on an empty config.
fn set_revision_floor(&self, revision: i64) {
let mut rev = self.revision.lock().unwrap();
Expand Down Expand Up @@ -373,7 +373,7 @@ impl<P: ConfigProvider> Supervisor<P> {
*rev = entry.revision;
}
}
// /health reads this — record the apply so `last_apply_age`
// /admin/v1/health reads this — record the apply so `last_apply_age`
// resets on every event we successfully process.
self.status.record_apply(entry.revision);
self.flush_cache();
Expand Down Expand Up @@ -447,7 +447,7 @@ impl<P: ConfigProvider> Supervisor<P> {
new
});
self.state.lock().unwrap().remove(key_str);
// Stamp /health freshness on a successful delete. We
// Stamp /admin/v1/health freshness on a successful delete. We
// don't have a per-event revision on the wire delete
// (the etcd watch revision is held at the cycle level);
// call record_apply with the current revision so age
Expand Down Expand Up @@ -481,7 +481,7 @@ impl<P: ConfigProvider> Supervisor<P> {
*rev = rev_val;
}
}
// /health: stamp freshness on every resync, even when the
// /admin/v1/health: stamp freshness on every resync, even when the
// resulting entry set is empty (record_apply with the current
// revision floor so the operator sees recent activity).
let cur_rev = *self.revision.lock().unwrap();
Expand Down Expand Up @@ -1079,7 +1079,7 @@ mod tests {
async fn watch_status_age_grows_when_no_events_arrive() {
// Pin the freshness signal: after an apply, the age is small;
// wait briefly and observe it has grown. This is what the
// /health handler reads to detect a wedged watch — without
// /admin/v1/health reads this to detect a wedged watch — without
// this signal the proxy could serve stale config indefinitely.
let provider = Arc::new(FakeProvider::new(vec![], 5));
let sup = Supervisor::new(provider, "/aisix");
Expand Down
108 changes: 108 additions & 0 deletions crates/aisix-proxy/src/health.rs
Original file line number Diff line number Diff line change
Expand Up @@ -14,8 +14,79 @@
//! under stress without waiting for a full outage.

use dashmap::DashMap;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::{AtomicU32, Ordering};

use axum::http::header::{HeaderName, HeaderValue, CONTENT_TYPE};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};

static X_CONTENT_TYPE_OPTIONS: HeaderName = HeaderName::from_static("x-content-type-options");
static NOSNIFF: HeaderValue = HeaderValue::from_static("nosniff");
static TEXT_PLAIN_UTF8: HeaderValue = HeaderValue::from_static("text/plain; charset=utf-8");

#[derive(Debug, Default)]
pub struct LivezState {
shutting_down: AtomicBool,
}

impl LivezState {
pub fn new() -> Self {
Self::default()
}

pub fn mark_shutting_down(&self) {
self.shutting_down.store(true, Ordering::Relaxed);
}

fn shutdown_check(&self) -> Result<(), &'static str> {
if self.shutting_down.load(Ordering::Relaxed) {
Err("process is shutting down")
} else {
Ok(())
}
}
}

pub fn livez_response(livez: &LivezState, verbose: bool) -> Response {
let mut body = String::new();
let mut failed = false;

body.push_str("[+]ping ok\n");
match livez.shutdown_check() {
Ok(()) => body.push_str("[+]shutdown ok\n"),
Err(_) => {
failed = true;
body.push_str("[-]shutdown failed: reason withheld\n");
Comment on lines +51 to +60
}
}
Comment on lines +51 to +62

let headers = [
(CONTENT_TYPE, TEXT_PLAIN_UTF8.clone()),
(X_CONTENT_TYPE_OPTIONS.clone(), NOSNIFF.clone()),
];

if failed {
return (
StatusCode::INTERNAL_SERVER_ERROR,
headers,
format!("{body}livez check failed"),
)
.into_response();
}

if !verbose {
return (StatusCode::OK, headers, "ok").into_response();
}

(
StatusCode::OK,
headers,
format!("{body}livez check passed\n"),
)
.into_response()
}

/// Numeric health level reported by the API.
#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)]
#[serde(into = "u8")]
Expand Down Expand Up @@ -143,6 +214,7 @@ impl HealthTracker {
#[cfg(test)]
mod tests {
use super::*;
use axum::body::to_bytes;

#[test]
fn new_model_is_healthy() {
Expand Down Expand Up @@ -194,4 +266,40 @@ mod tests {
t.record_success("m");
assert_eq!(t.all_levels().len(), 1);
}

#[tokio::test]
async fn livez_default_success_is_plain_ok() {
let state = LivezState::new();
let resp = livez_response(&state, false);

assert_eq!(resp.status(), StatusCode::OK);
let body = to_bytes(resp.into_body(), 1024).await.unwrap();
assert_eq!(std::str::from_utf8(&body).unwrap(), "ok");
}

#[tokio::test]
async fn livez_verbose_success_lists_checks() {
let state = LivezState::new();
let resp = livez_response(&state, true);

assert_eq!(resp.status(), StatusCode::OK);
let body = to_bytes(resp.into_body(), 1024).await.unwrap();
let text = std::str::from_utf8(&body).unwrap();
assert!(text.contains("[+]ping ok"));
assert!(text.contains("[+]shutdown ok"));
assert!(text.contains("livez check passed"));
}

#[tokio::test]
async fn livez_failure_returns_500_with_reason_withheld() {
let state = LivezState::new();
state.mark_shutting_down();
let resp = livez_response(&state, false);

assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
let body = to_bytes(resp.into_body(), 1024).await.unwrap();
let text = std::str::from_utf8(&body).unwrap();
assert!(text.contains("[-]shutdown failed: reason withheld"));
assert!(text.contains("livez check failed"));
}
}
Loading