From 38599ca0e803b9c8b551112559222010e6ed8841 Mon Sep 17 00:00:00 2001 From: un-lock-able Date: Fri, 30 Oct 2026 11:17:30 +0800 Subject: [PATCH] feat(agent): add agent registration, heartbeat and notification --- Cargo.lock | 76 ++++ netmito/Cargo.toml | 3 +- netmito/src/api/agents.rs | 108 +++++ netmito/src/api/groups.rs | 5 +- netmito/src/api/mod.rs | 4 + netmito/src/api/workers.rs | 6 +- netmito/src/api/ws.rs | 17 + netmito/src/client/http.rs | 55 +++ netmito/src/client/mod.rs | 36 ++ netmito/src/config/client/agents.rs | 72 +++ netmito/src/config/client/mod.rs | 4 + netmito/src/config/coordinator.rs | 44 ++ netmito/src/coordinator.rs | 35 +- netmito/src/lib.rs | 1 + netmito/src/schema/agent.rs | 171 +++++++ netmito/src/schema/mod.rs | 2 + netmito/src/service/agent/heartbeat.rs | 153 +++++++ netmito/src/service/agent/mod.rs | 608 +++++++++++++++++++++++++ netmito/src/service/auth/mod.rs | 49 +- netmito/src/service/auth/token.rs | 8 +- netmito/src/service/mod.rs | 1 + netmito/src/ws/connection.rs | 273 +++++++++++ netmito/src/ws/handler.rs | 117 +++++ netmito/src/ws/mod.rs | 12 + openapi.yaml | 435 ++++++++++++++++++ 25 files changed, 2280 insertions(+), 15 deletions(-) create mode 100644 netmito/src/api/agents.rs create mode 100644 netmito/src/api/ws.rs create mode 100644 netmito/src/config/client/agents.rs create mode 100644 netmito/src/schema/agent.rs create mode 100644 netmito/src/service/agent/heartbeat.rs create mode 100644 netmito/src/service/agent/mod.rs create mode 100644 netmito/src/ws/connection.rs create mode 100644 netmito/src/ws/handler.rs create mode 100644 netmito/src/ws/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 3b18f3f2..62a142cf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -563,6 +563,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8" dependencies = [ "axum-core", + "base64", "bytes", "form_urlencoded", "futures-util", @@ -581,8 +582,10 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_urlencoded", + "sha1", "sync_wrapper", "tokio", + "tokio-tungstenite", "tower", "tower-layer", "tower-service", @@ -1168,6 +1171,12 @@ dependencies = [ "syn", ] +[[package]] +name = "data-encoding" +version = "2.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4583a4551df46e2792f82ceeac45e850d2e2d5debba0b91f102385cda5b11f06" + [[package]] name = "der" version = "0.6.1" @@ -2476,6 +2485,15 @@ version = "2.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + [[package]] name = "mime" version = "0.3.17" @@ -2558,6 +2576,7 @@ dependencies = [ "sea-orm-migration", "serde", "serde_json", + "speedy", "thiserror 2.0.18", "time", "tokio", @@ -3999,6 +4018,28 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "speedy" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da1992073f0e55aab599f4483c460598219b4f9ff0affa124b33580ab511e25a" +dependencies = [ + "memoffset", + "speedy-derive", + "uuid", +] + +[[package]] +name = "speedy-derive" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "658f2ca5276b92c3dfd65fa88316b4e032ace68f88d7570b43967784c0bac5ac" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "spin" version = "0.9.8" @@ -4534,6 +4575,18 @@ dependencies = [ "xattr", ] +[[package]] +name = "tokio-tungstenite" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857" +dependencies = [ + "futures-util", + "log", + "tokio", + "tungstenite", +] + [[package]] name = "tokio-util" version = "0.7.18" @@ -4717,6 +4770,23 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "tungstenite" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8628dcc84e5a09eb3d8423d6cb682965dea9133204e8fb3efee74c2a0c259442" +dependencies = [ + "bytes", + "data-encoding", + "http 1.4.0", + "httparse", + "log", + "rand 0.9.2", + "sha1", + "thiserror 2.0.18", + "utf-8", +] + [[package]] name = "typenum" version = "1.19.0" @@ -4834,6 +4904,12 @@ dependencies = [ "serde_derive", ] +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + [[package]] name = "utf8_iter" version = "1.0.4" diff --git a/netmito/Cargo.toml b/netmito/Cargo.toml index 94c29d5c..4d8c246b 100644 --- a/netmito/Cargo.toml +++ b/netmito/Cargo.toml @@ -18,7 +18,7 @@ argon2 = { version = "0.5.3", features = ["std"] } async-compression = { version = "0.4.32", features = ["gzip", "tokio"] } async-stream = "0.3.6" aws-sdk-s3 = { version = "1.110", features = ["behavior-version-latest"] } -axum = { version = "0.8.6", features = ["http2"] } +axum = { version = "0.8.6", features = ["http2", "ws"] } axum-extra = { version = "0.12", features = ["typed-header", "query"] } base64 = "0.22.1" clap = { workspace = true } @@ -65,6 +65,7 @@ sea-orm = { version = "1.1.17", default-features = false, features = [ sea-orm-migration = "1.1.17" serde = { workspace = true } serde_json = { workspace = true } +speedy = { version = "0.8.7", features = ["uuid"] } thiserror = "2.0.17" time = { version = "0.3.44", features = ["serde-human-readable"] } tokio = { workspace = true } diff --git a/netmito/src/api/agents.rs b/netmito/src/api/agents.rs new file mode 100644 index 00000000..6532ec0b --- /dev/null +++ b/netmito/src/api/agents.rs @@ -0,0 +1,108 @@ +use axum::{ + extract::{Path, Query, State}, + middleware, + routing::{delete, post}, + Extension, Json, Router, +}; +use uuid::Uuid; + +use super::map_service_error; +use crate::{ + config::InfraPool, + error::ApiError, + schema::{ + AgentHeartbeatReq, AgentHeartbeatResp, AgentShutdownReq, AgentsQueryReq, AgentsQueryResp, + RegisterAgentReq, RegisterAgentResp, + }, + service::{ + self, + auth::{ + auth_middleware, + roles::{AGENT, USER}, + AuthAgent, AuthUser, + }, + }, +}; + +pub fn agents_router(st: InfraPool) -> Router { + let user_router = Router::new() + .route("/", post(register_agent)) + .route("/query", post(query_agents)) + .route("/{uuid}", delete(shutdown_agent)) + .route_layer(middleware::from_fn_with_state( + st.clone(), + auth_middleware::, + )) + .with_state(st.clone()); + + let agent_router = Router::new() + .route("/", delete(agent_exit)) + .route("/heartbeat", post(heartbeat)) + .route_layer(middleware::from_fn_with_state( + st.clone(), + auth_middleware::, + )) + .with_state(st.clone()); + + Router::new().merge(user_router).merge(agent_router) +} + +/// `POST /agents`: register an agent and issue its token. +async fn register_agent( + Extension(u): Extension, + State(pool): State, + Json(req): Json, +) -> Result, ApiError> { + let resp = service::agent::user_register_agent(u.id, &pool, req) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} + +/// `POST /agents/query` +async fn query_agents( + Extension(u): Extension, + State(pool): State, + Json(req): Json, +) -> Result, ApiError> { + let resp = service::agent::user_query_agents(u.id, &pool, req) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} + +/// `DELETE /agents/{uuid}?op=graceful|force`: shut the agent down. The agent row stays. +async fn shutdown_agent( + Extension(u): Extension, + State(pool): State, + Path(uuid): Path, + Query(req): Query, +) -> Result<(), ApiError> { + service::agent::user_shutdown_agent_by_uuid(u.id, uuid, req.op, &pool) + .await + .map_err(map_service_error)?; + Ok(()) +} + +/// `DELETE /agents`: the agent exits. Retire it now, so that it does not wait for its heartbeat times out. +async fn agent_exit( + Extension(a): Extension, + State(pool): State, +) -> Result<(), ApiError> { + service::agent::retire_agent(&pool, a.id, service::agent::RetireCause::SelfExit) + .await + .map_err(map_service_error)?; + Ok(()) +} + +/// `POST /agents/heartbeat`: record a heartbeat and return the notifications that the agent has not acknowledged. +async fn heartbeat( + Extension(a): Extension, + State(pool): State, + Json(req): Json, +) -> Result, ApiError> { + let resp = service::agent::agent_heartbeat(a.id, a.uuid, &pool, req) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} diff --git a/netmito/src/api/groups.rs b/netmito/src/api/groups.rs index 9000c1fe..8f1647f1 100644 --- a/netmito/src/api/groups.rs +++ b/netmito/src/api/groups.rs @@ -8,7 +8,7 @@ use axum_extra::extract::Query; use crate::{ config::InfraPool, - error::{ApiError, ApiResult}, + error::{ApiError, ApiResult, AuthError}, schema::*, service::{ self, @@ -195,6 +195,9 @@ pub async fn download_attachment( AuthIdentity::Worker(w) => { service::s3::worker_get_attachment(&pool, w.uuid, group_name, key).await } + // TODO: let an agent read the attachments of a group whose suite it runs a job for. + // Do this when agents run jobs. Before that, no job rows exist and the check always fails. + AuthIdentity::Agent(_) => Err(AuthError::PermissionDenied.into()), } .map_err(|e| match e { crate::error::Error::AuthError(err) => ApiError::AuthError(err), diff --git a/netmito/src/api/mod.rs b/netmito/src/api/mod.rs index cbb6bf1b..fdeca2df 100644 --- a/netmito/src/api/mod.rs +++ b/netmito/src/api/mod.rs @@ -1,9 +1,11 @@ pub mod admin; +pub mod agents; pub mod groups; pub mod suites; pub mod tasks; pub mod users; pub mod workers; +pub mod ws; #[cfg(feature = "debugging")] use std::net::SocketAddr; @@ -77,6 +79,8 @@ pub fn router(st: InfraPool, cancel_token: CancellationToken) -> Router { .nest("/workers", workers::workers_router(st.clone())) .nest("/tasks", tasks::tasks_router(st.clone())) .nest("/suites", suites::suites_router(st.clone())) + .nest("/agents", agents::agents_router(st.clone())) + .nest("/ws", ws::ws_router(st.clone())) .with_state(st) .layer(CorsLayer::permissive()) .layer(CatchPanicLayer::new()); diff --git a/netmito/src/api/workers.rs b/netmito/src/api/workers.rs index 8cdaf493..21efd945 100644 --- a/netmito/src/api/workers.rs +++ b/netmito/src/api/workers.rs @@ -16,7 +16,7 @@ use crate::{ auth::{ auth_middleware, roles::{USER, WORKER}, - token::generate_worker_token, + token::{generate_runner_token, TokenRole}, AuthUser, AuthWorker, }, }, @@ -78,8 +78,8 @@ pub async fn user_register_worker( ApiError::InternalServerError } })?; - let token = - generate_worker_token(uuid.clone().to_string(), 0, req.lifetime).map_err(|e| match e { + let token = generate_runner_token(uuid.clone().to_string(), 0, req.lifetime, TokenRole::Worker) + .map_err(|e| match e { crate::error::Error::AuthError(err) => ApiError::AuthError(err), crate::error::Error::ApiError(e) => e, _ => { diff --git a/netmito/src/api/ws.rs b/netmito/src/api/ws.rs new file mode 100644 index 00000000..b5f5f290 --- /dev/null +++ b/netmito/src/api/ws.rs @@ -0,0 +1,17 @@ +use axum::{middleware, routing::get, Router}; + +use crate::{ + config::InfraPool, + service::auth::{auth_middleware, roles::AGENT}, + ws::websocket_handler, +}; + +pub fn ws_router(st: InfraPool) -> Router { + Router::new() + .route("/agents", get(websocket_handler)) + .route_layer(middleware::from_fn_with_state( + st.clone(), + auth_middleware::, + )) + .with_state(st) +} diff --git a/netmito/src/client/http.rs b/netmito/src/client/http.rs index 94e0299f..7fcd62b8 100644 --- a/netmito/src/client/http.rs +++ b/netmito/src/client/http.rs @@ -886,6 +886,61 @@ impl MitoHttpClient { } } + pub async fn query_agents( + &mut self, + req: AgentsQueryReq, + ) -> crate::error::Result { + let credential = self + .credential_guard + .get_credential() + .ok_or(CredentialGuardError::CredentialNotFound)? + .token; + self.url.set_path("agents/query"); + let resp = self + .http_client + .post(self.url.as_str()) + .bearer_auth(credential) + .json(&req) + .send() + .await + .map_err(map_reqwest_err)?; + if resp.status().is_success() { + let resp = resp + .json::() + .await + .map_err(RequestError::from)?; + Ok(resp) + } else { + Err(get_error_from_resp(resp).await.into()) + } + } + + pub async fn shutdown_agent(&mut self, uuid: Uuid, force: bool) -> crate::error::Result<()> { + let credential = self + .credential_guard + .get_credential() + .ok_or(CredentialGuardError::CredentialNotFound)? + .token; + self.url.set_path(&format!("agents/{uuid}")); + if force { + self.url.set_query(Some("op=force")); + } + let resp = self + .http_client + .delete(self.url.as_str()) + .bearer_auth(credential) + .send() + .await + .map_err(map_reqwest_err)?; + // Other requests use `self.url` again, thus remove the query. + self.url.set_query(None); + if resp.status().is_success() { + Ok(()) + } else { + Err(get_error_from_resp(resp).await.into()) + } + } + pub async fn get_group_by_name( &mut self, group_name: &str, diff --git a/netmito/src/client/mod.rs b/netmito/src/client/mod.rs index 98ccd766..3954f2ae 100644 --- a/netmito/src/client/mod.rs +++ b/netmito/src/client/mod.rs @@ -2144,6 +2144,42 @@ impl MitoClient { self.suites_override_agents(args).await; } }, + ClientCommand::Agents(args) => match args.command { + AgentsCommands::Query(args) => { + let counted = args.count; + match self.http_client.query_agents(args.into()).await { + Ok(resp) => { + tracing::info!( + "Found {} agents in group {}", + resp.count, + resp.group_name + ); + if !counted { + for agent in resp.agents { + tracing::info!( + "{} ({}), tags: [{}], machine: {}, suite: {}", + agent.uuid, + agent.state, + agent.tags.join(", "), + agent.machine_code.unwrap_or("[unknown]".to_string()), + agent + .assigned_suite_uuid + .map(|u| u.to_string()) + .unwrap_or("[none]".to_string()), + ); + } + } + } + Err(e) => tracing::error!("{}", e), + } + } + AgentsCommands::Shutdown(args) => { + match self.http_client.shutdown_agent(args.uuid, args.force).await { + Ok(_) => tracing::info!("Agent {} asked to shut down", args.uuid), + Err(e) => tracing::error!("{}", e), + } + } + }, ClientCommand::Quit => { return false; } diff --git a/netmito/src/config/client/agents.rs b/netmito/src/config/client/agents.rs new file mode 100644 index 00000000..7d92f592 --- /dev/null +++ b/netmito/src/config/client/agents.rs @@ -0,0 +1,72 @@ +use clap::{Args, Subcommand}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +use crate::{entity::state::AgentState, schema::AgentsQueryReq}; + +#[derive(Serialize, Debug, Deserialize, Args, derive_more::From, Clone)] +pub struct AgentsArgs { + #[command(subcommand)] + pub command: AgentsCommands, +} + +#[derive(Subcommand, Serialize, Debug, Deserialize, Clone)] +pub enum AgentsCommands { + /// Query agents subject to a filter + Query(QueryAgentsArgs), + /// Shut an agent down. The agent record stays. + Shutdown(ShutdownAgentArgs), +} + +#[derive(Serialize, Debug, Deserialize, Args, Clone)] +pub struct QueryAgentsArgs { + /// Filter by group name. The default is your username. + #[arg(short, long)] + pub group: Option, + /// Filter by tags + #[arg(short, long, num_args = 0.., value_delimiter = ',')] + pub tags: Vec, + /// Filter by labels + #[arg(short, long, num_args = 0.., value_delimiter = ',')] + pub labels: Vec, + /// Filter by states + #[arg(long, num_args = 0.., value_delimiter = ',')] + pub states: Vec, + /// Filter by creator username + #[arg(long)] + pub creator: Option, + /// Maximum number of results to return + #[arg(long)] + pub limit: Option, + /// Number of results to skip + #[arg(long)] + pub offset: Option, + /// Only return the number of matching agents + #[arg(long)] + pub count: bool, +} + +impl From for AgentsQueryReq { + fn from(args: QueryAgentsArgs) -> Self { + Self { + group_name: args.group, + tags: (!args.tags.is_empty()).then(|| args.tags.into_iter().collect()), + labels: (!args.labels.is_empty()).then(|| args.labels.into_iter().collect()), + states: (!args.states.is_empty()).then(|| args.states.into_iter().collect()), + creator_username: args.creator, + limit: args.limit, + offset: args.offset, + count: args.count, + } + } +} + +#[derive(Serialize, Debug, Deserialize, Args, Clone)] +pub struct ShutdownAgentArgs { + /// The UUID of the agent + pub uuid: Uuid, + /// Mark the agent offline now. Without this flag, the agent finishes its + /// current work and then stops. + #[arg(short, long)] + pub force: bool, +} diff --git a/netmito/src/config/client/mod.rs b/netmito/src/config/client/mod.rs index 50485fc3..d9365e29 100644 --- a/netmito/src/config/client/mod.rs +++ b/netmito/src/config/client/mod.rs @@ -16,6 +16,7 @@ use uuid::Uuid; use super::coordinator::DEFAULT_COORDINATOR_ADDR; pub mod admin; +pub mod agents; pub mod artifacts; pub mod attachments; pub mod groups; @@ -24,6 +25,7 @@ pub mod tasks; pub mod users; pub mod workers; pub use admin::*; +pub use agents::*; pub use artifacts::*; pub use attachments::*; pub use groups::*; @@ -113,6 +115,8 @@ pub enum ClientCommand { Suites(SuitesArgs), /// Manage workers, including querying workers, cancel workers, etc. Workers(WorkersArgs), + /// Manage agents, including querying agents and shutting one down. + Agents(AgentsArgs), /// Run an external command Cmd(CmdArgs), /// Quit the client's interactive mode diff --git a/netmito/src/config/coordinator.rs b/netmito/src/config/coordinator.rs index bb23d9c8..a3601c71 100644 --- a/netmito/src/config/coordinator.rs +++ b/netmito/src/config/coordinator.rs @@ -23,7 +23,9 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, Layer}; use crate::{ error::Error, + service::agent::heartbeat::{AgentHeartbeatOp, AgentHeartbeatQueue}, service::worker::{HeartbeatOp, TaskDispatcher, TaskDispatcherOp, WorkerHeartbeatQueue}, + ws::{AgentWsRouter, RouterOp}, }; use super::TracingGuard; @@ -69,6 +71,10 @@ pub struct CoordinatorConfig { pub(crate) access_token_expires_in: std::time::Duration, #[serde(with = "humantime_serde")] pub(crate) heartbeat_timeout: std::time::Duration, + /// The interval of WebSocket pings to a connected agent. Without them, + /// proxies can close a notification socket that has no traffic. + #[serde(with = "humantime_serde", default = "default_ws_keepalive_interval")] + pub(crate) ws_keepalive_interval: std::time::Duration, pub(crate) log_path: Option, pub(crate) file_log: bool, } @@ -77,6 +83,10 @@ fn default_db_max_connections() -> u32 { 32 } +fn default_ws_keepalive_interval() -> std::time::Duration { + std::time::Duration::from_secs(30) +} + fn default_db_min_connections() -> u32 { 8 } @@ -174,6 +184,10 @@ pub struct CoordinatorConfigCli { #[arg(long)] #[serde(skip_serializing_if = "::std::option::Option::is_none")] pub heartbeat_timeout: Option, + /// The agent WebSocket keepalive interval, default to 30 seconds + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub ws_keepalive_interval: Option, /// The log file path. If not specified, then the default rolling log file path would be used. /// If specified, then the log file would be exactly at the path specified. #[arg(long)] @@ -217,6 +231,7 @@ impl Default for CoordinatorConfig { access_token_public_path: "public.pem".to_string().into(), access_token_expires_in: std::time::Duration::from_secs(60 * 60 * 24 * 7), heartbeat_timeout: std::time::Duration::from_secs(600), + ws_keepalive_interval: default_ws_keepalive_interval(), log_path: None, file_log: false, } @@ -261,6 +276,23 @@ impl CoordinatorConfig { WorkerHeartbeatQueue::new(cancel_token, self.heartbeat_timeout, pool, rx) } + pub fn build_agent_heartbeat_queue( + &self, + cancel_token: CancellationToken, + pool: InfraPool, + rx: AsyncRx, + ) -> AgentHeartbeatQueue { + AgentHeartbeatQueue::new(cancel_token, self.heartbeat_timeout, pool, rx) + } + + pub fn build_ws_router( + &self, + cancel_token: CancellationToken, + rx: AsyncRx, + ) -> AgentWsRouter { + AgentWsRouter::new(cancel_token, rx) + } + pub async fn build_redis_connection_info( &self, ) -> crate::error::Result> { @@ -325,6 +357,8 @@ impl CoordinatorConfig { &self, worker_task_queue_tx: MTx, worker_heartbeat_queue_tx: MTx, + agent_heartbeat_queue_tx: MTx, + ws_router_tx: MTx, ) -> crate::error::Result { let mut db_opt = sea_orm::ConnectOptions::new(self.db_url.clone()); db_opt @@ -352,6 +386,10 @@ impl CoordinatorConfig { attachments_bucket: self.attachments_bucket.clone(), worker_task_queue_tx, worker_heartbeat_queue_tx, + agent_heartbeat_queue_tx, + ws_router_tx, + boot_uuid: uuid::Uuid::new_v4(), + ws_keepalive_interval: self.ws_keepalive_interval, }) } @@ -470,6 +508,12 @@ pub struct InfraPool { pub attachments_bucket: String, pub worker_task_queue_tx: MTx, pub worker_heartbeat_queue_tx: MTx, + pub agent_heartbeat_queue_tx: MTx, + pub ws_router_tx: MTx, + /// Identifies this coordinator process. An agent compares it with the + /// `boot_id` of a `CounterSync` to find out that the coordinator restarted. + pub boot_uuid: uuid::Uuid, + pub ws_keepalive_interval: std::time::Duration, } #[derive(Debug)] diff --git a/netmito/src/coordinator.rs b/netmito/src/coordinator.rs index 7e961d6b..15822fb7 100644 --- a/netmito/src/coordinator.rs +++ b/netmito/src/coordinator.rs @@ -8,14 +8,18 @@ use tracing_subscriber::layer::SubscriberExt; use crate::api::router; use crate::config::{CoordinatorConfig, CoordinatorConfigCli, InfraPool}; use crate::migration::{Migrator, MigratorTrait}; +use crate::service::agent::{heartbeat::AgentHeartbeatQueue, restore_agents}; use crate::service::s3::setup_buckets; use crate::service::worker::{restore_workers, TaskDispatcher, WorkerHeartbeatQueue}; use crate::signal::shutdown_signal; +use crate::ws::AgentWsRouter; pub struct MitoCoordinator { pub infra_pool: InfraPool, pub worker_task_queue: TaskDispatcher, pub worker_heartbeat_queue: WorkerHeartbeatQueue, + pub agent_heartbeat_queue: AgentHeartbeatQueue, + pub ws_router: AgentWsRouter, pub cancel_token: CancellationToken, pub log_dir: PathBuf, } @@ -102,6 +106,9 @@ impl MitoCoordinator { let (worker_task_queue_tx, worker_task_queue_rx) = crossfire::mpsc::unbounded_async(); let (worker_heartbeat_queue_tx, worker_heartbeat_queue_rx) = crossfire::mpsc::unbounded_async(); + let (agent_heartbeat_queue_tx, agent_heartbeat_queue_rx) = + crossfire::mpsc::unbounded_async(); + let (ws_router_tx, ws_router_rx) = crossfire::mpsc::unbounded_async(); // Setup worker task queue let worker_task_queue = @@ -109,7 +116,12 @@ impl MitoCoordinator { // Setup infra pool let infra_pool = config - .build_infra_pool(worker_task_queue_tx, worker_heartbeat_queue_tx) + .build_infra_pool( + worker_task_queue_tx, + worker_heartbeat_queue_tx, + agent_heartbeat_queue_tx, + ws_router_tx, + ) .await?; // Setup worker heartbeat queue @@ -119,6 +131,14 @@ impl MitoCoordinator { worker_heartbeat_queue_rx, ); + // Setup agent heartbeat queue and notification router + let agent_heartbeat_queue = config.build_agent_heartbeat_queue( + cancel_token.clone(), + infra_pool.clone(), + agent_heartbeat_queue_rx, + ); + let ws_router = config.build_ws_router(cancel_token.clone(), ws_router_rx); + // Setup s3 storage // List all buckets and create if not exist setup_buckets( @@ -144,6 +164,8 @@ impl MitoCoordinator { infra_pool, worker_task_queue, worker_heartbeat_queue, + agent_heartbeat_queue, + ws_router, cancel_token, log_dir, }) @@ -155,6 +177,8 @@ impl MitoCoordinator { infra_pool, mut worker_task_queue, mut worker_heartbeat_queue, + mut agent_heartbeat_queue, + mut ws_router, cancel_token, .. } = self; @@ -173,7 +197,16 @@ impl MitoCoordinator { tracing::info!("Heartbeat queue stopped"); }); + task_tracker.spawn(async move { + agent_heartbeat_queue.run().await; + }); + + task_tracker.spawn(async move { + ws_router.run().await; + }); + restore_workers(&infra_pool).await?; + restore_agents(&infra_pool).await?; let app = router(infra_pool, cancel_token.clone()); let addr = crate::config::SERVER_CONFIG .get() diff --git a/netmito/src/lib.rs b/netmito/src/lib.rs index 20bb573c..eb483dbe 100644 --- a/netmito/src/lib.rs +++ b/netmito/src/lib.rs @@ -11,6 +11,7 @@ pub mod schema; pub mod service; pub mod signal; pub mod worker; +pub mod ws; pub mod reexports { pub use redis; pub use time; diff --git a/netmito/src/schema/agent.rs b/netmito/src/schema/agent.rs new file mode 100644 index 00000000..e3012b13 --- /dev/null +++ b/netmito/src/schema/agent.rs @@ -0,0 +1,171 @@ +//! Wire types for agents: the user routes under `/agents`, the agent routes +//! under `/agents`, and the notifications on `/ws/agents`. + +use std::collections::HashSet; + +use sea_orm::FromQueryResult; +use serde::{Deserialize, Serialize}; +use speedy::{Readable, Writable}; +use time::OffsetDateTime; +use uuid::Uuid; + +use crate::entity::state::AgentState; + +/// Request to register an agent. +/// +/// One machine has one agent row for all its life, because `machine_code` is +/// unique. If the machine is already known, registration uses its agent row +/// again: it replaces the tags, labels, groups and metadata, and it issues a +/// new token. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegisterAgentReq { + /// Tags that a suite must match to run on this agent (e.g. `["gpu", "linux"]`). + #[serde(default)] + pub tags: HashSet, + /// Labels for queries and filters (e.g. `["datacenter:us-west"]`). + #[serde(default)] + pub labels: HashSet, + /// The group that gets the Admin role on the agent. Only an Admin of this + /// group can shut the agent down, and the caller must be one. The default + /// is the personal group of the caller. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub admin_group: Option, + /// Groups that get the Write role on the agent. + #[serde(default)] + pub groups: HashSet, + /// The lifetime of the agent token. `None` means the token never expires. + #[serde(default, with = "humantime_serde")] + pub lifetime: Option, + /// A stable identifier of the machine that runs the agent. + pub machine_code: String, + /// Static data about the agent process, kept on the machine row. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub metadata: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegisterAgentResp { + pub agent_uuid: Uuid, + pub token: String, + /// The last notification id issued to this agent. The agent counts its + /// notifications from this value. + pub notification_counter: u64, + /// True if the machine was already known and its agent row is used again. + pub reused: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AgentMetadata { + /// Agent binary version (e.g. "0.6.8"). + pub version: String, + /// Long version string, with build data. + pub long_version: String, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct AgentsQueryReq { + pub group_name: Option, + pub tags: Option>, + pub labels: Option>, + pub states: Option>, + pub creator_username: Option, + pub limit: Option, + pub offset: Option, + pub count: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize, FromQueryResult)] +pub struct AgentInfo { + pub uuid: Uuid, + pub creator_username: String, + pub tags: Vec, + pub labels: Vec, + pub state: AgentState, + pub last_heartbeat: OffsetDateTime, + #[serde(skip_serializing_if = "Option::is_none")] + pub assigned_suite_uuid: Option, + pub created_at: OffsetDateTime, + pub updated_at: OffsetDateTime, + #[serde(skip_serializing_if = "Option::is_none")] + pub machine_code: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub metadata: Option, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AgentsQueryResp { + pub count: u64, + pub agents: Vec, + pub group_name: String, +} + +/// Query parameters of `DELETE /agents/{uuid}`. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct AgentShutdownReq { + #[serde(default)] + pub op: AgentShutdownOp, +} + +/// How to shut an agent down. The agent row stays in both cases. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq)] +pub enum AgentShutdownOp { + /// Tell the agent to stop. The agent finishes its current work and then + /// marks itself `Offline`. + #[default] + #[serde(alias = "graceful")] + Graceful, + /// Mark the agent `Offline` now, then tell it to stop. + #[serde(alias = "force")] + Force, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AgentHeartbeatReq { + /// The state the agent reports for itself. + pub state: AgentState, + /// The highest notification id the agent has processed. + #[serde(default)] + pub last_notification_id: u64, +} + +/// The notifications that the agent has not acknowledged yet. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct AgentHeartbeatResp { + pub notifications: Vec, +} + +/// A notification with its sequence id. The id tells an agent which +/// notifications it has already processed. +/// +/// `speedy` encodes it on the WebSocket. The heartbeat response sends the same +/// events as JSON, thus it also needs serde. +#[derive(Debug, Clone, Serialize, Deserialize, Readable, Writable)] +pub struct WsNotificationEvent { + /// Per-agent sequence id. It increases by one for each notification. + pub id: u64, + pub event: AgentNotification, +} + +/// A message from the coordinator that tells an agent to do something. +/// +/// Add new variants only at the end. `speedy` encodes the variant index, thus +/// a variant inserted in the middle changes the meaning of the variants after it. +#[derive(Debug, Clone, Serialize, Deserialize, Readable, Writable)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum AgentNotification { + /// Shut down. With `graceful`, finish the current work first. + Shutdown { graceful: bool }, + /// The coordinator restarted and its notification ids start again from + /// `counter`. `boot_id` identifies the new coordinator process. + CounterSync { counter: u64, boot_id: Uuid }, +} + +/// A message from an agent to the coordinator on the WebSocket. +#[derive(Debug, Clone, Serialize, Deserialize, Readable, Writable)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum AgentWsMessage { + /// The agent processed all notifications up to and including this id. + Ack { notification_id: u64 }, + /// The reply to a ping, with the agent's clock in Unix seconds. + Pong { client_time: i64 }, +} diff --git a/netmito/src/schema/mod.rs b/netmito/src/schema/mod.rs index 1877a16a..8e98ffbe 100644 --- a/netmito/src/schema/mod.rs +++ b/netmito/src/schema/mod.rs @@ -1,3 +1,4 @@ +mod agent; mod artifact; mod attachment; mod exec; @@ -7,6 +8,7 @@ mod task; mod user; mod worker; +pub use agent::*; pub use artifact::*; pub use attachment::*; pub use exec::*; diff --git a/netmito/src/service/agent/heartbeat.rs b/netmito/src/service/agent/heartbeat.rs new file mode 100644 index 00000000..11eea974 --- /dev/null +++ b/netmito/src/service/agent/heartbeat.rs @@ -0,0 +1,153 @@ +//! Heartbeat deadlines of agents. An agent that misses its deadline is retired. + +use std::{cmp::Reverse, time::Duration}; + +use crossfire::AsyncRx; +use priority_queue::PriorityQueue; +use sea_orm::prelude::*; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; + +use crate::{ + config::InfraPool, + entity::{agents as Agent, state::AgentState}, + service::agent::{retire_agent, RetireCause}, +}; + +/// The maximum time of one database call during timeout handling. +const DB_TIMEOUT: Duration = Duration::from_secs(10); + +pub enum AgentHeartbeatOp { + /// Record a heartbeat and set a new deadline for the agent. + Heartbeat(i64), + /// Stop tracking the agent. + Remove(i64), +} + +pub struct AgentHeartbeatQueue { + agents: PriorityQueue>, + cancel_token: CancellationToken, + heartbeat_timeout: Duration, + pool: InfraPool, + rx: AsyncRx, +} + +impl AgentHeartbeatQueue { + pub fn new( + cancel_token: CancellationToken, + heartbeat_timeout: Duration, + pool: InfraPool, + rx: AsyncRx, + ) -> Self { + Self { + agents: PriorityQueue::new(), + cancel_token, + heartbeat_timeout, + pool, + rx, + } + } + + fn handle_op(&mut self, op: AgentHeartbeatOp) { + match op { + AgentHeartbeatOp::Heartbeat(agent_id) => { + self.agents + .push(agent_id, Reverse(Instant::now() + self.heartbeat_timeout)); + } + AgentHeartbeatOp::Remove(agent_id) => { + self.agents.remove(&agent_id); + } + } + } + + /// The time until the earliest deadline. If no agent is tracked, a full timeout. + fn next_deadline(&self) -> Duration { + self.agents + .peek() + .map(|(_, deadline)| deadline.0.saturating_duration_since(Instant::now())) + .unwrap_or(self.heartbeat_timeout) + } + + async fn handle_timeout(&mut self) -> crate::error::Result<()> { + let expired = self + .agents + .peek() + .is_some_and(|(_, deadline)| deadline.0 <= Instant::now()); + if !expired { + return Ok(()); + } + let (agent_id, _) = self.agents.pop().unwrap(); + + let agent = match tokio::time::timeout( + DB_TIMEOUT, + Agent::Entity::find_by_id(agent_id).one(&self.pool.db), + ) + .await + { + Ok(Ok(Some(agent))) => agent, + Ok(Ok(None)) => { + tracing::debug!(agent_id, "Agent not found at its heartbeat timeout"); + return Ok(()); + } + Ok(Err(e)) => return Err(e.into()), + Err(_) => { + tracing::warn!(agent_id, "Agent lookup timed out during heartbeat check"); + return Ok(()); + } + }; + + if agent.state == AgentState::Offline { + return Ok(()); + } + + tracing::info!( + agent_id, + agent_uuid = %agent.uuid, + "Agent heartbeat timed out, retiring it" + ); + + // The agent is already off the queue, thus nothing tries again if this times out. + // `retire_agent` does all its writes in one transaction, so a timeout leaves them all done + // or none done. + match tokio::time::timeout( + DB_TIMEOUT, + retire_agent(&self.pool, agent_id, RetireCause::HeartbeatTimeout), + ) + .await + { + Ok(res) => res?, + Err(_) => tracing::warn!(agent_id, "Retiring the timed-out agent timed out"), + } + + Ok(()) + } + + pub async fn run(&mut self) { + tracing::info!("Agent heartbeat queue started"); + let mut timeout_duration = self.heartbeat_timeout; + loop { + tokio::select! { + biased; + _ = self.cancel_token.cancelled() => break, + op = self.rx.recv() => match op.ok() { + None => break, + Some(op) => { + self.handle_op(op); + timeout_duration = self.next_deadline(); + } + }, + _ = tokio::time::sleep(timeout_duration) => { + if let Err(e) = self.handle_timeout().await { + if self.cancel_token.is_cancelled() { + tracing::warn!("Agent timeout handling failed during shutdown: {e}"); + } else { + tracing::error!("Agent heartbeat timeout handling failed: {e}"); + } + } + timeout_duration = self.next_deadline(); + } + } + } + tracing::info!("Agent heartbeat queue stopped"); + } +} diff --git a/netmito/src/service/agent/mod.rs b/netmito/src/service/agent/mod.rs new file mode 100644 index 00000000..1635a49c --- /dev/null +++ b/netmito/src/service/agent/mod.rs @@ -0,0 +1,608 @@ +//! Agent registration, queries, heartbeats and shutdown. + +pub mod heartbeat; + +use std::collections::HashSet; + +use sea_orm::sea_query::extension::postgres::PgExpr; +use sea_orm::sea_query::{Alias, PgFunc, Query}; +use sea_orm::{prelude::*, FromQueryResult, QuerySelect, Set, TransactionTrait}; +use uuid::Uuid; + +use crate::config::InfraPool; +use crate::entity::{ + agents as Agent, group_agent as GroupAgent, groups as Group, machines as Machines, + role::{GroupAgentRole, UserGroupRole}, + state::AgentState, + task_suites as TaskSuites, user_group as UserGroup, users as User, +}; +use crate::error::{ApiError, AuthError, Error, Result}; +use crate::schema::{ + AgentHeartbeatReq, AgentHeartbeatResp, AgentInfo, AgentNotification, AgentShutdownOp, + AgentsQueryReq, AgentsQueryResp, CountQuery, RegisterAgentReq, RegisterAgentResp, +}; +use crate::service::auth::token::{generate_runner_token, TokenRole}; +use crate::ws::AgentWsRouter; + +use heartbeat::AgentHeartbeatOp; + +/// Register an agent. If the machine is already known, use its agent row again. +/// +/// `machines.machine_code` and `machines.agent_id` are both unique, thus one machine has one agent +/// row for all its life. For a known machine, this function updates that agent row: new tags, +/// labels, groups and metadata, and a new token. The caller must be an Admin of `admin_group`, and +/// must have at least the Write role in each group of `groups`. +/// +/// The request gives all the groups that can use the agent. Group access that is not in the request +/// is removed. +pub async fn user_register_agent( + user_id: i64, + pool: &InfraPool, + req: RegisterAgentReq, +) -> Result { + let now = TimeDateTimeWithTimeZone::now_utc(); + let tags: Vec = req.tags.into_iter().collect(); + let labels: Vec = req.labels.into_iter().collect(); + let groups: Vec = req.groups.into_iter().collect(); + let admin_group = req.admin_group.map(|g| g.trim().to_string()); + let machine_code = req.machine_code.trim().to_string(); + if machine_code.is_empty() { + return Err(Error::ApiError(ApiError::InvalidRequest( + "machine_code must not be empty".to_string(), + ))); + } + let metadata = req + .metadata + .as_ref() + .map(serde_json::to_value) + .transpose()?; + + let (agent_uuid, agent_id, reused, retired) = pool + .db + .transaction::<_, (Uuid, i64, bool, bool), Error>(move |txn| { + Box::pin(async move { + // Exactly one group has the Admin role on the agent, and only + // its Admins can shut the agent down. + let admin_group_name = match admin_group { + Some(name) => name, + None => { + User::Entity::find_by_id(user_id) + .one(txn) + .await? + .ok_or(Error::ApiError(ApiError::NotFound("User".to_string())))? + .username + } + }; + + let admin_group = Group::Entity::find() + .filter(Group::Column::GroupName.eq(admin_group_name.clone())) + .one(txn) + .await? + .ok_or(Error::ApiError(ApiError::NotFound(format!( + "Group {admin_group_name}" + ))))?; + let admin_membership = UserGroup::Entity::find() + .filter(UserGroup::Column::UserId.eq(user_id)) + .filter(UserGroup::Column::GroupId.eq(admin_group.id)) + .one(txn) + .await? + .ok_or(Error::AuthError(AuthError::PermissionDenied))?; + if admin_membership.role != UserGroupRole::Admin { + return Err(Error::AuthError(AuthError::PermissionDenied)); + } + + let mut group_ids = Vec::with_capacity(groups.len()); + for group_name in &groups { + let group = Group::Entity::find() + .filter(Group::Column::GroupName.eq(group_name.clone())) + .one(txn) + .await? + .ok_or(Error::ApiError(ApiError::NotFound(format!( + "Group {group_name}" + ))))?; + let user_group = UserGroup::Entity::find() + .filter(UserGroup::Column::UserId.eq(user_id)) + .filter(UserGroup::Column::GroupId.eq(group.id)) + .one(txn) + .await? + .ok_or(Error::AuthError(AuthError::PermissionDenied))?; + if user_group.role < UserGroupRole::Write { + return Err(Error::AuthError(AuthError::PermissionDenied)); + } + group_ids.push(group.id); + } + + let existing = Machines::Entity::find() + .filter(Machines::Column::MachineCode.eq(machine_code.clone())) + .one(txn) + .await?; + + let mut retired = false; + let (agent, reused) = match existing { + Some(machine) => { + let agent = Agent::Entity::find_by_id(machine.agent_id) + .one(txn) + .await? + .ok_or_else(|| { + Error::Custom(format!( + "Machine {machine_code} points at missing agent {}", + machine.agent_id + )) + })?; + // A new registration from the same machine means that + // the previous agent process has stopped, whatever its + // state shows. Retire it in this transaction, so that + // its work is not held until its heartbeat times out. + // This comes after the permission checks, so that a + // known machine code alone cannot retire an agent. + let agent_state = agent.state; + if agent_state != AgentState::Offline { + retire_agent_writes(txn, agent.id, now).await?; + retired = true; + } + let mut agent: Agent::ActiveModel = agent.into(); + agent.tags = Set(tags); + agent.labels = Set(labels); + agent.state = Set(AgentState::Idle); + agent.assigned_task_suite_id = Set(None); + agent.last_heartbeat = Set(now); + agent.updated_at = Set(now); + let agent = agent.update(txn).await?; + tracing::info!( + agent_uuid = %agent.uuid, + previous_state = %agent_state, + "Registering the known machine again with its agent row" + ); + + let mut machine: Machines::ActiveModel = machine.into(); + machine.metadata = Set(metadata); + machine.last_seen_at = Set(now); + machine.update(txn).await?; + + (agent, true) + } + None => { + let agent = Agent::ActiveModel { + uuid: Set(Uuid::new_v4()), + creator_id: Set(user_id), + tags: Set(tags), + labels: Set(labels), + state: Set(AgentState::Idle), + last_heartbeat: Set(now), + assigned_task_suite_id: Set(None), + created_at: Set(now), + updated_at: Set(now), + ..Default::default() + } + .insert(txn) + .await?; + + Machines::ActiveModel { + agent_id: Set(agent.id), + machine_code: Set(machine_code), + metadata: Set(metadata), + first_seen_at: Set(now), + last_seen_at: Set(now), + ..Default::default() + } + .insert(txn) + .await?; + + (agent, false) + } + }; + + GroupAgent::Entity::delete_many() + .filter(GroupAgent::Column::AgentId.eq(agent.id)) + .exec(txn) + .await?; + + // The admin group goes first. If it is also in `groups`, it + // keeps the Admin role and does not get a second Write row. + let mut granted = HashSet::with_capacity(group_ids.len() + 1); + let grants = std::iter::once((admin_group.id, GroupAgentRole::Admin)) + .chain(group_ids.into_iter().map(|id| (id, GroupAgentRole::Write))); + for (group_id, role) in grants { + if !granted.insert(group_id) { + continue; + } + GroupAgent::ActiveModel { + group_id: Set(group_id), + agent_id: Set(agent.id), + role: Set(role), + ..Default::default() + } + .insert(txn) + .await?; + } + + Ok((agent.uuid, agent.id, reused, retired)) + }) + }) + .await?; + + // This must come before the heartbeat below. It stops tracking the previous process, and the + // heartbeat then starts to track the new one. + if retired { + retire_agent_post_commit(pool, agent_id, agent_uuid, RetireCause::Reregistered); + } + + let token = generate_runner_token(agent_uuid.to_string(), 0, req.lifetime, TokenRole::Agent)?; + + let _ = pool + .agent_heartbeat_queue_tx + .send(AgentHeartbeatOp::Heartbeat(agent_id)); + + // Notifications queued for the previous process do not apply to this one. The router applies + // operations in order, thus the counter read below comes after the reset. + AgentWsRouter::reset(&pool.ws_router_tx, agent_uuid); + let notification_counter = AgentWsRouter::counter(&pool.ws_router_tx, agent_uuid) + .await + .unwrap_or_default(); + + Ok(RegisterAgentResp { + agent_uuid, + token, + notification_counter, + reused, + }) +} + +/// Query the agents of a group. The caller must have at least the Read role in the group. +pub async fn user_query_agents( + user_id: i64, + pool: &InfraPool, + mut query: AgentsQueryReq, +) -> Result { + if query.group_name.is_none() { + let user = User::Entity::find_by_id(user_id) + .one(&pool.db) + .await? + .ok_or(Error::ApiError(ApiError::NotFound("User".to_string())))?; + query.group_name = Some(user.username); + } + let group_name = query.group_name.clone().unwrap(); + + let group = Group::Entity::find() + .filter(Group::Column::GroupName.eq(&group_name)) + .one(&pool.db) + .await? + .ok_or(Error::ApiError(ApiError::NotFound(format!( + "Group {group_name}" + ))))?; + let authorized = UserGroup::Entity::find() + .filter(UserGroup::Column::UserId.eq(user_id)) + .filter(UserGroup::Column::GroupId.eq(group.id)) + .filter(UserGroup::Column::Role.gte(UserGroupRole::Read)) + .one(&pool.db) + .await? + .is_some(); + if !authorized { + return Err(Error::ApiError(ApiError::NotFound(format!( + "User doesn't have permission or group with name {group_name}" + )))); + } + + let mut stmt = Query::select(); + if query.count { + stmt.expr(Expr::col((Agent::Entity, Agent::Column::Uuid)).count()); + } else { + stmt.columns([ + (Agent::Entity, Agent::Column::Uuid), + (Agent::Entity, Agent::Column::Tags), + (Agent::Entity, Agent::Column::Labels), + (Agent::Entity, Agent::Column::State), + (Agent::Entity, Agent::Column::LastHeartbeat), + (Agent::Entity, Agent::Column::CreatedAt), + (Agent::Entity, Agent::Column::UpdatedAt), + ]) + .expr_as( + Expr::col((User::Entity, User::Column::Username)), + Alias::new("creator_username"), + ) + .expr_as( + Expr::col((TaskSuites::Entity, TaskSuites::Column::Uuid)), + Alias::new("assigned_suite_uuid"), + ) + .column((Machines::Entity, Machines::Column::MachineCode)) + .column((Machines::Entity, Machines::Column::Metadata)); + } + + stmt.from(Agent::Entity) + .join( + sea_orm::JoinType::Join, + GroupAgent::Entity, + Expr::col((GroupAgent::Entity, GroupAgent::Column::AgentId)) + .eq(Expr::col((Agent::Entity, Agent::Column::Id))), + ) + .join( + sea_orm::JoinType::Join, + User::Entity, + Expr::col((User::Entity, User::Column::Id)) + .eq(Expr::col((Agent::Entity, Agent::Column::CreatorId))), + ) + .join( + sea_orm::JoinType::LeftJoin, + Machines::Entity, + Expr::col((Machines::Entity, Machines::Column::AgentId)) + .eq(Expr::col((Agent::Entity, Agent::Column::Id))), + ) + .join( + sea_orm::JoinType::LeftJoin, + TaskSuites::Entity, + Expr::col((TaskSuites::Entity, TaskSuites::Column::Id)).eq(Expr::col(( + Agent::Entity, + Agent::Column::AssignedTaskSuiteId, + ))), + ) + .and_where(Expr::col((GroupAgent::Entity, GroupAgent::Column::GroupId)).eq(group.id)); + + if let Some(ref tags) = query.tags { + let tags: Vec = tags.iter().cloned().collect(); + stmt.and_where(Expr::col((Agent::Entity, Agent::Column::Tags)).contains(tags)); + } + if let Some(ref labels) = query.labels { + let labels: Vec = labels.iter().cloned().collect(); + stmt.and_where(Expr::col((Agent::Entity, Agent::Column::Labels)).contains(labels)); + } + if let Some(ref states) = query.states { + let states: Vec = states.iter().copied().collect(); + stmt.and_where(Expr::col((Agent::Entity, Agent::Column::State)).eq(PgFunc::any(states))); + } + if let Some(ref creator_username) = query.creator_username { + stmt.and_where( + Expr::col((User::Entity, User::Column::Username)).eq(creator_username.clone()), + ); + } + if let Some(limit) = query.limit { + stmt.limit(limit); + } + if let Some(offset) = query.offset { + stmt.offset(offset); + } + + let builder = pool.db.get_database_backend(); + if query.count { + let count = CountQuery::find_by_statement(builder.build(&stmt)) + .one(&pool.db) + .await? + .map(|c| c.count as u64) + .unwrap_or(0); + Ok(AgentsQueryResp { + count, + agents: vec![], + group_name, + }) + } else { + let agents = AgentInfo::find_by_statement(builder.build(&stmt)) + .all(&pool.db) + .await?; + Ok(AgentsQueryResp { + count: agents.len() as u64, + agents, + group_name, + }) + } +} + +/// Shut an agent down. The caller must be an Admin of the admin group of the agent. +/// +/// - `Graceful`: only tell the agent to stop. The agent finishes its current work and then retires +/// itself with `DELETE /agents`. Its heartbeat is still tracked, thus an agent that stops +/// responding still times out. +/// - `Force`: retire the agent now, then tell it to stop. +pub async fn user_shutdown_agent_by_uuid( + user_id: i64, + agent_uuid: Uuid, + op: AgentShutdownOp, + pool: &InfraPool, +) -> Result<()> { + let agent = Agent::Entity::find() + .join_rev(sea_orm::JoinType::Join, GroupAgent::Relation::Agents.def()) + .join(sea_orm::JoinType::Join, GroupAgent::Relation::Groups.def()) + .join(sea_orm::JoinType::Join, Group::Relation::UserGroup.def()) + .filter(Agent::Column::Uuid.eq(agent_uuid)) + .filter(GroupAgent::Column::Role.eq(GroupAgentRole::Admin)) + .filter(UserGroup::Column::UserId.eq(user_id)) + .filter(UserGroup::Column::Role.eq(UserGroupRole::Admin)) + .one(&pool.db) + .await? + .ok_or(Error::ApiError(ApiError::NotFound(format!( + "User doesn't have permission or agent with uuid {agent_uuid}" + ))))?; + + match op { + AgentShutdownOp::Force => { + retire_agent(pool, agent.id, RetireCause::ForceShutdown).await?; + } + AgentShutdownOp::Graceful => tracing::info!( + agent_uuid = %agent.uuid, + "Graceful shutdown requested; the agent stops after its current work" + ), + } + + // A retired agent cannot authenticate, thus it cannot read this from a heartbeat. It gets the + // notification only on this WebSocket that is already connected. If there is no websocket, the + // next heartbeat will return not authenticated and the agent will shut itself down. + AgentWsRouter::notify( + &pool.ws_router_tx, + agent.uuid, + AgentNotification::Shutdown { + graceful: matches!(op, AgentShutdownOp::Graceful), + }, + ); + + Ok(()) +} + +/// The reason to retire an agent. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RetireCause { + /// The agent did not send a heartbeat in time. + HeartbeatTimeout, + /// A user forced the agent to shut down. + ForceShutdown, + /// A new agent process registered from the same machine. + Reregistered, + /// The agent told the coordinator that it exits. + SelfExit, +} + +/// Retire an agent: mark it `Offline`, stop tracking its heartbeat, and remove its buffered +/// notifications. Its token stops working. +/// +/// This function does nothing for an agent that is already `Offline`, thus it is safe to call it +/// more than once. +pub async fn retire_agent(pool: &InfraPool, agent_id: i64, cause: RetireCause) -> Result<()> { + let now = TimeDateTimeWithTimeZone::now_utc(); + + let Some(agent) = Agent::Entity::find_by_id(agent_id).one(&pool.db).await? else { + tracing::debug!(agent_id, "Agent not found; nothing to retire"); + return Ok(()); + }; + if agent.state == AgentState::Offline { + return Ok(()); + } + let agent_uuid = agent.uuid; + + pool.db + .transaction::<_, (), Error>(move |txn| { + Box::pin(async move { retire_agent_writes(txn, agent_id, now).await }) + }) + .await?; + + retire_agent_post_commit(pool, agent_id, agent_uuid, cause); + Ok(()) +} + +/// The database writes of [`retire_agent`]. Registration also calls this in its own transaction, +/// because it retires the previous agent and updates the agent row in one step. +async fn retire_agent_writes( + txn: &C, + agent_id: i64, + now: TimeDateTimeWithTimeZone, +) -> Result<()> { + // TODO: when agents run jobs, end the in-flight jobs of the agent and put + // the tasks that it holds back to `Ready` here, in the same transaction. + Agent::Entity::update_many() + .col_expr(Agent::Column::State, Expr::value(AgentState::Offline)) + .col_expr(Agent::Column::AssignedTaskSuiteId, Expr::value(None::)) + .col_expr(Agent::Column::UpdatedAt, Expr::value(now)) + .filter(Agent::Column::Id.eq(agent_id)) + .exec(txn) + .await?; + Ok(()) +} + +/// The steps of [`retire_agent`] outside the database. They run after the commit, because a +/// rollback cannot undo them. +fn retire_agent_post_commit(pool: &InfraPool, agent_id: i64, agent_uuid: Uuid, cause: RetireCause) { + let _ = pool + .agent_heartbeat_queue_tx + .send(AgentHeartbeatOp::Remove(agent_id)); + AgentWsRouter::reset(&pool.ws_router_tx, agent_uuid); + tracing::info!(agent_id, %agent_uuid, ?cause, "Retired an agent"); +} + +/// Record a heartbeat and return the notifications that the agent has not acknowledged. +pub async fn agent_heartbeat( + agent_id: i64, + agent_uuid: Uuid, + pool: &InfraPool, + req: AgentHeartbeatReq, +) -> Result { + if req.state == AgentState::Offline { + return Err(Error::ApiError(ApiError::InvalidRequest( + "an agent exits with DELETE /agents, not with a heartbeat".to_string(), + ))); + } + let now = TimeDateTimeWithTimeZone::now_utc(); + + // The agent knows best what it does, thus its reported state replaces the recorded one. The + // filter on `Offline` prevents the heartbeat that is initiated from an agent that is in the process + // of retiring to re-activate the agent row. + let updated = Agent::Entity::update_many() + .col_expr(Agent::Column::State, Expr::value(req.state)) + .col_expr(Agent::Column::LastHeartbeat, Expr::value(now)) + .col_expr(Agent::Column::UpdatedAt, Expr::value(now)) + .filter(Agent::Column::Id.eq(agent_id)) + .filter(Agent::Column::State.ne(AgentState::Offline)) + .exec(&pool.db) + .await?; + if updated.rows_affected == 0 { + return Err(Error::AuthError(AuthError::WrongCredentials)); + } + + let _ = pool + .agent_heartbeat_queue_tx + .send(AgentHeartbeatOp::Heartbeat(agent_id)); + + // If the agent has processed a higher id than the coordinator has issued, the coordinator has + // restarted and its ids started again. Tell the agent the current counter and boot id. + if let Some(counter) = AgentWsRouter::counter(&pool.ws_router_tx, agent_uuid).await { + if req.last_notification_id > counter { + tracing::warn!( + agent_uuid = %agent_uuid, + agent_counter = req.last_notification_id, + coordinator_counter = counter, + "Notification counter mismatch; sending CounterSync" + ); + AgentWsRouter::notify( + &pool.ws_router_tx, + agent_uuid, + AgentNotification::CounterSync { + counter, + boot_id: pool.boot_uuid, + }, + ); + } + } + + let notifications = AgentWsRouter::pending_notifications( + &pool.ws_router_tx, + agent_uuid, + req.last_notification_id, + ) + .await; + + Ok(AgentHeartbeatResp { notifications }) +} + +/// At coordinator start, track the heartbeat of each agent that is not `Offline`, and send it a +/// `CounterSync` with the new boot id. +/// +/// An agent that stopped while the coordinator was down then times out. The `CounterSync` tells the +/// agent that the notification ids start again, so that it does not ignore new notifications with +/// low ids. +pub async fn restore_agents(pool: &InfraPool) -> Result<()> { + let agents = Agent::Entity::find() + .select_only() + .columns([Agent::Column::Id, Agent::Column::Uuid]) + .filter(Agent::Column::State.ne(AgentState::Offline)) + .into_tuple::<(i64, Uuid)>() + .all(&pool.db) + .await?; + tracing::info!( + agents = agents.len(), + boot_id = %pool.boot_uuid, + "Restoring agents after coordinator start" + ); + for (id, uuid) in agents { + if pool + .agent_heartbeat_queue_tx + .send(AgentHeartbeatOp::Heartbeat(id)) + .is_err() + { + return Err(Error::Custom("send agent heartbeat failed".to_string())); + } + AgentWsRouter::notify( + &pool.ws_router_tx, + uuid, + AgentNotification::CounterSync { + counter: 0, + boot_id: pool.boot_uuid, + }, + ); + } + Ok(()) +} diff --git a/netmito/src/service/auth/mod.rs b/netmito/src/service/auth/mod.rs index e73e6586..c623916e 100644 --- a/netmito/src/service/auth/mod.rs +++ b/netmito/src/service/auth/mod.rs @@ -18,7 +18,11 @@ use sea_orm::{entity::prelude::*, Set}; use crate::{ config::InfraPool, - entity::{state::UserState, users as User, workers as Worker}, + entity::{ + agents as Agent, + state::{AgentState, UserState}, + users as User, workers as Worker, + }, error::{ApiError, AuthError}, schema::{UserChangePasswordReq, UserLoginReq}, }; @@ -62,11 +66,18 @@ pub struct AuthWorker { pub uuid: Uuid, } +#[derive(Debug, Clone)] +pub struct AuthAgent { + pub id: i64, + pub uuid: Uuid, +} + /// The caller of a route open to more than one role. #[derive(Debug, Clone)] pub enum AuthIdentity { User(AuthUser), Worker(AuthWorker), + Agent(AuthAgent), } fn generate_auth_signature() -> i64 { @@ -320,9 +331,12 @@ pub async fn admin_change_password( Ok(()) } -/// Accept the token roles set in `ROLES` (a union of [`roles`] bits) and attach the -/// caller to the request: always as [`AuthIdentity`], and also as the role's own type -/// (`AuthUser`, `AuthWorker`) for handlers that serve a single role. +/// Check the token of the caller. `ROLES` is a set of [`roles`] bits. If the token +/// role is not in `ROLES`, reject the request. +/// +/// Add the caller to the request two times: as [`AuthIdentity`], and as the type of +/// its role (`AuthUser`, `AuthWorker` or `AuthAgent`). A handler for one role reads +/// the role type. A handler for more than one role reads `AuthIdentity`. pub async fn auth_middleware( State(pool): State, TypedHeader(Authorization(bearer)): TypedHeader>, @@ -349,8 +363,11 @@ pub async fn auth_middleware( req.extensions_mut().insert(auth_worker.clone()); AuthIdentity::Worker(auth_worker) } - // No agent tokens are issued yet. - TokenRole::Agent => return Err(AuthError::InvalidToken.into()), + TokenRole::Agent => { + let auth_agent = agent_auth(&pool.db, &claims).await?; + req.extensions_mut().insert(auth_agent.clone()); + AuthIdentity::Agent(auth_agent) + } }; req.extensions_mut().insert(identity); Ok(next.run(req).await) @@ -415,3 +432,23 @@ async fn worker_auth( uuid, }) } + +async fn agent_auth( + db: &DatabaseConnection, + claims: &TokenClaims<'_>, +) -> Result { + let uuid = Uuid::parse_str(&claims.sub).map_err(|_| AuthError::InvalidToken)?; + + let agent = Agent::Entity::find() + .filter(Agent::Column::Uuid.eq(uuid)) + .one(db) + .await + .map_err(|_| AuthError::WrongCredentials)? + .ok_or(AuthError::WrongCredentials)?; + // An agent row is never deleted, so a retired agent keeps its row. The + // Offline state is what makes its token stop working. + if agent.state == AgentState::Offline { + return Err(AuthError::WrongCredentials); + } + Ok(AuthAgent { id: agent.id, uuid }) +} diff --git a/netmito/src/service/auth/token.rs b/netmito/src/service/auth/token.rs index 3b8d8a31..e76048ad 100644 --- a/netmito/src/service/auth/token.rs +++ b/netmito/src/service/auth/token.rs @@ -77,11 +77,13 @@ where encode_token(&claims, encoding_key) } -/// Generate a worker token expiring after `lifetime`, or one that never expires if it is `None`. -pub fn generate_worker_token( +/// Generate a worker or agent token expiring after `lifetime`, or one that never expires if it +/// is `None`. +pub fn generate_runner_token( username: T, sign: i64, lifetime: Option, + role: TokenRole, ) -> crate::error::Result where T: AsRef, @@ -101,7 +103,7 @@ where sub: Cow::from(username.as_ref()), exp, sign, - role: TokenRole::Worker, + role, }; let encoding_key = crate::config::ENCODING_KEY diff --git a/netmito/src/service/mod.rs b/netmito/src/service/mod.rs index f3229faf..c479511f 100644 --- a/netmito/src/service/mod.rs +++ b/netmito/src/service/mod.rs @@ -20,6 +20,7 @@ //! row. `claim_candidates` locks only task rows with `FOR UPDATE SKIP LOCKED`, and no //! transaction waits for these locks. +pub mod agent; pub mod auth; pub mod group; pub mod s3; diff --git a/netmito/src/ws/connection.rs b/netmito/src/ws/connection.rs new file mode 100644 index 00000000..950bed45 --- /dev/null +++ b/netmito/src/ws/connection.rs @@ -0,0 +1,273 @@ +//! Per-agent notification sessions and the router actor that owns them. +//! +//! The service layer sends a [`RouterOp`] to [`AgentWsRouter`], and the router applies it. Each +//! agent has an [`AgentSession`]: a sequence counter, a buffer of notifications that are not +//! acknowledged, and the sender of its WebSocket while one is connected. +//! +//! A notification always goes into the buffer. If a socket is connected, the router also sends it +//! on the socket. The notification leaves the buffer only when the agent acknowledges it: on the +//! socket ([`RouterOp::AckBy`]), or with the `last_notification_id` of its next heartbeat. Thus a +//! session continues after its socket closes. A new socket gets the buffered notifications again, +//! and a heartbeat reads them. A session ends only with the agent process, at +//! [`RouterOp::ResetBuffer`]. + +use std::collections::{HashMap, VecDeque}; + +use axum::extract::ws::Message; +use crossfire::{AsyncRx, MTx, Tx}; +use speedy::Writable; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +use crate::schema::{AgentNotification, WsNotificationEvent}; + +#[derive(Debug)] +pub enum RouterOp { + /// A WebSocket connected. Attach its sender to the session of the agent. + Register { uuid: Uuid, sender: Tx }, + /// The WebSocket closed. The session and its buffer stay. + Unregister { uuid: Uuid }, + /// The process that had this uuid has stopped. Remove its buffered + /// notifications, because a `Shutdown` sent again to the next process + /// would stop the wrong process. The counter stays, thus the ids continue + /// to increase while this coordinator runs. + ResetBuffer { uuid: Uuid }, + /// The agent processed all notifications up to and including `id`. + AckBy { uuid: Uuid, id: u64 }, + /// Buffer a notification, and send it on the socket if one is connected. + Notify { + uuid: Uuid, + event: AgentNotification, + }, + /// For a heartbeat: acknowledge up to `ack_by_id`, then return all + /// notifications after it. The returned notifications stay in the buffer. + PendingNotifications { + uuid: Uuid, + ack_by_id: u64, + tx: Tx>, + }, + /// Read the sequence counter of the agent, `None` if it has no session. + GetCounter { uuid: Uuid, tx: Tx> }, +} + +struct AgentSession { + /// Sender of the connected WebSocket, if any. + sender: Option>, + /// The last sequence id given to a notification. + counter: u64, + /// Notifications that are not acknowledged, oldest first. + buffer: VecDeque, +} + +impl AgentSession { + fn new() -> Self { + Self { + sender: None, + counter: 0, + buffer: VecDeque::new(), + } + } + + fn push(&mut self, event: AgentNotification) -> WsNotificationEvent { + self.counter = self.counter.saturating_add(1); + let event = WsNotificationEvent { + id: self.counter, + event, + }; + self.buffer.push_back(event.clone()); + event + } + + fn ack(&mut self, ack_id: u64) { + while self.buffer.front().is_some_and(|e| e.id <= ack_id) { + self.buffer.pop_front(); + } + } + + fn pending_after(&self, after_id: u64) -> Vec { + self.buffer + .iter() + .filter(|e| e.id > after_id) + .cloned() + .collect() + } +} + +/// The actor that owns the notification sessions of all agents. +pub struct AgentWsRouter { + sessions: HashMap, + cancel_token: CancellationToken, + rx: AsyncRx, +} + +impl AgentWsRouter { + pub fn new(cancel_token: CancellationToken, rx: AsyncRx) -> Self { + Self { + sessions: HashMap::new(), + cancel_token, + rx, + } + } + + pub async fn run(&mut self) { + tracing::info!("Agent WebSocket router started"); + loop { + tokio::select! { + biased; + _ = self.cancel_token.cancelled() => break, + op = self.rx.recv() => match op.ok() { + None => break, + Some(op) => self.handle_op(op), + }, + } + } + tracing::info!("Agent WebSocket router stopped"); + } + + fn handle_op(&mut self, op: RouterOp) { + match op { + RouterOp::Register { uuid, sender } => { + let session = self.sessions.entry(uuid).or_insert_with(AgentSession::new); + session.sender = Some(sender); + // Send the buffered notifications again, so that the agent + // does not wait for its next heartbeat to get them. + let pending = session.pending_after(0); + if let Some(sender) = session.sender.as_ref() { + for event in &pending { + if !Self::push_to_socket(sender, event, uuid) { + break; + } + } + } + tracing::debug!( + agent_uuid = %uuid, + replayed = pending.len(), + sessions = self.sessions.len(), + "Agent WebSocket registered" + ); + } + RouterOp::Unregister { uuid } => { + if let Some(session) = self.sessions.get_mut(&uuid) { + session.sender = None; + } + tracing::debug!(agent_uuid = %uuid, "Agent WebSocket unregistered"); + } + RouterOp::ResetBuffer { uuid } => { + if let Some(session) = self.sessions.get_mut(&uuid) { + let dropped = session.buffer.len(); + session.buffer.clear(); + tracing::debug!( + agent_uuid = %uuid, + dropped, + "Dropped the notifications of a stopped agent process" + ); + } + } + RouterOp::AckBy { uuid, id } => { + if let Some(session) = self.sessions.get_mut(&uuid) { + session.ack(id); + } + } + RouterOp::Notify { uuid, event } => { + let session = self.sessions.entry(uuid).or_insert_with(AgentSession::new); + let event = session.push(event); + if let Some(sender) = session.sender.as_ref() { + Self::push_to_socket(sender, &event, uuid); + } + } + RouterOp::PendingNotifications { + uuid, + ack_by_id, + tx, + } => { + let pending = match self.sessions.get_mut(&uuid) { + Some(session) => { + session.ack(ack_by_id); + session.pending_after(ack_by_id) + } + None => Vec::new(), + }; + let _ = tx.send(pending); + } + RouterOp::GetCounter { uuid, tx } => { + let _ = tx.send(self.sessions.get(&uuid).map(|s| s.counter)); + } + } + } + + /// Return false if the socket did not accept the notification. The + /// notification stays in the buffer for the next heartbeat. + fn push_to_socket(sender: &Tx, event: &WsNotificationEvent, uuid: Uuid) -> bool { + let payload = match event.write_to_vec() { + Ok(payload) => payload, + Err(e) => { + tracing::error!(agent_uuid = %uuid, "Failed to serialize notification: {e}"); + return false; + } + }; + match sender.send(Message::Binary(payload.into())) { + Ok(()) => true, + Err(e) => { + tracing::debug!( + agent_uuid = %uuid, + notification_id = event.id, + "Notification not sent on WebSocket ({e}); kept for heartbeat" + ); + false + } + } + } + + pub fn register(tx: &MTx, uuid: Uuid, sender: Tx) { + let _ = tx.send(RouterOp::Register { uuid, sender }); + } + + pub fn unregister(tx: &MTx, uuid: Uuid) { + let _ = tx.send(RouterOp::Unregister { uuid }); + } + + pub fn reset(tx: &MTx, uuid: Uuid) { + let _ = tx.send(RouterOp::ResetBuffer { uuid }); + } + + pub fn ack(tx: &MTx, uuid: Uuid, id: u64) { + let _ = tx.send(RouterOp::AckBy { uuid, id }); + } + + /// Queue a notification for an agent. The router stops only during + /// shutdown, thus a failed send is logged and not returned to the request + /// handler. + pub fn notify(tx: &MTx, uuid: Uuid, event: AgentNotification) { + if tx.send(RouterOp::Notify { uuid, event }).is_err() { + tracing::debug!(agent_uuid = %uuid, "Notification dropped: WebSocket router is down"); + } + } + + /// Acknowledge up to `ack_by_id`, then read the notifications after it. + pub async fn pending_notifications( + tx: &MTx, + uuid: Uuid, + ack_by_id: u64, + ) -> Vec { + let (resp_tx, resp_rx) = crossfire::spsc::unbounded_async::>(); + if tx + .send(RouterOp::PendingNotifications { + uuid, + ack_by_id, + tx: resp_tx, + }) + .is_err() + { + return Vec::new(); + } + resp_rx.recv().await.unwrap_or_default() + } + + pub async fn counter(tx: &MTx, uuid: Uuid) -> Option { + let (resp_tx, resp_rx) = crossfire::spsc::unbounded_async::>(); + if tx.send(RouterOp::GetCounter { uuid, tx: resp_tx }).is_err() { + return None; + } + resp_rx.recv().await.ok().flatten() + } +} diff --git a/netmito/src/ws/handler.rs b/netmito/src/ws/handler.rs new file mode 100644 index 00000000..2e89677a --- /dev/null +++ b/netmito/src/ws/handler.rs @@ -0,0 +1,117 @@ +//! `GET /ws/agents`: the notification socket of an agent. +//! +//! The auth middleware checks the agent token before the upgrade. After the +//! upgrade, one task moves notifications from the router to the socket and +//! acknowledgements from the socket to the router. + +use axum::{ + extract::{ + ws::{Message, WebSocket}, + State, WebSocketUpgrade, + }, + response::IntoResponse, + Extension, +}; +use futures::{SinkExt, StreamExt}; +use speedy::Readable; +use uuid::Uuid; + +use crate::{ + config::InfraPool, schema::AgentWsMessage, service::auth::AuthAgent, + ws::connection::AgentWsRouter, +}; + +pub async fn websocket_handler( + ws: WebSocketUpgrade, + State(pool): State, + Extension(agent): Extension, +) -> impl IntoResponse { + let agent_uuid = agent.uuid; + tracing::debug!(agent_uuid = %agent_uuid, "Agent WebSocket upgrade accepted"); + ws.on_upgrade(move |socket| handle_agent_socket(socket, agent_uuid, pool)) +} + +async fn handle_agent_socket(socket: WebSocket, agent_uuid: Uuid, pool: InfraPool) { + let (mut sender, mut receiver) = socket.split(); + let (tx, rx) = crossfire::spsc::unbounded_async::(); + + AgentWsRouter::register(&pool.ws_router_tx, agent_uuid, tx); + + // Proxies between the agent and the coordinator can close a socket that + // has no traffic. The pings prevent this. + let mut keepalive = tokio::time::interval(pool.ws_keepalive_interval); + keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + // The first tick completes immediately. + keepalive.tick().await; + + loop { + tokio::select! { + incoming = receiver.next() => match incoming { + Some(Ok(msg)) => { + if !handle_ws_message(agent_uuid, &pool, msg) { + break; + } + } + Some(Err(e)) => { + tracing::debug!(agent_uuid = %agent_uuid, "Agent WebSocket error: {e}"); + break; + } + None => break, + }, + + outgoing = rx.recv() => match outgoing { + Ok(msg) => { + if let Err(e) = sender.send(msg).await { + tracing::debug!(agent_uuid = %agent_uuid, "Failed to send to agent: {e}"); + break; + } + } + Err(_) => break, + }, + + _ = keepalive.tick() => { + if let Err(e) = sender.send(Message::Ping(Vec::new().into())).await { + tracing::debug!(agent_uuid = %agent_uuid, "Keepalive ping failed: {e}"); + break; + } + } + } + } + + AgentWsRouter::unregister(&pool.ws_router_tx, agent_uuid); + tracing::debug!(agent_uuid = %agent_uuid, "Agent WebSocket closed"); +} + +/// Return false if the connection must close. +fn handle_ws_message(agent_uuid: Uuid, pool: &InfraPool, msg: Message) -> bool { + match msg { + Message::Binary(bytes) => match AgentWsMessage::read_from_buffer(&bytes) { + Ok(msg) => handle_agent_message(agent_uuid, msg, pool), + Err(e) => { + tracing::debug!(agent_uuid = %agent_uuid, "Unreadable agent message: {e}"); + } + }, + Message::Close(frame) => { + tracing::debug!(agent_uuid = %agent_uuid, ?frame, "Agent closed the WebSocket"); + return false; + } + // axum replies to pings. Agent messages are `speedy` binary frames. + Message::Ping(_) | Message::Pong(_) => {} + Message::Text(text) => { + tracing::debug!(agent_uuid = %agent_uuid, %text, "Ignoring an unexpected text frame"); + } + } + true +} + +fn handle_agent_message(agent_uuid: Uuid, message: AgentWsMessage, pool: &InfraPool) { + match message { + AgentWsMessage::Ack { notification_id } => { + AgentWsRouter::ack(&pool.ws_router_tx, agent_uuid, notification_id); + } + AgentWsMessage::Pong { client_time } => { + let latency = time::OffsetDateTime::now_utc().unix_timestamp() - client_time; + tracing::trace!(agent_uuid = %agent_uuid, latency_secs = latency, "Agent pong"); + } + } +} diff --git a/netmito/src/ws/mod.rs b/netmito/src/ws/mod.rs new file mode 100644 index 00000000..48030ccb --- /dev/null +++ b/netmito/src/ws/mod.rs @@ -0,0 +1,12 @@ +//! Notifications from the coordinator to agents on a WebSocket. +//! +//! The WebSocket only makes delivery faster. The router keeps each notification in a per-agent +//! buffer, and the next heartbeat response also returns it. Thus an agent without a connection +//! still gets all notifications. See [`connection::AgentWsRouter`] for the buffer and the +//! acknowledgements. + +pub mod connection; +pub mod handler; + +pub use connection::{AgentWsRouter, RouterOp}; +pub use handler::websocket_handler; diff --git a/openapi.yaml b/openapi.yaml index ded4a208..22d9a2da 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -1868,6 +1868,184 @@ paths: "404": $ref: "#/components/responses/NotFound" + /agents: + post: + summary: Register an agent + description: >- + Register the agent of a machine and issue its token. One machine has one + agent record for all its life, because `machine_code` is unique. If the + machine is already known, the coordinator uses its agent record again: + it replaces the tags, labels, groups and metadata, and it retires the + previous agent process. The caller must have the Admin role in + `admin_group`, and at least the Write role in each group of `groups`. + The request gives all the groups that can use the agent. Group access + that is not in the request is removed. + operationId: registerAgent + tags: + - Agents + security: + - bearerAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/RegisterAgentReq" + responses: + "200": + description: Agent registered + content: + application/json: + schema: + $ref: "#/components/schemas/RegisterAgentResp" + "400": + $ref: "#/components/responses/BadRequest" + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "404": + $ref: "#/components/responses/NotFound" + delete: + summary: Report that the agent exits + description: >- + The agent calls this when it exits. The coordinator marks the agent + Offline at once, and does not wait for its heartbeat to time out. The + agent record stays. After this call, the agent token does not work. + A second call has no effect. + operationId: agentExit + tags: + - Agents + security: + - agentAuth: [] + responses: + "200": + description: Agent marked Offline + "401": + $ref: "#/components/responses/Unauthorized" + + /agents/query: + post: + summary: Query agents + description: >- + List the agents of one group, with optional filters. The default group + is the personal group of the caller. The caller must have at least the + Read role in the group. + operationId: queryAgents + tags: + - Agents + security: + - bearerAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/AgentsQueryReq" + responses: + "200": + description: List of agents + content: + application/json: + schema: + $ref: "#/components/schemas/AgentsQueryResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /agents/{uuid}: + delete: + summary: Shut an agent down + description: >- + Stop an agent. The agent record stays. The caller must have the Admin + role in the admin group of the agent. If the caller does not, the + response is 404, the same as for an unknown uuid. + operationId: shutdownAgent + tags: + - Agents + security: + - bearerAuth: [] + parameters: + - name: uuid + in: path + required: true + schema: + type: string + format: uuid + - name: op + in: query + required: false + schema: + type: string + enum: + - graceful + - force + description: | + Shutdown mode: + - graceful (default): tell the agent to stop. The agent finishes its + current work and then marks itself Offline + - force: mark the agent Offline now, then tell it to stop + responses: + "200": + description: Shutdown requested + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /agents/heartbeat: + post: + summary: Agent heartbeat + description: >- + Record that the agent is alive, and return all notifications that the + agent has not acknowledged. `last_notification_id` acknowledges all + notifications up to that id. The coordinator returns the notifications + after it again until the agent acknowledges them. If an agent does not + send a heartbeat in time, the coordinator marks it Offline. + operationId: agentHeartbeat + tags: + - Agents + security: + - agentAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/AgentHeartbeatReq" + responses: + "200": + description: Notifications not acknowledged + content: + application/json: + schema: + $ref: "#/components/schemas/AgentHeartbeatResp" + "400": + $ref: "#/components/responses/BadRequest" + "401": + $ref: "#/components/responses/Unauthorized" + + /ws/agents: + get: + summary: Agent notification socket + description: >- + A WebSocket that sends notifications from the coordinator to the agent + as `speedy`-encoded binary `WsNotificationEvent` frames. The agent + acknowledges them with `AgentWsMessage` frames. The socket only makes + delivery faster: the next heartbeat also returns each notification, thus + an agent without a socket still gets all of them. + operationId: agentWebSocket + tags: + - Agents + security: + - agentAuth: [] + responses: + "101": + description: Switching protocols + "401": + $ref: "#/components/responses/Unauthorized" + /admin/users: post: summary: Create user (Admin) @@ -2195,6 +2373,11 @@ components: scheme: bearer bearerFormat: JWT description: JWT token obtained when registering worker + agentAuth: + type: http + scheme: bearer + bearerFormat: JWT + description: JWT token obtained when registering an agent responses: BadRequest: @@ -4834,6 +5017,256 @@ components: type: boolean description: Whether a manual override existed and was cleared (false = nothing to reset) + AgentState: + type: string + enum: + - Idle + - Provisioning + - Executing + - Cleaning + - Offline + + AgentMetadata: + type: object + required: + - version + - long_version + properties: + version: + type: string + long_version: + type: string + + RegisterAgentReq: + type: object + required: + - machine_code + properties: + tags: + type: array + items: + type: string + description: Tags that a suite must match to run on this agent + labels: + type: array + items: + type: string + description: Labels for queries and filters + admin_group: + type: string + description: >- + The group that gets the Admin role on the agent. Only an Admin of + this group can shut the agent down, and the caller must be one. The + default is the personal group of the caller. + groups: + type: array + items: + type: string + description: Groups that get the Write role on the agent + lifetime: + type: string + description: >- + The lifetime of the agent token, as a human duration. If it is not + given, the token never expires. + example: 30d + machine_code: + type: string + description: >- + A stable identifier of the machine. A restarted agent must send the + same value to use its agent record again. + metadata: + $ref: "#/components/schemas/AgentMetadata" + + RegisterAgentResp: + type: object + required: + - agent_uuid + - token + - notification_counter + - reused + properties: + agent_uuid: + type: string + format: uuid + token: + type: string + description: The JWT of the agent, for all agent-authenticated calls + notification_counter: + type: integer + format: uint64 + description: >- + The last notification id issued to this agent. The agent counts its + notifications from this value. + reused: + type: boolean + description: True if the machine was already known and its agent record is used again + + AgentInfo: + type: object + required: + - uuid + - creator_username + - tags + - labels + - state + - last_heartbeat + - created_at + - updated_at + properties: + uuid: + type: string + format: uuid + creator_username: + type: string + tags: + type: array + items: + type: string + labels: + type: array + items: + type: string + state: + $ref: "#/components/schemas/AgentState" + last_heartbeat: + type: string + format: date-time + assigned_suite_uuid: + type: string + format: uuid + created_at: + type: string + format: date-time + updated_at: + type: string + format: date-time + machine_code: + type: string + metadata: + type: object + + AgentsQueryReq: + type: object + required: + - count + properties: + group_name: + type: string + description: The default is the personal group of the caller + tags: + type: array + items: + type: string + labels: + type: array + items: + type: string + states: + type: array + items: + $ref: "#/components/schemas/AgentState" + creator_username: + type: string + limit: + type: integer + format: uint64 + offset: + type: integer + format: uint64 + count: + type: boolean + description: Return only the number of matching agents + + AgentsQueryResp: + type: object + required: + - count + - agents + - group_name + properties: + count: + type: integer + format: uint64 + agents: + type: array + items: + $ref: "#/components/schemas/AgentInfo" + group_name: + type: string + + AgentHeartbeatReq: + type: object + required: + - state + properties: + state: + $ref: "#/components/schemas/AgentState" + last_notification_id: + type: integer + format: uint64 + description: >- + The highest notification id the agent has processed. This + acknowledges all notifications up to this id. + + AgentHeartbeatResp: + type: object + required: + - notifications + properties: + notifications: + type: array + items: + $ref: "#/components/schemas/WsNotificationEvent" + + AgentNotification: + type: object + required: + - type + discriminator: + propertyName: type + properties: + type: + type: string + enum: + - shutdown + - counter_sync + description: | + A message from the coordinator that tells an agent to do something: + - shutdown: shut down; with `graceful`, finish the current work first + - counter_sync: the coordinator restarted and its notification ids + start again from `counter`; `boot_id` identifies the new process + + WsNotificationEvent: + type: object + required: + - id + - event + properties: + id: + type: integer + format: uint64 + description: Per-agent sequence id. It increases by one for each notification + event: + $ref: "#/components/schemas/AgentNotification" + + AgentWsMessage: + type: object + required: + - type + discriminator: + propertyName: type + properties: + type: + type: string + enum: + - ack + - pong + description: | + A message from an agent to the coordinator: + - ack: the agent processed all notifications up to and including + `notification_id` + - pong: the reply to a ping, with the agent clock in `client_time` + security: - bearerAuth: [] @@ -4850,6 +5283,8 @@ tags: description: Task submission and monitoring - name: Suites description: Task suite lifecycle and manual agent assignment + - name: Agents + description: Agent registration, heartbeat and shutdown - name: Admin description: Administrative operations - name: Health