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)| { 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/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/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 ae49303b..193c27f0 100644 Binary files a/opsqueue/opsqueue_example_database_schema.db and b/opsqueue/opsqueue_example_database_schema.db differ 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 69040f1a..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 @@ -748,8 +758,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 +793,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 +829,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 *; ", @@ -846,6 +856,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 }, @@ -908,11 +919,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 ) @@ -921,7 +935,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 ) @@ -930,33 +945,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 @@ -1065,6 +1087,7 @@ pub mod test { None, strategic_metadata.clone(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1157,6 +1180,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1167,6 +1191,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1177,6 +1202,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1187,6 +1213,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1219,6 +1246,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1229,6 +1257,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1239,6 +1268,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1288,6 +1318,8 @@ pub mod test { StrategicMetadataMap::default(), // chunk size ChunkSize::default(), + // paused + false, &mut conn, ) .await 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/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, notify_on_insert: Arc) -> Self { + Self { + pool, + config, + notify_on_submission_change: Arc::new(Notify::new()), + notify_on_insert, + 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<()> { + 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, +} + +#[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<(), 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. + // 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, paused = 0 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(); + // 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(); + } + 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))] +async fn job_kill( + State(state): State, + Json(request): Json>, +) -> 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))] +async fn job_return( + State(state): State, + Json(request): Json>, +) -> 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_paused = false; + + for task_id in &request { + let task_id = &task_id.0; + + // 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 WHERE jm_task_id = $1", + task_id, + ) + .execute(conn.get_inner()) + .await + .map_err(|e| { + tracing::error!(%task_id, "DB error pausing submission for return: {e:?}"); + StatusCode::INTERNAL_SERVER_ERROR + })? + .rows_affected(); + + if rows_affected > 0 { + 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 any_paused { + state.notify_on_submission_change.notify_one(); + } + + Ok(()) +} + +const JOBMACHINE_DELEGATION_BACKGROUND_LOOP: std::time::Duration = + std::time::Duration::from_secs(5); + +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, + }; + + match report_delegated_submissions(&state, jm_master_url, triggered_by_timeout).await { + Ok(()) => {} + Err(e) => tracing::error!("Error in delegation background loop: {e:?}"), + } + } + + Ok(()) +} + +/// Queries all terminal submissions with `jm_task_id IS NOT NULL`, reports them to the +/// 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, + 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?; + + // 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_complete: 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(); + + 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_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 mut conn = state.pool.writer_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(()) +} + +/// 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/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"); 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 cfc9ec4b..4902b57d 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.clone()).build_router(); + let delegation_routes = crate::delegation::server::ServerState::new(pool, config, notify_on_insert) + .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<_>| {