From 916c8bee5362ca12c18067fc725d3b03c75e18db Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 8 Jul 2026 17:29:53 +0200 Subject: [PATCH 1/7] Draft --- opsqueue/src/common/submission.rs | 1 + opsqueue/src/delegation/mod.rs | 2 + opsqueue/src/delegation/server.rs | 101 ++++++++++++++++++++++++++++++ opsqueue/src/lib.rs | 3 + 4 files changed, 107 insertions(+) create mode 100644 opsqueue/src/delegation/mod.rs create mode 100644 opsqueue/src/delegation/server.rs diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 69040f1a..1f0e389e 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -846,6 +846,7 @@ pub mod db { failure: String, mut conn: impl WriterConnection, ) -> sqlx::Result<()> { + // TODO: Update JM Master at /delegation/complete conn.transaction(move |mut tx| { Box::pin( async move { fail_submission_notx(id, failed_chunk_index, failure, &mut tx).await }, diff --git a/opsqueue/src/delegation/mod.rs b/opsqueue/src/delegation/mod.rs new file mode 100644 index 00000000..d46f6a49 --- /dev/null +++ b/opsqueue/src/delegation/mod.rs @@ -0,0 +1,2 @@ +#[cfg(feature = "server-logic")] +pub mod server; diff --git a/opsqueue/src/delegation/server.rs b/opsqueue/src/delegation/server.rs new file mode 100644 index 00000000..2afa8712 --- /dev/null +++ b/opsqueue/src/delegation/server.rs @@ -0,0 +1,101 @@ +use crate::common::submission::{Metadata, SubmissionId}; +use crate::common::{chunk, submission, StrategicMetadataMap}; +use crate::config::Config; +use crate::db; +use crate::db::DBPools; +use crate::producer::{ChunkContents, InsertSubmission}; +use axum::extract::State; +use axum::routing::post; +use axum::{Json, Router}; +use std::sync::Arc; +use tokio::select; +use tokio::sync::Notify; +use tokio_util::sync::CancellationToken; + +#[derive(Debug, Clone)] +pub struct ServerState { + pool: DBPools, + config: &'static Config, +} + +impl ServerState { + pub fn new(pool: DBPools, config: &'static Config) -> Self { + Self { pool, config } + } + + pub fn build_router(self: ServerState) -> Router<()> { + Router::new() + .route("/job/delegate", post(job_delegate)) + .route("/job/kill", post(job_kill)) + .route("/job/return", post(job_return)) + .with_state(self) + } +} + +#[derive(Debug, serde::Serialize, serde::Deserialize)] +#[serde(transparent)] +pub struct TaskId(pub String); + +#[derive(Debug, serde::Serialize, serde::Deserialize)] +pub struct JobDelegationPayload { + pub submission_id: SubmissionId, +} +#[derive(Debug, serde::Serialize, serde::Deserialize)] +pub struct JobDelegation { + // NOTE: We ignore most fields, since we do not care about them. + pub payload: JobDelegationPayload, + pub task_id: TaskId, +} + +// TODO(delegation): Add a `jm_task_id: Option` column to the `submission` and +// `submissions_{completed,failed,canceled}` table. Change the clean-up worker s.t. it never +// deletes a row +// that has `jm_task_id` set. These should be handled in the `run_in_background` function in this +// module instead. + +#[tracing::instrument(level = "debug", skip(state))] +async fn job_delegate( + State(state): State, + Json(request): Json, +) -> Result<(), ()> { + // TODO(delegation): implement this. + todo!("check that the submission exists and set the `jm_task_id` column if it is not populated yet."); +} + +#[tracing::instrument(level = "debug", skip(state))] +async fn job_kill( + State(state): State, + Json(request): Json>, +) -> Result<(), ()> { + todo!() +} + +#[tracing::instrument(level = "debug", skip(state))] +async fn job_return( + State(state): State, + Json(request): Json>, +) -> Result<(), ()> { + todo!() +} + +const JOBMACHINE_DELEGATION_BACKGROUND_LOOP: std::time::Duration = + std::time::Duration::from_secs(5); + +async fn run_in_background(notify_on_submission_change: Arc, cancellation_token: CancellationToken) -> Result<(), ()> { + // TODO(delegation): Implement this. + loop { + let triggered_by_timeout: bool = select! { + _ = cancellation_token.cancelled() => break, + _ = notify_on_submission_change.notified() => false, + _ = tokio::time::sleep(JOBMACHINE_DELEGATION_BACKGROUND_LOOP) => true, + }; + // Look for submissions that were completed/failed/cancelled that haven't been reported to JM yet. + // That is, rows in `submissions_{completed,failed,canceled}` where `jm_task_id IS NOT NULL`. + + // Make the necessary calls to Jobmachine. + + // If `triggered_by_timeout`, log a warning that there might be a missing call to `notify_on_submission_change`. + }; + + Ok(()) +} diff --git a/opsqueue/src/lib.rs b/opsqueue/src/lib.rs index 0389bbdb..919a269e 100644 --- a/opsqueue/src/lib.rs +++ b/opsqueue/src/lib.rs @@ -39,6 +39,9 @@ pub mod prometheus; #[cfg(feature = "server-logic")] pub mod config; +#[cfg(feature = "server-logic")] +pub mod delegation; + /// The Opsqueue library's semantic version /// as written in the Rust packages's `Cargo.toml` pub const VERSION_CARGO_SEMVER: &str = env!("CARGO_PKG_VERSION"); From 48585834f141c12a8fe0b560013151a0e147691d Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 9 Jul 2026 14:02:53 +0200 Subject: [PATCH 2/7] Vibe it --- opsqueue/Cargo.toml | 3 +- .../20260709130000_add_jm_task_id.down.sql | 4 + .../20260709130000_add_jm_task_id.up.sql | 4 + opsqueue/opsqueue_example_database_schema.db | Bin 102400 -> 102400 bytes opsqueue/src/common/submission.rs | 39 ++- opsqueue/src/config.rs | 9 + opsqueue/src/delegation/server.rs | 295 ++++++++++++++++-- opsqueue/src/server.rs | 8 +- 8 files changed, 324 insertions(+), 38 deletions(-) create mode 100644 opsqueue/migrations/20260709130000_add_jm_task_id.down.sql create mode 100644 opsqueue/migrations/20260709130000_add_jm_task_id.up.sql diff --git a/opsqueue/Cargo.toml b/opsqueue/Cargo.toml index d4f3e5e8..c341c40d 100644 --- a/opsqueue/Cargo.toml +++ b/opsqueue/Cargo.toml @@ -102,7 +102,8 @@ server-logic = [ "dep:tower-http", "dep:axum-prometheus", "dep:sentry", - "dep:sentry-tracing" + "dep:sentry-tracing", + "dep:reqwest", ] # Dependencies only in use by the client libraries: client-logic = [ diff --git a/opsqueue/migrations/20260709130000_add_jm_task_id.down.sql b/opsqueue/migrations/20260709130000_add_jm_task_id.down.sql new file mode 100644 index 00000000..3ace4725 --- /dev/null +++ b/opsqueue/migrations/20260709130000_add_jm_task_id.down.sql @@ -0,0 +1,4 @@ +ALTER TABLE submissions DROP COLUMN jm_task_id; +ALTER TABLE submissions_completed DROP COLUMN jm_task_id; +ALTER TABLE submissions_failed DROP COLUMN jm_task_id; +ALTER TABLE submissions_cancelled DROP COLUMN jm_task_id; diff --git a/opsqueue/migrations/20260709130000_add_jm_task_id.up.sql b/opsqueue/migrations/20260709130000_add_jm_task_id.up.sql new file mode 100644 index 00000000..8428c623 --- /dev/null +++ b/opsqueue/migrations/20260709130000_add_jm_task_id.up.sql @@ -0,0 +1,4 @@ +ALTER TABLE submissions ADD COLUMN jm_task_id TEXT; +ALTER TABLE submissions_completed ADD COLUMN jm_task_id TEXT; +ALTER TABLE submissions_failed ADD COLUMN jm_task_id TEXT; +ALTER TABLE submissions_cancelled ADD COLUMN jm_task_id TEXT; diff --git a/opsqueue/opsqueue_example_database_schema.db b/opsqueue/opsqueue_example_database_schema.db index ae49303b2d692551e8a9a018ee8a0af45eadaebf..e9d5a596d049016efe22e2b0467dfc39dd5be92a 100644 GIT binary patch delta 311 zcmZozz}B#UZGyC*Bm)Bj2*ZHJL>*&RNd~>L&c>9jDU1vGMVbFI@Ll1X%r467%^l0j zz{$w7oB997x?_x+4Fu$ww*MDq6k}rH@MdI~E-A+NhpWk$nO$62nXyH8a}O^Y3R$`FC5gq^@tG+KA+8Z2n!HU6+a^xj#MKhX&Mt0f$k-9PU0#e)n`!z>VMb9! zO!ZCCtnA|E=8P?&C5cHnsnhc%7)7@O2@OUeOvT%G%P@K~OK~#r#qoUQ{>Poi^@_cW xZ8J+OPbzaG(_%)>=?V&rFZr8X#2{v>L(EK61e$55$SBAth+*b-ecHS0FmQJHUIzs diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 1f0e389e..11685bdb 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -748,8 +748,8 @@ pub mod db { let submission_opt = query!( " INSERT INTO submissions_cancelled - (id, chunks_total, prefix, metadata, cancelled_at, chunks_done) - SELECT id, chunks_total, prefix, metadata, julianday($1), chunks_done FROM submissions WHERE id = $2; + (id, chunks_total, prefix, metadata, cancelled_at, chunks_done, jm_task_id) + SELECT id, chunks_total, prefix, metadata, julianday($1), chunks_done, jm_task_id FROM submissions WHERE id = $2; DELETE FROM submissions WHERE id = $3 RETURNING *; ", @@ -783,8 +783,8 @@ pub mod db { SAVEPOINT complete_submission_raw; INSERT INTO submissions_completed - (id, chunks_total, prefix, metadata, completed_at) - SELECT id, chunks_total, prefix, metadata, julianday($1) FROM submissions WHERE id = $2; + (id, chunks_total, prefix, metadata, completed_at, jm_task_id) + SELECT id, chunks_total, prefix, metadata, julianday($1), jm_task_id FROM submissions WHERE id = $2; DELETE FROM submissions WHERE id = $3 RETURNING *; @@ -819,8 +819,8 @@ pub mod db { query!( " INSERT INTO submissions_failed - (id, chunks_total, chunks_done, prefix, metadata, failed_at, failed_chunk_id) - SELECT id, chunks_total, chunks_done, prefix, metadata, julianday($1), $2 FROM submissions WHERE id = $3; + (id, chunks_total, chunks_done, prefix, metadata, failed_at, failed_chunk_id, jm_task_id) + SELECT id, chunks_total, chunks_done, prefix, metadata, julianday($1), $2, jm_task_id FROM submissions WHERE id = $3; DELETE FROM submissions WHERE id = $4 RETURNING *; ", @@ -909,11 +909,14 @@ pub mod db { tracing::info!("Cleaning up old completed/failed submissions..."); conn.transaction(move |mut tx| { Box::pin(async move { - // Clean up old submissions_metadata + // Clean up old submissions_metadata. + // Skip metadata for submissions with jm_task_id set; those are cleaned up by + // the delegation background loop after reporting to the JM master. query!( "DELETE FROM submissions_metadata WHERE submission_id IN ( - SELECT id FROM submissions_completed WHERE completed_at < julianday($1) + SELECT id FROM submissions_completed + WHERE completed_at < julianday($1) AND jm_task_id IS NULL );", older_than ) @@ -922,7 +925,8 @@ pub mod db { query!( "DELETE FROM submissions_metadata WHERE submission_id IN ( - SELECT id FROM submissions_failed WHERE failed_at < julianday($1) + SELECT id FROM submissions_failed + WHERE failed_at < julianday($1) AND jm_task_id IS NULL );", older_than ) @@ -931,33 +935,40 @@ pub mod db { query!( "DELETE FROM submissions_metadata WHERE submission_id IN ( - SELECT id FROM submissions_cancelled WHERE cancelled_at < julianday($1) + SELECT id FROM submissions_cancelled + WHERE cancelled_at < julianday($1) AND jm_task_id IS NULL );", older_than ) .execute(tx.get_inner()) .await?; - // Clean up old submissions: + // Clean up old submissions. Skip rows with jm_task_id set; the delegation + // background loop is responsible for reporting those to the JM master and + // deleting them afterwards. let n_submissions_completed = query!( - "DELETE FROM submissions_completed WHERE completed_at < julianday($1);", + "DELETE FROM submissions_completed + WHERE completed_at < julianday($1) AND jm_task_id IS NULL;", older_than ) .execute(tx.get_inner()) .await?.rows_affected(); let n_submissions_failed = query!( - "DELETE FROM submissions_failed WHERE failed_at < julianday($1);", + "DELETE FROM submissions_failed + WHERE failed_at < julianday($1) AND jm_task_id IS NULL;", older_than ) .execute(tx.get_inner()) .await?.rows_affected(); let n_submissions_cancelled = query!( - "DELETE FROM submissions_cancelled WHERE cancelled_at < julianday($1);", + "DELETE FROM submissions_cancelled + WHERE cancelled_at < julianday($1) AND jm_task_id IS NULL;", older_than ) .execute(tx.get_inner()) .await?.rows_affected(); + // Clean up old chunks. let n_chunks_completed = query!( "DELETE FROM chunks_completed WHERE completed_at < julianday($1);", older_than diff --git a/opsqueue/src/config.rs b/opsqueue/src/config.rs index 08acd428..bd98d556 100644 --- a/opsqueue/src/config.rs +++ b/opsqueue/src/config.rs @@ -77,6 +77,14 @@ pub struct Config { #[arg(long, default_value = "1 hour")] pub max_submission_age: humantime::Duration, + + /// Base URL of the Jobmachine master, used by the delegation server to report + /// completed/failed/cancelled tasks back to the master via `/delegation/complete`. + /// + /// When not set, the delegation endpoints (`/job/delegate`, `/job/kill`, `/job/return`) + /// still accept requests but completed tasks are never reported back to the master. + #[arg(long)] + pub jm_master_url: Option, } impl Default for Config { @@ -101,6 +109,7 @@ impl Default for Config { max_missable_heartbeats, max_chunk_retries, max_submission_age, + jm_master_url: None, } } } diff --git a/opsqueue/src/delegation/server.rs b/opsqueue/src/delegation/server.rs index 2afa8712..c2923ca6 100644 --- a/opsqueue/src/delegation/server.rs +++ b/opsqueue/src/delegation/server.rs @@ -1,10 +1,8 @@ -use crate::common::submission::{Metadata, SubmissionId}; -use crate::common::{chunk, submission, StrategicMetadataMap}; +use crate::common::submission::SubmissionId; use crate::config::Config; -use crate::db; -use crate::db::DBPools; -use crate::producer::{ChunkContents, InsertSubmission}; +use crate::db::{Connection as _, DBPools}; use axum::extract::State; +use axum::http::StatusCode; use axum::routing::post; use axum::{Json, Router}; use std::sync::Arc; @@ -16,11 +14,34 @@ use tokio_util::sync::CancellationToken; pub struct ServerState { pool: DBPools, config: &'static Config, + /// Notified whenever a submission transitions to a terminal state (completed/failed/cancelled), + /// so the background loop can promptly report it to the JM master. + pub notify_on_submission_change: Arc, + http_client: reqwest::Client, } impl ServerState { pub fn new(pool: DBPools, config: &'static Config) -> Self { - Self { pool, config } + Self { + pool, + config, + notify_on_submission_change: Arc::new(Notify::new()), + http_client: reqwest::Client::new(), + } + } + + pub fn run_background(self, cancellation_token: CancellationToken) -> Self { + let state = self.clone(); + tokio::spawn(async move { + run_in_background( + state.notify_on_submission_change.clone(), + state, + cancellation_token, + ) + .await + .ok(); + }); + self } pub fn build_router(self: ServerState) -> Router<()> { @@ -47,19 +68,127 @@ pub struct JobDelegation { pub task_id: TaskId, } -// TODO(delegation): Add a `jm_task_id: Option` column to the `submission` and -// `submissions_{completed,failed,canceled}` table. Change the clean-up worker s.t. it never -// deletes a row -// that has `jm_task_id` set. These should be handled in the `run_in_background` function in this -// module instead. +#[derive(serde::Serialize)] +#[serde(tag = "status", rename_all = "snake_case")] +enum JobCompletionStatus { + Success, + Failure { failure_reason: &'static str }, +} + +#[derive(serde::Serialize)] +struct DelegatedJobCompletion { + task_id: String, + completion: JobCompletionStatus, +} #[tracing::instrument(level = "debug", skip(state))] async fn job_delegate( State(state): State, Json(request): Json, -) -> Result<(), ()> { - // TODO(delegation): implement this. - todo!("check that the submission exists and set the `jm_task_id` column if it is not populated yet."); +) -> Result<(), StatusCode> { + let mut conn = state.pool.writer_conn().await.map_err(|e| { + tracing::error!("DB error acquiring writer connection: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let submission_id = request.payload.submission_id; + let task_id = &request.task_id.0; + + // Attempt to set jm_task_id on the submission, in whichever table it currently lives. + // The AND jm_task_id IS NULL guard makes this idempotent: if the master retries the + // same request after an already-successful response, we simply skip the update. + let n1 = sqlx::query!( + "UPDATE submissions SET jm_task_id = $1 WHERE id = $2 AND jm_task_id IS NULL", + task_id, + submission_id, + ) + .execute(conn.get_inner()) + .await + .map_err(|e| { + tracing::error!("DB error setting jm_task_id on submissions: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })? + .rows_affected(); + + let n2 = sqlx::query!( + "UPDATE submissions_completed SET jm_task_id = $1 WHERE id = $2 AND jm_task_id IS NULL", + task_id, + submission_id, + ) + .execute(conn.get_inner()) + .await + .map_err(|e| { + tracing::error!("DB error setting jm_task_id on submissions_completed: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })? + .rows_affected(); + + let n3 = sqlx::query!( + "UPDATE submissions_failed SET jm_task_id = $1 WHERE id = $2 AND jm_task_id IS NULL", + task_id, + submission_id, + ) + .execute(conn.get_inner()) + .await + .map_err(|e| { + tracing::error!("DB error setting jm_task_id on submissions_failed: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })? + .rows_affected(); + + let n4 = sqlx::query!( + "UPDATE submissions_cancelled SET jm_task_id = $1 WHERE id = $2 AND jm_task_id IS NULL", + task_id, + submission_id, + ) + .execute(conn.get_inner()) + .await + .map_err(|e| { + tracing::error!("DB error setting jm_task_id on submissions_cancelled: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })? + .rows_affected(); + + let updated = n1 + n2 + n3 + n4; + + if updated > 0 { + // Trigger the background loop in case the submission already reached a terminal state. + state.notify_on_submission_change.notify_one(); + return Ok(()); + } + + // No rows were updated. Either jm_task_id was already set (retry) or the submission + // doesn't exist. Check existence across all tables to distinguish the two cases. + let count = sqlx::query!( + r#"SELECT COUNT(1) as "count: i64" FROM ( + SELECT id FROM submissions WHERE id = $1 + UNION ALL + SELECT id FROM submissions_completed WHERE id = $2 + UNION ALL + SELECT id FROM submissions_failed WHERE id = $3 + UNION ALL + SELECT id FROM submissions_cancelled WHERE id = $4 + )"#, + submission_id, + submission_id, + submission_id, + submission_id, + ) + .fetch_one(conn.get_inner()) + .await + .map_err(|e| { + tracing::error!("DB error checking submission existence: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })? + .count; + + if count > 0 { + // Submission exists; jm_task_id was already set on a previous request. Idempotent OK. + Ok(()) + } else { + tracing::warn!(%submission_id, "Received job_delegate for unknown submission"); + Err(StatusCode::NOT_FOUND) + } } #[tracing::instrument(level = "debug", skip(state))] @@ -67,6 +196,7 @@ async fn job_kill( State(state): State, Json(request): Json>, ) -> Result<(), ()> { + let _ = (state, request); todo!() } @@ -75,27 +205,150 @@ async fn job_return( State(state): State, Json(request): Json>, ) -> Result<(), ()> { + let _ = (state, request); todo!() } const JOBMACHINE_DELEGATION_BACKGROUND_LOOP: std::time::Duration = std::time::Duration::from_secs(5); -async fn run_in_background(notify_on_submission_change: Arc, cancellation_token: CancellationToken) -> Result<(), ()> { - // TODO(delegation): Implement this. +async fn run_in_background( + notify_on_submission_change: Arc, + state: ServerState, + cancellation_token: CancellationToken, +) -> Result<(), ()> { + let Some(jm_master_url) = &state.config.jm_master_url else { + return Ok(()); + }; + + match report_delegated_submissions(&state, jm_master_url, false).await { + Ok(()) => {} + Err(e) => tracing::error!("Error in delegation background loop: {e:?}"), + } + loop { let triggered_by_timeout: bool = select! { _ = cancellation_token.cancelled() => break, _ = notify_on_submission_change.notified() => false, _ = tokio::time::sleep(JOBMACHINE_DELEGATION_BACKGROUND_LOOP) => true, }; - // Look for submissions that were completed/failed/cancelled that haven't been reported to JM yet. - // That is, rows in `submissions_{completed,failed,canceled}` where `jm_task_id IS NOT NULL`. - // Make the necessary calls to Jobmachine. + match report_delegated_submissions(&state, jm_master_url, triggered_by_timeout).await { + Ok(()) => {} + Err(e) => tracing::error!("Error in delegation background loop: {e:?}"), + } + } - // If `triggered_by_timeout`, log a warning that there might be a missing call to `notify_on_submission_change`. - }; + Ok(()) +} + +/// Queries all terminal submissions with `jm_task_id IS NOT NULL`, reports them to the +/// JM master via `/delegation/complete`, then deletes the rows (along with their +/// associated chunks and metadata) on success. +async fn report_delegated_submissions( + state: &ServerState, + jm_master_url: &url::Url, + triggered_by_timeout: bool, +) -> anyhow::Result<()> { + let mut conn = state.pool.reader_conn().await?; + + let completed = sqlx::query!( + r#"SELECT id AS "id: SubmissionId", jm_task_id + FROM submissions_completed WHERE jm_task_id IS NOT NULL"# + ) + .fetch_all(conn.get_inner()) + .await?; + let failed = sqlx::query!( + r#"SELECT id AS "id: SubmissionId", jm_task_id + FROM submissions_failed WHERE jm_task_id IS NOT NULL"# + ) + .fetch_all(conn.get_inner()) + .await?; + + let cancelled = sqlx::query!( + r#"SELECT id AS "id: SubmissionId", jm_task_id + FROM submissions_cancelled WHERE jm_task_id IS NOT NULL"# + ) + .fetch_all(conn.get_inner()) + .await?; + + drop(conn); + + let tasks_to_report: Vec = completed + .iter() + .map(|r| DelegatedJobCompletion { + task_id: r.jm_task_id.clone().unwrap(), + completion: JobCompletionStatus::Success, + }) + .chain(failed.iter().map(|r| DelegatedJobCompletion { + task_id: r.jm_task_id.clone().unwrap(), + completion: JobCompletionStatus::Failure { + // TODO(delegation): Distinguish between "known" and "unknown". + failure_reason: "unknown", + }, + })) + .chain(cancelled.iter().map(|r| DelegatedJobCompletion { + task_id: r.jm_task_id.clone().unwrap(), + completion: JobCompletionStatus::Failure { + // TODO(delegation): What should the failure reason for cancellations be? + failure_reason: "unknown", + }, + })) + .collect(); + + if tasks_to_report.is_empty() { + return Ok(()); + } + + if triggered_by_timeout { + tracing::warn!( + n_tasks = tasks_to_report.len(), + "Delegation background loop triggered by timeout with pending tasks; \ + possible missing notify_on_submission_change call" + ); + } + + let url = format!( + "{}/delegation/complete", + jm_master_url.as_str().trim_end_matches('/') + ); + state + .http_client + .put(&url) + .json(&tasks_to_report) + .send() + .await? + .error_for_status()?; + + // Report succeeded; clear jm_task_id on all reported rows so the periodic cleanup + // worker can delete them in due course. + let mut conn = state.pool.writer_conn().await?; + clear_jm_task_id("submissions_completed", &completed.iter().map(|r| r.id).collect::>(), &mut conn).await?; + clear_jm_task_id("submissions_failed", &failed.iter().map(|r| r.id).collect::>(), &mut conn).await?; + clear_jm_task_id("submissions_cancelled", &cancelled.iter().map(|r| r.id).collect::>(), &mut conn).await?; + + Ok(()) +} + +/// Sets `jm_task_id = NULL` for all given submission IDs in the specified table, +/// making those rows eligible for deletion by the periodic cleanup worker. +async fn clear_jm_task_id( + table: &str, + ids: &[SubmissionId], + conn: &mut impl crate::db::WriterConnection, +) -> sqlx::Result<()> { + if ids.is_empty() { + return Ok(()); + } + let mut qb = sqlx::QueryBuilder::new("UPDATE "); + qb.push(table); + qb.push(" SET jm_task_id = NULL WHERE id IN ("); + let mut sep = qb.separated(", "); + for id in ids { + sep.push_bind(*id); + } + qb.push(")"); + qb.build().execute(conn.get_inner()).await?; Ok(()) } diff --git a/opsqueue/src/server.rs b/opsqueue/src/server.rs index cfc9ec4b..37047933 100644 --- a/opsqueue/src/server.rs +++ b/opsqueue/src/server.rs @@ -81,11 +81,15 @@ pub fn build_router( .run_background() .build_router(); let producer_routes = - crate::producer::server::ServerState::new(pool, notify_on_insert).build_router(); + crate::producer::server::ServerState::new(pool.clone(), notify_on_insert).build_router(); + let delegation_routes = crate::delegation::server::ServerState::new(pool, config) + .run_background(cancellation_token.clone()) + .build_router(); let routes = Router::new() .nest("/producer", producer_routes) - .nest("/consumer", consumer_routes); + .nest("/consumer", consumer_routes) + .nest("/", delegation_routes); let tracing_middleware = tower_http::trace::TraceLayer::new_for_http() .make_span_with(|request: &http::Request<_>| { From 789a478829eb4e40e422325e0a0fce50da8d22b1 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 9 Jul 2026 16:57:17 +0200 Subject: [PATCH 3/7] Implement /job/kill --- ...20260709160000_jm_task_id_indexes.down.sql | 4 ++ .../20260709160000_jm_task_id_indexes.up.sql | 4 ++ opsqueue/opsqueue_example_database_schema.db | Bin 102400 -> 118784 bytes opsqueue/src/delegation/server.rs | 54 +++++++++++++++++- 4 files changed, 59 insertions(+), 3 deletions(-) create mode 100644 opsqueue/migrations/20260709160000_jm_task_id_indexes.down.sql create mode 100644 opsqueue/migrations/20260709160000_jm_task_id_indexes.up.sql diff --git a/opsqueue/migrations/20260709160000_jm_task_id_indexes.down.sql b/opsqueue/migrations/20260709160000_jm_task_id_indexes.down.sql new file mode 100644 index 00000000..c05312b2 --- /dev/null +++ b/opsqueue/migrations/20260709160000_jm_task_id_indexes.down.sql @@ -0,0 +1,4 @@ +DROP INDEX IF EXISTS submissions_by_jm_task_id; +DROP INDEX IF EXISTS submissions_completed_by_jm_task_id; +DROP INDEX IF EXISTS submissions_failed_by_jm_task_id; +DROP INDEX IF EXISTS submissions_cancelled_by_jm_task_id; diff --git a/opsqueue/migrations/20260709160000_jm_task_id_indexes.up.sql b/opsqueue/migrations/20260709160000_jm_task_id_indexes.up.sql new file mode 100644 index 00000000..9b3546f9 --- /dev/null +++ b/opsqueue/migrations/20260709160000_jm_task_id_indexes.up.sql @@ -0,0 +1,4 @@ +CREATE INDEX submissions_by_jm_task_id ON submissions (jm_task_id) WHERE jm_task_id IS NOT NULL; +CREATE INDEX submissions_completed_by_jm_task_id ON submissions_completed (jm_task_id) WHERE jm_task_id IS NOT NULL; +CREATE INDEX submissions_failed_by_jm_task_id ON submissions_failed (jm_task_id) WHERE jm_task_id IS NOT NULL; +CREATE INDEX submissions_cancelled_by_jm_task_id ON submissions_cancelled (jm_task_id) WHERE jm_task_id IS NOT NULL; diff --git a/opsqueue/opsqueue_example_database_schema.db b/opsqueue/opsqueue_example_database_schema.db index e9d5a596d049016efe22e2b0467dfc39dd5be92a..bd172f8694fbde943eeb4b0fb6c90e6cf6529dd8 100644 GIT binary patch delta 677 zcmZozz}B#UeS);0ECT}r2*ZHRL>*&B*^LSF`5F0|1f~fvvZ(Mfd|}|b!Z(>+l-HX( zmY0E(k!Lsaf2LT5FKp}B#3 z-!3Z7_<=)~p^Sm+6?+-mW|mlVnoAXf#1h=OOZf}ekgf?ueQ&vs5l#z~A!+{br5QzeQJge+B8wQKG{`N=f*5WAtJfml8PijB7$>kWFfcCMIB_l`&~|AQ z+XECCMfp%{Pnix9W0VEiFEu?@k�XRUQ~3)BP10#Y~8IM5B%ryST3*V`KW{jjR$p z$ex<6D9I?yC<6}p>2}(T>;lNLDIm)nfhx3#w~T8uqrrdqMGgXs91OrDkO4vf4{*(K delta 176 zcmZozz}~QcZGyC*Bm)Bj2*ZHJL>*&B$&Cr~`58Hz1f~fvvWRgq|7YO4!Z(>+l-HX( zmY0E(k!Lsa|BZFW7&jXT$TMyKFU%;x#KFiQZfMBZAvm2|lJVPiQE|o(9I}j@4194s zU%CHr=W)GaFJs%x63dgy+{m<;k&`K4;>73MITaZvF)}eQY@e&dc#&hWHrvnX={k&k Xn>Vp@Gc$5+W;FOOzsNygk%IvM%i=85 diff --git a/opsqueue/src/delegation/server.rs b/opsqueue/src/delegation/server.rs index c2923ca6..43f5f4b9 100644 --- a/opsqueue/src/delegation/server.rs +++ b/opsqueue/src/delegation/server.rs @@ -1,6 +1,7 @@ use crate::common::submission::SubmissionId; use crate::config::Config; use crate::db::{Connection as _, DBPools}; +use crate::common::errors::E; use axum::extract::State; use axum::http::StatusCode; use axum::routing::post; @@ -195,9 +196,56 @@ async fn job_delegate( async fn job_kill( State(state): State, Json(request): Json>, -) -> Result<(), ()> { - let _ = (state, request); - todo!() +) -> Result<(), StatusCode> { + let mut conn = state.pool.writer_conn().await.map_err(|e| { + tracing::error!("DB error acquiring writer connection: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let mut any_cancelled = false; + + for task_id in &request { + let task_id = &task_id.0; + // Find the submission ID for this task, if it is still active. + let row = sqlx::query!( + r#"SELECT id AS "id: SubmissionId" FROM submissions WHERE jm_task_id = $1"#, + task_id, + ) + .fetch_optional(conn.get_inner()) + .await + .map_err(|e| { + tracing::error!(%task_id, "DB error looking up submission for kill: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let Some(row) = row else { + // Either already in a terminal table (background loop will report it) or unknown. + // Both are fine — kill is best-effort. + tracing::debug!(%task_id, "job_kill: submission not active, skipping"); + continue; + }; + + match crate::common::submission::db::cancel_submission(row.id, &mut conn).await { + Ok(()) => { + tracing::info!(%task_id, submission_id = %row.id, "Cancelled submission on kill request"); + any_cancelled = true; + } + // Already in a terminal state — nothing to do. + Err(E::R(_)) => { + tracing::debug!(%task_id, submission_id = %row.id, "job_kill: submission already terminal"); + } + Err(E::L(e)) => { + tracing::error!(%task_id, submission_id = %row.id, "DB error cancelling submission: {e:?}"); + return Err(StatusCode::INTERNAL_SERVER_ERROR); + } + } + } + + if any_cancelled { + state.notify_on_submission_change.notify_one(); + } + + Ok(()) } #[tracing::instrument(level = "debug", skip(state))] From b44f28fea7e94401bff17edd0995f0646dd2a71d Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 9 Jul 2026 17:14:29 +0200 Subject: [PATCH 4/7] Implement submission pausing and job_return - Add `paused INTEGER NOT NULL DEFAULT 0` column to `submissions` table (migration 20260709170000_add_paused_to_submissions) - Add `paused: bool` field to `Submission` struct; thread through `insert_submission_raw`, `get_submission`, `submission_status`, and `insert_submission_from_chunks` (new `paused` parameter) - Add `#[serde(default)] paused: bool` to `InsertSubmission` so producers can create submissions in a paused state; default is `false` for backward compat - Filter paused submissions out of all chunk dispatch strategies: Oldest, Newest, and both halves of Random use `WHERE submission_id NOT IN (SELECT id FROM submissions WHERE paused)`; PreferDistinct inherits the filter via its CTE - Update strategy query-plan snapshots to reflect the new subquery - `job_delegate` now also clears `paused = 0` on the active submission and calls `notify_on_insert` to wake consumers; add `notify_on_insert: Arc` to delegation `ServerState` and thread it from `server.rs` - Implement `job_return`: for each task_id, find the still-paused submission, clear `jm_task_id` (so the cancelled row is not picked up by the background reporting loop), cancel the submission, then PUT the returned task IDs to `{jm_master_url}/delegation/return` Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- ...9170000_add_paused_to_submissions.down.sql | 1 + ...709170000_add_paused_to_submissions.up.sql | 1 + opsqueue/opsqueue_example_database_schema.db | Bin 118784 -> 118784 bytes opsqueue/src/common/chunk.rs | 1 + opsqueue/src/common/submission.rs | 24 +- opsqueue/src/consumer/strategy.rs | 318 ++++++++++-------- opsqueue/src/delegation/server.rs | 102 +++++- opsqueue/src/producer/client.rs | 2 + opsqueue/src/producer/common.rs | 4 + opsqueue/src/producer/server.rs | 1 + opsqueue/src/server.rs | 4 +- 11 files changed, 312 insertions(+), 146 deletions(-) create mode 100644 opsqueue/migrations/20260709170000_add_paused_to_submissions.down.sql create mode 100644 opsqueue/migrations/20260709170000_add_paused_to_submissions.up.sql diff --git a/opsqueue/migrations/20260709170000_add_paused_to_submissions.down.sql b/opsqueue/migrations/20260709170000_add_paused_to_submissions.down.sql new file mode 100644 index 00000000..6dfafa30 --- /dev/null +++ b/opsqueue/migrations/20260709170000_add_paused_to_submissions.down.sql @@ -0,0 +1 @@ +ALTER TABLE submissions DROP COLUMN paused; diff --git a/opsqueue/migrations/20260709170000_add_paused_to_submissions.up.sql b/opsqueue/migrations/20260709170000_add_paused_to_submissions.up.sql new file mode 100644 index 00000000..6c9af7be --- /dev/null +++ b/opsqueue/migrations/20260709170000_add_paused_to_submissions.up.sql @@ -0,0 +1 @@ +ALTER TABLE submissions ADD COLUMN paused INTEGER NOT NULL DEFAULT 0; diff --git a/opsqueue/opsqueue_example_database_schema.db b/opsqueue/opsqueue_example_database_schema.db index bd172f8694fbde943eeb4b0fb6c90e6cf6529dd8..193c27f0590e5f1b06c0782f7e811b606f863d93 100644 GIT binary patch delta 101 zcmZozz}~QceS);090LP`ED*zh?nE79M!Cj>tqF{4`59Mi=M!LLXJcX)oh~lM_>YCR yiDB}@iTzwnVa)8}%F2u_#oOh@8MPTXbQB5_ON&!erZdVgN;7J1XOv|8zYPE|(ibZL delta 90 zcmZozz}~QceS);0ECT}r2*ZHRL>*&B*~Wye35;v`8AZ4A2{5v=F>!cL7Z+pv$JJ!a m%r361%-Ev4xrdjHarzDkMpaf#-X?}^6DMxk&M3+Fe;WWF-4w|H diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index 1106f4ed..b97a0d9c 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -779,6 +779,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 11685bdb..9d36123c 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -152,6 +152,9 @@ pub struct Submission { pub chunk_size: ChunkSize, pub metadata: Option, pub otel_trace_carrier: String, + /// When `true`, the submission's chunks are not dispatched to consumers. + /// Used by the delegation protocol to hold a submission until a JM master delegates it. + pub paused: bool, } /// A submission that has been completed successfully. @@ -225,6 +228,7 @@ impl Submission { chunk_size: ChunkSize::default(), metadata: None, otel_trace_carrier, + paused: false, } } @@ -244,6 +248,7 @@ impl Submission { chunk_size, metadata, otel_trace_carrier, + paused: false, }; let chunks = chunks .into_iter() @@ -319,8 +324,8 @@ pub mod db { ) -> Result<(), DatabaseError> { sqlx::query!( " - INSERT INTO submissions (id, prefix, chunks_total, chunks_done, metadata, otel_trace_carrier, chunk_size) - VALUES ($1, $2, $3, $4, $5, $6, $7) + INSERT INTO submissions (id, prefix, chunks_total, chunks_done, metadata, otel_trace_carrier, chunk_size, paused) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8) ", submission.id, submission.prefix, @@ -329,6 +334,7 @@ pub mod db { submission.metadata, submission.otel_trace_carrier, submission.chunk_size.0, + submission.paused, ) .execute(conn.get_inner()) .await?; @@ -409,6 +415,7 @@ pub mod db { metadata: Option, strategic_metadata: StrategicMetadataMap, chunk_size: ChunkSize, + paused: bool, mut conn: impl WriterConnection, ) -> Result { let submission_id = SubmissionId::new(); @@ -422,6 +429,7 @@ pub mod db { chunk_size, metadata, otel_trace_carrier, + paused, }; let iter = chunks_contents .into_iter() @@ -471,6 +479,7 @@ pub mod db { , chunk_size AS "chunk_size: ChunkSize" , metadata , otel_trace_carrier + , paused AS "paused: bool" FROM submissions WHERE id = $1 "#, id @@ -546,6 +555,7 @@ pub mod db { , chunk_size AS "chunk_size: ChunkSize" , metadata , otel_trace_carrier + , paused AS "paused: bool" FROM submissions WHERE id = $1 "#, id @@ -1077,6 +1087,7 @@ pub mod test { None, strategic_metadata.clone(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1169,6 +1180,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1179,6 +1191,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1189,6 +1202,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1199,6 +1213,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1231,6 +1246,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1241,6 +1257,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1251,6 +1268,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1300,6 +1318,8 @@ pub mod test { StrategicMetadataMap::default(), // chunk size ChunkSize::default(), + // paused + false, &mut conn, ) .await diff --git a/opsqueue/src/consumer/strategy.rs b/opsqueue/src/consumer/strategy.rs index 7aeffb27..e1c5f6a8 100644 --- a/opsqueue/src/consumer/strategy.rs +++ b/opsqueue/src/consumer/strategy.rs @@ -41,14 +41,15 @@ impl Strategy { ) -> &'a mut QueryBuilder<'a, Sqlite> { use Strategy::*; match self { - Oldest => qb.push("SELECT * FROM chunks"), - Newest => qb.push("SELECT * FROM chunks"), + Oldest => qb.push("SELECT * FROM chunks WHERE submission_id NOT IN (SELECT id FROM submissions WHERE paused)"), + Newest => qb.push("SELECT * FROM chunks WHERE submission_id NOT IN (SELECT id FROM submissions WHERE paused)"), Random => { let random_offset: u16 = rand::random(); qb.push("SELECT * FROM chunks WHERE random_order >= ") .push_bind(random_offset) - .push(" UNION ALL SELECT * FROM chunks WHERE random_order < ") + .push(" AND submission_id NOT IN (SELECT id FROM submissions WHERE paused) UNION ALL SELECT * FROM chunks WHERE random_order < ") .push_bind(random_offset) + .push(" AND submission_id NOT IN (SELECT id FROM submissions WHERE paused)") } PreferDistinct { @@ -161,7 +162,10 @@ pub mod test { let explained = explain(qb, &mut conn).await; assert_streaming_query(qb, &explained); - assert_eq!(explained, "3, 0, SCAN chunks"); + assert_eq!( + explained, + "3, 0, SCAN chunks\n8, 0, LIST SUBQUERY 1\n10, 8, SCAN submissions" + ); } #[sqlx::test] @@ -174,7 +178,10 @@ pub mod test { let explained = explain(qb, &mut conn).await; assert_streaming_query(qb, &explained); - assert_eq!(explained, "3, 0, SCAN chunks"); + assert_eq!( + explained, + "3, 0, SCAN chunks\n8, 0, LIST SUBQUERY 1\n10, 8, SCAN submissions" + ); } #[sqlx::test] @@ -191,8 +198,12 @@ pub mod test { 1, 0, COMPOUND QUERY 2, 1, LEFT-MOST SUBQUERY 5, 2, SEARCH chunks USING INDEX random_chunks_order (random_order>?) - 22, 1, UNION ALL - 25, 22, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 16, 4, CORRELATED SCALAR SUBQUERY 5 - 20, 16, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 26, 16, LIST SUBQUERY 3 - 28, 26, SCAN json_each VIRTUAL TABLE INDEX 0: - 63, 3, UNION ALL - 66, 63, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 137, 125, CORRELATED SCALAR SUBQUERY 7 - 141, 137, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 147, 137, LIST SUBQUERY 3 - 149, 147, SCAN json_each VIRTUAL TABLE INDEX 0: - 184, 124, UNION ALL - 187, 184, SEARCH chunks USING INDEX random_chunks_order (random_order?) + 185, 171, LIST SUBQUERY 1 + 187, 185, SCAN submissions + 206, 171, CORRELATED SCALAR SUBQUERY 9 + 210, 206, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) + 216, 206, LIST SUBQUERY 5 + 218, 216, SCAN json_each VIRTUAL TABLE INDEX 0: + 253, 170, UNION ALL + 256, 253, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 18, 6, CORRELATED SCALAR SUBQUERY 5 - 22, 18, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 28, 18, LIST SUBQUERY 3 - 30, 28, SCAN json_each VIRTUAL TABLE INDEX 0: - 57, 6, CORRELATED SCALAR SUBQUERY 11 - 61, 57, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 67, 57, LIST SUBQUERY 9 - 69, 67, SCAN json_each VIRTUAL TABLE INDEX 0: - 104, 5, UNION ALL - 107, 104, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 217, 205, CORRELATED SCALAR SUBQUERY 7 - 221, 217, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 227, 217, LIST SUBQUERY 3 - 229, 227, SCAN json_each VIRTUAL TABLE INDEX 0: - 256, 205, CORRELATED SCALAR SUBQUERY 11 - 260, 256, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 266, 256, LIST SUBQUERY 9 - 268, 266, SCAN json_each VIRTUAL TABLE INDEX 0: - 303, 204, UNION ALL - 306, 303, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 418, 406, CORRELATED SCALAR SUBQUERY 5 - 422, 418, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 428, 418, LIST SUBQUERY 3 - 430, 428, SCAN json_each VIRTUAL TABLE INDEX 0: - 457, 406, CORRELATED SCALAR SUBQUERY 13 - 461, 457, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 467, 457, LIST SUBQUERY 9 - 469, 467, SCAN json_each VIRTUAL TABLE INDEX 0: - 504, 405, UNION ALL - 507, 504, SEARCH chunks USING INDEX random_chunks_order (random_order?) - 617, 605, CORRELATED SCALAR SUBQUERY 7 - 621, 617, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 627, 617, LIST SUBQUERY 3 - 629, 627, SCAN json_each VIRTUAL TABLE INDEX 0: - 656, 605, CORRELATED SCALAR SUBQUERY 13 - 660, 656, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) - 666, 656, LIST SUBQUERY 9 - 668, 666, SCAN json_each VIRTUAL TABLE INDEX 0: - 703, 604, UNION ALL - 706, 703, SEARCH chunks USING INDEX random_chunks_order (random_order?) + 265, 251, LIST SUBQUERY 1 + 267, 265, SCAN submissions + 286, 251, CORRELATED SCALAR SUBQUERY 9 + 290, 286, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) + 296, 286, LIST SUBQUERY 5 + 298, 296, SCAN json_each VIRTUAL TABLE INDEX 0: + 325, 251, CORRELATED SCALAR SUBQUERY 13 + 329, 325, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) + 335, 325, LIST SUBQUERY 11 + 337, 335, SCAN json_each VIRTUAL TABLE INDEX 0: + 372, 250, UNION ALL + 375, 372, SEARCH chunks USING INDEX random_chunks_order (random_order?) + 512, 498, LIST SUBQUERY 1 + 514, 512, SCAN submissions + 533, 498, CORRELATED SCALAR SUBQUERY 7 + 537, 533, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) + 543, 533, LIST SUBQUERY 5 + 545, 543, SCAN json_each VIRTUAL TABLE INDEX 0: + 572, 498, CORRELATED SCALAR SUBQUERY 15 + 576, 572, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) + 582, 572, LIST SUBQUERY 11 + 584, 582, SCAN json_each VIRTUAL TABLE INDEX 0: + 619, 497, UNION ALL + 622, 619, SEARCH chunks USING INDEX random_chunks_order (random_order?) + 757, 743, LIST SUBQUERY 1 + 759, 757, SCAN submissions + 778, 743, CORRELATED SCALAR SUBQUERY 9 + 782, 778, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) + 788, 778, LIST SUBQUERY 5 + 790, 788, SCAN json_each VIRTUAL TABLE INDEX 0: + 817, 743, CORRELATED SCALAR SUBQUERY 15 + 821, 817, SEARCH submissions_metadata USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=? AND submission_id=?) + 827, 817, LIST SUBQUERY 11 + 829, 827, SCAN json_each VIRTUAL TABLE INDEX 0: + 864, 742, UNION ALL + 867, 864, SEARCH chunks USING INDEX random_chunks_order (random_order, + /// Notified when new chunks become available for dispatch (e.g. after unpausing a submission). + pub notify_on_insert: Arc, http_client: reqwest::Client, } impl ServerState { - pub fn new(pool: DBPools, config: &'static Config) -> Self { + pub fn new(pool: DBPools, config: &'static Config, notify_on_insert: Arc) -> Self { Self { pool, config, notify_on_submission_change: Arc::new(Notify::new()), + notify_on_insert, http_client: reqwest::Client::new(), } } @@ -98,8 +101,9 @@ async fn job_delegate( // Attempt to set jm_task_id on the submission, in whichever table it currently lives. // The AND jm_task_id IS NULL guard makes this idempotent: if the master retries the // same request after an already-successful response, we simply skip the update. + // For active submissions, also clear the paused flag so consumers can pick up the chunks. let n1 = sqlx::query!( - "UPDATE submissions SET jm_task_id = $1 WHERE id = $2 AND jm_task_id IS NULL", + "UPDATE submissions SET jm_task_id = $1, paused = 0 WHERE id = $2 AND jm_task_id IS NULL", task_id, submission_id, ) @@ -155,6 +159,10 @@ async fn job_delegate( if updated > 0 { // Trigger the background loop in case the submission already reached a terminal state. state.notify_on_submission_change.notify_one(); + // If n1 > 0, an active submission was unpaused; wake consumers so they can pick up its chunks. + if n1 > 0 { + state.notify_on_insert.notify_waiters(); + } return Ok(()); } @@ -252,9 +260,93 @@ async fn job_kill( async fn job_return( State(state): State, Json(request): Json>, -) -> Result<(), ()> { - let _ = (state, request); - todo!() +) -> Result<(), StatusCode> { + let Some(jm_master_url) = &state.config.jm_master_url else { + tracing::warn!("job_return called but jm_master_url is not configured"); + return Err(StatusCode::SERVICE_UNAVAILABLE); + }; + + let mut conn = state.pool.writer_conn().await.map_err(|e| { + tracing::error!("DB error acquiring writer connection: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let mut returned_task_ids: Vec = Vec::new(); + + for task_id in &request { + let task_id = &task_id.0; + + // Only return submissions that are still paused (not yet delegated / running). + // This request is non-binding: skip tasks that are already active. + let row = sqlx::query!( + r#"SELECT id AS "id: SubmissionId" FROM submissions WHERE jm_task_id = $1 AND paused"#, + task_id, + ) + .fetch_optional(conn.get_inner()) + .await + .map_err(|e| { + tracing::error!(%task_id, "DB error looking up submission for return: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + let Some(row) = row else { + tracing::debug!(%task_id, "job_return: submission not paused or not found, skipping"); + continue; + }; + + // Clear jm_task_id before cancelling so the cancelled row is not picked up by the + // background reporting loop. + sqlx::query!( + "UPDATE submissions SET jm_task_id = NULL WHERE id = $1", + row.id, + ) + .execute(conn.get_inner()) + .await + .map_err(|e| { + tracing::error!(%task_id, submission_id = %row.id, "DB error clearing jm_task_id for return: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })?; + + match crate::common::submission::db::cancel_submission(row.id, &mut conn).await { + Ok(()) => { + tracing::info!(%task_id, submission_id = %row.id, "Returned paused submission"); + returned_task_ids.push(task_id.clone()); + } + Err(E::R(_)) => { + tracing::debug!(%task_id, submission_id = %row.id, "job_return: submission already terminal"); + } + Err(E::L(e)) => { + tracing::error!(%task_id, submission_id = %row.id, "DB error cancelling submission for return: {e:?}"); + return Err(StatusCode::INTERNAL_SERVER_ERROR); + } + } + } + + if returned_task_ids.is_empty() { + return Ok(()); + } + + let url = format!( + "{}/delegation/return", + jm_master_url.as_str().trim_end_matches('/') + ); + state + .http_client + .put(&url) + .json(&returned_task_ids) + .send() + .await + .map_err(|e| { + tracing::error!("HTTP error reporting returned tasks: {e:?}"); + StatusCode::BAD_GATEWAY + })? + .error_for_status() + .map_err(|e| { + tracing::error!("JM master rejected returned tasks: {e:?}"); + StatusCode::BAD_GATEWAY + })?; + + Ok(()) } const JOBMACHINE_DELEGATION_BACKGROUND_LOOP: std::time::Duration = diff --git a/opsqueue/src/producer/client.rs b/opsqueue/src/producer/client.rs index b09e20b0..be612e1a 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -359,6 +359,7 @@ mod tests { metadata: None, strategic_metadata: Default::default(), chunk_size: None, + paused: false, }; client .insert_submission(&submission, &Default::default()) @@ -402,6 +403,7 @@ mod tests { metadata: None, strategic_metadata: Default::default(), chunk_size: None, + paused: false, }; let submission_id = client .insert_submission(&submission, &Default::default()) diff --git a/opsqueue/src/producer/common.rs b/opsqueue/src/producer/common.rs index e4f80d15..0c52a354 100644 --- a/opsqueue/src/producer/common.rs +++ b/opsqueue/src/producer/common.rs @@ -9,6 +9,10 @@ pub struct InsertSubmission { pub metadata: Option, pub strategic_metadata: StrategicMetadataMap, pub chunk_size: Option, + /// When `true`, the submission is inserted in a paused state and chunks will not be + /// dispatched to consumers until it is unpaused (e.g. via `job_delegate`). + #[serde(default)] + pub paused: bool, } /// Either embedded chunk contents or a reference to object storage. diff --git a/opsqueue/src/producer/server.rs b/opsqueue/src/producer/server.rs index 0298a2b0..881ffa4b 100644 --- a/opsqueue/src/producer/server.rs +++ b/opsqueue/src/producer/server.rs @@ -152,6 +152,7 @@ async fn insert_submission( request.metadata, request.strategic_metadata, request.chunk_size.unwrap_or_default(), + request.paused, &mut conn, ) .await?; diff --git a/opsqueue/src/server.rs b/opsqueue/src/server.rs index 37047933..4902b57d 100644 --- a/opsqueue/src/server.rs +++ b/opsqueue/src/server.rs @@ -81,8 +81,8 @@ pub fn build_router( .run_background() .build_router(); let producer_routes = - crate::producer::server::ServerState::new(pool.clone(), notify_on_insert).build_router(); - let delegation_routes = crate::delegation::server::ServerState::new(pool, config) + crate::producer::server::ServerState::new(pool.clone(), notify_on_insert.clone()).build_router(); + let delegation_routes = crate::delegation::server::ServerState::new(pool, config, notify_on_insert) .run_background(cancellation_token.clone()) .build_router(); From 9f817e8262eba10f294cf218dd9bb2807dca42ac Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 9 Jul 2026 17:21:30 +0200 Subject: [PATCH 5/7] job_return: pause submission instead of cancelling Re-pause the active submission (SET paused = 1, jm_task_id = NULL) rather than cancelling it. The submission stays in the queue so it can be re-delegated later; consumers simply stop picking up its chunks. If the submission is not active (already terminal or unknown), it is silently skipped. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- opsqueue/src/delegation/server.rs | 49 ++++++++----------------------- 1 file changed, 13 insertions(+), 36 deletions(-) diff --git a/opsqueue/src/delegation/server.rs b/opsqueue/src/delegation/server.rs index 1b6c2dd7..2733816a 100644 --- a/opsqueue/src/delegation/server.rs +++ b/opsqueue/src/delegation/server.rs @@ -276,49 +276,26 @@ async fn job_return( for task_id in &request { let task_id = &task_id.0; - // Only return submissions that are still paused (not yet delegated / running). - // This request is non-binding: skip tasks that are already active. - let row = sqlx::query!( - r#"SELECT id AS "id: SubmissionId" FROM submissions WHERE jm_task_id = $1 AND paused"#, + // Re-pause the submission and clear jm_task_id so the background reporting loop + // does not pick it up, and so consumers stop dispatching its chunks. + // If the submission is not in the active table (already terminal or unknown), skip it. + let rows_affected = sqlx::query!( + "UPDATE submissions SET paused = 1, jm_task_id = NULL WHERE jm_task_id = $1", task_id, ) - .fetch_optional(conn.get_inner()) - .await - .map_err(|e| { - tracing::error!(%task_id, "DB error looking up submission for return: {e:?}"); - StatusCode::INTERNAL_SERVER_ERROR - })?; - - let Some(row) = row else { - tracing::debug!(%task_id, "job_return: submission not paused or not found, skipping"); - continue; - }; - - // Clear jm_task_id before cancelling so the cancelled row is not picked up by the - // background reporting loop. - sqlx::query!( - "UPDATE submissions SET jm_task_id = NULL WHERE id = $1", - row.id, - ) .execute(conn.get_inner()) .await .map_err(|e| { - tracing::error!(%task_id, submission_id = %row.id, "DB error clearing jm_task_id for return: {e:?}"); + tracing::error!(%task_id, "DB error pausing submission for return: {e:?}"); StatusCode::INTERNAL_SERVER_ERROR - })?; + })? + .rows_affected(); - match crate::common::submission::db::cancel_submission(row.id, &mut conn).await { - Ok(()) => { - tracing::info!(%task_id, submission_id = %row.id, "Returned paused submission"); - returned_task_ids.push(task_id.clone()); - } - Err(E::R(_)) => { - tracing::debug!(%task_id, submission_id = %row.id, "job_return: submission already terminal"); - } - Err(E::L(e)) => { - tracing::error!(%task_id, submission_id = %row.id, "DB error cancelling submission for return: {e:?}"); - return Err(StatusCode::INTERNAL_SERVER_ERROR); - } + if rows_affected > 0 { + tracing::info!(%task_id, "Returned (re-paused) submission"); + returned_task_ids.push(task_id.clone()); + } else { + tracing::debug!(%task_id, "job_return: submission not active or not found, skipping"); } } From 112045d781f1b6c4a2ec52648b164e67cadec536 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 9 Jul 2026 17:25:56 +0200 Subject: [PATCH 6/7] job_return: delegate reporting to background loop via paused + jm_task_id Instead of calling /delegation/return directly from the handler, job_return now only re-pauses the submission (SET paused = 1) while leaving jm_task_id set. The background loop in report_delegated_submissions then picks up these paused-with-jm_task_id submissions, PUTs their task IDs to /delegation/return, and clears jm_task_id afterwards (same pattern as /delegation/complete). This keeps all outbound JM master communication in one place and makes job_return resilient to transient HTTP failures. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- opsqueue/src/delegation/server.rs | 124 ++++++++++++++++-------------- 1 file changed, 68 insertions(+), 56 deletions(-) diff --git a/opsqueue/src/delegation/server.rs b/opsqueue/src/delegation/server.rs index 2733816a..7be69c37 100644 --- a/opsqueue/src/delegation/server.rs +++ b/opsqueue/src/delegation/server.rs @@ -159,7 +159,8 @@ async fn job_delegate( if updated > 0 { // Trigger the background loop in case the submission already reached a terminal state. state.notify_on_submission_change.notify_one(); - // If n1 > 0, an active submission was unpaused; wake consumers so they can pick up its chunks. + // If n1 > 0, an active submission might have been unpaused; wake consumers so they can pick + // up its chunks. if n1 > 0 { state.notify_on_insert.notify_waiters(); } @@ -261,26 +262,21 @@ async fn job_return( State(state): State, Json(request): Json>, ) -> Result<(), StatusCode> { - let Some(jm_master_url) = &state.config.jm_master_url else { - tracing::warn!("job_return called but jm_master_url is not configured"); - return Err(StatusCode::SERVICE_UNAVAILABLE); - }; - let mut conn = state.pool.writer_conn().await.map_err(|e| { tracing::error!("DB error acquiring writer connection: {e:?}"); StatusCode::INTERNAL_SERVER_ERROR })?; - let mut returned_task_ids: Vec = Vec::new(); + let mut any_paused = false; for task_id in &request { let task_id = &task_id.0; - // Re-pause the submission and clear jm_task_id so the background reporting loop - // does not pick it up, and so consumers stop dispatching its chunks. + // Re-pause the submission and keep jm_task_id set so the background loop can + // report it to the JM master via /delegation/return. // If the submission is not in the active table (already terminal or unknown), skip it. let rows_affected = sqlx::query!( - "UPDATE submissions SET paused = 1, jm_task_id = NULL WHERE jm_task_id = $1", + "UPDATE submissions SET paused = 1 WHERE jm_task_id = $1", task_id, ) .execute(conn.get_inner()) @@ -292,37 +288,17 @@ async fn job_return( .rows_affected(); if rows_affected > 0 { - tracing::info!(%task_id, "Returned (re-paused) submission"); - returned_task_ids.push(task_id.clone()); + tracing::info!(%task_id, "Queued submission for return (re-paused)"); + any_paused = true; } else { tracing::debug!(%task_id, "job_return: submission not active or not found, skipping"); } } - if returned_task_ids.is_empty() { - return Ok(()); + if any_paused { + state.notify_on_submission_change.notify_one(); } - let url = format!( - "{}/delegation/return", - jm_master_url.as_str().trim_end_matches('/') - ); - state - .http_client - .put(&url) - .json(&returned_task_ids) - .send() - .await - .map_err(|e| { - tracing::error!("HTTP error reporting returned tasks: {e:?}"); - StatusCode::BAD_GATEWAY - })? - .error_for_status() - .map_err(|e| { - tracing::error!("JM master rejected returned tasks: {e:?}"); - StatusCode::BAD_GATEWAY - })?; - Ok(()) } @@ -360,8 +336,11 @@ async fn run_in_background( } /// Queries all terminal submissions with `jm_task_id IS NOT NULL`, reports them to the -/// JM master via `/delegation/complete`, then deletes the rows (along with their -/// associated chunks and metadata) on success. +/// JM master via `/delegation/complete`, then clears `jm_task_id` on success. +/// +/// Also queries paused active submissions with `jm_task_id IS NOT NULL` (queued for return +/// via `job_return`), reports them to the JM master via `/delegation/return`, then clears +/// `jm_task_id` on success. async fn report_delegated_submissions( state: &ServerState, jm_master_url: &url::Url, @@ -390,9 +369,17 @@ async fn report_delegated_submissions( .fetch_all(conn.get_inner()) .await?; + // Paused active submissions queued for return via job_return. + let to_return = sqlx::query!( + r#"SELECT id AS "id: SubmissionId", jm_task_id + FROM submissions WHERE paused AND jm_task_id IS NOT NULL"# + ) + .fetch_all(conn.get_inner()) + .await?; + drop(conn); - let tasks_to_report: Vec = completed + let tasks_to_complete: Vec = completed .iter() .map(|r| DelegatedJobCompletion { task_id: r.jm_task_id.clone().unwrap(), @@ -414,36 +401,61 @@ async fn report_delegated_submissions( })) .collect(); - if tasks_to_report.is_empty() { + let tasks_to_return: Vec = to_return + .iter() + .map(|r| r.jm_task_id.clone().unwrap()) + .collect(); + + if tasks_to_complete.is_empty() && tasks_to_return.is_empty() { return Ok(()); } if triggered_by_timeout { tracing::warn!( - n_tasks = tasks_to_report.len(), + n_complete = tasks_to_complete.len(), + n_return = tasks_to_return.len(), "Delegation background loop triggered by timeout with pending tasks; \ possible missing notify_on_submission_change call" ); } - let url = format!( - "{}/delegation/complete", - jm_master_url.as_str().trim_end_matches('/') - ); - state - .http_client - .put(&url) - .json(&tasks_to_report) - .send() - .await? - .error_for_status()?; - - // Report succeeded; clear jm_task_id on all reported rows so the periodic cleanup - // worker can delete them in due course. let mut conn = state.pool.writer_conn().await?; - clear_jm_task_id("submissions_completed", &completed.iter().map(|r| r.id).collect::>(), &mut conn).await?; - clear_jm_task_id("submissions_failed", &failed.iter().map(|r| r.id).collect::>(), &mut conn).await?; - clear_jm_task_id("submissions_cancelled", &cancelled.iter().map(|r| r.id).collect::>(), &mut conn).await?; + + if !tasks_to_complete.is_empty() { + let url = format!( + "{}/delegation/complete", + jm_master_url.as_str().trim_end_matches('/') + ); + state + .http_client + .put(&url) + .json(&tasks_to_complete) + .send() + .await? + .error_for_status()?; + + // Report succeeded; clear jm_task_id so the periodic cleanup worker can delete them. + clear_jm_task_id("submissions_completed", &completed.iter().map(|r| r.id).collect::>(), &mut conn).await?; + clear_jm_task_id("submissions_failed", &failed.iter().map(|r| r.id).collect::>(), &mut conn).await?; + clear_jm_task_id("submissions_cancelled", &cancelled.iter().map(|r| r.id).collect::>(), &mut conn).await?; + } + + if !tasks_to_return.is_empty() { + let url = format!( + "{}/delegation/return", + jm_master_url.as_str().trim_end_matches('/') + ); + state + .http_client + .put(&url) + .json(&tasks_to_return) + .send() + .await? + .error_for_status()?; + + // Return reported; clear jm_task_id so these submissions are no longer tracked. + clear_jm_task_id("submissions", &to_return.iter().map(|r| r.id).collect::>(), &mut conn).await?; + } Ok(()) } From 0a003837732ee8e2cfd2d58511a5b129bea18eaf Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 9 Jul 2026 17:41:02 +0200 Subject: [PATCH 7/7] opsqueue_python: expose paused parameter for submission insertion Add a keyword-only `paused: bool = False` parameter to all submission insertion methods in both the Rust bindings and the Python wrappers: - insert_submission_direct (Rust) - insert_submission_chunks (Rust + Python) - run_submission_chunks (Rust + Python) - insert_submission (Python) - run_submission / async_run_submission (Python) - async_run_submission_chunks (Python) When paused=True, the submission is held in the queue but chunks are not dispatched to consumers until the submission is explicitly unpaused (e.g. via the delegation job_delegate endpoint). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../python/opsqueue/producer.py | 24 +++++++++++++++++++ libs/opsqueue_python/src/producer.rs | 12 +++++++--- 2 files changed, 33 insertions(+), 3 deletions(-) diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index 5c6a7f06..56c86598 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/libs/opsqueue_python/python/opsqueue/producer.py @@ -94,12 +94,16 @@ def run_submission( serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT, metadata: None | bytes = None, strategic_metadata: None | dict[str, int] = None, + paused: bool = False, ) -> Iterator[Any]: """ Inserts a submission into the queue, and blocks until it is completed. Chunking is done automatically, based on the provided chunk size. + If `paused` is True, the submission will not be dispatched to consumers until + it is explicitly unpaused (e.g. via the delegation `job_delegate` endpoint). + If the submission fails, an exception will be raised. (If opsqueue or the object storage cannot be reached, exceptions will also be raised). @@ -114,6 +118,7 @@ def run_submission( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + paused=paused, ) return _unchunk_iterator(results_iter, serialization_format) @@ -125,6 +130,7 @@ async def async_run_submission( serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT, metadata: None | bytes = None, strategic_metadata: None | dict[str, int] = None, + paused: bool = False, ) -> AsyncIterator[Any]: tracer = trace.get_tracer("opsqueue.producer") with tracer.start_as_current_span("run_submission"): @@ -133,6 +139,7 @@ async def async_run_submission( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + paused=paused, ) return _async_unchunk_iterator(results_iter, serialization_format) @@ -144,6 +151,7 @@ def insert_submission( serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT, metadata: None | bytes = None, strategic_metadata: None | dict[str, int] = None, + paused: bool = False, ) -> SubmissionId: """ Inserts a submission into the queue, @@ -151,6 +159,9 @@ def insert_submission( Chunking is done automatically, based on the provided chunk size. + If `paused` is True, the submission will not be dispatched to consumers until + it is explicitly unpaused (e.g. via the delegation `job_delegate` endpoint). + Raises: - `ChunkSizeIsZeroError` if passing an incorrect chunk size of zero; - `InternalProducerClientError` if there is a low-level internal error. @@ -160,6 +171,7 @@ def insert_submission( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + paused=paused, ) def blocking_stream_completed_submission( @@ -209,10 +221,14 @@ def run_submission_chunks( metadata: None | bytes = None, strategic_metadata: None | dict[str, int] = None, chunk_size: None | int = None, + paused: bool = False, ) -> Iterator[bytes]: """ Inserts an already-chunked submission into the queue, and blocks until it is completed. + If `paused` is True, the submission will not be dispatched to consumers until + it is explicitly unpaused (e.g. via the delegation `job_delegate` endpoint). + If the submission fails, an exception will be raised. (If opsqueue or the object storage cannot be reached, exceptions will also be raised). @@ -226,6 +242,7 @@ def run_submission_chunks( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + paused=paused, ) return self.blocking_stream_completed_submission_chunks(submission_id) @@ -236,6 +253,7 @@ async def async_run_submission_chunks( metadata: None | bytes = None, strategic_metadata: None | dict[str, int] = None, chunk_size: None | int = None, + paused: bool = False, ) -> AsyncIterator[bytes]: # NOTE: the insertion is not currently async. # Why? Simplicity. This is unlikely to be the bottleneck @@ -246,6 +264,7 @@ async def async_run_submission_chunks( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + paused=paused, ) return await self.async_stream_completed_submission_chunks(submission_id) @@ -257,11 +276,15 @@ def insert_submission_chunks( metadata: None | bytes = None, strategic_metadata: None | dict[str, int] = None, chunk_size: None | int = None, + paused: bool = False, ) -> SubmissionId: """ Inserts an already-chunked submission into the queue, returning an ID you can use to track the submission's progress afterwards. + If `paused` is True, the submission will not be dispatched to consumers until + it is explicitly unpaused (e.g. via the delegation `job_delegate` endpoint). + Raises: - `InternalProducerClientError` if there is a low-level internal error. """ @@ -272,6 +295,7 @@ def insert_submission_chunks( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + paused=paused, otel_trace_carrier=otel_trace_carrier, ) diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index 7d4266c6..6914df81 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -192,13 +192,14 @@ impl ProducerClient { /// Directly inserts a submission without sending the chunks to GCS /// (but immediately embedding them in the DB). /// NOTE: This does not support StrategicMetadata currently - #[pyo3(signature = (chunk_contents, metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default()))] + #[pyo3(signature = (chunk_contents, metadata=None, chunk_size=None, paused=false, otel_trace_carrier=CarrierMap::default()))] pub fn insert_submission_direct( &self, py: Python<'_>, chunk_contents: Vec, metadata: Option, chunk_size: Option, + paused: bool, otel_trace_carrier: CarrierMap, ) -> CPyResult> { let strategic_metadata = Default::default(); @@ -211,6 +212,7 @@ impl ProducerClient { }, metadata, strategic_metadata, + paused, }; self.block_unless_interrupted(async move { self.producer_client @@ -222,7 +224,7 @@ impl ProducerClient { }) } - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default()))] + #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, paused=false, otel_trace_carrier=CarrierMap::default()))] #[allow(clippy::type_complexity)] pub fn insert_submission_chunks( &self, @@ -231,6 +233,7 @@ impl ProducerClient { metadata: Option, strategic_metadata: Option, chunk_size: Option, + paused: bool, otel_trace_carrier: CarrierMap, ) -> CPyResult< SubmissionId, @@ -271,6 +274,7 @@ impl ProducerClient { }, metadata, strategic_metadata: strategic_metadata.unwrap_or_default(), + paused, }; self.producer_client .insert_submission(&submission, &otel_trace_carrier) @@ -312,7 +316,7 @@ impl ProducerClient { }) } - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default()))] + #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, paused=false, otel_trace_carrier=CarrierMap::default()))] #[allow(clippy::type_complexity)] pub fn run_submission_chunks( &self, @@ -321,6 +325,7 @@ impl ProducerClient { metadata: Option, strategic_metadata: Option, chunk_size: Option, + paused: bool, otel_trace_carrier: CarrierMap, ) -> CPyResult< PyChunksIter, @@ -338,6 +343,7 @@ impl ProducerClient { metadata, strategic_metadata, chunk_size, + paused, otel_trace_carrier, ) .map_err(|CError(e)| {