From ac61cb623a12c7d361881bd3c914c99b4a20deca Mon Sep 17 00:00:00 2001 From: un-lock-able Date: Fri, 30 Oct 2026 16:12:52 +0800 Subject: [PATCH 1/2] feat(agent): dispatch suite tasks to agents --- netmito/src/api/agents.rs | 91 +++- netmito/src/api/groups.rs | 12 +- netmito/src/api/suites.rs | 49 +- netmito/src/api/tasks.rs | 49 +- netmito/src/client/http.rs | 129 +++++ netmito/src/client/interactive.rs | 35 +- netmito/src/client/mod.rs | 85 +++- netmito/src/config/client/agents.rs | 15 + netmito/src/config/client/suites.rs | 78 ++- netmito/src/config/coordinator.rs | 15 + netmito/src/coordinator.rs | 38 ++ netmito/src/entity/state.rs | 22 +- netmito/src/executor/mod.rs | 119 ++--- netmito/src/executor/report.rs | 21 +- netmito/src/executor/upload.rs | 46 +- netmito/src/schema/agent.rs | 273 +++++++++- netmito/src/schema/suite.rs | 31 +- netmito/src/schema/task.rs | 12 + netmito/src/service/agent/hook.rs | 196 +++++++ netmito/src/service/agent/job.rs | 281 ++++++++++ netmito/src/service/agent/matching.rs | 236 +++++++++ netmito/src/service/agent/mod.rs | 704 ++++++++++++++++++++++++-- netmito/src/service/agent/task.rs | 334 ++++++++++++ netmito/src/service/mod.rs | 31 +- netmito/src/service/s3.rs | 240 +++++---- netmito/src/service/suite.rs | 411 +++++++++++++-- netmito/src/service/task.rs | 9 +- netmito/src/ws/connection.rs | 13 +- openapi.yaml | 666 +++++++++++++++++++++++- 29 files changed, 3882 insertions(+), 359 deletions(-) create mode 100644 netmito/src/service/agent/hook.rs create mode 100644 netmito/src/service/agent/job.rs create mode 100644 netmito/src/service/agent/matching.rs create mode 100644 netmito/src/service/agent/task.rs diff --git a/netmito/src/api/agents.rs b/netmito/src/api/agents.rs index 6532ec0..50282e6 100644 --- a/netmito/src/api/agents.rs +++ b/netmito/src/api/agents.rs @@ -1,7 +1,7 @@ use axum::{ extract::{Path, Query, State}, middleware, - routing::{delete, post}, + routing::{delete, get, post}, Extension, Json, Router, }; use uuid::Uuid; @@ -11,8 +11,9 @@ use crate::{ config::InfraPool, error::ApiError, schema::{ - AgentHeartbeatReq, AgentHeartbeatResp, AgentShutdownReq, AgentsQueryReq, AgentsQueryResp, - RegisterAgentReq, RegisterAgentResp, + AcceptSuiteReq, AcceptSuiteResp, AgentHeartbeatReq, AgentHeartbeatResp, AgentShutdownReq, + AgentsQueryReq, AgentsQueryResp, FetchTasksReq, FetchTasksResp, JobReportOp, JobReportReq, + JobReportResp, RegisterAgentReq, RegisterAgentResp, StopAgentJobReq, StopAgentJobResp, }, service::{ self, @@ -29,6 +30,7 @@ pub fn agents_router(st: InfraPool) -> Router { .route("/", post(register_agent)) .route("/query", post(query_agents)) .route("/{uuid}", delete(shutdown_agent)) + .route("/{uuid}/job", delete(stop_agent_job)) .route_layer(middleware::from_fn_with_state( st.clone(), auth_middleware::, @@ -38,6 +40,9 @@ pub fn agents_router(st: InfraPool) -> Router { let agent_router = Router::new() .route("/", delete(agent_exit)) .route("/heartbeat", post(heartbeat)) + .route("/suite", post(accept_suite)) + .route("/job", post(report_job)) + .route("/tasks", get(fetch_tasks)) .route_layer(middleware::from_fn_with_state( st.clone(), auth_middleware::, @@ -84,6 +89,20 @@ async fn shutdown_agent( Ok(()) } +/// `DELETE /agents/{uuid}/job?op=graceful|force`: stop the job that the agent runs now. The agent +/// stays up and picks a suite again. +async fn stop_agent_job( + Extension(u): Extension, + State(pool): State, + Path(uuid): Path, + Query(req): Query, +) -> Result, ApiError> { + let resp = service::agent::user_stop_agent_job(u.id, uuid, req.op, &pool) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} + /// `DELETE /agents`: the agent exits. Retire it now, so that it does not wait for its heartbeat times out. async fn agent_exit( Extension(a): Extension, @@ -106,3 +125,69 @@ async fn heartbeat( .map_err(map_service_error)?; Ok(Json(resp)) } + +/// `POST /agents/suite`: pick a suite, claim it, and open a job. +/// +/// The transaction picks the suite, thus the span records the suite and the job when they are +/// known. +#[tracing::instrument( + skip_all, + fields( + agent_uuid = %a.uuid, + suite_uuid = tracing::field::Empty, + job_id = tracing::field::Empty, + ) +)] +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`: move the job to the next phase. Only `Complete` returns data. +async fn report_job( + Extension(a): Extension, + State(pool): State, + Json(req): Json, +) -> Result, ApiError> { + let next_suite_available = match req.op { + JobReportOp::Start => { + service::agent::agent_start_job(a.id, &pool, req.job) + .await + .map_err(map_service_error)?; + None + } + JobReportOp::EnterCleanup => { + service::agent::agent_enter_cleanup(a.id, &pool, req.job) + .await + .map_err(map_service_error)?; + None + } + JobReportOp::Complete { outcome } => Some( + service::agent::agent_complete_job(a.id, a.uuid, &pool, req.job, outcome) + .await + .map_err(map_service_error)?, + ), + }; + Ok(Json(JobReportResp { + next_suite_available, + })) +} + +/// `GET /agents/tasks?suite_uuid=…`: claim ready tasks of the suite. +#[tracing::instrument(skip_all, fields(agent_uuid = %a.uuid, suite_uuid = %req.suite_uuid))] +async fn fetch_tasks( + Extension(a): Extension, + State(pool): State, + Query(req): Query, +) -> Result, ApiError> { + let resp = service::agent::task::agent_fetch_tasks(a.id, a.uuid, &pool, req.suite_uuid) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} diff --git a/netmito/src/api/groups.rs b/netmito/src/api/groups.rs index 8f1647f..a0ffd85 100644 --- a/netmito/src/api/groups.rs +++ b/netmito/src/api/groups.rs @@ -8,13 +8,13 @@ use axum_extra::extract::Query; use crate::{ config::InfraPool, - error::{ApiError, ApiResult, AuthError}, + error::{ApiError, ApiResult}, schema::*, service::{ self, auth::{ auth_middleware, - roles::{USER, WORKER}, + roles::{AGENT, USER, WORKER}, AuthIdentity, AuthUser, }, }, @@ -28,7 +28,7 @@ pub fn groups_router(st: InfraPool) -> Router { ) .route_layer(middleware::from_fn_with_state( st.clone(), - auth_middleware::<{ USER | WORKER }>, + auth_middleware::<{ USER | WORKER | AGENT }>, )); let user_router = Router::new() .route("/", post(create_group)) @@ -195,9 +195,9 @@ pub async fn download_attachment( AuthIdentity::Worker(w) => { service::s3::worker_get_attachment(&pool, w.uuid, group_name, key).await } - // TODO: let an agent read the attachments of a group whose suite it runs a job for. - // Do this when agents run jobs. Before that, no job rows exist and the check always fails. - AuthIdentity::Agent(_) => Err(AuthError::PermissionDenied.into()), + AuthIdentity::Agent(a) => { + service::s3::agent_get_attachment(&pool, a.id, group_name, key).await + } } .map_err(|e| match e { crate::error::Error::AuthError(err) => ApiError::AuthError(err), diff --git a/netmito/src/api/suites.rs b/netmito/src/api/suites.rs index dec0314..9b31024 100644 --- a/netmito/src/api/suites.rs +++ b/netmito/src/api/suites.rs @@ -12,8 +12,9 @@ use crate::{ error::ApiError, schema::{ CancelTaskSuiteParam, ChangeTaskSuiteReq, CreateTaskSuiteReq, CreateTaskSuiteResp, - SuiteAgentOverrideReq, SuiteAgentOverrideResp, TaskSuiteQueryResp, TaskSuitesQueryReq, - TaskSuitesQueryResp, + StopAgentJobReq, StopAgentJobResp, SuiteAgentOverrideReq, SuiteAgentOverrideResp, + SuiteJobQueryResp, SuiteJobsQueryReq, SuiteJobsQueryResp, TaskSuiteQueryResp, + TaskSuitesQueryReq, TaskSuitesQueryResp, }, service::{ self, @@ -33,6 +34,11 @@ pub fn suites_router(st: InfraPool) -> Router { ) .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).delete(stop_suite_job), + ) .route_layer(middleware::from_fn_with_state( st.clone(), auth_middleware::, @@ -119,3 +125,42 @@ pub async fn override_agents_for_suite( .map_err(map_service_error)?; Ok(Json(resp)) } + +/// `POST /suites/{uuid}/jobs/query`: the jobs of the suite, with a filter. +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 with its hooks. +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)) +} + +/// `DELETE /suites/{uuid}/jobs/{job_id}?op=graceful|force`: stop the job. The agent stays up and +/// picks a suite again, possibly this suite. +pub async fn stop_suite_job( + Extension(u): Extension, + State(pool): State, + Path((uuid, job_id)): Path<(Uuid, i32)>, + Query(req): Query, +) -> Result, ApiError> { + let resp = service::suite::user_stop_suite_job(u.id, &pool, uuid, job_id, req.op) + .await + .map_err(map_service_error)?; + Ok(Json(resp)) +} diff --git a/netmito/src/api/tasks.rs b/netmito/src/api/tasks.rs index a00d72d..4a6b249 100644 --- a/netmito/src/api/tasks.rs +++ b/netmito/src/api/tasks.rs @@ -9,14 +9,14 @@ use uuid::Uuid; use crate::{ config::InfraPool, entity::content::ArtifactContentType, - error::ApiError, + error::{ApiError, AuthError}, schema::*, service::{ self, auth::{ auth_middleware, - roles::{USER, WORKER}, - AuthIdentity, AuthUser, AuthWorker, + roles::{AGENT, USER, WORKER}, + AuthIdentity, AuthUser, }, }, }; @@ -30,13 +30,13 @@ pub fn tasks_router(st: InfraPool) -> Router { ) .route_layer(middleware::from_fn_with_state( st.clone(), - auth_middleware::<{ USER | WORKER }>, + auth_middleware::<{ USER | WORKER | AGENT }>, )); - let worker_router = Router::new() + let runner_router = Router::new() .route("/report", post(report_task)) .route_layer(middleware::from_fn_with_state( st.clone(), - auth_middleware::, + auth_middleware::<{ WORKER | AGENT }>, )); let user_router = Router::new() .route("/", post(submit_task)) @@ -67,7 +67,7 @@ pub fn tasks_router(st: InfraPool) -> Router { )); Router::new() .merge(shared_router) - .merge(worker_router) + .merge(runner_router) .merge(user_router) .with_state(st) } @@ -164,12 +164,41 @@ pub async fn query_task( Ok(Json(task)) } +/// `POST /tasks/report`: a worker reports its task, an agent reports a task or a hook of its job. pub async fn report_task( - Extension(w): Extension, + Extension(p): Extension, State(pool): State, - Json(ReportTaskReq { id, op }): Json, + Json(ReportTaskReq { id, kind, job, op }): Json, ) -> Result, ApiError> { - match service::worker::report_task(w.id, w.uuid, id, op, &pool).await { + let result = match (p, kind) { + (AuthIdentity::Worker(w), TaskKind::Task) => { + if job.is_some() { + return Err(ApiError::InvalidRequest( + "A worker task has no job".to_string(), + )); + } + service::worker::report_task(w.id, w.uuid, id, op, &pool).await + } + (AuthIdentity::Worker(_), TaskKind::Hook) => { + return Err(ApiError::AuthError(AuthError::PermissionDenied)); + } + (AuthIdentity::Agent(a), TaskKind::Task) => { + let Some(job) = job else { + return Err(ApiError::InvalidRequest( + "An agent task report must give its job".to_string(), + )); + }; + service::agent::task::agent_report_task(a.id, a.uuid, job, id, op, &pool).await + } + (AuthIdentity::Agent(a), TaskKind::Hook) => { + service::agent::hook::agent_report_hook(a.id, id, op, &pool).await + } + // The middleware of this route admits only workers and agents. + (AuthIdentity::User(_), _) => { + return Err(ApiError::AuthError(AuthError::PermissionDenied)); + } + }; + match result { Ok(url) => Ok(Json(ReportTaskResp { url })), Err(crate::error::Error::ApiError(e)) => Err(e), Err(e) => { diff --git a/netmito/src/client/http.rs b/netmito/src/client/http.rs index 7fcd62b..3ab736e 100644 --- a/netmito/src/client/http.rs +++ b/netmito/src/client/http.rs @@ -941,6 +941,135 @@ impl MitoHttpClient { } } + pub async fn stop_agent_job( + &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}/job")); + 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)?; + self.url.set_query(None); + 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_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()) + } + } + + 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 stop_suite_job( + &mut self, + suite_uuid: Uuid, + job_id: i32, + force: bool, + ) -> 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}")); + 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)?; + self.url.set_query(None); + 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 get_group_by_name( &mut self, group_name: &str, diff --git a/netmito/src/client/interactive.rs b/netmito/src/client/interactive.rs index c017a8b..c98bbdb 100644 --- a/netmito/src/client/interactive.rs +++ b/netmito/src/client/interactive.rs @@ -10,8 +10,8 @@ use crate::{ config::client::ClientInteractiveShell, entity::role::GroupWorkerRole, schema::{ - AdminChangePasswordReq, CreateUserReq, GroupQueryInfo, ParsedTaskQueryInfo, - ParsedTaskSuiteInfo, TaskQueryInfo, TaskSuiteInfo, UserChangePasswordReq, WorkerQueryInfo, + AdminChangePasswordReq, CreateUserReq, GroupQueryInfo, ParsedTaskQueryInfo, TaskQueryInfo, + TaskSuiteInfo, TaskSuiteQueryResp, UserChangePasswordReq, WorkerQueryInfo, }, service::auth::{get_and_prompt_password, get_and_prompt_username}, }; @@ -261,7 +261,12 @@ 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(resp: &TaskSuiteQueryResp) { + let TaskSuiteQueryResp { + info, + eligible_agents, + active_jobs, + } = resp; tracing::info!("Suite UUID: {}", info.uuid); if let Some(ref name) = info.name { tracing::info!("Name: {}", name); @@ -298,14 +303,30 @@ 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!("Eligible agents: None"); } else { - tracing::info!("Manually-included agents:"); - for agent in assigned_agents { + tracing::info!("Eligible agents:"); + for agent in eligible_agents { tracing::info!(" > {}", agent); } } + if active_jobs.is_empty() { + tracing::info!("Active jobs: None"); + } else { + tracing::info!("Active jobs:"); + for job in active_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, + ); + } + } } pub(crate) fn output_upload_artifact_resp(exist: bool) { diff --git a/netmito/src/client/mod.rs b/netmito/src/client/mod.rs index 3954f2a..8d9fc6f 100644 --- a/netmito/src/client/mod.rs +++ b/netmito/src/client/mod.rs @@ -901,6 +901,62 @@ impl MitoClient { } } + /// Show the jobs of a suite, or one job in full with `--job`. + 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 hooks"); + } + 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, @@ -2118,7 +2174,7 @@ impl MitoClient { } SuitesCommands::Get(args) => match self.suites_get(args).await { Ok(resp) => { - output_parsed_suite_info(&resp.info, &resp.eligible_agents); + output_parsed_suite_info(&resp); } Err(e) => { tracing::error!("{}", e); @@ -2143,6 +2199,17 @@ impl MitoClient { SuitesCommands::Override(args) => { self.suites_override_agents(args).await; } + SuitesCommands::Jobs(args) => self.suites_jobs(args).await, + SuitesCommands::StopJob(args) => { + match self + .http_client + .stop_suite_job(args.uuid, args.job, args.force) + .await + { + Ok(resp) => report_stopped_job(resp), + Err(e) => tracing::error!("{}", e), + } + } }, ClientCommand::Agents(args) => match args.command { AgentsCommands::Query(args) => { @@ -2179,6 +2246,12 @@ impl MitoClient { Err(e) => tracing::error!("{}", e), } } + AgentsCommands::StopJob(args) => { + match self.http_client.stop_agent_job(args.uuid, args.force).await { + Ok(resp) => report_stopped_job(resp), + Err(e) => tracing::error!("{}", e), + } + } }, ClientCommand::Quit => { return false; @@ -2215,3 +2288,13 @@ impl MitoClient { req } } + +/// Both stop-job commands use this. "Nothing to stop" is a normal answer, not an error. +fn report_stopped_job(resp: StopAgentJobResp) { + match (resp.stopped, resp.suite_uuid, resp.job_id) { + (true, Some(suite_uuid), Some(job_id)) => tracing::info!( + "Asked the agent to stop job {job_id} of suite {suite_uuid}; it will pick a suite again" + ), + _ => tracing::info!("No job was running to stop"), + } +} diff --git a/netmito/src/config/client/agents.rs b/netmito/src/config/client/agents.rs index 7d92f59..2fd6c34 100644 --- a/netmito/src/config/client/agents.rs +++ b/netmito/src/config/client/agents.rs @@ -16,6 +16,8 @@ pub enum AgentsCommands { Query(QueryAgentsArgs), /// Shut an agent down. The agent record stays. Shutdown(ShutdownAgentArgs), + /// Stop the job that an agent runs now. The agent then picks a suite again + StopJob(StopAgentJobArgs), } #[derive(Serialize, Debug, Deserialize, Args, Clone)] @@ -70,3 +72,16 @@ pub struct ShutdownAgentArgs { #[arg(short, long)] pub force: bool, } + +/// A stop moves an agent by hand: the agent ends the job and picks the suite with the highest +/// priority now. This can be the same suite, then the agent provisions it again. +#[derive(Serialize, Debug, Deserialize, Args, Clone)] +pub struct StopAgentJobArgs { + /// The UUID of the agent + pub uuid: Uuid, + /// Stop now: the job is killed without cleanup, and its tasks that are not committed go back to + /// Ready. Without this flag, the agent finishes and commits its running tasks, and runs the + /// cleanup hook first. + #[arg(short, long)] + pub force: bool, +} diff --git a/netmito/src/config/client/suites.rs b/netmito/src/config/client/suites.rs index cd3a08c..f751b87 100644 --- a/netmito/src/config/client/suites.rs +++ b/netmito/src/config/client/suites.rs @@ -5,10 +5,10 @@ use serde::{Deserialize, Serialize}; use uuid::Uuid; use crate::{ - entity::state::TaskSuiteState, + entity::state::{SuiteJobState, TaskSuiteState}, schema::{ ChangeTaskSuiteReq, CreateTaskSuiteReq, ExecHooks, ExecHooksUpdate, ExecSpec, - TaskSuitesQueryReq, UpdateOp, WorkerSchedulePlan, + SuiteJobsQueryReq, TaskSuitesQueryReq, UpdateOp, WorkerSchedulePlan, }, }; @@ -38,6 +38,10 @@ pub enum SuitesCommands { Cancel(CancelSuiteArgs), /// Set agent overrides on a suite in one batch: include, exclude, and/or clear agents Override(AgentsForSuiteOverrideArgs), + /// Show the jobs of the suite. A job is one run of the suite by one agent + Jobs(SuiteJobsArgs), + /// Stop a job of the suite. Its agent then picks a suite again + StopJob(StopSuiteJobArgs), } #[derive(Serialize, Debug, Deserialize, Args, Clone)] @@ -63,9 +67,10 @@ pub struct CreateSuiteArgs { /// Number of workers each agent spawns for this suite #[arg(short, long, default_value_t = 1)] pub workers: u32, - /// Number of tasks each worker prefetches locally - #[arg(long, default_value_t = 16)] - pub prefetch: u32, + /// Claim a task only when a worker is free. Without this flag, the agent keeps the next task + /// ready. Use it only for long tasks, because it adds a round trip for each task + #[arg(long, default_value_t = false)] + pub no_prefetch: bool, /// The provision hook, as a JSON exec spec, for example /// '{"args":["sh","-c","./setup.sh"],"terminal_output":true}'. It runs one time before the /// tasks. If it exits with a non-zero code, the job fails and its tasks do not start @@ -102,7 +107,7 @@ impl From for CreateTaskSuiteReq { worker_schedule: WorkerSchedulePlan::FixedWorkers { worker_count: args.workers, cpu_binding: None, - task_prefetch_count: args.prefetch, + prefetch: !args.no_prefetch, }, exec_hooks, } @@ -147,10 +152,10 @@ pub struct ChangeSuiteArgs { /// `--prefetch` #[arg(short, long, requires = "prefetch")] pub workers: Option, - /// The number of tasks that each worker gets before it needs them. You must also give - /// `--workers` + /// Whether the agent keeps the next task ready, or claims a task only when a worker is free. + /// You must also give `--workers` #[arg(long, requires = "workers")] - pub prefetch: Option, + pub prefetch: Option, /// The new provision hook, as a JSON exec spec #[arg(long, value_parser = parse_exec_spec, conflicts_with = "clear_provision")] pub provision: Option, @@ -202,7 +207,7 @@ impl From for ChangeTaskSuiteReq { (Some(worker_count), Some(prefetch)) => Some(WorkerSchedulePlan::FixedWorkers { worker_count, cpu_binding: None, - task_prefetch_count: prefetch, + prefetch, }), _ => None, }; @@ -304,3 +309,56 @@ 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, with its hooks + #[arg(long)] + pub job: Option, + /// Filter by job state + #[arg(long, num_args = 0.., value_delimiter = ',')] + pub states: Vec, + /// Show only the jobs of this agent + #[arg(long)] + pub agent: Option, + /// Maximum number of jobs to show + #[arg(long)] + pub limit: Option, + /// Number of jobs to skip + #[arg(long)] + pub offset: Option, + /// Only return the number of matching jobs + #[arg(long)] + pub count: bool, +} + +/// The same stop as `agents stop-job`, by the job number in the suite. The caller needs Write in +/// the group of the suite. The agent can accept this suite again: the command stops the job, it does +/// not keep the agent away. +#[derive(Serialize, Debug, Deserialize, Args, Clone)] +pub struct StopSuiteJobArgs { + /// The UUID of the suite + pub uuid: Uuid, + /// The job number in the suite, as `suites jobs` shows it + #[arg(long)] + pub job: i32, + /// Stop now: the job is killed without cleanup, and its tasks that are not committed go back to + /// Ready. Without this flag, the agent finishes and commits its running tasks, and runs the + /// cleanup hook first. + #[arg(short, long)] + pub force: 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 a3601c7..55daa36 100644 --- a/netmito/src/config/coordinator.rs +++ b/netmito/src/config/coordinator.rs @@ -75,6 +75,11 @@ pub struct CoordinatorConfig { /// proxies can close a notification socket that has no traffic. #[serde(with = "humantime_serde", default = "default_ws_keepalive_interval")] pub(crate) ws_keepalive_interval: std::time::Duration, + /// How long a suite can stay without a new task before the coordinator moves it out of + /// `Open`. It is also how long an agent keeps a job without tasks open for more work. A larger + /// value gives fewer provision hooks and more idle time of the machines. + #[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, } @@ -87,6 +92,10 @@ 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(60) +} + fn default_db_min_connections() -> u32 { 8 } @@ -188,6 +197,11 @@ pub struct CoordinatorConfigCli { #[arg(long)] #[serde(skip_serializing_if = "::std::option::Option::is_none")] pub ws_keepalive_interval: Option, + /// How long a suite can stay without a new task before it moves out of Open, and how long an + /// agent keeps a job without tasks. 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)] @@ -232,6 +246,7 @@ impl Default for CoordinatorConfig { 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, } diff --git a/netmito/src/coordinator.rs b/netmito/src/coordinator.rs index 15822fb..720789d 100644 --- a/netmito/src/coordinator.rs +++ b/netmito/src/coordinator.rs @@ -1,5 +1,6 @@ use std::net::SocketAddr; use std::path::PathBuf; +use std::time::Duration; use argon2::password_hash::rand_core::OsRng; use tokio_util::{sync::CancellationToken, task::TaskTracker}; @@ -10,10 +11,16 @@ use crate::config::{CoordinatorConfig, CoordinatorConfigCli, InfraPool}; use crate::migration::{Migrator, MigratorTrait}; use crate::service::agent::{heartbeat::AgentHeartbeatQueue, restore_agents}; use crate::service::s3::setup_buckets; +use crate::service::suite::sweep_inactive_suites; use crate::service::worker::{restore_workers, TaskDispatcher, WorkerHeartbeatQueue}; use crate::signal::shutdown_signal; use crate::ws::AgentWsRouter; +/// The period of the sweep of idle suites, for the configured idle window. +fn suite_sweep_period(idle_window: Duration) -> Duration { + (idle_window / 2).max(Duration::from_secs(1)) +} + pub struct MitoCoordinator { pub infra_pool: InfraPool, pub worker_task_queue: TaskDispatcher, @@ -22,6 +29,8 @@ pub struct MitoCoordinator { pub ws_router: AgentWsRouter, pub cancel_token: CancellationToken, pub log_dir: PathBuf, + /// How long a suite can stay without a new task before the sweep moves it out of `Open`. + pub suite_auto_close_timeout: Duration, } impl MitoCoordinator { @@ -138,6 +147,7 @@ impl MitoCoordinator { 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 @@ -168,6 +178,7 @@ impl MitoCoordinator { ws_router, cancel_token, log_dir, + suite_auto_close_timeout, }) } @@ -180,12 +191,39 @@ impl MitoCoordinator { 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(); + // Suite housekeeping: move the idle suites out of `Open`. + { + let pool = infra_pool.clone(); + let cancel = cancel_token.clone(); + let mut sweep = tokio::time::interval(suite_sweep_period(suite_auto_close_timeout)); + // A late tick is not done at once after the work that delayed it. There is no need to + // catch up. + sweep.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + task_tracker.spawn(async move { + // The first tick is immediate: it moves the suites that became idle while the + // coordinator was down. + loop { + tokio::select! { + biased; + _ = cancel.cancelled() => break, + _ = sweep.tick() => { + if let Err(e) = sweep_inactive_suites(&pool.db, suite_auto_close_timeout).await { + tracing::error!("Failed to sweep inactive suites: {e}"); + } + } + } + } + tracing::info!("Suite housekeeping stopped"); + }); + } + // Spawn background tasks using TaskTracker task_tracker.spawn(async move { worker_task_queue.run().await; diff --git a/netmito/src/entity/state.rs b/netmito/src/entity/state.rs index 35a68cd..6c28cb0 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. @@ -411,7 +419,13 @@ impl AgentState { } /// Lifecycle state of a hook task (provision / cleanup / background). -/// A hook task is reported on completion, so There is no `Running` state +/// +/// The coordinator creates the row as `Pending` when the agent accepts the suite. The agent +/// reports a hook like a task: `Finish` or `Cancel`, then `Commit` with the result. A hook has no +/// `Running` state, because the agent does not report the start of a hook. +/// +/// The row is closed when it has a result. When the job ends, the coordinator also changes the rows +/// that are still open: `Pending` becomes `Cancelled` and `Finished` becomes `Failed`. #[derive(EnumIter, DeriveActiveEnum, Clone, Debug, PartialEq, Eq, Serialize, Deserialize, Copy)] #[sea_orm(rs_type = "i32", db_type = "Integer")] pub enum HookExecState { @@ -421,6 +435,10 @@ pub enum HookExecState { Failed = 1, /// Hook was cancelled Cancelled = 2, + /// Hook is not reported yet. It may not have started + Pending = 3, + /// Hook process stopped, its result is not committed yet + Finished = 4, } impl Display for HookExecState { @@ -429,6 +447,8 @@ impl Display for HookExecState { HookExecState::Completed => write!(f, "Completed"), HookExecState::Failed => write!(f, "Failed"), HookExecState::Cancelled => write!(f, "Cancelled"), + HookExecState::Pending => write!(f, "Pending"), + HookExecState::Finished => write!(f, "Finished"), } } } diff --git a/netmito/src/executor/mod.rs b/netmito/src/executor/mod.rs index 7c36f46..5987d3e 100644 --- a/netmito/src/executor/mod.rs +++ b/netmito/src/executor/mod.rs @@ -91,19 +91,15 @@ impl Executor { _ = tokio::time::sleep(self.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, - }; - self.report_task(req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { + self.report_task(task.id, ReportTaskOp::Cancel).await?; + self.report_task( + task.id, + ReportTaskOp::Commit(TaskResultSpec { exit_status: 0, msg: Some(TaskResultMessage::FetchResourceTimeout), }), - }; - self.report_task(req).await?; + ) + .await?; self .announce_task_state_ex( &task.uuid, @@ -148,19 +144,15 @@ impl Executor { res = download_file(&self.http_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, - }; - self.report_task(req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { + self.report_task(task.id, ReportTaskOp::Cancel).await?; + self.report_task( + task.id, + ReportTaskOp::Commit(TaskResultSpec { exit_status: 0, msg: Some(TaskResultMessage::ResourceForbidden), }), - }; - self.report_task(req).await?; + ) + .await?; self .announce_task_state( &task.uuid, @@ -180,19 +172,15 @@ impl Executor { _ = self.task_cancel_token.cancelled() => return Ok(()), _ = tokio::time::sleep(self.single_download_budget) => { tracing::debug!("Fetching resource timeout, commit this task as canceled"); - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - self.report_task(req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { + self.report_task(task.id, ReportTaskOp::Cancel).await?; + self.report_task( + task.id, + ReportTaskOp::Commit(TaskResultSpec { exit_status: 0, msg: Some(TaskResultMessage::FetchResourceTimeout), }), - }; - self.report_task(req).await?; + ) + .await?; self .announce_task_state_ex( &task.uuid, @@ -207,19 +195,15 @@ impl Executor { } _ = 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, - }; - self.report_task(req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { + self.report_task(task.id, ReportTaskOp::Cancel).await?; + self.report_task( + task.id, + ReportTaskOp::Commit(TaskResultSpec { exit_status: 0, msg: Some(TaskResultMessage::FetchResourceTimeout), }), - }; - self.report_task(req).await?; + ) + .await?; self .announce_task_state_ex( &task.uuid, @@ -235,19 +219,15 @@ impl Executor { } } 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, - }; - self.report_task(req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { + self.report_task(task.id, ReportTaskOp::Cancel).await?; + self.report_task( + task.id, + ReportTaskOp::Commit(TaskResultSpec { exit_status: 0, msg: Some(TaskResultMessage::ResourceNotFound), }), - }; - self.report_task(req).await?; + ) + .await?; self.announce_task_state_ex( &task.uuid, TaskExecState::FetchResourceNotFound as i32, @@ -261,19 +241,15 @@ impl Executor { tracing::debug!( "Resource is forbidden to be fetched, commit this task as canceled" ); - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - self.report_task(req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { + self.report_task(task.id, ReportTaskOp::Cancel).await?; + self.report_task( + task.id, + ReportTaskOp::Commit(TaskResultSpec { exit_status: 0, msg: Some(TaskResultMessage::ResourceForbidden), }), - }; - self.report_task(req).await?; + ) + .await?; self.announce_task_state_ex( &task.uuid, TaskExecState::FetchResourceForbidden as i32, @@ -323,19 +299,15 @@ impl Executor { _ = tokio::time::sleep_until(timeout_until) => { tracing::debug!("Watching timeout, commit this task as canceled"); self.unsubscribe_task_exec_state(&watched_task_uuid).await; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Cancel, - }; - self.report_task(req).await?; - let req = ReportTaskReq { - id: task.id, - op: ReportTaskOp::Commit(TaskResultSpec { + self.report_task(task.id, ReportTaskOp::Cancel).await?; + self.report_task( + task.id, + ReportTaskOp::Commit(TaskResultSpec { exit_status: 0, msg: Some(TaskResultMessage::WatchTimeout), }), - }; - self.report_task(req).await?; + ) + .await?; self .announce_task_state_ex( &task.uuid, @@ -523,13 +495,6 @@ enum TaskResult { } 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, diff --git a/netmito/src/executor/report.rs b/netmito/src/executor/report.rs index 0563107..8db2743 100644 --- a/netmito/src/executor/report.rs +++ b/netmito/src/executor/report.rs @@ -8,7 +8,17 @@ use crate::schema::*; use super::Executor; impl Executor { - pub(super) async fn report_task(&mut self, req: ReportTaskReq) -> crate::error::Result<()> { + pub(super) async fn report_task( + &mut self, + id: i64, + op: ReportTaskOp, + ) -> crate::error::Result<()> { + let req = ReportTaskReq { + id, + kind: TaskKind::Task, + job: None, + op, + }; self.http_url.set_path("tasks/report"); loop { let resp = self @@ -101,11 +111,10 @@ impl Executor { task_id, submit_req.group_name ); - let req = ReportTaskReq { - id: task_id, - op: ReportTaskOp::Submit(Box::new(submit_req)), - }; - if let Err(e) = self.report_task(req).await { + if let Err(e) = self + .report_task(task_id, ReportTaskOp::Submit(Box::new(submit_req))) + .await + { tracing::warn!("Failed to submit new task: {}", e); return false; } diff --git a/netmito/src/executor/upload.rs b/netmito/src/executor/upload.rs index 90905d6..b9a41b1 100644 --- a/netmito/src/executor/upload.rs +++ b/netmito/src/executor/upload.rs @@ -25,24 +25,20 @@ impl Executor { uuid: Uuid, output: TaskResult, ) -> crate::error::Result<()> { - let (is_finished, exit_status) = output.state(); - let req = ReportTaskReq { - id, - op: if is_finished { - self.announce_task_state_ex(&uuid, TaskExecState::UploadFinishedResult as i32, 660) - .await; - ReportTaskOp::Finish - } else { - self.announce_task_state_ex( - &uuid, - TaskExecState::UploadCancelledResult as i32, - 660, - ) + let (is_finished, exit_status) = match &output { + TaskResult::Finish(output) => (true, output.get_exit_status()), + TaskResult::Timeout(output) => (false, output.get_exit_status()), + }; + let op = if is_finished { + self.announce_task_state_ex(&uuid, TaskExecState::UploadFinishedResult as i32, 660) .await; - ReportTaskOp::Cancel - }, + ReportTaskOp::Finish + } else { + self.announce_task_state_ex(&uuid, TaskExecState::UploadCancelledResult as i32, 660) + .await; + ReportTaskOp::Cancel }; - self.report_task(req).await?; + self.report_task(id, op).await?; // A spawned task archives each output, and this task uploads each archive when it is ready. // The channel sends the content type and size of each archive. let (tx, rx) = crossfire::spsc::unbounded_async::<(ArtifactContentType, u64)>(); @@ -232,6 +228,8 @@ impl Executor { while let Ok((content_type, content_length)) = rx.recv().await { let req = ReportTaskReq { id, + kind: TaskKind::Task, + job: None, op: ReportTaskOp::Upload { content_type, content_length, @@ -375,14 +373,14 @@ impl Executor { _ = tokio::time::sleep_until(timeout_until) => { tracing::warn!("Upload result timeout"); timeout_cancel_token.cancel(); - let req = ReportTaskReq { + self.report_task( id, - op: ReportTaskOp::Commit(TaskResultSpec { + ReportTaskOp::Commit(TaskResultSpec { exit_status: exit_status.into_raw(), msg: Some(TaskResultMessage::UploadResultTimeout), }), - }; - self.report_task(req).await?; + ) + .await?; self .announce_task_state_ex(&uuid, TaskExecState::UploadResultTimeout as i32, 60) .await; @@ -415,14 +413,14 @@ impl Executor { } else { Some(TaskResultMessage::ExecTimeout) }; - let req = ReportTaskReq { + self.report_task( id, - op: ReportTaskOp::Commit(TaskResultSpec { + ReportTaskOp::Commit(TaskResultSpec { exit_status: exit_status.into_raw(), msg }), - }; - self.report_task(req).await?; + ) + .await?; self .announce_task_state_ex(&uuid, TaskExecState::TaskCommitted as i32, 60) .await; diff --git a/netmito/src/schema/agent.rs b/netmito/src/schema/agent.rs index e3012b1..372062a 100644 --- a/netmito/src/schema/agent.rs +++ b/netmito/src/schema/agent.rs @@ -9,7 +9,14 @@ use speedy::{Readable, Writable}; use time::OffsetDateTime; use uuid::Uuid; -use crate::entity::state::AgentState; +use crate::entity::{ + hook_tasks::HookType, + state::{AgentState, HookExecState, SuiteJobState, TaskSuiteState}, +}; + +use super::exec::ExecHooks; +use super::suite::WorkerSchedulePlan; +use super::task::{TaskResultSpec, WorkerTaskResp}; /// Request to register an agent. /// @@ -119,6 +126,40 @@ pub enum AgentShutdownOp { Force, } +/// Query parameters of the stop-job endpoints. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct StopAgentJobReq { + #[serde(default)] + pub op: StopJobOp, +} + +/// How to stop a job. In both cases the agent stays up and then picks a suite again. Thus a stop +/// is how a user moves an agent to the suite with the highest priority now. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq)] +pub enum StopJobOp { + /// Claim no more tasks, let the running tasks finish, then run the cleanup hook. The agent + /// moves the job to `Completed` itself. + #[default] + #[serde(alias = "graceful")] + Graceful, + /// Stop now. The job becomes `Killed` and its tasks go back to `Ready`. The cleanup hook does + /// not run. + #[serde(alias = "force")] + Force, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StopAgentJobResp { + /// False if the agent had no job to stop. This is not an error. + pub stopped: bool, + /// The suite of the stopped job. + #[serde(skip_serializing_if = "Option::is_none")] + pub suite_uuid: Option, + /// The number of the stopped job in its suite. + #[serde(skip_serializing_if = "Option::is_none")] + pub job_id: Option, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct AgentHeartbeatReq { /// The state the agent reports for itself. @@ -134,6 +175,199 @@ pub struct AgentHeartbeatResp { pub notifications: Vec, } +/// Request body of `POST /agents/suite`. +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct AcceptSuiteReq { + /// The suite that the agent prefers. If this suite cannot run now, the coordinator picks the + /// best available suite. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub suite_uuid: Option, +} + +/// The suite that an agent runs, with all the data to run it. +#[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, +} + +/// A hook row of a job. The agent reports the hook with `id` on `POST /tasks/report`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobHook { + pub id: i64, + /// The key of the hook artifacts. + pub uuid: Uuid, + pub hook_type: HookType, +} + +/// Response of `POST /agents/suite`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AcceptSuiteResp { + /// False if the agent got no suite. This is not an error: no suite is available, or the agent + /// already runs a job. + pub accepted: bool, + /// The suite to run, if `accepted` is true. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub suite: Option, + /// The handle of the job (`suite_agent_jobs.id`). The agent sends it on each later call for + /// this job. + #[serde(skip_serializing_if = "Option::is_none")] + pub job: Option, + /// The number of the job in its suite. + #[serde(skip_serializing_if = "Option::is_none")] + pub job_id: Option, + /// One row for each hook of the suite. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + pub hooks: Vec, + /// Why the agent got no suite, if `accepted` is false. + #[serde(skip_serializing_if = "Option::is_none")] + pub reason: Option, +} + +/// Request body of `POST /agents/job`: the agent moves its job to the next phase. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobReportReq { + /// The job handle from `AcceptSuiteResp`. + pub job: i64, + pub op: JobReportOp, +} + +/// The phase that the agent starts. `Lost` and `Killed` are not here, because only the +/// coordinator writes them. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum JobReportOp { + /// The provision hook is done. The agent starts to run tasks. + Start, + /// The agent claims no more tasks. The cleanup hook starts. + EnterCleanup, + /// The job ended. `outcome` tells how. + Complete { outcome: SuiteJobOutcome }, +} + +/// How the job ended, as the agent reports it. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum SuiteJobOutcome { + /// All hooks completed successfully. + Completed, + /// The job failed. The hook rows keep the full output. + Failed { reason: JobFailureReason }, +} + +/// The phase in which a job failed. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum JobFailureKind { + ProvisionFailed, + /// The background hook stopped before the job ended. + BackgroundExited, + ExecutionError, + CleanupFailed, +} + +/// A short reason for a failed job. The job row has no column for it, thus the coordinator only +/// logs it. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct JobFailureReason { + pub kind: JobFailureKind, + pub message: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct JobReportResp { + /// For `Complete` only: whether the agent can accept another suite now. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub next_suite_available: Option, +} + +/// Query parameters of `GET /agents/tasks`. +/// +/// The coordinator calculates the number of tasks from the schedule of the suite. Thus the agent +/// does not give a number. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FetchTasksReq { + pub suite_uuid: Uuid, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct FetchTasksResp { + /// The claimed tasks. Empty if the suite has no ready task. + pub tasks: Vec, + /// True while the suite is `Open`. If `tasks` is empty and this is true, the agent keeps its + /// job and asks again later, because a new task can come before the suite closes. Thus the + /// provisioned environment is used again. + #[serde(default)] + pub hold_job_open: bool, +} + +/// Request body of `POST /suites/{uuid}/jobs/query`. All the fields must match. +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct SuiteJobsQueryReq { + #[serde(default, skip_serializing_if = "Option::is_none")] + pub states: Option>, + /// Only the jobs of 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 only the number of the matching jobs. + #[serde(default)] + pub count: bool, +} + +/// One job of a suite. +#[derive(Debug, Clone, Serialize, Deserialize, FromQueryResult)] +pub struct SuiteJobInfo { + /// The number of the job in its suite. + 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, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SuiteJobsQueryResp { + pub count: u64, + pub jobs: Vec, +} + +/// One hook of a job. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HookTaskInfo { + /// The key of the hook artifacts. + 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 of `GET /suites/{uuid}/jobs/{job_id}`. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct SuiteJobQueryResp { + pub info: SuiteJobInfo, + pub hooks: Vec, +} + /// A notification with its sequence id. The id tells an agent which /// notifications it has already processed. /// @@ -158,6 +392,43 @@ pub enum AgentNotification { /// The coordinator restarted and its notification ids start again from /// `counter`. `boot_id` identifies the new coordinator process. CounterSync { counter: u64, boot_id: Uuid }, + /// A suite that this agent can run has work. The event does not name the + /// suite: the coordinator picks the best suite when the agent accepts one, + /// because the best suite can change before the agent asks. + SuiteAvailable, + /// Stop the current suite and run a suite with a higher priority. + /// + /// The coordinator does not send this yet. The variant keeps its position + /// on the wire for a later change. + PreemptSuite { + new_suite_uuid: Uuid, + new_priority: i32, + current_suite_uuid: Uuid, + }, + /// The suite that the agent runs was cancelled. With `graceful`, claim no + /// more tasks and let the running tasks and the cleanup hook finish. + /// Otherwise stop the job now. + SuiteCancelled { + suite_uuid: Uuid, + reason: String, + graceful: bool, + }, + /// These tasks were cancelled. Stop them if they run. + TasksCancelled { task_uuids: Vec }, + /// Stop the job on `suite_uuid`. The agent stays up and picks a suite + /// again. An agent that runs a different suite ignores the event. + StopJob { suite_uuid: Uuid, graceful: bool }, +} + +impl AgentNotification { + /// Whether an event that waits in the buffer makes this event unnecessary. + /// + /// Only for the events that carry no data: the agent reads the current + /// state over HTTP, thus a second copy gives it nothing. A burst of + /// submissions also must not push older events out of the buffer. + pub fn coalesces_with(&self, other: &Self) -> bool { + matches!((self, other), (Self::SuiteAvailable, Self::SuiteAvailable)) + } } /// A message from an agent to the coordinator on the WebSocket. diff --git a/netmito/src/schema/suite.rs b/netmito/src/schema/suite.rs index c43654c..bb89f2d 100644 --- a/netmito/src/schema/suite.rs +++ b/netmito/src/schema/suite.rs @@ -7,6 +7,7 @@ use uuid::Uuid; use crate::entity::state::TaskSuiteState; +use super::agent::SuiteJobInfo; use super::exec::{ExecHooks, ExecSpec}; /// Request to create a new task suite @@ -138,19 +139,24 @@ impl ChangeTaskSuiteReq { #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(tag = "type", rename_all = "PascalCase")] pub enum WorkerSchedulePlan { - /// Fixed number of workers with optional CPU binding - /// This is the basic scheduling policy where a fixed number of workers - /// are spawned to process tasks from the suite - // TODO: should update it to our final design + /// Fixed number of workers with optional CPU binding This is the basic scheduling policy where + /// a fixed number of workers are spawned to process tasks from the suite + /// + /// TODO: should update it to our final design FixedWorkers { /// Number of workers to spawn (1-256) worker_count: u32, /// Optional CPU core binding strategy #[serde(default, skip_serializing_if = "Option::is_none")] cpu_binding: Option, - /// How many tasks to prefetch locally per worker (default: 16) - #[serde(default = "default_prefetch_count")] - task_prefetch_count: u32, + /// Whether the agent can hold claimed tasks in addition to the tasks that it runs. Then a + /// worker that becomes free gets its next task without a round trip. The agent holds at + /// most one more task for each worker. + /// + /// Turn it off only for suites whose tasks run much longer than a round trip. For these + /// suites, a task that waits on a busy agent is worse than a task that a free agent claims. + #[serde(default = "default_prefetch")] + prefetch: bool, }, // Future extensions: // AutoScale { min_workers, max_workers, scale_up_threshold, scale_down_threshold, ... } @@ -158,8 +164,8 @@ pub enum WorkerSchedulePlan { // Priority { high_priority_workers, low_priority_workers, ... } } -fn default_prefetch_count() -> u32 { - 16 +fn default_prefetch() -> bool { + true } /// CPU core binding configuration @@ -262,11 +268,16 @@ pub struct ParsedTaskSuiteInfo { pub completed_at: Option, } -/// Detailed suite response: the suite plus the UUIDs of its assigned agents +/// Detailed suite response: the suite, the UUIDs of the agents that can run it, and the jobs that +/// run it now. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct TaskSuiteQueryResp { pub info: ParsedTaskSuiteInfo, + /// The agents that match the tags or are included by hand, minus the agents that are excluded + /// by hand. The coordinator calculates this set for each query. pub eligible_agents: Vec, + /// The jobs of the suite that are not terminal, oldest first. + pub active_jobs: Vec, } /// Query parameter for `DELETE /suites/{uuid}` selecting the cancellation mode diff --git a/netmito/src/schema/task.rs b/netmito/src/schema/task.rs index 89b29d3..0616e88 100644 --- a/netmito/src/schema/task.rs +++ b/netmito/src/schema/task.rs @@ -23,10 +23,22 @@ pub struct WorkerTaskResp { #[derive(Debug, Serialize, Deserialize, Clone)] pub struct ReportTaskReq { + /// The id of the task, or of the hook if `kind` is `Hook`. pub id: i64, + pub kind: TaskKind, + /// The job handle of an agent task. Workers and hooks do not send it. + #[serde(default, skip_serializing_if = "Option::is_none")] + pub job: Option, pub op: ReportTaskOp, } +/// The kind of task that a report is for. Only agents report hooks. +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq)] +pub enum TaskKind { + Task, + Hook, +} + #[derive(Debug, Serialize, Deserialize, Clone)] pub enum ReportTaskOp { Finish, diff --git a/netmito/src/service/agent/hook.rs b/netmito/src/service/agent/hook.rs new file mode 100644 index 0000000..f4d2830 --- /dev/null +++ b/netmito/src/service/agent/hook.rs @@ -0,0 +1,196 @@ +//! The hook rows of a job (`hook_tasks`), and the hook reports on `POST /tasks/report`. +//! +//! The coordinator creates one `Pending` row for each hook of the suite when the agent accepts the +//! suite. The agent reports a hook like a task: `Finish` or `Cancel`, then `Upload` for each +//! artifact, then `Commit` with the result. The artifacts of a hook use the hook uuid as their key +//! in the `artifacts` table. +//! +//! A row is closed when it has a result. When the job ends, the coordinator closes the rows that +//! the agent did not finish ([`close_open_hooks`]). A row without a result still accepts `Upload` +//! and `Commit` after the job ends, because a hook can finish after the coordinator stopped the +//! job, and its result is still useful. + +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, JobHook, ReportTaskOp}; +use crate::service::agent::job; +use crate::service::s3::reserve_artifact_upload; + +/// Create one `Pending` row for each hook in `exec_hooks`. +pub async fn create_job_hooks( + db: &C, + job: i64, + exec_hooks: Option<&ExecHooks>, + now: TimeDateTimeWithTimeZone, +) -> Result> { + let Some(exec_hooks) = exec_hooks else { + return Ok(Vec::new()); + }; + let hooks = [ + (HookType::Provision, &exec_hooks.provision), + (HookType::Background, &exec_hooks.background), + (HookType::Cleanup, &exec_hooks.cleanup), + ]; + let mut created = Vec::new(); + for (hook_type, spec) in hooks { + let Some(spec) = spec else { continue }; + let row = HookTasks::ActiveModel { + uuid: Set(Uuid::new_v4()), + suite_agent_job_id: Set(job), + hook_type: Set(hook_type), + spec: Set(serde_json::to_value(spec)?), + state: Set(HookExecState::Pending), + created_at: Set(now), + updated_at: Set(now), + ..Default::default() + } + .insert(db) + .await?; + created.push(JobHook { + id: row.id, + uuid: row.uuid, + hook_type, + }); + } + Ok(created) +} + +/// Close the hook rows of these jobs that have no result: `Pending` becomes `Cancelled` (the hook +/// did not run) and `Finished` becomes `Failed` (the result did not come). +pub async fn close_still_open_hooks( + db: &C, + jobs: Vec, + now: TimeDateTimeWithTimeZone, +) -> Result<()> { + if jobs.is_empty() { + return Ok(()); + } + for (from, to) in [ + (HookExecState::Pending, HookExecState::Cancelled), + (HookExecState::Finished, HookExecState::Failed), + ] { + HookTasks::Entity::update_many() + .col_expr(HookTasks::Column::State, Expr::value(to)) + .col_expr(HookTasks::Column::UpdatedAt, Expr::value(now)) + .filter(HookTasks::Column::SuiteAgentJobId.is_in(jobs.clone())) + .filter(HookTasks::Column::State.eq(from)) + .filter(HookTasks::Column::Result.is_null()) + .exec(db) + .await?; + } + Ok(()) +} + +/// Apply a report of the agent to one of its hook rows. Returns the upload url for `Upload`. +/// +/// A missing row, a row of a different agent and a closed row give `NotFound`: the hook is no +/// longer the agent's. `Finish` and `Cancel` also give `NotFound` when the job is closed. +pub async fn agent_report_hook( + agent_id: i64, + hook_id: i64, + op: ReportTaskOp, + pool: &InfraPool, +) -> Result> { + let now = TimeDateTimeWithTimeZone::now_utc(); + + let hook = HookTasks::Entity::find_by_id(hook_id) + .one(&pool.db) + .await? + .ok_or_else(|| Error::ApiError(ApiError::NotFound(format!("Hook {hook_id}"))))?; + let job_row = job::load_validate_job(&pool.db, hook.suite_agent_job_id, agent_id).await?; + if hook.result.is_some() { + return Err(Error::ApiError(ApiError::NotFound(format!( + "Hook {} is closed with state {}", + hook.uuid, hook.state + )))); + } + + match op { + ReportTaskOp::Finish | ReportTaskOp::Cancel => { + job::reject_if_terminal(&job_row)?; + let state = match op { + ReportTaskOp::Finish => HookExecState::Finished, + _ => HookExecState::Cancelled, + }; + tracing::debug!(hook_uuid = %hook.uuid, %state, "Agent reported a hook"); + HookTasks::Entity::update_many() + .col_expr(HookTasks::Column::State, Expr::value(state)) + .col_expr(HookTasks::Column::UpdatedAt, Expr::value(now)) + .filter(HookTasks::Column::Id.eq(hook.id)) + .filter(HookTasks::Column::Result.is_null()) + .exec(&pool.db) + .await?; + } + + ReportTaskOp::Submit(_) => { + return Err(Error::ApiError(ApiError::InvalidRequest( + "A hook cannot submit a task".to_string(), + ))); + } + + ReportTaskOp::Upload { + content_type, + content_length, + } => { + let group_id = 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())))? + .group_id; + let hook_uuid = hook.uuid; + let s3_client = pool.s3.clone(); + let artifacts_bucket = pool.artifacts_bucket.clone(); + let (_, url) = pool + .db + .transaction::<_, (bool, String), Error>(|txn| { + Box::pin(async move { + reserve_artifact_upload( + txn, + &s3_client, + &artifacts_bucket, + group_id, + hook_uuid, + content_type, + content_length, + now, + ) + .await + }) + }) + .await?; + return Ok(Some(url)); + } + + ReportTaskOp::Commit(res) => { + let state = match hook.state { + HookExecState::Pending => { + return Err(Error::ApiError(ApiError::InvalidRequest( + "Hook must be Finished or Cancelled before Commit".to_string(), + ))); + } + HookExecState::Cancelled => HookExecState::Cancelled, + _ if res.exit_status == 0 => HookExecState::Completed, + _ => HookExecState::Failed, + }; + tracing::debug!(hook_uuid = %hook.uuid, %state, "Agent committed a hook"); + let result = serde_json::to_value(res)?; + HookTasks::Entity::update_many() + .col_expr(HookTasks::Column::State, Expr::value(state)) + .col_expr(HookTasks::Column::Result, Expr::value(result)) + .col_expr(HookTasks::Column::CompletedAt, Expr::value(now)) + .col_expr(HookTasks::Column::UpdatedAt, Expr::value(now)) + .filter(HookTasks::Column::Id.eq(hook.id)) + .filter(HookTasks::Column::Result.is_null()) + .exec(&pool.db) + .await?; + } + } + Ok(None) +} diff --git a/netmito/src/service/agent/job.rs b/netmito/src/service/agent/job.rs new file mode 100644 index 0000000..9f095dd --- /dev/null +++ b/netmito/src/service/agent/job.rs @@ -0,0 +1,281 @@ +//! The rows of `suite_agent_jobs`. +//! +//! A job is one run of a suite by one agent: accept, provision, execute, cleanup, then a terminal +//! state. These functions create the row, check the reports of an agent against it, and write the +//! terminal states that only the coordinator writes (`Lost`, `Killed`). +//! +//! A closed (terminal) job answers 404 to its agent: the job is no longer the agent's. A report +//! that comes in the wrong order answers 400. + +use sea_orm::{prelude::*, QueryOrder, QuerySelect, Set}; +use uuid::Uuid; + +use crate::entity::{ + active_tasks as ActiveTasks, agents as Agent, + state::{AgentState, SuiteJobState, TaskState}, + suite_agent_jobs as SuiteAgentJobs, task_suites as TaskSuites, +}; +use crate::error::{ApiError, Error, Result}; +use crate::service::agent::hook; + +/// The job states that are not terminal. +pub const IN_FLIGHT: [SuiteJobState; 3] = [ + SuiteJobState::Provisioning, + SuiteJobState::Executing, + SuiteJobState::Cleanup, +]; + +/// Create the job row of an agent that accepted a suite, in `Provisioning`. +/// +/// `job_id` is `max(job_id) + 1` in the suite. To calculate it safely, this function locks the +/// suite row, because accepts at the same time would otherwise get the same number. The lock comes +/// last, after all the checks, thus the caller holds no lock while it picks a suite. +/// +/// Call it inside a transaction, or the lock is released immediately. +pub async fn create_job( + db: &C, + suite_id: i64, + agent_id: i64, + now: TimeDateTimeWithTimeZone, +) -> Result { + TaskSuites::Entity::find_by_id(suite_id) + .lock_exclusive() + .one(db) + .await? + .ok_or_else(|| { + Error::ApiError(ApiError::NotFound(format!("Suite {suite_id} not found"))) + })?; + + 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?) +} + +/// Find a job by its handle and make sure that it belongs to this agent. A missing job and the job +/// of a different agent both give `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? + .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) +} + +/// Change a job of this agent to `to`, but only from one of the states in `from`. +/// +/// The state check and the write are one statement. Thus a terminal state that the coordinator +/// writes at the same time (`Lost`/`Killed`) either comes first, and this matches nothing, or comes +/// after. +/// +/// `Ok(None)` means that nothing matched. The caller decides if that is an error, and +/// [`explain_failed_transition`] makes the error. +pub async fn transition_job( + db: &C, + job: i64, + agent_id: i64, + from: &[SuiteJobState], + to: SuiteJobState, + now: TimeDateTimeWithTimeZone, +) -> Result> { + let updated = SuiteAgentJobs::Entity::update_many() + .col_expr(SuiteAgentJobs::Column::State, Expr::value(to)) + .col_expr(SuiteAgentJobs::Column::UpdatedAt, Expr::value(now)) + .filter(SuiteAgentJobs::Column::Id.eq(job)) + .filter(SuiteAgentJobs::Column::AgentId.eq(agent_id)) + .filter(SuiteAgentJobs::Column::State.is_in(from.iter().copied())) + .exec_with_returning(db) + .await?; + Ok(updated.into_iter().next()) +} + +/// The error for a [`transition_job`] that matched nothing: missing or not this agent's, or +/// terminal (`NotFound`), or in the wrong state that is not terminal (`InvalidRequest`). +/// `expected` is the source state of an ordered transition. `None` means that all the states that +/// are not terminal were accepted. +/// +/// The read has no lock and only makes the message, thus the successful path does not pay for it. +pub async fn explain_failed_transition( + db: &C, + job: i64, + agent_id: i64, + expected: Option, +) -> Error { + let row = match load_validate_job(db, job, agent_id).await { + Ok(row) => row, + Err(e) => return e, + }; + let rejection = match expected { + Some(expected) => expect_state(&row, expected), + None => reject_if_terminal(&row), + }; + match rejection { + Err(e) => e, + // The state changed again between the update and this read. + Ok(()) => Error::ApiError(ApiError::InvalidRequest(format!( + "Job {} changed state concurrently", + row.job_id + ))), + } +} + +/// Reject a report for a terminal job with `NotFound`: the job is no longer the agent's. +pub fn reject_if_terminal(job: &SuiteAgentJobs::Model) -> Result<()> { + if job.state.is_terminal() { + return Err(Error::ApiError(ApiError::NotFound(format!( + "Job {} is closed with state {}", + job.job_id, job.state + )))); + } + Ok(()) +} + +/// Make sure that a job is in the source state of an ordered transition: +/// - in `expected`: `Ok`. +/// - terminal: `NotFound`, the job is no longer the agent's. +/// - a different state that is not terminal: `InvalidRequest`. This is a protocol error, and it +/// must not look like a closed job, because the agent must not release a job that still runs. +pub fn expect_state(job: &SuiteAgentJobs::Model, expected: SuiteJobState) -> Result<()> { + if job.state == expected { + return Ok(()); + } + reject_if_terminal(job)?; + Err(Error::ApiError(ApiError::InvalidRequest(format!( + "Job {} is in state {}, expected {} for this transition", + job.job_id, job.state, expected + )))) +} + +/// Write `state` (`Lost` or `Killed`) to each job of the agent that is not terminal, and close +/// their open hooks. Returns the suite ids of these jobs. +pub async fn terminate_agent_jobs( + db: &C, + agent_id: i64, + state: SuiteJobState, + now: TimeDateTimeWithTimeZone, +) -> Result> { + let terminated = 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_with_returning(db) + .await?; + hook::close_still_open_hooks(db, terminated.iter().map(|j| j.id).collect(), now).await?; + Ok(terminated.into_iter().map(|j| j.task_suite_id).collect()) +} + +/// Stop each job of a cancelled suite that is not terminal (`Killed`, no cleanup), close their +/// open hooks, and release each agent that is still assigned to the suite. Returns the uuids of +/// the released agents. +/// +/// The agent cannot release itself: the coordinator rejects its `Complete` for a `Killed` job. +/// Call it in a transaction, so that the kill and the release occur together. +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?; + hook::close_still_open_hooks(db, killed.iter().map(|j| j.id).collect(), now).await?; + + // The filter is on the suite, not on the agents of the killed jobs. The suite is cancelled, + // thus each assignment to it is old, also the assignment of a job that ended without the + // release. + 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::AssignedTaskSuiteId.eq(suite_id)) + .exec_with_returning(db) + .await?; + + Ok(released.into_iter().map(|agent| agent.uuid).collect()) +} + +/// The uuids of the agents that run a job of the suite that is not terminal. +pub async fn agents_running_suite(db: &C, suite_id: i64) -> Result> { + Ok(SuiteAgentJobs::Entity::find() + .select_only() + .column(Agent::Column::Uuid) + .join( + sea_orm::JoinType::InnerJoin, + SuiteAgentJobs::Relation::Agents.def(), + ) + .filter(SuiteAgentJobs::Column::TaskSuiteId.eq(suite_id)) + .filter(SuiteAgentJobs::Column::State.is_in(IN_FLIGHT)) + .into_tuple::() + .all(db) + .await?) +} + +/// Put the tasks of an agent that are not committed back to `Ready`, without a runner, so that a +/// different agent runs them again. This includes `Finished` and `Cancelled` tasks: the result of a +/// task comes only with its `Commit`, and nothing archives a task in these states. +/// +/// `suite_id` limits the change: `None` takes all the tasks of the agent, `Some(id)` only its tasks +/// in that suite. Returns `(task id, suite id, priority)` for each task. +pub async fn reclaim_agent_tasks( + db: &C, + agent_uuid: Uuid, + suite_id: Option, + now: TimeDateTimeWithTimeZone, +) -> Result, i32)>> { + let mut query = 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, + TaskState::Cancelled, + ])); + if let Some(suite_id) = suite_id { + query = query.filter(ActiveTasks::Column::TaskSuiteId.eq(suite_id)); + } + let reclaimed = query.exec_with_returning(db).await?; + if !reclaimed.is_empty() { + tracing::info!( + agent_uuid = %agent_uuid, + reclaimed = reclaimed.len(), + "Reclaimed uncommitted tasks from an agent" + ); + } + Ok(reclaimed + .into_iter() + .map(|task| (task.id, task.task_suite_id, task.priority)) + .collect()) +} diff --git a/netmito/src/service/agent/matching.rs b/netmito/src/service/agent/matching.rs new file mode 100644 index 0000000..2b8ceda --- /dev/null +++ b/netmito/src/service/agent/matching.rs @@ -0,0 +1,236 @@ +//! Which agents can run which suites. +//! +//! The set of agents that can run a suite is +//! +//! ```text +//! (tag-matched ∪ UserIncluded) − UserExcluded only agents that the group of the suite can write to +//! ``` +//! +//! `task_suite_agent` keeps only the overrides that a user sets by hand. The tag-matched part is +//! not stored: each query calculates it in SQL. `agents.tags @> task_suites.tags` tells if an agent +//! has all the tags of the suite, and both `tags` columns have GIN indexes. +//! +//! All the eligibility checks of the agent service use the predicates here, thus the rule is in +//! one location. + +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, TaskState, TaskSuiteState}; +use crate::entity::task_suite_agent::SuiteAgentOverrideType; +use crate::entity::{ + active_tasks as ActiveTasks, 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)) +} + +/// An override row of type `kind` exists for the pair `(suite, agent)`. +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 group of the suite has at least the Write role on the agent. +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 pair `(suite, agent)`. The query that uses it must have +/// both `task_suites` and `agents`. +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 has work for an agent if it is `Open` or `Closed` and has a `Ready` task. +/// +/// `incomplete_tasks > 0` is not sufficient: it also counts the tasks that other agents run. A +/// suite whose remaining tasks all run would then get agents that find nothing to claim. +pub fn suite_has_work() -> SimpleExpr { + let has_claimable_task = Expr::exists( + Query::select() + .expr(Expr::val(1)) + .from(ActiveTasks::Entity) + .and_where( + Expr::col((ActiveTasks::Entity, ActiveTasks::Column::TaskSuiteId)) + .equals((TaskSuites::Entity, TaskSuites::Column::Id)), + ) + .and_where( + Expr::col((ActiveTasks::Entity, ActiveTasks::Column::State)).eq(TaskState::Ready), + ) + .to_owned(), + ); + + suite_col(TaskSuites::Column::State) + .is_in([TaskSuiteState::Open, TaskSuiteState::Closed]) + .and(has_claimable_task) +} + +/// `SELECT … FROM task_suites, agents WHERE `, for the callers to add their columns and +/// filters. +fn eligible_pairs() -> SelectStatement { + Query::select() + .from(TaskSuites::Entity) + .from(Agent::Entity) + .and_where(eligibility_predicate()) + .to_owned() +} + +/// Whether this agent can run this suite. It does not check if the suite 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()) +} + +/// Whether this agent has a suite with work now. +pub async fn agent_has_available_suite( + db: &C, + agent_id: i64, + blocked: &[i64], +) -> crate::error::Result { + Ok(best_available_suite_id(db, agent_id, blocked) + .await? + .is_some()) +} + +/// The suite that this agent should run next: the highest priority first, then the oldest. `None` +/// if no suite is available. +/// +/// The query excludes the suites in `blocked`. A filter after the query would be wrong: a blocked +/// suite at the top would hide the next suite. +pub async fn best_available_suite_id( + db: &C, + agent_id: i64, + blocked: &[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()); + if !blocked.is_empty() { + stmt.and_where(suite_col(TaskSuites::Column::Id).is_not_in(blocked.to_vec())); + } + stmt.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 all the agents that can run a suite. +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 the `Idle` agents that can run a suite: the agents that can start it now. +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 index 1635a49..1e18da6 100644 --- a/netmito/src/service/agent/mod.rs +++ b/netmito/src/service/agent/mod.rs @@ -1,25 +1,52 @@ -//! Agent registration, queries, heartbeats and shutdown. +//! Agents: registration, queries, heartbeats and shutdown for users, and the job loop for agents. +//! +//! ## The job loop +//! +//! ```text +//! register ──▶ heartbeat ──▶ accept suite ──▶ start ──▶ …tasks… ──▶ cleanup ──▶ complete +//! ▲ │ +//! └─────────────────────── Idle ◀───────────────────────────────────┘ +//! ``` +//! +//! Only `accept` claims a suite. One transaction picks the suite, checks it again, assigns it to +//! the agent, and creates a `suite_agent_jobs` row and its hook rows. The agent sends the job +//! handle on each later call. `start`, `cleanup` and `complete` move the job row forward. +//! +//! ## No automatic preemption +//! +//! An agent runs the suite that it accepted until the job ends. The priority of a suite only +//! changes which suite [`matching::best_available_suite_id`] picks when an agent looks for work. +//! To move an agent to a suite with a higher priority, a user stops its job +//! ([`user_stop_agent_job`]). The agent then picks a suite again. 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 sea_orm::{prelude::*, FromQueryResult, QueryOrder, 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, + active_tasks as ActiveTasks, agents as Agent, group_agent as GroupAgent, groups as Group, + machines as Machines, role::{GroupAgentRole, UserGroupRole}, - state::AgentState, - task_suites as TaskSuites, user_group as UserGroup, users as User, + state::{AgentState, SuiteJobState, TaskState, 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::{ - AgentHeartbeatReq, AgentHeartbeatResp, AgentInfo, AgentNotification, AgentShutdownOp, - AgentsQueryReq, AgentsQueryResp, CountQuery, RegisterAgentReq, RegisterAgentResp, + AcceptSuiteReq, AcceptSuiteResp, AgentHeartbeatReq, AgentHeartbeatResp, AgentInfo, + AgentNotification, AgentShutdownOp, AgentsQueryReq, AgentsQueryResp, CountQuery, ExecHooks, + JobHook, RegisterAgentReq, RegisterAgentResp, StopAgentJobResp, StopJobOp, SuiteJobOutcome, + TaskSuiteSpec, WorkerSchedulePlan, }; use crate::service::auth::token::{generate_runner_token, TokenRole}; use crate::ws::AgentWsRouter; @@ -59,7 +86,7 @@ pub async fn user_register_agent( let (agent_uuid, agent_id, reused, retired) = pool .db - .transaction::<_, (Uuid, i64, bool, bool), Error>(move |txn| { + .transaction::<_, (Uuid, i64, bool, Option), Error>(move |txn| { Box::pin(async move { // Exactly one group has the Admin role on the agent, and only // its Admins can shut the agent down. @@ -117,7 +144,7 @@ pub async fn user_register_agent( .one(txn) .await?; - let mut retired = false; + let mut retired = None; let (agent, reused) = match existing { Some(machine) => { let agent = Agent::Entity::find_by_id(machine.agent_id) @@ -137,8 +164,16 @@ pub async fn user_register_agent( // known machine code alone cannot retire an agent. let agent_state = agent.state; if agent_state != AgentState::Offline { - retire_agent_writes(txn, agent.id, now).await?; - retired = true; + retired = Some( + retire_agent_writes( + txn, + agent.id, + agent.uuid, + RetireCause::Reregistered, + now, + ) + .await?, + ); } let mut agent: Agent::ActiveModel = agent.into(); agent.tags = Set(tags); @@ -223,8 +258,15 @@ pub async fn user_register_agent( // This must come before the heartbeat below. It stops tracking the previous process, and the // heartbeat then starts to track the new one. - if retired { - retire_agent_post_commit(pool, agent_id, agent_uuid, RetireCause::Reregistered); + if let Some(retired) = retired { + retire_agent_post_commit( + pool, + agent_id, + agent_uuid, + RetireCause::Reregistered, + retired, + ) + .await; } let token = generate_runner_token(agent_uuid.to_string(), 0, req.lifetime, TokenRole::Agent)?; @@ -435,7 +477,7 @@ pub async fn user_shutdown_agent_by_uuid( Ok(()) } -/// The reason to retire an agent. +/// The reason to retire an agent. It also gives the terminal state of the jobs of the agent. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum RetireCause { /// The agent did not send a heartbeat in time. @@ -448,8 +490,25 @@ pub enum RetireCause { SelfExit, } -/// Retire an agent: mark it `Offline`, stop tracking its heartbeat, and remove its buffered -/// notifications. Its token stops working. +impl RetireCause { + fn into_job_terminal(self) -> SuiteJobState { + match self { + // The coordinator stopped the jobs. + Self::ForceShutdown => SuiteJobState::Killed, + // An agent that exits normally completed its job before. `Lost` labels only the jobs + // that it did not complete. + Self::HeartbeatTimeout | Self::Reregistered | Self::SelfExit => SuiteJobState::Lost, + } + } +} + +/// Retire an agent: mark it `Offline` without a suite, write a terminal state to its jobs that are +/// not terminal, put the tasks that it did not commit back to `Ready`, stop tracking its heartbeat, +/// remove its buffered notifications, and offer the returned work to other agents. Its token stops +/// working. +/// +/// These steps occur together, because each step without the others leaves work that nothing can +/// claim. /// /// This function does nothing for an agent that is already `Offline`, thus it is safe to call it /// more than once. @@ -465,25 +524,35 @@ pub async fn retire_agent(pool: &InfraPool, agent_id: i64, cause: RetireCause) - } let agent_uuid = agent.uuid; - pool.db - .transaction::<_, (), Error>(move |txn| { - Box::pin(async move { retire_agent_writes(txn, agent_id, now).await }) + let retired = pool + .db + .transaction::<_, RetiredWork, Error>(move |txn| { + Box::pin( + async move { retire_agent_writes(txn, agent_id, agent_uuid, cause, now).await }, + ) }) .await?; - retire_agent_post_commit(pool, agent_id, agent_uuid, cause); + retire_agent_post_commit(pool, agent_id, agent_uuid, cause, retired).await; Ok(()) } +/// The work that a retirement freed: the suites of the ended jobs, and the returned tasks as +/// `(task id, suite id, priority)`. +type RetiredWork = (Vec, Vec<(i64, Option, i32)>); + /// The database writes of [`retire_agent`]. Registration also calls this in its own transaction, /// because it retires the previous agent and updates the agent row in one step. async fn retire_agent_writes( txn: &C, agent_id: i64, + agent_uuid: Uuid, + cause: RetireCause, now: TimeDateTimeWithTimeZone, -) -> Result<()> { - // TODO: when agents run jobs, end the in-flight jobs of the agent and put - // the tasks that it holds back to `Ready` here, in the same transaction. +) -> Result { + let suite_ids = + job::terminate_agent_jobs(txn, agent_id, cause.into_job_terminal(), now).await?; + Agent::Entity::update_many() .col_expr(Agent::Column::State, Expr::value(AgentState::Offline)) .col_expr(Agent::Column::AssignedTaskSuiteId, Expr::value(None::)) @@ -491,17 +560,196 @@ async fn retire_agent_writes( .filter(Agent::Column::Id.eq(agent_id)) .exec(txn) .await?; - Ok(()) + + // All the suites: the agent leaves each suite at the same time. + let reclaimed = job::reclaim_agent_tasks(txn, agent_uuid, None, now).await?; + + Ok((suite_ids, reclaimed)) } /// The steps of [`retire_agent`] outside the database. They run after the commit, because a /// rollback cannot undo them. -fn retire_agent_post_commit(pool: &InfraPool, agent_id: i64, agent_uuid: Uuid, cause: RetireCause) { +async fn retire_agent_post_commit( + pool: &InfraPool, + agent_id: i64, + agent_uuid: Uuid, + cause: RetireCause, + (suite_ids, reclaimed): RetiredWork, +) { let _ = pool .agent_heartbeat_queue_tx .send(AgentHeartbeatOp::Remove(agent_id)); AgentWsRouter::reset(&pool.ws_router_tx, agent_uuid); - tracing::info!(agent_id, %agent_uuid, ?cause, "Retired an agent"); + tracing::info!( + agent_id, + %agent_uuid, + ?cause, + jobs = suite_ids.len(), + reclaimed = reclaimed.len(), + "Retired an agent" + ); + + // Offer the work again for each suite. This includes the suites whose jobs held no task, + // because the end of a job can let a different agent take the suite. + let mut suites: HashSet = suite_ids.into_iter().collect(); + suites.extend( + reclaimed + .into_iter() + .filter_map(|(_, suite_id, _)| suite_id), + ); + for suite_id in suites { + notify_suite_available(pool, suite_id).await; + } +} + +/// Stop the job that the agent runs now. The agent stays up. The caller must be an Admin of the +/// admin group of the agent, the same as for a shutdown. +pub async fn user_stop_agent_job( + user_id: i64, + agent_uuid: Uuid, + op: StopJobOp, + pool: &InfraPool, +) -> Result { + let agent = Agent::Entity::find() + .join_rev(sea_orm::JoinType::Join, GroupAgent::Relation::Agents.def()) + .join(sea_orm::JoinType::Join, GroupAgent::Relation::Groups.def()) + .join(sea_orm::JoinType::Join, Group::Relation::UserGroup.def()) + .filter(Agent::Column::Uuid.eq(agent_uuid)) + .filter(GroupAgent::Column::Role.eq(GroupAgentRole::Admin)) + .filter(UserGroup::Column::UserId.eq(user_id)) + .filter(UserGroup::Column::Role.eq(UserGroupRole::Admin)) + .one(&pool.db) + .await? + .ok_or(Error::ApiError(ApiError::NotFound(format!( + "User doesn't have permission or agent with uuid {agent_uuid}" + ))))?; + + // `accept` refuses a busy agent, thus there is at most one job. The order only makes the + // answer stable if there are more. + let job = SuiteAgentJobs::Entity::find() + .filter(SuiteAgentJobs::Column::AgentId.eq(agent.id)) + .filter(SuiteAgentJobs::Column::State.is_in(job::IN_FLIGHT)) + .order_by_desc(SuiteAgentJobs::Column::Id) + .one(&pool.db) + .await?; + let Some(job) = job else { + return Ok(StopAgentJobResp { + stopped: false, + suite_uuid: None, + job_id: None, + }); + }; + + let suite_uuid = TaskSuites::Entity::find_by_id(job.task_suite_id) + .one(&pool.db) + .await? + .map(|s| s.uuid) + .ok_or_else(|| Error::Custom(format!("Job {} points at a missing suite", job.id)))?; + + stop_agent_job(pool, &agent, &job, suite_uuid, op).await +} + +/// Stop one job that is not terminal. Both stop endpoints use this function. +/// +/// The agent stays `Idle`, not `Offline`: after the stop, it picks a suite again. Thus a user can +/// move an agent to a suite with a higher priority. +/// +/// `Force` writes the terminal state here, thus the coordinator rejects later reports for the job. +/// `Graceful` writes nothing: the agent moves the job to `Completed` itself, thus its running tasks +/// can commit. +pub(crate) async fn stop_agent_job( + pool: &InfraPool, + agent: &Agent::Model, + job: &SuiteAgentJobs::Model, + suite_uuid: Uuid, + op: StopJobOp, +) -> Result { + let now = TimeDateTimeWithTimeZone::now_utc(); + let (agent_id, agent_uuid) = (agent.id, agent.uuid); + let (job_handle, job_id) = (job.id, job.job_id); + + if op == StopJobOp::Force { + let stopped = pool + .db + .transaction::<_, Option, i32)>>, Error>(move |txn| { + Box::pin(async move { + // One statement, thus the terminal check and the write cannot occur on the two + // sides of a `Complete` of the agent. + let killed = job::transition_job( + txn, + job_handle, + agent_id, + &job::IN_FLIGHT, + SuiteJobState::Killed, + now, + ) + .await?; + let Some(job_row) = killed else { + // A job that is missing or of a different agent is an error. A job that is + // already terminal is not: it ended first, thus there is nothing to stop. + job::load_validate_job(txn, job_handle, agent_id).await?; + return Ok(None); + }; + hook::close_still_open_hooks(txn, vec![job_handle], now).await?; + + let suite_id = job_row.task_suite_id; + + // Release the agent only while it is still assigned to the suite of this job, + // the same as `agent_complete_job`. + 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?; + + // After the kill, these tasks will never commit. Only this suite, the same as + // `Complete`. + let reclaimed = + job::reclaim_agent_tasks(txn, agent_uuid, Some(suite_id), now).await?; + + Ok(Some(reclaimed)) + }) + }) + .await?; + + let Some(reclaimed) = stopped else { + return Ok(StopAgentJobResp { + stopped: false, + suite_uuid: None, + job_id: None, + }); + }; + + if !reclaimed.is_empty() { + notify_suite_available(pool, job.task_suite_id).await; + } + } + + tracing::info!( + agent_uuid = %agent_uuid, + job_id, + %suite_uuid, + graceful = op == StopJobOp::Graceful, + "Stopping a job of an agent on request" + ); + + AgentWsRouter::notify( + &pool.ws_router_tx, + agent_uuid, + AgentNotification::StopJob { + suite_uuid, + graceful: op == StopJobOp::Graceful, + }, + ); + + Ok(StopAgentJobResp { + stopped: true, + suite_uuid: Some(suite_uuid), + job_id: Some(job_id), + }) } /// Record a heartbeat and return the notifications that the agent has not acknowledged. @@ -521,22 +769,37 @@ pub async fn agent_heartbeat( // The agent knows best what it does, thus its reported state replaces the recorded one. The // filter on `Offline` prevents the heartbeat that is initiated from an agent that is in the process // of retiring to re-activate the agent row. - let updated = Agent::Entity::update_many() + let Some(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)) .filter(Agent::Column::State.ne(AgentState::Offline)) - .exec(&pool.db) - .await?; - if updated.rows_affected == 0 { + .exec_with_returning(&pool.db) + .await? + .into_iter() + .next() + else { return Err(Error::AuthError(AuthError::WrongCredentials)); - } + }; let _ = pool .agent_heartbeat_queue_tx .send(AgentHeartbeatOp::Heartbeat(agent_id)); + // An idle agent without a suite gets a new `SuiteAvailable` if a suite has work, in case it + // lost the earlier one. + if req.state == AgentState::Idle + && agent.assigned_task_suite_id.is_none() + && matching::agent_has_available_suite(&pool.db, agent_id, &[]).await? + { + AgentWsRouter::notify( + &pool.ws_router_tx, + agent_uuid, + AgentNotification::SuiteAvailable, + ); + } + // If the agent has processed a higher id than the coordinator has issued, the coordinator has // restarted and its ids started again. Tell the agent the current counter and boot id. if let Some(counter) = AgentWsRouter::counter(&pool.ws_router_tx, agent_uuid).await { @@ -568,6 +831,385 @@ pub async fn agent_heartbeat( Ok(AgentHeartbeatResp { notifications }) } +/// Tell the `Idle` agents that can run a suite that it has work. +/// +/// The event does not name the suite. The agent gets the best suite for it when it accepts one. +pub async fn notify_suite_available(pool: &InfraPool, suite_id: i64) { + 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, + ); + } +} + +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: assign it to the agent, and create the job row and +/// its hook rows. The pick occurs under the lock of the suite row, thus no other agent can take the +/// suite between the pick and the claim. +/// +/// `req.suite_uuid` is a preference. If that suite cannot run now (no ready task, cancelled, not +/// for this agent, or missing), the coordinator picks the best available suite. +/// +/// When the agent gets no suite (no suite is available, or the agent already runs a job), the +/// response has `accepted: false` and a reason. This is 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, Vec), String>, Error>( + |txn| { + Box::pin(async move { + // One suite at a time. A busy agent that asks again is not an error. + 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 + ))); + } + + // The suite that the agent asked for, if it can still run. + let requested = match req.suite_uuid { + Some(suite_uuid) => { + runnable_suite(txn, agent_id, SuiteTarget::Uuid(suite_uuid)).await? + } + None => None, + }; + // If not, the best available suite. + let suite = match requested { + Some(suite) => suite, + None => { + let picked = + match matching::best_available_suite_id(txn, agent_id, &[]).await? { + Some(suite_id) => { + runnable_suite(txn, agent_id, SuiteTarget::Id(suite_id)) + .await? + } + None => None, + }; + // `None` after a pick: the suite lost its work between the pick and + // the check. + let Some(suite) = picked else { + return Ok(Err("No suite is available for this agent".to_string())); + }; + suite + } + }; + + let suite_id = suite.id; + // Lock order: the job row locks the suite row, before the agent row. + let job = job::create_job(txn, suite_id, agent_id, now).await?; + let spec = suite_to_spec(txn, suite).await?; + let hooks = + hook::create_job_hooks(txn, job.id, spec.exec_hooks.as_ref(), now).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?; + + Ok(Ok((spec, job.id, job.job_id, hooks))) + }) + }, + ) + .await?; + + match outcome { + Ok((suite, job, job_id, hooks)) => { + let span = tracing::Span::current(); + span.record("suite_uuid", tracing::field::display(suite.uuid)); + span.record("job_id", job_id); + Ok(AcceptSuiteResp { + accepted: true, + suite: Some(suite), + job: Some(job), + job_id: Some(job_id), + hooks, + reason: None, + }) + } + Err(reason) => Ok(AcceptSuiteResp { + accepted: false, + suite: None, + job: None, + job_id: None, + hooks: Vec::new(), + reason: Some(reason), + }), + } +} + +/// How the agent got a candidate suite: it named the suite, or the coordinator picked it. +enum SuiteTarget { + Uuid(Uuid), + Id(i64), +} + +/// Return the suite if this agent can run it now. `None` for each reason that it cannot: missing, +/// not for this agent, in the wrong state, or no ready task. The caller then tries the next +/// candidate. +/// +/// This function takes no lock. A lock taken before the checks stays when the function returns +/// `None`. Then an agent whose preferred suite cannot run holds that suite while it locks a +/// different one, and two agents can deadlock. The checks are only advice: the hard rule on this +/// path is the unique index on the job id, and `create_job` locks the suite at the end. +async fn 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)), + } + .one(txn) + .await?; + + let Some(suite) = suite else { + return Ok(None); + }; + // The same rule as `matching::suite_has_work`, which the pick uses. + if !matches!(suite.state, TaskSuiteState::Open | TaskSuiteState::Closed) { + return Ok(None); + } + let claimable = ActiveTasks::Entity::find() + .filter(ActiveTasks::Column::TaskSuiteId.eq(suite.id)) + .filter(ActiveTasks::Column::State.eq(TaskState::Ready)) + .count(txn) + .await?; + if claimable == 0 { + return Ok(None); + } + if !matching::is_agent_eligible(txn, suite.id, agent_id).await? { + return Ok(None); + } + Ok(Some(suite)) +} + +/// The provision hook is done: 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 +} + +/// The agent claims no more tasks: 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 and the agent state that goes with it, in one transaction. +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 source state is in the `WHERE`. Thus a terminal state that the coordinator + // writes at the same time cannot give the agent of a closed job a live state. + let advanced = + job::transition_job(txn, job_handle, agent_id, &[from], to, now).await?; + let Some(job_row) = advanced else { + return Err(job::explain_failed_transition( + txn, + job_handle, + agent_id, + Some(from), + ) + .await); + }; + + // Do not change an agent that a retirement already released. + 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(job_row.task_suite_id)) + .exec(txn) + .await?; + + Ok(()) + }) + }) + .await?; + + Ok(()) +} + +/// End the job and release the agent to `Idle`. Returns whether a different suite is available +/// for this agent now. +/// +/// The agent reports only `Completed` or `Failed`. Only the coordinator writes `Lost` and `Killed`. +/// A job that is already terminal answers `NotFound`, thus the agent knows that the job was closed. +/// +/// One transaction writes the terminal state, closes the open hooks, releases the agent, and puts +/// the tasks that the agent did not commit in this suite back to `Ready`. If the release does not +/// occur, the agent stays assigned to a closed suite, and `accept` rejects it. The tasks must go +/// back, because nothing else moves a task that ran without a commit. +pub async fn agent_complete_job( + agent_id: i64, + agent_uuid: Uuid, + pool: &InfraPool, + job_handle: i64, + outcome: SuiteJobOutcome, +) -> Result { + let now = TimeDateTimeWithTimeZone::now_utc(); + + let (job_id, terminal, released, reclaimed, suite_id) = pool + .db + .transaction::<_, (i32, SuiteJobState, bool, Vec<(i64, Option, i32)>, i64), Error>( + move |txn| { + Box::pin(async move { + let terminal = match &outcome { + SuiteJobOutcome::Completed => SuiteJobState::Completed, + SuiteJobOutcome::Failed { .. } => SuiteJobState::Failed, + }; + + // One statement, thus the terminal check and the write cannot occur on the two + // sides of a `Killed` or `Lost` from a retirement. + let finished = job::transition_job( + txn, + job_handle, + agent_id, + &job::IN_FLIGHT, + terminal, + now, + ) + .await?; + let Some(job_row) = finished else { + return Err(job::explain_failed_transition( + txn, job_handle, agent_id, None, + ) + .await); + }; + hook::close_still_open_hooks(txn, vec![job_handle], now).await?; + + if let SuiteJobOutcome::Failed { reason } = &outcome { + tracing::warn!( + job = job_handle, + job_id = job_row.job_id, + kind = ?reason.kind, + "Agent reported a failed job: {}", + reason.message + ); + } + + let suite_id = job_row.task_suite_id; + let job_id = job_row.job_id; + + // Release the agent only while it is still assigned to the suite of this job. + // A retirement makes the agent `Offline` without a suite. Without this filter, + // a completion at the same time would make the agent `Idle` again, while + // nothing tracks its heartbeat. + 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; + + // Only this suite. + let reclaimed = + job::reclaim_agent_tasks(txn, agent_uuid, Some(suite_id), now).await?; + + Ok((job_id, terminal, released, reclaimed, suite_id)) + }) + }, + ) + .await?; + + tracing::info!(job = job_handle, job_id, %terminal, "Agent finished a job"); + if !released { + tracing::info!( + agent_id, + job_id, + "Agent was no longer assigned to the suite of the job; left its state alone" + ); + } + + if !reclaimed.is_empty() { + tracing::warn!( + agent_id, + job_id, + reclaimed = reclaimed.len(), + "Agent completed its job with uncommitted tasks; they went back to Ready and run \ + again, and the results of this run are lost" + ); + notify_suite_available(pool, suite_id).await; + } + + matching::agent_has_available_suite(&pool.db, agent_id, &[]).await +} + /// At coordinator start, track the heartbeat of each agent that is not `Offline`, and send it a /// `CounterSync` with the new boot id. /// diff --git a/netmito/src/service/agent/task.rs b/netmito/src/service/agent/task.rs new file mode 100644 index 0000000..501cb88 --- /dev/null +++ b/netmito/src/service/agent/task.rs @@ -0,0 +1,334 @@ +//! Task claims and task reports of agents. +//! +//! ## Claims +//! +//! Agents claim from the database. One transaction selects the `Ready` tasks of the suite with the +//! highest priority `FOR UPDATE SKIP LOCKED` and changes them to `Running`. Because of `SKIP +//! LOCKED`, agents at the same time step over the rows that the others locked. Thus two agents +//! cannot claim the same task, and no agent waits. +//! +//! ## Reports +//! +//! Agents send the same [`ReportTaskOp`] as workers on `POST /tasks/report`, with the job handle. +//! The job must exist, belong to the agent, and not be terminal. A terminal job answers 404: the +//! job is closed and the task is no longer the agent's. +//! +//! | Op | Result | +//! |---|---| +//! | `Finish` | task → `Finished`, not archived yet | +//! | `Cancel` | task → `Cancelled`, the agent sends `Commit` after it | +//! | `Upload` | a presigned S3 url for an artifact. Records the artifact and charges the group | +//! | `Commit` | archives the task with its result, decreases `incomplete_tasks` of the suite, and makes the downstream task `Ready` | +//! | `Submit` | submits a downstream task in the group of the parent, in a suite of that group or in no suite | + +use sea_orm::{prelude::*, QueryOrder, QuerySelect, Set, TransactionTrait}; +use uuid::Uuid; + +use crate::{ + config::InfraPool, + entity::{ + active_tasks as ActiveTasks, archived_tasks as ArchivedTasks, groups as Group, + state::{TaskState, TaskSuiteState}, + task_suites as TaskSuites, StoredTaskModel, + }, + error::{ApiError, Error, Result}, + schema::{ + ExecSpec, FetchTasksResp, ReportTaskOp, TaskExecOptions, WorkerSchedulePlan, WorkerTaskResp, + }, + service::{agent::job, agent::matching, s3::group_upload_artifact}, +}; + +/// The maximum number of tasks in one claim, so that a large `worker_count` cannot lock too many +/// rows. +const MAX_FETCH_BATCH: u32 = 256; + +/// The maximum number of tasks that one job can hold: the tasks that its workers run, and with +/// `prefetch`, one more task for each worker. Then a worker that becomes free gets its next task +/// without a round trip. +pub fn task_budget(schedule: &WorkerSchedulePlan) -> u32 { + match schedule { + WorkerSchedulePlan::FixedWorkers { + worker_count, + prefetch, + .. + } => ((*worker_count).max(1)) + .saturating_mul(1 + u32::from(*prefetch)) + .min(MAX_FETCH_BATCH), + } +} + +/// Claim the ready tasks of a suite, up to the free room of the job of this agent. +/// +/// The agent does not give a number. The schedule of the suite gives the budget of one job, and +/// the tasks that the agent holds in the suite use part of it. +pub async fn agent_fetch_tasks( + agent_id: i64, + agent_uuid: Uuid, + pool: &InfraPool, + suite_uuid: Uuid, +) -> Result { + 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::InvalidRequest(format!( + "Suite {suite_uuid} is {} and gives no more tasks", + suite.state + )))); + } + + let now = TimeDateTimeWithTimeZone::now_utc(); + let suite_id = suite.id; + + // The state comes from the same read as the checks above. If the sweep closed the suite just + // now, a wrong `true` costs only one more fetch. + let hold_job_open = matches!(suite.state, TaskSuiteState::Open); + + let budget = task_budget( + &serde_json::from_value(suite.worker_schedule).inspect_err(|e| { + tracing::error!("Stored worker schedule is unreadable: {e}"); + })?, + ); + let group_name = Group::Entity::find_by_id(suite.group_id) + .one(&pool.db) + .await? + .ok_or_else(|| Error::ApiError(ApiError::NotFound("Group of the suite".to_string())))? + .group_name; + + let claimed = pool + .db + .transaction::<_, Vec, Error>(|txn| { + Box::pin(async move { + // The tasks that the agent holds and has not committed use part of the budget. + let held = ActiveTasks::Entity::find() + .filter(ActiveTasks::Column::TaskSuiteId.eq(suite_id)) + .filter(ActiveTasks::Column::RunnerUuid.eq(agent_uuid)) + .count(txn) + .await?; + let room = (budget as u64).saturating_sub(held); + if room == 0 { + return Ok(Vec::new()); + } + 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(room) + .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, + group_name: group_name.clone(), + spec, + exec_options, + }); + } + + tracing::debug!(count = tasks.len(), "Agent claimed tasks from a suite"); + Ok(FetchTasksResp { + tasks, + hold_job_open, + }) +} + +/// Apply a report of the agent to a task that it 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}"))))?; + + // Only the agent that claimed the task can report it. A reclaimed task has no runner, thus a + // late report from the old agent gets `NotFound` and does not change the new 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"); + // If the user cancelled the task first, the report only confirms it. + 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; + let suite_id = task.task_suite_id; + pool.db + .transaction::<_, (), Error>(|txn| { + Box::pin(async move { + // The suite row before the task rows, the lock order of all writers. + if let Some(suite_id) = suite_id { + crate::service::suite::decrement_incomplete_tasks( + txn, suite_id, 1, now, + ) + .await?; + } + archived.insert(txn).await?; + ActiveTasks::Entity::delete_by_id(inner_task_id) + .exec(txn) + .await?; + Ok(()) + }) + }) + .await?; + + // An earlier `Submit` made a `Pending` downstream task. The parent committed, thus the + // downstream task becomes `Ready`. + 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 it submits a task".to_string(), + ))); + } + // The downstream task stays in the group of the parent, in a suite of that group or in + // no suite. The submit path checks this 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) +} diff --git a/netmito/src/service/mod.rs b/netmito/src/service/mod.rs index c479511..d29dff6 100644 --- a/netmito/src/service/mod.rs +++ b/netmito/src/service/mod.rs @@ -1,14 +1,29 @@ -//! # Lock order: suite before task +//! # Lock order: suite → job → hook → agent → task //! -//! A transaction that writes a `task_suites` row and `active_tasks` rows must write the suite -//! row first. +//! A transaction that locks rows of more than one of `task_suites`, `suite_agent_jobs`, +//! `hook_tasks`, `agents` and `active_tasks` must lock them in this order. A transaction can skip +//! a table, but it must not go back to an earlier table. //! //! An `UPDATE` or `DELETE` keeps its row locks until the transaction commits, not until the //! statement ends. Thus, a write in one statement does not solve the order problem. A //! transaction keeps all of its locks. If two transactions lock the same two rows in opposite //! order, they deadlock. Postgres finds the deadlock only after `deadlock_timeout`. //! -//! There are two correct ways to keep this order: +//! The pairs that occur now: +//! +//! - suite → job → hook → agent: `accept` locks the suite in `create_job`, inserts the job and its +//! hooks, then assigns the agent. +//! - job → hook → agent → task: `complete`, a force stop and a retirement end the job, close its +//! hooks, release the agent, then put its tasks back to `Ready`. No transaction locks a task and +//! then an agent. A commit that also changed its agent would deadlock with `complete`. +//! - suite → task: each commit and each cancel. +//! +//! **A foreign-key write locks the row that it refers to.** An insert of a job, or a change of +//! `agents.assigned_task_suite_id` to a suite, takes `FOR KEY SHARE` on the suite row. This lock +//! conflicts with the `FOR UPDATE` of `create_job`. Thus such a write locks the suite, and it must +//! come before all job, hook, agent and task locks. A change of a reference to NULL checks nothing. +//! +//! There are two correct ways to keep suite → task: //! //! - The suite id is known before the write, for example on each commit and on a single-task //! cancel. Write the suite row first, then the task rows. @@ -16,9 +31,11 @@ //! in the same statement that removes the tasks, and remove a task only after its suite is //! locked. //! -//! This rule does not apply to a transaction that locks task rows and does not write a suite -//! row. `claim_candidates` locks only task rows with `FOR UPDATE SKIP LOCKED`, and no -//! transaction waits for these locks. +//! These do not need the order, because no transaction waits for their locks while they wait: +//! +//! - The task claim of `agent_fetch_tasks` locks only task rows with `FOR UPDATE SKIP LOCKED`. +//! - A single autocommit statement that locks one row, for example the agent update of a +//! heartbeat. pub mod agent; pub mod auth; diff --git a/netmito/src/service/s3.rs b/netmito/src/service/s3.rs index 5221f51..574de17 100644 --- a/netmito/src/service/s3.rs +++ b/netmito/src/service/s3.rs @@ -23,7 +23,8 @@ 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, suite_agent_jobs as SuiteAgentJobs, task_suites as TaskSuites, + user_group as UserGroup, }, error::AuthError, }; @@ -42,7 +43,7 @@ use crate::{ AttachmentsDownloadByFilterReq, AttachmentsDownloadByKeysReq, AttachmentsDownloadListResp, AttachmentsQueryResp, CountQuery, RemoteResourceDownloadResp, }, - service::task::query_tasks_by_filter, + service::{agent::job, task::query_tasks_by_filter}, }; pub(crate) const UPLOAD_VALID_SECS: u64 = 3600; @@ -412,21 +413,10 @@ pub(crate) async fn group_upload_artifact( content_length: u64, ) -> Result<(bool, String), crate::error::Error> { let now = TimeDateTimeWithTimeZone::now_utc(); - let content_length = content_length as i64; let (uuid, group_id) = match task { 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(); // This returns (exist, url) @@ -453,85 +443,122 @@ 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, + &s3_client, + &artifacts_bucket, + group_id, + uuid, + content_type, + content_length, + now, + ) + .await }) }) .await?; Ok(resp) } +/// Record the upload of an artifact in the `artifacts` table, charge the storage of the group, and +/// return `(exist, url)`. `uuid` is the key of the artifact: the uuid of a task or of a hook. +#[allow(clippy::too_many_arguments)] +pub(crate) async fn reserve_artifact_upload( + txn: &C, + s3_client: &Client, + artifacts_bucket: &str, + group_id: i64, + uuid: Uuid, + content_type: ArtifactContentType, + content_length: u64, + now: TimeDateTimeWithTimeZone, +) -> Result<(bool, String), crate::error::Error> { + let content_length = content_length as i64; + // Check if group is active and has enough storage quota + let group = + Group::Entity::find_by_id(group_id) + .one(txn) + .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 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)) +} + pub async fn user_upload_artifact( pool: &InfraPool, user_id: i64, @@ -579,6 +606,47 @@ pub async fn worker_get_attachment( .one(&pool.db) .await? .ok_or(AuthError::PermissionDenied)?; + presign_group_attachment(pool, group_id, group_name, key).await +} + +/// An agent may read a group's attachments only while it runs a job on one of that group's suites. +pub async fn agent_get_attachment( + pool: &InfraPool, + agent_id: i64, + group_name: String, + key: String, +) -> Result { + let group_id = Group::Entity::find() + .filter(Group::Column::GroupName.eq(group_name.clone())) + .one(&pool.db) + .await? + .ok_or(crate::error::ApiError::NotFound(format!( + "Attachment of group {group_name} and key {key}" + )))? + .id; + SuiteAgentJobs::Entity::find() + .select_only() + .column(SuiteAgentJobs::Column::Id) + .join( + sea_orm::JoinType::InnerJoin, + SuiteAgentJobs::Relation::TaskSuites.def(), + ) + .filter(TaskSuites::Column::GroupId.eq(group_id)) + .filter(SuiteAgentJobs::Column::AgentId.eq(agent_id)) + .filter(SuiteAgentJobs::Column::State.is_in(job::IN_FLIGHT)) + .into_tuple::() + .one(&pool.db) + .await? + .ok_or(AuthError::PermissionDenied)?; + presign_group_attachment(pool, group_id, group_name, key).await +} + +async fn presign_group_attachment( + pool: &InfraPool, + group_id: i64, + group_name: String, + 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 8407de6..b554203 100644 --- a/netmito/src/service/suite.rs +++ b/netmito/src/service/suite.rs @@ -2,7 +2,7 @@ use std::collections::HashMap; use sea_orm::sea_query::extension::postgres::PgExpr; use sea_orm::sea_query::{ - Alias, CommonTableExpression, DeleteStatement, InsertStatement, PgFunc, Query, WithClause, + Alias, CommonTableExpression, DeleteStatement, Func, InsertStatement, PgFunc, Query, WithClause, }; use sea_orm::{prelude::*, ConnectionTrait, FromQueryResult, Set, TransactionTrait}; use uuid::Uuid; @@ -10,21 +10,23 @@ 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, ChangeTaskSuiteReq, CountQuery, CreateTaskSuiteReq, CreateTaskSuiteResp, - ExecHooks, ParsedTaskSuiteInfo, SuiteAgentOverrideReq, SuiteAgentOverrideResp, - TaskResultMessage, TaskResultSpec, TaskSuiteInfo, TaskSuiteQueryResp, TaskSuitesQueryReq, - TaskSuitesQueryResp, UpdateOp, WorkerSchedulePlan, + AgentNotification, CancelTaskSuiteOp, ChangeTaskSuiteReq, CountQuery, CreateTaskSuiteReq, + CreateTaskSuiteResp, ExecHooks, HookTaskInfo, ParsedTaskSuiteInfo, StopAgentJobResp, StopJobOp, + SuiteAgentOverrideReq, SuiteAgentOverrideResp, SuiteJobInfo, SuiteJobQueryResp, + SuiteJobsQueryReq, SuiteJobsQueryResp, TaskResultMessage, TaskResultSpec, TaskSuiteInfo, + TaskSuiteQueryResp, TaskSuitesQueryReq, TaskSuitesQueryResp, UpdateOp, WorkerSchedulePlan, }; +use crate::service::agent; use crate::service::task::{check_exec_spec, parse_operators_with_number, OperatorWithNumber}; +use crate::ws::AgentWsRouter; #[derive(FromQueryResult)] struct GroupIdResult { @@ -36,6 +38,77 @@ struct TaskUuidResult { uuid: Uuid, } +/// Move the suites without a new task out of `Open`. +/// +/// A suite stays `Open` for `timeout` after its last submission, also when it has no more tasks. +/// During this time, the agent that runs it keeps its job and can run a late task without a new +/// provision. Then: +/// +/// | from | condition | to | +/// |---|---|---| +/// | `Open` | idle, tasks remain | `Closed` | +/// | `Open` | idle, no task remains | `Complete` | +/// | `Closed` | no task remains | `Complete` | +/// +/// The last rule is a backup: `decrement_incomplete_tasks` makes a `Closed` suite `Complete` when its +/// last task commits. This rule catches the suites that were closed after their last task. +/// +/// The idle time starts at the last submission, or at the creation if the suite has no task. Thus a +/// suite that never got a task also closes. +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 idle. An `Open` suite must be idle for the full window. + 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) +} + /// Decrease `incomplete_tasks` of a suite by `count`, because `count` of its tasks are finished. /// If this makes the counter of a `Closed` suite zero, the suite becomes `Complete`. /// @@ -346,10 +419,10 @@ pub async fn user_change_task_suite( // If the change has new tags or a new priority, tell the agents that the suite is available. // New tags can let an idle agent run this suite. A new priority can change which suite an idle // agent selects or preempt an active agent. - let _needs_notify = req.affects_scheduling(); + let needs_notify = req.affects_scheduling(); let now = TimeDateTimeWithTimeZone::now_utc(); - let _suite_id = pool + let suite_id = pool .db .transaction::<_, i64, Error>(|txn| { Box::pin(async move { @@ -460,8 +533,9 @@ pub async fn user_change_task_suite( // Each query reads the eligibility and the pick order from the row, so there is nothing to // rebuild. A notified agent does a new fetch and gets the latest values. - // TODO: once agents can be notified, `if _needs_notify` call - // `notify_suite_available(pool, _suite_id)` here. + if needs_notify { + agent::notify_suite_available(pool, suite_id).await; + } Ok(()) } @@ -647,16 +721,8 @@ struct SuiteDetailResult { completed_at: Option, } -#[derive(FromQueryResult)] -struct AgentUuidResult { - uuid: Uuid, -} - -/// Get the details of one suite, with the UUIDs of the agents that can run it. The caller must -/// have at least `Read` in the group of the suite. -/// -/// At this time, the result shows only the agents that a user included in the suite. -/// TODO: read the in-memory list to show the agents that match the tags. +/// Get the details of one suite, with the UUIDs of the agents that can run it and the jobs that run +/// it now. The caller must have at least `Read` in the group of the suite. pub async fn user_get_task_suite_by_uuid( user_id: i64, pool: &InfraPool, @@ -724,37 +790,47 @@ 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) + // The tag-matched agents and the agents included by hand, without the agents excluded by hand. + // Refer to `service::agent::matching`. + let eligible_agents = agent::matching::eligible_agent_uuids(&pool.db, suite.id).await?; + + // The jobs of the suite that are not terminal: the agents that run it now. + let jobs_stmt = Query::select() + .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"), + ) + .from(SuiteAgentJobs::Entity) + // `agent_id` is nullable, thus a left join. .join( - sea_orm::JoinType::Join, - TaskSuiteAgent::Entity, - Expr::col((TaskSuiteAgent::Entity, TaskSuiteAgent::Column::AgentId)) - .eq(Expr::col((Agent::Entity, Agent::Column::Id))), + 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((TaskSuiteAgent::Entity, TaskSuiteAgent::Column::TaskSuiteId)).eq(suite.id), + Expr::col((SuiteAgentJobs::Entity, SuiteAgentJobs::Column::TaskSuiteId)).eq(suite.id), ) .and_where( - Expr::col((TaskSuiteAgent::Entity, TaskSuiteAgent::Column::OverrideType)) - .eq(SuiteAgentOverrideType::UserIncluded), + Expr::col((SuiteAgentJobs::Entity, SuiteAgentJobs::Column::State)) + .is_in(agent::job::IN_FLIGHT), + ) + .order_by_expr( + Expr::col((SuiteAgentJobs::Entity, SuiteAgentJobs::Column::JobId)).into(), + sea_orm::Order::Asc, ) .to_owned(); - - let eligible_agents = AgentUuidResult::find_by_statement(builder.build(&agent_stmt)) + let active_jobs = SuiteJobInfo::find_by_statement(builder.build(&jobs_stmt)) .all(&pool.db) - .await? - .into_iter() - .map(|m| m.uuid) - .collect(); + .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()?; @@ -780,6 +856,7 @@ pub async fn user_get_task_suite_by_uuid( completed_at: suite.completed_at, }, eligible_agents, + active_jobs, }) } @@ -846,12 +923,12 @@ 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(); let builder = pool.db.get_database_backend(); - let (_suite_id, _cancelled_task_uuids) = pool + let (suite_id, cancelled_task_uuids) = pool .db .transaction::<_, (i64, Vec), Error>(|txn| { Box::pin(async move { @@ -1020,9 +1097,41 @@ pub async fn user_cancel_task_suite( }) .await?; - // TODO: when agents run suites, stop the jobs of `_suite_id`. `Force` kills the jobs, and - // `Graceful` lets them finish their current work. Send `SuiteCancelled` to each agent, and - // send `TasksCancelled` for `_cancelled_task_uuids`. + // Stop the jobs that still run. `Force` writes the terminal state, thus the coordinator rejects + // later reports of the agents. `Graceful` does not change the jobs: each agent runs the cleanup + // hook and moves its job to `Completed`. + let graceful = matches!(op, CancelTaskSuiteOp::Graceful); + let running_agents = match op { + CancelTaskSuiteOp::Force => { + pool.db + .transaction::<_, Vec, Error>(move |txn| { + Box::pin(async move { agent::job::kill_suite_jobs(txn, suite_id, now).await }) + }) + .await? + } + CancelTaskSuiteOp::Graceful => agent::job::agents_running_suite(&pool.db, suite_id).await?, + }; + + for agent_uuid in running_agents { + AgentWsRouter::notify( + &pool.ws_router_tx, + agent_uuid, + AgentNotification::SuiteCancelled { + suite_uuid, + reason: "Suite was cancelled by a user".to_string(), + graceful, + }, + ); + if !cancelled_task_uuids.is_empty() { + AgentWsRouter::notify( + &pool.ws_router_tx, + agent_uuid, + AgentNotification::TasksCancelled { + task_uuids: cancelled_task_uuids.clone(), + }, + ); + } + } Ok(()) } @@ -1039,9 +1148,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::<_, (i64, 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 = @@ -1072,13 +1181,203 @@ pub async fn user_override_agents_for_suite( } } - Ok(errors) + Ok((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 overrides changed the set of agents that can run the suite. + agent::notify_suite_available(pool, suite_id).await; Ok(SuiteAgentOverrideResp { errors }) } + +/// Query the jobs of a suite, newest first. The caller must have at least `Read` in the group of +/// the suite. +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) + // `agent_id` is nullable, thus a left join. + .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, + }) + } +} + +/// Get one job of a suite with its hooks. The artifacts of a hook use the hook uuid on the artifact +/// endpoints. The caller must have at least `Read` in the group of the suite. +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, + }) +} + +/// Stop one job of a suite, by its number in the suite. The caller must have at least `Write` in +/// the group of the suite. +/// +/// The result is the same as the stop on the agent: the agent stays up and picks a suite again, +/// possibly this suite. +pub async fn user_stop_suite_job( + user_id: i64, + pool: &InfraPool, + suite_uuid: Uuid, + job_id: i32, + op: StopJobOp, +) -> Result { + let suite = match authorize_suite(&pool.db, user_id, suite_uuid, UserGroupRole::Write).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}" + ))))?; + + // A terminal job, or a job without an agent, has nothing to stop. + let agent = match (job.state.is_terminal(), job.agent_id) { + (false, Some(agent_id)) => Agent::Entity::find_by_id(agent_id).one(&pool.db).await?, + _ => None, + }; + let Some(agent) = agent else { + return Ok(StopAgentJobResp { + stopped: false, + suite_uuid: None, + job_id: None, + }); + }; + + agent::stop_agent_job(pool, &agent, &job, suite_uuid, op).await +} diff --git a/netmito/src/service/task.rs b/netmito/src/service/task.rs index 243157a..02c29ef 100644 --- a/netmito/src/service/task.rs +++ b/netmito/src/service/task.rs @@ -403,8 +403,8 @@ 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) => { + crate::service::agent::notify_suite_available(pool, suite.id).await; } None => { let builder = pool.db.get_database_backend(); @@ -496,9 +496,8 @@ pub async fn worker_trigger_pending_task(pool: &InfraPool, uuid: Uuid) -> crate: let task = task.update(&pool.db).await?; // A suite task does not go into the worker queues. Agents get it from its suite. Thus, wake the // idle agents that can run the suite. - if let Some(_suite_id) = task.task_suite_id { - // TODO: when agents can get notifications, call `notify_suite_available(pool, _suite_id)` - // here. + 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 diff --git a/netmito/src/ws/connection.rs b/netmito/src/ws/connection.rs index 950bed4..eb522a1 100644 --- a/netmito/src/ws/connection.rs +++ b/netmito/src/ws/connection.rs @@ -68,14 +68,19 @@ impl AgentSession { } } - fn push(&mut self, event: AgentNotification) -> WsNotificationEvent { + /// Buffer an event. `None` if an event in the buffer makes it unnecessary: then the router + /// does not send it. + fn push(&mut self, event: AgentNotification) -> Option { + if self.buffer.iter().any(|e| e.event.coalesces_with(&event)) { + return None; + } self.counter = self.counter.saturating_add(1); let event = WsNotificationEvent { id: self.counter, event, }; self.buffer.push_back(event.clone()); - event + Some(event) } fn ack(&mut self, ack_id: u64) { @@ -170,7 +175,9 @@ impl AgentWsRouter { } RouterOp::Notify { uuid, event } => { let session = self.sessions.entry(uuid).or_insert_with(AgentSession::new); - let event = session.push(event); + let Some(event) = session.push(event) else { + return; + }; if let Some(sender) = session.sender.as_ref() { Self::push_to_socket(sender, &event, uuid); } diff --git a/openapi.yaml b/openapi.yaml index 22d9a2d..d702ee0 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -496,13 +496,15 @@ paths: summary: Download attachment description: | Get download URL for an attachment. A user must be a member of the group; a - worker must be running a task of the group. + worker must be running a task of the group; an agent must be running a job + on a suite of the group. operationId: downloadAttachment tags: - Groups security: - bearerAuth: [] - workerAuth: [] + - agentAuth: [] parameters: - name: group_name in: path @@ -1303,12 +1305,18 @@ paths: /tasks/report: post: summary: Report task status - description: Worker reports task execution status + description: >- + A worker reports a task that it runs. An agent reports a task or a hook + of its job. `kind` tells which: a worker that sends `Hook` gets 403. An + agent task report must give `job`. A report for a task or a hook that is + no longer the runner's (a reclaimed task, a closed job, a committed + hook) gets 404. operationId: reportTask tags: - Tasks security: - workerAuth: [] + - agentAuth: [] requestBody: required: true content: @@ -1322,8 +1330,14 @@ paths: 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" /tasks/{uuid}: get: @@ -1335,6 +1349,7 @@ paths: security: - bearerAuth: [] - workerAuth: [] + - agentAuth: [] parameters: - name: uuid in: path @@ -1525,6 +1540,7 @@ paths: security: - bearerAuth: [] - workerAuth: [] + - agentAuth: [] parameters: - name: uuid in: path @@ -1868,6 +1884,123 @@ paths: "404": $ref: "#/components/responses/NotFound" + /suites/{uuid}/jobs/query: + post: + summary: Query the jobs of a suite + description: >- + List the jobs of a suite, newest first. A job is one run of the suite + by one agent. The caller must have at least the Read role in the group + of the suite. + operationId: querySuiteJobs + tags: + - Suites + security: + - bearerAuth: [] + parameters: + - name: uuid + in: path + required: true + schema: + type: string + format: uuid + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/SuiteJobsQueryReq" + responses: + "200": + description: The matching jobs + content: + application/json: + schema: + $ref: "#/components/schemas/SuiteJobsQueryResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /suites/{uuid}/jobs/{job_id}: + get: + summary: Get a job of a suite + description: >- + One job of a suite with its hooks. The artifacts of a hook use the hook + uuid on the artifact endpoints. The caller must have at least the Read + role in the group of the suite. + operationId: getSuiteJob + tags: + - Suites + security: + - bearerAuth: [] + parameters: + - name: uuid + in: path + required: true + schema: + type: string + format: uuid + - name: job_id + in: path + required: true + schema: + type: integer + format: int32 + description: The job number in the suite + responses: + "200": + description: The job + content: + application/json: + schema: + $ref: "#/components/schemas/SuiteJobQueryResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + delete: + summary: Stop a job of a suite + description: >- + Stop one job of the suite. The agent stays up and picks a suite again, + possibly this suite. The caller must have at least the Write role in the + group of the suite. A job that is already terminal gives + `stopped: false`. + operationId: stopSuiteJob + tags: + - Suites + security: + - bearerAuth: [] + parameters: + - name: uuid + in: path + required: true + schema: + type: string + format: uuid + - name: job_id + in: path + required: true + schema: + type: integer + format: int32 + description: The job number in the suite + - name: op + in: query + required: false + schema: + $ref: "#/components/schemas/StopJobOp" + responses: + "200": + description: Stop requested, or nothing to stop + content: + application/json: + schema: + $ref: "#/components/schemas/StopAgentJobResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + /agents: post: summary: Register an agent @@ -1994,6 +2127,44 @@ paths: "404": $ref: "#/components/responses/NotFound" + /agents/{uuid}/job: + delete: + summary: Stop the job of an agent + description: >- + Stop the job that the agent runs now. The agent stays up and picks a + suite again, thus a user can move an agent to the suite with the highest + priority now. The caller must have the Admin role in the admin group of + the agent. If the caller does not, the response is 404, the same as for + an unknown uuid. An agent without a job gives `stopped: false`. + operationId: stopAgentJob + tags: + - Agents + security: + - bearerAuth: [] + parameters: + - name: uuid + in: path + required: true + schema: + type: string + format: uuid + - name: op + in: query + required: false + schema: + $ref: "#/components/schemas/StopJobOp" + responses: + "200": + description: Stop requested, or nothing to stop + content: + application/json: + schema: + $ref: "#/components/schemas/StopAgentJobResp" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + /agents/heartbeat: post: summary: Agent heartbeat @@ -2026,6 +2197,106 @@ paths: "401": $ref: "#/components/responses/Unauthorized" + /agents/suite: + post: + summary: Accept a suite + description: >- + Pick a suite for the agent and claim it in one transaction: assign it to + the agent, and create the job and one `Pending` row for each hook of the + suite. `suite_uuid` is a preference: if that suite cannot run now, the + coordinator picks the best available suite (the highest priority, then + the oldest). If no suite is available, or the agent already runs a job, + the response has `accepted: false` and a reason. + operationId: acceptSuite + tags: + - Agents + security: + - agentAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/AcceptSuiteReq" + responses: + "200": + description: The claimed suite, or why nothing was claimed + content: + application/json: + schema: + $ref: "#/components/schemas/AcceptSuiteResp" + "401": + $ref: "#/components/responses/Unauthorized" + + /agents/job: + post: + summary: Move the job to the next phase + description: >- + `Start` (the provision hook is done), `EnterCleanup` (the agent claims no + more tasks), or `Complete` (the job ended; the agent becomes Idle). A + closed job gives 404. A report in the wrong order gives 400. When the + job completes, the hooks without a result are closed and the tasks that + the agent did not commit go back to Ready. + operationId: reportJob + tags: + - Agents + security: + - agentAuth: [] + requestBody: + required: true + content: + application/json: + schema: + $ref: "#/components/schemas/JobReportReq" + responses: + "200": + description: Transition recorded + content: + application/json: + schema: + $ref: "#/components/schemas/JobReportResp" + "400": + $ref: "#/components/responses/BadRequest" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + + /agents/tasks: + get: + summary: Claim tasks of a suite + description: >- + Claim the ready tasks of the suite, highest priority first, up to the + free room of the job. The budget of a job is `worker_count`, two times + that with `prefetch`; the tasks that the agent holds and has not + committed use part of it. An empty list with `hold_job_open` means that + the agent keeps its job and asks again later. + operationId: fetchAgentTasks + tags: + - Agents + security: + - agentAuth: [] + parameters: + - name: suite_uuid + in: query + required: true + schema: + type: string + format: uuid + responses: + "200": + description: The claimed tasks + content: + application/json: + schema: + $ref: "#/components/schemas/FetchTasksResp" + "400": + $ref: "#/components/responses/BadRequest" + "401": + $ref: "#/components/responses/Unauthorized" + "404": + $ref: "#/components/responses/NotFound" + /ws/agents: get: summary: Agent notification socket @@ -3834,12 +4105,23 @@ components: type: object required: - id + - kind - op properties: id: type: integer format: int64 - description: Internal task ID + description: The internal id of the task, or of the hook if `kind` is `Hook` + kind: + type: string + enum: + - Task + - Hook + description: The kind of task. Only agents report hooks + job: + type: integer + format: int64 + description: The job handle of an agent task. Workers and hooks do not send it op: $ref: "#/components/schemas/ReportTaskOp" @@ -4583,12 +4865,13 @@ components: example: 8 cpu_binding: $ref: "#/components/schemas/CpuBinding" - task_prefetch_count: - type: integer - format: uint32 - default: 16 - description: How many tasks each worker prefetches locally - example: 16 + prefetch: + type: boolean + default: true + description: >- + Whether the agent can hold one more claimed task for each worker, so + that a worker that becomes free gets its next task without a round + trip. Turn it off only for tasks that run much longer than a round trip ExecHooks: type: object @@ -4961,16 +5244,24 @@ components: type: object required: - info - - assigned_agents + - eligible_agents + - active_jobs properties: info: $ref: "#/components/schemas/ParsedTaskSuiteInfo" - assigned_agents: + eligible_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) + description: >- + The agents that match the tags or are included by hand, without the + agents that are excluded by hand + active_jobs: + type: array + items: + $ref: "#/components/schemas/SuiteJobInfo" + description: The jobs of the suite that are not terminal, oldest first CancelSuiteResp: type: object @@ -5230,11 +5521,24 @@ components: enum: - shutdown - counter_sync + - suite_available + - preempt_suite + - suite_cancelled + - tasks_cancelled + - stop_job description: | A message from the coordinator that tells an agent to do something: - shutdown: shut down; with `graceful`, finish the current work first - counter_sync: the coordinator restarted and its notification ids start again from `counter`; `boot_id` identifies the new process + - suite_available: a suite that the agent can run has work; the agent + accepts a suite to learn which + - preempt_suite: not sent yet + - suite_cancelled: the suite of `suite_uuid` was cancelled; with + `graceful`, claim no more tasks and run the cleanup hook + - tasks_cancelled: stop the tasks of `task_uuids` if they run + - stop_job: stop the job on `suite_uuid` and pick a suite again; with + `graceful`, finish the running tasks and run the cleanup hook first WsNotificationEvent: type: object @@ -5267,6 +5571,342 @@ components: `notification_id` - pong: the reply to a ping, with the agent clock in `client_time` + StopJobOp: + type: string + enum: + - graceful + - force + default: graceful + description: | + How to stop a job. The agent stays up in both cases: + - graceful: claim no more tasks, let the running tasks finish, then run + the cleanup hook. The agent moves the job to Completed + - force: the job becomes Killed now and its tasks go back to Ready. The + cleanup hook does not run + + StopAgentJobResp: + type: object + required: + - stopped + properties: + stopped: + type: boolean + description: False if there was no job to stop + suite_uuid: + type: string + format: uuid + job_id: + type: integer + format: int32 + description: The number of the stopped job in its suite + + SuiteJobState: + type: string + enum: + - Provisioning + - Executing + - Cleanup + - Completed + - Failed + - Lost + - Killed + + HookType: + type: string + enum: + - Provision + - Cleanup + - Background + + HookExecState: + type: string + enum: + - Completed + - Failed + - Cancelled + - Pending + - Finished + description: | + - Pending: not reported yet + - Finished: the process stopped, the result is not committed yet + - Completed, Failed, Cancelled: the final state. When the job ends, a + Pending hook becomes Cancelled and a Finished hook becomes Failed + + AcceptSuiteReq: + type: object + properties: + suite_uuid: + type: string + format: uuid + description: The suite that the agent prefers + + TaskSuiteSpec: + type: object + required: + - uuid + - group_name + - tags + - labels + - priority + - worker_schedule + - state + - total_tasks + - incomplete_tasks + properties: + uuid: + type: string + format: uuid + name: + type: string + description: + type: string + group_name: + 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" + total_tasks: + type: integer + format: int32 + incomplete_tasks: + type: integer + format: int32 + + JobHook: + type: object + required: + - id + - uuid + - hook_type + properties: + id: + type: integer + format: int64 + description: The id to report the hook with on /tasks/report + uuid: + type: string + format: uuid + description: The key of the hook artifacts + hook_type: + $ref: "#/components/schemas/HookType" + + AcceptSuiteResp: + type: object + required: + - accepted + properties: + accepted: + type: boolean + suite: + $ref: "#/components/schemas/TaskSuiteSpec" + job: + type: integer + format: int64 + description: The job handle. The agent sends it on each later call for the job + job_id: + type: integer + format: int32 + description: The number of the job in its suite + hooks: + type: array + items: + $ref: "#/components/schemas/JobHook" + reason: + type: string + description: Why nothing was claimed, if `accepted` is false + + JobFailureReason: + type: object + required: + - kind + - message + properties: + kind: + type: string + enum: + - ProvisionFailed + - BackgroundExited + - ExecutionError + - CleanupFailed + message: + type: string + + JobReportReq: + type: object + required: + - job + - op + properties: + job: + type: integer + format: int64 + description: The job handle from AcceptSuiteResp + op: + oneOf: + - type: string + enum: + - Start + - type: string + enum: + - EnterCleanup + - type: object + required: + - Complete + properties: + Complete: + type: object + required: + - outcome + properties: + outcome: + oneOf: + - type: string + enum: + - Completed + - type: object + required: + - Failed + properties: + Failed: + type: object + required: + - reason + properties: + reason: + $ref: "#/components/schemas/JobFailureReason" + + JobReportResp: + type: object + properties: + next_suite_available: + type: boolean + description: For Complete only, whether the agent can accept another suite now + + FetchTasksResp: + type: object + required: + - tasks + properties: + tasks: + type: array + items: + $ref: "#/components/schemas/WorkerTaskResp" + hold_job_open: + type: boolean + description: >- + True while the suite is Open. With an empty `tasks`, the agent keeps + its job and asks again later + + SuiteJobsQueryReq: + type: object + properties: + states: + type: array + items: + $ref: "#/components/schemas/SuiteJobState" + agent_uuid: + type: string + format: uuid + limit: + type: integer + format: uint64 + offset: + type: integer + format: uint64 + count: + type: boolean + default: false + description: Return only the number of matching jobs + + SuiteJobInfo: + type: object + required: + - job_id + - state + - created_at + - updated_at + properties: + job_id: + type: integer + format: int32 + description: The number of the job in its suite + 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 + description: The key of the hook artifacts + hook_type: + $ref: "#/components/schemas/HookType" + state: + $ref: "#/components/schemas/HookExecState" + result: + $ref: "#/components/schemas/TaskResultSpec" + started_at: + type: string + format: date-time + completed_at: + type: string + format: date-time + + SuiteJobQueryResp: + type: object + required: + - info + - hooks + properties: + info: + $ref: "#/components/schemas/SuiteJobInfo" + hooks: + type: array + items: + $ref: "#/components/schemas/HookTaskInfo" + security: - bearerAuth: [] @@ -5284,7 +5924,7 @@ tags: - name: Suites description: Task suite lifecycle and manual agent assignment - name: Agents - description: Agent registration, heartbeat and shutdown + description: Agent registration, heartbeat, shutdown and the job loop - name: Admin description: Administrative operations - name: Health From c822d714863563068027aaf9c47286db4853cad4 Mon Sep 17 00:00:00 2001 From: un-lock-able Date: Fri, 30 Oct 2026 17:12:16 +0800 Subject: [PATCH 2/2] feat(agent): offer new suite tasks to running jobs first --- netmito/src/api/agents.rs | 2 +- netmito/src/config/coordinator.rs | 21 + netmito/src/coordinator.rs | 25 +- netmito/src/schema/agent.rs | 12 +- netmito/src/service/agent/matching.rs | 7 +- netmito/src/service/agent/mod.rs | 187 +++++- netmito/src/service/agent/queue.rs | 831 ++++++++++++++++++++++++++ netmito/src/service/agent/task.rs | 113 ++-- netmito/src/service/suite.rs | 8 +- netmito/src/service/task.rs | 214 +++++-- openapi.yaml | 3 + 11 files changed, 1282 insertions(+), 141 deletions(-) create mode 100644 netmito/src/service/agent/queue.rs diff --git a/netmito/src/api/agents.rs b/netmito/src/api/agents.rs index 50282e6..1633dd0 100644 --- a/netmito/src/api/agents.rs +++ b/netmito/src/api/agents.rs @@ -143,7 +143,7 @@ async fn accept_suite( State(pool): State, Json(req): Json, ) -> Result, ApiError> { - let resp = service::agent::agent_accept_suite(a.id, &pool, req) + let resp = service::agent::agent_accept_suite(a.id, a.uuid, &pool, req) .await .map_err(map_service_error)?; Ok(Json(resp)) diff --git a/netmito/src/config/coordinator.rs b/netmito/src/config/coordinator.rs index 55daa36..a7c562f 100644 --- a/netmito/src/config/coordinator.rs +++ b/netmito/src/config/coordinator.rs @@ -24,6 +24,7 @@ use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, Layer}; use crate::{ error::Error, service::agent::heartbeat::{AgentHeartbeatOp, AgentHeartbeatQueue}, + service::agent::queue::SuiteQueues, service::worker::{HeartbeatOp, TaskDispatcher, TaskDispatcherOp, WorkerHeartbeatQueue}, ws::{AgentWsRouter, RouterOp}, }; @@ -80,6 +81,13 @@ pub struct CoordinatorConfig { /// value gives fewer provision hooks and more idle time of the machines. #[serde(with = "humantime_serde", default = "default_suite_auto_close_timeout")] pub(crate) suite_auto_close_timeout: std::time::Duration, + /// The interval of the check of the in-memory suite queues against the database. A queued task + /// that two checks in sequence do not find as ready goes out of the queue. + #[serde( + with = "humantime_serde", + default = "default_suite_queue_reconcile_interval" + )] + pub(crate) suite_queue_reconcile_interval: std::time::Duration, pub(crate) log_path: Option, pub(crate) file_log: bool, } @@ -96,6 +104,10 @@ fn default_suite_auto_close_timeout() -> std::time::Duration { std::time::Duration::from_secs(60) } +fn default_suite_queue_reconcile_interval() -> std::time::Duration { + std::time::Duration::from_secs(60) +} + fn default_db_min_connections() -> u32 { 8 } @@ -202,6 +214,11 @@ pub struct CoordinatorConfigCli { #[arg(long)] #[serde(skip_serializing_if = "::std::option::Option::is_none")] pub suite_auto_close_timeout: Option, + /// The interval of the check of the in-memory suite queues against the database. Default 60 + /// seconds + #[arg(long)] + #[serde(skip_serializing_if = "::std::option::Option::is_none")] + pub suite_queue_reconcile_interval: Option, /// The log file path. If not specified, then the default rolling log file path would be used. /// If specified, then the log file would be exactly at the path specified. #[arg(long)] @@ -247,6 +264,7 @@ impl Default for CoordinatorConfig { heartbeat_timeout: std::time::Duration::from_secs(600), ws_keepalive_interval: default_ws_keepalive_interval(), suite_auto_close_timeout: default_suite_auto_close_timeout(), + suite_queue_reconcile_interval: default_suite_queue_reconcile_interval(), log_path: None, file_log: false, } @@ -403,6 +421,7 @@ impl CoordinatorConfig { worker_heartbeat_queue_tx, agent_heartbeat_queue_tx, ws_router_tx, + suite_queues: SuiteQueues::new(), boot_uuid: uuid::Uuid::new_v4(), ws_keepalive_interval: self.ws_keepalive_interval, }) @@ -525,6 +544,8 @@ pub struct InfraPool { pub worker_heartbeat_queue_tx: MTx, pub agent_heartbeat_queue_tx: MTx, pub ws_router_tx: MTx, + /// The dispatch state of the suites that have a job in flight. + pub suite_queues: SuiteQueues, /// Identifies this coordinator process. An agent compares it with the /// `boot_id` of a `CounterSync` to find out that the coordinator restarted. pub boot_uuid: uuid::Uuid, diff --git a/netmito/src/coordinator.rs b/netmito/src/coordinator.rs index 720789d..d9cf8ec 100644 --- a/netmito/src/coordinator.rs +++ b/netmito/src/coordinator.rs @@ -9,7 +9,7 @@ use tracing_subscriber::layer::SubscriberExt; use crate::api::router; use crate::config::{CoordinatorConfig, CoordinatorConfigCli, InfraPool}; use crate::migration::{Migrator, MigratorTrait}; -use crate::service::agent::{heartbeat::AgentHeartbeatQueue, restore_agents}; +use crate::service::agent::{heartbeat::AgentHeartbeatQueue, queue, restore_agents}; use crate::service::s3::setup_buckets; use crate::service::suite::sweep_inactive_suites; use crate::service::worker::{restore_workers, TaskDispatcher, WorkerHeartbeatQueue}; @@ -31,6 +31,8 @@ pub struct MitoCoordinator { pub log_dir: PathBuf, /// How long a suite can stay without a new task before the sweep moves it out of `Open`. pub suite_auto_close_timeout: Duration, + /// The interval of the check of the in-memory suite queues against the database. + pub suite_queue_reconcile_interval: Duration, } impl MitoCoordinator { @@ -148,6 +150,7 @@ impl MitoCoordinator { ); let ws_router = config.build_ws_router(cancel_token.clone(), ws_router_rx); let suite_auto_close_timeout = config.suite_auto_close_timeout; + let suite_queue_reconcile_interval = config.suite_queue_reconcile_interval; // Setup s3 storage // List all buckets and create if not exist @@ -179,6 +182,7 @@ impl MitoCoordinator { cancel_token, log_dir, suite_auto_close_timeout, + suite_queue_reconcile_interval, }) } @@ -192,23 +196,28 @@ impl MitoCoordinator { mut ws_router, cancel_token, suite_auto_close_timeout, + suite_queue_reconcile_interval, .. } = self; // Create TaskTracker to manage background tasks let task_tracker = TaskTracker::new(); - // Suite housekeeping: move the idle suites out of `Open`. + // Suite housekeeping: move the idle suites out of `Open`, and remove from the in-memory + // queues the tasks that are not ready. { let pool = infra_pool.clone(); let cancel = cancel_token.clone(); let mut sweep = tokio::time::interval(suite_sweep_period(suite_auto_close_timeout)); + let mut reconcile = tokio::time::interval(suite_queue_reconcile_interval); // A late tick is not done at once after the work that delayed it. There is no need to // catch up. sweep.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + reconcile.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); task_tracker.spawn(async move { - // The first tick is immediate: it moves the suites that became idle while the - // coordinator was down. + // The first ticks are immediate. The first sweep moves the suites that became idle + // while the coordinator was down. The first reconcile removes nothing, because a + // task goes only after two reconciles. loop { tokio::select! { biased; @@ -218,6 +227,11 @@ impl MitoCoordinator { tracing::error!("Failed to sweep inactive suites: {e}"); } } + _ = reconcile.tick() => { + if let Err(e) = queue::reconcile(&pool).await { + tracing::error!("Failed to reconcile the suite queues: {e}"); + } + } } } tracing::info!("Suite housekeeping stopped"); @@ -244,6 +258,9 @@ impl MitoCoordinator { }); restore_workers(&infra_pool).await?; + // The jobs of agents continue through a coordinator restart. Add them to the suite queues + // before the first submit. + queue::restore(&infra_pool).await?; restore_agents(&infra_pool).await?; let app = router(infra_pool, cancel_token.clone()); let addr = crate::config::SERVER_CONFIG diff --git a/netmito/src/schema/agent.rs b/netmito/src/schema/agent.rs index 372062a..c8636e4 100644 --- a/netmito/src/schema/agent.rs +++ b/netmito/src/schema/agent.rs @@ -418,6 +418,10 @@ pub enum AgentNotification { /// Stop the job on `suite_uuid`. The agent stays up and picks a suite /// again. An agent that runs a different suite ignores the event. StopJob { suite_uuid: Uuid, graceful: bool }, + /// The job of the agent on `suite_uuid` has ready tasks. Fetch now, and do not wait for the + /// next poll. For a short time, the suite is reserved for the notified jobs. An agent that runs + /// a different suite ignores the event. + TasksAvailable { suite_uuid: Uuid }, } impl AgentNotification { @@ -427,7 +431,13 @@ impl AgentNotification { /// state over HTTP, thus a second copy gives it nothing. A burst of /// submissions also must not push older events out of the buffer. pub fn coalesces_with(&self, other: &Self) -> bool { - matches!((self, other), (Self::SuiteAvailable, Self::SuiteAvailable)) + match (self, other) { + (Self::SuiteAvailable, Self::SuiteAvailable) => true, + (Self::TasksAvailable { suite_uuid: a }, Self::TasksAvailable { suite_uuid: b }) => { + a == b + } + _ => false, + } } } diff --git a/netmito/src/service/agent/matching.rs b/netmito/src/service/agent/matching.rs index 2b8ceda..58adc47 100644 --- a/netmito/src/service/agent/matching.rs +++ b/netmito/src/service/agent/matching.rs @@ -163,8 +163,11 @@ pub async fn agent_has_available_suite( /// The suite that this agent should run next: the highest priority first, then the oldest. `None` /// if no suite is available. /// -/// The query excludes the suites in `blocked`. A filter after the query would be wrong: a blocked -/// suite at the top would hide the next suite. +/// The query excludes the suites in `blocked` ([`SuiteQueues::blocked_for`]: the suites reserved +/// for their jobs). A filter after the query would be wrong: a blocked suite at the top would hide +/// the next suite. +/// +/// [`SuiteQueues::blocked_for`]: crate::service::agent::queue::SuiteQueues::blocked_for pub async fn best_available_suite_id( db: &C, agent_id: i64, diff --git a/netmito/src/service/agent/mod.rs b/netmito/src/service/agent/mod.rs index 1e18da6..88d47f0 100644 --- a/netmito/src/service/agent/mod.rs +++ b/netmito/src/service/agent/mod.rs @@ -12,6 +12,14 @@ //! the agent, and creates a `suite_agent_jobs` row and its hook rows. The agent sends the job //! handle on each later call. `start`, `cleanup` and `complete` move the job row forward. //! +//! ## How work is offered +//! +//! When a suite gets ready tasks, [`notify_suite_available`] first notifies the jobs on the suite. +//! For a short time, the suite is reserved for these jobs. Idle agents get a notification only for +//! the tasks that these jobs cannot take, or for a suite without a job. This notification does not +//! name the suite: [`agent_accept_suite`] picks the best suite when the agent accepts one. +//! [`queue`] keeps the in-memory state for this. +//! //! ## No automatic preemption //! //! An agent runs the suite that it accepted until the job ends. The priority of a suite only @@ -23,9 +31,10 @@ pub mod heartbeat; pub mod hook; pub mod job; pub mod matching; +pub mod queue; pub mod task; -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use sea_orm::sea_query::extension::postgres::PgExpr; use sea_orm::sea_query::{Alias, PgFunc, Query}; @@ -52,6 +61,7 @@ use crate::service::auth::token::{generate_runner_token, TokenRole}; use crate::ws::AgentWsRouter; use heartbeat::AgentHeartbeatOp; +use queue::SuiteQueues; /// Register an agent. If the machine is already known, use its agent row again. /// @@ -84,6 +94,7 @@ pub async fn user_register_agent( .map(serde_json::to_value) .transpose()?; + let queues = pool.suite_queues.clone(); let (agent_uuid, agent_id, reused, retired) = pool .db .transaction::<_, (Uuid, i64, bool, Option), Error>(move |txn| { @@ -167,6 +178,7 @@ pub async fn user_register_agent( retired = Some( retire_agent_writes( txn, + &queues, agent.id, agent.uuid, RetireCause::Reregistered, @@ -524,12 +536,13 @@ pub async fn retire_agent(pool: &InfraPool, agent_id: i64, cause: RetireCause) - } let agent_uuid = agent.uuid; + let queues = pool.suite_queues.clone(); let retired = pool .db .transaction::<_, RetiredWork, Error>(move |txn| { - Box::pin( - async move { retire_agent_writes(txn, agent_id, agent_uuid, cause, now).await }, - ) + Box::pin(async move { + retire_agent_writes(txn, &queues, agent_id, agent_uuid, cause, now).await + }) }) .await?; @@ -545,6 +558,7 @@ type RetiredWork = (Vec, Vec<(i64, Option, i32)>); /// because it retires the previous agent and updates the agent row in one step. async fn retire_agent_writes( txn: &C, + queues: &SuiteQueues, agent_id: i64, agent_uuid: Uuid, cause: RetireCause, @@ -563,6 +577,8 @@ async fn retire_agent_writes( // All the suites: the agent leaves each suite at the same time. let reclaimed = job::reclaim_agent_tasks(txn, agent_uuid, None, now).await?; + // Add the tasks before the commit. See `SuiteQueues::push_ready`. + queues.push_reclaimed(&reclaimed); Ok((suite_ids, reclaimed)) } @@ -576,6 +592,16 @@ async fn retire_agent_post_commit( cause: RetireCause, (suite_ids, reclaimed): RetiredWork, ) { + // Count the returned tasks for each suite. `close` compares this count with the count of the + // queue. + let mut by_suite: HashMap = HashMap::new(); + for suite_id in reclaimed.iter().filter_map(|(_, suite_id, _)| *suite_id) { + *by_suite.entry(suite_id).or_default() += 1; + } + for suite_id in &suite_ids { + let count = by_suite.get(suite_id).copied().unwrap_or(0); + pool.suite_queues.close(*suite_id, agent_id, count); + } let _ = pool .agent_heartbeat_queue_tx .send(AgentHeartbeatOp::Remove(agent_id)); @@ -592,11 +618,7 @@ async fn retire_agent_post_commit( // Offer the work again for each suite. This includes the suites whose jobs held no task, // because the end of a job can let a different agent take the suite. let mut suites: HashSet = suite_ids.into_iter().collect(); - suites.extend( - reclaimed - .into_iter() - .filter_map(|(_, suite_id, _)| suite_id), - ); + suites.extend(by_suite.into_keys()); for suite_id in suites { notify_suite_available(pool, suite_id).await; } @@ -669,6 +691,7 @@ pub(crate) async fn stop_agent_job( let (job_handle, job_id) = (job.id, job.job_id); if op == StopJobOp::Force { + let queues = pool.suite_queues.clone(); let stopped = pool .db .transaction::<_, Option, i32)>>, Error>(move |txn| { @@ -709,6 +732,8 @@ pub(crate) async fn stop_agent_job( // `Complete`. let reclaimed = job::reclaim_agent_tasks(txn, agent_uuid, Some(suite_id), now).await?; + // Add the tasks before the commit. See `SuiteQueues::push_ready`. + queues.push_reclaimed(&reclaimed); Ok(Some(reclaimed)) }) @@ -723,8 +748,10 @@ pub(crate) async fn stop_agent_job( }); }; + let suite_id = job.task_suite_id; + pool.suite_queues.close(suite_id, agent_id, reclaimed.len()); if !reclaimed.is_empty() { - notify_suite_available(pool, job.task_suite_id).await; + notify_suite_available(pool, suite_id).await; } } @@ -791,7 +818,12 @@ pub async fn agent_heartbeat( // lost the earlier one. if req.state == AgentState::Idle && agent.assigned_task_suite_id.is_none() - && matching::agent_has_available_suite(&pool.db, agent_id, &[]).await? + && matching::agent_has_available_suite( + &pool.db, + agent_id, + &pool.suite_queues.blocked_for(agent_id), + ) + .await? { AgentWsRouter::notify( &pool.ws_router_tx, @@ -831,10 +863,57 @@ pub async fn agent_heartbeat( Ok(AgentHeartbeatResp { notifications }) } -/// Tell the `Idle` agents that can run a suite that it has work. +/// Offer the ready tasks of a suite. First notify the jobs on the suite, then the `Idle` agents for +/// the tasks that these jobs cannot take. +/// +/// A job on the suite is already provisioned. Thus it gets `TasksAvailable` before an idle agent +/// runs a provision hook. [`SuiteQueues::follow_up_offer`] selects the jobs. If they have room for +/// all the ready tasks, the suite is reserved for them ([`queue::RESERVATION_WINDOW`]), and no idle +/// agent gets a notification. If not, the reservation ends, and the idle agents get +/// `SuiteAvailable`. /// -/// The event does not name the suite. The agent gets the best suite for it when it accepts one. +/// `SuiteAvailable` does not name the suite. The agent gets the best suite for it when it accepts +/// one. pub async fn notify_suite_available(pool: &InfraPool, suite_id: i64) { + let offer = pool.suite_queues.follow_up_offer(suite_id); + if !offer.agents.is_empty() { + let suite_uuid = match TaskSuites::Entity::find_by_id(suite_id).one(&pool.db).await { + Ok(Some(suite)) => suite.uuid, + Ok(None) => return, + Err(e) => { + tracing::error!(suite_id, "Failed to load suite for agent notification: {e}"); + return; + } + }; + if offer.covered { + pool.suite_queues.reserve( + suite_id, + offer.agents.iter().map(|(agent_id, _)| *agent_id).collect(), + queue::RESERVATION_WINDOW, + ); + } + tracing::debug!( + suite_id, + jobs = offer.agents.len(), + covered = offer.covered, + "Offering the ready tasks of the suite to its jobs" + ); + for (_, agent_uuid) in &offer.agents { + AgentWsRouter::notify( + &pool.ws_router_tx, + *agent_uuid, + AgentNotification::TasksAvailable { suite_uuid }, + ); + } + if offer.covered { + return; + } + } + + // The jobs cannot take all the tasks, or the suite has no job. Thus a new job is necessary. End + // the reservation of an earlier, smaller offer. + pool.suite_queues.release_reservation(suite_id); + let agents = match matching::idle_eligible_agent_uuids(&pool.db, suite_id).await { Ok(agents) => agents, Err(e) => { @@ -889,14 +968,18 @@ async fn suite_to_spec( /// response has `accepted: false` and a reason. This is not an HTTP error. pub async fn agent_accept_suite( agent_id: i64, + agent_uuid: Uuid, pool: &InfraPool, req: AcceptSuiteReq, ) -> Result { let now = TimeDateTimeWithTimeZone::now_utc(); + // Read this before the transaction. A reservation is only a preference, thus a value from a + // short time before is sufficient. + let blocked = pool.suite_queues.blocked_for(agent_id); let outcome = pool .db - .transaction::<_, std::result::Result<(TaskSuiteSpec, i64, i32, Vec), String>, Error>( + .transaction::<_, std::result::Result<(TaskSuiteSpec, i64, i32, i64, Vec), String>, Error>( |txn| { Box::pin(async move { // One suite at a time. A busy agent that asks again is not an error. @@ -914,7 +997,8 @@ pub async fn agent_accept_suite( // The suite that the agent asked for, if it can still run. let requested = match req.suite_uuid { Some(suite_uuid) => { - runnable_suite(txn, agent_id, SuiteTarget::Uuid(suite_uuid)).await? + runnable_suite(txn, agent_id, SuiteTarget::Uuid(suite_uuid), &blocked) + .await? } None => None, }; @@ -922,14 +1006,22 @@ pub async fn agent_accept_suite( let suite = match requested { Some(suite) => suite, None => { - let picked = - match matching::best_available_suite_id(txn, agent_id, &[]).await? { - Some(suite_id) => { - runnable_suite(txn, agent_id, SuiteTarget::Id(suite_id)) - .await? - } - None => None, - }; + let picked = match matching::best_available_suite_id( + txn, agent_id, &blocked, + ) + .await? + { + Some(suite_id) => { + runnable_suite( + txn, + agent_id, + SuiteTarget::Id(suite_id), + &blocked, + ) + .await? + } + None => None, + }; // `None` after a pick: the suite lost its work between the pick and // the check. let Some(suite) = picked else { @@ -951,17 +1043,41 @@ pub async fn agent_accept_suite( agent.updated_at = Set(now); agent.update(txn).await?; - Ok(Ok((spec, job.id, job.job_id, hooks))) + Ok(Ok((spec, job.id, job.job_id, suite_id, hooks))) }) }, ) .await?; match outcome { - Ok((suite, job, job_id, hooks)) => { + Ok((suite, job, job_id, suite_id, hooks)) => { let span = tracing::Span::current(); span.record("suite_uuid", tracing::field::display(suite.uuid)); span.record("job_id", job_id); + + // The first job of the suite creates the queue entry. Read the ready tasks from the + // database, because the queue ignored them before the entry existed. + if pool.suite_queues.open( + suite_id, + agent_id, + agent_uuid, + job_id, + queue::task_budget(&suite.worker_schedule), + 0, + ) { + match queue::ready_tasks(&pool.db, vec![suite_id]).await { + Ok(tasks) => pool.suite_queues.push_ready( + suite_id, + tasks.into_iter().map(|(_, id, priority)| (id, priority)), + ), + // The job can still claim tasks, because a claim reads the database. Only the + // offers count too few ready tasks. + Err(e) => tracing::error!( + suite_id, + "Failed to read the ready tasks of a suite for its first job: {e}" + ), + } + } Ok(AcceptSuiteResp { accepted: true, suite: Some(suite), @@ -989,8 +1105,9 @@ enum SuiteTarget { } /// Return the suite if this agent can run it now. `None` for each reason that it cannot: missing, -/// not for this agent, in the wrong state, or no ready task. The caller then tries the next -/// candidate. +/// not for this agent, in the wrong state, no ready task, or reserved for other agents (`blocked`). +/// The caller then tries the next candidate. `best_available_suite_id` skips a reserved suite, but +/// an agent can still ask for one by its UUID. /// /// This function takes no lock. A lock taken before the checks stays when the function returns /// `None`. Then an agent whose preferred suite cannot run holds that suite while it locks a @@ -1000,6 +1117,7 @@ async fn runnable_suite( txn: &C, agent_id: i64, target: SuiteTarget, + blocked: &[i64], ) -> Result> { let query = TaskSuites::Entity::find(); let suite = match target { @@ -1012,6 +1130,9 @@ async fn runnable_suite( let Some(suite) = suite else { return Ok(None); }; + if blocked.contains(&suite.id) { + return Ok(None); + } // The same rule as `matching::suite_has_work`, which the pick uses. if !matches!(suite.state, TaskSuiteState::Open | TaskSuiteState::Closed) { return Ok(None); @@ -1120,6 +1241,7 @@ pub async fn agent_complete_job( ) -> Result { let now = TimeDateTimeWithTimeZone::now_utc(); + let queues = pool.suite_queues.clone(); let (job_id, terminal, released, reclaimed, suite_id) = pool .db .transaction::<_, (i32, SuiteJobState, bool, Vec<(i64, Option, i32)>, i64), Error>( @@ -1180,6 +1302,8 @@ pub async fn agent_complete_job( // Only this suite. let reclaimed = job::reclaim_agent_tasks(txn, agent_uuid, Some(suite_id), now).await?; + // Add the tasks before the commit. See `SuiteQueues::push_ready`. + queues.push_reclaimed(&reclaimed); Ok((job_id, terminal, released, reclaimed, suite_id)) }) @@ -1195,6 +1319,8 @@ pub async fn agent_complete_job( "Agent was no longer assigned to the suite of the job; left its state alone" ); } + // Remove the ended job, thus the suite does not offer tasks to it or keep a reservation for it. + pool.suite_queues.close(suite_id, agent_id, reclaimed.len()); if !reclaimed.is_empty() { tracing::warn!( @@ -1207,7 +1333,12 @@ pub async fn agent_complete_job( notify_suite_available(pool, suite_id).await; } - matching::agent_has_available_suite(&pool.db, agent_id, &[]).await + matching::agent_has_available_suite( + &pool.db, + agent_id, + &pool.suite_queues.blocked_for(agent_id), + ) + .await } /// At coordinator start, track the heartbeat of each agent that is not `Offline`, and send it a diff --git a/netmito/src/service/agent/queue.rs b/netmito/src/service/agent/queue.rs new file mode 100644 index 0000000..2db6dbd --- /dev/null +++ b/netmito/src/service/agent/queue.rs @@ -0,0 +1,831 @@ +//! The in-memory dispatch state of the coordinator. It has one entry for each suite that has a job +//! in flight. +//! +//! An entry has: +//! +//! - the jobs on the suite, by agent. For each job: its budget and the number of tasks that it +//! holds. The entry goes away when its last job ends. +//! - the ids of the ready tasks of the suite, in the order that a claim takes them. +//! - an optional reservation. While the reservation is active, only the agents that run the suite +//! can accept the suite or claim from it. +//! +//! The database decides each claim. This state only decides which agents get a notification, and +//! it counts the ready tasks of each suite. + +use std::{ + cmp::Ordering, + collections::{HashMap, HashSet}, + sync::{Arc, Mutex, MutexGuard}, + time::{Duration, Instant}, +}; + +use sea_orm::{prelude::*, QuerySelect}; +use uuid::Uuid; + +use crate::{ + config::InfraPool, + entity::{ + active_tasks as ActiveTasks, agents as Agent, state::TaskState, + suite_agent_jobs as SuiteAgentJobs, task_suites as TaskSuites, + }, + schema::WorkerSchedulePlan, + service::agent::job, +}; + +/// The maximum time that a suite stays reserved for its jobs. Usually a reservation ends earlier, +/// when the notified jobs fetch. This limit applies when a notified job does not fetch, for example +/// because its agent stopped. +pub const RESERVATION_WINDOW: Duration = Duration::from_secs(30); + +/// The maximum number of tasks in one claim, so that a large `worker_count` cannot lock too many +/// rows. +const MAX_FETCH_BATCH: u32 = 256; + +/// The maximum number of tasks that one job can hold: the tasks that its workers run, and with +/// `prefetch`, one more task for each worker. Then a worker that becomes free gets its next task +/// without a round trip. +pub fn task_budget(schedule: &WorkerSchedulePlan) -> u32 { + match schedule { + WorkerSchedulePlan::FixedWorkers { + worker_count, + prefetch, + .. + } => ((*worker_count).max(1)) + .saturating_mul(1 + u32::from(*prefetch)) + .min(MAX_FETCH_BATCH), + } +} + +/// The dispatch state of each suite that has a job in flight. All clones share one state. +#[derive(Debug, Clone, Default)] +pub struct SuiteQueues { + inner: Arc>>, +} + +#[derive(Debug, Default)] +struct SuiteQueue { + /// The jobs in flight on this suite, by agent. When this is empty, the entry goes away. + jobs: HashMap, + /// The ready tasks of the suite, in the order that a claim takes them. + ready: Vec, + reservation: Option, + /// Queued ids that the last reconcile did not find as ready in the database. The next + /// reconcile removes these ids if it also does not find them. + suspects: HashSet, +} + +/// A ready task. The order is the order of the claim query: the highest priority first, then the +/// lowest id. Thus the queued ids are the ids that the claim query selects. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct ReadyTask { + priority: i32, + id: i64, +} + +impl Ord for ReadyTask { + fn cmp(&self, other: &Self) -> Ordering { + other + .priority + .cmp(&self.priority) + .then(self.id.cmp(&other.id)) + } +} + +impl PartialOrd for ReadyTask { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.cmp(other)) + } +} + +#[derive(Debug, Default)] +struct JobState { + agent_uuid: Uuid, + /// The job id in the suite. If two jobs hold the same part of their budget, the older job gets + /// the offer first. + job_id: i32, + /// The maximum number of tasks that the job can hold ([`task_budget`]). The queue keeps a copy, + /// because the offer does not read the database. Each fetch updates it. + budget: u32, + /// The number of tasks that the job holds: claimed and not committed. A task leaves a job only + /// when the job commits it or when the job ends. + outstanding: u32, +} + +/// The jobs to notify about the ready tasks of a suite. +#[derive(Debug, Default, PartialEq, Eq)] +pub struct FollowUpOffer { + /// The jobs to notify. The job that holds the largest part of its budget is first. Empty if no + /// job has free room. + pub agents: Vec<(i64, Uuid)>, + /// True if these jobs have room for all the ready tasks. Only then does the suite get a + /// reservation. + pub covered: bool, +} + +#[derive(Debug)] +struct Reservation { + /// The agents that can claim from the suite: all the agents on the suite when the reservation + /// started. This includes agents that did not get a notification. + agents: Vec, + /// The notified agents that did not fetch yet. When this is empty, the reservation ends. + awaited: Vec, + until: Instant, +} + +impl SuiteQueue { + /// Remove a reservation after its end time. Then `Some` always means an active reservation. + fn expire(&mut self, now: Instant) { + if self.reservation.as_ref().is_some_and(|r| r.until <= now) { + self.reservation = None; + } + } + + fn reserved_against(&self, agent_id: i64) -> bool { + self.reservation + .as_ref() + .is_some_and(|r| !r.agents.contains(&agent_id)) + } + + /// End the reservation when all the notified jobs fetched, or when the queue is empty. An empty + /// queue does not prove that the suite has no ready task. But an early end only lets more + /// agents in, and that is safe. + fn settle_reservation(&mut self) { + let done = self + .reservation + .as_ref() + .is_some_and(|r| r.awaited.is_empty()); + if done || self.ready.is_empty() { + self.reservation = None; + } + } +} + +impl SuiteQueues { + pub fn new() -> Self { + Self::default() + } + + /// The lock is held only for map operations. Thus a poisoned lock means a panic in a different + /// thread, not bad data, and the data is safe to use. + fn lock(&self) -> MutexGuard<'_, HashMap> { + self.inner.lock().unwrap_or_else(|e| e.into_inner()) + } + + /// Add a job to the entry of its suite. Create the entry if it does not exist. + /// + /// A job that starts during a reservation does not join the reservation, because the + /// reservation is to keep such a job out. + /// + /// `outstanding` is the number of tasks that the job holds. It is zero for a new job. + /// + /// Returns true if this call created the entry. Then the caller must give the ready tasks of + /// the suite to [`SuiteQueues::push_ready`]. The queue ignores tasks of a suite without an + /// entry, thus the new entry does not have them. + pub fn open( + &self, + suite_id: i64, + agent_id: i64, + agent_uuid: Uuid, + job_id: i32, + budget: u32, + outstanding: u32, + ) -> bool { + let mut queues = self.lock(); + let opened = !queues.contains_key(&suite_id); + let queue = queues.entry(suite_id).or_default(); + queue.jobs.insert( + agent_id, + JobState { + agent_uuid, + job_id, + budget, + outstanding, + }, + ); + tracing::trace!( + suite_id, + agent_id, + job_id, + budget, + outstanding, + opened, + jobs = queue.jobs.len(), + ready = queue.ready.len(), + "A job joined the suite queue" + ); + opened + } + + /// Remove a job that ended. When the last job of the suite ends, the entry goes away. + /// + /// `reclaimed` is the number of tasks that the database took back from the job. It must be + /// equal to `outstanding`. A difference is a bug in the counts. This function only logs the + /// difference, because the job state goes away. + pub fn close(&self, suite_id: i64, agent_id: i64, reclaimed: usize) { + let mut queues = self.lock(); + match queues + .get_mut(&suite_id) + .and_then(|queue| queue.jobs.remove(&agent_id)) + { + Some(job) if job.outstanding as usize != reclaimed => tracing::warn!( + suite_id, + agent_id, + job_id = job.job_id, + counted = job.outstanding, + reclaimed, + "An ending job gave back a different number of tasks than the queue counted" + ), + None if reclaimed > 0 => tracing::warn!( + suite_id, + agent_id, + reclaimed, + "A job that the suite queue did not track gave back tasks" + ), + _ => {} + } + let Some(queue) = queues.get_mut(&suite_id) else { + return; + }; + if let Some(reservation) = queue.reservation.as_mut() { + reservation.agents.retain(|id| *id != agent_id); + reservation.awaited.retain(|id| *id != agent_id); + } + queue.settle_reservation(); + let (jobs, ready) = (queue.jobs.len(), queue.ready.len()); + if queue.jobs.is_empty() { + queues.remove(&suite_id); + } + tracing::trace!( + suite_id, + agent_id, + reclaimed, + jobs, + ready, + "A job left the suite queue" + ); + } + + /// Remove the entry of a suite when a cancel ends all its jobs at the same time. There is no + /// count check, because the cancel archives the tasks itself. + pub fn close_suite(&self, suite_id: i64) { + if let Some(queue) = self.lock().remove(&suite_id) { + tracing::trace!( + suite_id, + jobs = queue.jobs.len(), + ready = queue.ready.len(), + "Dropped the suite queue with all its jobs" + ); + } + } + + /// Add tasks of the suite that are now ready: new tasks, tasks that left `Pending`, or tasks + /// that an agent gave back. + /// + /// **Call this in the transaction that makes the tasks ready, before the commit.** A claim + /// cannot lock a row before the commit. Thus the [`SuiteQueues::served`] call that removes the + /// task always occurs after this call. If the push occurs after the commit, a claim can take + /// the task first. Then the push adds a running task, and nothing removes it. + /// + /// If the transaction rolls back, the ids stay in the queue until [`reconcile`] removes them. + /// + /// If no job runs the suite, this function does nothing. A new entry reads the ready tasks from + /// the database. A task that is already in the queue does not get a second position. + pub fn push_ready(&self, suite_id: i64, tasks: impl IntoIterator) { + let mut queues = self.lock(); + let Some(queue) = queues.get_mut(&suite_id) else { + return; + }; + let before = queue.ready.len(); + queue.ready.extend( + tasks + .into_iter() + .map(|(id, priority)| ReadyTask { id, priority }), + ); + queue.ready.sort_unstable(); + queue.ready.dedup(); + tracing::trace!( + suite_id, + added = queue.ready.len().saturating_sub(before), + ready = queue.ready.len(), + "Queued ready tasks before the commit that makes them ready" + ); + } + + /// Add the tasks of a reclaim, as [`job::reclaim_agent_tasks`] returns them. Tasks without a + /// suite are for workers, thus this function skips them. + /// + /// Call this in the reclaim transaction, for the same reason as [`SuiteQueues::push_ready`]. + /// + /// [`job::reclaim_agent_tasks`]: crate::service::agent::job::reclaim_agent_tasks + pub fn push_reclaimed(&self, reclaimed: &[(i64, Option, i32)]) { + let mut queues = self.lock(); + let mut touched: HashMap = HashMap::new(); + for (task_id, suite_id, priority) in reclaimed { + let Some(suite_id) = *suite_id else { continue }; + let Some(queue) = queues.get_mut(&suite_id) else { + continue; + }; + queue.ready.push(ReadyTask { + id: *task_id, + priority: *priority, + }); + *touched.entry(suite_id).or_default() += 1; + } + for (suite_id, pushed) in touched { + if let Some(queue) = queues.get_mut(&suite_id) { + queue.ready.sort_unstable(); + queue.ready.dedup(); + tracing::trace!( + suite_id, + pushed, + ready = queue.ready.len(), + "Queued the tasks that a reclaim took from an agent" + ); + } + } + } + + /// Remove tasks that a cancel archived. A claim removes its tasks with + /// [`SuiteQueues::served`]. + pub fn remove_ready(&self, suite_id: i64, task_ids: impl IntoIterator) { + let gone: HashSet = task_ids.into_iter().collect(); + if gone.is_empty() { + return; + } + let mut queues = self.lock(); + let Some(queue) = queues.get_mut(&suite_id) else { + return; + }; + let before = queue.ready.len(); + queue.ready.retain(|task| !gone.contains(&task.id)); + queue.settle_reservation(); + tracing::trace!( + suite_id, + removed = before.saturating_sub(queue.ready.len()), + ready = queue.ready.len(), + "Removed cancelled tasks from the queue" + ); + } + + /// The suites that have an entry. + pub fn tracked_suites(&self) -> Vec { + self.lock().keys().copied().collect() + } + + /// Remove the queued ids that the database does not show as ready. + /// + /// `claimable` has the ready tasks of each suite that was read. A suite with no ready task has + /// an empty set. An entry that is not in the map does not change. + /// + /// **An id goes only when two reads in sequence do not show it.** The queue gets a task before + /// the commit that makes it ready. Thus one read cannot find the difference between a + /// transaction that did not commit yet and a transaction that rolled back. At the second read, + /// the transaction has ended. + /// + /// This function does not add ids. A ready task that is not in the queue is a missing + /// `push_ready`, and `served` logs it. An add here can race with the claim that takes the task. + pub fn drop_unclaimable(&self, claimable: &HashMap>) { + let mut queues = self.lock(); + for (suite_id, claimable) in claimable { + let Some(queue) = queues.get_mut(suite_id) else { + continue; + }; + let missing: HashSet = queue + .ready + .iter() + .map(|task| task.id) + .filter(|id| !claimable.contains(id)) + .collect(); + let condemned: HashSet = missing.intersection(&queue.suspects).copied().collect(); + queue.suspects = missing.difference(&condemned).copied().collect(); + if condemned.is_empty() { + if !queue.suspects.is_empty() { + tracing::trace!( + suite_id, + suspects = queue.suspects.len(), + ready = queue.ready.len(), + "Queued tasks are not ready in the database yet; one read is not proof" + ); + } + continue; + } + queue.ready.retain(|task| !condemned.contains(&task.id)); + queue.settle_reservation(); + // Only a bug or a rollback causes this. Log the ids to find the tasks in + // `archived_tasks`. + tracing::warn!( + suite_id, + dropped = condemned.len(), + ?condemned, + suspects = queue.suspects.len(), + ready = queue.ready.len(), + "Removed queued tasks that two reads did not show as ready" + ); + } + } + + /// Select the jobs to notify about the ready tasks of the suite. + /// + /// Only jobs with free room. The job that holds the largest part of its budget is first. The + /// list stops when the jobs have room for all the ready tasks. If all the jobs together do not + /// have this room, `covered` is false. A suite with no ready task gives an empty offer. + pub fn follow_up_offer(&self, suite_id: i64) -> FollowUpOffer { + let queues = self.lock(); + let Some(queue) = queues.get(&suite_id) else { + return FollowUpOffer::default(); + }; + let backlog = queue.ready.len(); + if backlog == 0 { + return FollowUpOffer::default(); + } + + let mut ranked: Vec<(i64, &JobState)> = queue + .jobs + .iter() + .filter(|(_, job)| job.outstanding < job.budget) + .map(|(agent_id, job)| (*agent_id, job)) + .collect(); + // Compare `outstanding / budget` by cross-multiplication, because the jobs of one suite can + // have different budgets. + ranked.sort_by(|(_, a), (_, b)| { + (u64::from(b.outstanding) * u64::from(a.budget)) + .cmp(&(u64::from(a.outstanding) * u64::from(b.budget))) + .then(a.job_id.cmp(&b.job_id)) + }); + + let mut agents = Vec::new(); + let mut room = 0usize; + for (agent_id, job) in ranked { + if room >= backlog { + break; + } + room += job.budget.saturating_sub(job.outstanding) as usize; + agents.push((agent_id, job.agent_uuid)); + } + FollowUpOffer { + agents, + covered: room >= backlog, + } + } + + /// End the reservation of the suite. + pub fn release_reservation(&self, suite_id: i64) { + let mut queues = self.lock(); + if let Some(queue) = queues.get_mut(&suite_id) { + if queue.reservation.take().is_some() { + tracing::trace!( + suite_id, + ready = queue.ready.len(), + "Released the suite reservation: its jobs cannot take all its ready tasks" + ); + } + } + } + + /// Reserve the suite for its jobs until each agent in `notified` fetches or `window` ends. This + /// function does nothing for a suite without a job. + pub fn reserve(&self, suite_id: i64, notified: Vec, window: Duration) { + let mut queues = self.lock(); + let Some(queue) = queues.get_mut(&suite_id) else { + return; + }; + if queue.jobs.is_empty() { + return; + } + let awaited = notified.len(); + queue.reservation = Some(Reservation { + agents: queue.jobs.keys().copied().collect(), + awaited: notified, + until: Instant::now() + window, + }); + tracing::trace!( + suite_id, + agents = queue.jobs.len(), + awaited, + ready = queue.ready.len(), + ?window, + "Reserved the suite for its jobs" + ); + } + + /// The suites that are reserved for agents other than this agent. + pub fn blocked_for(&self, agent_id: i64) -> Vec { + let now = Instant::now(); + let mut queues = self.lock(); + queues + .iter_mut() + .filter_map(|(suite_id, queue)| { + queue.expire(now); + queue.reserved_against(agent_id).then_some(*suite_id) + }) + .collect() + } + + /// Get the queued ids that a claim tries first. `None` refuses the claim, because the suite is + /// reserved for other agents. + /// + /// The number of ids is at most the free room of the job. This function does not remove the + /// ids: [`SuiteQueues::served`] removes the ids that the claim locked. Thus two agents can get + /// the same ids, and `SKIP LOCKED` gives each row to one of them. + /// + /// A suite without an entry gives no ids and no refusal. This must not occur, because each job + /// in flight has an entry. + pub fn take(&self, suite_id: i64, agent_id: i64, budget: u32) -> Option> { + let now = Instant::now(); + let mut queues = self.lock(); + let Some(queue) = queues.get_mut(&suite_id) else { + tracing::debug!( + suite_id, + agent_id, + budget, + "A job fetched from a suite that the queue does not track" + ); + return Some(Vec::new()); + }; + queue.expire(now); + if queue.reserved_against(agent_id) { + tracing::trace!( + suite_id, + agent_id, + ready = queue.ready.len(), + "Refused a claim: the suite is reserved for its jobs" + ); + return None; + } + let want = queue.jobs.get_mut(&agent_id).map_or(budget, |job| { + job.budget = budget; + budget.saturating_sub(job.outstanding) + }); + let hinted: Vec = queue + .ready + .iter() + .take(want as usize) + .map(|task| task.id) + .collect(); + tracing::trace!( + suite_id, + agent_id, + budget, + want, + hinted = hinted.len(), + ready = queue.ready.len(), + "Gave the queued ids to a claim" + ); + Some(hinted) + } + + /// Record the tasks that a claim gave to a job. The tasks leave the queue, and the job holds + /// them until it commits them. + /// + /// `hinted` is the number of ids from [`SuiteQueues::take`]. If the claim got more tasks, a + /// code path possibly made a task ready without [`SuiteQueues::push_ready`]. Then the queue + /// counts too few ready tasks, and the offers are too small. A task that became ready between + /// the two calls gives the same log line and is not a bug. Thus one log line is not proof. + /// + /// For the reservation, each fetch of a job counts, also a refused or empty fetch. + pub fn served(&self, suite_id: i64, agent_id: i64, hinted: usize, claimed: &[i64]) { + let taken: HashSet = claimed.iter().copied().collect(); + let mut queues = self.lock(); + let Some(queue) = queues.get_mut(&suite_id) else { + if !claimed.is_empty() { + tracing::warn!( + suite_id, + agent_id, + claimed = claimed.len(), + "A claim was served for a suite that the queue does not track" + ); + } + return; + }; + if claimed.len() > hinted { + // The claimed ids that are not in the queue. If the list is empty, the tasks became + // ready between `take` and this call, and that is not a bug. + let queued: HashSet = queue.ready.iter().map(|task| task.id).collect(); + let unknown: Vec = claimed + .iter() + .copied() + .filter(|id| !queued.contains(id)) + .collect(); + tracing::warn!( + suite_id, + agent_id, + hinted, + claimed = claimed.len(), + unknown = unknown.len(), + unknown_ids = ?unknown, + ready = queue.ready.len(), + "A claim found more ready tasks of a suite than the queue had" + ); + } + let mut outstanding = 0; + if let Some(job) = queue.jobs.get_mut(&agent_id) { + job.outstanding = job.outstanding.saturating_add(claimed.len() as u32); + outstanding = job.outstanding; + } + queue.ready.retain(|task| !taken.contains(&task.id)); + if let Some(reservation) = queue.reservation.as_mut() { + reservation.awaited.retain(|id| *id != agent_id); + } + queue.settle_reservation(); + tracing::trace!( + suite_id, + agent_id, + claimed = claimed.len(), + outstanding, + ready = queue.ready.len(), + reserved = queue.reservation.is_some(), + "Recorded the tasks of a claim" + ); + } + + /// Move a queued task to the position for its new priority. If the task is not in the queue, + /// do nothing: a claim possibly took it, and the task is not ready anymore. + pub fn reprioritize(&self, suite_id: i64, task_id: i64, priority: i32) { + let mut queues = self.lock(); + let Some(queue) = queues.get_mut(&suite_id) else { + return; + }; + let Some(task) = queue.ready.iter_mut().find(|task| task.id == task_id) else { + return; + }; + if task.priority == priority { + return; + } + task.priority = priority; + queue.ready.sort_unstable(); + tracing::trace!( + suite_id, + task_id, + priority, + "Moved a queued task after a priority change" + ); + } + + /// Record that the job committed a task. + pub fn finished(&self, suite_id: i64, agent_id: i64) { + let mut queues = self.lock(); + if let Some(job) = queues + .get_mut(&suite_id) + .and_then(|queue| queue.jobs.get_mut(&agent_id)) + { + job.outstanding = job.outstanding.saturating_sub(1); + tracing::trace!( + suite_id, + agent_id, + outstanding = job.outstanding, + budget = job.budget, + "A job committed a task" + ); + } + } +} + +/// The ready tasks of these suites, as `(suite id, task id, priority)`. +pub async fn ready_tasks( + db: &C, + suite_ids: Vec, +) -> crate::error::Result> { + Ok(ActiveTasks::Entity::find() + .select_only() + .column(ActiveTasks::Column::TaskSuiteId) + .column(ActiveTasks::Column::Id) + .column(ActiveTasks::Column::Priority) + .filter(ActiveTasks::Column::TaskSuiteId.is_in(suite_ids)) + .filter(ActiveTasks::Column::State.eq(TaskState::Ready)) + .into_tuple::<(Option, i64, i32)>() + .all(db) + .await? + .into_iter() + .filter_map(|(suite_id, task_id, priority)| Some((suite_id?, task_id, priority))) + .collect()) +} + +/// Remove from each entry the ids that the database does not show as ready. +/// +/// A transaction that added ids and then rolled back leaves ids that no claim takes. These ids +/// make the offers too large. See [`SuiteQueues::drop_unclaimable`] for the two reads. +/// +/// This function has its own span, because no request starts it. +#[tracing::instrument( + name = "suite_queue_reconcile", + skip_all, + fields(suites = tracing::field::Empty) +)] +pub async fn reconcile(pool: &InfraPool) -> crate::error::Result<()> { + let suite_ids = pool.suite_queues.tracked_suites(); + if suite_ids.is_empty() { + return Ok(()); + } + tracing::Span::current().record("suites", suite_ids.len()); + // Put each suite that was read in the map, also a suite with no ready task. + let mut claimable: HashMap> = + suite_ids.iter().map(|id| (*id, HashSet::new())).collect(); + for (suite_id, task_id, _) in ready_tasks(&pool.db, suite_ids).await? { + claimable.entry(suite_id).or_default().insert(task_id); + } + tracing::debug!( + claimable = claimable.values().map(HashSet::len).sum::(), + "Read the ready tasks of the tracked suites" + ); + pool.suite_queues.drop_unclaimable(&claimable); + Ok(()) +} + +/// At coordinator start, create the entries for the jobs in flight. +pub async fn restore(pool: &InfraPool) -> crate::error::Result<()> { + let jobs: Vec<(i64, Option, i32)> = SuiteAgentJobs::Entity::find() + .select_only() + .column(SuiteAgentJobs::Column::TaskSuiteId) + .column(SuiteAgentJobs::Column::AgentId) + .column(SuiteAgentJobs::Column::JobId) + .filter(SuiteAgentJobs::Column::State.is_in(job::IN_FLIGHT)) + .into_tuple() + .all(&pool.db) + .await?; + if jobs.is_empty() { + return Ok(()); + } + + let agent_ids: Vec = jobs + .iter() + .filter_map(|(_, agent_id, _)| *agent_id) + .collect(); + let uuids: HashMap = Agent::Entity::find() + .select_only() + .column(Agent::Column::Id) + .column(Agent::Column::Uuid) + .filter(Agent::Column::Id.is_in(agent_ids)) + .into_tuple::<(i64, Uuid)>() + .all(&pool.db) + .await? + .into_iter() + .collect(); + + let suite_ids: Vec = jobs.iter().map(|(suite_id, _, _)| *suite_id).collect(); + let budgets: HashMap = TaskSuites::Entity::find() + .select_only() + .column(TaskSuites::Column::Id) + .column(TaskSuites::Column::WorkerSchedule) + .filter(TaskSuites::Column::Id.is_in(suite_ids.clone())) + .into_tuple::<(i64, serde_json::Value)>() + .all(&pool.db) + .await? + .into_iter() + .filter_map(|(suite_id, schedule)| { + serde_json::from_value::(schedule) + .inspect_err(|e| { + tracing::error!(suite_id, "Stored worker schedule is unreadable: {e}"); + }) + .ok() + .map(|plan| (suite_id, task_budget(&plan))) + }) + .collect(); + + // Count the tasks that each job holds. If the count is zero, the queue offers a full budget to + // a job that already holds tasks. + let mut held: HashMap<(i64, Uuid), u32> = HashMap::new(); + for (suite_id, runner) in ActiveTasks::Entity::find() + .select_only() + .column(ActiveTasks::Column::TaskSuiteId) + .column(ActiveTasks::Column::RunnerUuid) + .filter(ActiveTasks::Column::TaskSuiteId.is_in(suite_ids.clone())) + .filter(ActiveTasks::Column::RunnerUuid.is_not_null()) + .into_tuple::<(Option, Option)>() + .all(&pool.db) + .await? + { + let (Some(suite_id), Some(runner)) = (suite_id, runner) else { + continue; + }; + *held.entry((suite_id, runner)).or_default() += 1; + } + + let mut claimable: HashMap> = HashMap::new(); + for (suite_id, task_id, priority) in ready_tasks(&pool.db, suite_ids).await? { + claimable + .entry(suite_id) + .or_default() + .push((task_id, priority)); + } + + let mut restored = 0; + for (suite_id, agent_id, job_id) in jobs { + let Some(agent_id) = agent_id else { continue }; + let Some(uuid) = uuids.get(&agent_id) else { + continue; + }; + // If the schedule cannot be read, the budget is zero until the next fetch sets it. + let budget = budgets.get(&suite_id).copied().unwrap_or(0); + let outstanding = held.get(&(suite_id, *uuid)).copied().unwrap_or(0); + if pool + .suite_queues + .open(suite_id, agent_id, *uuid, job_id, budget, outstanding) + { + pool.suite_queues + .push_ready(suite_id, claimable.remove(&suite_id).unwrap_or_default()); + } + restored += 1; + } + tracing::info!(jobs = restored, "Restored the suite jobs in flight"); + Ok(()) +} diff --git a/netmito/src/service/agent/task.rs b/netmito/src/service/agent/task.rs index 501cb88..c2ac8c3 100644 --- a/netmito/src/service/agent/task.rs +++ b/netmito/src/service/agent/task.rs @@ -7,6 +7,11 @@ //! LOCKED`, agents at the same time step over the rows that the others locked. Thus two agents //! cannot claim the same task, and no agent waits. //! +//! [`SuiteQueues`] gives the queued ids of the suite to try first, and refuses the claim while the +//! suite is reserved for other jobs. The claim transaction decides what the agent gets. +//! +//! [`SuiteQueues`]: crate::service::agent::queue::SuiteQueues +//! //! ## Reports //! //! Agents send the same [`ReportTaskOp`] as workers on `POST /tasks/report`, with the job handle. @@ -32,31 +37,10 @@ use crate::{ task_suites as TaskSuites, StoredTaskModel, }, error::{ApiError, Error, Result}, - schema::{ - ExecSpec, FetchTasksResp, ReportTaskOp, TaskExecOptions, WorkerSchedulePlan, WorkerTaskResp, - }, - service::{agent::job, agent::matching, s3::group_upload_artifact}, + schema::{ExecSpec, FetchTasksResp, ReportTaskOp, TaskExecOptions, WorkerTaskResp}, + service::{agent::job, agent::matching, agent::queue, s3::group_upload_artifact}, }; -/// The maximum number of tasks in one claim, so that a large `worker_count` cannot lock too many -/// rows. -const MAX_FETCH_BATCH: u32 = 256; - -/// The maximum number of tasks that one job can hold: the tasks that its workers run, and with -/// `prefetch`, one more task for each worker. Then a worker that becomes free gets its next task -/// without a round trip. -pub fn task_budget(schedule: &WorkerSchedulePlan) -> u32 { - match schedule { - WorkerSchedulePlan::FixedWorkers { - worker_count, - prefetch, - .. - } => ((*worker_count).max(1)) - .saturating_mul(1 + u32::from(*prefetch)) - .min(MAX_FETCH_BATCH), - } -} - /// Claim the ready tasks of a suite, up to the free room of the job of this agent. /// /// The agent does not give a number. The schedule of the suite gives the budget of one job, and @@ -92,11 +76,24 @@ pub async fn agent_fetch_tasks( // now, a wrong `true` costs only one more fetch. let hold_job_open = matches!(suite.state, TaskSuiteState::Open); - let budget = task_budget( - &serde_json::from_value(suite.worker_schedule).inspect_err(|e| { + let budget = queue::task_budget(&serde_json::from_value(suite.worker_schedule).inspect_err( + |e| { tracing::error!("Stored worker schedule is unreadable: {e}"); - })?, - ); + }, + )?); + + // While the suite is reserved for other jobs, this agent gets no task. `hold_job_open` does not + // change, thus an agent that holds its job fetches again later. + let Some(hinted) = pool.suite_queues.take(suite_id, agent_id, budget) else { + tracing::debug!("Suite is reserved for its jobs; gave no task"); + pool.suite_queues.served(suite_id, agent_id, 0, &[]); + return Ok(FetchTasksResp { + tasks: Vec::new(), + hold_job_open, + }); + }; + let hinted_count = hinted.len(); + let group_name = Group::Entity::find_by_id(suite.group_id) .one(&pool.db) .await? @@ -117,18 +114,22 @@ pub async fn agent_fetch_tasks( if room == 0 { return Ok(Vec::new()); } - 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(room) - .lock_with_behavior( - sea_orm::sea_query::LockType::Update, - sea_orm::sea_query::LockBehavior::SkipLocked, - ) - .all(txn) - .await?; + // Try the queued ids first. They are only a hint: the claim skips an id that is + // not `Ready` anymore. Then the query in priority order fills the remaining room. + let mut candidates = if hinted.is_empty() { + Vec::new() + } else { + claim_candidates(txn, suite_id, Some(hinted), room).await? + }; + if (candidates.len() as u64) < room { + let want = room - candidates.len() as u64; + let claimed: Vec = candidates.iter().map(|t| t.id).collect(); + for task in claim_candidates(txn, suite_id, None, want).await? { + if !claimed.contains(&task.id) { + candidates.push(task); + } + } + } if candidates.is_empty() { return Ok(Vec::new()); } @@ -147,6 +148,10 @@ pub async fn agent_fetch_tasks( }) .await?; + let claimed_ids: Vec = claimed.iter().map(|task| task.id).collect(); + pool.suite_queues + .served(suite_id, agent_id, hinted_count, &claimed_ids); + let mut tasks = Vec::with_capacity(claimed.len()); for task in claimed { let spec: ExecSpec = serde_json::from_value(task.spec).inspect_err(|e| { @@ -176,6 +181,32 @@ pub async fn agent_fetch_tasks( }) } +/// Lock up to `limit` ready tasks of a suite `FOR UPDATE SKIP LOCKED`, in priority order. With +/// `ids`, only from these tasks. +async fn claim_candidates( + txn: &C, + suite_id: i64, + ids: Option>, + limit: u64, +) -> Result> { + let mut query = ActiveTasks::Entity::find() + .filter(ActiveTasks::Column::TaskSuiteId.eq(suite_id)) + .filter(ActiveTasks::Column::State.eq(TaskState::Ready)); + if let Some(ids) = ids { + query = query.filter(ActiveTasks::Column::Id.is_in(ids)); + } + Ok(query + .order_by_desc(ActiveTasks::Column::Priority) + .order_by_asc(ActiveTasks::Column::Id) + .limit(limit) + .lock_with_behavior( + sea_orm::sea_query::LockType::Update, + sea_orm::sea_query::LockBehavior::SkipLocked, + ) + .all(txn) + .await?) +} + /// Apply a report of the agent to a task that it claimed. pub async fn agent_report_task( agent_id: i64, @@ -280,6 +311,10 @@ pub async fn agent_report_task( }) .await?; + if let Some(suite_id) = suite_id { + pool.suite_queues.finished(suite_id, agent_id); + } + // An earlier `Submit` made a `Pending` downstream task. The parent committed, thus the // downstream task becomes `Ready`. if let Some(downstream_uuid) = task.downstream_task_uuid { diff --git a/netmito/src/service/suite.rs b/netmito/src/service/suite.rs index b554203..a15af1c 100644 --- a/netmito/src/service/suite.rs +++ b/netmito/src/service/suite.rs @@ -1103,11 +1103,15 @@ pub async fn user_cancel_task_suite( let graceful = matches!(op, CancelTaskSuiteOp::Graceful); let running_agents = match op { CancelTaskSuiteOp::Force => { - pool.db + let killed = pool + .db .transaction::<_, Vec, Error>(move |txn| { Box::pin(async move { agent::job::kill_suite_jobs(txn, suite_id, now).await }) }) - .await? + .await?; + // All the jobs of the suite ended, and the suite has no more tasks to give. + pool.suite_queues.close_suite(suite_id); + killed } CancelTaskSuiteOp::Graceful => agent::job::agents_running_suite(&pool.db, suite_id).await?, }; diff --git a/netmito/src/service/task.rs b/netmito/src/service/task.rs index 02c29ef..c76fb70 100644 --- a/netmito/src/service/task.rs +++ b/netmito/src/service/task.rs @@ -56,6 +56,7 @@ enum Submitter { }, } +/// Insert a task. [`dispatch_submitted_task`] tells the runners about it. async fn internal_submit_task( pool: &InfraPool, creator_id: i64, @@ -69,7 +70,7 @@ async fn internal_submit_task( spec, exec_options, }: SubmitTaskReq, -) -> crate::error::Result { +) -> crate::error::Result { let tags = Vec::from_iter(tags); let labels = Vec::from_iter(labels); let now = TimeDateTimeWithTimeZone::now_utc(); @@ -92,6 +93,7 @@ async fn internal_submit_task( .map(serde_json::to_value) .transpose()?; + let queues = pool.suite_queues.clone(); let (task, suite) = pool .db .transaction::<_, (ActiveTasks::Model, Option), crate::error::Error>( @@ -383,9 +385,9 @@ async fn internal_submit_task( // commit and the push. The claim takes the task, and `served` finds nothing to // remove. Then the push adds a running task to the backlog of the suite, and no // process removes it. - // - // TODO: when suites have a ready queue, push a `Ready` suite task to it here - // with `queues.push_ready(suite.id, [(task.id, task.priority)])`. + if let (Some(suite), TaskState::Ready) = (&suite, task.state) { + queues.push_ready(suite.id, [(task.id, task.priority)]); + } Ok((task, suite)) }) @@ -393,18 +395,46 @@ async fn internal_submit_task( ) .await?; - // If task is pending, then we have done + Ok(SubmittedTask { + task, + suite_id: suite.map(|suite| suite.id), + }) +} + +/// A new task that the runners do not know about yet. +struct SubmittedTask { + task: ActiveTasks::Model, + suite_id: Option, +} + +impl SubmittedTask { + fn resp(&self) -> SubmitTaskResp { + SubmitTaskResp { + task_id: self.task.task_id, + uuid: self.task.uuid, + } + } +} + +/// Tell the runners about a submitted task: the agents of its suite, or the workers that match a +/// task without a suite. +/// +/// This is separate from the insert, thus a batch can insert all its tasks and then make one offer +/// for each suite. +async fn dispatch_submitted_task( + pool: &InfraPool, + submitted: SubmittedTask, +) -> crate::error::Result<()> { + let SubmittedTask { task, suite_id } = submitted; + + // A pending task does not run until it is triggered. if matches!(task.state, TaskState::Pending) { - return Ok(SubmitTaskResp { - task_id: task.task_id, - uuid: task.uuid, - }); + return Ok(()); }; - // Not pending, notify agent/worker depending on whether it belongs to a suite - match suite { - Some(suite) => { - crate::service::agent::notify_suite_available(pool, suite.id).await; + match suite_id { + Some(suite_id) => { + crate::service::agent::notify_suite_available(pool, suite_id).await; } None => { let builder = pool.db.get_database_backend(); @@ -442,10 +472,7 @@ async fn internal_submit_task( } } } - Ok(SubmitTaskResp { - task_id: task.task_id, - uuid: task.uuid, - }) + Ok(()) } pub async fn user_submit_task( @@ -453,7 +480,10 @@ pub async fn user_submit_task( creator_id: i64, req: SubmitTaskReq, ) -> crate::error::Result { - internal_submit_task(pool, creator_id, Submitter::User, req).await + let submitted = internal_submit_task(pool, creator_id, Submitter::User, req).await?; + let resp = submitted.resp(); + dispatch_submitted_task(pool, submitted).await?; + Ok(resp) } /// Submit a downstream task of a running task, for the worker or agent that runs the parent task. @@ -468,7 +498,7 @@ pub async fn worker_submit_pending_task( parent_group_id: i64, req: SubmitTaskReq, ) -> crate::error::Result { - internal_submit_task( + let submitted = internal_submit_task( pool, creator_id, Submitter::Task { @@ -477,23 +507,40 @@ pub async fn worker_submit_pending_task( }, req, ) - .await + .await?; + let resp = submitted.resp(); + dispatch_submitted_task(pool, submitted).await?; + Ok(resp) } pub async fn worker_trigger_pending_task(pool: &InfraPool, uuid: Uuid) -> crate::error::Result<()> { tracing::debug!("Trigger pending task {uuid}"); - let task = ActiveTasks::Entity::find() - .filter(ActiveTasks::Column::Uuid.eq(uuid)) - .filter(ActiveTasks::Column::State.eq(TaskState::Pending)) - .one(&pool.db) - .await? - .ok_or(Error::ApiError(crate::error::ApiError::NotFound(format!( - "Pending task with uuid {uuid}" - ))))?; - let mut task: ActiveTasks::ActiveModel = task.into(); - task.state = Set(TaskState::Ready); - task.updated_at = Set(TimeDateTimeWithTimeZone::now_utc()); - let task = task.update(&pool.db).await?; + // Use a transaction, thus a suite task goes into the queue before the commit that makes it + // ready. See `SuiteQueues::push_ready`. + let queues = pool.suite_queues.clone(); + let task = pool + .db + .transaction::<_, ActiveTasks::Model, crate::error::Error>(move |txn| { + Box::pin(async move { + let task = ActiveTasks::Entity::find() + .filter(ActiveTasks::Column::Uuid.eq(uuid)) + .filter(ActiveTasks::Column::State.eq(TaskState::Pending)) + .one(txn) + .await? + .ok_or(Error::ApiError(crate::error::ApiError::NotFound(format!( + "Pending task with uuid {uuid}" + ))))?; + let mut task: ActiveTasks::ActiveModel = task.into(); + task.state = Set(TaskState::Ready); + task.updated_at = Set(TimeDateTimeWithTimeZone::now_utc()); + let task = task.update(txn).await?; + if let Some(suite_id) = task.task_suite_id { + queues.push_ready(suite_id, [(task.id, task.priority)]); + } + Ok(task) + }) + }) + .await?; // A suite task does not go into the worker queues. Agents get it from its suite. Thus, wake the // idle agents that can run the suite. if let Some(suite_id) = task.task_suite_id { @@ -601,10 +648,11 @@ pub async fn user_change_task( }) }) .await?; - // Suite-bound tasks never live in the worker task queues, so there is nothing to - // re-dispatch after a change. - // TODO: notify change to agents if this task belongs to a task suite - if task.task_suite_id.is_some() { + // A suite task is not in the worker queues, thus there is nothing to dispatch again. The suite + // queue keeps its tasks in priority order, thus it must get the new priority. + if let Some(suite_id) = task.task_suite_id { + pool.suite_queues + .reprioritize(suite_id, task.id, task.priority); return Ok(()); } let builder = pool.db.get_database_backend(); @@ -724,9 +772,9 @@ pub async fn user_cancel_task( user_id: i64, uuid: Uuid, ) -> crate::error::Result<()> { - let task_id = pool + let (task_id, suite_id) = pool .db - .transaction::<_, i64, crate::error::Error>(|txn| { + .transaction::<_, (i64, Option), crate::error::Error>(|txn| { Box::pin(async move { let task = ActiveTasks::Entity::find() .filter(ActiveTasks::Column::Uuid.eq(uuid)) @@ -786,13 +834,16 @@ pub async fn user_cancel_task( } archived_task.insert(txn).await?; ActiveTasks::Entity::delete_by_id(task.id).exec(txn).await?; - Ok(task.id) + Ok((task.id, suite_id)) }) }) .await?; let _ = remove_task(task_id, pool) .inspect_err(|e| tracing::warn!("Failed to remove task {}: {:?}", task_id, e)); - // TODO: when suites have a ready queue, also remove a cancelled suite task from it. + // The task is not ready anymore, thus it also leaves the queue of its suite. + if let Some(suite_id) = suite_id { + pool.suite_queues.remove_ready(suite_id, [task_id]); + } Ok(()) } @@ -1531,6 +1582,23 @@ struct IdResult { task_suite_id: Option, } +/// Remove cancelled tasks from the queues of their suites, as `(suite id, task id)`. A cancel +/// archives a task where it is, thus no claim removes it from the queue. Tasks without a suite are +/// for workers, and `remove_task` handles them. +fn remove_cancelled_from_suite_queues( + pool: &InfraPool, + cancelled: impl IntoIterator, i64)>, +) { + let mut by_suite: HashMap> = HashMap::new(); + for (suite_id, task_id) in cancelled { + let Some(suite_id) = suite_id else { continue }; + by_suite.entry(suite_id).or_default().push(task_id); + } + for (suite_id, task_ids) in by_suite { + pool.suite_queues.remove_ready(suite_id, task_ids); + } +} + /// Decrease `incomplete_tasks` by one for each cancelled task, for each suite that a batch cancel /// changed. /// @@ -1702,9 +1770,9 @@ pub async fn cancel_tasks_by_filter( // Execute delete and insert in a single transaction // Total: 1 query using CTE (DELETE RETURNING + INSERT SELECT) regardless of task count or column count - let task_ids = pool + let cancelled = pool .db - .transaction::<_, Vec, crate::error::Error>(|txn| { + .transaction::<_, Vec, crate::error::Error>(|txn| { let with_ctes = with_ctes.clone(); Box::pin(async move { let now = TimeDateTimeWithTimeZone::now_utc(); @@ -1779,22 +1847,20 @@ pub async fn cancel_tasks_by_filter( ) .await?; - let task_ids: Vec = cancelled.into_iter().map(|r| r.id).collect(); - - Ok(task_ids) + Ok(cancelled) }) }) .await?; // Remove tasks from dispatch queues - for task_id in &task_ids { - let _ = remove_task(*task_id, pool) - .inspect_err(|e| tracing::warn!("Failed to remove task {}: {:?}", task_id, e)); + for task in &cancelled { + let _ = remove_task(task.id, pool) + .inspect_err(|e| tracing::warn!("Failed to remove task {}: {:?}", task.id, e)); } - // TODO: when suites have a ready queue, also remove the cancelled suite tasks from it. + remove_cancelled_from_suite_queues(pool, cancelled.iter().map(|r| (r.task_suite_id, r.id))); Ok(TasksCancelByFilterResp { - cancelled_count: task_ids.len() as u64, + cancelled_count: cancelled.len() as u64, group_name, }) } @@ -1936,7 +2002,7 @@ pub async fn cancel_tasks_by_uuids( // Total: 1 query using CTE (DELETE RETURNING + INSERT SELECT) per chunk let (task_ids, found_uuids) = pool .db - .transaction::<_, (Vec, HashSet), crate::error::Error>(|txn| { + .transaction::<_, (Vec<(Option, i64)>, HashSet), Error>(|txn| { let with_ctes = with_ctes.clone(); Box::pin(async move { let now = TimeDateTimeWithTimeZone::now_utc(); @@ -2013,7 +2079,8 @@ pub async fn cancel_tasks_by_uuids( ) .await?; - let task_ids: Vec = results.iter().map(|r| r.id).collect(); + let task_ids: Vec<(Option, i64)> = + results.iter().map(|r| (r.task_suite_id, r.id)).collect(); let found_uuids: HashSet = results.into_iter().map(|r| r.uuid).collect(); Ok((task_ids, found_uuids)) @@ -2027,11 +2094,11 @@ pub async fn cancel_tasks_by_uuids( } // Remove tasks from dispatch queues - for task_id in &all_task_ids { + for (_, task_id) in &all_task_ids { let _ = remove_task(*task_id, pool) .inspect_err(|e| tracing::warn!("Failed to remove task {}: {:?}", task_id, e)); } - // TODO: when suites have a ready queue, also remove the cancelled suite tasks from it. + remove_cancelled_from_suite_queues(pool, all_task_ids.iter().copied()); // Determine which UUIDs failed (not found or no permission) let failed_uuids: Vec = req @@ -2059,6 +2126,10 @@ impl From for i64 { /// Batch submit multiple tasks. /// Submits each task in the list and returns individual results for each (including failures). +/// +/// Each task has its own transaction, thus one rejected task does not stop the others. Each task +/// goes into the queue of its suite in its transaction. A suite gets one offer for all its new +/// tasks after the inserts, not one offer for each task. pub async fn user_batch_submit_tasks( pool: &InfraPool, user_id: i64, @@ -2071,21 +2142,36 @@ pub async fn user_batch_submit_tasks( } let mut results = Vec::with_capacity(req.tasks.len()); + let mut offer: HashSet = HashSet::new(); for task_req in req.tasks { - let result = user_submit_task(pool, user_id, task_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 result = async { + let submitted = internal_submit_task(pool, user_id, Submitter::User, task_req).await?; + let resp = submitted.resp(); + match submitted.suite_id { + Some(suite_id) if !matches!(submitted.task.state, TaskState::Pending) => { + offer.insert(suite_id); } - }) - .map_err(ErrorMsg::from); + _ => dispatch_submitted_task(pool, submitted).await?, + } + Ok::<_, Error>(resp) + } + .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 + } + }) + .map_err(ErrorMsg::from); results.push(result); } + for suite_id in offer { + crate::service::agent::notify_suite_available(pool, suite_id).await; + } + Ok(crate::schema::TasksSubmitResp { results }) } diff --git a/openapi.yaml b/openapi.yaml index d702ee0..1bdef39 100644 --- a/openapi.yaml +++ b/openapi.yaml @@ -5526,6 +5526,7 @@ components: - suite_cancelled - tasks_cancelled - stop_job + - tasks_available description: | A message from the coordinator that tells an agent to do something: - shutdown: shut down; with `graceful`, finish the current work first @@ -5539,6 +5540,8 @@ components: - tasks_cancelled: stop the tasks of `task_uuids` if they run - stop_job: stop the job on `suite_uuid` and pick a suite again; with `graceful`, finish the running tasks and run the cleanup hook first + - tasks_available: the job on `suite_uuid` has ready tasks; fetch now. + For a short time, the suite is reserved for the notified jobs WsNotificationEvent: type: object