#![deny(missing_docs)]
#![deny(rustdoc::broken_intra_doc_links)]
pub mod error;
pub mod middleware;
use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use std::time::Duration;
use af_context::{PlatformContextProvider, RequestContext, RequestMetadata};
use axum::extract::{DefaultBodyLimit, State};
use axum::http::HeaderValue;
use axum::routing::get;
use axum::{Json, Router};
use serde_json::{json, Value};
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tonic::{Request, Status};
use tower::limit::ConcurrencyLimitLayer;
use tower_http::cors::{AllowOrigin, CorsLayer};
use tower_http::timeout::TimeoutLayer;
use tower_http::trace::TraceLayer;
pub use af_context::RequestId;
pub use error::ApiError;
pub use middleware::{latency_slo, rate_limit, request_id, RateLimiter, SLO_BUDGET_MS};
#[derive(Clone)]
pub struct EdgeConfig {
pub allowed_origins: Vec<HeaderValue>,
pub max_body_bytes: usize,
pub request_timeout: Duration,
pub max_concurrency: usize,
pub requests_per_window: u64,
pub rate_window: Duration,
}
impl Default for EdgeConfig {
fn default() -> Self {
Self {
allowed_origins: Vec::new(),
max_body_bytes: 1024 * 1024,
request_timeout: Duration::from_secs(30),
max_concurrency: 128,
requests_per_window: 120,
rate_window: Duration::from_secs(1),
}
}
}
#[derive(Clone, Default)]
pub struct Readiness(Arc<AtomicBool>);
impl Readiness {
pub fn set(&self, ready: bool) {
self.0.store(ready, Ordering::Release);
}
pub fn is_ready(&self) -> bool {
self.0.load(Ordering::Acquire)
}
}
pub async fn health() -> Json<Value> {
Json(json!({ "status": "ok" }))
}
pub async fn liveness() -> Json<Value> {
health().await
}
pub async fn readiness(State(readiness): State<Readiness>) -> impl axum::response::IntoResponse {
let status = if readiness.is_ready() {
axum::http::StatusCode::OK
} else {
axum::http::StatusCode::SERVICE_UNAVAILABLE
};
(
status,
Json(json!({ "status": if readiness.is_ready() { "ready" } else { "not_ready" } })),
)
}
pub async fn metrics() -> ([(&'static str, &'static str); 1], &'static str) {
(
[("content-type", "text/plain; version=0.0.4")],
"agent_factory_up 1\n",
)
}
pub fn with_standard_middleware<S>(router: Router<S>) -> Router<S>
where
S: Clone + Send + Sync + 'static,
{
router
.layer(axum::middleware::from_fn(latency_slo))
.layer(axum::middleware::from_fn(request_id))
}
pub fn with_edge_middleware(router: Router, config: EdgeConfig) -> Router {
let cors = if config.allowed_origins.is_empty() {
CorsLayer::new()
} else {
CorsLayer::new().allow_origin(AllowOrigin::list(config.allowed_origins))
};
let limiter = RateLimiter::new(config.requests_per_window, config.rate_window);
with_standard_middleware(router)
.layer(TraceLayer::new_for_http())
.layer(ConcurrencyLimitLayer::new(config.max_concurrency.max(1)))
.layer(TimeoutLayer::with_status_code(
axum::http::StatusCode::REQUEST_TIMEOUT,
config.request_timeout,
))
.layer(DefaultBodyLimit::max(config.max_body_bytes))
.layer(axum::middleware::from_fn_with_state(limiter, rate_limit))
.layer(cors)
}
pub fn base_router<S>() -> Router<S>
where
S: Clone + Send + Sync + 'static,
{
Router::new()
.route("/health", get(health))
.route("/metrics", get(metrics))
}
pub fn operations_router(state: Readiness) -> Router {
Router::new()
.route("/health", get(health))
.route("/live", get(liveness))
.route("/ready", get(readiness))
.route("/metrics", get(metrics))
.with_state(state)
}
pub async fn resolve_tonic_context<T>(
provider: &dyn PlatformContextProvider,
request: &Request<T>,
) -> Result<RequestContext, Status> {
let metadata = request
.metadata()
.iter()
.filter_map(|entry| match entry {
tonic::metadata::KeyAndValueRef::Ascii(key, value) => value
.to_str()
.ok()
.map(|value| (key.as_str().to_string(), value.to_string())),
tonic::metadata::KeyAndValueRef::Binary(_, _) => None,
})
.collect::<BTreeMap<_, _>>();
provider
.resolve(&RequestMetadata(metadata))
.await
.map_err(|error| Status::unauthenticated(error.to_string()))
}
pub async fn shutdown_signal(cancel: CancellationToken) {
#[cfg(unix)]
{
if let Ok(mut terminate) =
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
{
tokio::select! {
_ = tokio::signal::ctrl_c() => {}
_ = terminate.recv() => {}
}
} else {
let _ = tokio::signal::ctrl_c().await;
}
}
#[cfg(not(unix))]
let _ = tokio::signal::ctrl_c().await;
cancel.cancel();
}
pub async fn join_tasks(tasks: Vec<JoinHandle<()>>, grace: Duration) -> bool {
let aborts: Vec<_> = tasks.iter().map(JoinHandle::abort_handle).collect();
let joined = tokio::time::timeout(grace, futures_join(tasks))
.await
.unwrap_or(false);
if !joined {
for task in aborts {
task.abort();
}
}
joined
}
pub async fn supervise_tasks(
mut tasks: Vec<JoinHandle<()>>,
readiness: Readiness,
cancel: CancellationToken,
grace: Duration,
) -> bool {
use std::future::Future;
let requested = tokio::select! {
biased;
_ = cancel.cancelled() => true,
(index, outcome) = std::future::poll_fn(|cx| {
for (index, task) in tasks.iter_mut().enumerate() {
if let std::task::Poll::Ready(outcome) = std::pin::Pin::new(task).poll(cx) {
return std::task::Poll::Ready((index, outcome));
}
}
std::task::Poll::Pending
}) => {
drop(tasks.swap_remove(index));
tracing::error!(exit_reason = if outcome.is_err() { "task_panicked" } else { "task_exited" }, "host worker stopped unexpectedly");
false
}
};
readiness.set(false);
cancel.cancel();
let joined = join_tasks(tasks, grace).await;
tracing::info!(
exit_reason = if requested && joined {
"shutdown"
} else {
"worker_failure"
},
"host stopped"
);
requested && joined
}
async fn futures_join(tasks: Vec<JoinHandle<()>>) -> bool {
for task in tasks {
if task.await.is_err() {
return false;
}
}
true
}