diff --git a/Cargo.lock b/Cargo.lock index 390ddd2b..00b1fcd1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -548,6 +548,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8" dependencies = [ "axum-core", + "base64", "bytes", "form_urlencoded", "futures-util", @@ -566,8 +567,10 @@ dependencies = [ "serde_json", "serde_path_to_error", "serde_urlencoded", + "sha1", "sync_wrapper", "tokio", + "tokio-tungstenite", "tower", "tower-layer", "tower-service", @@ -761,7 +764,7 @@ dependencies = [ "iana-time-zone", "num-traits", "serde", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -1147,6 +1150,12 @@ dependencies = [ "syn", ] +[[package]] +name = "data-encoding" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4ae5f15dda3c708c0ade84bfee31ccab44a3da4f88015ed22f63732abe300c8" + [[package]] name = "der" version = "0.6.1" @@ -2408,6 +2417,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" @@ -2456,6 +2474,7 @@ dependencies = [ "async-compression", "async-once-cell", "async-stream", + "async-trait", "aws-sdk-s3", "axum", "axum-extra", @@ -2489,10 +2508,12 @@ dependencies = [ "sea-orm-migration", "serde", "serde_json", + "speedy", "thiserror 2.0.18", "time", "tokio", "tokio-tar", + "tokio-tungstenite", "tokio-util", "tower-http", "tracing", @@ -2694,7 +2715,7 @@ dependencies = [ "libc", "redox_syscall 0.5.18", "smallvec", - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -3631,9 +3652,9 @@ dependencies = [ [[package]] name = "security-framework" -version = "3.7.0" +version = "3.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +checksum = "80fb1d92c5028aa318b4b8bd7302a5bfcf48be96a37fc6fc790f806b0004ee0c" dependencies = [ "bitflags 2.11.0", "core-foundation 0.10.1", @@ -3644,9 +3665,9 @@ dependencies = [ [[package]] name = "security-framework-sys" -version = "2.17.0" +version = "2.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +checksum = "49db231d56a190491cb4aeda9527f1ad45345af50b0851622a7adb8c03b01c32" dependencies = [ "core-foundation-sys", "libc", @@ -3906,6 +3927,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" @@ -4429,6 +4472,22 @@ 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", + "rustls 0.23.37", + "rustls-native-certs", + "rustls-pki-types", + "tokio", + "tokio-rustls 0.26.4", + "tungstenite", +] + [[package]] name = "tokio-util" version = "0.7.18" @@ -4612,6 +4671,25 @@ 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", + "rustls 0.23.37", + "rustls-pki-types", + "sha1", + "thiserror 2.0.18", + "utf-8", +] + [[package]] name = "typenum" version = "1.19.0" @@ -4729,6 +4807,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" @@ -5009,7 +5093,7 @@ checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" dependencies = [ "windows-implement", "windows-interface", - "windows-link", + "windows-link 0.2.1", "windows-result", "windows-strings", ] @@ -5036,6 +5120,12 @@ dependencies = [ "syn", ] +[[package]] +name = "windows-link" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" + [[package]] name = "windows-link" version = "0.2.1" @@ -5048,7 +5138,7 @@ version = "0.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" dependencies = [ - "windows-link", + "windows-link 0.2.1", "windows-result", "windows-strings", ] @@ -5059,7 +5149,7 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" dependencies = [ - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -5068,7 +5158,7 @@ version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" dependencies = [ - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -5104,7 +5194,7 @@ version = "0.60.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" dependencies = [ - "windows-targets 0.53.5", + "windows-targets 0.53.3", ] [[package]] @@ -5113,7 +5203,7 @@ version = "0.61.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" dependencies = [ - "windows-link", + "windows-link 0.2.1", ] [[package]] @@ -5149,19 +5239,19 @@ dependencies = [ [[package]] name = "windows-targets" -version = "0.53.5" +version = "0.53.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +checksum = "d5fe6031c4041849d7c496a8ded650796e7b6ecc19df1a431c1a363342e5dc91" dependencies = [ - "windows-link", - "windows_aarch64_gnullvm 0.53.1", - "windows_aarch64_msvc 0.53.1", - "windows_i686_gnu 0.53.1", - "windows_i686_gnullvm 0.53.1", - "windows_i686_msvc 0.53.1", - "windows_x86_64_gnu 0.53.1", - "windows_x86_64_gnullvm 0.53.1", - "windows_x86_64_msvc 0.53.1", + "windows-link 0.1.3", + "windows_aarch64_gnullvm 0.53.0", + "windows_aarch64_msvc 0.53.0", + "windows_i686_gnu 0.53.0", + "windows_i686_gnullvm 0.53.0", + "windows_i686_msvc 0.53.0", + "windows_x86_64_gnu 0.53.0", + "windows_x86_64_gnullvm 0.53.0", + "windows_x86_64_msvc 0.53.0", ] [[package]] @@ -5178,9 +5268,9 @@ checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" [[package]] name = "windows_aarch64_gnullvm" -version = "0.53.1" +version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" +checksum = "86b8d5f90ddd19cb4a147a5fa63ca848db3df085e25fee3cc10b39b6eebae764" [[package]] name = "windows_aarch64_msvc" @@ -5196,9 +5286,9 @@ checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" [[package]] name = "windows_aarch64_msvc" -version = "0.53.1" +version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" +checksum = "c7651a1f62a11b8cbd5e0d42526e55f2c99886c77e007179efff86c2b137e66c" [[package]] name = "windows_i686_gnu" @@ -5214,9 +5304,9 @@ checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" [[package]] name = "windows_i686_gnu" -version = "0.53.1" +version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" +checksum = "c1dc67659d35f387f5f6c479dc4e28f1d4bb90ddd1a5d3da2e5d97b42d6272c3" [[package]] name = "windows_i686_gnullvm" @@ -5226,9 +5316,9 @@ checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" [[package]] name = "windows_i686_gnullvm" -version = "0.53.1" +version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" +checksum = "9ce6ccbdedbf6d6354471319e781c0dfef054c81fbc7cf83f338a4296c0cae11" [[package]] name = "windows_i686_msvc" @@ -5244,9 +5334,9 @@ checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" [[package]] name = "windows_i686_msvc" -version = "0.53.1" +version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" +checksum = "581fee95406bb13382d2f65cd4a908ca7b1e4c2f1917f143ba16efe98a589b5d" [[package]] name = "windows_x86_64_gnu" @@ -5262,9 +5352,9 @@ checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" [[package]] name = "windows_x86_64_gnu" -version = "0.53.1" +version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" +checksum = "2e55b5ac9ea33f2fc1716d1742db15574fd6fc8dadc51caab1c16a3d3b4190ba" [[package]] name = "windows_x86_64_gnullvm" @@ -5280,9 +5370,9 @@ checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" [[package]] name = "windows_x86_64_gnullvm" -version = "0.53.1" +version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" +checksum = "0a6e035dd0599267ce1ee132e51c27dd29437f63325753051e71dd9e42406c57" [[package]] name = "windows_x86_64_msvc" @@ -5298,9 +5388,9 @@ checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" [[package]] name = "windows_x86_64_msvc" -version = "0.53.1" +version = "0.53.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" +checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486" [[package]] name = "winnow" diff --git a/config.example.toml b/config.example.toml index b6c4b802..2ee1da36 100644 --- a/config.example.toml +++ b/config.example.toml @@ -24,6 +24,8 @@ access_token_public_path = "public.pem" # chosen by the Worker itself, via its own lifetime setting. access_token_expires_in = "7d" heartbeat_timeout = "600s" +# ws_keepalive_interval is the keepalive cadence on agent notification sockets +ws_keepalive_interval = "30s" file_log = false # log_path is not set. It will use the default rolling log file path if file_log is set to true @@ -47,6 +49,42 @@ file_log = false # - If shared_log is enabled and log_path is not set, it will use workers.log in cache directory # - If shared_log is disabled and log_path is not set, it will use {worker_uuid}.log in cache directory +[agent] +coordinator_addr = "http://127.0.0.1:5000" +heartbeat_interval = "1m" +# connect_retry_interval is the gap between retries of a coordinator call that +# could not connect. +connect_retry_interval = "30s" +lifetime = "30d" +# idle_poll_interval is not set. When set, an agent holding no suite asks the +# coordinator for one that often, instead of waiting to be told about one. +# Pickup is signal-driven without it: a notification, or the nudge that every +# heartbeat answer carries — so with no_ws and a slow heartbeat, waiting for +# work costs up to a whole heartbeat_interval. Each poll costs the coordinator +# a lookup, so keep it off unless the latency matters. +# ws_reconnect_interval is how long to wait before reopening a notification +# WebSocket that dropped. Ignored when no_ws is set. +ws_reconnect_interval = "5s" +# no_ws is not set, default to false. When enabled the agent never opens the +# notification WebSocket and takes its notifications from heartbeat answers +# alone. Nothing is lost — the coordinator buffers each notification until it +# is acknowledged — only delivery latency changes. For hosts where the socket +# cannot be opened or held, a proxy in the way being the usual reason. +# credential_path is not set +# user is not set +# password is not set +# admin_group is not set, default to the registering user's own group. It is the +# one group that can shut the agent down, and the user must be its Admin. +# groups are not set, default to the user's group. Each listed group gains Write +# access to the agent. Registration rewrites group access outright, so a group +# left out of a re-registration loses the access it had. +# tags are not set. An agent is eligible for a suite when it carries every tag +# the suite asks for. +# labels are not set. Labels are for querying only, never for matching. +# machine_code is not set. It is resolved from the cached value, then +# /etc/machine-id, then a generated UUID. It keys the agent's record, so a +# machine that comes back with a different code strands its old record. + [client] # user = "mitosis_admin" # password = "mitosis_admin" diff --git a/netmito/Cargo.toml b/netmito/Cargo.toml index dd675e1b..45fce463 100644 --- a/netmito/Cargo.toml +++ b/netmito/Cargo.toml @@ -17,8 +17,9 @@ categories.workspace = true argon2 = { version = "0.5.3", features = ["std"] } async-compression = { version = "0.4.32", features = ["gzip", "tokio"] } async-stream = "0.3.6" +async-trait = "0.1.89" 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 } @@ -64,10 +65,12 @@ 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 } tokio-tar = "0.3.1" +tokio-tungstenite = { version = "0.28", features = ["rustls-tls-native-roots"] } tokio-util = { version = "0.7.16", features = ["rt"] } tower-http = { version = "0.6.6", features = ["cors", "catch-panic"] } tracing = { workspace = true } diff --git a/netmito/src/agent.rs b/netmito/src/agent.rs new file mode 100644 index 00000000..d44b5e86 --- /dev/null +++ b/netmito/src/agent.rs @@ -0,0 +1,1560 @@ +//! The agent: an orchestrator process that claims task suites from the +//! coordinator and runs them. +//! +//! ```text +//! register ─▶ ws connect (unless --no-ws) ─┐ +//! heartbeat every N ───────────────────────┼─▶ main loop ─▶ (idle + work) ─▶ accept suite +//! idle poll every N (if configured) ───────┘ │ +//! ▼ +//! complete ◀─ cleanup hook ◀─ cleanup ◀─ tasks ◀─ start ◀─ provision hook +//! ``` +//! +//! The main loop only ever services heartbeats and notifications; a claimed +//! suite runs in a spawned [`SuiteRunner`] so neither starves the other. + +use std::path::PathBuf; +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use futures::{SinkExt, StreamExt}; +use reqwest::StatusCode; +use speedy::{Readable, Writable}; +use tokio::sync::mpsc; +use tokio::task::JoinHandle; +use tokio_tungstenite::tungstenite::Message as WsMessage; +use tokio_util::sync::CancellationToken; +use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; +use url::Url; +use uuid::Uuid; + +use crate::config::{AgentConfig, AgentConfigCli}; +use crate::entity::content::ArtifactContentType; +use crate::entity::hook_tasks::HookType; +use crate::entity::state::{AgentState, TaskExecState}; +use crate::error::{self, Result}; +use crate::executor::{execute, reset_workspace, ExecClient, Executor, UploadTarget}; +use crate::schema::*; +use crate::service::auth::{ + cred::get_user_credential, credential_guard::CredentialGuard, get_and_prompt_username, +}; + +/// How many task slots a job runs at a time. A slot is a working directory plus +/// an [`Executor`]; two concurrent tasks must never share one, since the +/// process's `result/` directory is what becomes its artifact. +/// +/// This is the suite's `FixedWorkers.worker_count` — the plan asks for that many +/// workers, and the agent gives it that many slots. A count of zero would leave +/// the job with nothing to drain it, so it is floored at one. +fn task_slots(suite: &TaskSuiteSpec) -> usize { + match suite.worker_schedule { + WorkerSchedulePlan::FixedWorkers { worker_count, .. } => (worker_count as usize).max(1), + } +} + +/// How often a poll-based `watch` re-asks the coordinator. +const WATCH_POLL_INTERVAL: Duration = Duration::from_secs(30); + +/// How often a job with nothing to do re-checks its suite for late work. This is +/// the latency a task submitted into a warm job pays before it starts. +const HOLD_POLL_INTERVAL: Duration = Duration::from_secs(1); + +pub struct MitoAgent; + +/// What the agent should claim next, recorded by a notification and consumed in +/// the idle branch of the main loop. Distinct from "nothing pending", so an idle +/// agent never polls unless the coordinator signalled work. +#[derive(Debug, Clone, Copy)] +enum PendingSuite { + /// Work exists but was not named — take whatever the coordinator offers. + Any, + /// A specific suite was named; target it directly. + Specific(Uuid), +} + +struct AgentClient { + coordinator_addr: Url, + token: String, + /// Highest notification id processed; echoed on every heartbeat + notification_counter: u64, + /// Which coordinator boot our counter belongs to. A different one means the + /// coordinator restarted and our sequence is meaningless. + coordinator_boot_id: Option, + state: AgentState, + assigned_suite_uuid: Option, + pending_suite: Option, + heartbeat_interval: Duration, + connect_retry_interval: Duration, + /// How often an idle agent asks for a suite unprompted; `None` waits for a + /// notification instead. + idle_poll_interval: Option, + /// How long to wait before reopening a notification socket that dropped. + ws_reconnect_interval: Duration, + /// Skip the notification socket and take everything from the heartbeat. + no_ws: bool, + /// Root of this agent's working directories; one subtree per job. + cache_path: PathBuf, + http_client: reqwest::Client, + /// Cancels the running suite (notification-driven); `None` while idle. + job_token: Option, + /// Releases a *held* suite without aborting its wind-down; `None` while idle. + drain_token: Option, + /// The suite in flight. Resolves to `complete`'s answer to "is there more + /// for this agent". + current_run: Option>, + /// Cancelling this exits the main loop. + shutdown_token: CancellationToken, + /// Set by a graceful `Shutdown` notification or `--run-once`: finish the + /// current suite, then stop instead of taking another. + stop_after_current: bool, + run_once: bool, +} + +impl MitoAgent { + pub async fn main(cli: AgentConfigCli) { + tracing_subscriber::registry() + .with( + tracing_subscriber::EnvFilter::try_from_default_env() + .unwrap_or_else(|_| "netmito=info".into()), + ) + .with(tracing_subscriber::fmt::layer()) + .init(); + + let run_once = cli.run_once; + match AgentConfig::new(&cli) { + Ok(config) => { + if let Err(e) = Self::run(config, run_once).await { + tracing::error!("Agent failed: {e}"); + std::process::exit(1); + } + } + Err(e) => { + tracing::error!("{e}"); + std::process::exit(1); + } + } + } + + /// Register with the coordinator and drive the agent loop until shutdown. + /// Public so tests (and embedders) can run an agent in-process. + pub async fn run(mut config: AgentConfig, run_once: bool) -> Result<()> { + tracing::info!("Starting agent against {}", config.coordinator_addr); + tracing::info!( + "groups={:?} tags={:?} labels={:?}", + config.groups, + config.tags, + config.labels + ); + + let http_client = reqwest::Client::new(); + let mut credential_guard = CredentialGuard::new( + config + .credential_path + .as_ref() + .map(|credential_path| credential_path.relative()), + &config.coordinator_addr, + ) + .await; + let username = match &config.user { + Some(name) => name.to_string(), + None => get_and_prompt_username(None, "Please input username")?, + }; + // Like the worker: an agent restart must not rotate the user's auth + // signature, which would invalidate every token minted before it. + let (_, user_credential) = get_user_credential( + &mut credential_guard, + &http_client, + config.coordinator_addr.clone(), + username, + config.password.take(), + false, + ) + .await?; + + let machine_code = resolve_machine_code(config.machine_code.take()); + tracing::info!("Machine code: {machine_code}"); + + let mut register_url = config.coordinator_addr.clone(); + register_url.set_path("agents"); + let req = RegisterAgentReq { + tags: config.tags.clone(), + labels: config.labels.clone(), + admin_group: config.admin_group.clone(), + groups: config.groups.clone(), + lifetime: config.lifetime, + machine_code, + metadata: Some(AgentMetadata { + version: env!("CARGO_PKG_VERSION").to_string(), + long_version: env!("CARGO_PKG_VERSION").to_string(), + }), + }; + let resp = http_client + .post(register_url.as_str()) + .bearer_auth(&user_credential) + .json(&req) + .send() + .await + .map_err(|e| { + if e.is_request() && e.is_connect() { + error::RequestError::ConnectionError(config.coordinator_addr.to_string()) + } else { + e.into() + } + })?; + let register: RegisterAgentResp = parse_json(resp, "register agent").await?; + + tracing::info!( + "Registered as agent {} ({})", + register.agent_uuid, + if register.reused { + "re-adopted this machine's existing agent" + } else { + "new agent" + } + ); + + let mut coordinator_addr = config.coordinator_addr; + coordinator_addr.set_path(""); + + let mut cache_path = + dirs::cache_dir().ok_or(error::Error::Custom("Cache dir not found".to_string()))?; + cache_path.push("mitosis"); + cache_path.push("agent"); + cache_path.push(register.agent_uuid.to_string()); + tokio::fs::create_dir_all(&cache_path).await?; + tracing::info!("Working directory: {}", cache_path.display()); + + let mut client = AgentClient { + coordinator_addr, + token: register.token, + notification_counter: register.notification_counter, + coordinator_boot_id: None, + state: AgentState::Idle, + assigned_suite_uuid: None, + // Registration is itself a reason to look for work: the coordinator + // could not have notified us before we existed. + pending_suite: Some(PendingSuite::Any), + heartbeat_interval: config.heartbeat_interval, + connect_retry_interval: config.connect_retry_interval, + idle_poll_interval: config.idle_poll_interval, + ws_reconnect_interval: config.ws_reconnect_interval, + no_ws: config.no_ws, + cache_path, + http_client, + job_token: None, + drain_token: None, + current_run: None, + shutdown_token: CancellationToken::new(), + stop_after_current: false, + run_once, + }; + + client.run_loop().await + } +} + +/// Resolve a stable machine code for this host. +/// +/// Precedence: explicit config → the cached value under the mitosis config dir → +/// `/etc/machine-id` → a fresh UUID. Whatever is resolved (unless explicitly +/// overridden) is cached, so the identity survives restarts even where +/// `/etc/machine-id` does not exist — containers, non-Linux hosts. The identity +/// matters: the coordinator keys the agent row on it, and a machine that comes +/// back with a different code strands its old row. +fn resolve_machine_code(explicit: Option) -> String { + if let Some(code) = explicit + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + { + return code; + } + + let cache_path = dirs::config_dir().map(|mut p| { + p.push("mitosis"); + p.push("machine-id"); + p + }); + if let Some(ref path) = cache_path { + if let Ok(cached) = std::fs::read_to_string(path) { + let cached = cached.trim().to_string(); + if !cached.is_empty() { + return cached; + } + } + } + + let code = std::fs::read_to_string("/etc/machine-id") + .ok() + .map(|s| s.trim().to_string()) + .filter(|s| !s.is_empty()) + .unwrap_or_else(|| { + tracing::info!("No /etc/machine-id available; generating a machine code"); + Uuid::new_v4().simple().to_string() + }); + + match cache_path { + Some(path) => { + let res = path + .parent() + .map(std::fs::create_dir_all) + .unwrap_or(Ok(())) + .and_then(|_| std::fs::write(&path, &code)); + if let Err(e) = res { + tracing::warn!( + "Failed to cache the machine code at {}: {e}; a new one may be generated next start", + path.display() + ); + } + } + None => tracing::warn!( + "No config directory to cache the machine code in; a new one may be generated next start" + ), + } + + code +} + +impl AgentClient { + fn api_url(&self, path: &str) -> Url { + let mut url = self.coordinator_addr.clone(); + url.set_path(path); + url + } + + async fn run_loop(&mut self) -> Result<()> { + let cancel_token = self.shutdown_token.clone(); + + let signal_token = cancel_token.clone(); + tokio::spawn(async move { + tokio::signal::ctrl_c().await.ok(); + tracing::info!("Received SIGINT, shutting down"); + signal_token.cancel(); + }); + + // `no_ws` opens no socket at all and disables the branch that reads it. + // The channel is built either way: `select!` evaluates a branch's + // expression even when its guard is false, so `ws_rx` has to exist. + let ws_enabled = !self.no_ws; + let (ws_tx, mut ws_rx) = mpsc::channel::(32); + let ws_handle = if ws_enabled { + Some(self.spawn_websocket_client(ws_tx, cancel_token.clone())) + } else { + tracing::info!("WebSocket disabled; Notifications rely purely on heartbeat"); + None + }; + + let mut heartbeat_timer = tokio::time::interval(self.heartbeat_interval); + heartbeat_timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + + let (idle_poll_enabled, mut idle_poll_timer) = match self.idle_poll_interval { + Some(interval) => { + tracing::info!( + "Idle agent will poll for a suite every {}s", + interval.as_secs() + ); + let mut timer = tokio::time::interval(interval); + timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + (true, timer) + } + None => { + // When idle_pool_enabled is false, the timer's select branch will be disabled so + // this timer never fires. + let mut timer = tokio::time::interval(self.heartbeat_interval); + timer.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + (false, timer) + } + }; + + tracing::info!("Agent entering its main loop"); + loop { + tokio::select! { + biased; + _ = cancel_token.cancelled() => { + tracing::info!("Shutdown signalled, leaving the main loop"); + // Stop the in-flight suite so its reports cease; whatever it + // was holding is reclaimed once our heartbeat lapses. + if let Some(token) = &self.job_token { + token.cancel(); + } + break; + } + + Some(event) = ws_rx.recv(), if ws_enabled => { + self.notification_counter = self.notification_counter.max(event.id); + self.handle_notification(event.event); + } + + _ = heartbeat_timer.tick() => { + if let Err(e) = self.send_heartbeat().await { + tracing::error!("Heartbeat failed: {e}"); + } + } + + _ = idle_poll_timer.tick(), if idle_poll_enabled => { + // This arm fires and exit the select, entering the self.advance that fetches a new + // suite. + // + // When the agent is in busy state, the self.advance returns almost immediately, + // which shouldn't cause any trouble. + if self.state == AgentState::Idle && !self.stop_after_current { + self.note_pending_suite(None); + } + } + } + + if let Err(e) = self.advance().await { + tracing::error!("Agent state handling failed: {e}"); + } + if self.stop_after_current && self.current_run.is_none() { + tracing::info!("Nothing left to run and a stop was requested; exiting"); + cancel_token.cancel(); + break; + } + } + + // Let the in-flight suite drain so its last reports land. + if let Some(handle) = self.current_run.take() { + let _ = handle.await; + } + if let Some(handle) = ws_handle { + let _ = handle.await; + } + tracing::info!("Agent stopped"); + Ok(()) + } + + // ── notifications ── + + fn spawn_websocket_client( + &self, + notification_tx: mpsc::Sender, + cancel_token: CancellationToken, + ) -> JoinHandle<()> { + let mut url = self.coordinator_addr.clone(); + let _ = url.set_scheme(if url.scheme() == "https" { "wss" } else { "ws" }); + url.set_path("ws/agents"); + let ws_url = url.to_string(); + let token = self.token.clone(); + let reconnect_interval = self.ws_reconnect_interval; + + tokio::spawn(async move { + while !cancel_token.is_cancelled() { + tracing::debug!("Connecting to {ws_url}"); + match websocket_session(&ws_url, &token, ¬ification_tx, &cancel_token).await { + Ok(()) => tracing::debug!("Agent WebSocket closed"), + // Notifications are an optimization — the heartbeat carries + // the same events — so a failure here is never fatal. + Err(e) => tracing::warn!("Agent WebSocket error: {e}"), + } + tokio::select! { + _ = cancel_token.cancelled() => break, + _ = tokio::time::sleep(reconnect_interval) => {} + } + } + tracing::debug!("Agent WebSocket task stopped"); + }) + } + + fn note_pending_suite(&mut self, suite_uuid: Option) { + match suite_uuid { + // A named suite always wins: a specific target beats the generic hint. + Some(uuid) => self.pending_suite = Some(PendingSuite::Specific(uuid)), + None => { + if self.pending_suite.is_none() { + self.pending_suite = Some(PendingSuite::Any); + } + } + } + } + + fn handle_notification(&mut self, notification: AgentNotification) { + match notification { + AgentNotification::SuiteAvailable { + suite_uuid, + priority, + } => { + tracing::debug!(?suite_uuid, priority, "Suite available"); + self.note_pending_suite(suite_uuid); + } + AgentNotification::PreemptSuite { + new_suite_uuid, + current_suite_uuid, + .. + } => { + // The coordinator does not emit this today (a running agent is + // never interrupted). Handled anyway, so turning preemption on + // is a coordinator-side change only. + if self.assigned_suite_uuid == Some(current_suite_uuid) { + tracing::info!("Preempting suite {current_suite_uuid} for {new_suite_uuid}"); + if let Some(token) = &self.job_token { + token.cancel(); + } + self.note_pending_suite(Some(new_suite_uuid)); + } + } + AgentNotification::SuiteCancelled { suite_uuid, reason } => { + if self.assigned_suite_uuid == Some(suite_uuid) { + tracing::warn!("Suite {suite_uuid} was cancelled: {reason}"); + if let Some(token) = &self.job_token { + token.cancel(); + } + } + } + AgentNotification::TasksCancelled { task_uuids } => { + // No client-side bookkeeping needed: a report against a task + // that is gone answers 404 and the runner moves on. + tracing::debug!(count = task_uuids.len(), "Tasks cancelled"); + } + AgentNotification::Shutdown { graceful } => { + tracing::warn!("Coordinator asked this agent to shut down (graceful={graceful})"); + if graceful { + self.stop_after_current = true; + self.pending_suite = None; + // Stop waiting for more work, but let the job finish its + // cleanup — that is exactly what `drain_token` separates. + if let Some(token) = &self.drain_token { + token.cancel(); + } + } else { + if let Some(token) = &self.job_token { + token.cancel(); + } + self.shutdown_token.cancel(); + } + } + AgentNotification::Ping { .. } => {} + AgentNotification::CounterSync { counter, boot_id } => { + if self.coordinator_boot_id != Some(boot_id) { + if let Some(previous) = self.coordinator_boot_id { + tracing::warn!("Coordinator restarted ({previous} → {boot_id})"); + } + self.coordinator_boot_id = Some(boot_id); + self.notification_counter = counter; + } else if counter > self.notification_counter { + self.notification_counter = counter; + } + } + }; + } + + async fn send_heartbeat(&mut self) -> Result<()> { + let req = AgentHeartbeatReq { + state: self.state, + assigned_suite_uuid: self.assigned_suite_uuid, + last_notification_id: self.notification_counter, + metrics: None, + }; + let resp = self + .http_client + .post(self.api_url("agents/heartbeat").as_str()) + .bearer_auth(&self.token) + .json(&req) + .send() + .await + .map_err(error::map_reqwest_err)?; + let resp: AgentHeartbeatResp = parse_json(resp, "heartbeat").await?; + + // Catch-up path: whatever the WebSocket did not deliver arrives here. + for event in resp.notifications { + self.notification_counter = self.notification_counter.max(event.id); + self.handle_notification(event.event); + } + Ok(()) + } + + // ── state machine ── + + async fn advance(&mut self) -> Result<()> { + match self.state { + AgentState::Idle => { + if self.stop_after_current { + return Ok(()); + } + // Act only on a signal: a notification, the re-check the + // coordinator runs on every heartbeat, or the idle poll. + let target = match self.pending_suite.take() { + None => return Ok(()), + Some(PendingSuite::Any) => None, + Some(PendingSuite::Specific(uuid)) => Some(uuid), + }; + + // One attempt per signal: the slot was cleared above, so + // "nothing available" leaves us idle until the next + // notification rather than spinning. + let Some((suite, job, job_id)) = self.fetch_suite(target).await? else { + return Ok(()); + }; + tracing::info!("Accepted suite {} as job {job_id}", suite.uuid); + + self.assigned_suite_uuid = Some(suite.uuid); + let job_token = CancellationToken::new(); + self.job_token = Some(job_token.clone()); + let drain_token = CancellationToken::new(); + self.drain_token = Some(drain_token.clone()); + // `--run-once` exits after this suite, so it never holds a + // drained one for the idle window. + if self.run_once { + drain_token.cancel(); + } + let runner = SuiteRunner { + http_client: self.http_client.clone(), + token: self.token.clone(), + coordinator_addr: self.coordinator_addr.clone(), + connect_retry_interval: self.connect_retry_interval, + job_id, + cache_path: self.cache_path.join("job"), + job_token, + drain_token, + }; + self.current_run = Some(tokio::spawn(runner.run(suite, job))); + self.state = AgentState::Executing; + } + AgentState::Executing => { + // Reap without blocking the loop, so heartbeats keep flowing. + let finished = self + .current_run + .as_ref() + .map(JoinHandle::is_finished) + .unwrap_or(true); + if finished { + if let Some(handle) = self.current_run.take() { + match handle.await { + // `complete` said more work is waiting for us. + Ok(true) => self.note_pending_suite(None), + Ok(false) => {} + Err(e) => tracing::error!("The suite runner task failed: {e}"), + } + } + self.assigned_suite_uuid = None; + self.job_token = None; + self.drain_token = None; + self.state = AgentState::Idle; + if self.run_once { + tracing::info!("--run-once: one suite done, stopping"); + self.stop_after_current = true; + } + tracing::info!("Agent is idle again"); + } + } + // Coordinator-owned: the agent only ever reports Idle/Executing, + // its other phases being bracketed by the start/cleanup calls. + AgentState::Provisioning | AgentState::Cleaning | AgentState::Offline => {} + } + Ok(()) + } + + /// Ask for a suite and claim it in one call, naming the one we were + /// notified about if there was one. `None` means nothing was available. + /// + /// The suite handed back need not be the one asked for: a stale hint is + /// answered with the best available instead. + async fn fetch_suite( + &self, + suite_uuid: Option, + ) -> Result> { + let resp = self + .http_client + .post(self.api_url("agents/suite").as_str()) + .bearer_auth(&self.token) + .json(&AcceptSuiteReq { suite_uuid }) + .send() + .await + .map_err(error::map_reqwest_err)?; + let resp: AcceptSuiteResp = parse_json(resp, "accept suite").await?; + if !resp.accepted { + tracing::debug!("No suite accepted: {}", resp.reason.unwrap_or_default()); + return Ok(None); + } + let Some(suite) = resp.suite else { + return Err(error::Error::Custom( + "the coordinator accepted a suite but sent no spec".to_string(), + )); + }; + Ok(resp.job.zip(resp.job_id).map(|(job, id)| (suite, job, id))) + } +} + +/// Drives one accepted suite from provision through to a terminal job state. +#[derive(Clone)] +struct SuiteRunner { + http_client: reqwest::Client, + token: String, + coordinator_addr: Url, + connect_retry_interval: Duration, + /// The suite's own job number, and the only one that belongs in a log or on + /// screen: it is what `suites jobs --job N` takes. The `job` handle + /// threaded through the calls below is `suite_agent_jobs.id`, an internal + /// key the user never sees. + job_id: i32, + /// This job's working subtree + cache_path: PathBuf, + /// Cancelled when the suite is cancelled, preempted, or the agent stops. + /// Everything the job runs (tasks, hooks, cleanup) runs under it. + job_token: CancellationToken, + /// Cancelled to stop waiting without stopping the wind-down: a graceful + /// shutdown releases a held job, but its cleanup hook still has to run, and + /// that runs under `job_token`. + drain_token: CancellationToken, +} + +impl SuiteRunner { + fn api_url(&self, path: &str) -> Url { + let mut url = self.coordinator_addr.clone(); + url.set_path(path); + url + } + + /// Every phase logs and continues on error rather than bailing: only the + /// agent can walk the job to a terminal state and release itself. + /// + /// Returns whether `complete` said this agent has more work waiting. + async fn run(self, suite: TaskSuiteSpec, job: i64) -> bool { + let suite_uuid = suite.uuid; + let mut failure: Option = None; + + let provisioned = match self.prepare_workspace().await { + Ok(()) => self.run_hook(job, &suite, HookType::Provision).await, + Err(e) => Err(e), + }; + if let Err(e) = provisioned { + tracing::error!("Provisioning failed for suite {suite_uuid}: {e}"); + failure = Some(JobFailureReason { + kind: JobFailureKind::ProvisionFailed, + message: e.to_string(), + }); + } + + if let Err(e) = self + .post_empty("agents/job/start", &StartJobReq { job }) + .await + { + tracing::error!("Failed to start suite {suite_uuid}: {e}"); + } + + let background_token = CancellationToken::new(); + let mut background = None; + + if failure.is_none() { + background = self.spawn_background_hook(job, &suite, background_token.clone()); + if let Err(e) = self + .execute_tasks(suite_uuid, job, task_slots(&suite)) + .await + { + tracing::error!("Task execution failed for suite {suite_uuid}: {e}"); + failure = Some(JobFailureReason { + kind: JobFailureKind::ExecutionError, + message: e.to_string(), + }); + } + } + + if let Some(handle) = background { + // Finishing before we ask it to is the failure mode a background + // hook has: it was supposed to outlive the tasks it serves. + let exited_early = handle.is_finished(); + background_token.cancel(); + if let Err(e) = handle.await { + tracing::error!("The background hook task failed for suite {suite_uuid}: {e}"); + } + if exited_early { + failure.get_or_insert(JobFailureReason { + kind: JobFailureKind::BackgroundExited, + message: "the background hook exited before the tasks drained".to_string(), + }); + } else { + self.record_stopped_hook(job, suite_uuid, HookType::Background) + .await; + } + } + + if let Err(e) = self + .post_empty("agents/job/cleanup", &EnterCleanupReq { job }) + .await + { + tracing::error!("Failed to enter cleanup for suite {suite_uuid}: {e}"); + } + + if let Err(e) = self.run_hook(job, &suite, HookType::Cleanup).await { + tracing::error!("Cleanup hook failed for suite {suite_uuid}: {e}"); + failure.get_or_insert(JobFailureReason { + kind: JobFailureKind::CleanupFailed, + message: e.to_string(), + }); + } + + let outcome = match failure { + None => SuiteJobOutcome::Completed, + Some(reason) => SuiteJobOutcome::Failed { reason }, + }; + let next_available = match self.report_complete(job, outcome).await { + Ok(next_available) => { + tracing::info!( + "Finished suite {suite_uuid} (job {}); more work waiting: {next_available}", + self.job_id + ); + next_available + } + // A `complete` that never landed says nothing about what is next. + Err(e) => { + tracing::error!("Failed to complete suite {suite_uuid}: {e}"); + false + } + }; + + // Everything worth keeping has been uploaded by now + let _ = tokio::fs::remove_dir_all(&self.cache_path).await; + next_available + } + + /// The directory the provision hook, the tasks and the cleanup hook all see + /// as `MITO_SUITE_SHARED`. + fn shared_path(&self) -> PathBuf { + self.cache_path.join("share") + } + + /// Create a job's file tree + async fn prepare_workspace(&self) -> Result<()> { + let _ = tokio::fs::remove_dir_all(&self.cache_path).await; + tokio::fs::create_dir_all(self.shared_path()).await?; + Ok(()) + } + + /// Claim and run the suite's tasks across `slots` task slots. + /// + /// One slot is one worker of the suite's `FixedWorkers` plan: its own + /// working directory and its own claim loop, so a slow task holds up only + /// the slot running it. + /// + /// TODO: `task_prefetch_count` is not honored yet — a slot claims one task + /// at a time rather than a local batch. + async fn execute_tasks(&self, suite_uuid: Uuid, job: i64, slots: usize) -> Result<()> { + tracing::info!("Running the tasks of suite {suite_uuid} across {slots} slot(s)"); + let mut running = tokio::task::JoinSet::new(); + for slot in 0..slots { + let runner = self.clone(); + running.spawn(async move { runner.run_slot(suite_uuid, job, slot).await }); + } + + // Every slot is joined before the job moves on, so one that fails never + // strands the task another still has in flight. + let mut failure: Option = None; + while let Some(joined) = running.join_next().await { + match joined { + Ok(Ok(())) => {} + Ok(Err(e)) => { + tracing::error!("A task slot of suite {suite_uuid} failed: {e}"); + failure.get_or_insert(e); + } + Err(e) => { + tracing::error!("A task slot of suite {suite_uuid} panicked: {e}"); + failure.get_or_insert(error::Error::Custom(format!("task slot panicked: {e}"))); + } + } + } + match failure { + Some(e) => Err(e), + None => Ok(()), + } + } + + /// One task slot: claim a task, run it in `task-{slot}`, repeat until the + /// suite drains and stops holding the job open. + async fn run_slot(&self, suite_uuid: Uuid, job: i64, slot: usize) -> Result<()> { + let dir = self.cache_path.join(format!("task-{slot}")); + let mut holding = false; + loop { + if self.job_token.is_cancelled() { + tracing::info!("Suite {suite_uuid} cancelled; stopping task slot {slot}"); + return Ok(()); + } + + let resp = self.fetch_tasks(suite_uuid, 1).await?; + if resp.tasks.is_empty() { + if !resp.hold_job_open { + tracing::info!( + "Suite {suite_uuid} is idle and has no more tasks for slot {slot}" + ); + return Ok(()); + } + if !holding { + tracing::info!( + "Suite {suite_uuid} is drained but still open; slot {slot} holds job {} open for more work", + self.job_id + ); + holding = true; + } + // The coordinator only notifies *idle* agents of a submission, + // and we are Executing, so the poll is what finds late work. + tokio::select! { + _ = self.job_token.cancelled() => return Ok(()), + // Stops the wait, not the wind-down: cleanup still runs. + _ = self.drain_token.cancelled() => { + tracing::info!( + "Asked to wind down; releasing slot {slot} of job {}", + self.job_id + ); + return Ok(()); + } + _ = tokio::time::sleep(HOLD_POLL_INTERVAL) => {} + } + continue; + } + if holding { + tracing::info!( + "Suite {suite_uuid} has work again; slot {slot} of job {} resumes", + self.job_id + ); + holding = false; + } + + for task in resp.tasks { + if self.job_token.is_cancelled() { + return Ok(()); + } + let uuid = task.uuid; + if let Err(e) = self.run_task(job, task, &dir).await { + // One task must not take the job down: the coordinator + // reclaims anything left uncommitted. + tracing::error!("Failed to run task {uuid}: {e}"); + } + } + } + } + + /// Run one task, reporting it through the agent's task endpoint. + async fn run_task(&self, job: i64, task: WorkerTaskResp, slot: &std::path::Path) -> Result<()> { + tracing::info!("Running task {} (args {:?})", task.uuid, task.spec.args); + reset_workspace(slot).await?; + let mut executor = Executor { + cancel_token: self.job_token.clone(), + polling_interval: self.connect_retry_interval, + cache_path: slot.to_path_buf(), + http_client: self.http_client.clone(), + client: Box::new(AgentTaskClient { + http: self.connection(), + job, + task_id: task.id, + task_uuid: task.uuid, + upstream_task_uuid: task.upstream_task_uuid, + shared_path: self.shared_path(), + }), + }; + execute(&mut executor, task.spec, task.exec_options.as_ref()).await + } + + /// Run a suite hook, reporting it through the hook endpoint. An unconfigured + /// hook is not a failed one — it just leaves no record. + async fn run_hook(&self, job: i64, suite: &TaskSuiteSpec, hook_type: HookType) -> Result<()> { + let Some(spec) = hook_spec(suite, hook_type) else { + tracing::debug!("No {hook_type} hook configured for suite {}", suite.uuid); + return Ok(()); + }; + self.run_hook_spec(job, suite, hook_type, spec, self.job_token.clone()) + .await + } + + async fn run_hook_spec( + &self, + job: i64, + suite: &TaskSuiteSpec, + hook_type: HookType, + spec: ExecSpec, + cancel_token: CancellationToken, + ) -> Result<()> { + tracing::info!("Running the {hook_type} hook of suite {}", suite.uuid); + let dir = self.cache_path.join(format!("hook-{hook_type}")); + reset_workspace(&dir).await?; + let outcome = HookOutcome::default(); + let mut executor = Executor { + cancel_token, + polling_interval: self.connect_retry_interval, + cache_path: dir, + http_client: self.http_client.clone(), + client: Box::new(AgentHookClient { + http: self.connection(), + job, + hook_type, + suite_uuid: suite.uuid, + outcome: outcome.clone(), + shared_path: self.shared_path(), + }), + }; + execute(&mut executor, spec, None).await?; + match outcome.exit_status() { + Some(0) | None => Ok(()), + Some(status) => Err(error::Error::Custom(format!( + "the {hook_type} hook exited with status {status}" + ))), + } + } + + /// Start the background hook, if the suite has one. Not awaited: it runs for + /// as long as the tasks do. + fn spawn_background_hook( + &self, + job: i64, + suite: &TaskSuiteSpec, + cancel_token: CancellationToken, + ) -> Option> { + let spec = hook_spec(suite, HookType::Background)?; + let runner = self.clone(); + let suite = suite.clone(); + Some(tokio::spawn(async move { + if let Err(e) = runner + .run_hook_spec(job, &suite, HookType::Background, spec, cancel_token) + .await + { + tracing::error!("Background hook failed for suite {}: {e}", suite.uuid); + } + })) + } + + /// Record a hook that ended by our own hand rather than by exiting — the + /// background hook, once the tasks it served have drained. The execution + /// core reports nothing on a cancellation, so without this the hook would + /// leave no row at all. + async fn record_stopped_hook(&self, job: i64, suite_uuid: Uuid, hook_type: HookType) { + let mut client = AgentHookClient { + http: self.connection(), + job, + hook_type, + suite_uuid, + outcome: HookOutcome::default(), + shared_path: self.shared_path(), + }; + let result = TaskResultSpec { + exit_status: 0, + msg: None, + }; + if let Err(e) = client.report_commit(result).await { + tracing::error!("Failed to record the stopped {hook_type} hook: {e}"); + } + } + + /// The connection context every per-unit client is built from. + fn connection(&self) -> AgentConnection { + AgentConnection { + http_client: self.http_client.clone(), + token: self.token.clone(), + coordinator_addr: self.coordinator_addr.clone(), + connect_retry_interval: self.connect_retry_interval, + job_token: self.job_token.clone(), + } + } + + async fn fetch_tasks(&self, suite_uuid: Uuid, max_count: u32) -> Result { + let resp = self + .http_client + .post(self.api_url("agents/tasks/fetch").as_str()) + .bearer_auth(&self.token) + .json(&FetchTasksReq { + suite_uuid, + max_count, + }) + .send() + .await + .map_err(error::map_reqwest_err)?; + // The suite going terminal under us is an ordinary end of run, not a + // failure: stop the loop and let cleanup proceed. + if resp.status() == StatusCode::CONFLICT { + tracing::info!("Suite {suite_uuid} stopped handing out tasks"); + self.job_token.cancel(); + return Ok(FetchTasksResp { + tasks: Vec::new(), + hold_job_open: false, + }); + } + parse_json(resp, "fetch tasks").await + } + + async fn report_complete(&self, job: i64, outcome: SuiteJobOutcome) -> Result { + let resp = self + .http_client + .post(self.api_url("agents/job/complete").as_str()) + .bearer_auth(&self.token) + .json(&CompleteJobReq { job, outcome }) + .send() + .await + .map_err(error::map_reqwest_err)?; + let resp: CompleteJobResp = parse_json(resp, "complete job").await?; + Ok(resp.next_suite_available) + } + + /// POST a request whose success carries no body. + async fn post_empty(&self, path: &str, req: &T) -> Result<()> { + let resp = self + .http_client + .post(self.api_url(path).as_str()) + .bearer_auth(&self.token) + .json(req) + .send() + .await + .map_err(error::map_reqwest_err)?; + if resp.status().is_success() { + return Ok(()); + } + Err(error::Error::Custom(format!( + "{path} failed: {}", + error::get_error_from_resp(resp).await + ))) + } +} + +/// The hook's `ExecSpec` from the suite definition, if it has one. +fn hook_spec(suite: &TaskSuiteSpec, hook_type: HookType) -> Option { + let hooks = suite.exec_hooks.as_ref()?; + match hook_type { + HookType::Provision => hooks.provision.clone(), + HookType::Cleanup => hooks.cleanup.clone(), + HookType::Background => hooks.background.clone(), + } +} + +/// The result a hook committed, handed back out of the execution core so the +/// runner can turn a non-zero exit into a job failure. Empty when the hook never +/// got as far as committing (it was cancelled, or its inputs were unavailable). +#[derive(Clone, Default)] +struct HookOutcome(Arc>>); + +impl HookOutcome { + fn record(&self, result: &TaskResultSpec) { + if let Ok(mut slot) = self.0.lock() { + *slot = Some(result.clone()); + } + } + + fn exit_status(&self) -> Option { + self.0.lock().ok()?.as_ref().map(|r| r.exit_status) + } +} + +// ── the agent's half of the execution seam ── + +/// What every agent-side [`ExecClient`] needs to reach the coordinator. Cloned +/// per unit; the clones are all cheap handles. +#[derive(Clone)] +struct AgentConnection { + http_client: reqwest::Client, + token: String, + coordinator_addr: Url, + connect_retry_interval: Duration, + /// The job's token. Cancelling it stops every unit of this job at once, + /// which is how a 409 ("this job is closed") propagates. + job_token: CancellationToken, +} + +impl AgentConnection { + fn api_url(&self, path: &str) -> Url { + let mut url = self.coordinator_addr.clone(); + url.set_path(path); + url + } + + fn get(&self, path: &str) -> reqwest::RequestBuilder { + self.http_client + .get(self.api_url(path).as_str()) + .bearer_auth(&self.token) + } + + /// POST a report, retrying connection failures until the job is cancelled. + /// `None` means the coordinator answered something that ends the reporting: + /// 409 (the job is closed — which also stops the job) or 404 (it is gone). + async fn post_report( + &self, + path: &str, + req: &Req, + what: &str, + ) -> Result> { + let url = self.api_url(path); + loop { + let resp = self + .http_client + .post(url.as_str()) + .bearer_auth(&self.token) + .json(req) + .send() + .await; + let resp = match resp { + Ok(resp) => resp, + Err(e) if e.is_connect() && e.is_request() => { + tracing::warn!( + "{what} failed to connect ({e}); retrying in {:?}", + self.connect_retry_interval + ); + tokio::select! { + _ = self.job_token.cancelled() => return Ok(None), + _ = tokio::time::sleep(self.connect_retry_interval) => {} + } + continue; + } + Err(e) => return Err(error::RequestError::from(e).into()), + }; + + if resp.status() == StatusCode::CONFLICT { + tracing::info!("The job is closed; stopping this run"); + self.job_token.cancel(); + return Ok(None); + } + if resp.status() == StatusCode::NOT_FOUND { + tracing::debug!("{what}: the target is gone; skipping the report"); + return Ok(None); + } + return parse_json(resp, what).await.map(Some); + } + } +} + +/// A task the agent runs on a suite's behalf. Same protocol as the worker's, +/// addressed to `/agents/tasks/report` with the job handle attached. +struct AgentTaskClient { + http: AgentConnection, + job: i64, + task_id: i64, + task_uuid: Uuid, + upstream_task_uuid: Option, + /// The job's `share/`, exported as `MITO_SUITE_SHARED`. + shared_path: PathBuf, +} + +impl AgentTaskClient { + async fn report(&self, op: ReportTaskOp) -> Result> { + self.http + .post_report( + "agents/tasks/report", + &ReportAgentTaskReq { + job: self.job, + id: self.task_id, + op, + }, + "report task", + ) + .await + } +} + +#[async_trait::async_trait] +impl ExecClient for AgentTaskClient { + fn describe(&self) -> String { + format!("task {}", self.task_uuid) + } + + fn exec_env(&self) -> Vec<(&'static str, String)> { + let mut env = vec![ + ("MITO_TASK_UUID", self.task_uuid.to_string()), + ( + "MITO_SUITE_SHARED", + self.shared_path.to_string_lossy().into_owned(), + ), + ]; + if let Some(uuid) = self.upstream_task_uuid { + env.push(("MITO_UPSTREAM_TASK_UUID", uuid.to_string())); + } + env + } + + fn supports_child_tasks(&self) -> bool { + true + } + + async fn report_finish(&mut self, finished: bool, _result: &TaskResultSpec) -> Result<()> { + let op = if finished { + ReportTaskOp::Finish + } else { + ReportTaskOp::Cancel + }; + self.report(op).await.map(|_| ()) + } + + async fn request_upload( + &mut self, + content_type: ArtifactContentType, + content_length: u64, + ) -> Result { + let resp = self + .report(ReportTaskOp::Upload { + content_type, + content_length, + }) + .await?; + Ok(match resp.and_then(|resp| resp.url) { + Some(url) => UploadTarget::Url(url), + None => UploadTarget::Skip, + }) + } + + async fn report_commit(&mut self, result: TaskResultSpec) -> Result<()> { + self.report(ReportTaskOp::Commit(result)).await.map(|_| ()) + } + + async fn submit_child_task(&mut self, req: SubmitTaskReq) -> Result<()> { + self.report(ReportTaskOp::Submit(Box::new(req))) + .await + .map(|_| ()) + } + + fn artifact_download_req( + &self, + uuid: Uuid, + content_type: ArtifactContentType, + ) -> reqwest::RequestBuilder { + self.http.get(&artifact_path(uuid, content_type)) + } + + fn attachment_download_req(&self, key: &str) -> reqwest::RequestBuilder { + let uuid = self.task_uuid; + self.http + .get(&format!("agents/tasks/{uuid}/attachments/{key}")) + } + + /// No redis on the agent side yet, so an agent task's fine-grained states go + /// unpublished. + async fn announce_state(&mut self, _state: TaskExecState, _ex: Option) {} + + fn can_watch(&self) -> bool { + true + } + + /// Poll-only, which resolves the coarse milestones a `watch` asks for (has + /// the other task committed?) but not the intra-execution states. + async fn watch(&mut self, uuid: &Uuid, target: TaskExecState) { + tracing::debug!("Watch task: {} -> {:?}", uuid, target); + loop { + let resp = self.http.get(&format!("agents/tasks/{uuid}")).send().await; + match resp { + Ok(resp) if resp.status().is_success() => { + match resp.json::().await { + Ok(task) => { + if task.info.state.is_reach(&target, task.info.result) { + return; + } + } + Err(e) => tracing::warn!("Unreadable watched task {uuid}: {e}"), + } + } + Ok(resp) => tracing::warn!( + "Watching task {uuid} failed: {}", + error::get_error_from_resp(resp).await + ), + Err(e) => tracing::warn!("Watching task {uuid} failed: {e}"), + } + tokio::select! { + _ = self.http.job_token.cancelled() => return, + _ = tokio::time::sleep(WATCH_POLL_INTERVAL) => {} + } + } + } +} + +/// A suite hook. Keyed by `{job, hook_type}` rather than by a task, and reported +/// to `/agents/job/hook`, whose `Result` op records the outcome and mints the +/// uuid its artifacts hang off — hence the outcome is reported before the +/// uploads. +struct AgentHookClient { + http: AgentConnection, + job: i64, + hook_type: HookType, + suite_uuid: Uuid, + outcome: HookOutcome, + /// The job's `share/`, exported as `MITO_SUITE_SHARED`. + shared_path: PathBuf, +} + +impl AgentHookClient { + async fn report(&self, op: HookReportOp) -> Result> { + self.http + .post_report( + "agents/job/hook", + &HookReportReq { + job: self.job, + hook_type: self.hook_type, + op, + }, + "report hook", + ) + .await + } +} + +#[async_trait::async_trait] +impl ExecClient for AgentHookClient { + fn describe(&self) -> String { + format!("the {} hook of job {}", self.hook_type, self.job) + } + + fn exec_env(&self) -> Vec<(&'static str, String)> { + vec![ + ("MITO_HOOK_TYPE", self.hook_type.to_string()), + ("MITO_SUITE_UUID", self.suite_uuid.to_string()), + ( + "MITO_SUITE_SHARED", + self.shared_path.to_string_lossy().into_owned(), + ), + ] + } + + /// The hook report endpoint has no submit operation, so `MITO_NEW_TASK` is + /// not exported to a hook. + fn supports_child_tasks(&self) -> bool { + false + } + + /// Writes the hook row. Its uuid is what the uploads that follow are keyed + /// by, so this must land before them. + async fn report_finish(&mut self, _finished: bool, result: &TaskResultSpec) -> Result<()> { + self.report(HookReportOp::Result(result.clone())) + .await + .map(|_| ()) + } + + async fn request_upload( + &mut self, + content_type: ArtifactContentType, + content_length: u64, + ) -> Result { + let resp = self + .report(HookReportOp::Upload { + content_type, + content_length, + }) + .await?; + Ok(match resp.and_then(|resp| resp.url) { + Some(url) => UploadTarget::Url(url), + None => UploadTarget::Skip, + }) + } + + /// Re-reports the row with the final message. The endpoint upserts on + /// `{job, hook_type}` and never rewrites the uuid, so the artifacts uploaded + /// in between stay attached. + async fn report_commit(&mut self, result: TaskResultSpec) -> Result<()> { + self.outcome.record(&result); + self.report(HookReportOp::Result(result)).await.map(|_| ()) + } + + async fn submit_child_task(&mut self, _req: SubmitTaskReq) -> Result<()> { + Err(error::Error::Custom( + "a suite hook cannot submit a task".to_string(), + )) + } + + fn artifact_download_req( + &self, + uuid: Uuid, + content_type: ArtifactContentType, + ) -> reqwest::RequestBuilder { + self.http.get(&artifact_path(uuid, content_type)) + } + + /// Through the suite, not a task: a hook has no task uuid to resolve the + /// owning group by. + fn attachment_download_req(&self, key: &str) -> reqwest::RequestBuilder { + let uuid = self.suite_uuid; + self.http + .get(&format!("agents/suites/{uuid}/attachments/{key}")) + } + + async fn announce_state(&mut self, _state: TaskExecState, _ex: Option) {} +} + +/// The artifact download path both agent clients use. Keyed by the artifact's +/// owning uuid, which the service resolves without caring whose it is. +fn artifact_path(uuid: Uuid, content_type: ArtifactContentType) -> String { + let content_type = serde_json::to_value(content_type) + .ok() + .and_then(|v| v.as_str().map(str::to_string)) + .unwrap_or_else(|| "result".to_string()); + format!("agents/tasks/{uuid}/artifacts/{content_type}") +} + +/// One WebSocket session: read notifications, hand each to the main loop, then +/// acknowledge it. Returns when the socket closes or shutdown is signalled. +async fn websocket_session( + ws_url: &str, + token: &str, + notification_tx: &mpsc::Sender, + cancel_token: &CancellationToken, +) -> Result<()> { + let host = Url::parse(ws_url) + .ok() + .and_then(|u| { + u.host_str().map(|h| match u.port() { + Some(port) => format!("{h}:{port}"), + None => h.to_string(), + }) + }) + .unwrap_or_default(); + let request = tokio_tungstenite::tungstenite::http::Request::builder() + .uri(ws_url) + .header("Host", host) + .header("Authorization", format!("Bearer {token}")) + .header("Sec-WebSocket-Version", "13") + .header("Connection", "Upgrade") + .header("Upgrade", "websocket") + .header( + "Sec-WebSocket-Key", + tokio_tungstenite::tungstenite::handshake::client::generate_key(), + ) + .body(()) + .map_err(|e| error::Error::Custom(format!("Invalid WebSocket request: {e}")))?; + + let (ws_stream, _) = tokio_tungstenite::connect_async(request) + .await + .map_err(|e| error::Error::Custom(format!("WebSocket connect failed: {e}")))?; + tracing::info!("Agent WebSocket connected"); + + let (mut ws_write, mut ws_read) = ws_stream.split(); + loop { + let msg = tokio::select! { + _ = cancel_token.cancelled() => return Ok(()), + msg = ws_read.next() => msg, + }; + let Some(msg) = msg else { return Ok(()) }; + match msg { + Ok(WsMessage::Binary(bytes)) => { + let event = match WsNotificationEvent::read_from_buffer(&bytes) { + Ok(event) => event, + Err(e) => { + tracing::warn!("Unparseable notification: {e}"); + continue; + } + }; + let id = event.id; + if notification_tx.send(event).await.is_err() { + return Ok(()); + } + // Acknowledge only once the main loop has taken it, so an event + // never leaves the coordinator's replay buffer before we own it. + let ack = AgentWsMessage::Ack { + notification_id: id, + }; + if let Ok(payload) = ack.write_to_vec() { + let _ = ws_write.send(WsMessage::Binary(payload.into())).await; + } + } + Ok(WsMessage::Ping(_)) | Ok(WsMessage::Pong(_)) | Ok(WsMessage::Frame(_)) => {} + Ok(WsMessage::Text(text)) => { + tracing::debug!(%text, "Ignoring an unexpected text frame") + } + Ok(WsMessage::Close(frame)) => { + tracing::debug!(?frame, "Coordinator closed the WebSocket"); + return Ok(()); + } + Err(e) => return Err(error::Error::Custom(format!("WebSocket error: {e}"))), + } + } +} + +/// Decode a successful JSON response, turning a non-2xx into a descriptive error. +async fn parse_json( + resp: reqwest::Response, + what: &str, +) -> Result { + if !resp.status().is_success() { + return Err(error::Error::Custom(format!( + "{what} failed: {}", + error::get_error_from_resp(resp).await + ))); + } + resp.json::() + .await + .map_err(|e| error::Error::Custom(format!("Unreadable {what} response: {e}"))) +} diff --git a/netmito/src/api/agents.rs b/netmito/src/api/agents.rs new file mode 100644 index 00000000..190ff7a1 --- /dev/null +++ b/netmito/src/api/agents.rs @@ -0,0 +1,256 @@ +use axum::{ + extract::{Path, Query, State}, + middleware, + routing::{delete, get, post}, + Extension, Json, Router, +}; +use uuid::Uuid; + +use super::map_service_error; +use crate::{ + config::InfraPool, + entity::content::ArtifactContentType, + error::ApiError, + schema::{ + AcceptSuiteReq, AcceptSuiteResp, AgentHeartbeatReq, AgentHeartbeatResp, AgentShutdownReq, + AgentsQueryReq, AgentsQueryResp, CompleteJobReq, CompleteJobResp, EnterCleanupReq, + FetchTasksReq, FetchTasksResp, HookReportReq, HookReportResp, RegisterAgentReq, + RegisterAgentResp, RemoteResourceDownloadResp, ReportAgentTaskReq, ReportTaskResp, + StartJobReq, TaskQueryResp, + }, + service::{ + self, + auth::{agent_auth_middleware, user_auth_middleware, 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(), + user_auth_middleware, + )) + .with_state(st.clone()); + + let agent_router = Router::new() + .route("/heartbeat", post(heartbeat)) + .route("/suite", post(accept_suite)) + .route("/job/start", post(start_job)) + .route("/job/cleanup", post(enter_cleanup)) + .route("/job/complete", post(complete_job)) + .route("/job/hook", post(report_hook)) + .route("/tasks/fetch", post(fetch_tasks)) + .route("/tasks/report", post(report_task)) + .route("/tasks/{uuid}", get(query_task)) + .route( + "/tasks/{uuid}/artifacts/{content_type}", + get(download_artifact), + ) + .route("/tasks/{uuid}/attachments/{*key}", get(download_attachment)) + .route( + "/suites/{uuid}/attachments/{*key}", + get(download_suite_attachment), + ) + .route_layer(middleware::from_fn_with_state( + st.clone(), + agent_auth_middleware, + )) + .with_state(st.clone()); + + Router::new().merge(user_router).merge(agent_router) +} + +// ── fleet management (user-authed) ── + +/// `POST /agents` — register, or re-adopt the agent already bound to this +/// machine, and mint 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 db row +/// persists +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(()) +} + +// ── execution loop (agent-authed) ── + +/// `POST /agents/heartbeat` — liveness plus the notifications the agent missed. +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)) +} + +/// `POST /agents/suite` — pick a suite and claim it, opening a job. The body's +/// optional `suite_uuid` is a preference; a stale one falls back to the best +/// available. +async fn accept_suite( + Extension(a): Extension, + State(pool): State, + Json(req): Json, +) -> Result, ApiError> { + let resp = service::agent::agent_accept_suite(a.id, &pool, req) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} + +/// `POST /agents/job/start` — provisioning done, execution starting. +async fn start_job( + Extension(a): Extension, + State(pool): State, + Json(req): Json, +) -> Result<(), ApiError> { + service::agent::agent_start_job(a.id, &pool, req.job) + .await + .map_err(map_service_error)?; + Ok(()) +} + +/// `POST /agents/job/cleanup` — tasks drained, cleanup starting. +async fn enter_cleanup( + Extension(a): Extension, + State(pool): State, + Json(req): Json, +) -> Result<(), ApiError> { + service::agent::agent_enter_cleanup(a.id, &pool, req.job) + .await + .map_err(map_service_error)?; + Ok(()) +} + +/// `POST /agents/job/complete` — job terminal, agent back to idle. +async fn complete_job( + Extension(a): Extension, + State(pool): State, + Json(req): Json, +) -> Result, ApiError> { + let resp = service::agent::agent_complete_job(a.id, &pool, req) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} + +/// `POST /agents/job/hook` — record a hook result or presign its artifacts. +async fn report_hook( + Extension(a): Extension, + State(pool): State, + Json(req): Json, +) -> Result, ApiError> { + let resp = service::agent::hook::agent_report_hook(a.id, req.job, req.hook_type, req.op, &pool) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} + +/// `POST /agents/tasks/fetch` — claim a batch of the suite's ready tasks. +async fn fetch_tasks( + Extension(a): Extension, + State(pool): State, + Json(req): Json, +) -> Result, ApiError> { + let resp = + service::agent::task::agent_fetch_tasks(a.id, a.uuid, &pool, req.suite_uuid, req.max_count) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} + +/// `POST /agents/tasks/report` — mirrors the worker report, plus the job handle. +async fn report_task( + Extension(a): Extension, + State(pool): State, + Json(req): Json, +) -> Result, ApiError> { + let url = service::agent::task::agent_report_task(a.id, a.uuid, req.job, req.id, req.op, &pool) + .await + .map_err(map_service_error)?; + Ok(Json(ReportTaskResp { url })) +} + +/// `GET /agents/tasks/{uuid}` — a task's state and result, for `watch` +/// dependencies. Reuses the service; agent auth is just the gate. +async fn query_task( + Extension(_): Extension, + State(pool): State, + Path(uuid): Path, +) -> Result, ApiError> { + let task = service::task::get_task_by_uuid(&pool, uuid) + .await + .map_err(map_service_error)?; + Ok(Json(task)) +} + +/// `GET /agents/tasks/{uuid}/artifacts/{content_type}` — presigned input download. +async fn download_artifact( + Extension(_): Extension, + State(pool): State, + Path((uuid, content_type)): Path<(Uuid, ArtifactContentType)>, +) -> Result, ApiError> { + let artifact = service::s3::download_artifact_by_uuid(&pool, uuid, content_type) + .await + .map_err(map_service_error)?; + Ok(Json(artifact)) +} + +/// `GET /agents/tasks/{uuid}/attachments/{*key}` — presigned input download. +async fn download_attachment( + Extension(_): Extension, + State(pool): State, + Path((uuid, key)): Path<(Uuid, String)>, +) -> Result, ApiError> { + let attachment = service::s3::worker_download_attachment(&pool, uuid, key) + .await + .map_err(map_service_error)?; + Ok(Json(attachment)) +} + +/// `GET /agents/suites/{uuid}/attachments/{*key}` — presigned input download for +/// a hook, whose inputs belong to the suite rather than to any one task. +async fn download_suite_attachment( + Extension(_): Extension, + State(pool): State, + Path((uuid, key)): Path<(Uuid, String)>, +) -> Result, ApiError> { + let attachment = service::s3::suite_download_attachment(&pool, uuid, key) + .await + .map_err(map_service_error)?; + Ok(Json(attachment)) +} diff --git a/netmito/src/api/mod.rs b/netmito/src/api/mod.rs index c26fba53..ad78c405 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()); @@ -92,8 +96,8 @@ pub fn router(st: InfraPool, cancel_token: CancellationToken) -> Router { } /// Map a service-layer error onto the API error surface. -// TODO: let modules other than suites using this function. -fn map_service_error(e: crate::error::Error) -> ApiError { +// TODO: let modules other than suites/agents using this function. +pub(crate) fn map_service_error(e: crate::error::Error) -> ApiError { match e { crate::error::Error::AuthError(err) => ApiError::AuthError(err), crate::error::Error::ApiError(e) => e, diff --git a/netmito/src/api/suites.rs b/netmito/src/api/suites.rs index 56cdea7a..bb88cc5f 100644 --- a/netmito/src/api/suites.rs +++ b/netmito/src/api/suites.rs @@ -12,7 +12,8 @@ use crate::{ error::ApiError, schema::{ CancelTaskSuiteParam, CreateTaskSuiteReq, CreateTaskSuiteResp, SuiteAgentOverrideReq, - SuiteAgentOverrideResp, TaskSuiteQueryResp, TaskSuitesQueryReq, TaskSuitesQueryResp, + SuiteAgentOverrideResp, SuiteJobQueryResp, SuiteJobsQueryReq, SuiteJobsQueryResp, + TaskSuiteQueryResp, TaskSuitesQueryReq, TaskSuitesQueryResp, }, service::{ self, @@ -27,6 +28,8 @@ pub fn suites_router(st: InfraPool) -> Router { .route("/{uuid}", get(get_suite_details).delete(cancel_suite)) .route("/{uuid}/close", post(close_suite)) .route("/{uuid}/agents/override", post(override_agents_for_suite)) + .route("/{uuid}/jobs/query", post(query_suite_jobs)) + .route("/{uuid}/jobs/{job_id}", get(get_suite_job)) .route_layer(middleware::from_fn_with_state( st.clone(), user_auth_middleware, @@ -101,3 +104,28 @@ pub async fn override_agents_for_suite( .map_err(map_service_error)?; Ok(Json(resp)) } + +/// `POST /suites/{uuid}/jobs/query` — the suite's job history, filtered. +pub async fn query_suite_jobs( + Extension(u): Extension, + State(pool): State, + Path(uuid): Path, + Json(req): Json, +) -> Result, ApiError> { + let resp = service::suite::user_query_suite_jobs(u.id, &pool, uuid, req) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} + +/// `GET /suites/{uuid}/jobs/{job_id}` — one job plus its hook executions. +pub async fn get_suite_job( + Extension(u): Extension, + State(pool): State, + Path((uuid, job_id)): Path<(Uuid, i32)>, +) -> Result, ApiError> { + let resp = service::suite::user_get_suite_job(u.id, &pool, uuid, job_id) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} diff --git a/netmito/src/api/ws.rs b/netmito/src/api/ws.rs new file mode 100644 index 00000000..12b9c99c --- /dev/null +++ b/netmito/src/api/ws.rs @@ -0,0 +1,15 @@ +//! `/ws` — the agent notification socket. + +use axum::{middleware, routing::get, Router}; + +use crate::{config::InfraPool, service::auth::agent_auth_middleware, 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(), + agent_auth_middleware, + )) + .with_state(st) +} diff --git a/netmito/src/channel.rs b/netmito/src/channel.rs new file mode 100644 index 00000000..49225299 --- /dev/null +++ b/netmito/src/channel.rs @@ -0,0 +1,46 @@ +//! Feature-agnostic unbounded MPSC channel. +//! +//! The worker-side actors predate this module and each carry two +//! `#[cfg(feature = "crossfire-channel")]` copies of their plumbing. The +//! agent-side actors use these aliases instead so their code is written once; +//! the feature still selects the same underlying implementation. + +/// Cloneable sender half. `send` never blocks (the channel is unbounded) and +/// only fails once every receiver is gone. +#[cfg(not(feature = "crossfire-channel"))] +pub type MTx = tokio::sync::mpsc::UnboundedSender; +#[cfg(feature = "crossfire-channel")] +pub type MTx = crossfire::MTx; + +/// Receiver half. Wrapped so `recv` has one signature across both backends +/// (tokio yields `Option`, crossfire yields `Result`). +pub struct MRx { + #[cfg(not(feature = "crossfire-channel"))] + inner: tokio::sync::mpsc::UnboundedReceiver, + #[cfg(feature = "crossfire-channel")] + inner: crossfire::AsyncRx, +} + +impl MRx { + /// Receive the next message, or `None` once the channel is closed and drained. + pub async fn recv(&mut self) -> Option { + #[cfg(not(feature = "crossfire-channel"))] + { + self.inner.recv().await + } + #[cfg(feature = "crossfire-channel")] + { + self.inner.recv().await.ok() + } + } +} + +/// Create an unbounded channel. `Unpin` is required by the crossfire backend +/// and stated unconditionally so both builds accept the same message types. +pub fn unbounded() -> (MTx, MRx) { + #[cfg(not(feature = "crossfire-channel"))] + let (tx, inner) = tokio::sync::mpsc::unbounded_channel(); + #[cfg(feature = "crossfire-channel")] + let (tx, inner) = crossfire::mpsc::unbounded_async(); + (tx, MRx { inner }) +} diff --git a/netmito/src/client/http.rs b/netmito/src/client/http.rs index 96277804..c29788e1 100644 --- a/netmito/src/client/http.rs +++ b/netmito/src/client/http.rs @@ -1759,4 +1759,124 @@ impl MitoHttpClient { Err(get_error_from_resp(resp).await.into()) } } + + /// List a suite's jobs — one row per agent attempt at running it. + pub async fn query_suite_jobs( + &mut self, + suite_uuid: Uuid, + req: SuiteJobsQueryReq, + ) -> crate::error::Result { + let credential = self + .credential_guard + .get_credential() + .ok_or(CredentialGuardError::CredentialNotFound)? + .token; + self.url + .set_path(&format!("suites/{suite_uuid}/jobs/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()) + } + } + + /// One job of a suite, with its hook executions embedded. + pub async fn get_suite_job( + &mut self, + suite_uuid: Uuid, + job_id: i32, + ) -> crate::error::Result { + let credential = self + .credential_guard + .get_credential() + .ok_or(CredentialGuardError::CredentialNotFound)? + .token; + self.url + .set_path(&format!("suites/{suite_uuid}/jobs/{job_id}")); + let resp = self + .http_client + .get(self.url.as_str()) + .bearer_auth(credential) + .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 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()) + } + } + + /// Shut an agent down. The agent row is never deleted — it is marked + /// `Offline`; see `service::agent::user_shutdown_agent_by_uuid`. + 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)?; + // Clear the query so it does not leak into later requests reusing self.url. + self.url.set_query(None); + if resp.status().is_success() { + Ok(()) + } else { + Err(get_error_from_resp(resp).await.into()) + } + } } diff --git a/netmito/src/client/interactive.rs b/netmito/src/client/interactive.rs index d3e88988..13e62949 100644 --- a/netmito/src/client/interactive.rs +++ b/netmito/src/client/interactive.rs @@ -255,7 +255,7 @@ pub(crate) fn output_suite_info(info: &TaskSuiteInfo) { ); } -pub(crate) fn output_parsed_suite_info(info: &ParsedTaskSuiteInfo, assigned_agents: &[uuid::Uuid]) { +pub(crate) fn output_parsed_suite_info(info: &ParsedTaskSuiteInfo, eligible_agents: &[uuid::Uuid]) { tracing::info!("Suite UUID: {}", info.uuid); if let Some(ref name) = info.name { tracing::info!("Name: {}", name); @@ -292,11 +292,11 @@ pub(crate) fn output_parsed_suite_info(info: &ParsedTaskSuiteInfo, assigned_agen if let Some(completed) = info.completed_at { tracing::info!("Completed at {}", completed); } - if assigned_agents.is_empty() { - tracing::info!("Manually-included agents: None"); + if eligible_agents.is_empty() { + tracing::info!("Currently eligible agents: None"); } else { - tracing::info!("Manually-included agents:"); - for agent in assigned_agents { + tracing::info!("Currently eligible agents:"); + for agent in eligible_agents { tracing::info!(" > {}", agent); } } diff --git a/netmito/src/client/mod.rs b/netmito/src/client/mod.rs index b8180058..66c54610 100644 --- a/netmito/src/client/mod.rs +++ b/netmito/src/client/mod.rs @@ -895,6 +895,62 @@ impl MitoClient { } } + /// List a suite's jobs, or show one in full when `--job` is given. + async fn suites_jobs(&mut self, args: SuiteJobsArgs) { + let suite = args.uuid; + if let Some(job_id) = args.job { + match self.http_client.get_suite_job(suite, job_id).await { + Ok(resp) => { + tracing::info!( + "Job {} ({}), agent: {}, created {}, updated {}", + resp.info.job_id, + resp.info.state, + resp.info + .agent_uuid + .map(|u| u.to_string()) + .unwrap_or("[detached]".to_string()), + resp.info.created_at, + resp.info.updated_at, + ); + if resp.hooks.is_empty() { + tracing::info!("No hook executions recorded"); + } + for hook in resp.hooks { + tracing::info!( + " hook {} ({}): {}, exit status {}", + hook.hook_type, + hook.uuid, + hook.state, + hook.result + .map(|r| r.exit_status.to_string()) + .unwrap_or("[none]".to_string()), + ); + } + } + Err(e) => tracing::error!("{}", e), + } + return; + } + let req = SuiteJobsQueryReq::from(&args); + match self.http_client.query_suite_jobs(suite, req).await { + Ok(resp) => { + tracing::info!("Found {} jobs for suite {suite}", resp.count); + for job in resp.jobs { + tracing::info!( + "Job {} ({}), agent: {}, updated {}", + job.job_id, + job.state, + job.agent_uuid + .map(|u| u.to_string()) + .unwrap_or("[detached]".to_string()), + job.updated_at, + ); + } + } + Err(e) => tracing::error!("{}", e), + } + } + pub async fn tasks_batch_cancel( &mut self, args: CancelTasksArgs, @@ -2051,7 +2107,7 @@ impl MitoClient { }, }, ClientCommand::Suites(args) => match args.command { - SuitesCommands::Create(args) => match self.suites_create(args).await { + SuitesCommands::Create(args) => match self.suites_create(*args).await { Ok(resp) => { tracing::info!("Suite created with uuid {}", resp.uuid); } @@ -2118,6 +2174,43 @@ impl MitoClient { SuitesCommands::Override(args) => { self.suites_override_agents(args).await; } + SuitesCommands::Jobs(args) => self.suites_jobs(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/agent.rs b/netmito/src/config/agent.rs new file mode 100644 index 00000000..f015cd8f --- /dev/null +++ b/netmito/src/config/agent.rs @@ -0,0 +1,181 @@ +use std::collections::HashSet; +use std::ops::Not; +use std::time::Duration; + +use clap::Args; +use figment::{ + providers::{Env, Format, Serialized, Toml}, + value::magic::RelativePathBuf, + Figment, +}; +use serde::{Deserialize, Serialize}; +use url::Url; + +use super::coordinator::DEFAULT_COORDINATOR_ADDR; + +#[derive(Deserialize, Serialize, Debug)] +pub struct AgentConfig { + pub(crate) coordinator_addr: Url, + pub(crate) credential_path: Option, + pub(crate) user: Option, + pub(crate) password: Option, + /// Group granted Admin over the agent — the role that may shut it down. The + /// registering user must be an Admin of it. Defaults to their own group. + #[serde(default)] + pub(crate) admin_group: Option, + /// Groups the agent joins; each gains Write access to it. + pub(crate) groups: HashSet, + /// Tags a suite is matched against. A suite is eligible when the agent + /// carries every tag the suite asks for. + pub(crate) tags: HashSet, + /// Free-form labels for querying; not used for matching. + pub(crate) labels: HashSet, + #[serde(with = "humantime_serde")] + pub(crate) heartbeat_interval: Duration, + /// How long to wait before retrying a coordinator call that could not + /// connect. + #[serde(with = "humantime_serde")] + pub(crate) connect_retry_interval: Duration, + /// How often an idle agent asks the coordinator for a suite rather than + /// waiting to be told about one. Unset disables polling and agent only + /// picks up suite after a heartbeat or upon a WebSocket notification. + #[serde(with = "humantime_serde")] + pub(crate) idle_poll_interval: Option, + /// How long to wait before reopening a notification WebSocket that dropped. + #[serde(with = "humantime_serde")] + pub(crate) ws_reconnect_interval: Duration, + #[serde(default)] + pub(crate) no_ws: bool, + /// The lifetime of the agent token. `None` means the token never expires. + #[serde(default, with = "humantime_serde")] + pub(crate) lifetime: Option, + /// Explicit machine code. When unset the agent resolves one from its cache, + /// `/etc/machine-id`, or a freshly generated value. + #[serde(default)] + pub(crate) machine_code: Option, +} + +#[derive(Args, Debug, Serialize, Default, Clone)] +#[command(rename_all = "kebab-case")] +pub struct AgentConfigCli { + /// The path of the config file + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub config: Option, + /// The address of the coordinator + #[arg(short, long = "coordinator")] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub coordinator_addr: Option, + /// The path of the user credential file + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub credential_path: Option, + /// The username of the user + #[arg(short, long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub user: Option, + /// The password of the user + #[arg(short, long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub password: Option, + /// The group granted Admin over the agent, which may shut it down. Defaults + /// to the registering user's own group + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub admin_group: Option, + /// The groups to join + #[arg(short, long, num_args = 0.., value_delimiter = ',')] + #[serde(skip_serializing_if = "::std::vec::Vec::is_empty")] + pub groups: Vec, + /// The tags used to match task suites + #[arg(short, long, num_args = 0.., value_delimiter = ',')] + #[serde(skip_serializing_if = "::std::vec::Vec::is_empty")] + pub tags: Vec, + /// The labels for identification + #[arg(short, long, num_args = 0.., value_delimiter = ',')] + #[serde(skip_serializing_if = "::std::vec::Vec::is_empty")] + pub labels: Vec, + /// The interval to send heartbeats (e.g. "30s", "1m") + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub heartbeat_interval: Option, + /// The interval between retries of a coordinator call that could not + /// connect (e.g. "30s") + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub connect_retry_interval: Option, + /// The interval for an idle agent to poll for a suite (e.g. "5s"). Waits to + /// be notified of one instead if unset + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub idle_poll_interval: Option, + /// The interval before reopening a dropped notification WebSocket (e.g. "5s") + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub ws_reconnect_interval: Option, + /// Whether to take notifications from the heartbeat only, without opening + /// the notification WebSocket + #[arg(long)] + #[serde(skip_serializing_if = "<&bool>::not")] + pub no_ws: bool, + /// The lifetime of the agent token (e.g., 7d, 1year). If not given, the agent token is valid forever + #[arg(long, value_parser = humantime_serde::re::humantime::parse_duration)] + #[serde( + with = "humantime_serde", + skip_serializing_if = "::std::option::Option::is_none" + )] + pub lifetime: Option, + /// Explicit machine code (overrides /etc/machine-id auto-detection) + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub machine_code: Option, + /// Run one suite and exit instead of staying up for more work. Intended for + /// tests and one-shot batches. + #[arg(long)] + #[serde(skip_serializing_if = "<&bool>::not")] + pub run_once: bool, +} + +impl Default for AgentConfig { + fn default() -> Self { + Self { + coordinator_addr: Url::parse(&format!("http://{DEFAULT_COORDINATOR_ADDR}")).unwrap(), + credential_path: None, + user: None, + password: None, + admin_group: None, + groups: HashSet::new(), + tags: HashSet::new(), + labels: HashSet::new(), + heartbeat_interval: Duration::from_secs(60), + connect_retry_interval: Duration::from_secs(30), + idle_poll_interval: None, + ws_reconnect_interval: Duration::from_secs(5), + no_ws: false, + lifetime: None, + machine_code: None, + } + } +} + +impl AgentConfig { + pub fn new(cli: &AgentConfigCli) -> crate::error::Result { + let global_config = dirs::config_dir().map(|mut p| { + p.push("mitosis"); + p.push("config.toml"); + p + }); + let mut figment = Figment::new().merge(Serialized::from(Self::default(), "agent")); + if let Some(global_config) = global_config { + if global_config.exists() { + figment = figment.merge(Toml::file(global_config).nested()); + } + } + figment = figment + .merge(Toml::file(cli.config.as_deref().unwrap_or("config.toml")).nested()) + .merge(Env::prefixed("MITO_").profile("agent")) + .merge(Serialized::from(cli, "agent")) + .select("agent"); + Ok(figment.extract()?) + } +} diff --git a/netmito/src/config/client/agents.rs b/netmito/src/config/client/agents.rs new file mode 100644 index 00000000..68728e7d --- /dev/null +++ b/netmito/src/config/client/agents.rs @@ -0,0 +1,73 @@ +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 itself is kept) + Shutdown(ShutdownAgentArgs), +} + +#[derive(Serialize, Debug, Deserialize, Args, Clone)] +pub struct QueryAgentsArgs { + /// Filter by group name (defaults to your username when omitted) + #[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 (for pagination) + #[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, + /// Stop now: the agent's in-flight job is killed without cleanup and its + /// uncommitted tasks are reclaimed. Without this the agent finishes its + /// current job first. + #[arg(short, long)] + pub force: bool, +} diff --git a/netmito/src/config/client/mod.rs b/netmito/src/config/client/mod.rs index 0d221916..390afe60 100644 --- a/netmito/src/config/client/mod.rs +++ b/netmito/src/config/client/mod.rs @@ -9,13 +9,14 @@ use std::ops::Not; use url::Url; use crate::entity::content::ArtifactContentType; -use crate::schema::{RemoteResource, RemoteResourceDownload}; +use crate::schema::{ExecSpec, RemoteResource, RemoteResourceDownload}; use std::path::PathBuf; 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::*; @@ -111,6 +113,8 @@ pub enum ClientCommand { Tasks(TasksArgs), /// Manage task suites, including creating, querying, and assigning agents. Suites(SuitesArgs), + /// Manage agents, including querying them and shutting one down. + Agents(AgentsArgs), /// Manage workers, including querying workers, cancel workers, etc. Workers(WorkersArgs), /// Run an external command @@ -158,6 +162,20 @@ fn parse_artifact_content_type( } } +/// Parse a hook's `ExecSpec` from a JSON object, the same shape `POST /suites` +/// takes: `{"args":["sh","-c","./setup.sh"],"terminal_output":true}`. Only +/// `args` is required. +fn parse_exec_spec( + s: &str, +) -> Result> { + let spec: ExecSpec = serde_json::from_str(s) + .map_err(|e| format!("invalid hook spec: {e}. Expected a JSON object such as {{\"args\":[\"sh\",\"-c\",\"./setup.sh\"],\"terminal_output\":true}}"))?; + if spec.args.is_empty() { + return Err("invalid hook spec: `args` must not be empty".into()); + } + Ok(spec) +} + /// Parse a resource string into RemoteResourceDownload fn parse_resources( s: &str, diff --git a/netmito/src/config/client/suites.rs b/netmito/src/config/client/suites.rs index e18ddc4a..a1db25ca 100644 --- a/netmito/src/config/client/suites.rs +++ b/netmito/src/config/client/suites.rs @@ -3,10 +3,15 @@ use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::{ - entity::state::TaskSuiteState, - schema::{CreateTaskSuiteReq, TaskSuitesQueryReq, WorkerSchedulePlan}, + entity::state::{SuiteJobState, TaskSuiteState}, + schema::{ + CreateTaskSuiteReq, ExecHooks, ExecSpec, SuiteJobsQueryReq, TaskSuitesQueryReq, + WorkerSchedulePlan, + }, }; +use super::parse_exec_spec; + #[derive(Serialize, Debug, Deserialize, Args, derive_more::From, Clone)] pub struct SuitesArgs { #[command(subcommand)] @@ -16,7 +21,8 @@ pub struct SuitesArgs { #[derive(Subcommand, Serialize, Debug, Deserialize, Clone)] pub enum SuitesCommands { /// Create a new task suite - Create(CreateSuiteArgs), + // Boxed: the three optional hook specs make this variant far larger than the rest. + Create(Box), /// Query task suites subject to a filter Query(QuerySuitesArgs), /// Get the details of a task suite @@ -27,6 +33,8 @@ pub enum SuitesCommands { Cancel(CancelSuiteArgs), /// Set agent overrides on a suite in one batch: include, exclude, and/or clear agents Override(AgentsForSuiteOverrideArgs), + /// Inspect the suite's jobs — one per agent attempt at running it + Jobs(SuiteJobsArgs), } #[derive(Serialize, Debug, Deserialize, Args, Clone)] @@ -55,10 +63,30 @@ pub struct CreateSuiteArgs { /// Number of tasks each worker prefetches locally #[arg(long, default_value_t = 16)] pub prefetch: u32, + /// Provision hook, as a JSON exec spec: '{"args":["sh","-c","./setup.sh"],"terminal_output":true}'. + /// Runs once before any task; a non-zero exit fails the job and its tasks never start + #[arg(long, value_parser = parse_exec_spec)] + pub provision: Option, + /// Cleanup hook, same JSON shape. Runs once after the tasks drain, even when the job already failed + #[arg(long, value_parser = parse_exec_spec)] + pub cleanup: Option, + /// Background hook, same JSON shape. Runs alongside the tasks and is expected to outlive them + #[arg(long, value_parser = parse_exec_spec)] + pub background: Option, } impl From for CreateTaskSuiteReq { fn from(args: CreateSuiteArgs) -> Self { + // A suite with no hook at all keeps `exec_hooks` absent rather than + // carrying three nulls. + let exec_hooks = match (args.provision, args.cleanup, args.background) { + (None, None, None) => None, + (provision, cleanup, background) => Some(ExecHooks { + provision, + cleanup, + background, + }), + }; Self { name: args.name, description: args.description, @@ -71,7 +99,7 @@ impl From for CreateTaskSuiteReq { cpu_binding: None, task_prefetch_count: args.prefetch, }, - exec_hooks: None, + exec_hooks, } } } @@ -157,3 +185,39 @@ pub struct AgentsForSuiteOverrideArgs { #[arg(long, num_args = 0.., value_delimiter = ',')] pub clear: Vec, } + +#[derive(Serialize, Debug, Deserialize, Args, Clone)] +pub struct SuiteJobsArgs { + /// The UUID of the suite + pub uuid: Uuid, + /// Show one job in full, including its hook executions + #[arg(long)] + pub job: Option, + /// Filter the listing by job state + #[arg(long, num_args = 0.., value_delimiter = ',')] + pub states: Vec, + /// Only list jobs run by this agent + #[arg(long)] + pub agent: Option, + /// Maximum number of jobs to list + #[arg(long)] + pub limit: Option, + /// Number of jobs to skip (for pagination) + #[arg(long)] + pub offset: Option, + /// Report the number of matching jobs instead of listing them + #[arg(long)] + pub count: bool, +} + +impl From<&SuiteJobsArgs> for SuiteJobsQueryReq { + fn from(args: &SuiteJobsArgs) -> Self { + Self { + states: (!args.states.is_empty()).then(|| args.states.iter().copied().collect()), + agent_uuid: args.agent, + limit: args.limit, + offset: args.offset, + count: args.count, + } + } +} diff --git a/netmito/src/config/coordinator.rs b/netmito/src/config/coordinator.rs index 1be8a1a3..8cf7e4bb 100644 --- a/netmito/src/config/coordinator.rs +++ b/netmito/src/config/coordinator.rs @@ -23,8 +23,11 @@ use tokio_util::sync::CancellationToken; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, Layer}; use crate::{ + channel::{MRx, MTx}, error::Error, + service::agent::heartbeat::{AgentHeartbeatOp, AgentHeartbeatQueue}, service::worker::{HeartbeatOp, HeartbeatQueue, TaskDispatcher, TaskDispatcherOp}, + ws::{AgentWsRouter, RouterOp}, }; use super::TracingGuard; @@ -60,10 +63,28 @@ pub struct CoordinatorConfig { pub(crate) access_token_expires_in: std::time::Duration, #[serde(with = "humantime_serde")] pub(crate) heartbeat_timeout: std::time::Duration, + /// How often to send a WebSocket keepalive to a connected agent, so idle + /// notification sockets are not culled by intermediate proxies. + #[serde(with = "humantime_serde", default = "default_ws_keepalive_interval")] + pub(crate) ws_keepalive_interval: std::time::Duration, + /// How long a suite may go without a new task before the coordinator sweeps + /// it out of `Open`. It is also how long an agent holds a drained job open + /// waiting for more work, so raising it trades idle machine time for fewer + /// provision hooks. + #[serde(with = "humantime_serde", default = "default_suite_auto_close_timeout")] + pub(crate) suite_auto_close_timeout: std::time::Duration, pub(crate) log_path: Option, pub(crate) file_log: bool, } +fn default_ws_keepalive_interval() -> std::time::Duration { + std::time::Duration::from_secs(30) +} + +fn default_suite_auto_close_timeout() -> std::time::Duration { + std::time::Duration::from_secs(120) +} + fn default_mitosis_region() -> String { "mitosis".to_string() } @@ -149,6 +170,15 @@ 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, + /// How long a suite may go without a new task before it is swept out of + /// Open, and how long an agent holds a drained job. Default 60 seconds + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub suite_auto_close_timeout: 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)] @@ -190,6 +220,8 @@ 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(), + suite_auto_close_timeout: default_suite_auto_close_timeout(), log_path: None, file_log: false, } @@ -236,6 +268,23 @@ impl CoordinatorConfig { HeartbeatQueue::new(cancel_token, self.heartbeat_timeout, pool, rx) } + pub fn build_agent_heartbeat_queue( + &self, + cancel_token: CancellationToken, + pool: InfraPool, + rx: MRx, + ) -> AgentHeartbeatQueue { + AgentHeartbeatQueue::new(cancel_token, self.heartbeat_timeout, pool, rx) + } + + pub fn build_ws_router( + &self, + cancel_token: CancellationToken, + rx: MRx, + ) -> AgentWsRouter { + AgentWsRouter::new(cancel_token, rx) + } + pub async fn build_redis_connection_info( &self, ) -> crate::error::Result> { @@ -310,6 +359,8 @@ impl CoordinatorConfig { #[cfg(feature = "crossfire-channel")] worker_heartbeat_queue_tx: crossfire::MTx< HeartbeatOp, >, + agent_heartbeat_queue_tx: MTx, + ws_router_tx: MTx, ) -> crate::error::Result { let db = sea_orm::Database::connect(&self.db_url).await?; let credential = Credentials::new( @@ -333,6 +384,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, }) } @@ -446,6 +501,13 @@ pub struct InfraPool { pub worker_heartbeat_queue_tx: UnboundedSender, #[cfg(feature = "crossfire-channel")] pub worker_heartbeat_queue_tx: crossfire::MTx, + pub agent_heartbeat_queue_tx: MTx, + pub ws_router_tx: MTx, + /// Identifies this coordinator process. Agents compare it against the + /// `boot_id` in a `CounterSync` to notice a restart and reset the + /// notification sequence they are tracking. + pub boot_uuid: uuid::Uuid, + pub ws_keepalive_interval: std::time::Duration, } #[derive(Debug)] diff --git a/netmito/src/config/mod.rs b/netmito/src/config/mod.rs index 748e374a..d116d9ae 100644 --- a/netmito/src/config/mod.rs +++ b/netmito/src/config/mod.rs @@ -1,7 +1,9 @@ +pub mod agent; pub mod client; pub mod coordinator; pub mod manager; pub mod worker; +pub use agent::{AgentConfig, AgentConfigCli}; pub use client::{ClientConfig, ClientConfigCli}; pub use coordinator::{CoordinatorConfig, CoordinatorConfigCli, InfraPool}; pub(crate) use coordinator::{ diff --git a/netmito/src/coordinator.rs b/netmito/src/coordinator.rs index 7c99b0bb..221aec6d 100644 --- a/netmito/src/coordinator.rs +++ b/netmito/src/coordinator.rs @@ -8,16 +8,38 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use crate::api::router; use crate::config::{CoordinatorConfig, CoordinatorConfigCli, InfraPool}; use crate::migration::{Migrator, MigratorTrait}; +use crate::service::agent::heartbeat::AgentHeartbeatQueue; use crate::service::s3::setup_buckets; +use crate::service::suite::sweep_inactive_suites; use crate::service::worker::{restore_workers, HeartbeatQueue, TaskDispatcher}; use crate::signal::shutdown_signal; +use crate::ws::AgentWsRouter; + +/// How often the idle-suite sweep runs, given the configured idle window. +/// +/// Half the window, so a drained suite lingers in `Open` at most about one and a +/// half windows past its last task — and an agent holds its job for no longer +/// than that. Scaling with the window is what keeps a small one meaningful: a +/// fixed period would swamp a 5-second window and waste queries on an hour-long +/// one. Clamped at both ends so neither extreme misbehaves. +fn suite_sweep_period(idle_window: std::time::Duration) -> std::time::Duration { + (idle_window / 2).clamp( + std::time::Duration::from_secs(1), + std::time::Duration::from_secs(30), + ) +} pub struct MitoCoordinator { pub infra_pool: InfraPool, pub worker_task_queue: TaskDispatcher, pub worker_heartbeat_queue: HeartbeatQueue, + pub agent_heartbeat_queue: AgentHeartbeatQueue, + pub ws_router: AgentWsRouter, pub cancel_token: CancellationToken, pub log_dir: PathBuf, + /// How long a suite may go without a new task before the sweep settles it + /// out of `Open`. + pub suite_auto_close_timeout: std::time::Duration, } impl MitoCoordinator { @@ -96,13 +118,21 @@ impl MitoCoordinator { let (worker_heartbeat_queue_tx, worker_heartbeat_queue_rx) = crossfire::mpsc::unbounded_async(); + let (agent_heartbeat_queue_tx, agent_heartbeat_queue_rx) = crate::channel::unbounded(); + let (ws_router_tx, ws_router_rx) = crate::channel::unbounded(); + // Setup worker task queue let worker_task_queue = config.build_worker_task_queue(cancel_token.clone(), worker_task_queue_rx); // 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 @@ -112,6 +142,15 @@ impl MitoCoordinator { worker_heartbeat_queue_rx, ); + // Setup the agent-side actors: liveness tracking and the 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); + let suite_auto_close_timeout = config.suite_auto_close_timeout; + // Setup s3 storage // List all buckets and create if not exist setup_buckets( @@ -137,8 +176,11 @@ impl MitoCoordinator { infra_pool, worker_task_queue, worker_heartbeat_queue, + agent_heartbeat_queue, + ws_router, cancel_token, log_dir, + suite_auto_close_timeout, }) } @@ -148,13 +190,45 @@ impl MitoCoordinator { infra_pool, mut worker_task_queue, mut worker_heartbeat_queue, + mut agent_heartbeat_queue, + mut ws_router, cancel_token, + suite_auto_close_timeout, .. } = self; // Create TaskTracker to manage background tasks let task_tracker = TaskTracker::new(); + // Settle idle suites out of `Open`. Jittered so a fleet of coordinators + // against one database does not sweep in lockstep. + { + let db = infra_pool.db.clone(); + let cancel = cancel_token.clone(); + let period = suite_sweep_period(suite_auto_close_timeout); + task_tracker.spawn(async move { + loop { + // Up to half a period of jitter, so several coordinators on + // one database do not all sweep on the same tick. + let delay = period + + std::time::Duration::from_millis(rand::Rng::random_range( + &mut rand::rng(), + 0..=(period.as_millis() as u64 / 2), + )); + tokio::select! { + biased; + _ = cancel.cancelled() => break, + _ = tokio::time::sleep(delay) => { + if let Err(e) = sweep_inactive_suites(&db, suite_auto_close_timeout).await { + tracing::error!("Failed to sweep inactive suites: {e}"); + } + } + } + } + tracing::info!("Suite sweep stopped"); + }); + } + // Spawn background tasks using TaskTracker task_tracker.spawn(async move { worker_task_queue.run().await; @@ -166,7 +240,18 @@ impl MitoCoordinator { tracing::info!("Heartbeat queue stopped"); }); + task_tracker.spawn(async move { + ws_router.run().await; + }); + + task_tracker.spawn(async move { + agent_heartbeat_queue.run().await; + }); + restore_workers(&infra_pool).await?; + // Agents keep their rows across a coordinator restart, so tell them this + // is a new boot and their notification sequence starts over. + crate::service::agent::notify_agents_of_restart(&infra_pool).await?; let app = router(infra_pool, cancel_token.clone()); let addr = crate::config::SERVER_CONFIG .get() diff --git a/netmito/src/entity/state.rs b/netmito/src/entity/state.rs index 35a68cd4..c3771074 100644 --- a/netmito/src/entity/state.rs +++ b/netmito/src/entity/state.rs @@ -347,6 +347,14 @@ impl TaskSuiteState { !matches!(self, Self::Cancelled) } + /// Returns true if the suite allows its tasks to be executed. + /// - Open, Closed: tasks keep running + /// - Complete: nothing is left to run, but asking is not an error + /// - Cancelled: terminal state, no further execution + pub fn allows_task_execution(&self) -> bool { + !matches!(self, Self::Cancelled) + } + // TODO: this method might be removed as we should do an idempotent update to state /// Returns true if the suite needs to be reopened before accepting tasks. /// This is true for Closed and Complete states. diff --git a/netmito/src/error.rs b/netmito/src/error.rs index 2e3b0c10..4bcba35f 100644 --- a/netmito/src/error.rs +++ b/netmito/src/error.rs @@ -135,6 +135,12 @@ pub enum ApiError { AlreadyExists(String), #[error("{0} not found")] NotFound(String), + /// The request is well-formed but the target has already moved past the + /// state it addressed. The agent protocol leans on this: a report against a + /// terminal job answers 409, which the agent reads as "the job is closed, + /// stop reporting and go idle". + #[error("Conflicting request: {0}")] + Conflict(String), #[error("Resource quota exceeded")] QuotaExceeded, #[error(transparent)] @@ -201,6 +207,7 @@ impl GetStatusCode for ApiError { ApiError::InvalidRequest(_) => StatusCode::BAD_REQUEST, ApiError::AlreadyExists(_) => StatusCode::CONFLICT, ApiError::NotFound(_) => StatusCode::NOT_FOUND, + ApiError::Conflict(_) => StatusCode::CONFLICT, ApiError::QuotaExceeded => StatusCode::FORBIDDEN, ApiError::PresignS3Error(_) => StatusCode::BAD_REQUEST, } diff --git a/netmito/src/executor.rs b/netmito/src/executor.rs new file mode 100644 index 00000000..1e0a40e0 --- /dev/null +++ b/netmito/src/executor.rs @@ -0,0 +1,921 @@ +//! The execution core: one flow, three report targets. +//! +//! Everything a running process needs from the outside world — input downloads, +//! dependency waits, state announcements and outcome reports — sits behind +//! [`ExecClient`]. [`execute`] itself knows nothing about workers, agents, tasks +//! or hooks, so the same code drives all three of: +//! +//! | impl | lives in | reports to | +//! |---|---|---| +//! | `WorkerTaskClient` | `worker.rs` | `POST /workers/tasks` | +//! | `AgentTaskClient` | `agent.rs` | `POST /agents/tasks/report` | +//! | `AgentHookClient` | `agent.rs` | `POST /agents/job/hook` | + +use std::os::unix::process::ExitStatusExt; +use std::path::PathBuf; +use std::process::ExitStatus; + +use async_compression::tokio::write::GzipEncoder; +use nix::sys::signal::{self, Signal}; +use nix::unistd::Pid; +use reqwest::header::CONTENT_LENGTH; +use reqwest::{Client, StatusCode}; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::process::Command; +use tokio::sync::mpsc; +use tokio::time::Instant; +use tokio_tar::{Builder, Header}; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +use crate::entity::content::ArtifactContentType; +use crate::entity::state::TaskExecState; +use crate::error::{Error, ErrorMsg, RequestError}; +use crate::schema::{ + ExecSpec, RemoteResource, RemoteResourceDownload, RemoteResourceDownloadResp, SubmitTaskReq, + TaskExecOptions, TaskResultMessage, TaskResultSpec, +}; +use crate::service::s3::download_file; + +/// How long the whole input-fetch phase may take. +const RESOURCE_FETCH_BUDGET: std::time::Duration = std::time::Duration::from_secs(1800); +/// How long a single input download may take. +const SINGLE_DOWNLOAD_BUDGET: std::time::Duration = std::time::Duration::from_secs(120); +/// How long the archive-and-upload phase may take once the process has exited. +const UPLOAD_BUDGET: std::time::Duration = std::time::Duration::from_secs(600); + +/// One unit of execution's coordinator-facing half. +/// +/// An impl bundles who is talking to the coordinator (which credential, which +/// endpoints) with what is being reported (which task, or which suite hook). It +/// is built per unit and owns that unit's identity, which is what lets a hook — +/// identified only by `{job, hook_type}` — use the same core as a task. +#[async_trait::async_trait] +pub trait ExecClient: Send { + /// How this unit is named in logs. + fn describe(&self) -> String; + + /// Identity exported into the process's environment, on top of the standard + /// `MITO_*` directory variables the core always sets. + fn exec_env(&self) -> Vec<(&'static str, String)>; + + /// Whether the process may hand back a child task through `MITO_NEW_TASK`. + /// False for hooks: the hook report endpoint has no submit operation. + fn supports_child_tasks(&self) -> bool { + false + } + + /// The process is done and its artifacts are about to be uploaded. + /// `finished` distinguishes a clean exit from a cancellation or timeout. + /// `result` is the outcome so far — the final message is not known yet, so + /// impls that can only report once should expect [`report_commit`] to refine + /// this. + /// + /// [`report_commit`]: ExecClient::report_commit + async fn report_finish( + &mut self, + finished: bool, + result: &TaskResultSpec, + ) -> crate::error::Result<()>; + + /// Presign an upload for one produced artifact. + async fn request_upload( + &mut self, + content_type: ArtifactContentType, + content_length: u64, + ) -> crate::error::Result; + + /// The unit's final result. Ends it. + async fn report_commit(&mut self, result: TaskResultSpec) -> crate::error::Result<()>; + + /// Submit a task the process asked to spawn. Only ever called when + /// [`supports_child_tasks`] is true. + /// + /// [`supports_child_tasks`]: ExecClient::supports_child_tasks + async fn submit_child_task(&mut self, req: SubmitTaskReq) -> crate::error::Result<()>; + + /// Authed request for an input artifact's presigned download info. + fn artifact_download_req( + &self, + uuid: Uuid, + content_type: ArtifactContentType, + ) -> reqwest::RequestBuilder; + + /// Authed request for an input attachment's presigned download info. + fn attachment_download_req(&self, key: &str) -> reqwest::RequestBuilder; + + /// Publish this unit's fine-grained exec state (worker: redis set+publish; + /// agent: nothing to publish to yet). `ex` is an expiry in seconds. + async fn announce_state(&mut self, state: TaskExecState, ex: Option); + + /// Whether a `watch` dependency can be resolved at all. When false the wait + /// is skipped and the unit runs immediately. + fn can_watch(&self) -> bool { + false + } + + /// Wait until `uuid` reaches `target`. Only called when [`can_watch`] is true. + /// + /// [`can_watch`]: ExecClient::can_watch + async fn watch(&mut self, uuid: &Uuid, target: TaskExecState) { + let _ = (uuid, target); + } +} + +/// Where one produced artifact should go. +pub enum UploadTarget { + /// PUT the artifact here. + Url(String), + /// Nothing to upload to; skip this artifact and carry on with the rest. + Skip, + /// The report target is closed, gone, or refusing writes. Abandon the + /// remaining uploads and go straight to committing the result. + Stop, +} + +/// The per-unit execution context: a working directory, a cancellation token, +/// and the client that reports it. +pub struct Executor { + /// Cancelled when the owning worker/agent is shutting down, or when the + /// coordinator has told us this unit's owner is closed. + pub cancel_token: CancellationToken, + /// How long to wait before retrying a failed coordinator call. + pub polling_interval: std::time::Duration, + /// The process's working set: `result/`, `exec/` and `resource/` live here, + /// as do the archives built from them. Never shared between concurrent + /// units. + pub cache_path: PathBuf, + /// Plain HTTP client for presigned S3 transfers. Carries no credential — + /// the coordinator-facing one lives in `client`. + pub http_client: Client, + pub client: Box, +} + +impl Executor { + async fn announce(&mut self, state: TaskExecState) { + self.client.announce_state(state, None).await + } + + async fn announce_ex(&mut self, state: TaskExecState, ex: u64) { + self.client.announce_state(state, Some(ex)).await + } + + /// Report a unit that never ran: unsuccessful, zero exit status, and a + /// message saying why. The caller announces the states around it. + async fn report_abandoned(&mut self, msg: TaskResultMessage) -> crate::error::Result<()> { + let result = TaskResultSpec { + exit_status: 0, + msg: Some(msg), + }; + self.client.report_finish(false, &result).await?; + self.client.report_commit(result).await + } +} + +/// Create (or recreate, discarding whatever the last unit left) the working +/// directories a unit's process expects. +pub async fn reset_workspace(cache_path: &std::path::Path) -> crate::error::Result<()> { + let _ = tokio::fs::remove_dir_all(cache_path).await; + tokio::fs::create_dir_all(cache_path.join("result")).await?; + tokio::fs::create_dir_all(cache_path.join("exec")).await?; + tokio::fs::create_dir_all(cache_path.join("resource")).await?; + Ok(()) +} + +enum ProcessOutput { + WithLog { + stdout: Vec, + stderr: Vec, + exit_status: ExitStatus, + }, + WithoutLog { + exit_status: ExitStatus, + }, +} + +impl ProcessOutput { + fn get_exit_status(&self) -> ExitStatus { + match self { + ProcessOutput::WithLog { exit_status, .. } => *exit_status, + ProcessOutput::WithoutLog { exit_status } => *exit_status, + } + } +} + +enum ExecResult { + Finish(ProcessOutput), + Timeout(ProcessOutput), +} + +impl ExecResult { + fn state(&self) -> (bool, ExitStatus) { + match self { + ExecResult::Finish(output) => (true, output.get_exit_status()), + ExecResult::Timeout(output) => (false, output.get_exit_status()), + } + } + + fn get_output(self) -> ProcessOutput { + match self { + ExecResult::Finish(output) => output, + ExecResult::Timeout(output) => output, + } + } +} + +enum ResourceError { + /// 404: the resource does not exist. + NotFound, + /// 403: the resource is forbidden. + ForbiddenStatus, + /// The presigned S3 download itself failed. + DownloadFailed, + /// Exceeded the per-download or the overall fetch budget. + Timeout, + /// Cancelled via the shutdown token. + Cancelled, + /// Connection error or unexpected status — propagate to the caller. + Other(Error), +} + +/// Fetch one input to its local path, with connection retry, per-download and +/// overall budgets, and cancellation. Touches no unit lifecycle: every failure +/// comes back as a [`ResourceError`] for [`execute`] to turn into an outcome. +async fn download_resource( + executor: &mut Executor, + resource: RemoteResourceDownload, + timeout_until: Instant, +) -> std::result::Result<(), ResourceError> { + let resp = loop { + // The request (URL + credential) is the only worker/agent-specific part; + // the retry, timeout and status handling below are shared. + let req = match &resource.remote_file { + RemoteResource::Artifact { uuid, content_type } => { + executor.client.artifact_download_req(*uuid, *content_type) + } + RemoteResource::Attachment { key } => executor.client.attachment_download_req(key), + }; + match req.send().await { + Ok(resp) => break resp, + Err(e) => { + if e.is_connect() && e.is_request() { + tracing::error!( + "Fetch resource info failed with connection error: {}. Retry after {:?}", + e, + executor.polling_interval + ); + tokio::select! { + biased; + _ = executor.cancel_token.cancelled() => return Err(ResourceError::Cancelled), + _ = tokio::time::sleep(executor.polling_interval) => {}, + _ = tokio::time::sleep_until(timeout_until) => { + return Err(ResourceError::Timeout); + } + } + continue; + } else { + return Err(ResourceError::Other(RequestError::from(e).into())); + } + } + } + }; + if resp.status().is_success() { + let download_resp = resp + .json::() + .await + .map_err(|e| ResourceError::Other(RequestError::from(e).into()))?; + let local_path = executor + .cache_path + .join("resource") + .join(resource.local_path); + tokio::select! { + biased; + res = download_file(&executor.http_client, &download_resp, local_path, false) => { + if let Err(e) = res { + tracing::error!("Failed to download resource: {}", e); + return Err(ResourceError::DownloadFailed); + } + Ok(()) + } + _ = executor.cancel_token.cancelled() => Err(ResourceError::Cancelled), + _ = tokio::time::sleep(SINGLE_DOWNLOAD_BUDGET) => Err(ResourceError::Timeout), + _ = tokio::time::sleep_until(timeout_until) => Err(ResourceError::Timeout), + } + } else if resp.status() == StatusCode::NOT_FOUND { + Err(ResourceError::NotFound) + } else if resp.status() == StatusCode::FORBIDDEN { + Err(ResourceError::ForbiddenStatus) + } else { + let resp: ErrorMsg = resp + .json() + .await + .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); + Err(ResourceError::Other(Error::Custom(format!( + "Fetch resource info failed with error: {}", + resp.msg + )))) + } +} + +/// Run one unit to completion and report it. +/// +/// `Ok(())` means "the unit was handled", including the cases where it was +/// abandoned (a missing input, a timeout) and committed as such. An `Err` is an +/// infrastructure failure the caller should treat as fatal for its own loop. +pub async fn execute( + executor: &mut Executor, + spec: ExecSpec, + exec_options: Option<&TaskExecOptions>, +) -> crate::error::Result<()> { + executor + .announce_ex(TaskExecState::FetchResource, 360) + .await; + let timeout_until = tokio::time::Instant::now() + RESOURCE_FETCH_BUDGET; + for resource in spec.resources { + match download_resource(executor, resource, timeout_until).await { + Ok(()) => {} + // Shutdown while fetching — just stop; teardown reclaims the unit. + Err(ResourceError::Cancelled) => return Ok(()), + // Connection error or unexpected status — surface as a hard error. + Err(ResourceError::Other(e)) => { + executor + .announce_ex(TaskExecState::FetchResourceError, 60) + .await; + return Err(e); + } + Err(kind) => { + let (state, msg) = match kind { + ResourceError::NotFound => ( + TaskExecState::FetchResourceNotFound, + TaskResultMessage::ResourceNotFound, + ), + ResourceError::ForbiddenStatus | ResourceError::DownloadFailed => ( + TaskExecState::FetchResourceForbidden, + TaskResultMessage::ResourceForbidden, + ), + _ => ( + TaskExecState::FetchResourceTimeout, + TaskResultMessage::FetchResourceTimeout, + ), + }; + tracing::debug!( + "Input unavailable for {}, commit it as cancelled", + executor.client.describe() + ); + executor.report_abandoned(msg).await?; + executor.announce_ex(state, 60).await; + executor.announce_ex(TaskExecState::TaskCommitted, 60).await; + return Ok(()); + } + } + } + + if let Some((watched_uuid, watched_state)) = exec_options.and_then(|opts| opts.watch) { + // Wait for another task to reach a state before running this one. + if executor.client.can_watch() { + executor.announce(TaskExecState::Watch).await; + let cancel_token = executor.cancel_token.clone(); + tokio::select! { + biased; + _ = cancel_token.cancelled() => { + tracing::info!("Watching interrupted by shutdown signal"); + executor.announce_ex(TaskExecState::WorkerExited, 60).await; + return Ok(()); + }, + _ = executor.client.watch(&watched_uuid, watched_state) => {}, + _ = tokio::time::sleep_until(timeout_until) => { + tracing::debug!("Watching timeout, commit this task as cancelled"); + executor.report_abandoned(TaskResultMessage::WatchTimeout).await?; + executor.announce_ex(TaskExecState::WatchTimeout, 60).await; + executor.announce_ex(TaskExecState::TaskCommitted, 60).await; + return Ok(()); + } + } + } + } + + // No timeout means no deadline + let exec_timeout = spec + .timeout + .and_then(|t| u64::try_from(t).ok()) + .map(std::time::Duration::from_secs); + match exec_timeout { + // Outlive the deadline by a little, so a watcher sees `ExecPending` for + // the whole run rather than losing it just before the timeout fires. + Some(timeout) => { + executor + .announce_ex(TaskExecState::ExecPending, timeout.as_secs() + 60) + .await + } + // Nothing to expire against: the state stands until the next announce. + None => executor.announce(TaskExecState::ExecPending).await, + } + + // Setup the child-task hand-off file and clean up any stale one + let new_task_path = executor.cache_path.join("new_task.json"); + let _ = tokio::fs::remove_file(&new_task_path).await; // Ignore errors if file doesn't exist + + let mut command = Command::new("/usr/bin/env"); + command + .args(spec.args) + .envs(spec.envs) + .envs(executor.client.exec_env()) + .env("MITO_RESULT_DIR", executor.cache_path.join("result")) + .env("MITO_EXEC_DIR", executor.cache_path.join("exec")) + .env("MITO_RESOURCE_DIR", executor.cache_path.join("resource")) + .stdin(std::process::Stdio::null()); + if executor.client.supports_child_tasks() { + command.env("MITO_NEW_TASK", &new_task_path); + } + if spec.terminal_output { + command + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()); + } else { + command + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()); + } + + let mut child = command.spawn().inspect_err(|e| { + tracing::error!("Failed to spawn {}: {}", executor.client.describe(), e); + })?; + executor.announce(TaskExecState::ExecSpawned).await; + // A deadline only when one was asked for; otherwise a future that never + // resolves, so the arm below simply never fires. + let exec_deadline = async move { + match exec_timeout { + Some(timeout) => tokio::time::sleep(timeout).await, + None => std::future::pending().await, + } + }; + let process_output = async { + if spec.terminal_output { + let process_output = async { + let mut stdout_buf = Vec::new(); + let mut stdout = child.stdout.take().unwrap(); + let mut stderr_buf = Vec::new(); + let mut stderr = child.stderr.take().unwrap(); + tokio::try_join!( + stdout.read_to_end(&mut stdout_buf), + stderr.read_to_end(&mut stderr_buf), + child.wait() + ) + .map(|(_, _, exit_status)| ProcessOutput::WithLog { + stdout: stdout_buf, + stderr: stderr_buf, + exit_status, + }) + }; + process_output.await + } else { + child + .wait() + .await + .map(|exit_status| ProcessOutput::WithoutLog { exit_status }) + } + }; + + let output = tokio::select! { + biased; + _ = executor.cancel_token.cancelled() => { + tracing::info!("Execution interrupted by shutdown signal"); + child.kill().await.inspect_err(|e| { + tracing::error!("Failed to kill the process: {}", e); + })?; + executor.announce_ex(TaskExecState::WorkerExited, 60).await; + return Ok(()); + }, + output = process_output => { + executor.announce_ex(TaskExecState::ExecFinished, 660).await; + output.map(ExecResult::Finish) + }, + _ = exec_deadline => { + tracing::debug!("Execution timeout"); + executor.announce_ex(TaskExecState::ExecTimeout, 60).await; + if let Some(id) = child.id() { + // TODO: we may change this when once the `linux_pidfd` is stabilized in standard library + // Tracking issue for std lib: [rust-lang/rust #82971](https://github.com/rust-lang/rust/issues/82971) + // Tracking issue for tokio: [tokio-rs/tokio #6281](https://github.com/tokio-rs/tokio/issues/6281) + let _ = signal::kill(Pid::from_raw(id as i32), Signal::SIGTERM).inspect_err(|e| { + tracing::error!("Failed to send SIGTERM to the process: {}", e); + }); + } + tokio::select! { + biased; + _ = child.wait() => {}, + _ = executor.cancel_token.cancelled() => { + child.kill().await.inspect_err(|e| { + tracing::error!("Failed to kill the process: {}", e); + })?; + return Ok(()); + }, + _ = tokio::time::sleep(std::time::Duration::from_secs(10)) => { + child.kill().await.inspect_err(|e| { + tracing::error!("Failed to kill the process: {}", e); + })?; + }, + } + if spec.terminal_output { + let output = child.wait_with_output().await?; + Ok(ExecResult::Timeout(ProcessOutput::WithLog { + stdout: output.stdout, + stderr: output.stderr, + exit_status: output.status, + })) + } else { + let exit_status = child.wait().await?; + Ok(ExecResult::Timeout(ProcessOutput::WithoutLog { + exit_status, + })) + } + }, + }?; + tracing::debug!("Execution of {} finished", executor.client.describe()); + executor.announce_ex(TaskExecState::UploadResult, 660).await; + process_exec_result(executor, output).await?; + Ok(()) +} + +/// Archive whatever the process produced, upload it, then commit the outcome. +async fn process_exec_result( + executor: &mut Executor, + output: ExecResult, +) -> crate::error::Result<()> { + let (is_finished, exit_status) = output.state(); + executor + .announce_ex( + if is_finished { + TaskExecState::UploadFinishedResult + } else { + TaskExecState::UploadCancelledResult + }, + 660, + ) + .await; + executor + .client + .report_finish( + is_finished, + &TaskResultSpec { + exit_status: exit_status.into_raw(), + msg: (!is_finished).then_some(TaskResultMessage::ExecTimeout), + }, + ) + .await?; + + // Compress possible output and upload + let (tx, mut rx) = mpsc::channel::<(ArtifactContentType, u64)>(3); + // Spawn a task to archive the output + let timeout_cancel_token = CancellationToken::new(); + let archive_timeout_cancel_token = timeout_cancel_token.clone(); + let archive_cancel_token = executor.cancel_token.clone(); + let archive_cache_path = executor.cache_path.clone(); + let archive_hd = tokio::spawn(async move { + let result_dir = archive_cache_path.join("result"); + if !result_dir + .read_dir() + .map(|mut dir| dir.next().is_none()) + .unwrap_or(true) + { + let tar_file = + tokio::fs::File::create(archive_cache_path.join("result.tar.gz")).await?; + let encoder = GzipEncoder::new(tar_file); + let mut ar = Builder::new(encoder); + let compress_task = async { + ar.append_dir_all("result", result_dir).await?; + std::io::Result::Ok(()) + }; + tokio::select! { + biased; + _ = archive_cancel_token.cancelled() => { + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + tokio::fs::remove_file(archive_cache_path.join("result.tar.gz")).await?; + tracing::info!("Output generation interrupted by shutdown signal"); + return Ok(()); + } + _ = archive_timeout_cancel_token.cancelled() => { + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + tokio::fs::remove_file(archive_cache_path.join("result.tar.gz")).await?; + tracing::warn!("Output generation timeout"); + return Ok(()); + } + res = compress_task => { + match res { + Ok(_) => { + ar.finish().await?; + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + let file = encoder.into_inner(); + let size = file.metadata().await?.len(); + if let Err(e) = tx.send((ArtifactContentType::Result, size)).await { + tracing::error!("Failed to send result size: {}", e); + archive_cancel_token.cancel(); + return Ok(()); + } + } + Err(e) => { + tracing::error!("Failed to compress result: {}", e); + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + archive_cancel_token.cancel(); + return Err(e); + } + } + + } + } + } + let exec_log_dir = archive_cache_path.join("exec"); + if !exec_log_dir + .read_dir() + .map(|mut dir| dir.next().is_none()) + .unwrap_or(true) + { + let file_name = ArtifactContentType::ExecLog.to_string(); + let tar_file = tokio::fs::File::create(archive_cache_path.join(&file_name)).await?; + let encoder = GzipEncoder::new(tar_file); + let mut ar = Builder::new(encoder); + let compress_task = async { + ar.append_dir_all("exec-log", exec_log_dir).await?; + std::io::Result::Ok(()) + }; + tokio::select! { + biased; + _ = archive_cancel_token.cancelled() => { + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + tokio::fs::remove_file(archive_cache_path.join(&file_name)).await?; + tracing::info!("Output generation interrupted by shutdown signal"); + return Ok(()); + } + _ = archive_timeout_cancel_token.cancelled() => { + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + tokio::fs::remove_file(archive_cache_path.join(&file_name)).await?; + tracing::warn!("Output generation timeout"); + return Ok(()); + } + res = compress_task => { + match res { + Ok(_) => { + ar.finish().await?; + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + let file = encoder.into_inner(); + let size = file.metadata().await?.len(); + if let Err(e) = tx.send((ArtifactContentType::ExecLog, size)).await { + tracing::error!("Failed to compress exec log: {}", e); + archive_cancel_token.cancel(); + return Ok(()); + } + } + Err(e) => { + tracing::error!("Failed to compress exec log: {}", e); + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + archive_cancel_token.cancel(); + return Err(e); + } + } + + } + } + } + if let ProcessOutput::WithLog { stdout, stderr, .. } = output.get_output() { + let tar_file = + tokio::fs::File::create(archive_cache_path.join("std-log.tar.gz")).await?; + let encoder = GzipEncoder::new(tar_file); + let mut ar = Builder::new(encoder); + let compress_task = async { + let mut header = Header::new_gnu(); + header.set_cksum(); + header.set_mode(436); + header.set_size(stdout.len() as u64); + ar.append_data(&mut header, "std-log/stdout.log", &*stdout) + .await?; + header.set_size(stderr.len() as u64); + ar.append_data(&mut header, "std-log/stderr.log", &*stderr) + .await?; + std::io::Result::Ok(()) + }; + tokio::select! { + biased; + _ = archive_cancel_token.cancelled() => { + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + tokio::fs::remove_file(archive_cache_path.join("std-log.tar.gz")).await?; + tracing::info!("Output generation interrupted by shutdown signal"); + return Ok(()); + } + _ = archive_timeout_cancel_token.cancelled() => { + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + tokio::fs::remove_file(archive_cache_path.join("std-log.tar.gz")).await?; + tracing::warn!("Output generation timeout"); + return Ok(()); + } + res = compress_task => { + match res { + Ok(_) => { + ar.finish().await?; + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + let file = encoder.into_inner(); + let size = file.metadata().await?.len(); + if let Err(e) = tx.send((ArtifactContentType::StdLog, size)).await { + tracing::error!("Failed to compress std log: {}", e); + archive_cancel_token.cancel(); + return Ok(()); + } + } + Err(e) => { + tracing::error!("Failed to compress std log: {}", e); + let mut encoder = ar.into_inner().await?; + encoder.shutdown().await?; + archive_cancel_token.cancel(); + return Err(e); + } + } + + } + } + } + Ok(()) + }); + let upload_artifact_fut = async { + while let Some((content_type, content_length)) = rx.recv().await { + let url = match executor + .client + .request_upload(content_type, content_length) + .await? + { + UploadTarget::Url(url) => url, + UploadTarget::Skip => continue, + UploadTarget::Stop => return Ok(()), + }; + loop { + let file = + tokio::fs::File::open(executor.cache_path.join(content_type.to_string())) + .await?; + let upload_file = executor + .http_client + .put(url.as_str()) + .header(CONTENT_LENGTH, content_length) + .body(file) + .send(); + let resp = tokio::select! { + biased; + _ = executor.cancel_token.cancelled() => { + tracing::info!("Upload failed with shutdown signal"); + return Ok(()); + } + _ = timeout_cancel_token.cancelled() => { + tracing::warn!("Upload failed with timeout"); + return Ok(()); + } + resp = upload_file => resp + }; + match resp { + Ok(resp) => { + if resp.status().is_success() { + break; + } else { + let status = resp.status(); + return Err(Error::Custom(format!( + "Upload failed with status code: {status}" + ))); + } + } + Err(e) => { + if e.is_connect() && e.is_request() { + tracing::error!( + "Upload failed with connection error: {}. Retry after {:?}", + e, + executor.polling_interval + ); + tokio::select! { + biased; + _ = executor.cancel_token.cancelled() => return Ok(()), + _ = timeout_cancel_token.cancelled() => { + tracing::warn!("Upload failed with timeout"); + return Ok(()); + } + _ = tokio::time::sleep(executor.polling_interval) => {}, + } + continue; + } else { + return Err(RequestError::from(e).into()); + } + } + } + } + } + crate::error::Result::Ok(()) + }; + let timeout_until = tokio::time::Instant::now() + UPLOAD_BUDGET; + tokio::select! { + biased; + _ = tokio::time::sleep_until(timeout_until) => { + tracing::warn!("Upload result timeout"); + timeout_cancel_token.cancel(); + executor + .client + .report_commit(TaskResultSpec { + exit_status: exit_status.into_raw(), + msg: Some(TaskResultMessage::UploadResultTimeout), + }) + .await?; + executor.announce_ex(TaskExecState::UploadResultTimeout, 60).await; + executor.announce_ex(TaskExecState::TaskCommitted, 60).await; + archive_hd.await??; + } + res = upload_artifact_fut => { + res?; + archive_hd.await??; + if executor.cancel_token.is_cancelled() { + tracing::info!("Execution interrupted by shutdown signal"); + executor.announce_ex(TaskExecState::WorkerExited, 60).await; + return Ok(()); + } + executor.announce_ex(TaskExecState::UploadResultFinished, 60).await; + + let child_task_submitted = submit_child_task_if_present(executor).await; + let msg = if is_finished { + if child_task_submitted { + None + } else { + Some(TaskResultMessage::SubmitNewTaskFailed) + } + } else { + Some(TaskResultMessage::ExecTimeout) + }; + executor + .client + .report_commit(TaskResultSpec { + exit_status: exit_status.into_raw(), + msg, + }) + .await?; + executor.announce_ex(TaskExecState::TaskCommitted, 60).await; + } + } + + // Uploaded and committed: leave a clean directory for the next unit. + reset_workspace(&executor.cache_path).await?; + Ok(()) +} + +/// Submit the task the process left behind in `new_task.json`, if it left one. +/// A false answer becomes a `SubmitNewTaskFailed` message on the unit's result. +async fn submit_child_task_if_present(executor: &mut Executor) -> bool { + if !executor.client.supports_child_tasks() { + return true; + } + let new_task_path = executor.cache_path.join("new_task.json"); + if !new_task_path.exists() { + return true; // No new task to submit + } + + // Read and parse the file + let new_task_content = match tokio::fs::read_to_string(&new_task_path).await { + Ok(content) if !content.trim().is_empty() => content, + Ok(_) => { + tracing::debug!("New task file exists but is empty, ignoring"); + let _ = tokio::fs::remove_file(&new_task_path).await; + return true; + } + Err(e) => { + tracing::warn!("Failed to read new task file: {}", e); + let _ = tokio::fs::remove_file(&new_task_path).await; + return false; + } + }; + + let submit_req: SubmitTaskReq = match serde_json::from_str(&new_task_content) { + Ok(req) => req, + Err(e) => { + tracing::warn!("Failed to parse new task JSON: {}", e); + let _ = tokio::fs::remove_file(&new_task_path).await; + return false; + } + }; + + tracing::debug!( + "Submitting a new task from {} to group '{}'", + executor.client.describe(), + submit_req.group_name + ); + let submitted = executor.client.submit_child_task(submit_req).await; + // Always clean up the file after processing (success or failure) + let _ = tokio::fs::remove_file(&new_task_path).await; + match submitted { + Ok(()) => true, + Err(e) => { + tracing::warn!("Failed to submit new task: {}", e); + false + } + } +} diff --git a/netmito/src/lib.rs b/netmito/src/lib.rs index 1f2f8ff0..21c6458f 100644 --- a/netmito/src/lib.rs +++ b/netmito/src/lib.rs @@ -1,15 +1,19 @@ +pub mod agent; pub mod api; +pub mod channel; pub mod client; pub mod config; pub mod coordinator; pub mod entity; pub mod error; +pub mod executor; pub mod manager; pub mod migration; 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..5a35fda3 --- /dev/null +++ b/netmito/src/schema/agent.rs @@ -0,0 +1,495 @@ +//! Wire types for the agent surface: fleet management (`/agents`, user-authed) +//! and the execution loop (`/agents/...`, agent-authed), plus the coordinator → +//! agent notification events carried over `/ws/agents`. +//! +//! Ported from `../mitosis-dev` with the run→job rename applied and the fields +//! our slim `suite_agent_jobs` row does not carry dropped. + +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::{ + content::ArtifactContentType, + hook_tasks::HookType, + state::{AgentState, HookExecState, SuiteJobState, TaskSuiteState}, +}; + +use super::exec::ExecHooks; +use super::suite::WorkerSchedulePlan; +use super::task::{ReportTaskOp, TaskResultSpec, WorkerTaskResp}; + +// ============================================================================ +// Fleet management (user-authed) +// ============================================================================ + +/// Request to register an agent. +/// +/// Registration is an **upsert keyed by `machine_code`**: our `machines` table +/// holds the FK back to `agents` and both `machine_code` and `agent_id` are +/// unique, so one machine has exactly one agent row for its whole life. A +/// re-registering machine (agent restart) reuses that row — its tags, labels, +/// groups and metadata are refreshed and a fresh token is minted. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegisterAgentReq { + /// Tags for suite matching (e.g., `["gpu", "linux", "cuda:11.8"]`) + #[serde(default)] + pub tags: HashSet, + /// Labels for querying/filtering (e.g., `["datacenter:us-west"]`) + #[serde(default)] + pub labels: HashSet, + /// The group granted Admin over the agent — the role that may shut it down. + /// The caller must be an Admin of it. Defaults to the registering user's + /// personal group, the one group every user is guaranteed to administer. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub admin_group: Option, + /// Groups to associate with (the agent gets the Write role in each) + #[serde(default)] + pub groups: HashSet, + /// Optional token lifetime. When unset the token never expires. + #[serde(default, with = "humantime_serde")] + pub lifetime: Option, + /// Stable identifier for the machine running this agent. The agent client + /// resolves one from: config override → cached value → `/etc/machine-id` → + /// a generated UUID (persisted to the cache). + pub machine_code: String, + /// Static metadata about the agent process (stored on the machine row). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub metadata: Option, +} + +/// Response after registering an agent +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RegisterAgentResp { + pub agent_uuid: Uuid, + pub token: String, + /// Current notification counter for this agent (start of sequence) + pub notification_counter: u64, + /// True if an existing agent row was reused (same `machine_code`). + pub reused: bool, +} + +/// Static metadata reported by the agent at registration time +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AgentMetadata { + /// Agent binary version (e.g. "0.6.8") + pub version: String, + /// Long version string (e.g. includes build metadata) + pub long_version: String, +} + +/// Query parameters for listing agents +#[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, +} + +/// Information about an agent +#[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, + /// Machine code of the host this agent runs on (from the `machines` row). + #[serde(skip_serializing_if = "Option::is_none")] + pub machine_code: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub metadata: Option, +} + +/// Response for agent query +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AgentsQueryResp { + pub count: u64, + pub agents: Vec, + pub group_name: String, +} + +/// Query parameters for `DELETE /agents/{uuid}` selecting the shutdown mode +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct AgentShutdownReq { + #[serde(default)] + pub op: AgentShutdownOp, +} + +/// Shutdown operation type. +/// +/// Neither variant deletes the agent row — an agent is a durable identity and +/// every FK to it is `RESTRICT`. Both mark the agent `Offline`; they differ in +/// what happens to an in-flight job. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq)] +pub enum AgentShutdownOp { + /// Ask the agent to stop after its current job finishes cleanly. + #[default] + #[serde(alias = "graceful")] + Graceful, + /// Stop now: in-flight jobs go `Killed` and their tasks are reclaimed. + #[serde(alias = "force")] + Force, +} + +// ============================================================================ +// Execution loop (agent-authed) +// ============================================================================ + +/// Request for agent heartbeat +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AgentHeartbeatReq { + /// Current agent state + pub state: AgentState, + /// Currently assigned suite UUID (if any) + #[serde(default, skip_serializing_if = "Option::is_none")] + pub assigned_suite_uuid: Option, + /// The last notification ID the agent has processed + #[serde(default)] + pub last_notification_id: u64, + /// Optional metrics + #[serde(default, skip_serializing_if = "Option::is_none")] + pub metrics: Option, +} + +/// Metrics reported by an agent alongside its heartbeat +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AgentMetrics { + pub active_workers: u32, + pub tasks_completed: u64, + pub tasks_failed: u64, +} + +/// Response for agent heartbeat: the notifications the agent has not seen yet. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct AgentHeartbeatResp { + pub notifications: Vec, +} + +/// Request body for `POST /agents/suite` +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct AcceptSuiteReq { + /// The suite the agent was notified about, if any. A preference, not a + /// demand: one that is gone, drained or no longer this agent's to run falls + /// back to the best available. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub suite_uuid: Option, +} + +/// Full specification of a task suite handed to an agent for execution +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TaskSuiteSpec { + pub uuid: Uuid, + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + pub group_name: String, + pub tags: Vec, + pub labels: Vec, + pub priority: i32, + pub worker_schedule: WorkerSchedulePlan, + #[serde(skip_serializing_if = "Option::is_none")] + pub exec_hooks: Option, + pub state: TaskSuiteState, + pub total_tasks: i32, + pub incomplete_tasks: i32, +} + +/// Response to `POST /agents/suite`: the claim and everything needed to run it. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AcceptSuiteResp { + /// Whether a suite was claimed. False is an ordinary answer (nothing + /// available, or this agent is already busy), not an error. + pub accepted: bool, + /// What to run. Present exactly when `accepted` is true. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub suite: Option, + /// Opaque job handle (`suite_agent_jobs.id`) the agent echoes on every + /// later job-scoped call. Present only when `accepted` is true. + #[serde(skip_serializing_if = "Option::is_none")] + pub job: Option, + /// Per-suite job number, for display/inspection. Present with `job`. + #[serde(skip_serializing_if = "Option::is_none")] + pub job_id: Option, + /// Why nothing was claimed, when `accepted` is false. + #[serde(skip_serializing_if = "Option::is_none")] + pub reason: Option, +} + +/// Request to report that provisioning finished and execution is starting +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StartJobReq { + /// Opaque job handle from `AcceptSuiteResp` + pub job: i64, +} + +/// Request to report the agent entering the cleanup phase +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct EnterCleanupReq { + /// Opaque job handle from `AcceptSuiteResp` + pub job: i64, +} + +/// Request to report job completion +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CompleteJobReq { + /// Opaque job handle from `AcceptSuiteResp` + pub job: i64, + /// What the agent did: finished cleanly, or failed with a reason. + pub outcome: SuiteJobOutcome, +} + +/// The agent's report of how its job ended. By design the agent reports only +/// what *it* did — `Lost`/`Killed` are coordinator decisions and are never +/// agent outcomes. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum SuiteJobOutcome { + /// Provision, execution, and cleanup all succeeded. + Completed, + /// The job failed; `reason` summarizes the failing phase and cause. Full + /// hook output lives in the corresponding `hook_tasks` row. + Failed { reason: JobFailureReason }, +} + +/// Machine-readable category for why a job terminated abnormally +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum JobFailureKind { + /// Provision hook failed + ProvisionFailed, + /// Background hook exited before the job finished + BackgroundExited, + /// Task execution phase failed + ExecutionError, + /// Cleanup hook failed + CleanupFailed, +} + +/// One-line, job-level summary of abnormal termination. Our slim job row has no +/// `failure_reason` column, so this is logged by the coordinator rather than +/// stored; full hook stdout/stderr lives in `hook_tasks.result`. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct JobFailureReason { + pub kind: JobFailureKind, + pub message: String, +} + +/// Response after completing a job +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CompleteJobResp { + /// Whether another suite is available for this agent immediately + pub next_suite_available: bool, +} + +/// Request to report a suite hook execution (`POST /agents/job/hook`). +/// Append-only: accepted even on a terminal job (a cleanup hook may legitimately +/// finish after the coordinator terminated the job). +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HookReportReq { + /// Opaque job handle from `AcceptSuiteResp`. + pub job: i64, + /// Which hook this report is for (provision / cleanup / background). + pub hook_type: HookType, + /// What to do: record the hook's result, or presign a log upload. + pub op: HookReportOp, +} + +/// Operation for a hook report. `Result` writes the `hook_tasks` row and must +/// precede `Upload`, which presigns an S3 PUT for that row's log. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum HookReportOp { + /// Record the hook's final result (state derived from `exit_status`). + Result(TaskResultSpec), + /// Presign an S3 upload for a hook artifact/log; returns the URL. Requires + /// the hook's `Result` to have been reported first. + Upload { + content_type: ArtifactContentType, + content_length: u64, + }, +} + +/// Response to a hook report: the hook's uuid (its artifact key), plus the +/// presigned URL for the `Upload` op. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HookReportResp { + pub hook_uuid: Uuid, + #[serde(skip_serializing_if = "Option::is_none")] + pub url: Option, +} + +/// Request to claim tasks from a suite for execution +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FetchTasksReq { + /// Suite UUID to fetch tasks from + pub suite_uuid: Uuid, + /// Maximum number of tasks to claim in this batch + #[serde(default = "default_fetch_count")] + pub max_count: u32, +} + +fn default_fetch_count() -> u32 { + 1 +} + +/// Response containing the claimed tasks (empty if the suite has none ready) +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FetchTasksResp { + pub tasks: Vec, + /// Whether the agent should keep waiting on this job rather than wind it + /// down. True while the suite is `Open`, which covers "drained, but not idle + /// long enough to be sure". An empty batch with this set means "come back + /// and ask again", so the provisioned environment stays warm. + #[serde(default)] + pub hold_job_open: bool, +} + +/// Request to report a task result. Mirrors the worker's `ReportTaskReq{id, op}` +/// with the suite `job` handle as the only agent-specific addition. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ReportAgentTaskReq { + /// Opaque job handle from `AcceptSuiteResp` — the job this task belongs to. + pub job: i64, + /// Internal id of the task being reported (from the fetched `WorkerTaskResp`). + pub id: i64, + /// Operation to perform + pub op: ReportTaskOp, +} + +/// Filter for `POST /suites/{uuid}/jobs/query`. +/// +/// The relationship between the fields is AND. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct SuiteJobsQueryReq { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub states: Option>, + /// Only jobs run by this agent. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub agent_uuid: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub limit: Option, + #[serde(default, skip_serializing_if = "Option::is_none")] + pub offset: Option, + /// Return the total number of matching jobs instead of the jobs themselves. + #[serde(default)] + pub count: bool, +} + +/// One job of a suite +#[derive(Debug, Clone, Serialize, Deserialize, FromQueryResult)] +pub struct SuiteJobInfo { + /// Per-suite job number (the user-facing key) + pub job_id: i32, + pub state: SuiteJobState, + #[serde(skip_serializing_if = "Option::is_none")] + pub agent_uuid: Option, + pub created_at: OffsetDateTime, + pub updated_at: OffsetDateTime, +} + +/// Response for `GET /suites/{uuid}/jobs` +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SuiteJobsQueryResp { + pub count: u64, + pub jobs: Vec, +} + +/// One hook execution of a job +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HookTaskInfo { + /// Also the key its artifacts are stored under in the `artifacts` table. + pub uuid: Uuid, + pub hook_type: HookType, + pub state: HookExecState, + #[serde(skip_serializing_if = "Option::is_none")] + pub result: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub started_at: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub completed_at: Option, +} + +/// Response for `GET /suites/{uuid}/jobs/{job_id}`: the job plus its hooks. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SuiteJobQueryResp { + pub info: SuiteJobInfo, + pub hooks: Vec, +} + +// ============================================================================ +// Coordinator → agent notifications (WebSocket, `speedy` binary frames) +// ============================================================================ + +/// A sequenced notification. The `id` lets an agent that reconnects (or falls +/// back to heartbeat catch-up) tell what it has already seen. +/// +/// `speedy` carries it over the socket; serde is still needed because the same +/// events come back as JSON in the heartbeat response. +#[derive(Debug, Clone, Serialize, Deserialize, Readable, Writable)] +pub struct WsNotificationEvent { + /// Monotonically increasing sequence ID, per agent + pub id: u64, + pub event: AgentNotification, +} + +/// Lightweight push notification prompting the agent to act over HTTP. +#[derive(Debug, Clone, Serialize, Deserialize, Readable, Writable)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum AgentNotification { + /// A suite is available for this agent to execute. + SuiteAvailable { + /// Optional hint about which suite (the agent still fetches) + suite_uuid: Option, + priority: i32, + }, + + /// Stop the current suite and pick up a higher-priority one. + /// + /// **Not emitted yet.** A running agent is never interrupted for now. The + /// variant exists so the signal can be turned on later without a wire + /// change; the agent already handles it by cancelling its job and targeting + /// the new suite. + PreemptSuite { + new_suite_uuid: Uuid, + new_priority: i32, + current_suite_uuid: Uuid, + }, + + /// The suite the agent is running was cancelled; stop and clean up. + SuiteCancelled { suite_uuid: Uuid, reason: String }, + + /// Specific tasks were cancelled; stop executing them if in progress. + TasksCancelled { task_uuids: Vec }, + + /// The agent should shut down. + Shutdown { graceful: bool }, + + /// Keepalive. + Ping { server_time: i64 }, + + /// Resync the notification counter (coordinator restart / wrap-around). + CounterSync { counter: u64, boot_id: Uuid }, +} + +/// Message from an agent to the coordinator over the WebSocket. +#[derive(Debug, Clone, Serialize, Deserialize, Readable, Writable)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum AgentWsMessage { + /// Acknowledge receipt of a notification (drops it from the replay buffer). + Ack { notification_id: u64 }, + /// Response to a `Ping`. + 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/schema/suite.rs b/netmito/src/schema/suite.rs index d97a9164..6ded0e36 100644 --- a/netmito/src/schema/suite.rs +++ b/netmito/src/schema/suite.rs @@ -172,7 +172,9 @@ pub struct ParsedTaskSuiteInfo { pub completed_at: Option, } -/// Detailed suite response: the suite plus the UUIDs of its assigned agents +/// Detailed suite response: the suite plus the UUIDs of the agents currently +/// eligible to run it — tag-matched plus manual includes, minus manual excludes, +/// computed at query time rather than stored #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TaskSuiteQueryResp { pub info: ParsedTaskSuiteInfo, diff --git a/netmito/src/service/agent/heartbeat.rs b/netmito/src/service/agent/heartbeat.rs new file mode 100644 index 00000000..efb1b994 --- /dev/null +++ b/netmito/src/service/agent/heartbeat.rs @@ -0,0 +1,195 @@ +//! Liveness tracking for agents. + +use std::{cmp::Reverse, time::Duration}; + +use priority_queue::PriorityQueue; +use sea_orm::prelude::*; +use tokio::time::Instant; +use tokio_util::sync::CancellationToken; + +use crate::{ + channel::MRx, + config::InfraPool, + entity::{ + agents as Agent, + state::{AgentState, SuiteJobState}, + }, + service::agent::job, +}; + +/// Ceiling on any single database call made while handling a timeout, so a +/// stalled database cannot wedge the actor. +const DB_TIMEOUT: Duration = Duration::from_secs(10); + +pub enum AgentHeartbeatOp { + /// Record a heartbeat, resetting the agent's deadline. + Heartbeat(i64), + /// Stop tracking the agent (it shut down, or was already marked offline). + Remove(i64), +} + +pub struct AgentHeartbeatQueue { + agents: PriorityQueue>, + cancel_token: CancellationToken, + heartbeat_timeout: Duration, + pool: InfraPool, + rx: MRx, +} + +impl AgentHeartbeatQueue { + pub fn new( + cancel_token: CancellationToken, + heartbeat_timeout: Duration, + pool: InfraPool, + rx: MRx, + ) -> 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); + } + } + } + + /// Time until the earliest deadline, or a full timeout when idle. + 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 gone before its heartbeat timeout fired"); + return Ok(()); + } + Ok(Err(e)) => return Err(e.into()), + Err(_) => { + tracing::warn!(agent_id, "Agent lookup timed out during heartbeat check"); + return Ok(()); + } + }; + + // An agent already parked Offline (e.g. a user shutdown) has had its + // jobs and tasks handled; nothing to redo. + if agent.state == AgentState::Offline { + return Ok(()); + } + + tracing::info!( + agent_id, + agent_uuid = %agent.uuid, + "Agent heartbeat timed out — marking offline and reclaiming its work" + ); + + let now = TimeDateTimeWithTimeZone::now_utc(); + match tokio::time::timeout(DB_TIMEOUT, mark_offline(&self.pool, agent_id, now)).await { + Ok(res) => res?, + Err(_) => { + tracing::warn!(agent_id, "Marking the agent offline timed out"); + return Ok(()); + } + } + + match tokio::time::timeout( + DB_TIMEOUT, + job::terminate_agent_jobs(&self.pool.db, agent_id, SuiteJobState::Lost, now), + ) + .await + { + Ok(res) => { + res?; + } + Err(_) => tracing::warn!(agent_id, "Marking the agent's jobs Lost timed out"), + } + + match tokio::time::timeout( + DB_TIMEOUT, + job::reclaim_agent_tasks(&self.pool, agent.uuid, now), + ) + .await + { + Ok(res) => { + res?; + } + Err(_) => tracing::warn!(agent_id, "Reclaiming the agent's tasks 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 { + 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"); + } +} + +/// Park an agent `Offline` and drop its suite assignment. +pub(crate) async fn mark_offline( + pool: &InfraPool, + agent_id: i64, + now: TimeDateTimeWithTimeZone, +) -> crate::error::Result<()> { + 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(&pool.db) + .await?; + Ok(()) +} diff --git a/netmito/src/service/agent/hook.rs b/netmito/src/service/agent/hook.rs new file mode 100644 index 00000000..2d3034da --- /dev/null +++ b/netmito/src/service/agent/hook.rs @@ -0,0 +1,163 @@ +//! `POST /agents/job/hook` — provision / cleanup / background hook reports. +//! +//! Append-only, and deliberately **not** guarded against a terminal job: the job +//! existing and belonging to this agent is checked, but a cleanup hook may +//! legitimately finish after the coordinator already terminated the job (a +//! cancel, for instance), and losing that record helps nobody. +//! +//! `Result` writes the `hook_tasks` row; `Upload` presigns an S3 PUT for a log +//! or artifact of a hook whose result is already recorded. Hook artifacts live +//! in the shared `artifacts` table keyed by `hook_tasks.uuid` — there is no +//! separate hook-artifact table. + +use sea_orm::sea_query::OnConflict; +use sea_orm::{prelude::*, Set, TransactionTrait}; + +use crate::config::InfraPool; +use crate::entity::{ + hook_tasks::{self as HookTasks, HookType}, + state::HookExecState, + task_suites as TaskSuites, +}; +use crate::error::{ApiError, Error, Result}; +use crate::schema::{ExecHooks, HookReportOp, HookReportResp}; +use crate::service::agent::job; +use crate::service::s3::reserve_artifact_upload; + +pub async fn agent_report_hook( + agent_id: i64, + job_handle: i64, + hook_type: HookType, + op: HookReportOp, + pool: &InfraPool, +) -> Result { + let now = TimeDateTimeWithTimeZone::now_utc(); + + let job_row = job::load_validate_job(&pool.db, job_handle, agent_id).await?; + let suite = TaskSuites::Entity::find_by_id(job_row.task_suite_id) + .one(&pool.db) + .await? + .ok_or_else(|| Error::ApiError(ApiError::NotFound("Suite of the job".to_string())))?; + + match op { + HookReportOp::Result(result_spec) => { + let state = if result_spec.exit_status == 0 { + HookExecState::Completed + } else { + HookExecState::Failed + }; + let result_json = serde_json::to_value(&result_spec)?; + + let row = HookTasks::ActiveModel { + uuid: Set(Uuid::new_v4()), + suite_agent_job_id: Set(job_handle), + hook_type: Set(hook_type), + spec: Set(hook_spec_snapshot(&suite, hook_type)), + state: Set(state), + result: Set(Some(result_json)), + started_at: Set(None), + completed_at: Set(Some(now)), + created_at: Set(now), + updated_at: Set(now), + ..Default::default() + }; + // Upsert on (job, hook_type): re-reporting a hook overwrites its + // outcome. + // `uuid` is deliberately absent from the update list: it is the key + // this hook's artifacts are already filed under. + let row = HookTasks::Entity::insert(row) + .on_conflict( + OnConflict::columns([ + HookTasks::Column::SuiteAgentJobId, + HookTasks::Column::HookType, + ]) + .update_columns([ + HookTasks::Column::Spec, + HookTasks::Column::State, + HookTasks::Column::Result, + HookTasks::Column::CompletedAt, + HookTasks::Column::UpdatedAt, + ]) + .to_owned(), + ) + .exec_with_returning(&pool.db) + .await?; + + Ok(HookReportResp { + hook_uuid: row.uuid, + url: None, + }) + } + + HookReportOp::Upload { + content_type, + content_length, + } => { + // The result must land first: it is what mints the hook uuid the + // artifact is keyed by. + let hook = HookTasks::Entity::find() + .filter(HookTasks::Column::SuiteAgentJobId.eq(job_handle)) + .filter(HookTasks::Column::HookType.eq(hook_type)) + .one(&pool.db) + .await? + .ok_or_else(|| { + Error::ApiError(ApiError::InvalidRequest( + "Should report the hook Result before uploading its artifacts".to_string(), + )) + })?; + + let hook_uuid = hook.uuid; + let group_id = suite.group_id; + let content_length = content_length as i64; + let pool_cloned = pool.clone(); + + let (_, url) = pool + .db + .transaction::<_, (bool, String), Error>(|txn| { + Box::pin(async move { + reserve_artifact_upload( + txn, + &pool_cloned, + group_id, + hook_uuid, + content_type, + content_length, + now, + ) + .await + }) + }) + .await?; + + Ok(HookReportResp { + hook_uuid, + url: Some(url), + }) + } + } +} + +/// Snapshot of the hook's spec from the suite definition, stored on the row for +/// diagnostics. Falls back to JSON `null` when absent or unparseable — the +/// column is NOT NULL, but jsonb accepts `null`, and a missing snapshot must +/// never cost us the outcome record. +fn hook_spec_snapshot(suite: &TaskSuites::Model, hook_type: HookType) -> serde_json::Value { + let Some(raw) = suite.exec_hooks.as_ref() else { + return serde_json::Value::Null; + }; + let hooks: ExecHooks = match serde_json::from_value(raw.clone()) { + Ok(hooks) => hooks, + Err(e) => { + tracing::warn!("Failed to parse the suite's exec_hooks for a hook snapshot: {e}"); + return serde_json::Value::Null; + } + }; + let chosen = match hook_type { + HookType::Provision => hooks.provision, + HookType::Cleanup => hooks.cleanup, + HookType::Background => hooks.background, + }; + chosen + .map(|spec| serde_json::to_value(spec).unwrap_or(serde_json::Value::Null)) + .unwrap_or(serde_json::Value::Null) +} diff --git a/netmito/src/service/agent/job.rs b/netmito/src/service/agent/job.rs new file mode 100644 index 00000000..b74c6739 --- /dev/null +++ b/netmito/src/service/agent/job.rs @@ -0,0 +1,234 @@ +//! `suite_agent_jobs` lifecycle helpers. +//! +//! A job row is the single authority for one attempt of an agent running a task +//! suite (accept → provision → execute → cleanup → terminal). These helpers +//! create the row, validate job-scoped agent reports against it, and write the +//! coordinator-owned terminals (`Lost`, `Killed`). + +use sea_orm::{prelude::*, QueryOrder, QuerySelect, Set}; +use uuid::Uuid; + +use crate::config::InfraPool; +use crate::entity::{ + active_tasks as ActiveTasks, + state::{SuiteJobState, TaskState}, + suite_agent_jobs as SuiteAgentJobs, +}; +use crate::error::{ApiError, Error, Result}; + +/// The non-terminal job states (a job is "in flight"). +pub const IN_FLIGHT: [SuiteJobState; 3] = [ + SuiteJobState::Provisioning, + SuiteJobState::Executing, + SuiteJobState::Cleanup, +]; + +/// Create the job row for an agent that just accepted a suite (`Provisioning`). +/// +/// `job_id` is allocated as `max(job_id) + 1` scoped to the suite. The caller +/// MUST hold an exclusive lock on the suite row (`.lock_exclusive()` in the same +/// transaction) so concurrent accepts serialize and the `(task_suite_id, job_id)` +/// unique index never trips. +pub async fn create_job( + db: &C, + suite_id: i64, + agent_id: i64, + now: TimeDateTimeWithTimeZone, +) -> Result { + let next_job_id = SuiteAgentJobs::Entity::find() + .filter(SuiteAgentJobs::Column::TaskSuiteId.eq(suite_id)) + .order_by_desc(SuiteAgentJobs::Column::JobId) + .one(db) + .await? + .map(|m| m.job_id + 1) + .unwrap_or(1); + + let job = SuiteAgentJobs::ActiveModel { + task_suite_id: Set(suite_id), + job_id: Set(next_job_id), + agent_id: Set(Some(agent_id)), + state: Set(SuiteJobState::Provisioning), + created_at: Set(now), + updated_at: Set(now), + ..Default::default() + }; + Ok(job.insert(db).await?) +} + +/// Look up a job by its opaque internal `id` and verify it belongs to the +/// authenticated agent. +/// +/// A missing job and one owned by another agent both surface as `NotFound` +pub async fn load_validate_job( + db: &C, + job: i64, + agent_id: i64, +) -> Result { + let model = SuiteAgentJobs::Entity::find_by_id(job).one(db).await?; + validate_job_owner(model, job, agent_id) +} + +/// [`load_validate_job`] with the row locked `FOR UPDATE`. +/// +/// The caller must be inside a transaction: the lock is what turns a later +/// check-then-write (`reject_if_terminal` plus an update) into something atomic +/// rather than advisory. Outside a transaction every statement commits on its +/// own and the lock is released immediately, which only *looks* like protection +/// — hence the separate entry point instead of a flag on the plain loader. +pub async fn load_validate_job_locked( + db: &C, + job: i64, + agent_id: i64, +) -> Result { + let model = SuiteAgentJobs::Entity::find_by_id(job) + .lock_exclusive() + .one(db) + .await?; + validate_job_owner(model, job, agent_id) +} + +fn validate_job_owner( + model: Option, + job: i64, + agent_id: i64, +) -> Result { + let model = + model.ok_or_else(|| Error::ApiError(ApiError::NotFound(format!("Job {job} not found"))))?; + if model.agent_id != Some(agent_id) { + return Err(Error::ApiError(ApiError::NotFound(format!( + "Job {job} not found" + )))); + } + Ok(model) +} + +/// Reject a state-mutating report against an already-terminal job with `409 +/// Conflict`. The agent reads that as "job already closed → I'm free"; teardown +/// already owns any stranded task's fate. Used where any non-terminal state is a +/// valid source (`/complete`, task reports). +pub fn reject_if_terminal(job: &SuiteAgentJobs::Model) -> Result<()> { + if job.state.is_terminal() { + return Err(Error::ApiError(ApiError::Conflict(format!( + "Job {} is already in terminal state {}", + job.job_id, job.state + )))); + } + Ok(()) +} + +/// Validate that a job is in the exact `expected` source state for an ordered +/// transition (`/start` expects `Provisioning`, `/cleanup` expects `Executing`): +/// - already in `expected` → `Ok`. +/// - terminal → `409 Conflict` (the agent frees itself). +/// - any other non-terminal state → `400` — an out-of-order protocol error, and +/// deliberately *not* the "closed" signal, so the agent does not release +/// itself off a still-live job. +pub fn expect_state(job: &SuiteAgentJobs::Model, expected: SuiteJobState) -> Result<()> { + if job.state == expected { + return Ok(()); + } + if job.state.is_terminal() { + return Err(Error::ApiError(ApiError::Conflict(format!( + "Job {} is already in terminal state {}", + job.job_id, job.state + )))); + } + Err(Error::ApiError(ApiError::InvalidRequest(format!( + "Job {} is in state {}, expected {} for this transition", + job.job_id, job.state, expected + )))) +} + +/// Coordinator-written terminal: mark every in-flight job owned by `agent_id` +/// with `state` (`Lost` on heartbeat timeout, `Killed` on a force shutdown). +/// Returns the number of jobs affected. +pub async fn terminate_agent_jobs( + db: &C, + agent_id: i64, + state: SuiteJobState, + now: TimeDateTimeWithTimeZone, +) -> Result { + let res = SuiteAgentJobs::Entity::update_many() + .col_expr(SuiteAgentJobs::Column::State, Expr::value(state)) + .col_expr(SuiteAgentJobs::Column::UpdatedAt, Expr::value(now)) + .filter(SuiteAgentJobs::Column::AgentId.eq(agent_id)) + .filter(SuiteAgentJobs::Column::State.is_in(IN_FLIGHT)) + .exec(db) + .await?; + Ok(res.rows_affected) +} + +/// Coordinator-written terminal: force-stop every in-flight job of a suite +/// (`Killed`, no cleanup). Returns the agent ids whose jobs were stopped, so the +/// caller can notify them. +pub async fn kill_suite_jobs( + db: &C, + suite_id: i64, + now: TimeDateTimeWithTimeZone, +) -> Result> { + let killed = SuiteAgentJobs::Entity::update_many() + .col_expr( + SuiteAgentJobs::Column::State, + Expr::value(SuiteJobState::Killed), + ) + .col_expr(SuiteAgentJobs::Column::UpdatedAt, Expr::value(now)) + .filter(SuiteAgentJobs::Column::TaskSuiteId.eq(suite_id)) + .filter(SuiteAgentJobs::Column::State.is_in(IN_FLIGHT)) + .exec_with_returning(db) + .await?; + Ok(killed.into_iter().filter_map(|j| j.agent_id).collect()) +} + +/// The agent ids currently running an in-flight job of a suite. +pub async fn agents_running_suite(db: &C, suite_id: i64) -> Result> { + let jobs = SuiteAgentJobs::Entity::find() + .select_only() + .column(SuiteAgentJobs::Column::AgentId) + .filter(SuiteAgentJobs::Column::TaskSuiteId.eq(suite_id)) + .filter(SuiteAgentJobs::Column::State.is_in(IN_FLIGHT)) + .into_tuple::>() + .all(db) + .await?; + Ok(jobs.into_iter().flatten().collect()) +} + +/// Whether the agent owns any job that has not reached a terminal state. +/// +/// Used to tell a live suite assignment from a stale one: `accept` writes the +/// assignment and the job row in the same transaction, so an agent that is +/// genuinely running something always has an in-flight job here. +pub async fn agent_has_in_flight_job(db: &C, agent_id: i64) -> Result { + let count = SuiteAgentJobs::Entity::find() + .filter(SuiteAgentJobs::Column::AgentId.eq(agent_id)) + .filter(SuiteAgentJobs::Column::State.is_in(IN_FLIGHT)) + .count(db) + .await?; + Ok(count > 0) +} + +/// Reclaim an agent's executed-but-uncommitted tasks: every `Running` **and** +/// `Finished` task still owned by `agent_uuid` goes back to `Ready` with its +/// `runner_uuid` cleared, so another agent re-runs it. A `Finished` task counts +/// because its result was never committed. Returns how many were reclaimed. +pub async fn reclaim_agent_tasks( + pool: &InfraPool, + agent_uuid: Uuid, + now: TimeDateTimeWithTimeZone, +) -> Result { + let reclaimed = ActiveTasks::Entity::update_many() + .col_expr(ActiveTasks::Column::State, Expr::value(TaskState::Ready)) + .col_expr(ActiveTasks::Column::RunnerUuid, Expr::value(None::)) + .col_expr(ActiveTasks::Column::UpdatedAt, Expr::value(now)) + .filter(ActiveTasks::Column::RunnerUuid.eq(agent_uuid)) + .filter(ActiveTasks::Column::State.is_in([TaskState::Running, TaskState::Finished])) + .exec_with_returning(&pool.db) + .await?; + if !reclaimed.is_empty() { + tracing::info!( + agent_uuid = %agent_uuid, + reclaimed = reclaimed.len(), + "Reclaimed uncommitted tasks from an agent" + ); + } + Ok(reclaimed.len()) +} diff --git a/netmito/src/service/agent/matching.rs b/netmito/src/service/agent/matching.rs new file mode 100644 index 00000000..b9f0634d --- /dev/null +++ b/netmito/src/service/agent/matching.rs @@ -0,0 +1,221 @@ +//! Which agents may run which suites. +//! +//! The effective set for a suite is +//! +//! ```text +//! (tag-matched ∪ UserIncluded) − UserExcluded restricted to agents the +//! suite's group may write to +//! ``` +//! +//! `task_suite_agent` persists **only** the user's manual overrides +//! (`[[suite-agent-table-manual-only]]`); the tag-matched half is derived, never +//! stored. Dev derives it in a scheduler actor holding an in-memory index. We +//! derive it per query instead, in SQL: Postgres array containment +//! (`agents.tags @> task_suites.tags`) answers "does this agent carry every tag +//! the suite asks for?" directly, and both `tags` columns already have GIN +//! indexes. Same semantics, no second source of truth to keep coherent. +//! +//! Every eligibility check in the agent layer goes through the predicates here, +//! so the rule lives in exactly one place — including for the day the in-memory +//! scheduler does arrive. + +use sea_orm::sea_query::{ + extension::postgres::PgExpr, Alias, Expr, Order, Query, SelectStatement, SimpleExpr, +}; +use sea_orm::{ConnectionTrait, FromQueryResult}; +use uuid::Uuid; + +use crate::entity::role::GroupAgentRole; +use crate::entity::state::{AgentState, TaskSuiteState}; +use crate::entity::task_suite_agent::SuiteAgentOverrideType; +use crate::entity::{ + agents as Agent, group_agent as GroupAgent, task_suite_agent as TaskSuiteAgent, + task_suites as TaskSuites, +}; + +fn suite_col(col: TaskSuites::Column) -> Expr { + Expr::col((TaskSuites::Entity, col)) +} + +fn agent_col(col: Agent::Column) -> Expr { + Expr::col((Agent::Entity, col)) +} + +/// A manual override row exists for this `(suite, agent)` pair with `kind`. +fn has_override(kind: SuiteAgentOverrideType) -> SimpleExpr { + Expr::exists( + Query::select() + .expr(Expr::val(1)) + .from(TaskSuiteAgent::Entity) + .and_where( + Expr::col((TaskSuiteAgent::Entity, TaskSuiteAgent::Column::TaskSuiteId)) + .equals((TaskSuites::Entity, TaskSuites::Column::Id)), + ) + .and_where( + Expr::col((TaskSuiteAgent::Entity, TaskSuiteAgent::Column::AgentId)) + .equals((Agent::Entity, Agent::Column::Id)), + ) + .and_where( + Expr::col((TaskSuiteAgent::Entity, TaskSuiteAgent::Column::OverrideType)).eq(kind), + ) + .to_owned(), + ) +} + +/// The suite's owning group grants this agent at least Write. +fn group_grants_access() -> SimpleExpr { + Expr::exists( + Query::select() + .expr(Expr::val(1)) + .from(GroupAgent::Entity) + .and_where( + Expr::col((GroupAgent::Entity, GroupAgent::Column::GroupId)) + .equals((TaskSuites::Entity, TaskSuites::Column::GroupId)), + ) + .and_where( + Expr::col((GroupAgent::Entity, GroupAgent::Column::AgentId)) + .equals((Agent::Entity, Agent::Column::Id)), + ) + .and_where( + Expr::col((GroupAgent::Entity, GroupAgent::Column::Role)) + .gte(GroupAgentRole::Write), + ) + .to_owned(), + ) +} + +/// The full eligibility predicate for a `(suite, agent)` pair, in terms of the +/// `task_suites` and `agents` columns of the surrounding query. Both tables must +/// be in scope wherever this is used. +pub fn eligibility_predicate() -> SimpleExpr { + let tag_matched = agent_col(Agent::Column::Tags).contains(suite_col(TaskSuites::Column::Tags)); + let included = has_override(SuiteAgentOverrideType::UserIncluded); + let excluded = has_override(SuiteAgentOverrideType::UserExcluded); + + group_grants_access() + .and(tag_matched.or(included)) + .and(excluded.not()) +} + +/// A suite is worth handing to an agent when it can still take work: not +/// cancelled, and with at least one task left to finish. `Complete` suites are +/// excluded — they have nothing to run until a new task reopens them. +pub fn suite_has_work() -> SimpleExpr { + suite_col(TaskSuites::Column::State) + .is_in([TaskSuiteState::Open, TaskSuiteState::Closed]) + .and(suite_col(TaskSuites::Column::IncompleteTasks).gt(0)) +} + +/// Base `SELECT … FROM task_suites, agents WHERE ` that callers extend +/// with their own projection and filters. +fn eligible_pairs() -> SelectStatement { + Query::select() + .from(TaskSuites::Entity) + .from(Agent::Entity) + .and_where(eligibility_predicate()) + .to_owned() +} + +/// Is this agent eligible to run this suite (ignoring whether the suite +/// currently has work)? +pub async fn is_agent_eligible( + db: &C, + suite_id: i64, + agent_id: i64, +) -> crate::error::Result { + #[derive(FromQueryResult)] + struct Exists { + #[allow(dead_code)] + eligible: i32, + } + + let mut stmt = eligible_pairs(); + stmt.expr_as(Expr::val(1), Alias::new("eligible")) + .and_where(suite_col(TaskSuites::Column::Id).eq(suite_id)) + .and_where(agent_col(Agent::Column::Id).eq(agent_id)) + .limit(1); + + let builder = db.get_database_backend(); + Ok(Exists::find_by_statement(builder.build(&stmt)) + .one(db) + .await? + .is_some()) +} + +/// Does this agent have any suite with pending work right now? +pub async fn agent_has_available_suite( + db: &C, + agent_id: i64, +) -> crate::error::Result { + Ok(best_available_suite_id(db, agent_id).await?.is_some()) +} + +/// The suite this agent should pick up next: highest priority first, then +/// oldest. `None` when nothing is available. +pub async fn best_available_suite_id( + db: &C, + agent_id: i64, +) -> crate::error::Result> { + #[derive(FromQueryResult)] + struct SuiteId { + suite_id: i64, + } + + let mut stmt = eligible_pairs(); + stmt.expr_as(suite_col(TaskSuites::Column::Id), Alias::new("suite_id")) + .and_where(agent_col(Agent::Column::Id).eq(agent_id)) + .and_where(suite_has_work()) + .order_by_expr(suite_col(TaskSuites::Column::Priority).into(), Order::Desc) + .order_by_expr(suite_col(TaskSuites::Column::CreatedAt).into(), Order::Asc) + .limit(1); + + let builder = db.get_database_backend(); + Ok(SuiteId::find_by_statement(builder.build(&stmt)) + .one(db) + .await? + .map(|r| r.suite_id)) +} + +#[derive(FromQueryResult)] +struct AgentUuid { + uuid: Uuid, +} + +/// The uuids of every agent eligible for a suite — what a suite's detail view +/// reports and who gets considered when the suite gains work. +pub async fn eligible_agent_uuids( + db: &C, + suite_id: i64, +) -> crate::error::Result> { + let mut stmt = eligible_pairs(); + stmt.expr_as(agent_col(Agent::Column::Uuid), Alias::new("uuid")) + .and_where(suite_col(TaskSuites::Column::Id).eq(suite_id)); + + let builder = db.get_database_backend(); + Ok(AgentUuid::find_by_statement(builder.build(&stmt)) + .all(db) + .await? + .into_iter() + .map(|r| r.uuid) + .collect()) +} + +/// The uuids of eligible agents that are **idle**, i.e. the ones that could +/// start this suite right now. Used to target `SuiteAvailable` pushes. +pub async fn idle_eligible_agent_uuids( + db: &C, + suite_id: i64, +) -> crate::error::Result> { + let mut stmt = eligible_pairs(); + stmt.expr_as(agent_col(Agent::Column::Uuid), Alias::new("uuid")) + .and_where(suite_col(TaskSuites::Column::Id).eq(suite_id)) + .and_where(agent_col(Agent::Column::State).eq(AgentState::Idle)); + + let builder = db.get_database_backend(); + Ok(AgentUuid::find_by_statement(builder.build(&stmt)) + .all(db) + .await? + .into_iter() + .map(|r| r.uuid) + .collect()) +} diff --git a/netmito/src/service/agent/mod.rs b/netmito/src/service/agent/mod.rs new file mode 100644 index 00000000..91b7f4bd --- /dev/null +++ b/netmito/src/service/agent/mod.rs @@ -0,0 +1,990 @@ +//! Agent lifecycle: fleet management for users, and the execution loop for +//! agents. +//! +//! ## The loop +//! +//! ```text +//! register ──▶ heartbeat ──▶ accept suite ──▶ start ──▶ …tasks… ──▶ cleanup ──▶ complete +//! ▲ │ +//! └─────────────────────── Idle ◀───────────────────────────────────┘ +//! ``` +//! +//! `accept` is the only transition that claims anything, and it both chooses +//! the suite and takes it: one transaction picks a candidate, locks it, +//! re-checks eligibility, moves the agent `Idle → Provisioning`, and opens a +//! `suite_agent_jobs` row whose opaque id the agent echoes on every later call. +//! `start`/`cleanup`/`complete` advance that row and the agent state in step. +//! +//! ## No preemption (yet) +//! +//! An agent that has accepted a suite runs it to completion; nothing takes it +//! away for a higher-priority suite. Priority only orders the *choice* made in +//! [`matching::best_available_suite_id`]. The pieces a future preemption needs +//! are already in place and unused: the `PreemptSuite` notification, the agent's +//! handling of it, and the fact that `accept` records the job so the coordinator +//! can address a specific in-flight run. Turning it on means emitting that +//! notification here — no schema or protocol change. + +pub mod heartbeat; +pub mod hook; +pub mod job; +pub mod matching; +pub mod task; + +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, SuiteJobState, TaskSuiteState}, + suite_agent_jobs as SuiteAgentJobs, task_suites as TaskSuites, user_group as UserGroup, + users as User, +}; +use crate::error::{ApiError, AuthError, Error, Result}; +use crate::schema::{ + AcceptSuiteReq, AcceptSuiteResp, AgentHeartbeatReq, AgentHeartbeatResp, AgentInfo, + AgentNotification, AgentShutdownOp, AgentsQueryReq, AgentsQueryResp, CompleteJobReq, + CompleteJobResp, CountQuery, ExecHooks, RegisterAgentReq, RegisterAgentResp, SuiteJobOutcome, + TaskSuiteSpec, WorkerSchedulePlan, +}; +use crate::service::auth::token::generate_worker_token; +use crate::ws::AgentWsRouter; + +use heartbeat::AgentHeartbeatOp; + +// ───────────────────────────────────────────────────────────────────────────── +// Fleet management (user-authed) +// ───────────────────────────────────────────────────────────────────────────── + +/// Register an agent, or re-adopt the one already bound to this machine. +/// +/// `machines.machine_code` and `machines.agent_id` are both unique, so a machine +/// has exactly one agent row for its whole life. Rather than fail a restarting +/// agent on the unique index, registration is an upsert: an existing machine +/// hands back its agent, refreshed with the new tags/labels/groups/metadata and +/// a newly minted token. The caller needs Write or Admin in every group listed, +/// and Admin in `admin_group`. +/// +/// Access is **rewritten**, not merged: the request is the whole truth about +/// which groups may reach the agent, so a re-registration that drops a group +/// drops its access too. +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) = pool + .db + .transaction::<_, (Uuid, i64, bool), Error>(|txn| { + Box::pin(async move { + // Exactly one group holds Admin over the agent, and that role is + // the only way it is ever shut down. Unless the caller names a + // group, it is their personal one + let admin_group_name = match admin_group { + Some(name) => name, + _ => { + 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}" + ))))?; + // The registering user must be the admin of the admin_group + 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)); + } + + // Resolve every requested group up front + // + // The registering user must have at least Write permission in each group. + 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 (agent, reused) = match existing { + Some(machine) => { + let agent = Agent::Entity::find_by_id(machine.agent_id) + .one(txn) + .await? + .ok_or_else(|| { + // The FK guarantees this cannot happen. + Error::Custom(format!( + "Machine {machine_code} points at missing agent {}", + machine.agent_id + )) + })?; + let agent_state = agent.state; + let mut agent: Agent::ActiveModel = agent.into(); + agent.tags = Set(tags); + agent.labels = Set(labels); + // A restarting agent comes back Idle. If the previous + // process died mid-job, that job is still in flight; + // the heartbeat queue will time it out and reclaim its + // tasks. Clearing the assignment here keeps the fresh + // process from being handed a suite it never accepted. + 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, + "Re-adopting the agent already registered for this machine" + ); + + 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) + } + }; + + // Grant group access to the agent. A re-registration rewrites all group access + GroupAgent::Entity::delete_many() + .filter(GroupAgent::Column::AgentId.eq(agent.id)) + .exec(txn) + .await?; + + // Admin first, so an admin group that also appears in `groups` + // keeps Admin instead of being written as Write. + 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)) + }) + }) + .await?; + + // An agent is a long-lived daemon: without an explicit lifetime its token + // never expires. + let token = generate_worker_token(agent_uuid.to_string(), 0, req.lifetime)?; + + // Start (or restart) liveness tracking for this agent. + let _ = pool + .agent_heartbeat_queue_tx + .send(AgentHeartbeatOp::Heartbeat(agent_id)); + + let notification_counter = AgentWsRouter::counter(&pool.ws_router_tx, agent_uuid) + .await + .unwrap_or_default(); + + Ok(RegisterAgentResp { + agent_uuid, + token, + notification_counter, + reused, + }) +} + +/// Query agents visible to the user. Requires at least Read 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. +/// +/// - `Graceful`: ask the agent to stop. An idle agent is parked `Offline` now; a +/// busy one keeps its job and finishes it, going `Offline` when its heartbeat +/// stops. +/// - `Force`: park it `Offline` immediately, kill its in-flight jobs, and +/// reclaim its uncommitted tasks so other agents re-run them. +pub async fn user_shutdown_agent_by_uuid( + user_id: i64, + agent_uuid: Uuid, + op: AgentShutdownOp, + pool: &InfraPool, +) -> Result<()> { + let now = TimeDateTimeWithTimeZone::now_utc(); + + // Fetch and authorize in one lookup + 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 => { + heartbeat::mark_offline(pool, agent.id, now).await?; + let killed = + job::terminate_agent_jobs(&pool.db, agent.id, SuiteJobState::Killed, now).await?; + job::reclaim_agent_tasks(pool, agent.uuid, now).await?; + let _ = pool + .agent_heartbeat_queue_tx + .send(AgentHeartbeatOp::Remove(agent.id)); + tracing::info!(agent_uuid = %agent.uuid, killed_jobs = killed, "Agent force shutdown"); + } + AgentShutdownOp::Graceful => { + if agent.assigned_task_suite_id.is_none() { + heartbeat::mark_offline(pool, agent.id, now).await?; + let _ = pool + .agent_heartbeat_queue_tx + .send(AgentHeartbeatOp::Remove(agent.id)); + } else { + // Busy: let it finish. Its own shutdown handling stops it from taking a + // next suite; liveness tracking stays on so a stall still times out. + tracing::info!( + agent_uuid = %agent.uuid, + "Graceful shutdown requested while busy — the agent will stop after its current job" + ); + } + } + } + + AgentWsRouter::notify( + &pool.ws_router_tx, + agent.uuid, + AgentNotification::Shutdown { + graceful: matches!(op, AgentShutdownOp::Graceful), + }, + ); + + Ok(()) +} + +// ───────────────────────────────────────────────────────────────────────────── +// Execution loop (agent-authed) +// ───────────────────────────────────────────────────────────────────────────── + +/// Record a heartbeat and hand back everything the agent has not seen. +/// +/// +/// Heartbeat is used both to re-notify the agent of pending tasks, and also +/// use for syncing the counters so that if the agent has a counter ahead of +/// the coordinator, the coordinator must have been restarted, and we should +/// notify the agent to resync. +pub async fn agent_heartbeat( + agent_id: i64, + agent_uuid: Uuid, + pool: &InfraPool, + req: AgentHeartbeatReq, +) -> Result { + let now = TimeDateTimeWithTimeZone::now_utc(); + + let _ = pool + .agent_heartbeat_queue_tx + .send(AgentHeartbeatOp::Heartbeat(agent_id)); + + // The agent's own state wins: it is the authority on what it is doing. + let agent = 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)) + .exec_with_returning(&pool.db) + .await? + .into_iter() + .next() + .ok_or(Error::ApiError(ApiError::NotFound("Agent".to_string())))?; + + // An idle agent with nothing assigned and work waiting gets a fresh nudge. + if req.state == AgentState::Idle && agent.assigned_task_suite_id.is_none() { + if let Some(suite_id) = matching::best_available_suite_id(&pool.db, agent_id).await? { + if let Some(suite) = TaskSuites::Entity::find_by_id(suite_id) + .one(&pool.db) + .await? + { + AgentWsRouter::notify( + &pool.ws_router_tx, + agent_uuid, + AgentNotification::SuiteAvailable { + suite_uuid: Some(suite.uuid), + priority: suite.priority, + }, + ); + } + } + } + + // The agent claiming a higher notification id than we ever issued means our + // sequence restarted with the process; hand it our boot id and counter. + 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 desync — sending CounterSync" + ); + AgentWsRouter::notify( + &pool.ws_router_tx, + agent_uuid, + AgentNotification::CounterSync { + counter, + boot_id: pool.boot_uuid, + }, + ); + } + } + + // Any unacked messages get sent in batch in heartbeat response + let notifications = AgentWsRouter::pending_notifications( + &pool.ws_router_tx, + agent_uuid, + req.last_notification_id, + ) + .await; + + Ok(AgentHeartbeatResp { notifications }) +} + +/// Tell every agent that could take this suite that it has work. Only idle +/// agents are targeted +/// +/// TODO: check according to agent scheduling strategy, and notify some agents +/// running low-priority suites to stop the job and switch to this suite +pub async fn notify_suite_available(pool: &InfraPool, suite_id: i64) { + let suite = match TaskSuites::Entity::find_by_id(suite_id).one(&pool.db).await { + Ok(Some(suite)) => suite, + Ok(None) => return, + Err(e) => { + tracing::error!(suite_id, "Failed to load suite for agent notification: {e}"); + return; + } + }; + let agents = match matching::idle_eligible_agent_uuids(&pool.db, suite_id).await { + Ok(agents) => agents, + Err(e) => { + tracing::error!(suite_id, "Failed to resolve eligible agents: {e}"); + return; + } + }; + for agent_uuid in agents { + AgentWsRouter::notify( + &pool.ws_router_tx, + agent_uuid, + AgentNotification::SuiteAvailable { + suite_uuid: Some(suite.uuid), + priority: suite.priority, + }, + ); + } +} + +async fn suite_to_spec( + db: &C, + suite: TaskSuites::Model, +) -> Result { + let group = Group::Entity::find_by_id(suite.group_id) + .one(db) + .await? + .ok_or_else(|| Error::ApiError(ApiError::NotFound("Group of the suite".to_string())))?; + let worker_schedule: WorkerSchedulePlan = serde_json::from_value(suite.worker_schedule)?; + let exec_hooks: Option = suite.exec_hooks.map(serde_json::from_value).transpose()?; + + Ok(TaskSuiteSpec { + uuid: suite.uuid, + name: suite.name, + description: suite.description, + group_name: group.group_name, + tags: suite.tags, + labels: suite.labels, + priority: suite.priority, + worker_schedule, + exec_hooks, + state: suite.state, + total_tasks: suite.total_tasks, + incomplete_tasks: suite.incomplete_tasks, + }) +} + +/// Pick a suite and claim it in one transaction: agent `Idle → Provisioning`, a +/// fresh job row, and the spec to run. Choosing under the suite's row lock is +/// what leaves no window for another agent to take it in between. +/// +/// `req.suite_uuid` is a preference. A hint that no longer points at runnable +/// work — drained, cancelled, no longer this agent's, or simply gone — falls +/// through to the best available suite rather than answering "nothing". +/// +/// Rejections that are the coordinator's decision rather than a client error +/// (nothing available, agent already busy) come back as `accepted: false` with a +/// reason, not an HTTP error. +pub async fn agent_accept_suite( + agent_id: i64, + pool: &InfraPool, + req: AcceptSuiteReq, +) -> Result { + let now = TimeDateTimeWithTimeZone::now_utc(); + + let outcome = pool + .db + .transaction::<_, std::result::Result<(TaskSuiteSpec, i64, i32), String>, Error>(|txn| { + Box::pin(async move { + // One suite at a time, checked before anything is locked. A busy + // agent that asks again is racing its own state, not erroring. + + let agent = Agent::Entity::find_by_id(agent_id) + .one(txn) + .await? + .ok_or(Error::ApiError(ApiError::NotFound("Agent".to_string())))?; + if agent.assigned_task_suite_id.is_some() || agent.state.is_busy() { + return Ok(Err(format!( + "Agent is already {} and cannot accept another suite", + agent.state + ))); + } + + // What the agent asked for, if it is still worth running. + let requested = match req.suite_uuid { + Some(suite_uuid) => { + lock_runnable_suite(txn, agent_id, SuiteTarget::Uuid(suite_uuid)).await? + } + None => None, + }; + // Otherwise — or if the hint went stale — the best on offer. + let suite = match requested { + Some(suite) => suite, + None => { + match matching::best_available_suite_id(txn, agent_id).await? { + Some(suite_id) => { + match lock_runnable_suite(txn, agent_id, SuiteTarget::Id(suite_id)) + .await? + { + Some(suite) => suite, + // Drained between the pick and the lock. + None => { + return Ok(Err( + "No suite is available for this agent".to_string() + )) + } + } + } + None => { + return Ok(Err("No suite is available for this agent".to_string())) + } + } + } + }; + + let suite_id = suite.id; + let spec = suite_to_spec(txn, suite).await?; + + let mut agent: Agent::ActiveModel = agent.into(); + agent.assigned_task_suite_id = Set(Some(suite_id)); + agent.updated_at = Set(now); + agent.update(txn).await?; + + let job = job::create_job(txn, suite_id, agent_id, now).await?; + Ok(Ok((spec, job.id, job.job_id))) + }) + }) + .await?; + + match outcome { + Ok((suite, job, job_id)) => Ok(AcceptSuiteResp { + accepted: true, + suite: Some(suite), + job: Some(job), + job_id: Some(job_id), + reason: None, + }), + Err(reason) => Ok(AcceptSuiteResp { + accepted: false, + suite: None, + job: None, + job_id: None, + reason: Some(reason), + }), + } +} + +/// How a candidate suite was arrived at: named by the agent, or picked for it. +enum SuiteTarget { + Uuid(Uuid), + Id(i64), +} + +/// Lock one candidate suite and answer whether this agent may run it now. +/// +/// `None` covers every way a candidate can fail — gone, not this agent's, +/// terminal, or drained — since the caller falls through to the next candidate +/// for all of them. The checks are re-run under the lock even for a suite +/// `best_available_suite_id` just returned: that query takes no locks, so its +/// answer can go stale before we take one. +async fn lock_runnable_suite( + txn: &C, + agent_id: i64, + target: SuiteTarget, +) -> Result> { + let query = TaskSuites::Entity::find(); + let suite = match target { + SuiteTarget::Uuid(uuid) => query.filter(TaskSuites::Column::Uuid.eq(uuid)), + SuiteTarget::Id(id) => query.filter(TaskSuites::Column::Id.eq(id)), + } + .lock_exclusive() + .one(txn) + .await?; + + let Some(suite) = suite else { + return Ok(None); + }; + // Same predicate as `matching::suite_has_work`, which is what the picker + // filters on. + if !matches!(suite.state, TaskSuiteState::Open | TaskSuiteState::Closed) + || suite.incomplete_tasks == 0 + { + return Ok(None); + } + if !matching::is_agent_eligible(txn, suite.id, agent_id).await? { + return Ok(None); + } + Ok(Some(suite)) +} + +/// Provisioning finished: job `Provisioning → Executing`, agent `Executing`. +pub async fn agent_start_job(agent_id: i64, pool: &InfraPool, job_handle: i64) -> Result<()> { + advance_job( + agent_id, + pool, + job_handle, + SuiteJobState::Provisioning, + SuiteJobState::Executing, + AgentState::Executing, + ) + .await +} + +/// Tasks drained: job `Executing → Cleanup`, agent `Cleaning`. +pub async fn agent_enter_cleanup(agent_id: i64, pool: &InfraPool, job_handle: i64) -> Result<()> { + advance_job( + agent_id, + pool, + job_handle, + SuiteJobState::Executing, + SuiteJobState::Cleanup, + AgentState::Cleaning, + ) + .await +} + +/// One ordered job transition plus the agent state that goes with it, in one +/// transaction so the pair cannot be observed (or crash) half-applied. +async fn advance_job( + agent_id: i64, + pool: &InfraPool, + job_handle: i64, + from: SuiteJobState, + to: SuiteJobState, + agent_state: AgentState, +) -> Result<()> { + let now = TimeDateTimeWithTimeZone::now_utc(); + + pool.db + .transaction::<_, (), Error>(|txn| { + Box::pin(async move { + // The lock is what makes `expect_state` binding: a coordinator + // terminal written between the read and the write would + // otherwise drag a dead job's agent into a live-looking state. + let job_row = job::load_validate_job_locked(txn, job_handle, agent_id).await?; + job::expect_state(&job_row, from)?; + + let suite_id = job_row.task_suite_id; + let mut job: SuiteAgentJobs::ActiveModel = job_row.into(); + job.state = Set(to); + job.updated_at = Set(now); + job.update(txn).await?; + + // Same guard as `agent_complete_job`: never write over an agent + // that teardown has already unassigned. + Agent::Entity::update_many() + .col_expr(Agent::Column::State, Expr::value(agent_state)) + .col_expr(Agent::Column::UpdatedAt, Expr::value(now)) + .filter(Agent::Column::Id.eq(agent_id)) + .filter(Agent::Column::AssignedTaskSuiteId.eq(suite_id)) + .exec(txn) + .await?; + + Ok(()) + }) + }) + .await?; + + Ok(()) +} + +/// Finish the job and release the agent back to `Idle`. +/// +/// The agent reports only what it did — `Completed` or `Failed` with a reason. +/// `Lost` and `Killed` are the coordinator's to write. A job that is already +/// terminal answers 409 so the agent knows it was torn down under it. +/// +/// Both writes — the job's terminal and the agent's release — land in one +/// transaction: a half-applied pair would leave the agent pointing at a finished +/// suite, which `agent_accept_suite` reads as "still busy" and which nothing +/// short of a restart would clear. +pub async fn agent_complete_job( + agent_id: i64, + pool: &InfraPool, + req: CompleteJobReq, +) -> Result { + let now = TimeDateTimeWithTimeZone::now_utc(); + let job_handle = req.job; + + let (job_id, terminal, released) = pool + .db + .transaction::<_, (i32, SuiteJobState, bool), Error>(|txn| { + Box::pin(async move { + // Locked, so the terminal check and the write below cannot + // straddle a `Killed`/`Lost` written by teardown. + let job_row = job::load_validate_job_locked(txn, job_handle, agent_id).await?; + job::reject_if_terminal(&job_row)?; + + let terminal = match &req.outcome { + SuiteJobOutcome::Completed => SuiteJobState::Completed, + SuiteJobOutcome::Failed { reason } => { + tracing::warn!( + job = job_handle, + job_id = job_row.job_id, + kind = ?reason.kind, + "Agent reported a failed job: {}", + reason.message + ); + SuiteJobState::Failed + } + }; + + let suite_id = job_row.task_suite_id; + let job_id = job_row.job_id; + let mut job_row: SuiteAgentJobs::ActiveModel = job_row.into(); + job_row.state = Set(terminal); + job_row.updated_at = Set(now); + job_row.update(txn).await?; + + // Release the agent to Idle only while it is still bound to *this* + // job's suite. A force shutdown or a heartbeat timeout parks it + // `Offline` and clears the assignment; without this filter a + // completion racing that teardown would resurrect the agent as + // `Idle` with nothing tracking its liveness. + let released = Agent::Entity::update_many() + .col_expr(Agent::Column::State, Expr::value(AgentState::Idle)) + .col_expr(Agent::Column::AssignedTaskSuiteId, Expr::value(None::)) + .col_expr(Agent::Column::UpdatedAt, Expr::value(now)) + .filter(Agent::Column::Id.eq(agent_id)) + .filter(Agent::Column::AssignedTaskSuiteId.eq(suite_id)) + .exec(txn) + .await? + .rows_affected + > 0; + + Ok((job_id, terminal, released)) + }) + }) + .await?; + + tracing::info!( + "Agent finished a suite job: job_id {job_id} (handle {job_handle}), \ + terminal state: {terminal}" + ); + if !released { + tracing::info!( + agent_id, + job_id, + "Agent was no longer assigned to the completed job's suite; left its state alone" + ); + } + + let next_suite_available = matching::agent_has_available_suite(&pool.db, agent_id).await?; + Ok(CompleteJobResp { + next_suite_available, + }) +} + +/// Push a notification to every agent id in `agent_ids` (resolving uuids first). +pub(crate) async fn notify_agents_by_id( + pool: &InfraPool, + agent_ids: &[i64], + event: AgentNotification, +) { + if agent_ids.is_empty() { + return; + } + let uuids = Agent::Entity::find() + .select_only() + .column(Agent::Column::Uuid) + .filter(Agent::Column::Id.is_in(agent_ids.to_vec())) + .into_tuple::() + .all(&pool.db) + .await; + match uuids { + Ok(uuids) => { + for uuid in uuids { + AgentWsRouter::notify(&pool.ws_router_tx, uuid, event.clone()); + } + } + Err(e) => tracing::error!("Failed to resolve agent uuids for notification: {e}"), + } +} + +/// On coordinator start, tell the router what boot it is on. Agents learn the +/// new `boot_id` from the `CounterSync` and reset their own sequence, so a +/// restart does not leave them ignoring notifications whose ids look stale. +pub async fn notify_agents_of_restart(pool: &InfraPool) -> Result<()> { + let agents = Agent::Entity::find() + .select_only() + .column(Agent::Column::Uuid) + .into_tuple::() + .all(&pool.db) + .await?; + tracing::info!( + agents = agents.len(), + boot_id = %pool.boot_uuid, + "Announcing coordinator restart to known agents" + ); + for uuid in agents { + AgentWsRouter::notify( + &pool.ws_router_tx, + uuid, + AgentNotification::CounterSync { + counter: 0, + boot_id: pool.boot_uuid, + }, + ); + } + Ok(()) +} diff --git a/netmito/src/service/agent/task.rs b/netmito/src/service/agent/task.rs new file mode 100644 index 00000000..ebe1994c --- /dev/null +++ b/netmito/src/service/agent/task.rs @@ -0,0 +1,349 @@ +//! Task claim and result reporting for agents. +//! +//! ## Claiming +//! +//! Agents claim straight from the database: one transaction selects the +//! highest-priority `Ready` tasks of the suite `FOR UPDATE SKIP LOCKED` and +//! flips them to `Running`. `SKIP LOCKED` is what makes it safe — concurrent +//! agents step over each other's locked rows instead of contending, so no two +//! agents can claim the same task and none of them block. +//! +//! ## Reporting +//! +//! Agents use the same [`ReportTaskOp`] variants as workers, with the suite job +//! handle attached. Every report is validated against the job first: it must +//! exist, belong to this agent, and be non-terminal — a terminal job answers +//! `409 Conflict`, which the agent reads as "the job is closed, stop". +//! +//! | Op | Behaviour | +//! |---|---| +//! | `Finish` | task → `Finished` (not yet archived) | +//! | `Cancel` | task → `Cancelled`; the agent follows up with `Commit` | +//! | `Upload` | request for presigned S3 URL for an artifact (records metadata, charges group quota) | +//! | `Commit` | archives the task with its result, decrements the suite's `incomplete_tasks` (completing the suite if that empties an already-`Closed` one), triggers the downstream task if one was spawned | +//! | `Submit` | spawns a downstream child in the parent's group, in any suite that group owns or none | + +use sea_orm::{prelude::*, QueryOrder, QuerySelect, Set, TransactionTrait}; +use uuid::Uuid; + +use crate::{ + config::InfraPool, + entity::{ + active_tasks as ActiveTasks, archived_tasks as ArchivedTasks, + state::{TaskState, TaskSuiteState}, + task_suites as TaskSuites, StoredTaskModel, + }, + error::{ApiError, Error, Result}, + schema::{ExecSpec, FetchTasksResp, ReportTaskOp, TaskExecOptions, WorkerTaskResp}, + service::{agent::job, agent::matching, s3::group_upload_artifact}, +}; + +/// Upper bound on a single claim batch, so a malformed `max_count` cannot ask +/// the coordinator to lock an unbounded number of rows. +const MAX_FETCH_BATCH: u32 = 256; + +/// Claim up to `max_count` ready tasks of a suite for this agent. +pub async fn agent_fetch_tasks( + agent_id: i64, + agent_uuid: Uuid, + pool: &InfraPool, + suite_uuid: Uuid, + max_count: u32, +) -> Result { + let max_count = max_count.clamp(1, MAX_FETCH_BATCH); + + let suite = TaskSuites::Entity::find() + .filter(TaskSuites::Column::Uuid.eq(suite_uuid)) + .one(&pool.db) + .await? + .ok_or_else(|| Error::ApiError(ApiError::NotFound(format!("Suite {suite_uuid}"))))?; + + if !matching::is_agent_eligible(&pool.db, suite.id, agent_id).await? { + return Err(Error::ApiError(ApiError::NotFound(format!( + "Suite {suite_uuid}" + )))); + } + if !suite.state.allows_task_execution() { + return Err(Error::ApiError(ApiError::Conflict(format!( + "Suite {suite_uuid} is {} and hands out no more tasks", + suite.state + )))); + } + + let now = TimeDateTimeWithTimeZone::now_utc(); + let suite_id = suite.id; + let claimed = pool + .db + .transaction::<_, Vec, Error>(|txn| { + Box::pin(async move { + // TODO: serve claims from a per-suite in-memory priority queue + // refilled from the database in batches, falling back to this + // query + let candidates = ActiveTasks::Entity::find() + .filter(ActiveTasks::Column::TaskSuiteId.eq(suite_id)) + .filter(ActiveTasks::Column::State.eq(TaskState::Ready)) + .order_by_desc(ActiveTasks::Column::Priority) + .order_by_asc(ActiveTasks::Column::Id) + .limit(max_count as u64) + .lock_with_behavior( + sea_orm::sea_query::LockType::Update, + sea_orm::sea_query::LockBehavior::SkipLocked, + ) + .all(txn) + .await?; + if candidates.is_empty() { + return Ok(Vec::new()); + } + + let ids: Vec = candidates.iter().map(|t| t.id).collect(); + ActiveTasks::Entity::update_many() + .col_expr(ActiveTasks::Column::State, Expr::value(TaskState::Running)) + .col_expr(ActiveTasks::Column::RunnerUuid, Expr::value(agent_uuid)) + .col_expr(ActiveTasks::Column::UpdatedAt, Expr::value(now)) + .filter(ActiveTasks::Column::Id.is_in(ids)) + .exec(txn) + .await?; + + Ok(candidates) + }) + }) + .await?; + + let mut tasks = Vec::with_capacity(claimed.len()); + for task in claimed { + let spec: ExecSpec = serde_json::from_value(task.spec).inspect_err(|e| { + tracing::error!(task_uuid = %task.uuid, "Stored task spec is unreadable: {e}"); + })?; + let exec_options: Option = task + .exec_options + .map(serde_json::from_value) + .transpose() + .inspect_err(|e| { + tracing::error!(task_uuid = %task.uuid, "Stored exec_options are unreadable: {e}"); + })?; + tasks.push(WorkerTaskResp { + id: task.id, + uuid: task.uuid, + upstream_task_uuid: task.upstream_task_uuid, + spec, + exec_options, + }); + } + + tracing::debug!( + agent_uuid = %agent_uuid, + suite_uuid = %suite_uuid, + count = tasks.len(), + "Agent claimed tasks from a suite" + ); + // Read from the same snapshot as the eligibility check above, which is at + // most a claim-transaction old. Erring on the side of `true` for a suite the + // sweep has just settled only costs one more poll. + let hold_job_open = matches!(suite.state, TaskSuiteState::Open); + Ok(FetchTasksResp { + tasks, + hold_job_open, + }) +} + +/// Report the result of a task this agent claimed. +pub async fn agent_report_task( + agent_id: i64, + agent_uuid: Uuid, + job_handle: i64, + task_id: i64, + op: ReportTaskOp, + pool: &InfraPool, +) -> Result> { + let now = TimeDateTimeWithTimeZone::now_utc(); + + let job_row = job::load_validate_job(&pool.db, job_handle, agent_id).await?; + job::reject_if_terminal(&job_row)?; + + let task = ActiveTasks::Entity::find_by_id(task_id) + .one(&pool.db) + .await? + .ok_or_else(|| Error::ApiError(ApiError::NotFound(format!("Task {task_id}"))))?; + + // Ownership: only the agent the task was handed to may report on it. A task + // reclaimed after a heartbeat lapse has had its runner cleared, so a late + // report from the old owner lands here as "not found" rather than + // overwriting the re-run. + if task.runner_uuid != Some(agent_uuid) { + return Err(Error::ApiError(ApiError::NotFound(format!( + "Task {task_id}" + )))); + } + + match op { + ReportTaskOp::Finish => { + tracing::debug!(agent_uuid = %agent_uuid, task_uuid = %task.uuid, "Agent finished a task"); + ActiveTasks::Entity::update_many() + .col_expr(ActiveTasks::Column::State, Expr::value(TaskState::Finished)) + .col_expr(ActiveTasks::Column::UpdatedAt, Expr::value(now)) + .filter(ActiveTasks::Column::Id.eq(task.id)) + .exec(&pool.db) + .await?; + } + + ReportTaskOp::Cancel => { + tracing::debug!(agent_uuid = %agent_uuid, task_uuid = %task.uuid, "Agent cancelled a task"); + // Already cancelled (the user got there first) is an acknowledgement, + // not an error. + if task.state != TaskState::Cancelled { + ActiveTasks::Entity::update_many() + .col_expr( + ActiveTasks::Column::State, + Expr::value(TaskState::Cancelled), + ) + .col_expr(ActiveTasks::Column::UpdatedAt, Expr::value(now)) + .filter(ActiveTasks::Column::Id.eq(task.id)) + .exec(&pool.db) + .await?; + } + } + + ReportTaskOp::Commit(res) => { + tracing::debug!(agent_uuid = %agent_uuid, task_uuid = %task.uuid, "Agent committed a task"); + if task.state != TaskState::Finished && task.state != TaskState::Cancelled { + return Err(Error::ApiError(ApiError::InvalidRequest( + "Task must be Finished or Cancelled before Commit".to_string(), + ))); + } + let result = serde_json::to_value(res)?; + let archived = ArchivedTasks::ActiveModel { + id: Set(task.id), + creator_id: Set(task.creator_id), + group_id: Set(task.group_id), + task_id: Set(task.task_id), + uuid: Set(task.uuid), + tags: Set(task.tags.clone()), + labels: Set(task.labels.clone()), + created_at: Set(task.created_at), + updated_at: Set(now), + state: Set(task.state), + runner_uuid: Set(task.runner_uuid), + priority: Set(task.priority), + spec: Set(task.spec.clone()), + exec_options: Set(task.exec_options.clone()), + result: Set(Some(result)), + upstream_task_uuid: Set(task.upstream_task_uuid), + downstream_task_uuid: Set(task.downstream_task_uuid), + task_suite_id: Set(task.task_suite_id), + }; + + let inner_task_id = task.id; + pool.db + .transaction::<_, (), Error>(|txn| { + Box::pin(async move { + archived.insert(txn).await?; + ActiveTasks::Entity::delete_by_id(inner_task_id) + .exec(txn) + .await?; + Ok(()) + }) + }) + .await?; + + if let Some(suite_id) = task.task_suite_id { + commit_suite_task(pool, suite_id, now).await?; + } + + // Task chaining: a child registered by an earlier Submit goes + // Pending → Ready now that its parent has committed. + if let Some(downstream_uuid) = task.downstream_task_uuid { + crate::service::task::worker_trigger_pending_task(pool, downstream_uuid).await?; + } + } + + ReportTaskOp::Upload { + content_type, + content_length, + } => { + let (_, url) = group_upload_artifact( + pool, + StoredTaskModel::Active(task), + content_type, + content_length, + ) + .await?; + return Ok(Some(url)); + } + + ReportTaskOp::Submit(req) => { + if task.state != TaskState::Finished && task.state != TaskState::Cancelled { + return Err(Error::ApiError(ApiError::InvalidRequest( + "Task must be Finished or Cancelled before spawning a new task".to_string(), + ))); + } + // The child stays in the parent's group, but is free to name any + // suite that group owns — the parent's own suite, a sibling one, or + // none at all for a task the workers pick up. The submit path + // enforces the group rule. + let resp = crate::service::task::worker_submit_pending_task( + pool, + task.creator_id, + task.uuid, + task.group_id, + *req, + ) + .await?; + + ActiveTasks::Entity::update_many() + .col_expr( + ActiveTasks::Column::DownstreamTaskUuid, + Expr::value(Some(resp.uuid)), + ) + .col_expr(ActiveTasks::Column::UpdatedAt, Expr::value(now)) + .filter(ActiveTasks::Column::Id.eq(task.id)) + .exec(&pool.db) + .await?; + } + } + + Ok(None) +} + +/// Tick a suite's `incomplete_tasks` down after a commit, and complete it if +/// that was the last task of an already-`Closed` suite. +/// +/// Draining an `Open` suite does not complete it: it stays `Open` until it has +/// also been quiet for `suite_auto_close_timeout`, which is what lets the agent +/// hold its job — and its provisioned environment — open for a late arrival. +/// `service::suite::sweep_inactive_suites` settles it afterwards. A `Closed` +/// suite has no such window to wait out, so it is completed here. +async fn commit_suite_task( + pool: &InfraPool, + suite_id: i64, + now: TimeDateTimeWithTimeZone, +) -> Result<()> { + TaskSuites::Entity::update_many() + .col_expr( + TaskSuites::Column::IncompleteTasks, + Expr::col(TaskSuites::Column::IncompleteTasks).sub(1), + ) + .col_expr(TaskSuites::Column::UpdatedAt, Expr::value(now)) + .filter(TaskSuites::Column::Id.eq(suite_id)) + .filter(TaskSuites::Column::IncompleteTasks.gt(0)) + .exec(&pool.db) + .await?; + + // Conditional on the state and the count, which is what makes a second + // statement safe without a transaction around the pair: a submission racing + // it either lands first and leaves a non-zero count, or lands after and sets + // `Open` itself. + TaskSuites::Entity::update_many() + .col_expr( + TaskSuites::Column::State, + Expr::value(TaskSuiteState::Complete), + ) + .col_expr(TaskSuites::Column::CompletedAt, Expr::value(Some(now))) + .col_expr(TaskSuites::Column::UpdatedAt, Expr::value(now)) + .filter(TaskSuites::Column::Id.eq(suite_id)) + .filter(TaskSuites::Column::State.eq(TaskSuiteState::Closed)) + .filter(TaskSuites::Column::IncompleteTasks.eq(0)) + .exec(&pool.db) + .await?; + + Ok(()) +} diff --git a/netmito/src/service/auth/credential_guard.rs b/netmito/src/service/auth/credential_guard.rs index a6eaffad..22c222f3 100644 --- a/netmito/src/service/auth/credential_guard.rs +++ b/netmito/src/service/auth/credential_guard.rs @@ -44,9 +44,9 @@ impl CredentialGuard { /// If the file isn't fully read/write-able, we warn it, but we will still record user's token /// in-memory in self.active_credential. /// - /// Upon new(), we will immediately try to parse the file and try to load the credential into - /// self.active_credential. If the credential is not found, the active_credential just stays - /// `None`. + /// Nothing is read from the file yet: a credential is stored per user, and the user is not + /// known until [`load_credential`](Self::load_credential) names one. That call is what fills + /// `self.credential`, and it stays `None` until then. pub(crate) async fn new(credential_path: Option, coordinator_url: &Url) -> Self { let credential_path = credential_path.or_else(|| { dirs::config_dir().map(|mut path| { @@ -56,16 +56,12 @@ impl CredentialGuard { }) }); - let mut credential_guard = Self { + Self { credential_path, origin: normalize_origin(coordinator_url), username: None, credential: None, - }; - - credential_guard.credential = credential_guard.load_credential_file().await; - - credential_guard + } } /// The caller wants to access the credential for current username diff --git a/netmito/src/service/auth/mod.rs b/netmito/src/service/auth/mod.rs index 20985ad9..69802f1e 100644 --- a/netmito/src/service/auth/mod.rs +++ b/netmito/src/service/auth/mod.rs @@ -18,7 +18,7 @@ use sea_orm::{entity::prelude::*, Set}; use crate::{ config::InfraPool, - entity::{state::UserState, users as User, workers as Worker}, + entity::{agents as Agent, state::UserState, users as User, workers as Worker}, error::{ApiError, AuthError}, schema::{UserChangePasswordReq, UserLoginReq}, }; @@ -26,7 +26,6 @@ use token::{generate_token, verify_token}; // TODO: should check if the structure and logic is high cohesive and low coupling enough, and // refactor if necessary. -// TODO: check if we allow non-expiry lifetime #[derive(Debug, Clone)] pub struct AuthUser { @@ -50,6 +49,12 @@ pub struct AuthWorker { pub uuid: Uuid, } +#[derive(Debug, Clone)] +pub struct AuthAgent { + pub id: i64, + pub uuid: Uuid, +} + fn generate_auth_signature() -> i64 { StdRng::from_os_rng().next_u32() as i64 + 1 } @@ -413,3 +418,29 @@ async fn worker_auth(db: &DatabaseConnection, bearer: &Bearer) -> Result, + TypedHeader(Authorization(bearer)): TypedHeader>, + mut req: Request, + next: Next, +) -> Result { + let auth_agent = agent_auth(&pool.db, &bearer).await?; + req.extensions_mut().insert(auth_agent); + Ok(next.run(req).await) +} + +async fn agent_auth(db: &DatabaseConnection, bearer: &Bearer) -> Result { + let token = bearer.token(); + let claims = verify_token(token).map_err(|_| AuthError::InvalidToken)?; + 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)?; + + Ok(AuthAgent { id: agent.id, uuid }) +} diff --git a/netmito/src/service/mod.rs b/netmito/src/service/mod.rs index c7b4edda..607d011b 100644 --- a/netmito/src/service/mod.rs +++ b/netmito/src/service/mod.rs @@ -1,3 +1,4 @@ +pub mod agent; pub mod auth; pub mod group; pub mod s3; diff --git a/netmito/src/service/s3.rs b/netmito/src/service/s3.rs index 4c01ce79..f990c3e4 100644 --- a/netmito/src/service/s3.rs +++ b/netmito/src/service/s3.rs @@ -8,7 +8,7 @@ use reqwest::{header::CONTENT_LENGTH, Response}; use sea_orm::{ prelude::*, sea_query::{Expr, Query}, - FromQueryResult, Set, TransactionTrait, + FromQueryResult, QuerySelect, Set, TransactionTrait, }; use tokio::io::AsyncWriteExt; use tokio_util::io::ReaderStream; @@ -23,7 +23,7 @@ use crate::{ entity::{ active_tasks as ActiveTask, archived_tasks as ArchivedTask, artifacts as Artifact, attachments as Attachment, content::ArtifactContentType, groups as Group, - role::UserGroupRole, user_group as UserGroup, + role::UserGroupRole, task_suites as TaskSuite, user_group as UserGroup, }, error::AuthError, }; @@ -413,18 +413,9 @@ pub(crate) async fn group_upload_artifact( StoredTaskModel::Active(ref task) => (task.uuid, task.group_id), StoredTaskModel::Archived(ref task) => (task.uuid, task.group_id), }; - // Check if group is active and has enough storage quota - let group = Group::Entity::find_by_id(group_id) - .one(&pool.db) - .await? - .ok_or(ApiError::InvalidRequest( - "Group for the task not found".to_string(), - ))?; - if group.state != GroupState::Active { - return Err(ApiError::InvalidRequest("Group is not active".to_string()).into()); - } - let s3_client = pool.s3.clone(); - let artifacts_bucket = pool.artifacts_bucket.clone(); + // The group is loaded, checked and quota-allocated inside the transaction + // by `reserve_artifact_upload`. + let pool_cloned = pool.clone(); // This returns (exist, url) let resp = pool .db @@ -449,85 +440,123 @@ pub(crate) async fn group_upload_artifact( updated_task.update(txn).await?; } } - let artifact = Artifact::Entity::find() - .filter(Artifact::Column::TaskId.eq(uuid)) - .filter(Artifact::Column::ContentType.eq(content_type)) - .one(txn) - .await?; - let s3_object_key = format!("{uuid}/{content_type}"); - let url: String; - // Check group storage quota and allocate storage for the artifact - let exist = match artifact { - Some(artifact) => { - let recorded_content_length = content_length.max(artifact.size); - let new_storage_used = - group.storage_used + (recorded_content_length - artifact.size); - if new_storage_used > group.storage_quota { - return Err(ApiError::QuotaExceeded.into()); - } - url = get_presigned_upload_link( - &s3_client, - &artifacts_bucket, - s3_object_key, - content_length, - ) - .await - .map_err(ApiError::from)?; - let artifact = Artifact::ActiveModel { - id: Set(artifact.id), - size: Set(recorded_content_length), - updated_at: Set(now), - ..Default::default() - }; - artifact.update(txn).await?; - let group = Group::ActiveModel { - id: Set(group_id), - storage_used: Set(new_storage_used), - updated_at: Set(now), - ..Default::default() - }; - group.update(txn).await?; - true - } - None => { - let new_storage_used = group.storage_used + content_length; - if new_storage_used > group.storage_quota { - return Err(ApiError::QuotaExceeded.into()); - } - url = get_presigned_upload_link( - &s3_client, - &artifacts_bucket, - s3_object_key, - content_length, - ) - .await - .map_err(ApiError::from)?; - let artifact = Artifact::ActiveModel { - task_id: Set(uuid), - content_type: Set(content_type), - size: Set(content_length), - created_at: Set(now), - updated_at: Set(now), - ..Default::default() - }; - artifact.insert(txn).await?; - let group = Group::ActiveModel { - id: Set(group_id), - storage_used: Set(new_storage_used), - updated_at: Set(now), - ..Default::default() - }; - group.update(txn).await?; - false - } - }; - Ok((exist, url)) + reserve_artifact_upload( + txn, + &pool_cloned, + group_id, + uuid, + content_type, + content_length, + now, + ) + .await }) }) .await?; Ok(resp) } +/// Account for one artifact upload and presign its PUT, inside the caller's +/// transaction. Returns `(already_existed, url)`. +/// +/// `owner_uuid` is what the artifact is filed under in `artifacts.task_id`. +/// +/// Check for group active, check for group quota and signs a s3 upload link +pub(crate) async fn reserve_artifact_upload( + txn: &C, + pool: &InfraPool, + group_id: i64, + owner_uuid: Uuid, + content_type: ArtifactContentType, + content_length: i64, + now: TimeDateTimeWithTimeZone, +) -> Result<(bool, String), Error> { + let s3_client = &pool.s3; + let artifacts_bucket = &pool.artifacts_bucket; + let group = Group::Entity::find_by_id(group_id) + .lock_exclusive() + .one(txn) + .await? + .ok_or(ApiError::InvalidRequest( + "Group for the artifact not found".to_string(), + ))?; + if group.state != GroupState::Active { + return Err(ApiError::InvalidRequest("Group is not active".to_string()).into()); + } + let artifact = Artifact::Entity::find() + .filter(Artifact::Column::TaskId.eq(owner_uuid)) + .filter(Artifact::Column::ContentType.eq(content_type)) + .one(txn) + .await?; + let s3_object_key = format!("{owner_uuid}/{content_type}"); + let url: String; + // Check group storage quota and allocate storage for the artifact + let exist = match artifact { + Some(artifact) => { + let recorded_content_length = content_length.max(artifact.size); + let new_storage_used = group.storage_used + (recorded_content_length - artifact.size); + if new_storage_used > group.storage_quota { + return Err(ApiError::QuotaExceeded.into()); + } + url = get_presigned_upload_link( + s3_client, + artifacts_bucket, + s3_object_key, + content_length, + ) + .await + .map_err(ApiError::from)?; + let artifact = Artifact::ActiveModel { + id: Set(artifact.id), + size: Set(recorded_content_length), + updated_at: Set(now), + ..Default::default() + }; + artifact.update(txn).await?; + let group = Group::ActiveModel { + id: Set(group.id), + storage_used: Set(new_storage_used), + updated_at: Set(now), + ..Default::default() + }; + group.update(txn).await?; + true + } + None => { + let new_storage_used = group.storage_used + content_length; + if new_storage_used > group.storage_quota { + return Err(ApiError::QuotaExceeded.into()); + } + url = get_presigned_upload_link( + s3_client, + artifacts_bucket, + s3_object_key, + content_length, + ) + .await + .map_err(ApiError::from)?; + let artifact = Artifact::ActiveModel { + task_id: Set(owner_uuid), + content_type: Set(content_type), + size: Set(content_length), + created_at: Set(now), + updated_at: Set(now), + ..Default::default() + }; + artifact.insert(txn).await?; + let group = Group::ActiveModel { + id: Set(group.id), + storage_used: Set(new_storage_used), + updated_at: Set(now), + ..Default::default() + }; + group.update(txn).await?; + false + } + }; + Ok((exist, url)) +} + pub async fn user_upload_artifact( pool: &InfraPool, user_id: i64, @@ -609,6 +638,52 @@ pub async fn worker_download_attachment( } } }; + presign_group_attachment(pool, group_id, &group_name, key).await +} + +/// Presign an attachment of the group that owns `suite_uuid`. +/// +/// The hook half of [`worker_download_attachment`]: a hook's inputs are named +/// exactly like a task's, but a hook has no task uuid to resolve the group +/// through — it belongs to a suite. Identity-free like its sibling; the route is +/// what gates access. +pub async fn suite_download_attachment( + pool: &InfraPool, + suite_uuid: Uuid, + key: String, +) -> Result { + let builder = pool.db.get_database_backend(); + let stmt = Query::select() + .column((Group::Entity, Group::Column::Id)) + .column((Group::Entity, Group::Column::GroupName)) + .from(Group::Entity) + .join( + sea_orm::JoinType::Join, + TaskSuite::Entity, + Expr::col((TaskSuite::Entity, TaskSuite::Column::GroupId)) + .eq(Expr::col((Group::Entity, Group::Column::Id))), + ) + .and_where(Expr::col((TaskSuite::Entity, TaskSuite::Column::Uuid)).eq(suite_uuid)) + .limit(1) + .to_owned(); + let GroupInfo { + id: group_id, + group_name, + } = GroupInfo::find_by_statement(builder.build(&stmt)) + .one(&pool.db) + .await? + .ok_or(crate::error::ApiError::NotFound(format!( + "Task suite with uuid {suite_uuid}" + )))?; + presign_group_attachment(pool, group_id, &group_name, key).await +} + +async fn presign_group_attachment( + pool: &InfraPool, + group_id: i64, + group_name: &str, + key: String, +) -> Result { let attachment = Attachment::Entity::find() .filter(Attachment::Column::GroupId.eq(group_id)) .filter(Attachment::Column::Key.eq(key.clone())) diff --git a/netmito/src/service/suite.rs b/netmito/src/service/suite.rs index d68ae442..39144b7a 100644 --- a/netmito/src/service/suite.rs +++ b/netmito/src/service/suite.rs @@ -1,24 +1,24 @@ use std::collections::HashMap; use sea_orm::sea_query::extension::postgres::PgExpr; -use sea_orm::sea_query::{Alias, PgFunc, Query}; +use sea_orm::sea_query::{Alias, Func, PgFunc, Query}; use sea_orm::{prelude::*, ConnectionTrait, FromQueryResult, Set, TransactionTrait}; use uuid::Uuid; use super::suite_agent::{apply_suite_agent_override, authorize_suite, resolve_agent}; use crate::config::InfraPool; use crate::entity::role::UserGroupRole; -use crate::entity::state::{TaskState, TaskSuiteState}; -use crate::entity::task_suite_agent::SuiteAgentOverrideType; +use crate::entity::state::{SuiteJobState, TaskState, TaskSuiteState}; use crate::entity::{ active_tasks as ActiveTasks, agents as Agent, archived_tasks as ArchivedTasks, groups as Group, - task_suite_agent as TaskSuiteAgent, task_suites as TaskSuites, user_group as UserGroup, - users as User, + hook_tasks as HookTasks, suite_agent_jobs as SuiteAgentJobs, task_suites as TaskSuites, + user_group as UserGroup, users as User, }; use crate::error::{ApiError, Error, ResolveError, Result}; use crate::schema::{ - CancelTaskSuiteOp, CountQuery, CreateTaskSuiteReq, CreateTaskSuiteResp, ExecHooks, - ParsedTaskSuiteInfo, SuiteAgentOverrideReq, SuiteAgentOverrideResp, TaskResultMessage, + AgentNotification, CancelTaskSuiteOp, CountQuery, CreateTaskSuiteReq, CreateTaskSuiteResp, + ExecHooks, HookTaskInfo, ParsedTaskSuiteInfo, SuiteAgentOverrideReq, SuiteAgentOverrideResp, + SuiteJobInfo, SuiteJobQueryResp, SuiteJobsQueryReq, SuiteJobsQueryResp, TaskResultMessage, TaskResultSpec, TaskSuiteInfo, TaskSuiteQueryResp, TaskSuitesQueryReq, TaskSuitesQueryResp, WorkerSchedulePlan, }; @@ -29,6 +29,81 @@ struct GroupIdResult { id: i64, } +/// Sweep suites that have gone quiet out of `Open`. +/// +/// A suite stays `Open` from its last submission until `timeout` has passed, +/// even after its tasks have all drained, which is what lets the agent running +/// it hold its job open and pick up a late arrival without provisioning again. +/// The states this settles it into: +/// +/// | from | condition | to | +/// |---|---|---| +/// | `Open` | idle, work remains | `Closed` | +/// | `Open` | idle, nothing left | `Complete` | +/// | `Closed` | nothing left | `Complete` | +/// +/// The last rule is a backstop: `agent::task::commit_suite_task` completes a +/// `Closed` suite as its last task commits, leaving this to catch the ones +/// closed after their work had already drained. +/// +/// Idleness is measured from the last submission, falling back to creation, so a +/// suite that never received a task closes too rather than sitting `Open` +/// forever. +pub async fn sweep_inactive_suites( + db: &DatabaseConnection, + timeout: std::time::Duration, +) -> Result { + let now = TimeDateTimeWithTimeZone::now_utc(); + let threshold = now - time::Duration::seconds(timeout.as_secs() as i64); + let idle_since = Expr::expr(Func::coalesce([ + Expr::col(TaskSuites::Column::LastTaskSubmittedAt).into(), + Expr::col(TaskSuites::Column::CreatedAt).into(), + ])); + + let completed = TaskSuites::Entity::update_many() + .col_expr( + TaskSuites::Column::State, + Expr::value(TaskSuiteState::Complete), + ) + .col_expr(TaskSuites::Column::CompletedAt, Expr::value(Some(now))) + .col_expr(TaskSuites::Column::UpdatedAt, Expr::value(now)) + .filter(TaskSuites::Column::IncompleteTasks.eq(0)) + .filter( + // A `Closed` suite is already known to be idle; an `Open` one has to + // have gone quiet for the whole window first. + TaskSuites::Column::State + .eq(TaskSuiteState::Closed) + .or(TaskSuites::Column::State + .eq(TaskSuiteState::Open) + .and(idle_since.clone().lt(threshold))), + ) + .exec(db) + .await? + .rows_affected; + + let closed = TaskSuites::Entity::update_many() + .col_expr( + TaskSuites::Column::State, + Expr::value(TaskSuiteState::Closed), + ) + .col_expr(TaskSuites::Column::UpdatedAt, Expr::value(now)) + .filter(TaskSuites::Column::State.eq(TaskSuiteState::Open)) + .filter(TaskSuites::Column::IncompleteTasks.gt(0)) + .filter(idle_since.lt(threshold)) + .exec(db) + .await? + .rows_affected; + + if completed + closed > 0 { + tracing::debug!( + completed, + closed, + "Swept suites idle for more than {timeout:?}" + ); + } + Ok(completed + closed) +} + pub async fn user_create_task_suite( user_id: i64, pool: &InfraPool, @@ -383,11 +458,6 @@ struct SuiteDetailResult { completed_at: Option, } -#[derive(FromQueryResult)] -struct AgentUuidResult { - uuid: Uuid, -} - /// Get a single suite's details, including the uuids of agents allowed to execute this suite. /// The caller must have at least `Read` in the suite's group. /// @@ -460,37 +530,10 @@ pub async fn user_get_task_suite_by_uuid( "User doesn't have permission or suite with uuid {suite_uuid}" ))))?; - // Only the manually-included agents are persisted; excluded rows must not be - // reported as assigned. - // - // TODO: this is only the manual half. The *effective* assigned set is - // `(in-memory tag-matched ∪ UserIncluded) − UserExcluded`. Once the agent scheduler - // exists, merge its in-memory tag-matched set here (and subtract UserExcluded) to - // produce the real assigned-agent list. - let agent_stmt = Query::select() - .column((Agent::Entity, Agent::Column::Uuid)) - .from(Agent::Entity) - .join( - sea_orm::JoinType::Join, - TaskSuiteAgent::Entity, - Expr::col((TaskSuiteAgent::Entity, TaskSuiteAgent::Column::AgentId)) - .eq(Expr::col((Agent::Entity, Agent::Column::Id))), - ) - .and_where( - Expr::col((TaskSuiteAgent::Entity, TaskSuiteAgent::Column::TaskSuiteId)).eq(suite.id), - ) - .and_where( - Expr::col((TaskSuiteAgent::Entity, TaskSuiteAgent::Column::OverrideType)) - .eq(SuiteAgentOverrideType::UserIncluded), - ) - .to_owned(); - - let eligible_agents = AgentUuidResult::find_by_statement(builder.build(&agent_stmt)) - .all(&pool.db) - .await? - .into_iter() - .map(|m| m.uuid) - .collect(); + // The effective set: `(tag-matched ∪ UserIncluded) − UserExcluded`, gated by + // the suite group's access to each agent. See `service::agent::matching`. + let eligible_agents = + crate::service::agent::matching::eligible_agent_uuids(&pool.db, suite.id).await?; let worker_schedule: WorkerSchedulePlan = serde_json::from_value(suite.worker_schedule)?; let exec_hooks: Option = suite.exec_hooks.map(serde_json::from_value).transpose()?; @@ -558,20 +601,24 @@ pub async fn user_close_task_suite(user_id: i64, pool: &InfraPool, suite_uuid: U /// Cancel a suite (`* → Cancelled`, terminal). Archives every non-terminal task /// of the suite as `Cancelled`. Requires Write/Admin in the suite's group. /// -/// `op` (`Graceful`/`Force`) currently has no differentiated effect: its only -/// distinction is how running agents/jobs tear down and get notified, which lives -/// in the agent layer that is not ported yet. The parameter is accepted now so the -/// API stays stable; wire its behavior in with the agent work (see the seams below). +/// `op` decides what happens to agents mid-run: +/// - `Graceful` — in-flight jobs keep their state and run their cleanup hook; +/// each agent is told the suite was cancelled and drives its own job to +/// `Completed`. No task the agent is holding is archived out from under it +/// without notice: it also gets the list of task uuids that were cancelled. +/// - `Force` — in-flight jobs are written `Killed` on the spot (no cleanup) and +/// their agents are told to stop. pub async fn user_cancel_task_suite( user_id: i64, pool: &InfraPool, suite_uuid: Uuid, - _op: CancelTaskSuiteOp, + op: CancelTaskSuiteOp, ) -> Result<()> { let now = TimeDateTimeWithTimeZone::now_utc(); - pool.db - .transaction::<_, u64, Error>(|txn| { + let (suite_id, cancelled_task_uuids) = pool + .db + .transaction::<_, (i64, Vec), Error>(|txn| { Box::pin(async move { // Resolve the suite and authorize the caller's Write role in one join. let suite = @@ -606,19 +653,16 @@ pub async fn user_cancel_task_suite( .exec_with_returning(txn) .await?; - // The inflight subset is the only set tied to an executing agent: each - // carries the `runner_uuid` of the agent running it. A Force cancel must - // signal those agents to stop so they don't commit tasks we just archived. - // Collected now and intentionally left unused. - // TODO: wire this to a per-agent shutdown/TasksCancelled - // push once the agent/notification layer exists - let _inflight_signals: Vec<(Option, Uuid)> = tasks + // The in-flight subset is the only one an agent is holding right + // now. Those agents are told which task uuids vanished, so they + // stop rather than reporting against rows we just archived. + let cancelled_task_uuids: Vec = tasks .iter() .filter(|t| matches!(t.state, TaskState::Running | TaskState::Finished)) - .map(|t| (t.runner_uuid, t.uuid)) + .map(|t| t.uuid) .collect(); - let cancelled_count = tasks.len() as u64; + let suite_id = suite.id; if !tasks.is_empty() { let archived: Vec = tasks .into_iter() @@ -655,14 +699,42 @@ pub async fn user_cancel_task_suite( suite.completed_at = Set(Some(now)); suite.update(txn).await?; - Ok(cancelled_count) + Ok((suite_id, cancelled_task_uuids)) }) }) .await?; - // TODO: notify assigned/executing agents that this suite was cancelled (and, for - // Force, push per-agent TasksCancelled built from `_running_signals` above), then - // drop the suite's in-memory task buffer in the dispatcher. + // Tear down whatever is still running. `Force` writes the terminal itself so + // no further agent report is accepted; `Graceful` leaves the job alone and + // lets the agent walk it through cleanup to `Completed`. + let running_agents = match op { + CancelTaskSuiteOp::Force => { + crate::service::agent::job::kill_suite_jobs(&pool.db, suite_id, now).await? + } + CancelTaskSuiteOp::Graceful => { + crate::service::agent::job::agents_running_suite(&pool.db, suite_id).await? + } + }; + + crate::service::agent::notify_agents_by_id( + pool, + &running_agents, + AgentNotification::SuiteCancelled { + suite_uuid, + reason: "Suite was cancelled by a user".to_string(), + }, + ) + .await; + if !cancelled_task_uuids.is_empty() { + crate::service::agent::notify_agents_by_id( + pool, + &running_agents, + AgentNotification::TasksCancelled { + task_uuids: cancelled_task_uuids, + }, + ) + .await; + } Ok(()) } @@ -679,9 +751,9 @@ pub async fn user_override_agents_for_suite( ) -> Result { let now = TimeDateTimeWithTimeZone::now_utc(); - let errors = pool + let (suite_id, errors) = pool .db - .transaction::<_, HashMap, Error>(|txn| { + .transaction::<_, (Option, HashMap), Error>(|txn| { Box::pin(async move { // The suite is fixed for the whole batch, so a suite-level failure is request-level let suite = @@ -692,6 +764,7 @@ pub async fn user_override_agents_for_suite( } Err(ResolveError::Fatal(e)) => return Err(e), }; + let suite_id = suite.id; let mut errors = HashMap::new(); for (agent_uuid, action) in req.overrides { @@ -712,13 +785,161 @@ pub async fn user_override_agents_for_suite( } } - Ok(errors) + Ok((Some(suite_id), errors)) }) }) .await?; - // TODO: the manual overrides just changed the effective agent set. Once the - // scheduler exists, trigger a recompute for this suite. + // The effective agent set just changed. Notify the agents + if let Some(suite_id) = suite_id { + crate::service::agent::notify_suite_available(pool, suite_id).await; + } Ok(SuiteAgentOverrideResp { errors }) } + +/// Query a suite's jobs, newest first. +pub async fn user_query_suite_jobs( + user_id: i64, + pool: &InfraPool, + suite_uuid: Uuid, + query: SuiteJobsQueryReq, +) -> Result { + let suite = match authorize_suite(&pool.db, user_id, suite_uuid, UserGroupRole::Read).await { + Ok(suite) => suite, + Err(ResolveError::Item(e)) => return Err(Error::ApiError(ApiError::NotFound(e.msg))), + Err(ResolveError::Fatal(e)) => return Err(e), + }; + + let mut stmt = Query::select(); + if query.count { + stmt.expr(Expr::col((SuiteAgentJobs::Entity, SuiteAgentJobs::Column::Id)).count()); + } else { + stmt.columns([ + (SuiteAgentJobs::Entity, SuiteAgentJobs::Column::JobId), + (SuiteAgentJobs::Entity, SuiteAgentJobs::Column::State), + (SuiteAgentJobs::Entity, SuiteAgentJobs::Column::CreatedAt), + (SuiteAgentJobs::Entity, SuiteAgentJobs::Column::UpdatedAt), + ]) + .expr_as( + Expr::col((Agent::Entity, Agent::Column::Uuid)), + Alias::new("agent_uuid"), + ); + } + + stmt.from(SuiteAgentJobs::Entity) + // Left join: `agent_id` is nullable to leave room for a future detach. + .join( + sea_orm::JoinType::LeftJoin, + Agent::Entity, + Expr::col((Agent::Entity, Agent::Column::Id)).eq(Expr::col(( + SuiteAgentJobs::Entity, + SuiteAgentJobs::Column::AgentId, + ))), + ) + .and_where( + Expr::col((SuiteAgentJobs::Entity, SuiteAgentJobs::Column::TaskSuiteId)).eq(suite.id), + ); + + if let Some(ref states) = query.states { + let states_vec: Vec = states.iter().copied().collect(); + stmt.and_where( + Expr::col((SuiteAgentJobs::Entity, SuiteAgentJobs::Column::State)) + .eq(PgFunc::any(states_vec)), + ); + } + if let Some(agent_uuid) = query.agent_uuid { + stmt.and_where(Expr::col((Agent::Entity, Agent::Column::Uuid)).eq(agent_uuid)); + } + 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(SuiteJobsQueryResp { + count, + jobs: vec![], + }) + } else { + stmt.order_by_expr( + Expr::col((SuiteAgentJobs::Entity, SuiteAgentJobs::Column::JobId)).into(), + sea_orm::Order::Desc, + ); + let jobs = SuiteJobInfo::find_by_statement(builder.build(&stmt)) + .all(&pool.db) + .await?; + Ok(SuiteJobsQueryResp { + count: jobs.len() as u64, + jobs, + }) + } +} + +/// One job of a suite, with its hook executions embedded. A hook's artifacts +/// are downloaded through the existing artifact endpoints using its `uuid`. +pub async fn user_get_suite_job( + user_id: i64, + pool: &InfraPool, + suite_uuid: Uuid, + job_id: i32, +) -> Result { + let suite = match authorize_suite(&pool.db, user_id, suite_uuid, UserGroupRole::Read).await { + Ok(suite) => suite, + Err(ResolveError::Item(e)) => return Err(Error::ApiError(ApiError::NotFound(e.msg))), + Err(ResolveError::Fatal(e)) => return Err(e), + }; + + let job = SuiteAgentJobs::Entity::find() + .filter(SuiteAgentJobs::Column::TaskSuiteId.eq(suite.id)) + .filter(SuiteAgentJobs::Column::JobId.eq(job_id)) + .one(&pool.db) + .await? + .ok_or(Error::ApiError(ApiError::NotFound(format!( + "Job {job_id} of suite {suite_uuid}" + ))))?; + + let agent_uuid = match job.agent_id { + Some(agent_id) => Agent::Entity::find_by_id(agent_id) + .one(&pool.db) + .await? + .map(|a| a.uuid), + None => None, + }; + + let hooks = HookTasks::Entity::find() + .filter(HookTasks::Column::SuiteAgentJobId.eq(job.id)) + .all(&pool.db) + .await? + .into_iter() + .map(|hook| { + Ok(HookTaskInfo { + uuid: hook.uuid, + hook_type: hook.hook_type, + state: hook.state, + result: hook.result.map(serde_json::from_value).transpose()?, + started_at: hook.started_at, + completed_at: hook.completed_at, + }) + }) + .collect::>>()?; + + Ok(SuiteJobQueryResp { + info: SuiteJobInfo { + job_id: job.job_id, + state: job.state, + agent_uuid, + created_at: job.created_at, + updated_at: job.updated_at, + }, + hooks, + }) +} diff --git a/netmito/src/service/task.rs b/netmito/src/service/task.rs index af925024..c14507e2 100644 --- a/netmito/src/service/task.rs +++ b/netmito/src/service/task.rs @@ -45,9 +45,15 @@ fn check_exec_spec(spec: &ExecSpec) -> crate::error::Result<()> { Ok(()) } +#[derive(Clone, Copy)] enum Submitter { User, - Worker(Uuid), + /// A running task spawning a downstream child, identified by the parent's + /// uuid and the group the parent itself lives in. + Task { + upstream_task_uuid: Uuid, + parent_group_id: i64, + }, } async fn internal_submit_task( @@ -71,7 +77,9 @@ async fn internal_submit_task( // Used later when creating active tasks active model let (state, upstream_task_uuid) = match submitter { Submitter::User => (Set(crate::entity::state::TaskState::Ready), NotSet), - Submitter::Worker(upstream_task_uuid) => ( + Submitter::Task { + upstream_task_uuid, .. + } => ( Set(crate::entity::state::TaskState::Pending), Set(Some(upstream_task_uuid)), ), @@ -89,51 +97,18 @@ async fn internal_submit_task( .transaction::<_, (ActiveTasks::Model, Option), crate::error::Error>( |txn| { Box::pin(async move { - // Determine the owning group and (optionally) the suite, then - // atomically bump the group's task counter. When a suite is designated it - // is authoritative: the task's group is the suite's owning group, and the - // request's group_name must match it. - let (group_id, task_id, suite) = match suite_uuid { - None => { - // No suite: resolve the user-specified group, verify the caller's - // membership, and bump its task counter in one statement. - let mut upd = Group::Entity::update_many() - .col_expr( - Group::Column::TaskCount, - Expr::col((Group::Entity, Group::Column::TaskCount)).add(1), - ) - .col_expr(Group::Column::UpdatedAt, Expr::value(now)) - .filter(Group::Column::GroupName.eq(&group_name)) - .filter( - Expr::col((UserGroup::Entity, UserGroup::Column::UserId)) - .eq(creator_id), - ) - // TODO: when there is 'exec' access level, enforce access check here. - // - // .filter( - // Expr::col((UserGroup::Entity, UserGroup::Column::Role)) - // .gte(UserGroupRole::Write), - // ) - .filter(Expr::col((Group::Entity, Group::Column::Id)).eq( - Expr::col((UserGroup::Entity, UserGroup::Column::GroupId)), - )); - upd.query().from(UserGroup::Entity); - let group = upd - .exec_with_returning(txn) - .await? - .into_iter() - .next() - .ok_or_else(|| { - Error::ApiError(crate::error::ApiError::NotFound(format!( - "User doesn't have permission or group with name {}", - group_name - ))) - })?; - (group.id, group.task_count, None) - } - Some(suite_uuid) => { - // Suite designated: resolve it (and the caller's membership in its - // owning group) in one join. The suite is authoritative for the group. + // Resolve the designated suite, if any. How it is authorized + // depends on the submitter: a user reaches a suite through their + // own group memberships, while a task spawning a child may target + // any suite owned by the group the parent task itself lives in. + // + // The parent's creator is deliberately not re-checked, so a long-running suite + // keeps spawning work after that user leaves the group. + let suite = match (suite_uuid, submitter) { + (None, _) => None, + (Some(suite_uuid), Submitter::User) => { + // Resolve the suite and the caller's membership in its + // owning group in one join. let builder = txn.get_database_backend(); let suite_stmt = Query::select() .columns([ @@ -186,47 +161,140 @@ async fn internal_submit_task( suite_uuid ))) })?; - // The suite must be able to accept new tasks. - if !suite.state.can_accept_tasks() { - return Err(Error::ApiError( - crate::error::ApiError::InvalidRequest(format!( - "Suite is in {} state and cannot accept new tasks", - suite.state - )), - )); - } - // Derive the owning group from the suite and bump its counter. + Some(suite) + } + ( + Some(suite_uuid), + Submitter::Task { + parent_group_id, .. + }, + ) => Some( + TaskSuites::Entity::find() + .filter(TaskSuites::Column::Uuid.eq(suite_uuid)) + .filter(TaskSuites::Column::GroupId.eq(parent_group_id)) + .one(txn) + .await? + .ok_or_else(|| { + // A suite owned by another group is reported the + // same way as one that does not exist. + Error::ApiError(crate::error::ApiError::NotFound(format!( + "Suite with uuid {} in the parent task's group", + suite_uuid + ))) + })?, + ), + }; + + // The suite must be able to accept new tasks. + if let Some(suite) = &suite { + if !suite.state.can_accept_tasks() { + return Err(Error::ApiError(crate::error::ApiError::InvalidRequest( + format!( + "Suite is in {} state and cannot accept new tasks", + suite.state + ), + ))); + } + } + + // The owning group, when it is already pinned by something other + // than `group_name`: a designated suite owns the task, and a + // task-spawned child always stays in its parent's group. `None` + // is the plain user path, where the group is resolved from + // `group_name` and the caller's membership. + let pinned_group_id = match (&suite, submitter) { + (Some(suite), _) => Some(suite.group_id), + ( + None, + Submitter::Task { + parent_group_id, .. + }, + ) => Some(parent_group_id), + (None, Submitter::User) => None, + }; + + // Bump the owning group's task counter atomically. + let (group_id, task_id) = match pinned_group_id { + Some(pinned_group_id) => { let group = Group::Entity::update_many() .col_expr( Group::Column::TaskCount, Expr::col(Group::Column::TaskCount).add(1), ) .col_expr(Group::Column::UpdatedAt, Expr::value(now)) - .filter(Group::Column::Id.eq(suite.group_id)) + .filter(Group::Column::Id.eq(pinned_group_id)) .exec_with_returning(txn) .await? .into_iter() .next() .ok_or_else(|| { - // The suite's group row must exist (FK) - tracing::error!( - "Owning group {} of suite {} not found", - suite.group_id, - suite_uuid - ); + // The row must exist: both the suite and the + // parent task hold an FK to it. + tracing::error!("Owning group {} not found", pinned_group_id); Error::ApiError(crate::error::ApiError::InternalServerError) })?; - // The request's group_name must match the suite's owning group. - // Safe to reveal the real group: access to the suite is proven. + // The request's group_name must match if group.group_name != group_name { + let owner = match &suite { + Some(suite) => { + format!( + "Suite {} belongs to group {}", + suite.uuid, group.group_name + ) + } + None => format!( + "The parent task belongs to group {}", + group.group_name + ), + }; return Err(Error::ApiError( crate::error::ApiError::InvalidRequest(format!( - "Suite {} belongs to group {}, not {}", - suite_uuid, group.group_name, group_name + "{}, not {}", + owner, group_name )), )); } - (group.id, group.task_count, Some(suite)) + (group.id, group.task_count) + } + None => { + // Resolve the user-specified group, verify the caller's + // membership, and bump its task counter in one statement. + // Membership goes through a subquery rather than a + // joined `FROM user_group`: sea-orm emits an + // unqualified `RETURNING "id", …`, which Postgres + // rejects as ambiguous the moment a second table + // with an `id` column is in scope. Every suite-less + // `POST /tasks` 500s without this. + let member_of = Query::select() + .column(UserGroup::Column::GroupId) + .from(UserGroup::Entity) + .and_where(Expr::col(UserGroup::Column::UserId).eq(creator_id)) + // TODO: when there is 'exec' access level, enforce access check here. + // + // .and_where( + // Expr::col(UserGroup::Column::Role) + // .gte(UserGroupRole::Write), + // ) + .to_owned(); + let group = Group::Entity::update_many() + .col_expr( + Group::Column::TaskCount, + Expr::col((Group::Entity, Group::Column::TaskCount)).add(1), + ) + .col_expr(Group::Column::UpdatedAt, Expr::value(now)) + .filter(Group::Column::GroupName.eq(&group_name)) + .filter(Group::Column::Id.in_subquery(member_of)) + .exec_with_returning(txn) + .await? + .into_iter() + .next() + .ok_or_else(|| { + Error::ApiError(crate::error::ApiError::NotFound(format!( + "User doesn't have permission or group with name {}", + group_name + ))) + })?; + (group.id, group.task_count) } }; @@ -264,6 +332,7 @@ async fn internal_submit_task( exec_options: Set(exec_options_json), result: Set(None), upstream_task_uuid, + task_suite_id: Set(suite.as_ref().map(|s| s.id)), ..Default::default() }; let task = task.insert(txn).await?; @@ -283,8 +352,11 @@ async fn internal_submit_task( // Not pending, notify agent/worker depending on whether it belongs to a suite match suite { - Some(_suite) => { - // TODO: If this task belongs to a suite, notify agents + Some(suite) => { + // Suite tasks are pulled by agents from the suite, not pushed into + // worker queues; all this does is wake the idle eligible agents so + // one picks the suite up without waiting for its next heartbeat. + crate::service::agent::notify_suite_available(pool, suite.id).await; } None => { let builder = pool.db.get_database_backend(); @@ -336,13 +408,29 @@ pub async fn user_submit_task( internal_submit_task(pool, creator_id, Submitter::User, req).await } +/// Submit the downstream child of a running task, on behalf of the worker or +/// agent executing it. +/// +/// The child stays in the parent's group, and `req.suite_uuid` may name any +/// suite that group owns — the parent's own suite, a sibling one, or `None` for +/// a suite-less task dispatched to workers. pub async fn worker_submit_pending_task( pool: &InfraPool, creator_id: i64, upstream_task_uuid: Uuid, + parent_group_id: i64, req: SubmitTaskReq, ) -> crate::error::Result { - internal_submit_task(pool, creator_id, Submitter::Worker(upstream_task_uuid), req).await + internal_submit_task( + pool, + creator_id, + Submitter::Task { + upstream_task_uuid, + parent_group_id, + }, + req, + ) + .await } pub async fn worker_trigger_pending_task(pool: &InfraPool, uuid: Uuid) -> crate::error::Result<()> { @@ -359,6 +447,12 @@ pub async fn worker_trigger_pending_task(pool: &InfraPool, uuid: Uuid) -> crate: task.state = Set(TaskState::Ready); task.updated_at = Set(TimeDateTimeWithTimeZone::now_utc()); let task = task.update(&pool.db).await?; + // A suite task never enters the worker queues — agents pull it from its + // suite. Wake the idle eligible agents instead. + if let Some(suite_id) = task.task_suite_id { + crate::service::agent::notify_suite_available(pool, suite_id).await; + return Ok(()); + } // Batch add task to worker task queues let builder = pool.db.get_database_backend(); let tasks_stmt = Query::select() diff --git a/netmito/src/service/worker/mod.rs b/netmito/src/service/worker/mod.rs index 314d8bbd..aa0d288a 100644 --- a/netmito/src/service/worker/mod.rs +++ b/netmito/src/service/worker/mod.rs @@ -630,19 +630,9 @@ pub async fn report_task( if user.state != UserState::Active { return Err(AuthError::PermissionDenied.into()); } - // Load group information to verify group name matches - let group = Group::Entity::find_by_id(task.group_id) - .one(&pool.db) - .await? - .ok_or(ApiError::NotFound("Group not found".to_string()))?; - if req.group_name != group.group_name { - return Err(ApiError::InvalidRequest(format!( - "Group name mismatch: expected '{}', got '{}'", - group.group_name, req.group_name - )) - .into()); - } - // Verify worker has permission to access this group + // Verify worker has permission to access this group. The child's own + // group and suite are pinned to the parent's by the submit path, + // which also rejects a mismatched `req.group_name`. GroupWorker::Entity::find() .filter(GroupWorker::Column::WorkerId.eq(worker_id)) .filter(GroupWorker::Column::GroupId.eq(task.group_id)) @@ -652,17 +642,22 @@ pub async fn report_task( .ok_or(ApiError::AuthError( crate::error::AuthError::PermissionDenied, ))?; - let resp = - service::task::worker_submit_pending_task(pool, task.creator_id, task.uuid, req) - .await - .map_err(|e| match e { - crate::error::Error::AuthError(err) => ApiError::AuthError(err), - crate::error::Error::ApiError(e) => e, - _ => { - tracing::error!("{}", e); - ApiError::InternalServerError - } - })?; + let resp = service::task::worker_submit_pending_task( + pool, + task.creator_id, + task.uuid, + task.group_id, + req, + ) + .await + .map_err(|e| match e { + crate::error::Error::AuthError(err) => ApiError::AuthError(err), + crate::error::Error::ApiError(e) => e, + _ => { + tracing::error!("{}", e); + ApiError::InternalServerError + } + })?; let task = ActiveTask::ActiveModel { id: Set(task_id), updated_at: Set(now), diff --git a/netmito/src/worker.rs b/netmito/src/worker.rs index 39527a39..44a16b1b 100644 --- a/netmito/src/worker.rs +++ b/netmito/src/worker.rs @@ -1,22 +1,12 @@ -use std::os::unix::process::ExitStatusExt; use std::path::PathBuf; -use std::process::ExitStatus; use std::sync::atomic::AtomicBool; use std::sync::Arc; -use async_compression::tokio::write::GzipEncoder; use futures::StreamExt; -use nix::sys::signal::{self, Signal}; -use nix::unistd::Pid; -use redis::aio::{MultiplexedConnection, PubSub}; +use redis::aio::MultiplexedConnection; use redis::AsyncCommands; -use reqwest::header::CONTENT_LENGTH; use reqwest::{Client, StatusCode}; -use tokio::io::{AsyncReadExt, AsyncWriteExt}; -use tokio::process::Command; -use tokio::sync::mpsc; use tokio::time::Instant; -use tokio_tar::{Builder, Header}; use tokio_util::sync::CancellationToken; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt}; use url::Url; @@ -26,9 +16,9 @@ use crate::config::TracingGuard; use crate::entity::content::ArtifactContentType; use crate::entity::state::TaskExecState; use crate::error::RequestError; +use crate::executor::{execute, ExecClient, Executor, UploadTarget}; use crate::schema::*; use crate::service::auth::get_and_prompt_username; -use crate::service::s3::download_file; use crate::{ config::{WorkerConfig, WorkerConfigCli}, error::{Error, ErrorMsg}, @@ -47,179 +37,319 @@ pub struct MitoWorker { redis_client: Option, } -pub struct TaskExecutor { - pub(crate) task_client: Client, - pub(crate) task_credential: String, - pub(crate) task_url: Url, - pub(crate) task_cancel_token: CancellationToken, - pub(crate) coordinator_force_exit: Arc, - pub(crate) polling_interval: std::time::Duration, - pub(crate) task_cache_path: PathBuf, - pub(crate) task_redis_conn: Option, - pub(crate) task_redis_pubsub: Option, +/// The worker's half of the execution seam: one task, reported to +/// `POST /workers/tasks` with the worker's own credential, its exec state +/// announced over redis. +struct WorkerTaskClient { + http_client: Client, + credential: String, + coordinator_addr: Url, + cancel_token: CancellationToken, + coordinator_force_exit: Arc, + polling_interval: std::time::Duration, + task_id: i64, + task_uuid: Uuid, + upstream_task_uuid: Option, + /// A clone of the worker's single multiplexed connection, not a new one. + redis_conn: Option, + /// Kept whole so `watch` can open its own pubsub connection for the wait and + /// drop it afterwards. + redis_client: Option, } -impl TaskExecutor { - async fn set_task_state_ex(&mut self, uuid: &Uuid, state: i32, ex: u64) { - if let Some(ref mut conn) = self.task_redis_conn { - tracing::trace!("Set task state: {} -> {}", uuid, state); - let _: Result = conn.set_ex(format!("task:{uuid}"), state, ex).await; +impl WorkerTaskClient { + fn api_url(&self, path: &str) -> Url { + let mut url = self.coordinator_addr.clone(); + url.set_path(path); + url + } + + /// POST one report, retrying connection failures. Returns the presigned + /// upload URL for an `Upload` op, and `None` when the coordinator answered + /// something that means "stop bothering with this task". + async fn report(&mut self, op: ReportTaskOp) -> crate::error::Result> { + // Only an upload can legitimately be refused: the group's storage quota + // is enforced there. Any other 403 is a real error. + let upload = matches!(op, ReportTaskOp::Upload { .. }); + let req = ReportTaskReq { + id: self.task_id, + op, + }; + let url = self.api_url("workers/tasks"); + loop { + let resp = self + .http_client + .post(url.as_str()) + .json(&req) + .bearer_auth(&self.credential) + .send() + .await; + match resp { + Ok(resp) => { + if resp.status().is_success() { + return resp + .json::() + .await + .map(Some) + .map_err(|e| RequestError::from(e).into()); + } else if resp.status() == StatusCode::UNAUTHORIZED { + tracing::info!("Report task failed with coordinator force exit"); + self.coordinator_force_exit + .store(true, std::sync::atomic::Ordering::Release); + self.cancel_token.cancel(); + return Ok(None); + } else if resp.status() == StatusCode::NOT_FOUND { + tracing::debug!("Task not found, ignore and go on for next cycle"); + return Ok(None); + } else if upload && resp.status() == StatusCode::FORBIDDEN { + let resp: ErrorMsg = resp + .json() + .await + .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); + tracing::info!( + "Request upload url failed with permission denied: {}", + resp.msg + ); + return Ok(None); + } else { + let resp: ErrorMsg = resp + .json() + .await + .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); + return Err(Error::Custom(format!( + "Report task failed with error: {}", + resp.msg + ))); + } + } + Err(e) => { + if e.is_connect() && e.is_request() { + tracing::error!( + "Report task failed with connection error: {}. Retry after {:?}", + e, + self.polling_interval + ); + tokio::select! { + biased; + _ = self.cancel_token.cancelled() => return Ok(None), + _ = tokio::time::sleep(self.polling_interval) => {}, + } + continue; + } else { + return Err(RequestError::from(e).into()); + } + } + } } } - async fn set_task_state(&mut self, uuid: &Uuid, state: i32) { - if let Some(ref mut conn) = self.task_redis_conn { - tracing::trace!("Set task state: {} -> {}", uuid, state); - let _: Result = conn.set(format!("task:{uuid}"), state).await; + async fn cached_state(&mut self) -> Option { + let conn = self.redis_conn.as_mut()?; + tracing::trace!("Get task state: {}", self.task_uuid); + let state: Result = conn.get(format!("task:{}", self.task_uuid)).await; + state.ok().map(TaskExecState::from) + } + + /// Ask the coordinator directly whether the watched task has got there. The + /// redis mirror can be missing or stale, so the poll is the source of truth. + async fn poll_watched_task(&self, uuid: &Uuid, target: TaskExecState) -> bool { + let resp = self + .http_client + .get(self.api_url(&format!("workers/tasks/{uuid}")).as_str()) + .bearer_auth(&self.credential) + .send() + .await; + match resp { + Ok(resp) => { + if resp.status().is_success() { + match resp.json::().await { + Ok(task) => task.info.state.is_reach(&target, task.info.result), + Err(_) => false, + } + } else { + let resp: ErrorMsg = resp + .json() + .await + .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); + tracing::error!("Get Task failed with error: {}", resp.msg); + false + } + } + Err(e) => { + tracing::error!("Get task failed with error: {}", e); + false + } } } +} - async fn get_task_state(&mut self, uuid: &Uuid) -> Option { - if let Some(ref mut conn) = self.task_redis_conn { - tracing::trace!("Get task state: {}", uuid); - let state: Result = conn.get(format!("task:{uuid}")).await; - state.ok().map(TaskExecState::from) - } else { - None +#[async_trait::async_trait] +impl ExecClient for WorkerTaskClient { + fn describe(&self) -> String { + format!("task {}", self.task_uuid) + } + + fn exec_env(&self) -> Vec<(&'static str, String)> { + let mut env = vec![("MITO_TASK_UUID", self.task_uuid.to_string())]; + if let Some(uuid) = self.upstream_task_uuid { + env.push(("MITO_UPSTREAM_TASK_UUID", uuid.to_string())); } + env + } + + fn supports_child_tasks(&self) -> bool { + true + } + + async fn report_finish( + &mut self, + finished: bool, + _result: &TaskResultSpec, + ) -> crate::error::Result<()> { + let op = if finished { + ReportTaskOp::Finish + } else { + ReportTaskOp::Cancel + }; + self.report(op).await.map(|_| ()) } - async fn publish_state(&mut self, uuid: &Uuid, state: i32) { - if let Some(ref mut conn) = self.task_redis_conn { - tracing::trace!("Publish task state: {} -> {}", uuid, state); - let _: Result = conn.publish(format!("task:{uuid}"), state).await; + async fn request_upload( + &mut self, + content_type: ArtifactContentType, + content_length: u64, + ) -> crate::error::Result { + match self + .report(ReportTaskOp::Upload { + content_type, + content_length, + }) + .await? + { + Some(resp) => Ok(match resp.url { + Some(url) => UploadTarget::Url(url), + None => UploadTarget::Skip, + }), + None => Ok(UploadTarget::Stop), } } - async fn announce_task_state(&mut self, uuid: &Uuid, state: i32) { - self.set_task_state(uuid, state).await; - self.publish_state(uuid, state).await; + async fn report_commit(&mut self, result: TaskResultSpec) -> crate::error::Result<()> { + self.report(ReportTaskOp::Commit(result)).await.map(|_| ()) } - async fn announce_task_state_ex(&mut self, uuid: &Uuid, state: i32, ex: u64) { - self.set_task_state_ex(uuid, state, ex).await; - self.publish_state(uuid, state).await; + async fn submit_child_task(&mut self, req: SubmitTaskReq) -> crate::error::Result<()> { + self.report(ReportTaskOp::Submit(Box::new(req))) + .await + .map(|_| ()) } - async fn watch_task(&mut self, uuid: &Uuid, state: TaskExecState) { - tracing::debug!("Watch task: {} -> {:?}", uuid, state); - let mut wait_until = Instant::now(); - if let Some(pubsub) = self.task_redis_pubsub.as_mut() { - let channel_name = format!("task:{uuid}"); - let _ = pubsub.subscribe(&channel_name).await; - let mut stream = pubsub.on_message(); - loop { - tokio::select! { - biased; - msg = stream.next() => { - if let Some(msg) = msg { - if msg.get_channel_name() == channel_name { - if let Ok(task_state) = msg.get_payload::() { - let cur_state = TaskExecState::from(task_state); - if cur_state.is_reach(&state) { - break; - } - } - } - } - }, - _ = tokio::time::sleep_until(wait_until) => { - wait_until = Instant::now() + std::time::Duration::from_secs(30); - let cur_state = if let Some(ref mut conn) = self.task_redis_conn { - tracing::trace!("Get task state: {}", uuid); - let state: Result = conn.get(format!("task:{uuid}")).await; - state.ok().map(TaskExecState::from) - } else { - None - }; - if let Some(cur_state) = cur_state { - if cur_state.is_reach(&state) { - break; - } - } - self.task_url - .set_path(format!("workers/tasks/{uuid}").as_str()); - let resp = self - .task_client - .get(self.task_url.as_str()) - .bearer_auth(&self.task_credential) - .send() - .await; - match resp { - Ok(resp) => { - if resp.status().is_success() { - if let Ok(task) = resp.json::().await { - if task.info.state.is_reach(&state, task.info.result) { - break; - } - } - } else { - let resp: ErrorMsg = resp - .json() - .await - .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); - tracing::error!("Get Task failed with error: {}", resp.msg); - } - } - Err(e) => { - tracing::error!("Get task failed with error: {}", e); - } - } - }, - } + fn artifact_download_req( + &self, + uuid: Uuid, + content_type: ArtifactContentType, + ) -> reqwest::RequestBuilder { + let content_type = serde_json::to_value(content_type) + .ok() + .and_then(|v| v.as_str().map(str::to_string)) + .unwrap_or_else(|| "result".to_string()); + self.http_client + .get( + self.api_url(&format!("workers/tasks/{uuid}/artifacts/{content_type}")) + .as_str(), + ) + .bearer_auth(&self.credential) + } + + fn attachment_download_req(&self, key: &str) -> reqwest::RequestBuilder { + let uuid = self.task_uuid; + self.http_client + .get( + self.api_url(&format!("workers/tasks/{uuid}/attachments/{key}")) + .as_str(), + ) + .bearer_auth(&self.credential) + } + + async fn announce_state(&mut self, state: TaskExecState, ex: Option) { + let uuid = self.task_uuid; + let state = state as i32; + let Some(conn) = self.redis_conn.as_mut() else { + return; + }; + tracing::trace!("Set task state: {} -> {}", uuid, state); + match ex { + Some(ex) => { + let _: Result = conn.set_ex(format!("task:{uuid}"), state, ex).await; } - } else { - loop { - tokio::time::sleep_until(wait_until).await; - wait_until = Instant::now() + std::time::Duration::from_secs(30); - if let Some(cur_state) = self.get_task_state(uuid).await { - if cur_state.is_reach(&state) { - break; - } + None => { + let _: Result = conn.set(format!("task:{uuid}"), state).await; + } + } + tracing::trace!("Publish task state: {} -> {}", uuid, state); + let _: Result = conn.publish(format!("task:{uuid}"), state).await; + } + + fn can_watch(&self) -> bool { + self.redis_conn.is_some() && self.redis_client.is_some() + } + + /// Subscribe to the watched task's channel and, in parallel, re-poll every + /// 30s in case the notification was missed or predates the subscription. + async fn watch(&mut self, uuid: &Uuid, target: TaskExecState) { + tracing::debug!("Watch task: {} -> {:?}", uuid, target); + let channel_name = format!("task:{uuid}"); + // Opened for this wait only; dropping it on the way out is what used to + // need an explicit unsubscribe. + let mut pubsub = match self.redis_client.as_ref() { + Some(client) => client + .get_async_pubsub() + .await + .inspect_err(|e| tracing::warn!("Cannot open a redis pubsub connection: {}", e)) + .ok(), + None => None, + }; + if let Some(pubsub) = pubsub.as_mut() { + let _ = pubsub.subscribe(&channel_name).await; + } + let mut stream = pubsub.as_mut().map(|pubsub| pubsub.on_message()); + + let mut wait_until = Instant::now(); + loop { + let published = async { + match stream.as_mut() { + Some(stream) => stream.next().await, + // No pubsub: fall through to the poll arm forever. + None => std::future::pending().await, } - self.task_url - .set_path(format!("workers/tasks/{uuid}").as_str()); - let resp = self - .task_client - .get(self.task_url.as_str()) - .bearer_auth(&self.task_credential) - .send() - .await; - match resp { - Ok(resp) => { - if resp.status().is_success() { - if let Ok(task) = resp.json::().await { - if task.info.state.is_reach(&state, task.info.result) { + }; + tokio::select! { + biased; + msg = published => { + if let Some(msg) = msg { + if msg.get_channel_name() == channel_name { + if let Ok(state) = msg.get_payload::() { + if TaskExecState::from(state).is_reach(&target) { break; } } - } else { - let resp: ErrorMsg = resp - .json() - .await - .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); - tracing::error!("Get Task failed with error: {}", resp.msg); } } - Err(e) => { - tracing::error!("Get task failed with error: {}", e); + }, + _ = tokio::time::sleep_until(wait_until) => { + wait_until = Instant::now() + std::time::Duration::from_secs(30); + if let Some(state) = self.cached_state().await { + if state.is_reach(&target) { + break; + } } - } + if self.poll_watched_task(uuid, target).await { + break; + } + }, } } } - - pub async fn subscribe_task_exec_state(&mut self, uuid: &Uuid) { - if let Some(pubsub) = self.task_redis_pubsub.as_mut() { - let _ = pubsub.subscribe(format!("task:{uuid}")).await; - } - } - - pub async fn unsubscribe_task_exec_state(&mut self, uuid: &Uuid) { - if let Some(pubsub) = self.task_redis_pubsub.as_mut() { - let _ = pubsub.unsubscribe(format!("task:{uuid}")).await; - } - } } impl MitoWorker { @@ -337,38 +467,6 @@ impl MitoWorker { } } - pub async fn get_task_executor(&self) -> TaskExecutor { - let task_redis_conn = if let Some(ref client) = self.redis_client { - client - .get_multiplexed_tokio_connection() - .await - .inspect_err(|e| tracing::warn!("{}", e)) - .ok() - } else { - None - }; - let task_redis_pubsub = if let Some(ref client) = self.redis_client { - client - .get_async_pubsub() - .await - .inspect_err(|e| tracing::warn!("{}", e)) - .ok() - } else { - None - }; - TaskExecutor { - task_client: self.http_client.clone(), - task_credential: self.credential.clone(), - task_url: self.config.coordinator_addr.clone(), - task_cancel_token: self.cancel_token.clone(), - coordinator_force_exit: self.coordinator_force_exit.clone(), - polling_interval: self.config.polling_interval, - task_cache_path: self.cache_path.clone(), - task_redis_conn, - task_redis_pubsub, - } - } - pub async fn run(&mut self) -> crate::error::Result<()> { tracing::info!("Worker is running"); let mut heartbeat_url = self.config.coordinator_addr.clone(); @@ -425,88 +523,66 @@ impl MitoWorker { } } }); - let mut task_executor = self.get_task_executor().await; + let mut fetcher = TaskFetcher { + http_client: self.http_client.clone(), + credential: self.credential.clone(), + fetch_url: { + let mut url = self.config.coordinator_addr.clone(); + url.set_path("workers/tasks"); + url + }, + cancel_token: self.cancel_token.clone(), + coordinator_force_exit: self.coordinator_force_exit.clone(), + polling_interval: self.config.polling_interval, + }; + let redis_conn = match self.redis_client { + Some(ref client) => client + .get_multiplexed_tokio_connection() + .await + .inspect_err(|e| tracing::warn!("{}", e)) + .ok(), + None => None, + }; + let coordinator_addr = self.config.coordinator_addr.clone(); + let redis_client = self.redis_client.clone(); + let cache_path = self.cache_path.clone(); + let http_client = self.http_client.clone(); + let credential = self.credential.clone(); + let polling_interval = self.config.polling_interval; + let cancel_token = self.cancel_token.clone(); + let coordinator_force_exit = self.coordinator_force_exit.clone(); let task_hd = tokio::spawn(async move { loop { - if task_executor.task_cancel_token.is_cancelled() { + let task = match fetcher.next_task().await { + FetchOutcome::Task(task) => *task, + FetchOutcome::Idle => continue, + FetchOutcome::Stop => break, + }; + let mut executor = Executor { + cancel_token: cancel_token.clone(), + polling_interval, + cache_path: cache_path.clone(), + http_client: http_client.clone(), + client: Box::new(WorkerTaskClient { + http_client: http_client.clone(), + credential: credential.clone(), + coordinator_addr: coordinator_addr.clone(), + cancel_token: cancel_token.clone(), + coordinator_force_exit: coordinator_force_exit.clone(), + polling_interval, + task_id: task.id, + task_uuid: task.uuid, + upstream_task_uuid: task.upstream_task_uuid, + redis_conn: redis_conn.clone(), + redis_client: redis_client.clone(), + }), + }; + if let Err(e) = execute(&mut executor, task.spec, task.exec_options.as_ref()).await + { + tracing::error!("Task execution failed: {}", e); + cancel_token.cancel(); break; } - task_executor.task_url.set_path("workers/tasks"); - let resp = task_executor - .task_client - .get(task_executor.task_url.as_str()) - .bearer_auth(&task_executor.task_credential) - .send() - .await; - match resp { - Ok(resp) => { - if resp.status().is_success() { - match resp.json::>().await { - Ok(task) => match task { - Some(task) => { - match execute_task(task, &mut task_executor).await { - Ok(_) => {} - Err(e) => { - tracing::error!("Task execution failed: {}", e); - task_executor.task_cancel_token.cancel(); - } - } - } - None => { - tracing::debug!( - "No task fetched. Next fetch after {:?}", - task_executor.polling_interval - ); - tokio::select! { - biased; - _ = task_executor.task_cancel_token.cancelled() => break, - _ = tokio::time::sleep(task_executor.polling_interval) => {}, - } - } - }, - Err(e) => { - tracing::error!("Failed to parse task specification: {}", e); - task_executor.task_cancel_token.cancel(); - break; - } - } - } else if resp.status() == StatusCode::UNAUTHORIZED { - tracing::info!("Task fetch failed with coordinator force exit"); - task_executor - .coordinator_force_exit - .store(true, std::sync::atomic::Ordering::Release); - task_executor.task_cancel_token.cancel(); - break; - } else { - let resp: ErrorMsg = resp - .json() - .await - .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); - tracing::error!("Task fetch failed with error: {}", resp.msg); - task_executor.task_cancel_token.cancel(); - break; - } - } - Err(e) => { - if e.is_connect() && e.is_request() { - tracing::error!( - "Fetching task failed with connection error: {}. Retry after {:?}", - e, - task_executor.polling_interval - ); - tokio::select! { - biased; - _ = task_executor.task_cancel_token.cancelled() => break, - _ = tokio::time::sleep(task_executor.polling_interval) => {}, - } - continue; - } else { - tracing::error!("Fetching task failed with error: {}", e); - task_executor.task_cancel_token.cancel(); - break; - } - } - } } }); tokio::select! { @@ -545,1025 +621,94 @@ impl MitoWorker { } } -enum ProcessOutput { - WithLog { - stdout: Vec, - stderr: Vec, - exit_status: ExitStatus, - }, - WithoutLog { - exit_status: ExitStatus, - }, +/// What one poll of `GET /workers/tasks` came back with. +enum FetchOutcome { + Task(Box), + /// Nothing to run right now; the caller should poll again. + Idle, + /// The worker is done — shut down or told to exit. + Stop, } -impl ProcessOutput { - fn get_exit_status(&self) -> ExitStatus { - match self { - ProcessOutput::WithLog { exit_status, .. } => *exit_status, - ProcessOutput::WithoutLog { exit_status } => *exit_status, - } - } -} - -enum TaskResult { - Finish(ProcessOutput), - Timeout(ProcessOutput), -} - -impl TaskResult { - fn state(&self) -> (bool, ExitStatus) { - match self { - TaskResult::Finish(output) => (true, output.get_exit_status()), - TaskResult::Timeout(output) => (false, output.get_exit_status()), - } - } - - fn get_output(self) -> ProcessOutput { - match self { - TaskResult::Finish(output) => output, - TaskResult::Timeout(output) => output, - } - } -} - -async fn execute_task( - task: WorkerTaskResp, - task_executor: &mut TaskExecutor, -) -> crate::error::Result<()> { - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::FetchResource as i32, 360) - .await; - // Allow downloading resources for at most 30 minutes - let timeout_until = tokio::time::Instant::now() + std::time::Duration::from_secs(1800); - for resource in task.spec.resources { - match resource.remote_file { - RemoteResource::Artifact { uuid, content_type } => { - let content_serde_val = serde_json::to_value(content_type)?; - let content_serde_str = content_serde_val.as_str().unwrap_or("result"); - task_executor.task_url.set_path(&format!( - "workers/tasks/{uuid}/artifacts/{content_serde_str}" - )); - } - RemoteResource::Attachment { key } => { - let uuid = task.uuid; - task_executor - .task_url - .set_path(&format!("workers/tasks/{uuid}/attachments/{key}")); - } - }; - let resp = loop { - match task_executor - .task_client - .get(task_executor.task_url.as_str()) - .bearer_auth(&task_executor.task_credential) - .send() - .await - { - Ok(resp) => break resp, - Err(e) => { - if e.is_connect() && e.is_request() { - tracing::error!( - "Fetch resource info failed with connection error: {}. Retry after {:?}", - e, - task_executor.polling_interval - ); - tokio::select! { - biased; - _ = task_executor.task_cancel_token.cancelled() => return Ok(()), - _ = tokio::time::sleep(task_executor.polling_interval) => {}, - _ = tokio::time::sleep_until(timeout_until) => { - tracing::debug!("Fetching resource timeout, commit this task as canceled"); - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - report_task(task_executor, req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { - exit_status: 0, - msg: Some(TaskResultMessage::FetchResourceTimeout), - }), - }; - report_task(task_executor, req).await?; - task_executor - .announce_task_state_ex( - &task.uuid, - TaskExecState::FetchResourceTimeout as i32, - 60, - ) - .await; - task_executor - .announce_task_state_ex( - &task.uuid, - TaskExecState::TaskCommitted as i32, - 60, - ) - .await; - return Ok(()); - } - } - continue; - } else { - task_executor - .announce_task_state_ex( - &task.uuid, - TaskExecState::FetchResourceError as i32, - 60, - ) - .await; - return Err(RequestError::from(e).into()); - } - } - } - }; - if resp.status().is_success() { - let download_resp = resp - .json::() - .await - .map_err(RequestError::from)?; - let local_path = task_executor - .task_cache_path - .join("resource") - .join(resource.local_path); - tokio::select! { - biased; - res = download_file(&task_executor.task_client, &download_resp, local_path, false) => { - if let Err(e) = res { - tracing::error!("Failed to download resource: {}", e); - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - report_task(task_executor, req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { - exit_status: 0, - msg: Some(TaskResultMessage::ResourceForbidden), - }), - }; - report_task(task_executor, req).await?; - task_executor - .announce_task_state( - &task.uuid, - TaskExecState::FetchResourceForbidden as i32, - ) - .await; - task_executor - .announce_task_state_ex( - &task.uuid, - TaskExecState::TaskCommitted as i32, - 60, - ) - .await; - return Ok(()); - } - } - _ = task_executor.task_cancel_token.cancelled() => return Ok(()), - _ = tokio::time::sleep(std::time::Duration::from_secs(120)) => { - tracing::debug!("Fetching resource timeout, commit this task as canceled"); - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - report_task(task_executor, req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { - exit_status: 0, - msg: Some(TaskResultMessage::FetchResourceTimeout), - }), - }; - report_task(task_executor, req).await?; - task_executor - .announce_task_state_ex( - &task.uuid, - TaskExecState::FetchResourceTimeout as i32, - 60, - ) - .await; - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::TaskCommitted as i32, 60) - .await; - return Ok(()); - } - _ = tokio::time::sleep_until(timeout_until) => { - tracing::debug!("Fetching resource timeout, commit this task as canceled"); - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - report_task(task_executor, req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { - exit_status: 0, - msg: Some(TaskResultMessage::FetchResourceTimeout), - }), - }; - report_task(task_executor, req).await?; - task_executor - .announce_task_state_ex( - &task.uuid, - TaskExecState::FetchResourceTimeout as i32, - 60, - ) - .await; - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::TaskCommitted as i32, 60) - .await; - return Ok(()); - } - } - } else if resp.status() == StatusCode::NOT_FOUND { - tracing::debug!("Resource not found, commit this task as canceled"); - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - report_task(task_executor, req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { - exit_status: 0, - msg: Some(TaskResultMessage::ResourceNotFound), - }), - }; - report_task(task_executor, req).await?; - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::FetchResourceNotFound as i32, 60) - .await; - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::TaskCommitted as i32, 60) - .await; - return Ok(()); - } else if resp.status() == StatusCode::FORBIDDEN { - tracing::debug!("Resource is forbidden to be fetched, commit this task as canceled"); - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - report_task(task_executor, req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { - exit_status: 0, - msg: Some(TaskResultMessage::ResourceForbidden), - }), - }; - report_task(task_executor, req).await?; - task_executor - .announce_task_state_ex( - &task.uuid, - TaskExecState::FetchResourceForbidden as i32, - 60, - ) - .await; - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::TaskCommitted as i32, 60) - .await; - return Ok(()); - } else { - let resp: ErrorMsg = resp - .json() - .await - .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::FetchResourceError as i32, 60) - .await; - return Err(Error::Custom(format!( - "Fetch resource info failed with error: {}", - resp.msg - ))); - } - } - - if let Some((watched_task_uuid, watched_task_state)) = - task.exec_options.as_ref().and_then(|o| o.watch) - { - // Watch other tasks to specified state to trigger this task - if task_executor.task_redis_conn.is_some() && task_executor.task_redis_pubsub.is_some() { - task_executor - .announce_task_state(&task.uuid, TaskExecState::Watch as i32) - .await; - let tmp_cancel_token = task_executor.task_cancel_token.clone(); - tokio::select! { - biased; - _ = tmp_cancel_token.cancelled() => { - tracing::info!("Task watching interrupted by shutdown signal"); - task_executor.unsubscribe_task_exec_state(&watched_task_uuid).await; - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::WorkerExited as i32, 60) - .await; - return Ok(()); - }, - _ = task_executor.watch_task(&watched_task_uuid, watched_task_state) => {}, - _ = tokio::time::sleep_until(timeout_until) => { - tracing::debug!("Watching timeout, commit this task as canceled"); - task_executor.unsubscribe_task_exec_state(&watched_task_uuid).await; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - report_task(task_executor, req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { - exit_status: 0, - msg: Some(TaskResultMessage::WatchTimeout), - }), - }; - report_task(task_executor, req).await?; - task_executor - .announce_task_state_ex( - &task.uuid, - TaskExecState::WatchTimeout as i32, - 60, - ) - .await; - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::TaskCommitted as i32, 60) - .await; - return Ok(()); - } - } - } - } - - // TODO: Support unlimited exec time when exec_timeout is none - let exec_timeout = task - .spec - .timeout - .and_then(|t| u64::try_from(t).ok()) - .unwrap_or(600); - let exec_timeout = std::time::Duration::from_secs(exec_timeout); - task_executor - .announce_task_state_ex( - &task.uuid, - TaskExecState::ExecPending as i32, - exec_timeout.as_secs() + 60, - ) - .await; - let timeout_until = tokio::time::Instant::now() + exec_timeout; - - // Setup new task file path and clean up any stale file - let new_task_path = task_executor.task_cache_path.join("new_task.json"); - let _ = tokio::fs::remove_file(&new_task_path).await; // Ignore errors if file doesn't exist - - let mut command = Command::new("/usr/bin/env"); - command - .args(task.spec.args) - .envs(task.spec.envs) - .env( - "MITO_RESULT_DIR", - task_executor.task_cache_path.join("result"), - ) - .env("MITO_EXEC_DIR", task_executor.task_cache_path.join("exec")) - .env( - "MITO_RESOURCE_DIR", - task_executor.task_cache_path.join("resource"), - ) - .env("MITO_TASK_UUID", task.uuid.to_string()) - .env("MITO_NEW_TASK", &new_task_path) - .stdin(std::process::Stdio::null()); - if let Some(uuid) = task.upstream_task_uuid { - command.env("MITO_UPSTREAM_TASK_UUID", uuid.to_string()); - } - if task.spec.terminal_output { - command - .stdout(std::process::Stdio::piped()) - .stderr(std::process::Stdio::piped()); - } else { - command - .stdout(std::process::Stdio::null()) - .stderr(std::process::Stdio::null()); - } - - let mut child = command.spawn().inspect_err(|e| { - tracing::error!("Failed to spawn task: {}", e); - })?; - task_executor - .announce_task_state(&task.uuid, TaskExecState::ExecSpawned as i32) - .await; - let process_output = async { - if task.spec.terminal_output { - let process_output = async { - let mut stdout_buf = Vec::new(); - let mut stdout = child.stdout.take().unwrap(); - let mut stderr_buf = Vec::new(); - let mut stderr = child.stderr.take().unwrap(); - tokio::try_join!( - stdout.read_to_end(&mut stdout_buf), - stderr.read_to_end(&mut stderr_buf), - child.wait() - ) - .map(|(_, _, exit_status)| ProcessOutput::WithLog { - stdout: stdout_buf, - stderr: stderr_buf, - exit_status, - }) - }; - process_output.await - } else { - child - .wait() - .await - .map(|exit_status| ProcessOutput::WithoutLog { exit_status }) - } - }; - - let output = tokio::select! { - biased; - _ = task_executor.task_cancel_token.cancelled() => { - tracing::info!("Task execution interrupted by shutdown signal"); - child.kill().await.inspect_err(|e| { - tracing::error!("Failed to kill task: {}", e); - })?; - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::WorkerExited as i32, 60) - .await; - return Ok(()); - }, - output = process_output => { - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::ExecFinished as i32, 660) - .await; - output.map(TaskResult::Finish) - }, - _ = tokio::time::sleep_until(timeout_until) => { - tracing::debug!("Task execution timeout"); - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::ExecTimeout as i32, 60) - .await; - if let Some(id) = child.id() { - // TODO: we may change this when once the `linux_pidfd` is stabilized in standard library - // Tracking issue for std lib: [rust-lang/rust #82971](https://github.com/rust-lang/rust/issues/82971) - // Tracking issue for tokio: [tokio-rs/tokio #6281](https://github.com/tokio-rs/tokio/issues/6281) - let _ = signal::kill(Pid::from_raw(id as i32), Signal::SIGTERM).inspect_err(|e| { - tracing::error!("Failed to send SIGTERM to task: {}", e); - }); - } - tokio::select! { - biased; - _ = child.wait() => {}, - _ = task_executor.task_cancel_token.cancelled() => { - child.kill().await.inspect_err(|e| { - tracing::error!("Failed to kill task: {}", e); - })?; - return Ok(()); - }, - _ = tokio::time::sleep(std::time::Duration::from_secs(10)) => { - child.kill().await.inspect_err(|e| { - tracing::error!("Failed to kill task: {}", e); - })?; - }, - } - if task.spec.terminal_output { - let output = child.wait_with_output().await?; - Ok(TaskResult::Timeout(ProcessOutput::WithLog { - stdout: output.stdout, - stderr: output.stderr, - exit_status: output.status, - })) - } else { - let exit_status = child.wait().await?; - Ok(TaskResult::Timeout(ProcessOutput::WithoutLog { - exit_status, - })) - } - }, - }?; - tracing::debug!("Task execution finished"); - task_executor - .announce_task_state_ex(&task.uuid, TaskExecState::UploadResult as i32, 660) - .await; - process_task_result(task.id, task.uuid, task_executor, output).await?; - Ok(()) +/// The worker's own task source. Kept out of the execution seam on purpose: +/// pulling work is the one coordinator interaction a worker does *not* share +/// with the agent, which is handed its tasks a suite at a time. +struct TaskFetcher { + http_client: Client, + credential: String, + fetch_url: Url, + cancel_token: CancellationToken, + coordinator_force_exit: Arc, + polling_interval: std::time::Duration, } -async fn process_task_result( - id: i64, - uuid: Uuid, - task_executor: &mut TaskExecutor, - output: TaskResult, -) -> crate::error::Result<()> { - let (is_finished, exit_status) = output.state(); - let req = ReportTaskReq { - id, - op: if is_finished { - task_executor - .announce_task_state_ex(&uuid, TaskExecState::UploadFinishedResult as i32, 660) - .await; - ReportTaskOp::Finish - } else { - task_executor - .announce_task_state_ex(&uuid, TaskExecState::UploadCancelledResult as i32, 660) - .await; - ReportTaskOp::Cancel - }, - }; - report_task(task_executor, req).await?; - // Compress possible output and upload - let (tx, mut rx) = mpsc::channel::<(ArtifactContentType, u64)>(3); - // Spawn a task to archive the output - let timeout_cancel_token = CancellationToken::new(); - let archive_timeout_cancel_token = timeout_cancel_token.clone(); - let archive_cancel_token = task_executor.task_cancel_token.clone(); - let archive_cache_path = task_executor.task_cache_path.clone(); - let archive_hd = tokio::spawn(async move { - let result_dir = archive_cache_path.join("result"); - if !result_dir - .read_dir() - .map(|mut dir| dir.next().is_none()) - .unwrap_or(true) - { - let tar_file = - tokio::fs::File::create(archive_cache_path.join("result.tar.gz")).await?; - let encoder = GzipEncoder::new(tar_file); - let mut ar = Builder::new(encoder); - let compress_task = async { - ar.append_dir_all("result", result_dir).await?; - std::io::Result::Ok(()) - }; - tokio::select! { - biased; - _ = archive_cancel_token.cancelled() => { - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - tokio::fs::remove_file(archive_cache_path.join("result.tar.gz")).await?; - tracing::info!("Task output generation interrupted by shutdown signal"); - return Ok(()); - } - _ = archive_timeout_cancel_token.cancelled() => { - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - tokio::fs::remove_file(archive_cache_path.join("result.tar.gz")).await?; - tracing::warn!("Task output generation timeout"); - return Ok(()); - } - res = compress_task => { - match res { - Ok(_) => { - ar.finish().await?; - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - let file = encoder.into_inner(); - let size = file.metadata().await?.len(); - if let Err(e) = tx.send((ArtifactContentType::Result, size)).await { - tracing::error!("Failed to send result size: {}", e); - archive_cancel_token.cancel(); - return Ok(()); - } - } - Err(e) => { - tracing::error!("Failed to compress result: {}", e); - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - archive_cancel_token.cancel(); - return Err(e); - } - } - - } - } - } - let exec_log_dir = archive_cache_path.join("exec"); - if !exec_log_dir - .read_dir() - .map(|mut dir| dir.next().is_none()) - .unwrap_or(true) - { - let file_name = ArtifactContentType::ExecLog.to_string(); - let tar_file = tokio::fs::File::create(archive_cache_path.join(&file_name)).await?; - let encoder = GzipEncoder::new(tar_file); - let mut ar = Builder::new(encoder); - let compress_task = async { - ar.append_dir_all("exec-log", exec_log_dir).await?; - std::io::Result::Ok(()) - }; - tokio::select! { - biased; - _ = archive_cancel_token.cancelled() => { - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - tokio::fs::remove_file(archive_cache_path.join(&file_name)).await?; - tracing::info!("Task output generation interrupted by shutdown signal"); - return Ok(()); - } - _ = archive_timeout_cancel_token.cancelled() => { - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - tokio::fs::remove_file(archive_cache_path.join(&file_name)).await?; - tracing::warn!("Task output generation timeout"); - return Ok(()); - } - res = compress_task => { - match res { - Ok(_) => { - ar.finish().await?; - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - let file = encoder.into_inner(); - let size = file.metadata().await?.len(); - if let Err(e) = tx.send((ArtifactContentType::ExecLog, size)).await { - tracing::error!("Failed to compress exec log: {}", e); - archive_cancel_token.cancel(); - return Ok(()); - } - } - Err(e) => { - tracing::error!("Failed to compress exec log: {}", e); - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - archive_cancel_token.cancel(); - return Err(e); - } - } - - } - } - } - if let ProcessOutput::WithLog { stdout, stderr, .. } = output.get_output() { - let tar_file = - tokio::fs::File::create(archive_cache_path.join("std-log.tar.gz")).await?; - let encoder = GzipEncoder::new(tar_file); - let mut ar = Builder::new(encoder); - let compress_task = async { - let mut header = Header::new_gnu(); - header.set_cksum(); - header.set_mode(436); - header.set_size(stdout.len() as u64); - ar.append_data(&mut header, "std-log/stdout.log", &*stdout) - .await?; - header.set_size(stderr.len() as u64); - ar.append_data(&mut header, "std-log/stderr.log", &*stderr) - .await?; - std::io::Result::Ok(()) - }; - tokio::select! { - biased; - _ = archive_cancel_token.cancelled() => { - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - tokio::fs::remove_file(archive_cache_path.join("std-log.tar.gz")).await?; - tracing::info!("Task output generation interrupted by shutdown signal"); - return Ok(()); - } - _ = archive_timeout_cancel_token.cancelled() => { - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - tokio::fs::remove_file(archive_cache_path.join("std-log.tar.gz")).await?; - tracing::warn!("Task output generation timeout"); - return Ok(()); - } - res = compress_task => { - match res { - Ok(_) => { - ar.finish().await?; - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - let file = encoder.into_inner(); - let size = file.metadata().await?.len(); - if let Err(e) = tx.send((ArtifactContentType::StdLog, size)).await { - tracing::error!("Failed to compress std log: {}", e); - archive_cancel_token.cancel(); - return Ok(()); - } - } - Err(e) => { - tracing::error!("Failed to compress std log: {}", e); - let mut encoder = ar.into_inner().await?; - encoder.shutdown().await?; - archive_cancel_token.cancel(); - return Err(e); - } - } - - } - } - } - Ok(()) - }); - let upload_artifact_fut = async { - while let Some((content_type, content_length)) = rx.recv().await { - let req = ReportTaskReq { - id, - op: ReportTaskOp::Upload { - content_type, - content_length, - }, - }; - task_executor.task_url.set_path("workers/tasks"); - let resp = loop { - let resp = task_executor - .task_client - .post(task_executor.task_url.as_str()) - .json(&req) - .bearer_auth(&task_executor.task_credential) - .send() - .await; - match resp { - Ok(resp) => { - if resp.status().is_success() { - let resp = resp - .json::() - .await - .map_err(RequestError::from)?; - break resp; - } else if resp.status() == StatusCode::UNAUTHORIZED { - tracing::info!("Request upload url failed with coordinator force exit"); - task_executor - .coordinator_force_exit - .store(true, std::sync::atomic::Ordering::Release); - task_executor.task_cancel_token.cancel(); - return Ok(()); - } else if resp.status() == StatusCode::NOT_FOUND { - tracing::debug!("Task not found, ignore and go on for next cycle"); - return Ok(()); - } else if resp.status() == StatusCode::FORBIDDEN { - let resp: ErrorMsg = resp - .json() - .await - .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); - tracing::info!( - "Request upload url failed with permission denied: {}", - resp.msg - ); - return Ok(()); - } else { - let resp: ErrorMsg = resp - .json() - .await - .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); - task_executor.task_cancel_token.cancel(); - return Err(Error::Custom(format!( - "Request upload url failed with error: {}", - resp.msg - ))); - } - } - Err(e) => { - if e.is_connect() && e.is_request() { - tracing::error!( - "Request upload url failed with connection error: {}. Retry after {:?}", - e, - task_executor.polling_interval - ); - tokio::select! { - biased; - _ = task_executor.task_cancel_token.cancelled() => return Ok(()), - _ = tokio::time::sleep(task_executor.polling_interval) => {}, - } - continue; - } else { - task_executor.task_cancel_token.cancel(); - return Err(RequestError::from(e).into()); - } - } - } - }; - if let Some(url) = resp.url { - loop { - let file = tokio::fs::File::open( - task_executor.task_cache_path.join(content_type.to_string()), - ) - .await?; - let upload_file = task_executor - .task_client - .put(url.as_str()) - .header(CONTENT_LENGTH, content_length) - .body(file) - .send(); - let resp = tokio::select! { - biased; - _ = task_executor.task_cancel_token.cancelled() => { - tracing::info!("Upload failed with shutdown signal"); - return Ok(()); - } - _ = timeout_cancel_token.cancelled() => { - tracing::warn!("Upload failed with timeout"); - return Ok(()); - } - resp = upload_file => resp - }; - match resp { - Ok(resp) => { - if resp.status().is_success() { - break; - } else { - let status = resp.status(); - return Err(Error::Custom(format!( - "Upload failed with status code: {status}" - ))); - } - } - Err(e) => { - if e.is_connect() && e.is_request() { - tracing::error!( - "Upload failed with connection error: {}. Retry after {:?}", - e, - task_executor.polling_interval - ); - tokio::select! { - biased; - _ = task_executor.task_cancel_token.cancelled() => return Ok(()), - _ = timeout_cancel_token.cancelled() => { - tracing::warn!("Upload failed with timeout"); - return Ok(()); - } - _ = tokio::time::sleep(task_executor.polling_interval) => {}, - } - continue; - } else { - return Err(RequestError::from(e).into()); - } - } - } - } - } - } - crate::error::Result::Ok(()) - }; - let timeout_until = tokio::time::Instant::now() + std::time::Duration::from_secs(600); - tokio::select! { - biased; - _ = tokio::time::sleep_until(timeout_until) => { - tracing::warn!("Upload result timeout"); - timeout_cancel_token.cancel(); - // Commit the task result - let req = ReportTaskReq { - id, - op: ReportTaskOp::Commit(TaskResultSpec { - exit_status: exit_status.into_raw(), - msg: Some(TaskResultMessage::UploadResultTimeout), - }), - }; - report_task(task_executor, req).await?; - task_executor - .announce_task_state_ex(&uuid, TaskExecState::UploadResultTimeout as i32, 60) - .await; - task_executor - .announce_task_state_ex(&uuid, TaskExecState::TaskCommitted as i32, 60) - .await; - archive_hd.await??; +impl TaskFetcher { + async fn next_task(&mut self) -> FetchOutcome { + if self.cancel_token.is_cancelled() { + return FetchOutcome::Stop; } - res = upload_artifact_fut => { - res?; - archive_hd.await??; - if task_executor.task_cancel_token.is_cancelled() { - tracing::info!("Task execution interrupted by shutdown signal"); - task_executor - .announce_task_state_ex(&uuid, TaskExecState::WorkerExited as i32, 60) - .await; - return Ok(()); - } - task_executor - .announce_task_state_ex(&uuid, TaskExecState::UploadResultFinished as i32, 60) - .await; - - // Check for new task file and submit if present - let new_task_scceed = submit_new_task_if_present(id, task_executor).await; - let msg = if is_finished { - if new_task_scceed { - None - } else { - Some(TaskResultMessage::SubmitNewTaskFailed) - } - } else { - Some(TaskResultMessage::ExecTimeout) - }; - // Commit the task result - let req = ReportTaskReq { - id, - op: ReportTaskOp::Commit(TaskResultSpec { - exit_status: exit_status.into_raw(), - msg - }), - }; - report_task(task_executor, req).await?; - task_executor - .announce_task_state_ex(&uuid, TaskExecState::TaskCommitted as i32, 60) - .await; - - } - } - - // clean the directory after all the artifacts uploaded and the task committed - tokio::fs::remove_dir_all(&task_executor.task_cache_path).await?; - tokio::fs::create_dir_all(&task_executor.task_cache_path).await?; - tokio::fs::create_dir_all(&task_executor.task_cache_path.join("result")).await?; - tokio::fs::create_dir_all(&task_executor.task_cache_path.join("exec")).await?; - Ok(()) -} - -async fn report_task( - task_executor: &mut TaskExecutor, - req: ReportTaskReq, -) -> crate::error::Result<()> { - task_executor.task_url.set_path("workers/tasks"); - loop { - let resp = task_executor - .task_client - .post(task_executor.task_url.as_str()) - .json(&req) - .bearer_auth(&task_executor.task_credential) + let resp = self + .http_client + .get(self.fetch_url.as_str()) + .bearer_auth(&self.credential) .send() .await; match resp { Ok(resp) => { if resp.status().is_success() { - break; + match resp.json::>().await { + Ok(Some(task)) => FetchOutcome::Task(Box::new(task)), + Ok(None) => { + tracing::debug!( + "No task fetched. Next fetch after {:?}", + self.polling_interval + ); + self.backoff().await + } + Err(e) => { + tracing::error!("Failed to parse task specification: {}", e); + self.cancel_token.cancel(); + FetchOutcome::Stop + } + } } else if resp.status() == StatusCode::UNAUTHORIZED { - tracing::info!("Report task failed with coordinator force exit"); - task_executor - .coordinator_force_exit + tracing::info!("Task fetch failed with coordinator force exit"); + self.coordinator_force_exit .store(true, std::sync::atomic::Ordering::Release); - task_executor.task_cancel_token.cancel(); - return Ok(()); - } else if resp.status() == StatusCode::NOT_FOUND { - tracing::debug!("Task not found, ignore and go on for next cycle"); - return Ok(()); + self.cancel_token.cancel(); + FetchOutcome::Stop } else { let resp: ErrorMsg = resp .json() .await .unwrap_or_else(|e| ErrorMsg { msg: e.to_string() }); - return Err(Error::Custom(format!( - "Report task failed with error: {}", - resp.msg - ))); + tracing::error!("Task fetch failed with error: {}", resp.msg); + self.cancel_token.cancel(); + FetchOutcome::Stop } } Err(e) => { if e.is_connect() && e.is_request() { tracing::error!( - "Report task failed with connection error: {}. Retry after {:?}", + "Fetching task failed with connection error: {}. Retry after {:?}", e, - task_executor.polling_interval + self.polling_interval ); - tokio::select! { - biased; - _ = task_executor.task_cancel_token.cancelled() => break, - _ = tokio::time::sleep(task_executor.polling_interval) => {}, - } - continue; + self.backoff().await } else { - return Err(RequestError::from(e).into()); + tracing::error!("Fetching task failed with error: {}", e); + self.cancel_token.cancel(); + FetchOutcome::Stop } } } } - Ok(()) -} - -async fn submit_new_task_if_present(task_id: i64, task_executor: &mut TaskExecutor) -> bool { - let new_task_path = task_executor.task_cache_path.join("new_task.json"); - // Check if the new task file exists - if !new_task_path.exists() { - return true; // No new task to submit - } - // Read and parse the file - let new_task_content = match tokio::fs::read_to_string(&new_task_path).await { - Ok(content) if !content.trim().is_empty() => content, - Ok(_) => { - tracing::debug!("New task file exists but is empty, ignoring"); - let _ = tokio::fs::remove_file(&new_task_path).await; - return true; - } - Err(e) => { - tracing::warn!("Failed to read new task file: {}", e); - let _ = tokio::fs::remove_file(&new_task_path).await; - return false; - } - }; - - // Parse JSON to SubmitTaskReq - let submit_req: crate::schema::SubmitTaskReq = match serde_json::from_str(&new_task_content) { - Ok(req) => req, - Err(e) => { - tracing::warn!("Failed to parse new task JSON: {}", e); - let _ = tokio::fs::remove_file(&new_task_path).await; - return false; + async fn backoff(&self) -> FetchOutcome { + tokio::select! { + biased; + _ = self.cancel_token.cancelled() => FetchOutcome::Stop, + _ = tokio::time::sleep(self.polling_interval) => FetchOutcome::Idle, } - }; - - tracing::debug!( - "Submitting new task from completed task {} to group '{}'", - task_id, - submit_req.group_name - ); - let req = ReportTaskReq { - id: task_id, - op: ReportTaskOp::Submit(Box::new(submit_req)), - }; - if let Err(e) = report_task(task_executor, req).await { - tracing::warn!("Failed to submit new task: {}", e); - return false; } - - // Always clean up the file after processing (success or failure) - let _ = tokio::fs::remove_file(&new_task_path).await; - true } diff --git a/netmito/src/ws/connection.rs b/netmito/src/ws/connection.rs new file mode 100644 index 00000000..5822dec1 --- /dev/null +++ b/netmito/src/ws/connection.rs @@ -0,0 +1,285 @@ +//! Per-agent notification sessions and the router actor that owns them. +//! +//! Each agent gets an [`AgentSession`]: a monotonic counter, a replay buffer of +//! notifications it has not acknowledged, and (while connected) the sender half +//! of its WebSocket. Notifications are pushed over the socket when one is +//! attached and always land in the buffer; the buffer is what a heartbeat +//! catch-up reads and what a reconnect replays. +//! +//! **Buffer eviction is acknowledgement-driven.** An event leaves the buffer +//! only once the agent confirms it processed it — over the socket +//! (`AgentWsMessage::Ack`) or implicitly via the `last_notification_id` on its +//! next heartbeat. A heartbeat read is a *peek*, so a response lost in flight +//! is redelivered instead of dropped. + +use std::collections::{HashMap, VecDeque}; + +use axum::extract::ws::Message; +use speedy::Writable; +use tokio::sync::{mpsc, oneshot}; +use tokio_util::sync::CancellationToken; +use uuid::Uuid; + +use crate::channel::{MRx, MTx}; +use crate::schema::{AgentNotification, WsNotificationEvent}; + +/// Hard cap on unacknowledged notifications kept per agent. An agent that never +/// acknowledges (permanently dead) must not grow the coordinator's memory; past +/// this point the oldest events are dropped. Notifications are hints — the +/// agent re-derives the real state from the HTTP endpoints — so a dropped one +/// costs at most a delay until the next `SuiteAvailable`. +const MAX_BUFFERED_NOTIFICATIONS: usize = 256; + +/// Command for the [`AgentWsRouter`] actor. +#[derive(Debug)] +pub enum RouterOp { + /// A WebSocket connected; attach its sender to the agent's session. + Register { + uuid: Uuid, + sender: mpsc::Sender, + }, + /// The WebSocket dropped; the session and its buffer survive. + Unregister { uuid: Uuid }, + /// Forget the agent entirely (buffer included). + RemoveAgent { uuid: Uuid }, + /// The agent processed everything up to and including `id`. + AckBy { uuid: Uuid, id: u64 }, + /// Queue a notification (and push it over the socket if connected). + Notify { + uuid: Uuid, + event: AgentNotification, + }, + /// Heartbeat catch-up: acknowledge through `after_id`, then return + /// everything still buffered beyond it **without** dropping it. + PendingNotifications { + uuid: Uuid, + after_id: u64, + tx: oneshot::Sender>, + }, + /// Read the agent's current sequence counter (`None` if unknown). + GetCounter { + uuid: Uuid, + tx: oneshot::Sender>, + }, +} + +/// One agent's notification state. +struct AgentSession { + /// Sender of the live WebSocket, if any. + sender: Option>, + /// Last allocated sequence ID. + counter: u64, + /// Notifications not yet 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()); + while self.buffer.len() > MAX_BUFFERED_NOTIFICATIONS { + if let Some(dropped) = self.buffer.pop_front() { + tracing::warn!( + notification_id = dropped.id, + "Agent notification buffer full; dropping the oldest unacknowledged event" + ); + } + } + event + } + + /// Drop everything the agent has confirmed it processed. + 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() + } +} + +/// Actor owning every agent's notification session. +pub struct AgentWsRouter { + sessions: HashMap, + cancel_token: CancellationToken, + rx: MRx, +} + +impl AgentWsRouter { + pub fn new(cancel_token: CancellationToken, rx: MRx) -> 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 { + 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); + // Replay whatever is still unacknowledged so a reconnecting + // agent catches up without waiting for a heartbeat. + let pending = session.pending_after(0); + let sender = session.sender.clone(); + if let Some(sender) = sender { + for event in &pending { + if !Self::try_push(&sender, event, uuid) { + break; + } + } + } + tracing::debug!( + agent_uuid = %uuid, + replayed = pending.len(), + connected_agents = 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::RemoveAgent { uuid } => { + self.sessions.remove(&uuid); + tracing::debug!(agent_uuid = %uuid, "Agent removed from WebSocket router"); + } + 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.clone() { + Self::try_push(&sender, &event, uuid); + } + } + RouterOp::PendingNotifications { uuid, after_id, tx } => { + let pending = match self.sessions.get_mut(&uuid) { + Some(session) => { + session.ack(after_id); + session.pending_after(after_id) + } + None => Vec::new(), + }; + let _ = tx.send(pending); + } + RouterOp::GetCounter { uuid, tx } => { + let _ = tx.send(self.sessions.get(&uuid).map(|s| s.counter)); + } + } + } + + /// Non-blocking push over the socket. The actor must never await on a slow + /// consumer, so a full channel simply skips the push — the event stays + /// buffered and reaches the agent on its next heartbeat. + fn try_push(sender: &mpsc::Sender, 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.try_send(Message::Binary(payload.into())) { + Ok(()) => true, + Err(e) => { + tracing::debug!( + agent_uuid = %uuid, + notification_id = event.id, + "Notification not pushed over WebSocket ({e}); left buffered for heartbeat" + ); + false + } + } + } + + // ── convenience senders, used from the service layer ── + + pub fn register(tx: &MTx, uuid: Uuid, sender: mpsc::Sender) { + let _ = tx.send(RouterOp::Register { uuid, sender }); + } + + pub fn unregister(tx: &MTx, uuid: Uuid) { + let _ = tx.send(RouterOp::Unregister { uuid }); + } + + pub fn ack(tx: &MTx, uuid: Uuid, id: u64) { + let _ = tx.send(RouterOp::AckBy { uuid, id }); + } + + /// Queue a notification for an agent. Best effort: a closed router (only + /// possible during shutdown) is logged and ignored, never propagated into a + /// 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 through `after_id` and read back what is still pending. + pub async fn pending_notifications( + tx: &MTx, + uuid: Uuid, + after_id: u64, + ) -> Vec { + let (resp_tx, resp_rx) = oneshot::channel(); + if tx + .send(RouterOp::PendingNotifications { + uuid, + after_id, + tx: resp_tx, + }) + .is_err() + { + return Vec::new(); + } + resp_rx.await.unwrap_or_default() + } + + pub async fn counter(tx: &MTx, uuid: Uuid) -> Option { + let (resp_tx, resp_rx) = oneshot::channel(); + if tx.send(RouterOp::GetCounter { uuid, tx: resp_tx }).is_err() { + return None; + } + resp_rx.await.ok().flatten() + } +} diff --git a/netmito/src/ws/handler.rs b/netmito/src/ws/handler.rs new file mode 100644 index 00000000..063926fb --- /dev/null +++ b/netmito/src/ws/handler.rs @@ -0,0 +1,121 @@ +//! `GET /ws/agents` — the agent's notification socket. +//! +//! Authentication is the agent JWT, checked by `agent_auth_middleware` before +//! the upgrade. Once upgraded, the task bridges two directions: router → socket +//! (notifications) and socket → router (acknowledgements). + +use axum::{ + extract::{ + ws::{Message, WebSocket}, + State, WebSocketUpgrade, + }, + response::IntoResponse, + Extension, +}; +use futures::{SinkExt, StreamExt}; +use speedy::Readable; +use tokio::sync::mpsc; +use uuid::Uuid; + +use crate::{ + config::InfraPool, schema::AgentWsMessage, service::auth::AuthAgent, + ws::connection::AgentWsRouter, +}; + +/// How many notifications may sit in a single connection's outbound queue +/// before the router starts skipping pushes (the events stay buffered and are +/// delivered on the next heartbeat instead). +const OUTBOUND_QUEUE_LEN: usize = 128; + +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, mut rx) = mpsc::channel::(OUTBOUND_QUEUE_LEN); + + AgentWsRouter::register(&pool.ws_router_tx, agent_uuid, tx); + + // Keepalive so idle connections survive intermediate proxies. + let mut keepalive = tokio::time::interval(pool.ws_keepalive_interval); + keepalive.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + keepalive.tick().await; // the first tick fires immediately + + 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 { + Some(msg) => { + if let Err(e) = sender.send(msg).await { + tracing::debug!(agent_uuid = %agent_uuid, "Failed to push to agent: {e}"); + break; + } + } + None => 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"); +} + +/// Returns false when the connection should be torn down. +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, "Unparseable agent message: {e}"); + } + }, + Message::Close(frame) => { + tracing::debug!(agent_uuid = %agent_uuid, ?frame, "Agent closed the WebSocket"); + return false; + } + // axum answers pings itself; frames are `speedy` binary, never text. + 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..a0ba2d61 --- /dev/null +++ b/netmito/src/ws/mod.rs @@ -0,0 +1,16 @@ +//! Coordinator → agent push notifications over WebSocket. +//! +//! The WS layer is a pure latency optimization: every notification is also +//! buffered per agent and handed back on the next heartbeat, so an agent that +//! never connects (or whose connection drops) still makes progress. See +//! [`connection::AgentWsRouter`] for the buffer/ack model. +//! +//! Frames are `speedy`-encoded binary in both directions, matching dev. The +//! same notification types also travel as JSON in the heartbeat response, so +//! they carry both derives. + +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 8f218fb2..89665138 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -1754,10 +1754,6 @@ paths: responses: "200": description: Suite cancelled successfully - content: - application/json: - schema: - $ref: "#/components/schemas/CancelSuiteResp" "400": $ref: "#/components/responses/BadRequest" "401": @@ -1795,11 +1791,17 @@ paths: "404": $ref: "#/components/responses/NotFound" - /suites/{uuid}/agents/include: + /suites/{uuid}/agents/override: post: - summary: Include an agent in a suite - description: Manually pin an agent to the suite even if it does not tag-match. Upserts any existing opposite override. - operationId: includeSuiteAgent + summary: Set agent overrides on a suite + description: >- + Batch-apply manual agent overrides for one suite. An agent may be pinned + (`include`), blocked (`exclude`), or returned to the tag-match default + (`clear`). The effective agent set for a suite is + `(tag-matched ∪ included) − excluded`, restricted to agents the suite's + group has write access to; only the manual half is persisted. Per-agent + failures are reported in the response rather than aborting the batch. + operationId: overrideSuiteAgents tags: - Suites security: @@ -1816,28 +1818,29 @@ paths: content: application/json: schema: - $ref: "#/components/schemas/SuiteAgentReq" + $ref: "#/components/schemas/SuiteAgentOverrideReq" responses: "200": - description: Override applied + description: Batch applied; the body lists any entries that failed content: application/json: schema: - $ref: "#/components/schemas/SuiteAgentResp" + $ref: "#/components/schemas/SuiteAgentOverrideResp" "400": $ref: "#/components/responses/BadRequest" "401": $ref: "#/components/responses/Unauthorized" - "403": - $ref: "#/components/responses/Forbidden" "404": $ref: "#/components/responses/NotFound" - /suites/{uuid}/agents/exclude: + /suites/{uuid}/jobs/query: post: - summary: Exclude an agent from a suite - description: Manually block an agent from the suite even if it tag-matches. Upserts any existing opposite override. - operationId: excludeSuiteAgent + summary: Query a suite's jobs + description: >- + A job is one agent's attempt at running the suite. Read-only: jobs are + driven entirely by agents, and a user-initiated stop goes through the + suite's close/cancel instead. Requires the Read role in the suite's group. + operationId: querySuiteJobs tags: - Suites security: @@ -1854,28 +1857,29 @@ paths: content: application/json: schema: - $ref: "#/components/schemas/SuiteAgentReq" + $ref: "#/components/schemas/SuiteJobsQueryReq" responses: "200": - description: Override applied + description: The suite's jobs, newest first content: application/json: schema: - $ref: "#/components/schemas/SuiteAgentResp" + $ref: "#/components/schemas/SuiteJobsQueryResp" "400": $ref: "#/components/responses/BadRequest" "401": $ref: "#/components/responses/Unauthorized" - "403": - $ref: "#/components/responses/Forbidden" "404": $ref: "#/components/responses/NotFound" - /suites/{uuid}/agents/reset: - post: - summary: Reset an agent to the tag-match default - description: Clear a manual include/exclude override for an agent, reverting it to the tag-match default. Idempotent. - operationId: resetSuiteAgent + /suites/{uuid}/jobs/{job_id}: + get: + summary: Get one job of a suite + description: >- + One job plus its hook executions. A hook's artifacts live in the shared + artifact store keyed by the hook's own uuid, so they are downloaded + through the ordinary artifact endpoints. + operationId: getSuiteJob tags: - Suites security: @@ -1887,19 +1891,58 @@ paths: schema: type: string format: uuid + - name: job_id + in: path + required: true + description: The suite-scoped job number, ascending from 1 + schema: + type: integer + format: int32 + responses: + "200": + description: Job details with embedded hook executions + content: + application/json: + schema: + $ref: "#/components/schemas/SuiteJobQueryResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /agents: + post: + summary: Register an agent + description: >- + Register the agent for a machine and mint its token. Registration is an + upsert keyed by `machine_code`: a machine has exactly one agent record + for its whole life, so a restarting agent re-adopts that record (with + refreshed tags, labels, groups, and metadata) rather than creating a + second one. The caller needs the Write or Admin role in every group + listed, and the Admin role in `admin_group` — the single group that can + shut the agent down, defaulting to the caller's personal group. Group + access is rewritten on every registration, not merged: groups this + request omits lose the access a previous registration gave them. + operationId: registerAgent + tags: + - Agents + security: + - bearerAuth: [] requestBody: required: true content: application/json: schema: - $ref: "#/components/schemas/SuiteAgentReq" + $ref: "#/components/schemas/RegisterAgentReq" responses: "200": - description: Override cleared (or nothing to reset) + description: Agent registered (or re-adopted) content: application/json: schema: - $ref: "#/components/schemas/ResetSuiteAgentResp" + $ref: "#/components/schemas/RegisterAgentResp" + "400": + $ref: "#/components/responses/BadRequest" "401": $ref: "#/components/responses/Unauthorized" "403": @@ -1907,13 +1950,16 @@ paths: "404": $ref: "#/components/responses/NotFound" - /admin/users: + /agents/query: post: - summary: Create user (Admin) - description: Create a new user account (admin only) - operationId: adminCreateUser + summary: Query agents + description: >- + List agents with optional filtering, scoped to a single group. The group + defaults to the caller's personal group when omitted, and the caller must + have at least the Read role in it. + operationId: queryAgents tags: - - Admin + - Agents security: - bearerAuth: [] requestBody: @@ -1921,141 +1967,244 @@ paths: content: application/json: schema: - $ref: "#/components/schemas/CreateUserReq" + $ref: "#/components/schemas/AgentsQueryReq" responses: "200": - description: User created successfully + description: List of agents + content: + application/json: + schema: + $ref: "#/components/schemas/AgentsQueryResp" "401": $ref: "#/components/responses/Unauthorized" - "403": - $ref: "#/components/responses/Forbidden" - "409": - $ref: "#/components/responses/Conflict" + "404": + $ref: "#/components/responses/NotFound" - /admin/users/{username}: + /agents/{uuid}: delete: - summary: Delete user (Admin) - description: Mark user as deleted (admin only) - operationId: adminDeleteUser + summary: Shut an agent down + description: >- + Stop an agent. **The agent record is never deleted** — an agent is a + durable identity and every reference to it is restricted, so this marks + it Offline instead. Requires the caller to have the Admin role in a group + that is Admin over the agent; an agent that fails that check answers 404 + exactly as an unknown uuid does. + operationId: shutdownAgent tags: - - Admin + - Agents security: - bearerAuth: [] parameters: - - name: username + - 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): an idle agent stops now; a busy one finishes + its current job first + - force: the agent's in-flight jobs are killed without cleanup and + its uncommitted tasks are reclaimed for other agents responses: "200": - description: User deleted successfully + description: Shutdown requested + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /agents/heartbeat: + post: + summary: Agent heartbeat + description: >- + Report liveness and receive every notification the agent has not + acknowledged. `last_notification_id` acknowledges through that id, so + anything beyond it is redelivered until the agent confirms it. An agent + that misses its heartbeat is marked Offline, its in-flight jobs become + Lost, and its uncommitted tasks are reclaimed. + operationId: agentHeartbeat + tags: + - Agents + security: + - agentAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/AgentHeartbeatReq" + responses: + "200": + description: Pending notifications content: application/json: schema: - $ref: "#/components/schemas/UserStateResp" + $ref: "#/components/schemas/AgentHeartbeatResp" "401": $ref: "#/components/responses/Unauthorized" - "403": - $ref: "#/components/responses/Forbidden" "404": $ref: "#/components/responses/NotFound" - /admin/users/{username}/password: + /agents/suite: post: - summary: Change user password (Admin) - description: Change password for any user (admin only) - operationId: adminChangeUserPassword + summary: Claim a suite to run + description: >- + Select a suite and claim it in one transaction: the agent moves Idle → + Provisioning, a job record is opened whose opaque handle the agent echoes + on every later job-scoped call, and the suite's spec comes back with it. + Selecting and claiming are not separable, so no other agent can take the + suite in between. `suite_uuid` is a preference — a hint that no longer + points at runnable work falls back to the best available suite rather + than answering "nothing". A refusal the coordinator decides (nothing + available, agent already busy) is reported as `accepted: false` with a + reason rather than an error status. + operationId: agentAcceptSuite tags: - - Admin + - Agents security: - - bearerAuth: [] - parameters: - - name: username - in: path - required: true - schema: - type: string + - agentAuth: [] requestBody: required: true content: application/json: schema: - $ref: "#/components/schemas/AdminChangePasswordReq" + $ref: "#/components/schemas/AcceptSuiteReq" responses: "200": - description: Password changed successfully + description: The claimed suite and its job, or a declined answer with a reason + content: + application/json: + schema: + $ref: "#/components/schemas/AcceptSuiteResp" + "401": + $ref: "#/components/responses/Unauthorized" + + /agents/job/start: + post: + summary: Report provisioning complete + description: Job Provisioning → Executing; the agent becomes Executing. + operationId: agentStartJob + tags: + - Agents + security: + - agentAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/StartJobReq" + responses: + "200": + description: Transition applied + "400": + $ref: "#/components/responses/BadRequest" "401": $ref: "#/components/responses/Unauthorized" - "403": - $ref: "#/components/responses/Forbidden" "404": $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/Conflict" - /admin/users/{username}/group-quota: + /agents/job/cleanup: post: - summary: Change user group quota (Admin) - description: Change the group quota for a user (admin only) - operationId: adminChangeUserGroupQuota + summary: Report entering cleanup + description: Job Executing → Cleanup; the agent becomes Cleaning. + operationId: agentEnterCleanup tags: - - Admin + - Agents security: - - bearerAuth: [] - parameters: - - name: username - in: path - required: true - schema: - type: string - example: johndoe + - agentAuth: [] requestBody: required: true content: application/json: schema: - $ref: "#/components/schemas/ChangeUserGroupQuota" + $ref: "#/components/schemas/EnterCleanupReq" responses: "200": - description: User group quota changed successfully + description: Transition applied + "400": + $ref: "#/components/responses/BadRequest" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/Conflict" + + /agents/job/complete: + post: + summary: Report job completion + description: >- + Write the job's terminal state and release the agent back to Idle. The + agent reports only what it did — Completed, or Failed with a reason; + Lost and Killed are the coordinator's to write. A job that is already + terminal answers 409, which tells the agent it was torn down under it. + operationId: agentCompleteJob + tags: + - Agents + security: + - agentAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/CompleteJobReq" + responses: + "200": + description: Job closed content: application/json: schema: - $ref: "#/components/schemas/UserGroupQuotaResp" + $ref: "#/components/schemas/CompleteJobResp" "401": $ref: "#/components/responses/Unauthorized" - "403": - $ref: "#/components/responses/Forbidden" "404": $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/Conflict" - /admin/users/{username}/state: + /agents/job/hook: post: - summary: Change user state (Admin) - description: Change user state (active/locked/deleted) (admin only) - operationId: adminChangeUserState + summary: Report a job hook + description: >- + Record a provision/cleanup/background hook's outcome, or presign an + upload for its artifacts. Append-only: accepted even on a terminal job, + because a cleanup hook may legitimately finish after the coordinator + already ended the job. `Result` must precede `Upload` — it is what mints + the hook uuid the artifacts are keyed by. + operationId: agentReportHook tags: - - Admin + - Agents security: - - bearerAuth: [] - parameters: - - name: username - in: path - required: true - schema: - type: string + - agentAuth: [] requestBody: required: true content: application/json: schema: - $ref: "#/components/schemas/ChangeUserStateReq" + $ref: "#/components/schemas/HookReportReq" responses: "200": - description: User state changed successfully + description: Hook recorded; `url` is present for the Upload operation content: application/json: schema: - $ref: "#/components/schemas/UserStateResp" + $ref: "#/components/schemas/HookReportResp" + "400": + $ref: "#/components/responses/BadRequest" "401": $ref: "#/components/responses/Unauthorized" "403": @@ -2063,95 +2212,476 @@ paths: "404": $ref: "#/components/responses/NotFound" - /admin/groups/{group_name}/storage-quota: + /agents/tasks/fetch: post: - summary: Change group storage quota (Admin) - description: Update storage quota for a group (admin only) - operationId: adminChangeGroupStorageQuota + summary: Claim tasks from a suite + description: >- + Atomically claim up to `max_count` of the suite's ready tasks, highest + priority first, moving them to Running under this agent. Concurrent + agents never claim the same task. An empty list means nothing is ready + right now, not that the suite is finished. + operationId: agentFetchTasks tags: - - Admin + - Agents security: - - bearerAuth: [] - parameters: - - name: group_name - in: path - required: true - schema: - type: string + - agentAuth: [] requestBody: required: true content: application/json: schema: - $ref: "#/components/schemas/ChangeGroupStorageQuotaReq" + $ref: "#/components/schemas/FetchTasksReq" responses: "200": - description: Storage quota updated successfully + description: The claimed tasks content: application/json: schema: - $ref: "#/components/schemas/GroupStorageQuotaResp" + $ref: "#/components/schemas/FetchTasksResp" "401": $ref: "#/components/responses/Unauthorized" - "403": - $ref: "#/components/responses/Forbidden" "404": $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/Conflict" - /admin/workers/{uuid}/: - delete: - summary: Shutdown worker (Admin) - description: Force shutdown a worker (admin only) - operationId: adminShutdownWorker + /agents/tasks/report: + post: + summary: Report a task result + description: >- + The worker task-report operations with the suite job handle attached: + `Finish`, `Cancel`, `Upload` (returns a presigned artifact URL), + `Commit` (archives the task and ticks the suite's counter), and `Submit` + (spawns a downstream child in the parent's group, in whichever suite the + request names — the parent's own, a sibling, or none). A terminal job + answers 409 and a task that no longer exists answers 404; both tell the + agent to stop reporting on it. + operationId: agentReportTask tags: - - Admin + - 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: | - Worker shutdown operation: - - Graceful: Wait for current task to complete - - Force: Stop immediately + - agentAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ReportAgentTaskReq" responses: "200": - description: Worker shutdown initiated + description: Reported; `url` is present for the Upload operation + content: + application/json: + schema: + $ref: "#/components/schemas/ReportTaskResp" + "400": + $ref: "#/components/responses/BadRequest" "401": $ref: "#/components/responses/Unauthorized" "403": $ref: "#/components/responses/Forbidden" "404": $ref: "#/components/responses/NotFound" + "409": + $ref: "#/components/responses/Conflict" - /admin/groups/{group_name}/attachments/{key}: - delete: - summary: Delete group attachment (Admin) - description: Delete any group attachment (admin only) - operationId: adminDeleteAttachment + /agents/tasks/{uuid}: + get: + summary: Query a task as an agent + description: A task's state and result, for resolving watch dependencies. + operationId: agentQueryTask tags: - - Admin + - Agents security: - - bearerAuth: [] + - agentAuth: [] parameters: - - name: group_name + - name: uuid in: path required: true schema: type: string - - name: key + format: uuid + responses: + "200": + description: Task details + content: + application/json: + schema: + $ref: "#/components/schemas/TaskQueryResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /agents/tasks/{uuid}/artifacts/{content_type}: + get: + summary: Download a task artifact as an agent + description: Presigned download for a task input artifact. + operationId: agentDownloadTaskArtifact + tags: + - Agents + security: + - agentAuth: [] + parameters: + - name: uuid + in: path + required: true + schema: + type: string + format: uuid + - name: content_type + in: path + required: true + schema: + $ref: "#/components/schemas/ArtifactContentType" + responses: + "200": + description: Presigned download URL and size + content: + application/json: + schema: + $ref: "#/components/schemas/RemoteResourceDownloadResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /agents/tasks/{uuid}/attachments/{key}: + get: + summary: Download an attachment as an agent + description: Presigned download for a group attachment a task needs. + operationId: agentDownloadAttachment + tags: + - Agents + security: + - agentAuth: [] + parameters: + - name: uuid + in: path + required: true + schema: + type: string + format: uuid + - name: key + in: path + required: true + schema: + type: string + responses: + "200": + description: Presigned download URL and size + content: + application/json: + schema: + $ref: "#/components/schemas/RemoteResourceDownloadResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /agents/suites/{uuid}/attachments/{key}: + get: + summary: Download a suite attachment as an agent + description: >- + Presigned download for a group attachment a **hook** needs. The hook half + of `/agents/tasks/{uuid}/attachments/{key}`: a hook's inputs are named the + same way, but it belongs to a suite rather than to any one task, so the + owning group is resolved through the suite. + operationId: agentDownloadSuiteAttachment + tags: + - Agents + security: + - agentAuth: [] + parameters: + - name: uuid + in: path + required: true + description: Suite UUID + schema: + type: string + format: uuid + - name: key + in: path + required: true + schema: + type: string + responses: + "200": + description: Presigned download URL and size + content: + application/json: + schema: + $ref: "#/components/schemas/RemoteResourceDownloadResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /ws/agents: + get: + summary: Agent notification socket + description: >- + WebSocket upgrade carrying coordinator → agent notifications as JSON + `WsNotificationEvent` frames; the agent replies with `AgentWsMessage` + frames to acknowledge them. Purely an optimization — every notification + is also buffered per agent and returned by the next heartbeat — so an + agent that never connects still makes progress. + operationId: agentWebSocket + tags: + - Agents + security: + - agentAuth: [] + responses: + "101": + description: Switching protocols + "401": + $ref: "#/components/responses/Unauthorized" + + /admin/users: + post: + summary: Create user (Admin) + description: Create a new user account (admin only) + operationId: adminCreateUser + tags: + - Admin + security: + - bearerAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/CreateUserReq" + responses: + "200": + description: User created successfully + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "409": + $ref: "#/components/responses/Conflict" + + /admin/users/{username}: + delete: + summary: Delete user (Admin) + description: Mark user as deleted (admin only) + operationId: adminDeleteUser + tags: + - Admin + security: + - bearerAuth: [] + parameters: + - name: username + in: path + required: true + schema: + type: string + responses: + "200": + description: User deleted successfully + content: + application/json: + schema: + $ref: "#/components/schemas/UserStateResp" + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "404": + $ref: "#/components/responses/NotFound" + + /admin/users/{username}/password: + post: + summary: Change user password (Admin) + description: Change password for any user (admin only) + operationId: adminChangeUserPassword + tags: + - Admin + security: + - bearerAuth: [] + parameters: + - name: username + in: path + required: true + schema: + type: string + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/AdminChangePasswordReq" + responses: + "200": + description: Password changed successfully + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "404": + $ref: "#/components/responses/NotFound" + + /admin/users/{username}/group-quota: + post: + summary: Change user group quota (Admin) + description: Change the group quota for a user (admin only) + operationId: adminChangeUserGroupQuota + tags: + - Admin + security: + - bearerAuth: [] + parameters: + - name: username + in: path + required: true + schema: + type: string + example: johndoe + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ChangeUserGroupQuota" + responses: + "200": + description: User group quota changed successfully + content: + application/json: + schema: + $ref: "#/components/schemas/UserGroupQuotaResp" + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "404": + $ref: "#/components/responses/NotFound" + + /admin/users/{username}/state: + post: + summary: Change user state (Admin) + description: Change user state (active/locked/deleted) (admin only) + operationId: adminChangeUserState + tags: + - Admin + security: + - bearerAuth: [] + parameters: + - name: username + in: path + required: true + schema: + type: string + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ChangeUserStateReq" + responses: + "200": + description: User state changed successfully + content: + application/json: + schema: + $ref: "#/components/schemas/UserStateResp" + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "404": + $ref: "#/components/responses/NotFound" + + /admin/groups/{group_name}/storage-quota: + post: + summary: Change group storage quota (Admin) + description: Update storage quota for a group (admin only) + operationId: adminChangeGroupStorageQuota + tags: + - Admin + security: + - bearerAuth: [] + parameters: + - name: group_name + in: path + required: true + schema: + type: string + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/ChangeGroupStorageQuotaReq" + responses: + "200": + description: Storage quota updated successfully + content: + application/json: + schema: + $ref: "#/components/schemas/GroupStorageQuotaResp" + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "404": + $ref: "#/components/responses/NotFound" + + /admin/workers/{uuid}/: + delete: + summary: Shutdown worker (Admin) + description: Force shutdown a worker (admin only) + operationId: adminShutdownWorker + tags: + - Admin + 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: | + Worker shutdown operation: + - Graceful: Wait for current task to complete + - Force: Stop immediately + responses: + "200": + description: Worker shutdown initiated + "401": + $ref: "#/components/responses/Unauthorized" + "403": + $ref: "#/components/responses/Forbidden" + "404": + $ref: "#/components/responses/NotFound" + + /admin/groups/{group_name}/attachments/{key}: + delete: + summary: Delete group attachment (Admin) + description: Delete any group attachment (admin only) + operationId: adminDeleteAttachment + tags: + - Admin + security: + - bearerAuth: [] + parameters: + - name: group_name + in: path + required: true + schema: + type: string + - name: key in: path required: true schema: @@ -2234,6 +2764,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: @@ -2321,10 +2856,10 @@ components: minItems: 16 maxItems: 16 description: MD5 hash of password as byte array - retain: + refresh: type: boolean - default: true - description: Whether to retain existing login state. Defaults to true. Set to false to refresh the login state and invalidate previously issued tokens. + default: false + description: Whether to refresh the login state, invalidating previously issued tokens. Defaults to false, which keeps the existing login state. UserLoginResp: type: object @@ -3214,7 +3749,9 @@ components: description: >- Optional suite to assign the task to. When set, the suite is authoritative: the task is placed in the suite's owning group and - group_name must match that group. + group_name must match that group. For a child spawned by a running + task (worker or agent `Submit`), the suite must be owned by the + parent task's group; omit it to submit a suite-less task instead. example: 550e8400-e29b-41d4-a716-446655440000 tags: type: array @@ -4324,16 +4861,6 @@ components: - Complete: no incomplete tasks remain; a new task reopens it - Cancelled: terminal; cannot accept tasks - SuiteAgentSelectionType: - type: string - enum: - - UserIncluded - - UserExcluded - description: | - Manual agent override on a suite: - - UserIncluded: pin the agent to the suite even if it does not tag-match - - UserExcluded: block the agent even if it tag-matches - CpuBindingStrategy: type: string enum: @@ -4402,104 +4929,378 @@ components: type: object description: Hook run alongside workers - CreateTaskSuiteReq: + CreateTaskSuiteReq: + type: object + required: + - group_name + - worker_schedule + properties: + name: + type: string + nullable: true + description: Optional human-readable name (non-unique; if present, must not be blank) + example: nightly-eval + description: + type: string + nullable: true + description: Optional description + group_name: + type: string + description: Group that owns the suite (caller needs Write/Admin in it) + example: research-team + tags: + type: array + items: + type: string + description: Tags for agent matching + example: + - linux + - cuda:11 + labels: + type: array + items: + type: string + description: Labels for querying/filtering + example: + - project:cauldron + priority: + type: integer + format: int32 + default: 0 + description: Suite scheduling priority (higher = more important) + worker_schedule: + $ref: "#/components/schemas/WorkerSchedulePlan" + exec_hooks: + $ref: "#/components/schemas/ExecHooks" + + CreateTaskSuiteResp: + type: object + required: + - uuid + properties: + uuid: + type: string + format: uuid + description: Unique UUID of the created suite + + TaskSuitesQueryReq: + type: object + properties: + name: + type: string + nullable: true + description: Filter by exact suite name + description: + type: string + nullable: true + description: Filter by description substring + creator_usernames: + type: array + items: + type: string + nullable: true + description: Filter by suite creators + group_name: + type: string + nullable: true + description: Filter by group name (defaults to the caller's username when omitted) + tags: + type: array + items: + type: string + nullable: true + description: Filter by suite tags + labels: + type: array + items: + type: string + nullable: true + description: Filter by suite labels + states: + type: array + items: + $ref: "#/components/schemas/TaskSuiteState" + nullable: true + description: Filter by suite states + priority: + type: string + nullable: true + description: Filter by priority (e.g. ">5", "<=10") + example: ">5" + limit: + type: integer + format: uint64 + nullable: true + description: Maximum number of results + offset: + type: integer + format: uint64 + nullable: true + description: Offset for pagination + count: + type: boolean + default: false + description: When true, return only the total match count (suites is empty) + + TaskSuiteInfo: + type: object + required: + - uuid + - group_name + - creator_username + - tags + - labels + - priority + - worker_schedule + - state + - total_tasks + - incomplete_tasks + - created_at + - updated_at + properties: + uuid: + type: string + format: uuid + name: + type: string + nullable: true + description: + type: string + nullable: true + group_name: + type: string + creator_username: + type: string + tags: + type: array + items: + type: string + labels: + type: array + items: + type: string + priority: + type: integer + format: int32 + worker_schedule: + $ref: "#/components/schemas/WorkerSchedulePlan" + exec_hooks: + $ref: "#/components/schemas/ExecHooks" + state: + $ref: "#/components/schemas/TaskSuiteState" + last_task_submitted_at: + type: string + format: date-time + nullable: true + total_tasks: + type: integer + format: int32 + incomplete_tasks: + type: integer + format: int32 + created_at: + type: string + format: date-time + updated_at: + type: string + format: date-time + completed_at: + type: string + format: date-time + nullable: true + + TaskSuitesQueryResp: + type: object + required: + - count + - suites + - group_name + properties: + count: + type: integer + format: uint64 + description: Total match count (count mode), otherwise the number of suites on this page + suites: + type: array + items: + $ref: "#/components/schemas/TaskSuiteInfo" + group_name: + type: string + + ParsedTaskSuiteInfo: type: object required: + - uuid - group_name + - creator_username + - tags + - labels + - priority - worker_schedule + - state + - total_tasks + - incomplete_tasks + - created_at + - updated_at properties: + uuid: + type: string + format: uuid name: type: string nullable: true - description: Optional human-readable name (non-unique; if present, must not be blank) - example: nightly-eval description: type: string nullable: true - description: Optional description group_name: type: string - description: Group that owns the suite (caller needs Write/Admin in it) - example: research-team + creator_username: + type: string tags: type: array items: type: string - description: Tags for agent matching - example: - - linux - - cuda:11 labels: type: array items: type: string - description: Labels for querying/filtering - example: - - project:cauldron priority: type: integer format: int32 - default: 0 - description: Suite scheduling priority (higher = more important) worker_schedule: $ref: "#/components/schemas/WorkerSchedulePlan" exec_hooks: $ref: "#/components/schemas/ExecHooks" - - CreateTaskSuiteResp: - type: object - required: - - uuid - properties: - uuid: - type: string - format: uuid - description: Unique UUID of the created suite - - TaskSuitesQueryReq: - type: object - properties: - name: + state: + $ref: "#/components/schemas/TaskSuiteState" + last_task_submitted_at: type: string + format: date-time nullable: true - description: Filter by exact suite name - description: + total_tasks: + type: integer + format: int32 + incomplete_tasks: + type: integer + format: int32 + created_at: type: string - nullable: true - description: Filter by description substring - creator_usernames: - type: array - items: - type: string - nullable: true - description: Filter by suite creators - group_name: + format: date-time + updated_at: type: string + format: date-time + completed_at: + type: string + format: date-time nullable: true - description: Filter by group name (defaults to the caller's username when omitted) - tags: - type: array - items: - type: string - nullable: true - description: Filter by suite tags - labels: + + TaskSuiteQueryResp: + type: object + required: + - info + - eligible_agents + properties: + info: + $ref: "#/components/schemas/ParsedTaskSuiteInfo" + eligible_agents: type: array items: type: string - nullable: true - description: Filter by suite labels + format: uuid + description: >- + UUIDs of every agent that may run this suite: + `(tag-matched ∪ included) − excluded`, restricted to agents the + suite's group has write access to + + SuiteAgentOverrideAction: + type: string + enum: + - include + - exclude + - clear + description: | + What to do with one agent on a suite: + - include: pin it even if it does not tag-match + - exclude: block it even if it tag-matches + - clear: drop any manual override, back to the tag-match default + + SuiteAgentOverrideReq: + type: object + required: + - overrides + properties: + overrides: + type: object + description: Agent uuid → action + additionalProperties: + $ref: "#/components/schemas/SuiteAgentOverrideAction" + + SuiteAgentOverrideResp: + type: object + required: + - errors + properties: + errors: + type: object + description: >- + Entries that could not be applied, keyed by the same agent uuid as + the request. An empty object means every entry succeeded. + additionalProperties: + type: object + properties: + msg: + type: string + + SuiteJobState: + type: string + enum: + - Provisioning + - Executing + - Cleanup + - Completed + - Failed + - Lost + - Killed + description: | + One agent's attempt at running a suite: + - Provisioning → Executing → Cleanup are the active phases + - Completed: every hook succeeded (whether or not the suite drained) + - Failed: a hook failed + - Lost: the agent's heartbeat lapsed mid-job + - Killed: force-stopped, with no cleanup + + HookType: + type: string + enum: + - Provision + - Cleanup + - Background + + HookExecState: + type: string + enum: + - Completed + - Failed + - Cancelled + description: >- + A hook is recorded only when it finishes, so there is no running state. + + SuiteJobsQueryReq: + type: object + properties: states: type: array items: - $ref: "#/components/schemas/TaskSuiteState" + $ref: "#/components/schemas/SuiteJobState" nullable: true - description: Filter by suite states - priority: + description: Filter by job states + agent_uuid: type: string + format: uuid nullable: true - description: Filter by priority (e.g. ">5", "<=10") - example: ">5" + description: >- + Only jobs run by this agent. An agent that never ran this suite + yields an empty listing rather than a 404. limit: type: integer format: uint64 @@ -4513,35 +5314,182 @@ components: count: type: boolean default: false - description: When true, return only the total match count (suites is empty) + description: When true, return only the total match count (jobs is empty) - TaskSuiteInfo: + SuiteJobInfo: type: object required: - - uuid - - group_name - - creator_username - - tags - - labels - - priority - - worker_schedule + - job_id - state - - total_tasks - - incomplete_tasks - created_at - updated_at + properties: + job_id: + type: integer + format: int32 + description: Suite-scoped job number, ascending from 1 + state: + $ref: "#/components/schemas/SuiteJobState" + agent_uuid: + type: string + format: uuid + created_at: + type: string + format: date-time + updated_at: + type: string + format: date-time + + SuiteJobsQueryResp: + type: object + required: + - count + - jobs + properties: + count: + type: integer + format: uint64 + jobs: + type: array + items: + $ref: "#/components/schemas/SuiteJobInfo" + + HookTaskInfo: + type: object + required: + - uuid + - hook_type + - state properties: uuid: type: string format: uuid - name: + description: Also the key this hook's artifacts are stored under + hook_type: + $ref: "#/components/schemas/HookType" + state: + $ref: "#/components/schemas/HookExecState" + result: + $ref: "#/components/schemas/TaskResultSpec" + started_at: type: string - nullable: true - description: + format: date-time + completed_at: type: string - nullable: true - group_name: + format: date-time + + SuiteJobQueryResp: + type: object + required: + - info + - hooks + properties: + info: + $ref: "#/components/schemas/SuiteJobInfo" + hooks: + type: array + items: + $ref: "#/components/schemas/HookTaskInfo" + + 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 a suite is matched against. An agent is eligible for a suite + when it carries every tag the suite asks for. + labels: + type: array + items: + type: string + description: Free-form labels for querying; not used for matching + admin_group: + type: string + description: >- + The group granted Admin over the agent — the role that may shut it + down. The caller must be an Admin of it. Defaults to the caller's + personal group. + groups: + type: array + items: + type: string + description: Groups that gain Write access to the agent + lifetime: + type: string + description: >- + Token lifetime as a human duration. When omitted the token never + expires. + example: 30d + machine_code: + type: string + description: >- + Stable identifier for the host. Registration is keyed on it, so a + restarting agent must present the same value to re-adopt its record. + 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 agent's JWT, used for every agent-authenticated call + notification_counter: + type: integer + format: uint64 + description: Start of the agent's notification sequence + reused: + type: boolean + description: True if an existing agent record for this machine was re-adopted + + 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: @@ -4552,60 +5500,125 @@ components: type: array items: type: string - priority: - type: integer - format: int32 - worker_schedule: - $ref: "#/components/schemas/WorkerSchedulePlan" - exec_hooks: - $ref: "#/components/schemas/ExecHooks" state: - $ref: "#/components/schemas/TaskSuiteState" - last_task_submitted_at: + $ref: "#/components/schemas/AgentState" + last_heartbeat: type: string format: date-time - nullable: true - total_tasks: - type: integer - format: int32 - incomplete_tasks: - type: integer - format: int32 + assigned_suite_uuid: + type: string + format: uuid created_at: type: string format: date-time updated_at: type: string format: date-time - completed_at: + machine_code: type: string - format: date-time - nullable: true + metadata: + type: object - TaskSuitesQueryResp: + AgentsQueryReq: type: object required: - count - - suites + properties: + group_name: + type: string + description: Defaults to the caller's personal group + 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 - description: Total match count (count mode), otherwise the number of suites on this page - suites: + agents: type: array items: - $ref: "#/components/schemas/TaskSuiteInfo" + $ref: "#/components/schemas/AgentInfo" group_name: type: string - ParsedTaskSuiteInfo: + AgentMetrics: + type: object + required: + - active_workers + - tasks_completed + - tasks_failed + properties: + active_workers: + type: integer + format: uint32 + tasks_completed: + type: integer + format: uint64 + tasks_failed: + type: integer + format: uint64 + + AgentHeartbeatReq: + type: object + required: + - state + properties: + state: + $ref: "#/components/schemas/AgentState" + assigned_suite_uuid: + type: string + format: uuid + last_notification_id: + type: integer + format: uint64 + description: >- + The highest notification id the agent has processed. Acknowledges + through this id; anything beyond it is returned again. + metrics: + $ref: "#/components/schemas/AgentMetrics" + + AgentHeartbeatResp: + type: object + required: + - notifications + properties: + notifications: + type: array + items: + $ref: "#/components/schemas/WsNotificationEvent" + + TaskSuiteSpec: type: object required: - uuid - group_name - - creator_username - tags - labels - priority @@ -4613,22 +5626,16 @@ components: - state - total_tasks - incomplete_tasks - - created_at - - updated_at properties: uuid: type: string format: uuid name: type: string - nullable: true description: type: string - nullable: true group_name: type: string - creator_username: - type: string tags: type: array items: @@ -4646,86 +5653,300 @@ components: $ref: "#/components/schemas/ExecHooks" state: $ref: "#/components/schemas/TaskSuiteState" - last_task_submitted_at: - type: string - format: date-time - nullable: true total_tasks: type: integer format: int32 incomplete_tasks: type: integer format: int32 - created_at: + + AcceptSuiteReq: + type: object + properties: + suite_uuid: type: string - format: date-time - updated_at: + format: uuid + description: >- + The suite the agent was notified about, if any. A preference, not a + demand: one that is gone, drained, cancelled or no longer this + agent's falls back to the best available suite. + + AcceptSuiteResp: + type: object + required: + - accepted + properties: + accepted: + type: boolean + suite: + allOf: + - $ref: "#/components/schemas/TaskSuiteSpec" + description: What to run. Present exactly when `accepted` is true. + job: + type: integer + format: int64 + description: Opaque job handle, echoed on every later job-scoped call + job_id: + type: integer + format: int32 + description: Suite-scoped job number, for display + reason: type: string - format: date-time - completed_at: + description: Why nothing was claimed, when `accepted` is false + + StartJobReq: + type: object + required: + - job + properties: + job: + type: integer + format: int64 + + EnterCleanupReq: + type: object + required: + - job + properties: + job: + type: integer + format: int64 + + JobFailureKind: + type: string + enum: + - ProvisionFailed + - BackgroundExited + - ExecutionError + - CleanupFailed + + JobFailureReason: + type: object + required: + - kind + - message + properties: + kind: + $ref: "#/components/schemas/JobFailureKind" + message: type: string - format: date-time - nullable: true - TaskSuiteQueryResp: + SuiteJobOutcome: + oneOf: + - type: string + enum: + - Completed + - type: object + required: + - Failed + properties: + Failed: + type: object + required: + - reason + properties: + reason: + $ref: "#/components/schemas/JobFailureReason" + description: >- + What the agent did. Lost and Killed are coordinator decisions and are + never agent outcomes. + + CompleteJobReq: type: object required: - - info - - assigned_agents + - job + - outcome properties: - info: - $ref: "#/components/schemas/ParsedTaskSuiteInfo" - assigned_agents: - type: array - items: - type: string - format: uuid - description: UUIDs of the manually-included agents (tag-matched agents are computed in-memory and not listed here) + job: + type: integer + format: int64 + outcome: + $ref: "#/components/schemas/SuiteJobOutcome" + + CompleteJobResp: + type: object + required: + - next_suite_available + properties: + next_suite_available: + type: boolean + + HookReportOp: + oneOf: + - type: object + required: + - Result + properties: + Result: + $ref: "#/components/schemas/TaskResultSpec" + - type: object + required: + - Upload + properties: + Upload: + type: object + required: + - content_type + - content_length + properties: + content_type: + $ref: "#/components/schemas/ArtifactContentType" + content_length: + type: integer + format: uint64 - CancelSuiteResp: + HookReportReq: type: object required: - - cancelled_task_count + - job + - hook_type + - op properties: - cancelled_task_count: + job: type: integer - format: uint64 - description: Number of tasks that were cancelled + format: int64 + hook_type: + $ref: "#/components/schemas/HookType" + op: + $ref: "#/components/schemas/HookReportOp" - SuiteAgentReq: + HookReportResp: type: object required: - - agent_uuid + - hook_uuid properties: - agent_uuid: + hook_uuid: type: string format: uuid - description: The agent to include, exclude, or remove the override for + description: The key this hook's artifacts are stored under + url: + type: string + description: Presigned upload URL, for the Upload operation - SuiteAgentResp: + FetchTasksReq: type: object required: - suite_uuid - - agent_uuid - - selection properties: suite_uuid: type: string format: uuid - agent_uuid: - type: string - format: uuid - selection: - $ref: "#/components/schemas/SuiteAgentSelectionType" + max_count: + type: integer + format: uint32 + default: 1 - ResetSuiteAgentResp: + FetchTasksResp: type: object required: - - reset + - tasks properties: - reset: + tasks: + type: array + items: + $ref: "#/components/schemas/WorkerTaskResp" + hold_job_open: type: boolean - description: Whether a manual override existed and was cleared (false = nothing to reset) + default: false + description: >- + Whether the agent should keep waiting on this job rather than wind it + down — true while the suite is `Open`, which covers "drained, but not + idle long enough to be sure". An empty `tasks` with this set means + "come back and ask again", not "there is no more work", so a task + submitted moments later runs in the same job against the environment + its provision hook already built. It goes false once the + coordinator's idle sweep has settled the suite. + + ReportAgentTaskReq: + type: object + required: + - job + - id + - op + properties: + job: + type: integer + format: int64 + description: The job this task belongs to + id: + type: integer + format: int64 + description: Internal task id, from the fetched task + op: + $ref: "#/components/schemas/ReportTaskOp" + + AgentNotification: + type: object + required: + - type + discriminator: + propertyName: type + properties: + type: + type: string + enum: + - suite_available + - preempt_suite + - suite_cancelled + - tasks_cancelled + - shutdown + - ping + - counter_sync + description: | + A coordinator → agent hint prompting an HTTP action: + - suite_available: a suite is ready for this agent to claim + - preempt_suite: stop the current suite for a higher-priority one. + **Not emitted yet** — a running agent is never interrupted; the variant + exists so preemption can be enabled without a protocol change + - suite_cancelled: the suite being run was cancelled + - tasks_cancelled: specific tasks were cancelled + - shutdown: stop, gracefully or immediately + - ping: keepalive + - counter_sync: reset the notification sequence (coordinator restart) + + WsNotificationEvent: + type: object + required: + - id + - event + properties: + id: + type: integer + format: uint64 + description: Per-agent, monotonically increasing sequence id + event: + $ref: "#/components/schemas/AgentNotification" + + AgentWsMessage: + type: object + required: + - type + discriminator: + propertyName: type + properties: + type: + type: string + enum: + - ack + - pong + description: | + An agent → coordinator frame: + - ack: the agent processed a notification, identified by + `notification_id`; it then leaves the replay buffer + - pong: reply to a ping, carrying `client_time` + + RemoteResourceDownloadResp: + type: object + required: + - url + - size + properties: + url: + type: string + description: Presigned download URL + size: + type: integer + format: int64 security: - bearerAuth: [] @@ -4742,7 +5963,9 @@ tags: - name: Tasks description: Task submission and monitoring - name: Suites - description: Task suite lifecycle and manual agent assignment + description: Task suite lifecycle, agent assignment, and job inspection + - name: Agents + description: Agent fleet management and the suite execution loop - name: Admin description: Administrative operations - name: Health diff --git a/src/main.rs b/src/main.rs index 7100479e..ab87f121 100644 --- a/src/main.rs +++ b/src/main.rs @@ -1,8 +1,11 @@ use crate::build::CLAP_LONG_VERSION; use clap::{Parser, Subcommand}; use netmito::{ + agent::MitoAgent, client::MitoClient, - config::{ClientConfigCli, CoordinatorConfigCli, ManagerConfigCli, WorkerConfigCli}, + config::{ + AgentConfigCli, ClientConfigCli, CoordinatorConfigCli, ManagerConfigCli, WorkerConfigCli, + }, coordinator::MitoCoordinator, manager::MitoManager, worker::MitoWorker, @@ -26,6 +29,8 @@ enum Mode { Coordinator(CoordinatorConfigCli), /// Run a mitosis worker. Worker(WorkerConfigCli), + /// Run a mitosis agent, which executes task suites. + Agent(AgentConfigCli), /// Run a mitosis client. Client(ClientConfigCli), /// Manage mitosis workers. @@ -53,6 +58,17 @@ fn main() { MitoWorker::main(worker_cli).await; }); } + Mode::Agent(agent_cli) => { + // Multi-threaded: the agent runs its main loop, a WebSocket reader, + // and a suite runner concurrently. + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .unwrap() + .block_on(async { + MitoAgent::main(agent_cli).await; + }); + } Mode::Client(client_cli) => { tokio::runtime::Builder::new_current_thread() .enable_all()