use axum::{
extract::{Query, State},
Json,
};
use serde::{Deserialize, Serialize};
use crate::error::AppError;
use crate::services::runtime::{RegisterRuntimeRequest, RuntimeFilter, RuntimeService};
#[derive(Debug, Clone, Deserialize)]
pub struct DeregisterRequest {
pub kind: String,
pub name: String,
}
#[derive(Debug, Clone, Deserialize)]
pub struct HeartbeatRequest {
pub kind: String,
pub name: String,
}
#[derive(Debug, Clone, Serialize)]
pub struct RuntimeOperationResponse {
pub status: String,
pub message: String,
}
#[derive(Debug, Clone, Default, Deserialize)]
pub struct ListRuntimesQuery {
pub kind: Option<String>,
pub status: Option<String>,
pub name: Option<String>,
}
pub async fn register_pool(
service: State<RuntimeService>,
request: Json<RegisterRuntimeRequest>,
) -> Result<Json<crate::services::runtime::Runtime>, AppError> {
let started_at = std::time::Instant::now();
let result = register_pool_inner(service, request).await;
let status_label = if result.is_ok() { "ok" } else { "error" };
crate::metrics::record_write_request(
crate::metrics::endpoint::RUNTIME_REGISTER,
status_label,
started_at.elapsed().as_secs_f64(),
);
result
}
async fn register_pool_inner(
State(service): State<RuntimeService>,
Json(request): Json<RegisterRuntimeRequest>,
) -> Result<Json<crate::services::runtime::Runtime>, AppError> {
let runtime = service.register(&request).await?;
Ok(Json(runtime))
}
pub async fn deregister_pool(
State(service): State<RuntimeService>,
Json(request): Json<DeregisterRequest>,
) -> Result<Json<RuntimeOperationResponse>, AppError> {
service.deregister(&request.kind, &request.name).await?;
Ok(Json(RuntimeOperationResponse {
status: "ok".to_string(),
message: format!("Runtime {} {} deregistered", request.kind, request.name),
}))
}
pub async fn heartbeat(
service: State<RuntimeService>,
request: Json<HeartbeatRequest>,
) -> Result<Json<RuntimeOperationResponse>, AppError> {
let started_at = std::time::Instant::now();
let result = heartbeat_inner(service, request).await;
let status_label = if result.is_ok() { "ok" } else { "error" };
crate::metrics::record_write_request(
crate::metrics::endpoint::RUNTIME_HEARTBEAT,
status_label,
started_at.elapsed().as_secs_f64(),
);
result
}
async fn heartbeat_inner(
State(service): State<RuntimeService>,
Json(request): Json<HeartbeatRequest>,
) -> Result<Json<RuntimeOperationResponse>, AppError> {
service.heartbeat(&request.kind, &request.name).await?;
Ok(Json(RuntimeOperationResponse {
status: "ok".to_string(),
message: "Heartbeat recorded".to_string(),
}))
}
pub async fn list_pools(
State(service): State<RuntimeService>,
Query(query): Query<ListRuntimesQuery>,
) -> Result<Json<Vec<crate::services::runtime::Runtime>>, AppError> {
let filter = RuntimeFilter {
kind: query.kind.or(Some("worker_pool".to_string())),
status: query.status,
name: query.name,
};
let pools = service.list(&filter).await?;
Ok(Json(pools))
}
#[allow(dead_code)]
pub async fn list_all(
State(service): State<RuntimeService>,
Query(query): Query<ListRuntimesQuery>,
) -> Result<Json<Vec<crate::services::runtime::Runtime>>, AppError> {
let filter = RuntimeFilter {
kind: query.kind,
status: query.status,
name: query.name,
};
let runtimes = service.list(&filter).await?;
Ok(Json(runtimes))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_deregister_request_deserialization() {
let json = r#"{"kind": "worker_pool", "name": "worker-1"}"#;
let request: DeregisterRequest = serde_json::from_str(json).unwrap();
assert_eq!(request.kind, "worker_pool");
assert_eq!(request.name, "worker-1");
}
#[test]
fn test_heartbeat_request_deserialization() {
let json = r#"{"kind": "worker_pool", "name": "worker-1"}"#;
let request: HeartbeatRequest = serde_json::from_str(json).unwrap();
assert_eq!(request.kind, "worker_pool");
assert_eq!(request.name, "worker-1");
}
#[test]
fn test_operation_response_serialization() {
let response = RuntimeOperationResponse {
status: "ok".to_string(),
message: "Success".to_string(),
};
let json = serde_json::to_string(&response).unwrap();
assert!(json.contains("ok"));
assert!(json.contains("Success"));
}
}