From 0b490774fdddc12c68ae30c9f47ebaebaef46e56 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 16 Sep 2026 17:10:15 +0200 Subject: [PATCH 1/3] WIP --- opsqueue/src/delegation/server.rs | 437 ++++++++++++++---------------- 1 file changed, 197 insertions(+), 240 deletions(-) diff --git a/opsqueue/src/delegation/server.rs b/opsqueue/src/delegation/server.rs index 4a636ce..866c3ea 100644 --- a/opsqueue/src/delegation/server.rs +++ b/opsqueue/src/delegation/server.rs @@ -10,6 +10,7 @@ use std::sync::Arc; use tokio::select; use tokio::sync::Notify; use tokio_util::sync::CancellationToken; +// use tower::ServiceExt; #[cfg(test)] pub(crate) fn app_for_tests( @@ -28,8 +29,8 @@ pub(crate) fn app_for_tests( notify_on_insert, notify_on_submission_change, ) - .run_background() - .build_router(); + .run_background() + .build_router(); Router::new().nest("/job", router) } @@ -81,18 +82,15 @@ impl ServerState { state, cancellation_token, ) - .await - .ok(); + .await + .ok(); }); self } pub fn build_router(self: ServerState) -> Router<()> { Router::new() - .route("/delegate", post(job_delegate)) - .route("/kill", post(job_kill)) - .route("/return", post(job_return)) - // .route("/submit", post(submit)) + .route("/submit", post(submit)) .with_state(self) } } @@ -107,16 +105,16 @@ enum DelegatedJobStatus { Cancelled, } -// #[derive(Debug, serde::Deserialize)] -// #[serde(tag = "type", content = "contents")] -// enum WorkerDelegationEvent { -// #[serde(rename = "delegate")] -// Delegate(Vec), -// #[serde(rename = "kill")] -// Kill(Vec), -// #[serde(rename = "return")] -// Return(Vec), -// } +#[derive(Debug, serde::Deserialize)] +#[serde(tag = "type", content = "contents")] +enum WorkerDelegationEvent { + #[serde(rename = "delegate")] + Delegate(Vec), + #[serde(rename = "kill")] + Kill(Vec), + #[serde(rename = "return")] + Return(Vec), +} #[derive(Debug, serde::Serialize, serde::Deserialize)] struct DelegatedJob { task_id: String, @@ -128,14 +126,14 @@ struct DelegatedJobPayload { submission_id: SubmissionId, } -// #[derive(Debug, serde::Serialize)] -// #[serde(tag = "type", content = "contents")] -// enum MasterDelegationEvent<'a> { -// #[serde(rename = "updated")] -// Updated(Vec>), -// #[serde(rename = "completed")] -// Completed(Vec>), -// } +#[derive(Debug, serde::Serialize)] +#[serde(tag = "type", content = "contents")] +enum MasterDelegationEvent<'a> { + #[serde(rename = "updated")] + Updated(Vec>), + #[serde(rename = "completed")] + Completed(Vec>), +} #[derive(Debug, serde::Serialize)] struct DelegatedJobUpdate<'a> { @@ -172,77 +170,10 @@ enum FailureReason { Forced, } -// TODO(delegation): Switch to -// #[tracing::instrument(level = "debug", skip(state))] -// async fn submit( -// State(state): State, -// Json(events): Json>, -// ) -> Result { -// let mut conn = state.pool.writer_conn().await.map_err(|e| { -// tracing::error!("DB error acquiring writer connection: {e:?}"); -// StatusCode::INTERNAL_SERVER_ERROR -// })?; -// // TODO(delegation): Operate within a transaction. -// for event in events { -// match event { -// WorkerDelegationEvent::Delegate(delegations) => { -// for delegation in delegations { -// handle_delegate_event(&state, &mut conn, &delegation) -// .await -// .map_err(|e| { -// tracing::error!("Error handling delegate event: {e:?}"); -// e -// })?; -// } -// } -// WorkerDelegationEvent::Kill(task_ids) => { -// for task_id in task_ids { -// handle_kill_event(&state, &mut conn, &task_id) -// .await -// .map_err(|e| { -// tracing::error!( -// "Error handling kill event for task_id={task_id}: {e:?}" -// ); -// e -// })?; -// } -// } -// WorkerDelegationEvent::Return(_task_ids) => { -// tracing::info!( -// "Received 'return' delegation event, which is not yet implemented; ignoring." -// ); -// return Ok(StatusCode::ACCEPTED); -// } -// } -// } -// -// Ok(StatusCode::ACCEPTED) -// } - -#[tracing::instrument(level = "debug", skip(state))] -async fn job_delegate( - State(state): State, - Json(job): Json, -) -> Result { - let mut conn = state.pool.writer_conn().await.map_err(|e| { - tracing::error!("DB error acquiring writer connection: {e:?}"); - StatusCode::INTERNAL_SERVER_ERROR - })?; - handle_delegate_event(&mut conn, &job).await.map_err(|e| { - tracing::error!("DB error handling delegate event: {e:?}"); - StatusCode::INTERNAL_SERVER_ERROR - })?; - - state.notify_on_submission_change.notify_one(); - state.notify_on_insert.notify_waiters(); - - Ok(StatusCode::ACCEPTED) -} - #[tracing::instrument(level = "debug", skip(state))] -async fn job_kill( +async fn submit( State(state): State, - Json(task_ids): Json>, + Json(events): Json>, ) -> Result { let mut conn = state.pool.writer_conn().await.map_err(|e| { tracing::error!("DB error acquiring writer connection: {e:?}"); @@ -251,33 +182,107 @@ async fn job_kill( conn.transaction(move |mut tx| { Box::pin(async move { - for task_id in &task_ids { - handle_kill_event(&mut tx, task_id).await?; + for event in events { + match event { + WorkerDelegationEvent::Delegate(delegations) => { + for delegation in delegations { + handle_delegate_event(&mut tx, &delegation) + .await + .map_err(|e| { + tracing::error!("Error handling delegate event: {e:?}"); + e + })?; + } + } + WorkerDelegationEvent::Kill(task_ids) => { + for task_id in task_ids { + handle_kill_event(&mut tx, &task_id) + .await + .map_err(|e| { + tracing::error!( + "Error handling kill event for task_id={task_id}: {e:?}" + ); + e + })?; + } + } + WorkerDelegationEvent::Return(_task_ids) => { + tracing::info!( + "Received 'return' delegation event, which is not yet implemented; ignoring." + ); + } + } } - Ok::<(), sqlx::Error>(()) + Ok(()) }) - }) - .await - .map_err(|e| { + }).await.map_err(|e: sqlx::Error| { tracing::error!("DB error handling kill event: {e:?}"); StatusCode::INTERNAL_SERVER_ERROR })?; - state.notify_on_submission_change.notify_one(); - Ok(StatusCode::ACCEPTED) } -#[tracing::instrument(level = "debug", skip(_state))] -async fn job_return( - State(_state): State, - Json(task_ids): Json>, -) -> Result { - tracing::info!("Received 'return' delegation event, which is not yet implemented; ignoring."); - - Ok(StatusCode::ACCEPTED) -} +// #[tracing::instrument(level = "debug", skip(state))] +// async fn job_delegate( +// State(state): State, +// Json(job): Json, +// ) -> Result { +// let mut conn = state.pool.writer_conn().await.map_err(|e| { +// tracing::error!("DB error acquiring writer connection: {e:?}"); +// StatusCode::INTERNAL_SERVER_ERROR +// })?; +// handle_delegate_event(&mut conn, &job).await.map_err(|e| { +// tracing::error!("DB error handling delegate event: {e:?}"); +// StatusCode::INTERNAL_SERVER_ERROR +// })?; +// +// state.notify_on_submission_change.notify_one(); +// state.notify_on_insert.notify_waiters(); +// +// Ok(StatusCode::ACCEPTED) +// } +// +// #[tracing::instrument(level = "debug", skip(state))] +// async fn job_kill( +// State(state): State, +// Json(task_ids): Json>, +// ) -> Result { +// let mut conn = state.pool.writer_conn().await.map_err(|e| { +// tracing::error!("DB error acquiring writer connection: {e:?}"); +// StatusCode::INTERNAL_SERVER_ERROR +// })?; +// +// conn.transaction(move |mut tx| { +// Box::pin(async move { +// for task_id in &task_ids { +// handle_kill_event(&mut tx, task_id).await?; +// } +// +// Ok::<(), sqlx::Error>(()) +// }) +// }) +// .await +// .map_err(|e| { +// tracing::error!("DB error handling kill event: {e:?}"); +// StatusCode::INTERNAL_SERVER_ERROR +// })?; +// +// state.notify_on_submission_change.notify_one(); +// +// Ok(StatusCode::ACCEPTED) +// } +// +// #[tracing::instrument(level = "debug", skip(_state))] +// async fn job_return( +// State(_state): State, +// Json(task_ids): Json>, +// ) -> Result { +// tracing::info!("Received 'return' delegation event, which is not yet implemented; ignoring."); +// +// Ok(StatusCode::ACCEPTED) +// } #[tracing::instrument(level = "debug", skip(conn))] async fn handle_delegate_event( @@ -315,8 +320,8 @@ async fn handle_kill_event(conn: &mut impl WriterConnection, task_id: &str) -> s WHERE task_id = $1"#, task_id, ) - .fetch_optional(conn.get_inner()) - .await?; + .fetch_optional(conn.get_inner()) + .await?; let Some(submission_id) = submission_id else { tracing::warn!(%task_id, "Kill event for unknown task_id; ignoring"); @@ -421,38 +426,46 @@ async fn report_submission_status( } } - if !updates.is_empty() { - send_updates(state, &updates).await?; - let conn = state.pool.writer_conn().await?; - update_last_status_sent( - conn, - out_of_date_tasks - .iter() - .filter(|task| { - task.current_status == DelegatedJobStatus::Paused - || task.current_status == DelegatedJobStatus::InProgress - }) - .collect(), - ) - .await?; - } - - if !completions.is_empty() { - send_completions(state, &completions).await?; - let conn = state.pool.writer_conn().await?; - delete_external_tasks( - conn, - out_of_date_tasks - .iter() - .filter(|task| { - task.current_status == DelegatedJobStatus::Completed - || task.current_status == DelegatedJobStatus::Failed - || task.current_status == DelegatedJobStatus::Cancelled - }) - .collect(), - ) - .await?; - } + let events = vec![ + MasterDelegationEvent::Updated(updates), + MasterDelegationEvent::Completed(completions), + ]; + + send_events(state, &events).await?; + + let updated = out_of_date_tasks + .iter() + .filter(|task| { + task.current_status == DelegatedJobStatus::Paused + || task.current_status == DelegatedJobStatus::InProgress + }) + .collect(); + + let completed = out_of_date_tasks + .iter() + .filter(|task| { + task.current_status == DelegatedJobStatus::Completed + || task.current_status == DelegatedJobStatus::Failed + || task.current_status == DelegatedJobStatus::Cancelled + }) + .collect(); + + let mut conn = state.pool.writer_conn().await?; + conn.transaction(move |mut tx| { + Box::pin(async move { + update_last_status_sent( + &mut tx, + updated, + ) + .await?; + delete_external_tasks( + &mut tx, + completed, + ).await?; + + Ok::<(), sqlx::Error>(()) + }) + }).await?; } Ok(()) @@ -474,9 +487,9 @@ async fn insert_external_task( submission_id, task_id, ) - .execute(conn.get_inner()) - .await? - .rows_affected(); + .execute(conn.get_inner()) + .await? + .rows_affected(); Ok(rows_affected) } @@ -529,22 +542,15 @@ async fn update_last_status_sent( .map(|t| (t.current_status, t.task_id.clone())) .collect::>(); - conn.transaction(move |mut tx| { - Box::pin(async move { - for (current_status, task_id) in tasks { - sqlx::query!( - "UPDATE submissions_external_task SET last_status_sent = $1 WHERE task_id = $2", - current_status, - task_id, - ) - .execute(tx.get_inner()) - .await?; - } - - Ok::<_, sqlx::Error>(()) - }) - }) - .await?; + for (current_status, task_id) in tasks { + sqlx::query!( + "UPDATE submissions_external_task SET last_status_sent = $1 WHERE task_id = $2", + current_status, + task_id, + ) + .execute(conn.get_inner()) + .await?; + } Ok(()) } @@ -555,80 +561,31 @@ async fn delete_external_tasks( ) -> sqlx::Result<()> { let tasks = tasks.iter().map(|t| t.task_id.clone()).collect::>(); - conn.transaction(move |mut tx| { - Box::pin(async move { - for task_id in tasks { - sqlx::query!( - "DELETE FROM submissions_external_task WHERE task_id = $1", - task_id, - ) - .execute(tx.get_inner()) - .await?; - } - - Ok::<_, sqlx::Error>(()) - }) - }) - .await?; - - Ok(()) -} - -// TODO(delegation): Replace `send_updates` and `send_completions` with `send_events`, -// after https://github.com/channable/jobmachine/pull/2210 is merged. -// async fn send_events( -// state: &ServerState, -// events: &MasterDelegationEvent<'_>, -// ) -> reqwest::Result<()> { -// state -// .http_client -// .put( -// state -// .delegation_server_url -// .join("/delegation/submit") -// .unwrap(), -// ) -// .json(&events) -// .send() -// .await? -// .error_for_status()?; -// -// Ok(()) -// } - -async fn send_updates( - state: &ServerState, - updates: &[DelegatedJobUpdate<'_>], -) -> reqwest::Result<()> { - state - .http_client - .put( - state - .delegation_server_url - .join("/delegation/update") - .unwrap(), + for task_id in tasks { + sqlx::query!( + "DELETE FROM submissions_external_task WHERE task_id = $1", + task_id, ) - .json(updates) - .send() - .await? - .error_for_status()?; + .execute(conn.get_inner()) + .await?; + } Ok(()) } -async fn send_completions( +async fn send_events( state: &ServerState, - completions: &[DelegatedJobCompletion<'_>], + events: &Vec>, ) -> reqwest::Result<()> { state .http_client .put( state .delegation_server_url - .join("/delegation/complete") + .join("/delegation/submit") .unwrap(), ) - .json(completions) + .json(&events) .send() .await? .error_for_status()?; @@ -735,8 +692,8 @@ pub mod test { InitialSubmissionStatus::Paused, &mut conn, ) - .await - .unwrap() + .await + .unwrap() }; { @@ -770,7 +727,7 @@ pub mod test { submission_id: submission, }, }) - .unwrap(), + .unwrap(), )) .unwrap(), ) @@ -835,8 +792,8 @@ pub mod test { InitialSubmissionStatus::Paused, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await @@ -932,8 +889,8 @@ pub mod test { InitialSubmissionStatus::Paused, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await @@ -1007,8 +964,8 @@ pub mod test { InitialSubmissionStatus::InProgress, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await @@ -1084,8 +1041,8 @@ pub mod test { InitialSubmissionStatus::InProgress, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await @@ -1111,8 +1068,8 @@ pub mod test { &mut conn, 0, ) - .await - .unwrap(); + .await + .unwrap(); notify_on_submission_change.notify_one(); } @@ -1164,8 +1121,8 @@ pub mod test { InitialSubmissionStatus::InProgress, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await From 6a25922806d679919940956ec9a0694c37cdccc4 Mon Sep 17 00:00:00 2001 From: Marten Wijnja Date: Thu, 17 Sep 2026 18:05:55 +0200 Subject: [PATCH 2/3] WIP --- opsqueue/src/db/mod.rs | 3 +- opsqueue/src/delegation/server.rs | 76 +++++++++++++++++-------------- 2 files changed, 43 insertions(+), 36 deletions(-) diff --git a/opsqueue/src/db/mod.rs b/opsqueue/src/db/mod.rs index edc9574..01636f9 100644 --- a/opsqueue/src/db/mod.rs +++ b/opsqueue/src/db/mod.rs @@ -130,8 +130,7 @@ pub trait Connection { where for<'t> F: FnOnce(Conn>) -> BoxFuture<'t, Result> + Send - + Sync - + 't, + + Sync, O: Send, E: From + Send, { diff --git a/opsqueue/src/delegation/server.rs b/opsqueue/src/delegation/server.rs index 866c3ea..590a29b 100644 --- a/opsqueue/src/delegation/server.rs +++ b/opsqueue/src/delegation/server.rs @@ -6,6 +6,9 @@ use axum::extract::State; use axum::http::StatusCode; use axum::routing::post; use axum::{Json, Router}; +use futures::Stream; +use futures::stream::BoxStream; +use std::pin::Pin; use std::sync::Arc; use tokio::select; use tokio::sync::Notify; @@ -372,6 +375,16 @@ async fn run_in_background( Ok(()) } +async fn report_submission_status2(state: &ServerState, triggered_by_timeout: bool,) -> anyhow::Result<()> { + let mut conn = state.pool.reader_conn().await?; + let out_of_date_tasks = select_out_of_date_tasks(&mut conn).await; + for Ok(batch) in out_of_date_tasks.try_chunks(2048).await? { + + } + + Ok(()) +} + async fn report_submission_status( state: &ServerState, triggered_by_timeout: bool, @@ -433,31 +446,34 @@ async fn report_submission_status( send_events(state, &events).await?; - let updated = out_of_date_tasks - .iter() - .filter(|task| { - task.current_status == DelegatedJobStatus::Paused - || task.current_status == DelegatedJobStatus::InProgress - }) - .collect(); - - let completed = out_of_date_tasks - .iter() - .filter(|task| { - task.current_status == DelegatedJobStatus::Completed - || task.current_status == DelegatedJobStatus::Failed - || task.current_status == DelegatedJobStatus::Cancelled - }) - .collect(); + // let updated = batch + // .iter() + // .filter(|task| { + // task.current_status == DelegatedJobStatus::Paused + // || task.current_status == DelegatedJobStatus::InProgress + // }) + // .collect(); + + // let completed = batch + // .iter() + // .filter(|task| { + // task.current_status == DelegatedJobStatus::Completed + // || task.current_status == DelegatedJobStatus::Failed + // || task.current_status == DelegatedJobStatus::Cancelled + // }) + // .collect(); let mut conn = state.pool.writer_conn().await?; conn.transaction(move |mut tx| { Box::pin(async move { + let updated = batch.iter().filter(|task| task.current_status == DelegatedJobStatus::Paused || task.current_status == DelegatedJobStatus::InProgress); update_last_status_sent( &mut tx, updated, ) .await?; + + let completed = batch.iter().filter(|task| task.current_status == DelegatedJobStatus::Completed || task.current_status == DelegatedJobStatus::Failed || task.current_status == DelegatedJobStatus::Cancelled); delete_external_tasks( &mut tx, completed, @@ -501,8 +517,8 @@ struct OutOfDateTaskRow { } async fn select_out_of_date_tasks( - mut conn: impl Connection, -) -> sqlx::Result> { + conn: &mut impl Connection, +) -> BoxStream<'_, sqlx::Result> { sqlx::query_as!( OutOfDateTaskRow, r#"WITH out_of_date_tasks AS ( @@ -529,24 +545,18 @@ async fn select_out_of_date_tasks( ) AS "current_status!: DelegatedJobStatus" FROM out_of_date_tasks AS t "#) - .fetch_all(conn.get_inner()) - .await + .fetch(conn.get_inner()) } async fn update_last_status_sent( mut conn: impl WriterConnection, - tasks: Vec<&OutOfDateTaskRow>, + tasks: impl Iterator, ) -> sqlx::Result<()> { - let tasks = tasks - .iter() - .map(|t| (t.current_status, t.task_id.clone())) - .collect::>(); - - for (current_status, task_id) in tasks { + for task in tasks { sqlx::query!( "UPDATE submissions_external_task SET last_status_sent = $1 WHERE task_id = $2", - current_status, - task_id, + task.current_status, + task.task_id, ) .execute(conn.get_inner()) .await?; @@ -557,14 +567,12 @@ async fn update_last_status_sent( async fn delete_external_tasks( mut conn: impl WriterConnection, - tasks: Vec<&OutOfDateTaskRow>, + tasks: impl Iterator, ) -> sqlx::Result<()> { - let tasks = tasks.iter().map(|t| t.task_id.clone()).collect::>(); - - for task_id in tasks { + for task in tasks { sqlx::query!( "DELETE FROM submissions_external_task WHERE task_id = $1", - task_id, + task.task_id, ) .execute(conn.get_inner()) .await?; From 01462ebde657ccca6c78d6afe919c299820466d0 Mon Sep 17 00:00:00 2001 From: Marten Wijnja Date: Fri, 18 Sep 2026 09:53:34 +0200 Subject: [PATCH 3/3] Cleaned up --- opsqueue/src/db/mod.rs | 5 +- opsqueue/src/delegation/server.rs | 367 +++++++++++++++++++----------- 2 files changed, 241 insertions(+), 131 deletions(-) diff --git a/opsqueue/src/db/mod.rs b/opsqueue/src/db/mod.rs index 01636f9..f3db5e3 100644 --- a/opsqueue/src/db/mod.rs +++ b/opsqueue/src/db/mod.rs @@ -128,9 +128,8 @@ pub trait Connection { #[allow(async_fn_in_trait)] async fn transaction(&mut self, f: F) -> Result where - for<'t> F: FnOnce(Conn>) -> BoxFuture<'t, Result> - + Send - + Sync, + for<'t> F: + FnOnce(Conn>) -> BoxFuture<'t, Result> + Send + Sync, O: Send, E: From + Send, { diff --git a/opsqueue/src/delegation/server.rs b/opsqueue/src/delegation/server.rs index 590a29b..b53f0f5 100644 --- a/opsqueue/src/delegation/server.rs +++ b/opsqueue/src/delegation/server.rs @@ -6,8 +6,10 @@ use axum::extract::State; use axum::http::StatusCode; use axum::routing::post; use axum::{Json, Router}; -use futures::Stream; use futures::stream::BoxStream; +use futures::{FutureExt, Stream}; +use itertools::Either; +use itertools::Itertools; use std::pin::Pin; use std::sync::Arc; use tokio::select; @@ -32,8 +34,8 @@ pub(crate) fn app_for_tests( notify_on_insert, notify_on_submission_change, ) - .run_background() - .build_router(); + .run_background() + .build_router(); Router::new().nest("/job", router) } @@ -85,8 +87,8 @@ impl ServerState { state, cancellation_token, ) - .await - .ok(); + .await + .ok(); }); self } @@ -323,8 +325,8 @@ async fn handle_kill_event(conn: &mut impl WriterConnection, task_id: &str) -> s WHERE task_id = $1"#, task_id, ) - .fetch_optional(conn.get_inner()) - .await?; + .fetch_optional(conn.get_inner()) + .await?; let Some(submission_id) = submission_id else { tracing::warn!(%task_id, "Kill event for unknown task_id; ignoring"); @@ -375,118 +377,198 @@ async fn run_in_background( Ok(()) } -async fn report_submission_status2(state: &ServerState, triggered_by_timeout: bool,) -> anyhow::Result<()> { - let mut conn = state.pool.reader_conn().await?; - let out_of_date_tasks = select_out_of_date_tasks(&mut conn).await; - for Ok(batch) in out_of_date_tasks.try_chunks(2048).await? { - - } - - Ok(()) -} - async fn report_submission_status( state: &ServerState, triggered_by_timeout: bool, ) -> anyhow::Result<()> { - let out_of_date_tasks = { - let conn = state.pool.reader_conn().await?; - select_out_of_date_tasks(conn).await? - }; + use futures::{StreamExt, TryStreamExt}; + let mut conn = state.pool.reader_conn().await?; + let mut out_of_date_tasks = select_out_of_date_tasks(&mut conn).await.try_chunks(2048); - if out_of_date_tasks.is_empty() { - return Ok(()); - } + // TODO: triggered_by_timeout early-return ? - if triggered_by_timeout { - tracing::warn!( - n_out_of_date_tasks = out_of_date_tasks.len(), - "Delegation background loop triggered by timeout with pending tasks; \ - possible missing notify_on_submission_change call" - ); + while let Some(batch_or_error) = out_of_date_tasks.next().await { + let batch = batch_or_error?; + + send_status_update_events(state, batch.iter()).await?; + update_statuses_in_db(&state.pool, batch).await?; } + Ok(()) +} - for batch in out_of_date_tasks.chunks(2048) { - let mut updates = Vec::new(); - let mut completions = Vec::new(); - - for task in batch { - match task.current_status { - DelegatedJobStatus::Paused => updates.push(DelegatedJobUpdate { - task_id: &task.task_id, - status: DelegatedJobUpdateStatus::Queued, - }), - DelegatedJobStatus::InProgress => updates.push(DelegatedJobUpdate { - task_id: &task.task_id, - status: DelegatedJobUpdateStatus::Running, - }), - DelegatedJobStatus::Completed => completions.push(DelegatedJobCompletion { - task_id: &task.task_id, - completion: DelegatedJobCompletionStatus::Success, - }), - DelegatedJobStatus::Failed => completions.push(DelegatedJobCompletion { - task_id: &task.task_id, - completion: DelegatedJobCompletionStatus::Failure { - failure_reason: FailureReason::Unknown, - }, - }), - DelegatedJobStatus::Cancelled => completions.push(DelegatedJobCompletion { - task_id: &task.task_id, - completion: DelegatedJobCompletionStatus::Failure { - failure_reason: FailureReason::Forced, - }, - }), +async fn update_statuses_in_db(pool: &DBPools, tasks: Vec) -> anyhow::Result<()> { + let mut conn = pool.writer_conn().await?; + // NOTE: Instead we could also run many short transactions + // I don't think we need atomicity between task updates + // (swapping the nesting of `conn.transaction` and the `for task in batch`) + conn.transaction(|mut tx| { + async move { + for task in tasks { + match task.current_status { + DelegatedJobStatus::Paused | DelegatedJobStatus::InProgress => { + update_single_last_status_sent(&mut tx, &task).await?; + } + DelegatedJobStatus::Completed + | DelegatedJobStatus::Failed + | DelegatedJobStatus::Cancelled => { + delete_external_task(&mut tx, &task).await?; + } + } } + Ok::<(), sqlx::Error>(()) } + .boxed() + }) + .await?; + Ok(()) +} - let events = vec![ - MasterDelegationEvent::Updated(updates), - MasterDelegationEvent::Completed(completions), - ]; - - send_events(state, &events).await?; - - // let updated = batch - // .iter() - // .filter(|task| { - // task.current_status == DelegatedJobStatus::Paused - // || task.current_status == DelegatedJobStatus::InProgress - // }) - // .collect(); - - // let completed = batch - // .iter() - // .filter(|task| { - // task.current_status == DelegatedJobStatus::Completed - // || task.current_status == DelegatedJobStatus::Failed - // || task.current_status == DelegatedJobStatus::Cancelled - // }) - // .collect(); - - let mut conn = state.pool.writer_conn().await?; - conn.transaction(move |mut tx| { - Box::pin(async move { - let updated = batch.iter().filter(|task| task.current_status == DelegatedJobStatus::Paused || task.current_status == DelegatedJobStatus::InProgress); - update_last_status_sent( - &mut tx, - updated, - ) - .await?; - - let completed = batch.iter().filter(|task| task.current_status == DelegatedJobStatus::Completed || task.current_status == DelegatedJobStatus::Failed || task.current_status == DelegatedJobStatus::Cancelled); - delete_external_tasks( - &mut tx, - completed, - ).await?; - - Ok::<(), sqlx::Error>(()) - }) - }).await?; - } - +async fn send_status_update_events( + state: &ServerState, + tasks: impl Iterator, +) -> anyhow::Result<()> { + let (updates, completions): (Vec<_>, Vec<_>) = tasks.partition_map(status_update); + let events = vec![ + MasterDelegationEvent::Updated(updates), + MasterDelegationEvent::Completed(completions), + ]; + send_events(state, &events).await?; Ok(()) } +fn status_update( + task: &OutOfDateTaskRow, +) -> Either, DelegatedJobCompletion<'_>> { + match task.current_status { + DelegatedJobStatus::Paused => Either::Left(DelegatedJobUpdate { + task_id: &task.task_id, + status: DelegatedJobUpdateStatus::Queued, + }), + DelegatedJobStatus::InProgress => Either::Left(DelegatedJobUpdate { + task_id: &task.task_id, + status: DelegatedJobUpdateStatus::Running, + }), + DelegatedJobStatus::Completed => Either::Right(DelegatedJobCompletion { + task_id: &task.task_id, + completion: DelegatedJobCompletionStatus::Success, + }), + DelegatedJobStatus::Failed => Either::Right(DelegatedJobCompletion { + task_id: &task.task_id, + completion: DelegatedJobCompletionStatus::Failure { + failure_reason: FailureReason::Unknown, + }, + }), + DelegatedJobStatus::Cancelled => Either::Right(DelegatedJobCompletion { + task_id: &task.task_id, + completion: DelegatedJobCompletionStatus::Failure { + failure_reason: FailureReason::Forced, + }, + }), + } +} + +// async fn report_submission_status( +// state: &ServerState, +// triggered_by_timeout: bool, +// ) -> anyhow::Result<()> { +// let out_of_date_tasks = { +// let conn = state.pool.reader_conn().await?; +// select_out_of_date_tasks(&mut conn).await? +// }; + +// if out_of_date_tasks.is_empty() { +// return Ok(()); +// } + +// if triggered_by_timeout { +// tracing::warn!( +// n_out_of_date_tasks = out_of_date_tasks.len(), +// "Delegation background loop triggered by timeout with pending tasks; \ +// possible missing notify_on_submission_change call" +// ); +// } + +// for batch in out_of_date_tasks.chunks(2048) { +// let mut updates = Vec::new(); +// let mut completions = Vec::new(); + +// for task in batch { +// match task.current_status { +// DelegatedJobStatus::Paused => updates.push(DelegatedJobUpdate { +// task_id: &task.task_id, +// status: DelegatedJobUpdateStatus::Queued, +// }), +// DelegatedJobStatus::InProgress => updates.push(DelegatedJobUpdate { +// task_id: &task.task_id, +// status: DelegatedJobUpdateStatus::Running, +// }), +// DelegatedJobStatus::Completed => completions.push(DelegatedJobCompletion { +// task_id: &task.task_id, +// completion: DelegatedJobCompletionStatus::Success, +// }), +// DelegatedJobStatus::Failed => completions.push(DelegatedJobCompletion { +// task_id: &task.task_id, +// completion: DelegatedJobCompletionStatus::Failure { +// failure_reason: FailureReason::Unknown, +// }, +// }), +// DelegatedJobStatus::Cancelled => completions.push(DelegatedJobCompletion { +// task_id: &task.task_id, +// completion: DelegatedJobCompletionStatus::Failure { +// failure_reason: FailureReason::Forced, +// }, +// }), +// } +// } + +// let events = vec![ +// MasterDelegationEvent::Updated(updates), +// MasterDelegationEvent::Completed(completions), +// ]; + +// send_events(state, &events).await?; + +// // let updated = batch +// // .iter() +// // .filter(|task| { +// // task.current_status == DelegatedJobStatus::Paused +// // || task.current_status == DelegatedJobStatus::InProgress +// // }) +// // .collect(); + +// // let completed = batch +// // .iter() +// // .filter(|task| { +// // task.current_status == DelegatedJobStatus::Completed +// // || task.current_status == DelegatedJobStatus::Failed +// // || task.current_status == DelegatedJobStatus::Cancelled +// // }) +// // .collect(); + +// let mut conn = state.pool.writer_conn().await?; +// conn.transaction(move |mut tx| { +// Box::pin(async move { +// let updated = batch.iter().filter(|task| task.current_status == DelegatedJobStatus::Paused || task.current_status == DelegatedJobStatus::InProgress); +// update_last_status_sent( +// &mut tx, +// updated, +// ) +// .await?; + +// let completed = batch.iter().filter(|task| task.current_status == DelegatedJobStatus::Completed || task.current_status == DelegatedJobStatus::Failed || task.current_status == DelegatedJobStatus::Cancelled); +// delete_external_tasks( +// &mut tx, +// completed, +// ).await?; + +// Ok::<(), sqlx::Error>(()) +// }) +// }).await?; +// } + +// Ok(()) +// } + async fn insert_external_task( mut conn: impl Connection, submission_id: SubmissionId, @@ -503,9 +585,9 @@ async fn insert_external_task( submission_id, task_id, ) - .execute(conn.get_inner()) - .await? - .rows_affected(); + .execute(conn.get_inner()) + .await? + .rows_affected(); Ok(rows_affected) } @@ -518,7 +600,7 @@ struct OutOfDateTaskRow { async fn select_out_of_date_tasks( conn: &mut impl Connection, -) -> BoxStream<'_, sqlx::Result> { +) -> BoxStream<'_, sqlx::Result> { sqlx::query_as!( OutOfDateTaskRow, r#"WITH out_of_date_tasks AS ( @@ -558,13 +640,28 @@ async fn update_last_status_sent( task.current_status, task.task_id, ) - .execute(conn.get_inner()) - .await?; + .execute(conn.get_inner()) + .await?; } Ok(()) } +async fn update_single_last_status_sent( + mut conn: impl WriterConnection, + task: &OutOfDateTaskRow, +) -> sqlx::Result<()> { + sqlx::query!( + "UPDATE submissions_external_task SET last_status_sent = $1 WHERE task_id = $2", + task.current_status, + task.task_id, + ) + .execute(conn.get_inner()) + .await?; + + Ok(()) +} + async fn delete_external_tasks( mut conn: impl WriterConnection, tasks: impl Iterator, @@ -574,13 +671,27 @@ async fn delete_external_tasks( "DELETE FROM submissions_external_task WHERE task_id = $1", task.task_id, ) - .execute(conn.get_inner()) - .await?; + .execute(conn.get_inner()) + .await?; } Ok(()) } +async fn delete_external_task( + mut conn: impl WriterConnection, + task: &OutOfDateTaskRow, +) -> sqlx::Result<()> { + sqlx::query!( + "DELETE FROM submissions_external_task WHERE task_id = $1", + task.task_id, + ) + .execute(conn.get_inner()) + .await?; + + Ok(()) +} + async fn send_events( state: &ServerState, events: &Vec>, @@ -700,8 +811,8 @@ pub mod test { InitialSubmissionStatus::Paused, &mut conn, ) - .await - .unwrap() + .await + .unwrap() }; { @@ -735,7 +846,7 @@ pub mod test { submission_id: submission, }, }) - .unwrap(), + .unwrap(), )) .unwrap(), ) @@ -800,8 +911,8 @@ pub mod test { InitialSubmissionStatus::Paused, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await @@ -897,8 +1008,8 @@ pub mod test { InitialSubmissionStatus::Paused, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await @@ -972,8 +1083,8 @@ pub mod test { InitialSubmissionStatus::InProgress, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await @@ -1049,8 +1160,8 @@ pub mod test { InitialSubmissionStatus::InProgress, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await @@ -1076,8 +1187,8 @@ pub mod test { &mut conn, 0, ) - .await - .unwrap(); + .await + .unwrap(); notify_on_submission_change.notify_one(); } @@ -1129,8 +1240,8 @@ pub mod test { InitialSubmissionStatus::InProgress, &mut conn, ) - .await - .unwrap(); + .await + .unwrap(); insert_external_task(&mut conn, submission, "test") .await