From adc99b86900d95247e47c77f91d9a37d41e94813 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 22 Jul 2026 15:02:33 +0200 Subject: [PATCH 01/11] Use u64 in count_* functions u63 forces us to wrap literals in `u63::new`, and we need to convert to u64 at actual usage sites anyway. --- opsqueue/src/common/chunk.rs | 46 +++++++++++-------- opsqueue/src/common/submission.rs | 75 +++++++++++++------------------ opsqueue/src/producer/client.rs | 7 ++- opsqueue/src/producer/server.rs | 4 +- opsqueue/src/prometheus.rs | 4 +- 5 files changed, 64 insertions(+), 72 deletions(-) diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index 6a256192..403207c6 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -238,7 +238,6 @@ impl Chunk { pub mod db { use super::{ Chunk, ChunkCompleted, ChunkFailed, ChunkId, ChunkIndex, ChunkSize, DateTime, SubmissionId, - Utc, u63, }; use crate::common::errors::{ChunkNotFound, DatabaseError, E, SubmissionNotFound}; use crate::db::{Connection, True, WriterConnection}; @@ -632,13 +631,16 @@ pub mod db { /// # Errors /// /// Returns an error if the count query fails. + /// + /// # Panics + /// + /// Panics if `COUNT(*)` returns a negative value, which `SQLite` never does. #[tracing::instrument(skip(db))] - pub async fn count_chunks(mut db: impl Connection) -> sqlx::Result { + pub async fn count_chunks(mut db: impl Connection) -> sqlx::Result { let count = sqlx::query_scalar!("SELECT COUNT(1) as count FROM chunks;") .fetch_one(db.get_inner()) .await?; - let count = u63::new(count.cast_unsigned()); - Ok(count) + Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) } /// Count completed chunks. @@ -646,13 +648,16 @@ pub mod db { /// # Errors /// /// Returns an error if the count query fails. + /// + /// # Panics + /// + /// Panics if `COUNT(*)` returns a negative value, which `SQLite` never does. #[tracing::instrument(skip(db))] - pub async fn count_chunks_completed(mut db: impl Connection) -> sqlx::Result { + pub async fn count_chunks_completed(mut db: impl Connection) -> sqlx::Result { let count = sqlx::query_scalar!("SELECT COUNT(1) as count FROM chunks_completed;") .fetch_one(db.get_inner()) .await?; - let count = u63::new(count.cast_unsigned()); - Ok(count) + Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) } /// Count failed chunks. @@ -660,13 +665,16 @@ pub mod db { /// # Errors /// /// Returns an error if the count query fails. + /// + /// # Panics + /// + /// Panics if `COUNT(*)` returns a negative value, which `SQLite` never does. #[tracing::instrument(skip(db))] - pub async fn count_chunks_failed(mut db: impl Connection) -> sqlx::Result { + pub async fn count_chunks_failed(mut db: impl Connection) -> sqlx::Result { let count = sqlx::query_scalar!("SELECT COUNT(1) as count FROM chunks_failed;") .fetch_one(db.get_inner()) .await?; - let count = u63::new(count.cast_unsigned()); - Ok(count) + Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) } /// Looks up the number of operations in the backlog. @@ -704,11 +712,11 @@ pub mod test { vec![1, 2, 3, 4, 5].into(), ); - assert_eq!(count_chunks(&mut conn).await.unwrap(), u63::new(0)); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 0); insert_chunk(chunk.clone(), &mut conn) .await .expect("Insert chunk failed"); - assert_eq!(count_chunks(&mut conn).await.unwrap(), u63::new(1)); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 1); } #[sqlx::test(migrator = "crate::MIGRATOR")] @@ -764,12 +772,12 @@ pub mod test { .await .expect("complete chunk failed"); - assert_eq!(count_chunks(&mut conn).await.unwrap(), u63::new(0)); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 0); assert_eq!( count_chunks_completed(&mut conn).await.unwrap(), - u63::new(1) + 1 ); - assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), u63::new(0)); + assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), 0); } #[sqlx::test(migrator = "crate::MIGRATOR")] @@ -787,7 +795,7 @@ pub mod test { .await .unwrap(); - assert_eq!(count_chunks(&mut conn).await.unwrap(), u63::new(1)); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 1); conn.transaction(move |mut tx| { Box::pin(async move { @@ -836,11 +844,11 @@ pub mod test { .await .expect("Succeed chunk failed"); - assert_eq!(count_chunks(&mut conn).await.unwrap(), u63::new(0)); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 0); assert_eq!( count_chunks_completed(&mut conn).await.unwrap(), - u63::new(0) + 0 ); - assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), u63::new(1)); + assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), 1); } } diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index c36efe0a..9f586075 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -290,7 +290,6 @@ pub mod db { use axum_prometheus::metrics::{counter, histogram}; use chunk::ChunkSize; use sqlx::{QueryBuilder, Sqlite, query, query_scalar}; - use ux::u63; use super::{ Chunk, ChunkCount, ChunkIndex, DateTime, Duration, E, Metadata, Submission, @@ -1046,12 +1045,16 @@ pub mod db { /// # Errors /// /// Returns an error if the count query fails. + /// + /// # Panics + /// + /// Panics if `COUNT(*)` returns a negative value, which `SQLite` never does. #[tracing::instrument(skip(db))] - pub async fn count_submissions(mut db: impl Connection) -> sqlx::Result { + pub async fn count_submissions(mut db: impl Connection) -> sqlx::Result { let count = sqlx::query_scalar!("SELECT COUNT(1) as count FROM submissions;") .fetch_one(db.get_inner()) .await?; - Ok(u63::new(count.cast_unsigned())) + Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) } /// Count completed submissions. @@ -1059,12 +1062,16 @@ pub mod db { /// # Errors /// /// Returns an error if the count query fails. + /// + /// # Panics + /// + /// Panics if `COUNT(*)` returns a negative value, which `SQLite` never does. #[tracing::instrument(skip(db))] - pub async fn count_submissions_completed(mut db: impl Connection) -> sqlx::Result { + pub async fn count_submissions_completed(mut db: impl Connection) -> sqlx::Result { let count = sqlx::query_scalar!("SELECT COUNT(1) as count FROM submissions_completed;") .fetch_one(db.get_inner()) .await?; - Ok(u63::new(count.cast_unsigned())) + Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) } /// Count failed submissions. @@ -1072,12 +1079,16 @@ pub mod db { /// # Errors /// /// Returns an error if the count query fails. + /// + /// # Panics + /// + /// Panics if `COUNT(*)` returns a negative value, which `SQLite` never does. #[tracing::instrument(skip(db))] - pub async fn count_submissions_failed(mut db: impl Connection) -> sqlx::Result { + pub async fn count_submissions_failed(mut db: impl Connection) -> sqlx::Result { let count = sqlx::query_scalar!("SELECT COUNT(1) as count FROM submissions_failed;") .fetch_one(db.get_inner()) .await?; - Ok(u63::new(count.cast_unsigned())) + Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) } /// Transactionally removes all completed/failed submissions, @@ -1379,7 +1390,7 @@ pub mod test { let db = WriterPool::new(db); let mut conn = db.writer_conn().await.unwrap(); - assert_eq!(count_submissions(&mut conn).await.unwrap(), u63::new(0)); + assert_eq!(count_submissions(&mut conn).await.unwrap(), 0); let (submission, chunks) = Submission::from_vec( vec![Some("foo".into()), Some("bar".into()), Some("baz".into())], @@ -1391,7 +1402,7 @@ pub mod test { .await .expect("insertion failed"); - assert_eq!(count_submissions(&mut conn).await.unwrap(), u63::new(1)); + assert_eq!(count_submissions(&mut conn).await.unwrap(), 1); } #[sqlx::test(migrator = "crate::MIGRATOR")] @@ -1464,15 +1475,9 @@ pub mod test { .await .unwrap(); - assert_eq!(count_submissions(&mut conn).await.unwrap(), u63::new(0)); - assert_eq!( - count_submissions_completed(&mut conn).await.unwrap(), - u63::new(1) - ); - assert_eq!( - count_submissions_failed(&mut conn).await.unwrap(), - u63::new(0) - ); + assert_eq!(count_submissions(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_completed(&mut conn).await.unwrap(), 1); + assert_eq!(count_submissions_failed(&mut conn).await.unwrap(), 0); } #[sqlx::test(migrator = "crate::MIGRATOR")] @@ -1497,15 +1502,9 @@ pub mod test { ) .await .unwrap(); - assert_eq!(count_submissions(&mut conn).await.unwrap(), u63::new(0)); - assert_eq!( - count_submissions_completed(&mut conn).await.unwrap(), - u63::new(0) - ); - assert_eq!( - count_submissions_failed(&mut conn).await.unwrap(), - u63::new(1) - ); + assert_eq!(count_submissions(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_completed(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_failed(&mut conn).await.unwrap(), 1); } #[sqlx::test(migrator = "crate::MIGRATOR")] @@ -1624,18 +1623,12 @@ pub mod test { .await .unwrap(); - assert_eq!( - count_submissions_failed(&mut conn).await.unwrap(), - u63::new(5) - ); + assert_eq!(count_submissions_failed(&mut conn).await.unwrap(), 5); let mut conn2 = db.writer_conn().await.unwrap(); cleanup_old(&mut conn2, cutoff_timestamp).await.unwrap(); - assert_eq!( - count_submissions_failed(&mut conn).await.unwrap(), - u63::new(2) - ); + assert_eq!(count_submissions_failed(&mut conn).await.unwrap(), 2); let _sub1 = submission_status(old_four_unfailed, &mut conn) .await @@ -1666,15 +1659,9 @@ pub mod test { .await .expect("insertion failed"); - assert_eq!(count_submissions(&mut conn).await.unwrap(), u63::new(0)); - assert_eq!( - count_submissions_completed(&mut conn).await.unwrap(), - u63::new(1) - ); - assert_eq!( - count_submissions_failed(&mut conn).await.unwrap(), - u63::new(0) - ); + assert_eq!(count_submissions(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_completed(&mut conn).await.unwrap(), 1); + assert_eq!(count_submissions_failed(&mut conn).await.unwrap(), 0); } /// Removes the given top-level key from a JSON object, panicking if it was not present. diff --git a/opsqueue/src/producer/client.rs b/opsqueue/src/producer/client.rs index 4f1dc92d..1632050b 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -398,7 +398,6 @@ impl InternalProducerClientError { #[cfg(test)] #[cfg(feature = "server-logic")] mod tests { - use ux::u63; use crate::{ common::{ @@ -459,7 +458,7 @@ mod tests { let count = submission::db::count_submissions(&mut conn) .await .expect("Should be OK"); - assert_eq!(count, u63::new(0)); + assert_eq!(count, 0); let submission = InsertSubmission { chunk_contents: ChunkContents::Direct { @@ -477,7 +476,7 @@ mod tests { let count = submission::db::count_submissions(&mut conn) .await .expect("Should be OK"); - assert_eq!(count, u63::new(1)); + assert_eq!(count, 1); client .insert_submission(&submission, &std::collections::HashMap::default()) @@ -495,7 +494,7 @@ mod tests { let count = submission::db::count_submissions(&mut conn) .await .expect("Should be OK"); - assert_eq!(count, u63::new(4)); + assert_eq!(count, 4); } #[sqlx::test(migrator = "crate::MIGRATOR")] diff --git a/opsqueue/src/producer/server.rs b/opsqueue/src/producer/server.rs index 74b87075..1fa4d678 100644 --- a/opsqueue/src/producer/server.rs +++ b/opsqueue/src/producer/server.rs @@ -222,7 +222,7 @@ pub struct InsertSubmissionResponse { async fn submissions_count(State(state): State) -> Result, ServerError> { let mut conn = state.pool.reader_conn().await?; let count = submission::db::count_submissions(&mut conn).await?; - Ok(Json(u64::from(count))) + Ok(Json(count)) } async fn submissions_count_completed( @@ -230,5 +230,5 @@ async fn submissions_count_completed( ) -> Result, ServerError> { let mut conn = state.pool.reader_conn().await?; let count = submission::db::count_submissions_completed(&mut conn).await?; - Ok(Json(u64::from(count))) + Ok(Json(count)) } diff --git a/opsqueue/src/prometheus.rs b/opsqueue/src/prometheus.rs index 8bc7f023..d4b308a1 100644 --- a/opsqueue/src/prometheus.rs +++ b/opsqueue/src/prometheus.rs @@ -211,9 +211,7 @@ pub fn time_delta_as_f64(td: chrono::TimeDelta) -> f64 { #[allow(clippy::cast_precision_loss)] pub async fn calculate_scaling_metrics(db_pool: &DBPools) -> anyhow::Result<()> { let mut conn = db_pool.reader_conn().await?; - let chunks_backlog_count: u64 = crate::common::chunk::db::count_chunks(&mut conn) - .await? - .into(); + let chunks_backlog_count: u64 = crate::common::chunk::db::count_chunks(&mut conn).await?; gauge!(CHUNKS_BACKLOG_GAUGE).set(chunks_backlog_count as f64); let ops_backlog_count: f64 = crate::common::chunk::db::count_ops_in_backlog_estimate(&mut conn).await?; From 9c33f641c59e70bc7ec87d4cc3cd6d61c094fa13 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 22 Jul 2026 12:55:48 +0200 Subject: [PATCH 02/11] Make tests timeout properly, in preparation for showing failure logs properly --- .../python/opsqueue/producer.py | 16 +++- libs/opsqueue_python/src/errors.rs | 8 +- libs/opsqueue_python/src/producer.rs | 71 +++++------------ libs/opsqueue_python/tests/test_roundtrip.py | 79 ++++++++++++++++--- 4 files changed, 106 insertions(+), 68 deletions(-) diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index 82a877e3..dc807cc6 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/libs/opsqueue_python/python/opsqueue/producer.py @@ -96,6 +96,7 @@ def run_submission( serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT, metadata: None | bytes = None, strategic_metadata: None | dict[str, int] = None, + timeout: float | None = None, ) -> Iterator[Any]: """ Inserts a submission into the queue, and blocks until it is completed. @@ -116,6 +117,7 @@ def run_submission( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + timeout=timeout, ) return _unchunk_iterator(results_iter, serialization_format) @@ -169,6 +171,7 @@ def blocking_stream_completed_submission( submission_id: SubmissionId, *, serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT, + timeout: float | None = None, ) -> Iterator[Any]: """ Blocks until the submission is completed. @@ -181,7 +184,7 @@ def blocking_stream_completed_submission( (after retrying a consumer kept failing on one of the chunks) """ return _unchunk_iterator( - self.blocking_stream_completed_submission_chunks(submission_id), + self.blocking_stream_completed_submission_chunks(submission_id, timeout), serialization_format, ) @@ -211,6 +214,7 @@ def run_submission_chunks( metadata: None | bytes = None, strategic_metadata: None | dict[str, int] = None, chunk_size: None | int = None, + timeout: float | None = None, ) -> Iterator[bytes]: """ Inserts an already-chunked submission into the queue, and blocks until it is completed. @@ -229,7 +233,7 @@ def run_submission_chunks( strategic_metadata=strategic_metadata, chunk_size=chunk_size, ) - return self.blocking_stream_completed_submission_chunks(submission_id) + return self.blocking_stream_completed_submission_chunks(submission_id, timeout) async def async_run_submission_chunks( self, @@ -278,7 +282,9 @@ def insert_submission_chunks( ) def blocking_stream_completed_submission_chunks( - self, submission_id: SubmissionId + self, + submission_id: SubmissionId, + timeout: float | None = None, ) -> Iterator[bytes]: """ Blocks until the submission is completed, and returns an iterator that lazily @@ -289,7 +295,9 @@ def blocking_stream_completed_submission_chunks( - `SubmissionFailedError` if the submission failed permanently (after retrying a consumer kept failing on one of the chunks) """ - return self.inner.blocking_stream_completed_submission_chunks(submission_id) # type: ignore[no-any-return] + return self.inner.blocking_stream_completed_submission_chunks( # type: ignore[no-any-return] + submission_id, timeout + ) async def async_stream_completed_submission_chunks( self, submission_id: SubmissionId diff --git a/libs/opsqueue_python/src/errors.rs b/libs/opsqueue_python/src/errors.rs index 45f0f7da..110dded0 100644 --- a/libs/opsqueue_python/src/errors.rs +++ b/libs/opsqueue_python/src/errors.rs @@ -7,7 +7,7 @@ use opsqueue::common::errors::{ ChunkNotFound, E, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions, UnexpectedOpsqueueConsumerServerResponse, }; -use pyo3::exceptions::PyBaseException; +use pyo3::exceptions::{PyBaseException, PyTimeoutError}; use pyo3::{Bound, PyErr, Python, import_exception}; use crate::common; @@ -201,6 +201,12 @@ impl From> for PyErr { } } +impl From> for PyErr { + fn from(_value: CError) -> Self { + PyTimeoutError::new_err("timeout was reached") + } +} + impl From for CError> { fn from(value: PyErr) -> Self { CError(E::L(FatalPythonException(value))) diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index 3b917ada..45eb5f88 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -18,6 +18,7 @@ use opsqueue::{ producer::client::{Client as ActualClient, InternalProducerClientError}, tracing::CarrierMap, }; +use tokio::time::error::Elapsed; use ux::u63; use crate::{ @@ -380,57 +381,6 @@ impl ProducerClient { }) } - /// Submit chunks and then stream the completed output chunks. - /// - /// # Errors - /// - /// Returns an error if upload, submission creation, or streaming fails. - #[allow(clippy::result_large_err, clippy::type_complexity)] - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default()))] - pub fn run_submission_chunks( - &self, - py: Python<'_>, - chunk_contents: Py, - metadata: Option, - strategic_metadata: Option, - chunk_size: Option, - otel_trace_carrier: CarrierMap, - ) -> CPyResult< - PyChunksIter, - E![ - FatalPythonException, - errors::SubmissionFailed, - ChunksStorageError, - InternalProducerClientError, - ], - > { - let submission_id = self - .insert_submission_chunks( - py, - chunk_contents, - metadata, - strategic_metadata, - chunk_size, - otel_trace_carrier, - ) - .map_err(|CError(e)| { - CError(match e { - L(e) => L(e), - R(e) => R(R(e)), - }) - })?; - let res = self - .blocking_stream_completed_submission_chunks(py, submission_id) - .map_err(|CError(e)| { - CError(match e { - L(e) => L(e), - R(L(e)) => R(L(e)), - R(R(e)) => R(R(R(e))), - }) - })?; - Ok(res) - } - /// Blocks (and short-polls) until the submission is completed. /// /// We start with a small short-polling interval @@ -447,17 +397,34 @@ impl ProducerClient { &self, py: Python<'_>, submission_id: SubmissionId, + timeout: Option, ) -> CPyResult< PyChunksIter, E![ FatalPythonException, + Elapsed, errors::SubmissionFailed, InternalProducerClientError ], > { py.detach(|| { self.block_unless_interrupted(async move { - self.stream_completed_submission_chunks(submission_id).await + let fut = self.stream_completed_submission_chunks(submission_id); + match timeout { + Some(duration) => tokio::time::timeout(Duration::from_secs_f64(duration), fut) + .await + .map_err(|err| CError(R(L(err)))) + .and_then(|err| { + err.map_err(|err| match err.0 { + L(err) => CError(L(err)), + R(err) => CError(R(R(err))), + }) + }), + None => fut.await.map_err(|err| match err.0 { + L(err) => CError(L(err)), + R(err) => CError(R(R(err))), + }), + } }) }) } diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index 2e887ac4..9052823d 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -29,6 +29,8 @@ import logging import pytest +SUBMISSION_COMPLETED_TIMEOUT = 10.0 + def increment(data: int) -> int: return data + 1 @@ -56,7 +58,10 @@ def run_consumer() -> None: input_iter = range(0, 100) output_iter: Iterator[int] = producer_client.run_submission( - input_iter, chunk_size=20, strategic_metadata={"id": 42} + input_iter, + chunk_size=20, + strategic_metadata={"id": 42}, + timeout=SUBMISSION_COMPLETED_TIMEOUT, ) res = sum(output_iter) @@ -128,6 +133,7 @@ def run_consumer(_consumer_id: int) -> None: input_iter, chunk_size=chunk_size, strategic_metadata={"id": 42, "second_id": 69}, + timeout=SUBMISSION_COMPLETED_TIMEOUT, ) res = sum(output_iter) @@ -146,7 +152,9 @@ def test_empty_submission(opsqueue: OpsqueueProcess) -> None: input_iter: list[int] = [] output_iter: Iterator[int] = producer_client.run_submission( - input_iter, chunk_size=20 + input_iter, + chunk_size=20, + timeout=SUBMISSION_COMPLETED_TIMEOUT, ) res = sum(output_iter) assert res == 0 @@ -182,7 +190,10 @@ def run_consumer() -> None: input_iter = range(0, 100) output_iter: Iterator[int] = producer_client.run_submission( - input_iter, chunk_size=20, serialization_format=serialization_format + input_iter, + chunk_size=20, + serialization_format=serialization_format, + timeout=SUBMISSION_COMPLETED_TIMEOUT, ) res = sum(output_iter) @@ -225,7 +236,11 @@ def broken_increment(input: int) -> float: input_iter = range(0, 100) with pytest.raises(SubmissionFailedError) as exc_info: - producer_client.run_submission(input_iter, chunk_size=20) + producer_client.run_submission( + input_iter, + chunk_size=20, + timeout=SUBMISSION_COMPLETED_TIMEOUT, + ) # We expect the intended attributes to be there: assert isinstance(exc_info.value.failure, str) @@ -265,7 +280,10 @@ def increment_list(ints: Sequence[int], _chunk: Chunk) -> Sequence[int]: input_iter = map(lambda i: cbor2.dumps([i, i, i]), range(0, 10)) output_iter: Iterator[list[int]] = map( lambda c: cbor2.loads(c), - producer_client.run_submission_chunks(input_iter), + producer_client.run_submission_chunks( + input_iter, + timeout=SUBMISSION_COMPLETED_TIMEOUT, + ), ) import itertools @@ -304,7 +322,9 @@ def run_consumer(consumer_id: int) -> None: with multiple_background_processes(run_consumer, n_consumers) as _consumers: input_iter = range(0, 1000) output_iter: Iterator[int] = producer_client.run_submission( - input_iter, chunk_size=100 + input_iter, + chunk_size=100, + timeout=SUBMISSION_COMPLETED_TIMEOUT, ) res = sum(output_iter) @@ -379,7 +399,10 @@ def run_consumer() -> None: with background_process(run_consumer): # Wait for the submission to complete. - producer_client.blocking_stream_completed_submission(submission_id) + producer_client.blocking_stream_completed_submission( + submission_id, + timeout=SUBMISSION_COMPLETED_TIMEOUT, + ) submission = producer_client.get_submission_status(submission_id) assert submission is not None assert isinstance(submission.submission, SubmissionCompleted) @@ -423,7 +446,10 @@ def assert_submission_failed_has_metadata(x: SubmissionFailed) -> None: with pytest.raises(SubmissionFailedError) as exc_info: # Wait for the submission to fail. - producer_client.blocking_stream_completed_submission(submission_id) + producer_client.blocking_stream_completed_submission( + submission_id, + timeout=SUBMISSION_COMPLETED_TIMEOUT, + ) assert_submission_failed_has_metadata(exc_info.value.submission) submission = producer_client.get_submission_status(submission_id) @@ -511,7 +537,10 @@ def run_consumer() -> None: with background_process(run_consumer): # Wait for the submission to complete. - producer_client.blocking_stream_completed_submission(submission_id) + producer_client.blocking_stream_completed_submission( + submission_id, + timeout=SUBMISSION_COMPLETED_TIMEOUT, + ) submission = producer_client.get_submission_status(submission_id) assert submission is not None assert isinstance(submission.submission, SubmissionCompleted) @@ -544,7 +573,10 @@ def consume(x: int) -> None: with background_process(run_consumer): with pytest.raises(SubmissionFailedError): - producer_client.blocking_stream_completed_submission(submission_id) + producer_client.blocking_stream_completed_submission( + submission_id, + timeout=SUBMISSION_COMPLETED_TIMEOUT, + ) # Cancelling the failed submission should fail. with pytest.raises(SubmissionNotCancellableError) as exc_info: producer_client.cancel_submission(submission_id) @@ -576,7 +608,10 @@ def consume(x: int) -> int | None: with background_process(run_consumer): with pytest.raises(SubmissionFailedError) as exc_info: - producer_client.blocking_stream_completed_submission(submission_id) + producer_client.blocking_stream_completed_submission( + submission_id, + timeout=SUBMISSION_COMPLETED_TIMEOUT, + ) assert exc_info.value.submission.chunks_done == len(chunks) - 1 @@ -696,3 +731,25 @@ def test_prefer_distinct_strategy_fairness(opsqueue: OpsqueueProcess) -> None: [chunk] = consumer_client.reserve_chunks(strategy=strategy) reserved_company_order.append(company_id_per_submission[chunk.submission_id]) assert reserved_company_order == [1, 2, 3] * chunks_per_company + + +def test_run_submission_timeout(opsqueue: OpsqueueProcess) -> None: + url = "file:///tmp/opsqueue/test_run_submission_timeout" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + + def run_consumer() -> None: + consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url) + + def process_op(x: int) -> int: + time.sleep(2.0) + return x + + consumer_client.run_each_op(process_op) + + with background_process(run_consumer) as _consumer: + with pytest.raises(TimeoutError): + producer_client.run_submission( + [1], + chunk_size=1, + timeout=0.1, + ) From 88100c40e5dde05dc77cfecae1b99c2f72e7bc6b Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 23 Jul 2026 16:16:43 +0200 Subject: [PATCH 03/11] Make complete_chunk and fail_chunk not error when processing previously completed, failed, or cancelled chunks Because of the idempotency assumption for processing chunks, nothing should break if we just ignore the error. Besides, we were already ignoring the error accidentally. --- .../python/opsqueue/exceptions.py | 9 -- libs/opsqueue_python/src/errors.rs | 23 +--- opsqueue/src/common/chunk.rs | 127 ++++++++++-------- opsqueue/src/common/errors.rs | 6 +- 4 files changed, 76 insertions(+), 89 deletions(-) diff --git a/libs/opsqueue_python/python/opsqueue/exceptions.py b/libs/opsqueue_python/python/opsqueue/exceptions.py index c946f339..542aa7a6 100644 --- a/libs/opsqueue_python/python/opsqueue/exceptions.py +++ b/libs/opsqueue_python/python/opsqueue/exceptions.py @@ -92,15 +92,6 @@ class TryFromIntError(IncorrectUsageError): pass -class ChunkNotFoundError(IncorrectUsageError): - """ - Raised when a method is used to look up information about a chunk - but the chunk doesn't exist within the Opsqueue. - """ - - pass - - class SubmissionNotFoundError(IncorrectUsageError): """ Raised when a method is used to look up information about a submission diff --git a/libs/opsqueue_python/src/errors.rs b/libs/opsqueue_python/src/errors.rs index 110dded0..b0de8e7e 100644 --- a/libs/opsqueue_python/src/errors.rs +++ b/libs/opsqueue_python/src/errors.rs @@ -2,16 +2,14 @@ /// so we have nice IDE support for docs-on-hover and for 'go to definition'. use std::error::Error; -use opsqueue::common::chunk::ChunkId; use opsqueue::common::errors::{ - ChunkNotFound, E, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound, - TooManyMatchingSubmissions, UnexpectedOpsqueueConsumerServerResponse, + E, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions, + UnexpectedOpsqueueConsumerServerResponse, }; use pyo3::exceptions::{PyBaseException, PyTimeoutError}; use pyo3::{Bound, PyErr, Python, import_exception}; use crate::common; -use crate::common::{ChunkIndex, SubmissionId}; // Expected errors: import_exception!(opsqueue.exceptions, SubmissionFailedError); @@ -19,7 +17,6 @@ import_exception!(opsqueue.exceptions, SubmissionFailedError); // Incorrect usage errors: import_exception!(opsqueue.exceptions, IncorrectUsageError); import_exception!(opsqueue.exceptions, TryFromIntError); -import_exception!(opsqueue.exceptions, ChunkNotFoundError); import_exception!(opsqueue.exceptions, SubmissionNotFoundError); import_exception!(opsqueue.exceptions, SubmissionNotCancellableError); import_exception!(opsqueue.exceptions, TooManyMatchingSubmissionsError); @@ -173,22 +170,6 @@ impl From> for PyErr { } } -impl From> for PyErr { - fn from(value: CError) -> Self { - let ChunkId { - submission_id, - chunk_index, - } = value.0.0; - ChunkNotFoundError::new_err(( - value.0.to_string(), - ( - SubmissionId::from(submission_id), - ChunkIndex::from(chunk_index), - ), - )) - } -} - impl From> for PyErr { fn from(value: CError) -> Self { NewObjectStoreClientError::new_err(value.0.to_string()) diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index 403207c6..bc267b2f 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -237,13 +237,13 @@ impl Chunk { #[cfg(feature = "server-logic")] pub mod db { use super::{ - Chunk, ChunkCompleted, ChunkFailed, ChunkId, ChunkIndex, ChunkSize, DateTime, SubmissionId, + Chunk, ChunkCompleted, ChunkFailed, ChunkId, ChunkIndex, DateTime, SubmissionId, Utc, }; - use crate::common::errors::{ChunkNotFound, DatabaseError, E, SubmissionNotFound}; + use crate::common::errors::{DatabaseError, E, SubmissionNotFound}; use crate::db::{Connection, True, WriterConnection}; use axum_prometheus::metrics::{counter, gauge}; use sqlx::{QueryBuilder, Sqlite}; - use sqlx::{query, query_as}; + use sqlx::{query, query_as, query_scalar}; impl<'q> sqlx::Encode<'q, Sqlite> for super::ChunkIndex { fn encode_by_ref( @@ -312,25 +312,18 @@ pub mod db { chunk_id: ChunkId, output_content: Option>, mut conn: impl WriterConnection, - ) -> Result<(), E>> { - let _chunk_size: Result>> = - conn.transaction(move |mut tx| { - Box::pin(async move { - let completed_work = - complete_chunk_raw(chunk_id, output_content, &mut tx).await?; - crate::common::submission::db::maybe_complete_submission( - chunk_id.submission_id, - &mut tx, - ) - .await - .map_err(|e| match e { - E::L(e) => E::L(e), - E::R(e) => E::R(E::L(e)), - })?; - Ok(completed_work.unwrap_or_default()) - }) + ) -> Result<(), E> { + conn.transaction(move |mut tx| { + Box::pin(async move { + complete_chunk_raw(chunk_id, output_content, &mut tx).await?; + crate::common::submission::db::maybe_complete_submission( + chunk_id.submission_id, + &mut tx, + ) + .await }) - .await; + }) + .await?; counter!(crate::prometheus::CHUNKS_COMPLETED_COUNTER).increment(1); Ok(()) @@ -346,9 +339,9 @@ pub mod db { chunk_id: ChunkId, output_content: Option>, mut tx: impl WriterConnection, - ) -> sqlx::Result> { + ) -> sqlx::Result<()> { let now = chrono::prelude::Utc::now(); - query!( + let chunk_moved = query!( " INSERT INTO chunks_completed (submission_id, chunk_index, output_content, completed_at) @@ -365,26 +358,42 @@ pub mod db { chunk_id.submission_id, chunk_id.chunk_index, ) - .fetch_one(tx.get_inner()) - .await?; - // Defense in depth: Above query should never be called twice on the same chunk. - // If it _does_ happen, it means that either a consumer is attempting a chunk they didn't reserve, - // or we gave out the same reservation twice. + .fetch_optional(tx.get_inner()) + .await? + .is_some(); + // Defense in depth: Above query could be called twice on the same chunk. For instance, + // when the server was restarted and the reservations are forgotten, and the same chunk + // was reserved again. + // + // In addition, cancelling a submission while a chunk is reserved also results in the chunk + // not being in the `chunks` table. Which is fine, because cancelled submissions count as + // failed. + // + // By only updating `chunks_done` when we actually moved a chunk, we ensure that we never + // mess up the submission's `chunks_done` counter. // - // By returning early if the chunk was not found, - // we ensure that even in these situations - // we never mess up the submission's `chunks_done` counter. + // This does mean we potentially run the same chunk twice, but that is fine because we + // assume chunks to be processed idempotently. // // (Not doing that resulted in a hard-to-track-down bug in the past. // https://github.com/channable/opsqueue/issues/76 // ) - sqlx::query_scalar!( - "UPDATE submissions SET chunks_done = chunks_done + 1 WHERE submissions.id = $1 RETURNING submissions.chunk_size;", - chunk_id.submission_id, - ) - .fetch_one(tx.get_inner()) - .await - .map(|opt| opt.map(ChunkSize)) + if chunk_moved { + sqlx::query_scalar!( + "UPDATE submissions SET chunks_done = chunks_done + 1 WHERE submissions.id = $1 RETURNING submissions.chunk_size;", + chunk_id.submission_id, + ) + .fetch_one(tx.get_inner()) + .await?; + } else { + tracing::warn!( + "Could not complete chunk {:?} because it was either: \ + completed, failed, or cancelled before. Ignoring.", + chunk_id + ); + } + + Ok(()) } /// Increment retries for a chunk, or move it to failed state. @@ -406,7 +415,7 @@ pub mod db { submission_id, chunk_index, } = chunk_id; - let fields = query!( + let retries = query_scalar!( " UPDATE chunks SET retries = retries + 1 WHERE submission_id = $1 AND chunk_index = $2 @@ -415,23 +424,33 @@ pub mod db { submission_id, chunk_index ) - .fetch_one(tx.get_inner()) + .fetch_optional(tx.get_inner()) .await?; - tracing::trace!("Retries: {}", fields.retries); - if fields.retries >= max_retries.into() { - crate::common::submission::db::fail_submission_notx( - submission_id, - chunk_index, - failure, - &mut tx, - ) - .await?; - - Ok::<_, sqlx::Error>(true) + if let Some(retries) = retries { + tracing::trace!("Retries: {}", retries); + if retries >= max_retries.into() { + crate::common::submission::db::fail_submission_notx( + submission_id, + chunk_index, + failure, + &mut tx, + ) + .await?; + + Ok::<_, sqlx::Error>(true) + } else { + counter!(crate::prometheus::CHUNKS_RETRIED_COUNTER).increment(1); + // When retrying, the chunk re-enters ('stays') in the backlog, + // so we *don't* decrement the backlog gauge here. + Ok::<_, sqlx::Error>(false) + } } else { - counter!(crate::prometheus::CHUNKS_RETRIED_COUNTER).increment(1); - // When retrying, the chunk re-enters ('stays') in the backlog, - // so we *don't* decrement the backlog gauge here. + tracing::warn!( + "Could not fail chunk {:?} because it was either: \ + completed, failed, or cancelled before. Ignoring.", + chunk_id + ); + Ok::<_, sqlx::Error>(false) } }) diff --git a/opsqueue/src/common/errors.rs b/opsqueue/src/common/errors.rs index 6527810b..6504ecd5 100644 --- a/opsqueue/src/common/errors.rs +++ b/opsqueue/src/common/errors.rs @@ -12,7 +12,7 @@ use thiserror::Error; use crate::consumer::common::SyncServerToClientResponse; use super::{ - chunk::{ChunkFailed, ChunkId}, + chunk::ChunkFailed, submission::{SubmissionCancelled, SubmissionCompleted, SubmissionFailed, SubmissionId}, }; @@ -28,10 +28,6 @@ impl From for E { } } -#[derive(Error, Debug)] -#[error("Chunk not found for ID {0:?}")] -pub struct ChunkNotFound(pub ChunkId); - #[derive(Error, Debug, Deserialize, Serialize)] #[error("Submission not found for ID {0:?}")] pub struct SubmissionNotFound(pub SubmissionId); From 82f3023d61a1bd59cce83654bb91f4762da5bb28 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 23 Jul 2026 16:32:15 +0200 Subject: [PATCH 04/11] Allow submissions to be created in a paused state Introduce `submissions_paused` and `chunks_paused` tables (alongside the existing `submissions_{completed,failed,cancelled}` and `chunks_{completed,failed}` tables). A submission can now be created in a Paused state. It's then stored in `submissions_paused` and its chunks are stored in `chunks_paused`. Because paused chunks are not in the `chunks` table, the consumer dispatcher naturally skips them without any changes to the dispatch query. Unpausing moves the submission and the chunks to `submissions` and `chunks` and notifies waiting consumers. Paused submissions are cancellable; `cancel_submission` now handles the case where the submission is found in `submissions_paused`. We don't allow pausing submissions after creation. That proved to have too many edge cases we would need to resolve. --- .../python/opsqueue/producer.py | 21 +- libs/opsqueue_python/src/common.rs | 44 +- libs/opsqueue_python/src/lib.rs | 1 + libs/opsqueue_python/src/producer.rs | 40 +- libs/opsqueue_python/tests/test_roundtrip.py | 66 +++ .../20260805143000_pausing.down.sql | 2 + .../migrations/20260805143000_pausing.up.sql | 22 + opsqueue/src/common/chunk.rs | 115 ++++- opsqueue/src/common/submission.rs | 442 +++++++++++++++++- opsqueue/src/consumer/client.rs | 1 + opsqueue/src/consumer/strategy.rs | 1 + opsqueue/src/producer/client.rs | 125 ++++- opsqueue/src/producer/common.rs | 4 + opsqueue/src/producer/server.rs | 34 +- opsqueue/src/prometheus.rs | 12 + 15 files changed, 893 insertions(+), 37 deletions(-) create mode 100644 opsqueue/migrations/20260805143000_pausing.down.sql create mode 100644 opsqueue/migrations/20260805143000_pausing.up.sql diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index dc807cc6..9bb72d6c 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/libs/opsqueue_python/python/opsqueue/producer.py @@ -27,6 +27,7 @@ SubmissionFailed, ChunkFailed, SubmissionNotCancellable, + SubmissionPaused, ) __all__ = [ @@ -39,6 +40,7 @@ "SubmissionNotCancellable", "SubmissionNotCancellableError", "SubmissionNotFoundError", + "SubmissionPaused", "TooManyMatchingSubmissionsError", "ChunkFailed", ] @@ -148,6 +150,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, @@ -164,6 +167,7 @@ def insert_submission( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + paused=paused, ) def blocking_stream_completed_submission( @@ -263,6 +267,7 @@ 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, @@ -279,6 +284,7 @@ def insert_submission_chunks( strategic_metadata=strategic_metadata, chunk_size=chunk_size, otel_trace_carrier=otel_trace_carrier, + paused=paused, ) def blocking_stream_completed_submission_chunks( @@ -334,7 +340,7 @@ def count_submissions(self) -> int: def cancel_submission(self, submission_id: SubmissionId) -> None: """ - Cancel a specific submission that is in progress. + Cancel a specific submission that is in progress or paused. Returns None if the submission was successfully cancelled. @@ -345,6 +351,19 @@ def cancel_submission(self, submission_id: SubmissionId) -> None: """ self.inner.cancel_submission(submission_id) + def unpause_submission(self, submission_id: SubmissionId) -> None: + """ + Unpause a specific submission that is currently paused, + making it available to consumers. + + Returns None if the submission was successfully unpaused. + + Raises: + - `SubmissionNotFoundError` if the submission is not currently paused. + - `InternalProducerClientError` if there is a low-level internal error. + """ + self.inner.unpause_submission(submission_id) + def get_submission_status( self, submission_id: SubmissionId ) -> SubmissionStatus | None: diff --git a/libs/opsqueue_python/src/common.rs b/libs/opsqueue_python/src/common.rs index 0402c85a..e2f6d378 100644 --- a/libs/opsqueue_python/src/common.rs +++ b/libs/opsqueue_python/src/common.rs @@ -366,12 +366,15 @@ pub enum SubmissionStatus { Cancelled { submission: SubmissionCancelled, }, + Paused { + submission: SubmissionPaused, + }, } impl From for SubmissionStatus { fn from(value: opsqueue::common::submission::SubmissionStatus) -> Self { use opsqueue::common::submission::SubmissionStatus::{ - Cancelled, Completed, Failed, InProgress, + Cancelled, Completed, Failed, InProgress, Paused, }; match value { InProgress(s) => SubmissionStatus::InProgress { @@ -388,6 +391,9 @@ impl From for SubmissionStatus { Cancelled(s) => SubmissionStatus::Cancelled { submission: s.into(), }, + Paused(s) => SubmissionStatus::Paused { + submission: s.into(), + }, } } } @@ -512,6 +518,42 @@ pub struct SubmissionCancelled { pub cancelled_at: DateTime, } +#[pyclass(from_py_object, frozen, get_all, module = "opsqueue")] +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SubmissionPaused { + pub id: SubmissionId, + pub chunks_total: u64, + pub chunks_done: u64, + pub metadata: Option, + pub strategic_metadata: StrategicMetadataMap, +} + +impl From for SubmissionPaused { + fn from(value: opsqueue::common::submission::SubmissionPaused) -> Self { + Self { + id: value.id.into(), + chunks_total: value.chunks_total.into(), + chunks_done: value.chunks_done.into(), + metadata: value.metadata, + strategic_metadata: value.strategic_metadata, + } + } +} + +#[pymethods] +impl SubmissionPaused { + fn __repr__(&self) -> String { + format!( + "SubmissionPaused(id={0}, chunks_total={1}, chunks_done={2}, metadata={3:?}, strategic_metadata={4:?})", + self.id.__repr__(), + self.chunks_total, + self.chunks_done, + self.metadata, + self.strategic_metadata + ) + } +} + /// Submission could not be cancelled because it was already completed, failed /// or cancelled. #[pyclass(from_py_object, frozen, module = "opsqueue")] diff --git a/libs/opsqueue_python/src/lib.rs b/libs/opsqueue_python/src/lib.rs index b5f804f5..27f6835c 100644 --- a/libs/opsqueue_python/src/lib.rs +++ b/libs/opsqueue_python/src/lib.rs @@ -24,6 +24,7 @@ fn opsqueue_internal(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index 45eb5f88..33f42189 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -160,6 +160,36 @@ impl ProducerClient { }) } + /// Unpause a paused submission, making it available to consumers again. + /// + /// Will return an error if the submission is not currently paused. + /// + /// # Errors + /// + /// Returns an error if the submission is not found or if an internal client error occurs. + #[allow(clippy::result_large_err, clippy::type_complexity)] + pub fn unpause_submission( + &self, + py: Python<'_>, + id: SubmissionId, + ) -> CPyResult< + (), + E![ + FatalPythonException, + SubmissionNotFound, + InternalProducerClientError + ], + > { + py.detach(|| { + self.block_unless_interrupted(async { + self.client + .unpause_submission(id.into()) + .await + .map_err(|e| CError(R(e))) + }) + }) + } + /// Retrieve the status (in progress, completed or failed) of a specific submission. /// /// The returned `SubmissionStatus` object also includes the number of chunks finished so far, @@ -251,7 +281,7 @@ impl ProducerClient { /// # Errors /// /// Returns an error if submission insertion fails. - #[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, otel_trace_carrier=CarrierMap::default(), paused=false))] pub fn insert_submission_direct( &self, py: Python<'_>, @@ -260,6 +290,7 @@ impl ProducerClient { strategic_metadata: Option, chunk_size: Option, otel_trace_carrier: CarrierMap, + paused: bool, ) -> CPyResult> { py.detach(|| { let submission = opsqueue::producer::InsertSubmission { @@ -269,6 +300,7 @@ impl ProducerClient { }, metadata, strategic_metadata: strategic_metadata.unwrap_or_default(), + paused, }; self.block_unless_interrupted(async move { self.client @@ -285,8 +317,8 @@ impl ProducerClient { /// # Errors /// /// Returns an error if chunk upload or submission insertion fails. - #[allow(clippy::type_complexity)] - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default()))] + #[allow(clippy::type_complexity, clippy::too_many_arguments)] + #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), paused=false))] pub fn insert_submission_chunks( &self, py: Python<'_>, @@ -295,6 +327,7 @@ impl ProducerClient { strategic_metadata: Option, chunk_size: Option, otel_trace_carrier: CarrierMap, + paused: bool, ) -> CPyResult< SubmissionId, E![ @@ -334,6 +367,7 @@ impl ProducerClient { }, metadata, strategic_metadata: strategic_metadata.unwrap_or_default(), + paused, }; self.client .insert_submission(&submission, &otel_trace_carrier) diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index 9052823d..3d6dddb0 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -27,6 +27,7 @@ strategy_from_description, ) import logging +import time import pytest SUBMISSION_COMPLETED_TIMEOUT = 10.0 @@ -753,3 +754,68 @@ def process_op(x: int) -> int: chunk_size=1, timeout=0.1, ) + + +def test_unpause_and_complete(opsqueue: OpsqueueProcess) -> None: + """Unpausing a paused submission makes it available to consumers again, + and it can be completed normally afterwards.""" + url = "file:///tmp/opsqueue/test_unpause_and_complete" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission( + (1, 2, 3), chunk_size=1, paused=True + ) + + assert isinstance( + producer_client.get_submission_status(submission_id), SubmissionStatus.Paused + ) + + producer_client.unpause_submission(submission_id) + assert isinstance( + producer_client.get_submission_status(submission_id), + SubmissionStatus.InProgress, + ) + + def run_consumer() -> None: + consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url) + consumer_client.run_each_op(lambda x: x) + + with background_process(run_consumer): + producer_client.blocking_stream_completed_submission(submission_id) + assert isinstance( + producer_client.get_submission_status(submission_id), + SubmissionStatus.Completed, + ) + + +def test_unpause_not_found(opsqueue: OpsqueueProcess) -> None: + """Unpausing a submission that is not paused (e.g. in-progress) raises + SubmissionNotFoundError.""" + url = "file:///tmp/opsqueue/test_unpause_not_found" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission( + (1, 2, 3), chunk_size=1, paused=False + ) + assert isinstance( + producer_client.get_submission_status(submission_id), + SubmissionStatus.InProgress, + ) + with pytest.raises(SubmissionNotFoundError): + producer_client.unpause_submission(submission_id) + + +def test_cancel_paused(opsqueue: OpsqueueProcess) -> None: + """A paused submission can be cancelled; its status becomes Cancelled.""" + url = "file:///tmp/opsqueue/test_cancel_paused" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission( + (1, 2, 3), chunk_size=1, paused=True + ) + + assert isinstance( + producer_client.get_submission_status(submission_id), SubmissionStatus.Paused + ) + + producer_client.cancel_submission(submission_id) + assert isinstance( + producer_client.get_submission_status(submission_id), SubmissionStatus.Cancelled + ) diff --git a/opsqueue/migrations/20260805143000_pausing.down.sql b/opsqueue/migrations/20260805143000_pausing.down.sql new file mode 100644 index 00000000..eed0c27f --- /dev/null +++ b/opsqueue/migrations/20260805143000_pausing.down.sql @@ -0,0 +1,2 @@ +DROP TABLE chunks_paused; +DROP TABLE submissions_paused; diff --git a/opsqueue/migrations/20260805143000_pausing.up.sql b/opsqueue/migrations/20260805143000_pausing.up.sql new file mode 100644 index 00000000..2a3a60f4 --- /dev/null +++ b/opsqueue/migrations/20260805143000_pausing.up.sql @@ -0,0 +1,22 @@ +CREATE TABLE submissions_paused +( + id BIGINT PRIMARY KEY NOT NULL, + prefix TEXT, + chunks_total INTEGER NOT NULL DEFAULT 0, + chunks_done INTEGER NOT NULL DEFAULT 0, + metadata BLOB, + otel_trace_carrier TEXT NOT NULL DEFAULT '{}', + chunk_size INTEGER +); + +CREATE INDEX submissions_paused_prefix ON submissions_paused (prefix, id); + +CREATE TABLE chunks_paused +( + submission_id INTEGER NOT NULL, + chunk_index INTEGER NOT NULL, + input_content BLOB NULL, + retries INTEGER NOT NULL DEFAULT 0, + + PRIMARY KEY (submission_id, chunk_index) +) WITHOUT ROWID, STRICT; diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index bc267b2f..74de9ed4 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -613,6 +613,93 @@ pub mod db { Ok(()) } + /// # Errors + /// + /// Returns an error if a SQL query fails. + #[tracing::instrument(skip(chunks, conn))] + pub async fn insert_many_paused_chunks( + chunks: &[Chunk], + mut conn: impl WriterConnection, + ) -> sqlx::Result<()> { + const ROWS_PER_QUERY: usize = 1000; + + let mut iter = chunks.iter().peekable(); + while iter.peek().is_some() { + let query_chunks = iter.by_ref().take(ROWS_PER_QUERY); + + let mut query_builder: QueryBuilder = QueryBuilder::new( + "INSERT INTO chunks_paused (submission_id, chunk_index, input_content) ", + ); + query_builder.push_values(query_chunks, |mut b, chunk| { + b.push_bind(chunk.submission_id) + .push_bind(chunk.chunk_index) + .push_bind(chunk.input_content.clone()); + }); + let query = query_builder.build(); + + query.execute(conn.get_inner()).await?; + } + + Ok(()) + } + + /// Move all chunks of a paused submission from `chunks_paused` back to `chunks`. + /// + /// # Errors + /// + /// Returns an error if the SQL query fails. + #[tracing::instrument(skip(conn))] + pub async fn restore_paused_chunks( + submission_id: SubmissionId, + mut conn: impl WriterConnection, + ) -> sqlx::Result<()> { + sqlx::query!( + " + INSERT INTO chunks (submission_id, chunk_index, input_content, retries) + SELECT submission_id, chunk_index, input_content, retries FROM chunks_paused WHERE submission_id = $1; + + DELETE FROM chunks_paused WHERE submission_id = $2; + ", + submission_id, + submission_id, + ) + .execute(conn.get_inner()) + .await?; + Ok(()) + } + + /// Skip (cancel) all chunks of a paused submission by moving them from + /// `chunks_paused` to `chunks_failed` with `skipped = true`. + /// + /// # Errors + /// + /// Returns an error if the SQL query fails. + #[tracing::instrument(skip(conn))] + pub async fn skip_remaining_paused_chunks( + submission_id: SubmissionId, + mut conn: impl WriterConnection, + ) -> sqlx::Result<()> { + let now = chrono::prelude::Utc::now(); + + let query_res = sqlx::query!( + " + INSERT INTO chunks_failed + (submission_id, chunk_index, input_content, failure, skipped, failed_at) + SELECT submission_id, chunk_index, input_content, '', 1, julianday($1) FROM chunks_paused WHERE submission_id = $2; + + DELETE FROM chunks_paused WHERE submission_id = $3; + ", + now, + submission_id, + submission_id, + ) + .execute(conn.get_inner()) + .await?; + + counter!(crate::prometheus::CHUNKS_SKIPPED_COUNTER).increment(query_res.rows_affected()); + Ok(()) + } + /// Mark all remaining chunks for a submission as skipped/failed. /// /// # Errors @@ -696,6 +783,23 @@ pub mod db { Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) } + /// Count paused chunks. + /// + /// # Errors + /// + /// Returns an error if the count query fails. + /// + /// # Panics + /// + /// Panics if `COUNT(*)` returns a negative value, which `SQLite` never does. + #[tracing::instrument(skip(db))] + pub async fn count_chunks_paused(mut db: impl Connection) -> sqlx::Result { + let count = sqlx::query_scalar!("SELECT COUNT(1) as count FROM chunks_paused;") + .fetch_one(db.get_inner()) + .await?; + Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) + } + /// Looks up the number of operations in the backlog. /// /// An estimation that returns a slightly too high number, @@ -792,10 +896,7 @@ pub mod test { .expect("complete chunk failed"); assert_eq!(count_chunks(&mut conn).await.unwrap(), 0); - assert_eq!( - count_chunks_completed(&mut conn).await.unwrap(), - 1 - ); + assert_eq!(count_chunks_completed(&mut conn).await.unwrap(), 1); assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), 0); } @@ -809,6 +910,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -864,10 +966,7 @@ pub mod test { .expect("Succeed chunk failed"); assert_eq!(count_chunks(&mut conn).await.unwrap(), 0); - assert_eq!( - count_chunks_completed(&mut conn).await.unwrap(), - 0 - ); + assert_eq!(count_chunks_completed(&mut conn).await.unwrap(), 0); assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), 1); } } diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 9f586075..6333ef8b 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -212,12 +212,32 @@ pub struct SubmissionCancelled { pub cancelled_at: DateTime, } +/// A submission that has been paused. +/// +/// A submission can only be submitted in a paused state. We don't support pausing submissions +/// after submission. +/// +/// A paused submission can be unpaused or canceled. +#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] +pub struct SubmissionPaused { + pub id: SubmissionId, + pub prefix: Option, + pub chunks_total: ChunkCount, + pub chunks_done: ChunkCount, + pub chunk_size: ChunkSize, + pub metadata: Option, + #[serde(default)] + pub strategic_metadata: StrategicMetadataMap, + pub otel_trace_carrier: String, +} + #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum SubmissionStatus { InProgress(Submission), Completed(SubmissionCompleted), Failed(SubmissionFailed, ChunkFailed), Cancelled(SubmissionCancelled), + Paused(SubmissionPaused), } impl Default for Submission { @@ -284,6 +304,7 @@ pub mod db { DatabaseError, E, SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions, }, + submission::SubmissionPaused, }, db::{Connection, True, WriterConnection, WriterPool}, }; @@ -429,9 +450,123 @@ pub mod db { res } + #[tracing::instrument(skip(chunks, conn))] + pub(crate) async fn insert_paused_submission( + submission: Submission, + chunks: Vec, + mut conn: impl WriterConnection, + ) -> Result<(), DatabaseError> { + use axum_prometheus::metrics::counter; + use futures::FutureExt as _; + + let chunks_total = submission.chunks_total.into(); + tracing::debug!("Inserting paused submission {}", submission.id); + + let res = conn + .transaction(move |mut tx| { + async move { + insert_paused_submission_raw(&submission, &mut tx).await?; + insert_submission_metadata_raw( + &submission, + &submission.strategic_metadata, + &mut tx, + ) + .await?; + super::chunk::db::insert_many_paused_chunks(&chunks, &mut tx).await?; + Ok(()) + } + .boxed() + }) + .await; + + counter!(crate::prometheus::SUBMISSIONS_PAUSED_COUNTER).increment(1); + counter!(crate::prometheus::SUBMISSIONS_TOTAL_COUNTER).increment(1); + counter!(crate::prometheus::CHUNKS_TOTAL_COUNTER).increment(chunks_total); + res + } + + #[tracing::instrument(skip(conn))] + async fn insert_paused_submission_raw( + submission: &Submission, + mut conn: impl WriterConnection, + ) -> Result<(), DatabaseError> { + sqlx::query!( + " + INSERT INTO submissions_paused (id, prefix, chunks_total, chunks_done, metadata, otel_trace_carrier, chunk_size) + VALUES ($1, $2, $3, $4, $5, $6, $7) + ", + submission.id, + submission.prefix, + submission.chunks_total, + submission.chunks_done, + submission.metadata, + submission.otel_trace_carrier, + submission.chunk_size.0, + ) + .execute(conn.get_inner()) + .await?; + + Ok(()) + } + + /// Unpause a paused submission. Atomically moves it back from `submissions_paused` + /// to `submissions` and its chunks from `chunks_paused` to `chunks`. + /// + /// # Errors + /// + /// Returns [`DatabaseError`] if the transaction or any SQL query fails. + /// + /// Returns [`SubmissionNotFound`] if the submission is not currently paused. + #[tracing::instrument(skip(conn))] + pub async fn unpause_submission( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + conn.transaction(move |mut tx| { + Box::pin(async move { + unpause_submission_raw(id, &mut tx).await?; + super::chunk::db::restore_paused_chunks(id, &mut tx).await?; + Ok(()) + }) + }) + .await + } + + #[tracing::instrument(skip(conn))] + pub(super) async fn unpause_submission_raw( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + let row = query!( + " + INSERT INTO submissions + (id, chunks_total, chunks_done, prefix, metadata, otel_trace_carrier, chunk_size) + SELECT id, chunks_total, chunks_done, prefix, metadata, otel_trace_carrier, chunk_size + FROM submissions_paused WHERE id = $1; + + DELETE FROM submissions_paused WHERE id = $2 RETURNING *; + ", + id, + id, + ) + .fetch_optional(conn.get_inner()) + .await?; + if row.is_none() { + Err(E::R(SubmissionNotFound(id))) + } else { + counter!(crate::prometheus::SUBMISSIONS_UNPAUSED_COUNTER).increment(1); + Ok(()) + } + } + /// Creates a new submission with the given chunks and inserts it into the database. /// - /// If the number of chunks is 0, the submission is marked as completed immediately afterwards. + /// If `paused` is false and the number of chunks is 0, the submission is marked + /// as completed immediately afterwards. + /// + /// If `paused` is true, the submission is inserted directly into `submissions_paused` + /// (and its chunks into `chunks_paused`), so it won't be picked up by consumers + /// until explicitly unpaused. Zero-chunk paused submissions stay paused. /// /// # Panics /// @@ -447,6 +582,7 @@ pub mod db { metadata: Option, strategic_metadata: StrategicMetadataMap, chunk_size: ChunkSize, + paused: bool, mut conn: impl WriterConnection, ) -> Result { let submission_id = SubmissionId::new(); @@ -462,7 +598,7 @@ pub mod db { strategic_metadata, otel_trace_carrier, }; - let iter = chunks_contents + let chunks: Vec = chunks_contents .into_iter() .enumerate() .map(move |(chunk_index, uri)| { @@ -470,25 +606,30 @@ pub mod db { Chunk::new(submission_id, chunk_index.try_into().unwrap(), uri) }) .collect(); - insert_submission(submission, iter, &mut conn).await?; - // Empty submissions get special handling: we mark them as completed right away. - // See https://github.com/channable/opsqueue/issues/86 for rationale. - if len == 0 { - match maybe_complete_submission(submission_id, conn).await { - // Forward our database errors to the caller. - Err(E::L(e)) => return Err(e), - // If the submission ID can't be found, that's too bad, but it's not our problem anymore i guess. - Err(E::R(_)) => { - tracing::warn!(%submission_id, "Presumed zero-length submission not found"); - } - // If everything went OK, this *could* still indicate a bug in producer code, so let's just log it. - // Our future selves might thank us. - Ok(true) => { - tracing::debug!(%submission_id, "Zero-length submission marked as completed"); - } - // This should never happen. If it does, better log it. - Ok(false) => { - tracing::warn!(%submission_id, "Zero-length submission wasn't zero-length?!"); + + if paused { + insert_paused_submission(submission, chunks, &mut conn).await?; + } else { + insert_submission(submission, chunks, &mut conn).await?; + // Empty submissions get special handling: we mark them as completed right away. + // See https://github.com/channable/opsqueue/issues/86 for rationale. + if len == 0 { + match maybe_complete_submission(submission_id, conn).await { + // Forward our database errors to the caller. + Err(E::L(e)) => return Err(e), + // If the submission ID can't be found, that's too bad, but it's not our problem anymore i guess. + Err(E::R(_)) => { + tracing::warn!(%submission_id, "Presumed zero-length submission not found"); + } + // If everything went OK, this *could* still indicate a bug in producer code, so let's just log it. + // Our future selves might thank us. + Ok(true) => { + tracing::debug!(%submission_id, "Zero-length submission marked as completed"); + } + // This should never happen. If it does, better log it. + Ok(false) => { + tracing::warn!(%submission_id, "Zero-length submission wasn't zero-length?!"); + } } } } @@ -579,6 +720,8 @@ pub mod db { r#" SELECT id AS "id: SubmissionId" FROM submissions WHERE prefix = $1 UNION ALL + SELECT id AS "id: SubmissionId" FROM submissions_paused WHERE prefix = $1 + UNION ALL SELECT id AS "id: SubmissionId" FROM submissions_completed WHERE prefix = $1 UNION ALL SELECT id AS "id: SubmissionId" FROM submissions_failed WHERE prefix = $1 @@ -806,6 +949,40 @@ pub mod db { return Ok(Some(SubmissionStatus::Cancelled(cancelled_submission))); } + let paused_row_opt = query!( + r#" + SELECT + id AS "id: SubmissionId" + , prefix + , chunks_total AS "chunks_total: ChunkCount" + , chunks_done AS "chunks_done: ChunkCount" + , chunk_size AS "chunk_size!: ChunkSize" + , metadata + , ( SELECT json_group_object(metadata_key, metadata_value) + FROM submissions_metadata + WHERE submission_id = submissions_paused.id + ) AS "strategic_metadata!: sqlx::types::Json" + , otel_trace_carrier + FROM submissions_paused WHERE id = $1 + "#, + id + ) + .fetch_optional(conn.get_inner()) + .await?; + if let Some(row) = paused_row_opt { + let paused_submission = SubmissionPaused { + id: row.id, + prefix: row.prefix, + chunks_total: row.chunks_total, + chunks_done: row.chunks_done, + chunk_size: row.chunk_size, + metadata: row.metadata, + strategic_metadata: row.strategic_metadata.0, + otel_trace_carrier: row.otel_trace_carrier, + }; + return Ok(Some(SubmissionStatus::Paused(paused_submission))); + } + Ok(None) } @@ -875,6 +1052,15 @@ pub mod db { Ok(Some(SubmissionStatus::Cancelled(submission))) => { Err(E::R(E::R(SubmissionNotCancellable::Cancelled(submission)))) } + Ok(Some(SubmissionStatus::Paused(_))) => { + // Paused submissions are cancellable. + cancel_paused_submission_notx(id, &mut tx).await.map_err( + |e| match e { + E::L(db_err) => E::L(db_err), + E::R(not_found) => E::R(E::L(not_found)), + }, + ) + } Err(db_err) => Err(E::L(db_err)), } } @@ -898,6 +1084,22 @@ pub mod db { Ok(()) } + /// Do not call directly! Must be called inside a transaction. + /// + /// # Errors + /// + /// Returns [`DatabaseError`] if any SQL query fails. + /// + /// Returns [`SubmissionNotFound`] if the submission is not found in `submissions_paused`. + pub async fn cancel_paused_submission_notx( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + cancel_paused_submission_raw(id, &mut conn).await?; + super::chunk::db::skip_remaining_paused_chunks(id, conn).await?; + Ok(()) + } + #[tracing::instrument(skip(conn))] pub(super) async fn cancel_submission_raw( id: SubmissionId, @@ -930,6 +1132,38 @@ pub mod db { } } + #[tracing::instrument(skip(conn))] + pub(super) async fn cancel_paused_submission_raw( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + let now = chrono::prelude::Utc::now(); + + 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_paused WHERE id = $2; + + DELETE FROM submissions_paused WHERE id = $3 RETURNING *; + ", + now, + id, + id, + ) + .fetch_optional(conn.get_inner()) + .await?; + if submission_opt.is_none() { + Err(E::R(SubmissionNotFound(id))) + } else { + counter!(crate::prometheus::SUBMISSIONS_CANCELLED_COUNTER).increment(1); + histogram!(crate::prometheus::SUBMISSIONS_DURATION_CANCEL_HISTOGRAM).record( + crate::prometheus::time_delta_as_f64(Utc::now() - id.timestamp()), + ); + Ok(()) + } + } + #[tracing::instrument(skip(conn))] /// Do not call directly! MUST be called inside a transaction. pub(super) async fn complete_submission_raw( @@ -1091,6 +1325,40 @@ pub mod db { Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) } + /// Count paused submissions. + /// + /// # Errors + /// + /// Returns an error if the count query fails. + /// + /// # Panics + /// + /// Panics if `COUNT(*)` returns a negative value, which `SQLite` never does. + #[tracing::instrument(skip(db))] + pub async fn count_submissions_paused(mut db: impl Connection) -> sqlx::Result { + let count = sqlx::query_scalar!("SELECT COUNT(1) as count FROM submissions_paused;") + .fetch_one(db.get_inner()) + .await?; + Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) + } + + /// Count cancelled submissions. + /// + /// # Errors + /// + /// Returns an error if the count query fails. + /// + /// # Panics + /// + /// Panics if `COUNT(*)` returns a negative value, which `SQLite` never does. + #[tracing::instrument(skip(db))] + pub async fn count_submissions_cancelled(mut db: impl Connection) -> sqlx::Result { + let count = sqlx::query_scalar!("SELECT COUNT(1) as count FROM submissions_cancelled;") + .fetch_one(db.get_inner()) + .await?; + Ok(u64::try_from(count).expect("COUNT(*) is always non-negative")) + } + /// Transactionally removes all completed/failed submissions, /// including all their chunks and associated strategic metadata. /// @@ -1204,8 +1472,10 @@ pub mod test { use itertools::Itertools; use sqlformat::{FormatOptions, QueryParams, format}; use sqlx::{Row, SqliteConnection}; + use std::assert_matches; use crate::common::StrategicMetadataMap; + use crate::common::chunk::db::{count_chunks, count_chunks_failed, count_chunks_paused}; use crate::db::{Connection as _, WriterPool}; use super::db::*; @@ -1444,6 +1714,7 @@ pub mod test { None, strategic_metadata.clone(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1519,6 +1790,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1529,6 +1801,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1539,6 +1812,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1549,6 +1823,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1581,6 +1856,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1591,6 +1867,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1601,6 +1878,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1654,6 +1932,7 @@ pub mod test { StrategicMetadataMap::default(), // chunk size ChunkSize::default(), + false, &mut conn, ) .await @@ -1759,4 +2038,125 @@ pub mod test { let deserialized: SubmissionCancelled = serde_json::from_value(json).unwrap(); assert_eq!(deserialized, cancelled); } + + #[sqlx::test(migrator = "crate::MIGRATOR")] + pub async fn test_query_plan_submission_status_paused(db: sqlx::SqlitePool) { + let mut conn = db.acquire().await.unwrap(); + let query = r" + SELECT + id + , prefix + , chunks_total + , chunks_done + , chunk_size + , metadata + , ( SELECT json_group_object(metadata_key, metadata_value) + FROM submissions_metadata + WHERE submission_id = submissions_paused.id + ) AS strategic_metadata + , otel_trace_carrier + FROM submissions_paused WHERE id = 1 + "; + + let explained = explain_query_plan(query, &mut conn).await; + assert_non_regressing_query_plan(query, &explained); + insta::assert_snapshot!(explained, @r" + 3, 0, SEARCH submissions_paused USING INDEX sqlite_autoindex_submissions_paused_1 (id=?) + 15, 0, CORRELATED SCALAR SUBQUERY 1 + 20, 15, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) + "); + } + + #[sqlx::test(migrator = "crate::MIGRATOR")] + pub async fn test_unpause_submission(db: sqlx::SqlitePool) { + let db = WriterPool::new(db); + let mut conn = db.writer_conn().await.unwrap(); + let (submission, chunks) = Submission::from_vec( + vec![Some("foo".into()), Some("bar".into()), Some("baz".into())], + None, + ChunkSize::default(), + ) + .unwrap(); + insert_paused_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + assert_eq!(count_submissions(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_paused(&mut conn).await.unwrap(), 1); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 0); + assert_eq!(count_chunks_paused(&mut conn).await.unwrap(), 3); + + unpause_submission(submission.id, &mut conn).await.unwrap(); + assert_eq!(count_submissions(&mut conn).await.unwrap(), 1); + assert_eq!(count_submissions_paused(&mut conn).await.unwrap(), 0); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 3); + assert_eq!(count_chunks_paused(&mut conn).await.unwrap(), 0); + } + + #[sqlx::test(migrator = "crate::MIGRATOR")] + pub async fn test_cancel_paused_submission(db: sqlx::SqlitePool) { + let db = WriterPool::new(db); + let mut conn = db.writer_conn().await.unwrap(); + let (submission, chunks) = Submission::from_vec( + vec![Some("foo".into()), Some("bar".into()), Some("baz".into())], + None, + ChunkSize::default(), + ) + .unwrap(); + insert_paused_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + assert_eq!(count_submissions_paused(&mut conn).await.unwrap(), 1); + assert_eq!(count_chunks_paused(&mut conn).await.unwrap(), 3); + + cancel_submission(submission.id, &mut conn).await.unwrap(); + + assert_eq!(count_submissions(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_paused(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_completed(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_failed(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_cancelled(&mut conn).await.unwrap(), 1); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 0); + assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), 3); + assert_eq!(count_chunks_paused(&mut conn).await.unwrap(), 0); + } + + #[sqlx::test(migrator = "crate::MIGRATOR")] + pub async fn test_submission_status_paused(db: sqlx::SqlitePool) { + let db = WriterPool::new(db); + let mut conn = db.writer_conn().await.unwrap(); + let (submission, chunks) = Submission::from_vec( + vec![Some("foo".into()), Some("bar".into()), Some("baz".into())], + None, + ChunkSize::default(), + ) + .unwrap(); + insert_paused_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + let status = submission_status(submission.id, &mut conn) + .await + .unwrap() + .unwrap(); + assert_matches!(status, SubmissionStatus::Paused(_)); + } + + #[sqlx::test(migrator = "crate::MIGRATOR")] + /// Test that an empty submission inserted in the paused state stays paused + /// (unlike empty non-paused submissions which are auto-completed). + pub async fn insert_empty_paused_submission_stays_paused(db: sqlx::SqlitePool) { + let db = WriterPool::new(db); + let mut conn = db.writer_conn().await.unwrap(); + let (submission, chunks) = + Submission::from_vec(vec![], None, ChunkSize::default()).unwrap(); + insert_paused_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + assert_eq!(count_submissions(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_paused(&mut conn).await.unwrap(), 1); + assert_eq!(count_submissions_completed(&mut conn).await.unwrap(), 0); + } } diff --git a/opsqueue/src/consumer/client.rs b/opsqueue/src/consumer/client.rs index 94abc811..ac37557c 100644 --- a/opsqueue/src/consumer/client.rs +++ b/opsqueue/src/consumer/client.rs @@ -578,6 +578,7 @@ mod tests { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await diff --git a/opsqueue/src/consumer/strategy.rs b/opsqueue/src/consumer/strategy.rs index b4decf26..8339be4f 100644 --- a/opsqueue/src/consumer/strategy.rs +++ b/opsqueue/src/consumer/strategy.rs @@ -849,6 +849,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await diff --git a/opsqueue/src/producer/client.rs b/opsqueue/src/producer/client.rs index 1632050b..09e4ab5b 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -203,6 +203,49 @@ impl Client { .await } + /// Unpause a paused submission, making it available to consumers again. + /// + /// Returns an error if the submission is not currently paused. + /// + /// # Errors + /// + /// Returns an error if the HTTP request fails or the server returns an unexpected status. + pub async fn unpause_submission( + &self, + submission_id: SubmissionId, + ) -> Result<(), E![SubmissionNotFound, InternalProducerClientError]> { + (|| async { + let base_url = &self.base_url; + let response = self + .http_client + .post(format!("{base_url}/submissions/unpause/{submission_id}")) + .send() + .await + .map_err(|e| R(e.into()))?; + let status = response.status(); + match status { + StatusCode::OK => Ok(()), + StatusCode::NOT_FOUND => { + let not_found_err = response + .json::() + .await + .map_err(|e| R(e.into()))?; + Err(L(not_found_err)) + } + _ => Err(R(InternalProducerClientError::UnexpectedStatus(status))), + } + }) + .retry(retry_policy()) + .when(|e| match e { + L(_) => false, + R(client_err) => client_err.is_ephemeral(), + }) + .notify(|err, dur| { + tracing::debug!("retrying error {err:?} with sleeping {dur:?}"); + }) + .await + } + /// Get the status of an existing submission identified by its `submission_id`. /// /// This uses the GET `/producer/submissions` endpoint. @@ -438,6 +481,7 @@ mod tests { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -467,6 +511,7 @@ mod tests { metadata: None, strategic_metadata: StrategicMetadataMap::default(), chunk_size: None, + paused: false, }; client .insert_submission(&submission, &std::collections::HashMap::default()) @@ -510,6 +555,7 @@ mod tests { metadata: None, strategic_metadata: StrategicMetadataMap::default(), chunk_size: None, + paused: false, }; let submission_id = client .insert_submission(&submission, &std::collections::HashMap::default()) @@ -524,7 +570,8 @@ mod tests { match status { SubmissionStatus::Completed(_) | SubmissionStatus::Failed(_, _) - | SubmissionStatus::Cancelled(_) => { + | SubmissionStatus::Cancelled(_) + | SubmissionStatus::Paused(_) => { panic!("Expected a SubmissionStatus that is still Inprogress, got: {status:?}"); } SubmissionStatus::InProgress(submission) => { @@ -534,4 +581,80 @@ mod tests { } } } + + #[sqlx::test(migrator = "crate::MIGRATOR")] + async fn test_insert_paused_submission_and_unpause(pool: sqlx::SqlitePool) { + let url = "0.0.0.0:4003"; + start_server_in_background(&pool, url).await; + let client = Client::new(url); + + let pool = WriterPool::new(pool); + let mut conn = pool.writer_conn().await.unwrap(); + let count = submission::db::count_submissions(&mut conn) + .await + .expect("Should be OK"); + assert_eq!(count, 0); + + let submission = InsertSubmission { + chunk_contents: ChunkContents::Direct { + contents: vec![None, None, None], + }, + metadata: None, + strategic_metadata: StrategicMetadataMap::default(), + chunk_size: None, + paused: true, + }; + let submission_id = client + .insert_submission(&submission, &std::collections::HashMap::default()) + .await + .expect("Should be OK"); + + let count = submission::db::count_submissions_paused(&mut conn) + .await + .expect("Should be OK"); + assert_eq!(count, 1); + + let status: SubmissionStatus = client + .get_submission(submission_id) + .await + .expect("Should be OK") + .expect("Should be Some"); + match status { + SubmissionStatus::Completed(_) + | SubmissionStatus::Failed(_, _) + | SubmissionStatus::Cancelled(_) + | SubmissionStatus::InProgress(_) => { + panic!("Expected a SubmissionStatus that is Paused, got: {status:?}"); + } + SubmissionStatus::Paused(submission) => { + assert_eq!(submission.chunks_done, 0); + assert_eq!(submission.chunks_total, 3); + assert_eq!(submission.id, submission_id); + } + } + + client + .unpause_submission(submission_id) + .await + .expect("Should be OK"); + + let status: SubmissionStatus = client + .get_submission(submission_id) + .await + .expect("Should be OK") + .expect("Should be Some"); + match status { + SubmissionStatus::Completed(_) + | SubmissionStatus::Failed(_, _) + | SubmissionStatus::Cancelled(_) + | SubmissionStatus::Paused(_) => { + panic!("Expected a SubmissionStatus that is InProgress, got: {status:?}"); + } + SubmissionStatus::InProgress(submission) => { + assert_eq!(submission.chunks_done, 0); + assert_eq!(submission.chunks_total, 3); + assert_eq!(submission.id, submission_id); + } + } + } } diff --git a/opsqueue/src/producer/common.rs b/opsqueue/src/producer/common.rs index 41f3c8e9..f98b4b86 100644 --- a/opsqueue/src/producer/common.rs +++ b/opsqueue/src/producer/common.rs @@ -10,6 +10,10 @@ pub struct InsertSubmission { #[serde(default)] pub strategic_metadata: StrategicMetadataMap, pub chunk_size: Option, + /// When `true`, the submission is inserted in a paused state and will not + /// be dispatched to consumers until explicitly unpaused. + #[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 1fa4d678..7e603233 100644 --- a/opsqueue/src/producer/server.rs +++ b/opsqueue/src/producer/server.rs @@ -67,6 +67,10 @@ impl ServerState { "/submissions/cancel/{submission_id}", post(cancel_submission), ) + .route( + "/submissions/unpause/{submission_id}", + post(unpause_submission), + ) .route( "/submissions/count_completed", get(submissions_count_completed), @@ -138,6 +142,29 @@ async fn cancel_submission( } } +/// 200 if the submission was successfully unpaused. +/// 404 if the submission could not be found in the paused state. +/// 500 if a `DatabaseError` occurred. +async fn unpause_submission( + State(state): State, + Path(submission_id): Path, +) -> Result<(), Response> { + let mut conn = state + .pool + .writer_conn() + .await + .map_err(|e| ServerError(e.into()).into_response())?; + match submission::db::unpause_submission(submission_id, &mut conn).await { + Ok(()) => { + // Wake up any waiting consumers now that new chunks are available. + state.notify_on_insert.notify_waiters(); + Ok(()) + } + Err(L(db_err)) => Err(ServerError(db_err.into()).into_response()), + Err(R(not_found_err)) => Err((StatusCode::NOT_FOUND, Json(not_found_err)).into_response()), + } +} + async fn submission_status( State(state): State, Path(submission_id): Path, @@ -199,6 +226,7 @@ async fn insert_submission( request.metadata, request.strategic_metadata, request.chunk_size.unwrap_or_default(), + request.paused, &mut conn, ) .await?; @@ -208,8 +236,10 @@ async fn insert_submission( // this is the moment to perform an extra WAL checkpoint let _ = db::perform_explicit_wal_checkpoint(conn).await; - // We've done a new insert! Let's tell any waiting consumers! - state.notify_on_insert.notify_waiters(); + // Notify waiting consumers, but only for non-paused submissions. + if !request.paused { + state.notify_on_insert.notify_waiters(); + } Ok(Json(submission_id)) } diff --git a/opsqueue/src/prometheus.rs b/opsqueue/src/prometheus.rs index d4b308a1..29f63baf 100644 --- a/opsqueue/src/prometheus.rs +++ b/opsqueue/src/prometheus.rs @@ -19,6 +19,8 @@ pub const SUBMISSIONS_TOTAL_COUNTER: &str = "submissions_total_count"; pub const SUBMISSIONS_COMPLETED_COUNTER: &str = "submissions_completed_count"; pub const SUBMISSIONS_FAILED_COUNTER: &str = "submissions_failed_count"; pub const SUBMISSIONS_CANCELLED_COUNTER: &str = "submissions_cancelled_count"; +pub const SUBMISSIONS_PAUSED_COUNTER: &str = "submissions_paused_count"; +pub const SUBMISSIONS_UNPAUSED_COUNTER: &str = "submissions_unpaused_count"; pub const SUBMISSIONS_DURATION_COMPLETE_HISTOGRAM: &str = "submissions_complete_duration_seconds"; pub const SUBMISSIONS_DURATION_FAIL_HISTOGRAM: &str = "submissions_fail_duration_seconds"; pub const SUBMISSIONS_DURATION_CANCEL_HISTOGRAM: &str = "submissions_cancel_duration_seconds"; @@ -67,6 +69,16 @@ pub fn describe_metrics() { Unit::Count, "Number of submissions cancelled (client-requested cancellation, not failure) permanently" ); + describe_counter!( + SUBMISSIONS_PAUSED_COUNTER, + Unit::Count, + "Number of submissions paused" + ); + describe_counter!( + SUBMISSIONS_UNPAUSED_COUNTER, + Unit::Count, + "Number of submissions unpaused (resumed)" + ); describe_histogram!( SUBMISSIONS_DURATION_COMPLETE_HISTOGRAM, Unit::Seconds, From a6e23cecc5888eb6fb4be06d994e53386465de76 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 5 Aug 2026 12:06:01 +0200 Subject: [PATCH 05/11] Address feedback --- libs/opsqueue_python/src/producer.rs | 2 +- opsqueue/src/common/chunk.rs | 4 ++-- opsqueue/src/common/submission.rs | 32 ++++++++++++++-------------- opsqueue/src/producer/client.rs | 2 +- opsqueue/src/prometheus.rs | 2 +- 5 files changed, 21 insertions(+), 21 deletions(-) diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index 33f42189..d74f82d7 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -160,7 +160,7 @@ impl ProducerClient { }) } - /// Unpause a paused submission, making it available to consumers again. + /// Unpause a paused submission, making it available to consumers. /// /// Will return an error if the submission is not currently paused. /// diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index 74de9ed4..3036bc77 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -685,7 +685,7 @@ pub mod db { " INSERT INTO chunks_failed (submission_id, chunk_index, input_content, failure, skipped, failed_at) - SELECT submission_id, chunk_index, input_content, '', 1, julianday($1) FROM chunks_paused WHERE submission_id = $2; + SELECT submission_id, chunk_index, input_content, '', TRUE, julianday($1) FROM chunks_paused WHERE submission_id = $2; DELETE FROM chunks_paused WHERE submission_id = $3; ", @@ -717,7 +717,7 @@ pub mod db { INSERT INTO chunks_failed (submission_id, chunk_index, input_content, failure, skipped, failed_at) - SELECT submission_id, chunk_index, input_content, '', 1, julianday($1) FROM chunks WHERE chunks.submission_id = $2; + SELECT submission_id, chunk_index, input_content, '', TRUE, julianday($1) FROM chunks WHERE chunks.submission_id = $2; DELETE FROM chunks WHERE chunks.submission_id = $3; diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 6333ef8b..23bca9c0 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -537,21 +537,21 @@ pub mod db { id: SubmissionId, mut conn: impl WriterConnection, ) -> Result<(), E> { - let row = query!( + let res = query!( " INSERT INTO submissions (id, chunks_total, chunks_done, prefix, metadata, otel_trace_carrier, chunk_size) SELECT id, chunks_total, chunks_done, prefix, metadata, otel_trace_carrier, chunk_size FROM submissions_paused WHERE id = $1; - DELETE FROM submissions_paused WHERE id = $2 RETURNING *; + DELETE FROM submissions_paused WHERE id = $2; ", id, id, ) - .fetch_optional(conn.get_inner()) + .execute(conn.get_inner()) .await?; - if row.is_none() { + if res.rows_affected() == 0 { Err(E::R(SubmissionNotFound(id))) } else { counter!(crate::prometheus::SUBMISSIONS_UNPAUSED_COUNTER).increment(1); @@ -617,7 +617,7 @@ pub mod db { match maybe_complete_submission(submission_id, conn).await { // Forward our database errors to the caller. Err(E::L(e)) => return Err(e), - // If the submission ID can't be found, that's too bad, but it's not our problem anymore i guess. + // If the submission ID can't be found, that's too bad, but it's not our problem anymore I guess. Err(E::R(_)) => { tracing::warn!(%submission_id, "Presumed zero-length submission not found"); } @@ -1075,7 +1075,7 @@ pub mod db { /// # Errors /// /// Returns an error if cancellation or chunk skipping fails. - pub async fn cancel_submission_notx( + async fn cancel_submission_notx( id: SubmissionId, mut conn: impl WriterConnection, ) -> Result<(), E> { @@ -1091,7 +1091,7 @@ pub mod db { /// Returns [`DatabaseError`] if any SQL query fails. /// /// Returns [`SubmissionNotFound`] if the submission is not found in `submissions_paused`. - pub async fn cancel_paused_submission_notx( + async fn cancel_paused_submission_notx( id: SubmissionId, mut conn: impl WriterConnection, ) -> Result<(), E> { @@ -1107,21 +1107,21 @@ pub mod db { ) -> Result<(), E> { let now = chrono::prelude::Utc::now(); - let submission_opt = query!( + let res = 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; - DELETE FROM submissions WHERE id = $3 RETURNING *; + DELETE FROM submissions WHERE id = $3; ", now, id, id, ) - .fetch_optional(conn.get_inner()) + .execute(conn.get_inner()) .await?; - if submission_opt.is_none() { + if res.rows_affected() == 0 { Err(E::R(SubmissionNotFound(id))) } else { counter!(crate::prometheus::SUBMISSIONS_CANCELLED_COUNTER).increment(1); @@ -1139,21 +1139,21 @@ pub mod db { ) -> Result<(), E> { let now = chrono::prelude::Utc::now(); - let submission_opt = query!( + let res = 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_paused WHERE id = $2; - DELETE FROM submissions_paused WHERE id = $3 RETURNING *; + DELETE FROM submissions_paused WHERE id = $3; ", now, id, id, ) - .fetch_optional(conn.get_inner()) + .execute(conn.get_inner()) .await?; - if submission_opt.is_none() { + if res.rows_affected() == 0 { Err(E::R(SubmissionNotFound(id))) } else { counter!(crate::prometheus::SUBMISSIONS_CANCELLED_COUNTER).increment(1); @@ -1257,7 +1257,7 @@ pub mod db { /// # Errors /// /// Returns an error if submission/chunk failure transitions cannot be persisted. - pub async fn fail_submission_notx( + pub(crate) async fn fail_submission_notx( id: SubmissionId, failed_chunk_index: ChunkIndex, failure: String, diff --git a/opsqueue/src/producer/client.rs b/opsqueue/src/producer/client.rs index 09e4ab5b..17a94691 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -203,7 +203,7 @@ impl Client { .await } - /// Unpause a paused submission, making it available to consumers again. + /// Unpause a paused submission, making it available to consumers. /// /// Returns an error if the submission is not currently paused. /// diff --git a/opsqueue/src/prometheus.rs b/opsqueue/src/prometheus.rs index 29f63baf..99e76439 100644 --- a/opsqueue/src/prometheus.rs +++ b/opsqueue/src/prometheus.rs @@ -77,7 +77,7 @@ pub fn describe_metrics() { describe_counter!( SUBMISSIONS_UNPAUSED_COUNTER, Unit::Count, - "Number of submissions unpaused (resumed)" + "Number of submissions unpaused" ); describe_histogram!( SUBMISSIONS_DURATION_COMPLETE_HISTOGRAM, From 1384dd122ab18f8c4796ca751b071ad821dcc9fb Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 5 Aug 2026 13:06:38 +0200 Subject: [PATCH 06/11] Pull out queries and EXPLAIN the now shared query --- opsqueue/src/common/submission.rs | 487 ++++++++++++++++-------------- 1 file changed, 259 insertions(+), 228 deletions(-) diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 23bca9c0..bf56ba24 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -310,7 +310,7 @@ pub mod db { }; use axum_prometheus::metrics::{counter, histogram}; use chunk::ChunkSize; - use sqlx::{QueryBuilder, Sqlite, query, query_scalar}; + use sqlx::{Database, QueryBuilder, Sqlite, query, query_as, query_scalar}; use super::{ Chunk, ChunkCount, ChunkIndex, DateTime, Duration, E, Metadata, Submission, @@ -806,7 +806,122 @@ pub mod db { // NOTE: The order is important here; a concurrent writer could move a submission // from InProgress to Completed/Failed in-between the queries. - let submission_row = query!( + let submission_row = submission_status_in_progress_query(id) + .fetch_optional(conn.get_inner()) + .await?; + if let Some(row) = submission_row { + let submission = Submission { + id: row.id, + prefix: row.prefix, + chunks_total: row.chunks_total, + chunks_done: row.chunks_done, + chunk_size: row.chunk_size, + metadata: row.metadata, + strategic_metadata: row.strategic_metadata.0, + otel_trace_carrier: row.otel_trace_carrier, + }; + return Ok(Some(SubmissionStatus::InProgress(submission))); + } + + let completed_row_opt = submission_status_completed_query(id) + .fetch_optional(conn.get_inner()) + .await?; + if let Some(row) = completed_row_opt { + let submission_completed = SubmissionCompleted { + id: row.id, + prefix: row.prefix, + chunks_total: row.chunks_total, + chunk_size: row.chunk_size, + metadata: row.metadata, + strategic_metadata: row.strategic_metadata.0, + completed_at: row.completed_at, + otel_trace_carrier: row.otel_trace_carrier, + }; + return Ok(Some(SubmissionStatus::Completed(submission_completed))); + } + + let failed_row_opt = submission_status_failed_query(id) + .fetch_optional(conn.get_inner()) + .await?; + if let Some(row) = failed_row_opt { + let failed_submission = SubmissionFailed { + id: row.id, + prefix: row.prefix, + chunks_total: row.chunks_total, + chunks_done: row.chunks_done, + chunk_size: row.chunk_size, + metadata: row.metadata, + strategic_metadata: row.strategic_metadata.0, + failed_at: row.failed_at, + failed_chunk_id: row.failed_chunk_id, + otel_trace_carrier: row.otel_trace_carrier, + }; + let failed_chunk_id = (row.id, row.failed_chunk_id).into(); + let failed_chunk = super::chunk::db::get_chunk_failed(failed_chunk_id, conn).await?; + return Ok(Some(SubmissionStatus::Failed( + failed_submission, + failed_chunk, + ))); + } + + let cancelled_row_opt = submission_status_cancelled_query(id) + .fetch_optional(conn.get_inner()) + .await?; + if let Some(row) = cancelled_row_opt { + let cancelled_submission = SubmissionCancelled { + id: row.id, + prefix: row.prefix, + chunks_total: row.chunks_total, + chunks_done: row.chunks_done, + metadata: row.metadata, + strategic_metadata: row.strategic_metadata.0, + cancelled_at: row.cancelled_at, + }; + return Ok(Some(SubmissionStatus::Cancelled(cancelled_submission))); + } + + let paused_row_opt = submission_status_paused_query(id) + .fetch_optional(conn.get_inner()) + .await?; + if let Some(row) = paused_row_opt { + let paused_submission = SubmissionPaused { + id: row.id, + prefix: row.prefix, + chunks_total: row.chunks_total, + chunks_done: row.chunks_done, + chunk_size: row.chunk_size, + metadata: row.metadata, + strategic_metadata: row.strategic_metadata.0, + otel_trace_carrier: row.otel_trace_carrier, + }; + return Ok(Some(SubmissionStatus::Paused(paused_submission))); + } + + Ok(None) + } + + pub(crate) struct SubmissionStatusInProgressRow { + id: SubmissionId, + prefix: Option, + chunks_total: ChunkCount, + chunks_done: ChunkCount, + chunk_size: ChunkSize, + metadata: Option, + strategic_metadata: sqlx::types::Json, + otel_trace_carrier: String, + } + + #[allow(clippy::type_complexity)] + pub(crate) fn submission_status_in_progress_query( + id: SubmissionId, + ) -> query::Map< + 'static, + Sqlite, + fn(::Row) -> Result, + ::Arguments, + > { + query_as!( + SubmissionStatusInProgressRow, r#" SELECT id AS "id: SubmissionId" @@ -824,23 +939,30 @@ pub mod db { "#, id ) - .fetch_optional(conn.get_inner()) - .await?; - if let Some(row) = submission_row { - let submission = Submission { - id: row.id, - prefix: row.prefix, - chunks_total: row.chunks_total, - chunks_done: row.chunks_done, - chunk_size: row.chunk_size, - metadata: row.metadata, - strategic_metadata: row.strategic_metadata.0, - otel_trace_carrier: row.otel_trace_carrier, - }; - return Ok(Some(SubmissionStatus::InProgress(submission))); - } + } - let completed_row_opt = query!( + pub(crate) struct SubmissionStatusCompletedRow { + id: SubmissionId, + prefix: Option, + chunks_total: ChunkCount, + chunk_size: ChunkSize, + metadata: Option, + strategic_metadata: sqlx::types::Json, + completed_at: DateTime, + otel_trace_carrier: String, + } + + #[allow(clippy::type_complexity)] + pub(crate) fn submission_status_completed_query( + id: SubmissionId, + ) -> query::Map< + 'static, + Sqlite, + fn(::Row) -> Result, + ::Arguments, + > { + query_as!( + SubmissionStatusCompletedRow, r#" SELECT id AS "id: SubmissionId" @@ -858,23 +980,32 @@ pub mod db { "#, id ) - .fetch_optional(conn.get_inner()) - .await?; - if let Some(row) = completed_row_opt { - let submission_completed = SubmissionCompleted { - id: row.id, - prefix: row.prefix, - chunks_total: row.chunks_total, - chunk_size: row.chunk_size, - metadata: row.metadata, - strategic_metadata: row.strategic_metadata.0, - completed_at: row.completed_at, - otel_trace_carrier: row.otel_trace_carrier, - }; - return Ok(Some(SubmissionStatus::Completed(submission_completed))); - } + } - let failed_row_opt = query!( + pub(crate) struct SubmissionStatusFailedRow { + id: SubmissionId, + prefix: Option, + chunks_total: ChunkCount, + chunks_done: Option, + chunk_size: ChunkSize, + metadata: Option, + strategic_metadata: sqlx::types::Json, + failed_at: DateTime, + failed_chunk_id: ChunkIndex, + otel_trace_carrier: String, + } + + #[allow(clippy::type_complexity)] + pub(crate) fn submission_status_failed_query( + id: SubmissionId, + ) -> query::Map< + 'static, + Sqlite, + fn(::Row) -> Result, + ::Arguments, + > { + query_as!( + SubmissionStatusFailedRow, r#" SELECT id AS "id: SubmissionId" @@ -894,30 +1025,29 @@ pub mod db { "#, id ) - .fetch_optional(conn.get_inner()) - .await?; - if let Some(row) = failed_row_opt { - let failed_submission = SubmissionFailed { - id: row.id, - prefix: row.prefix, - chunks_total: row.chunks_total, - chunks_done: row.chunks_done, - chunk_size: row.chunk_size, - metadata: row.metadata, - strategic_metadata: row.strategic_metadata.0, - failed_at: row.failed_at, - failed_chunk_id: row.failed_chunk_id, - otel_trace_carrier: row.otel_trace_carrier, - }; - let failed_chunk_id = (row.id, row.failed_chunk_id).into(); - let failed_chunk = super::chunk::db::get_chunk_failed(failed_chunk_id, conn).await?; - return Ok(Some(SubmissionStatus::Failed( - failed_submission, - failed_chunk, - ))); - } + } + + pub(crate) struct SubmissionStatusCancelledRow { + id: SubmissionId, + prefix: Option, + chunks_total: ChunkCount, + chunks_done: ChunkCount, + metadata: Option, + strategic_metadata: sqlx::types::Json, + cancelled_at: DateTime, + } - let cancelled_row_opt = query!( + #[allow(clippy::type_complexity)] + pub(crate) fn submission_status_cancelled_query( + id: SubmissionId, + ) -> query::Map< + 'static, + Sqlite, + fn(::Row) -> Result, + ::Arguments, + > { + query_as!( + SubmissionStatusCancelledRow, r#" SELECT id AS "id: SubmissionId" @@ -934,22 +1064,30 @@ pub mod db { "#, id ) - .fetch_optional(conn.get_inner()) - .await?; - if let Some(row) = cancelled_row_opt { - let cancelled_submission = SubmissionCancelled { - id: row.id, - prefix: row.prefix, - chunks_total: row.chunks_total, - chunks_done: row.chunks_done, - metadata: row.metadata, - strategic_metadata: row.strategic_metadata.0, - cancelled_at: row.cancelled_at, - }; - return Ok(Some(SubmissionStatus::Cancelled(cancelled_submission))); - } + } - let paused_row_opt = query!( + pub(crate) struct SubmissionStatusPausedRow { + id: SubmissionId, + prefix: Option, + chunks_total: ChunkCount, + chunks_done: ChunkCount, + chunk_size: ChunkSize, + metadata: Option, + strategic_metadata: sqlx::types::Json, + otel_trace_carrier: String, + } + + #[allow(clippy::type_complexity)] + pub(crate) fn submission_status_paused_query( + id: SubmissionId, + ) -> query::Map< + 'static, + Sqlite, + fn(::Row) -> Result, + ::Arguments, + > { + query_as!( + SubmissionStatusPausedRow, r#" SELECT id AS "id: SubmissionId" @@ -967,23 +1105,6 @@ pub mod db { "#, id ) - .fetch_optional(conn.get_inner()) - .await?; - if let Some(row) = paused_row_opt { - let paused_submission = SubmissionPaused { - id: row.id, - prefix: row.prefix, - chunks_total: row.chunks_total, - chunks_done: row.chunks_done, - chunk_size: row.chunk_size, - metadata: row.metadata, - strategic_metadata: row.strategic_metadata.0, - otel_trace_carrier: row.otel_trace_carrier, - }; - return Ok(Some(SubmissionStatus::Paused(paused_submission))); - } - - Ok(None) } #[tracing::instrument(skip(conn))] @@ -1471,7 +1592,7 @@ pub mod test { use chunk::ChunkSize; use itertools::Itertools; use sqlformat::{FormatOptions, QueryParams, format}; - use sqlx::{Row, SqliteConnection}; + use sqlx::{Execute, Row, Sqlite}; use std::assert_matches; use crate::common::StrategicMetadataMap; @@ -1481,40 +1602,36 @@ pub mod test { use super::db::*; use super::*; - async fn explain_query_plan(query: &str, conn: &mut SqliteConnection) -> String { - sqlx::raw_sql(sqlx::AssertSqlSafe(format!("EXPLAIN QUERY PLAN {query}"))) - .fetch_all(&mut *conn) - .await - .unwrap_or_else(|_| panic!("Invalid query: \n{query}\n")) - .into_iter() - .map(|row| { - let id = row.get::("id"); - let parent = row.get::("parent"); - let detail = row.get::("detail"); - format!("{id}, {parent}, {detail}") - }) - .join("\n") - } - - fn assert_non_regressing_query_plan(query: &str, explained: &str) { - assert!( - !explained.contains("MATERIALIZED"), - "Query should contain no materialization, but it did.\n\nQuery: {query}\n\nPlan:\n\n{explained}" - ); - assert!( - !explained.contains("B-TREE"), - "Query should contain no temporary B-tree construction, but it did.\n\nQuery: {query}\n\nPlan:\n\n{explained}" - ); + async fn explain_query_plan<'q, Q: Execute<'q, Sqlite>>( + query: Q, + db: sqlx::SqlitePool, + ) -> String { + let mut conn = db.acquire().await.unwrap(); + let query = query.sql(); + let query_string = query.as_str(); + sqlx::raw_sql(sqlx::AssertSqlSafe(format!( + "EXPLAIN QUERY PLAN {query_string}" + ))) + .fetch_all(&mut *conn) + .await + .unwrap_or_else(|_| panic!("Invalid query: \n{query_string}\n")) + .into_iter() + .map(|row| { + let id = row.get::("id"); + let parent = row.get::("parent"); + let detail = row.get::("detail"); + format!("{id}, {parent}, {detail}") + }) + .join("\n") } #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_lookup_by_strategic_metadata(db: sqlx::SqlitePool) { - let mut conn = db.acquire().await.unwrap(); let strategic_metadata: StrategicMetadataMap = [("company_id".to_string(), 1), ("project_id".to_string(), 2)] .into_iter() .collect(); - let qb = lookup_ids_by_strategic_metadata_query(&strategic_metadata, 100_000); + let mut qb = lookup_ids_by_strategic_metadata_query(&strategic_metadata, 100_000); let options = FormatOptions::default(); let formatted_query = format(qb.sql().as_str(), &QueryParams::None, &options); insta::assert_snapshot!(formatted_query, @" @@ -1533,8 +1650,7 @@ pub mod test { LIMIT ? "); - let explained = explain_query_plan(&formatted_query, &mut conn).await; - assert_non_regressing_query_plan(&formatted_query, &explained); + let explained = explain_query_plan(qb.build_query_scalar::(), db).await; insta::assert_snapshot!(explained, @" 8, 0, SEARCH s0 USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=?) 16, 0, SEARCH submissions USING COVERING INDEX sqlite_autoindex_submissions_1 (id=?) @@ -1544,114 +1660,46 @@ pub mod test { #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_submission_status_in_progress(db: sqlx::SqlitePool) { - let mut conn = db.acquire().await.unwrap(); - let query = r" - SELECT - id - , prefix - , chunks_total - , chunks_done - , chunk_size - , metadata - , ( SELECT json_group_object(metadata_key, metadata_value) - FROM submissions_metadata - WHERE submission_id = submissions.id - ) AS strategic_metadata - , otel_trace_carrier - FROM submissions WHERE id = 1 - "; - - let explained = explain_query_plan(query, &mut conn).await; - assert_non_regressing_query_plan(query, &explained); - insta::assert_snapshot!(explained, @r" + let query = submission_status_in_progress_query(SubmissionId::new()); + let explained = explain_query_plan(query, db).await; + insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions USING INDEX sqlite_autoindex_submissions_1 (id=?) - 15, 0, CORRELATED SCALAR SUBQUERY 1 - 20, 15, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) + 17, 0, CORRELATED SCALAR SUBQUERY 1 + 22, 17, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) "); } #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_submission_status_completed(db: sqlx::SqlitePool) { - let mut conn = db.acquire().await.unwrap(); - let query = r" - SELECT - id - , prefix - , chunks_total - , chunk_size - , metadata - , ( SELECT json_group_object(metadata_key, metadata_value) - FROM submissions_metadata - WHERE submission_id = submissions_completed.id - ) AS strategic_metadata - , completed_at - , otel_trace_carrier - FROM submissions_completed WHERE id = 1 - "; + let query = submission_status_completed_query(SubmissionId::new()); - let explained = explain_query_plan(query, &mut conn).await; - assert_non_regressing_query_plan(query, &explained); - insta::assert_snapshot!(explained, @r" + let explained = explain_query_plan(query, db).await; + insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions_completed USING INDEX sqlite_autoindex_submissions_completed_1 (id=?) - 14, 0, CORRELATED SCALAR SUBQUERY 1 - 19, 14, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) + 16, 0, CORRELATED SCALAR SUBQUERY 1 + 21, 16, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) "); } #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_submission_status_failed(db: sqlx::SqlitePool) { - let mut conn = db.acquire().await.unwrap(); - let query = r" - SELECT - id - , prefix - , chunks_total - , chunks_done - , chunk_size - , metadata - , ( SELECT json_group_object(metadata_key, metadata_value) - FROM submissions_metadata - WHERE submission_id = submissions_failed.id - ) AS strategic_metadata - , failed_at - , failed_chunk_id - , otel_trace_carrier - FROM submissions_failed WHERE id = 1 - "; - - let explained = explain_query_plan(query, &mut conn).await; - assert_non_regressing_query_plan(query, &explained); - insta::assert_snapshot!(explained, @r" + let query = submission_status_failed_query(SubmissionId::new()); + let explained = explain_query_plan(query, db).await; + insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions_failed USING INDEX sqlite_autoindex_submissions_failed_1 (id=?) - 15, 0, CORRELATED SCALAR SUBQUERY 1 - 20, 15, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) + 17, 0, CORRELATED SCALAR SUBQUERY 1 + 22, 17, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) "); } #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_submission_status_cancelled(db: sqlx::SqlitePool) { - let mut conn = db.acquire().await.unwrap(); - let query = r" - SELECT - id - , prefix - , chunks_total - , chunks_done - , metadata - , ( SELECT json_group_object(metadata_key, metadata_value) - FROM submissions_metadata - WHERE submission_id = submissions_cancelled.id - ) AS strategic_metadata - , cancelled_at - FROM submissions_cancelled WHERE id = 1 - "; - - let explained = explain_query_plan(query, &mut conn).await; - assert_non_regressing_query_plan(query, &explained); - insta::assert_snapshot!(explained, @r" + let query = submission_status_cancelled_query(SubmissionId::new()); + let explained = explain_query_plan(query, db).await; + insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions_cancelled USING INDEX sqlite_autoindex_submissions_cancelled_1 (id=?) - 14, 0, CORRELATED SCALAR SUBQUERY 1 - 19, 14, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) + 16, 0, CORRELATED SCALAR SUBQUERY 1 + 21, 16, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) "); } @@ -2041,29 +2089,12 @@ pub mod test { #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_submission_status_paused(db: sqlx::SqlitePool) { - let mut conn = db.acquire().await.unwrap(); - let query = r" - SELECT - id - , prefix - , chunks_total - , chunks_done - , chunk_size - , metadata - , ( SELECT json_group_object(metadata_key, metadata_value) - FROM submissions_metadata - WHERE submission_id = submissions_paused.id - ) AS strategic_metadata - , otel_trace_carrier - FROM submissions_paused WHERE id = 1 - "; - - let explained = explain_query_plan(query, &mut conn).await; - assert_non_regressing_query_plan(query, &explained); - insta::assert_snapshot!(explained, @r" + let query = submission_status_paused_query(SubmissionId::new()); + let explained = explain_query_plan(query, db).await; + insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions_paused USING INDEX sqlite_autoindex_submissions_paused_1 (id=?) - 15, 0, CORRELATED SCALAR SUBQUERY 1 - 20, 15, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) + 17, 0, CORRELATED SCALAR SUBQUERY 1 + 22, 17, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) "); } From be4235c4ea8689eae3116763572ea9272d348e05 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 5 Aug 2026 15:26:08 +0200 Subject: [PATCH 07/11] Address Copilot's suppressed comments --- libs/opsqueue_python/src/errors.rs | 9 +- libs/opsqueue_python/src/producer.rs | 27 +++-- libs/opsqueue_python/tests/test_roundtrip.py | 7 +- opsqueue/src/common/chunk.rs | 120 +++++++++++++++---- opsqueue/src/common/submission.rs | 63 +++++++--- 5 files changed, 169 insertions(+), 57 deletions(-) diff --git a/libs/opsqueue_python/src/errors.rs b/libs/opsqueue_python/src/errors.rs index b0de8e7e..043ac11f 100644 --- a/libs/opsqueue_python/src/errors.rs +++ b/libs/opsqueue_python/src/errors.rs @@ -1,12 +1,13 @@ /// NOTE: We define the potentially raisable errors/exceptions in Python /// so we have nice IDE support for docs-on-hover and for 'go to definition'. use std::error::Error; +use std::time::TryFromFloatSecsError; use opsqueue::common::errors::{ E, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions, UnexpectedOpsqueueConsumerServerResponse, }; -use pyo3::exceptions::{PyBaseException, PyTimeoutError}; +use pyo3::exceptions::{PyBaseException, PyTimeoutError, PyValueError}; use pyo3::{Bound, PyErr, Python, import_exception}; use crate::common; @@ -188,6 +189,12 @@ impl From> for PyErr { } } +impl From> for PyErr { + fn from(value: CError) -> Self { + PyValueError::new_err(value.0.to_string()) + } +} + impl From for CError> { fn from(value: PyErr) -> Self { CError(E::L(FatalPythonException(value))) diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index d74f82d7..ed2ecae9 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -1,11 +1,11 @@ -use std::{future::IntoFuture, sync::Arc, time::Duration}; - use pyo3::{ create_exception, exceptions::{PyException, PyStopAsyncIteration}, prelude::*, types::PyIterator, }; +use std::time::TryFromFloatSecsError; +use std::{future::IntoFuture, sync::Arc, time::Duration}; use futures::{StreamExt, TryStreamExt, stream::BoxStream}; use opsqueue::{ @@ -436,6 +436,7 @@ impl ProducerClient { PyChunksIter, E![ FatalPythonException, + TryFromFloatSecsError, Elapsed, errors::SubmissionFailed, InternalProducerClientError @@ -445,18 +446,22 @@ impl ProducerClient { self.block_unless_interrupted(async move { let fut = self.stream_completed_submission_chunks(submission_id); match timeout { - Some(duration) => tokio::time::timeout(Duration::from_secs_f64(duration), fut) - .await - .map_err(|err| CError(R(L(err)))) - .and_then(|err| { - err.map_err(|err| match err.0 { - L(err) => CError(L(err)), - R(err) => CError(R(R(err))), + Some(duration) => { + let duration = Duration::try_from_secs_f64(duration) + .map_err(|err| CError(R(L(err))))?; + tokio::time::timeout(duration, fut) + .await + .map_err(|err| CError(R(R(L(err))))) + .and_then(|err| { + err.map_err(|err| match err.0 { + L(err) => CError(L(err)), + R(err) => CError(R(R(R(err)))), + }) }) - }), + } None => fut.await.map_err(|err| match err.0 { L(err) => CError(L(err)), - R(err) => CError(R(R(err))), + R(err) => CError(R(R(R(err)))), }), } }) diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index 3d6dddb0..44a68edf 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -757,7 +757,7 @@ def process_op(x: int) -> int: def test_unpause_and_complete(opsqueue: OpsqueueProcess) -> None: - """Unpausing a paused submission makes it available to consumers again, + """Unpausing a paused submission makes it available to consumers, and it can be completed normally afterwards.""" url = "file:///tmp/opsqueue/test_unpause_and_complete" producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) @@ -780,7 +780,10 @@ def run_consumer() -> None: consumer_client.run_each_op(lambda x: x) with background_process(run_consumer): - producer_client.blocking_stream_completed_submission(submission_id) + producer_client.blocking_stream_completed_submission( + submission_id, + timeout=SUBMISSION_COMPLETED_TIMEOUT, + ) assert isinstance( producer_client.get_submission_status(submission_id), SubmissionStatus.Completed, diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index 3036bc77..be7be9ea 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -313,19 +313,33 @@ pub mod db { output_content: Option>, mut conn: impl WriterConnection, ) -> Result<(), E> { - conn.transaction(move |mut tx| { - Box::pin(async move { - complete_chunk_raw(chunk_id, output_content, &mut tx).await?; - crate::common::submission::db::maybe_complete_submission( - chunk_id.submission_id, - &mut tx, - ) - .await + let chunks_moved = conn + .transaction(move |mut tx| { + Box::pin(async move { + let chunks_moved = + complete_chunk_raw(chunk_id, output_content, &mut tx).await?; + if chunks_moved { + crate::common::submission::db::maybe_complete_submission( + chunk_id.submission_id, + &mut tx, + ) + .await?; + } else { + tracing::warn!( + "Could not complete chunk {:?} because it was either: \ + completed, failed, or cancelled before. Ignoring.", + chunk_id + ); + } + + Result::>::Ok(chunks_moved) + }) }) - }) - .await?; + .await?; - counter!(crate::prometheus::CHUNKS_COMPLETED_COUNTER).increment(1); + if chunks_moved { + counter!(crate::prometheus::CHUNKS_COMPLETED_COUNTER).increment(1); + } Ok(()) } @@ -339,7 +353,7 @@ pub mod db { chunk_id: ChunkId, output_content: Option>, mut tx: impl WriterConnection, - ) -> sqlx::Result<()> { + ) -> sqlx::Result { let now = chrono::prelude::Utc::now(); let chunk_moved = query!( " @@ -348,8 +362,7 @@ pub mod db { SELECT submission_id, chunk_index, $1, julianday($2) FROM chunks WHERE chunks.submission_id = $3 AND chunks.chunk_index = $4; - DELETE FROM chunks WHERE chunks.submission_id = $5 AND chunks.chunk_index = $6 - RETURNING submission_id, chunk_index; + DELETE FROM chunks WHERE chunks.submission_id = $5 AND chunks.chunk_index = $6; ", output_content, now, @@ -358,9 +371,10 @@ pub mod db { chunk_id.submission_id, chunk_id.chunk_index, ) - .fetch_optional(tx.get_inner()) + .execute(tx.get_inner()) .await? - .is_some(); + .rows_affected() + > 0; // Defense in depth: Above query could be called twice on the same chunk. For instance, // when the server was restarted and the reservations are forgotten, and the same chunk // was reserved again. @@ -385,15 +399,8 @@ pub mod db { ) .fetch_one(tx.get_inner()) .await?; - } else { - tracing::warn!( - "Could not complete chunk {:?} because it was either: \ - completed, failed, or cancelled before. Ignoring.", - chunk_id - ); } - - Ok(()) + Ok(chunk_moved) } /// Increment retries for a chunk, or move it to failed state. @@ -450,7 +457,6 @@ pub mod db { completed, failed, or cancelled before. Ignoring.", chunk_id ); - Ok::<_, sqlx::Error>(false) } }) @@ -818,9 +824,10 @@ pub mod db { #[cfg(feature = "server-logic")] pub mod test { use crate::common::StrategicMetadataMap; - use crate::common::submission::db::insert_submission_raw; + use crate::common::submission::db::{insert_submission, insert_submission_raw}; use crate::common::submission::{Submission, SubmissionStatus}; use crate::db::{Connection as _, WriterPool}; + use std::assert_matches; use super::db::*; use super::*; @@ -944,6 +951,35 @@ pub mod test { } } + #[sqlx::test(migrator = "crate::MIGRATOR")] + pub async fn test_calling_complete_chunk_twice_for_same_chunk_does_not_error( + db: sqlx::SqlitePool, + ) { + let db = WriterPool::new(db); + let mut conn = db.writer_conn().await.unwrap(); + let (submission, chunks) = Submission::from_vec( + vec![Some("foo".into()), Some("bar".into()), Some("baz".into())], + None, + ChunkSize::default(), + ) + .unwrap(); + + let chunk_id = ChunkId { + submission_id: chunks[0].submission_id, + chunk_index: chunks[0].chunk_index, + }; + + insert_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + let res = complete_chunk(chunk_id, None, &mut conn).await; + assert_matches!(res, Ok(())); + + let res = complete_chunk(chunk_id, None, &mut conn).await; + assert_matches!(res, Ok(())); + } + #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_fail_chunk(db: sqlx::SqlitePool) { let db = WriterPool::new(db); @@ -969,4 +1005,36 @@ pub mod test { assert_eq!(count_chunks_completed(&mut conn).await.unwrap(), 0); assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), 1); } + + #[sqlx::test(migrator = "crate::MIGRATOR")] + pub async fn test_calling_fail_chunk_after_exceeding_retries_does_not_error( + db: sqlx::SqlitePool, + ) { + let db = WriterPool::new(db); + let mut conn = db.writer_conn().await.unwrap(); + let (submission, chunks) = Submission::from_vec( + vec![Some("foo".into()), Some("bar".into()), Some("baz".into())], + None, + ChunkSize::default(), + ) + .unwrap(); + + let chunk_id = ChunkId { + submission_id: chunks[0].submission_id, + chunk_index: chunks[0].chunk_index, + }; + + insert_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + let res = retry_or_fail_chunk(chunk_id, "kapot".into(), &mut conn, 2).await; + assert_matches!(res, Ok(false)); + + let res = retry_or_fail_chunk(chunk_id, "kapot".into(), &mut conn, 2).await; + assert_matches!(res, Ok(true)); + + let res = retry_or_fail_chunk(chunk_id, "kapot".into(), &mut conn, 2).await; + assert_matches!(res, Ok(false)); + } } diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index bf56ba24..db69662e 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -526,6 +526,9 @@ pub mod db { Box::pin(async move { unpause_submission_raw(id, &mut tx).await?; super::chunk::db::restore_paused_chunks(id, &mut tx).await?; + // NOTE: We need to check whether the submission is completed, because it might + // be the case that we are unpausing a 0-chunk submission. + maybe_complete_submission(id, &mut tx).await?; Ok(()) }) }) @@ -805,6 +808,24 @@ pub mod db { ) -> Result, DatabaseError> { // NOTE: The order is important here; a concurrent writer could move a submission // from InProgress to Completed/Failed in-between the queries. + // TODO: Rewrite the queries here into a single query using `UNION ALL`. + + let paused_row_opt = submission_status_paused_query(id) + .fetch_optional(conn.get_inner()) + .await?; + if let Some(row) = paused_row_opt { + let paused_submission = SubmissionPaused { + id: row.id, + prefix: row.prefix, + chunks_total: row.chunks_total, + chunks_done: row.chunks_done, + chunk_size: row.chunk_size, + metadata: row.metadata, + strategic_metadata: row.strategic_metadata.0, + otel_trace_carrier: row.otel_trace_carrier, + }; + return Ok(Some(SubmissionStatus::Paused(paused_submission))); + } let submission_row = submission_status_in_progress_query(id) .fetch_optional(conn.get_inner()) @@ -880,23 +901,6 @@ pub mod db { return Ok(Some(SubmissionStatus::Cancelled(cancelled_submission))); } - let paused_row_opt = submission_status_paused_query(id) - .fetch_optional(conn.get_inner()) - .await?; - if let Some(row) = paused_row_opt { - let paused_submission = SubmissionPaused { - id: row.id, - prefix: row.prefix, - chunks_total: row.chunks_total, - chunks_done: row.chunks_done, - chunk_size: row.chunk_size, - metadata: row.metadata, - strategic_metadata: row.strategic_metadata.0, - otel_trace_carrier: row.otel_trace_carrier, - }; - return Ok(Some(SubmissionStatus::Paused(paused_submission))); - } - Ok(None) } @@ -2124,6 +2128,31 @@ pub mod test { assert_eq!(count_chunks_paused(&mut conn).await.unwrap(), 0); } + #[sqlx::test(migrator = "crate::MIGRATOR")] + pub async fn test_unpausing_a_zero_chunk_submission_completes_it(db: sqlx::SqlitePool) { + let db = WriterPool::new(db); + let mut conn = db.writer_conn().await.unwrap(); + let (submission, chunks) = + Submission::from_vec(vec![], None, ChunkSize::default()).unwrap(); + insert_paused_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + assert_eq!(count_submissions(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_completed(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_paused(&mut conn).await.unwrap(), 1); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 0); + assert_eq!(count_chunks_paused(&mut conn).await.unwrap(), 0); + + unpause_submission(submission.id, &mut conn).await.unwrap(); + + assert_eq!(count_submissions(&mut conn).await.unwrap(), 0); + assert_eq!(count_submissions_completed(&mut conn).await.unwrap(), 1); + assert_eq!(count_submissions_paused(&mut conn).await.unwrap(), 0); + assert_eq!(count_chunks(&mut conn).await.unwrap(), 0); + assert_eq!(count_chunks_paused(&mut conn).await.unwrap(), 0); + } + #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_cancel_paused_submission(db: sqlx::SqlitePool) { let db = WriterPool::new(db); From f6a6f885acffc8ae9f15953cd6d8406c6cf9adcb Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Tue, 15 Sep 2026 15:45:42 +0200 Subject: [PATCH 08/11] Address feedback --- Cargo.lock | 4 +- Cargo.toml | 4 +- libs/opsqueue_python/src/producer.rs | 34 ++--- opsqueue/benches/chunks_select.rs | 1 + opsqueue/opsqueue_example_database_schema.db | Bin 102400 -> 110592 bytes opsqueue/src/common/chunk.rs | 21 ++- opsqueue/src/common/submission.rs | 133 +++++++++++-------- opsqueue/src/consumer/dispatcher/mod.rs | 1 + opsqueue/src/producer/client.rs | 1 - 9 files changed, 115 insertions(+), 84 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 70cb3ed6..90be89c1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1995,7 +1995,7 @@ dependencies = [ [[package]] name = "opsqueue" -version = "0.40.1" +version = "0.41.0" dependencies = [ "anyhow", "arc-swap", @@ -2048,7 +2048,7 @@ dependencies = [ [[package]] name = "opsqueue_python" -version = "0.40.1" +version = "0.41.0" dependencies = [ "anyhow", "chrono", diff --git a/Cargo.toml b/Cargo.toml index 69fea44e..e8cca97d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,7 +14,7 @@ hakari-package = "workspace-hack" [workspace.package] edition = "2024" -version = "0.40.1" +version = "0.41.0" [workspace.dependencies] anyhow = { version = "1.0.102", default-features = false } @@ -28,7 +28,7 @@ ciborium = { version = "0.2.2", default-features = false, features = ["std"] } clap = { version = "4.5.60", default-features = false, features = ["derive", "std"] } criterion = { version = "0.8", default-features = false, features = ["async_tokio"] } crossbeam-skiplist = { version = "0.1.3", default-features = false, features = ["std"] } -dashmap = { version = "6.1.0", default-features = false, features = ["serde"]} +dashmap = { version = "6.1.0", default-features = false, features = ["serde"] } either = { version = "1.13.0", default-features = false } futures = { version = "0.3.32", default-features = false } http = { version = "1.4.0", default-features = false } diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index ed2ecae9..e21a5f53 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -168,6 +168,7 @@ impl ProducerClient { /// /// Returns an error if the submission is not found or if an internal client error occurs. #[allow(clippy::result_large_err, clippy::type_complexity)] + #[pyo3(signature = (id))] pub fn unpause_submission( &self, py: Python<'_>, @@ -281,7 +282,9 @@ impl ProducerClient { /// # Errors /// /// Returns an error if submission insertion fails. - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), paused=false))] + #[allow(clippy::too_many_arguments)] + #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), paused=false) + )] pub fn insert_submission_direct( &self, py: Python<'_>, @@ -318,7 +321,8 @@ impl ProducerClient { /// /// Returns an error if chunk upload or submission insertion fails. #[allow(clippy::type_complexity, clippy::too_many_arguments)] - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), paused=false))] + #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), paused=false) + )] pub fn insert_submission_chunks( &self, py: Python<'_>, @@ -343,18 +347,18 @@ impl ProducerClient { let prefix = uuid::Uuid::now_v7().to_string(); tracing::debug!("Uploading submission chunks to object store subfolder {prefix}..."); let chunk_count = self.block_unless_interrupted(async { - let chunk_contents = std::iter::from_fn(move || { - Python::attach(|py| - chunk_contents.bind(py).clone().next() - .map(|item| item.and_then( - |item| item.extract()).map_err(Into::into))) - }); - let stream = futures::stream::iter(chunk_contents); - self.object_store_client - .store_chunks(&prefix, ChunkType::Input, stream) - .await - .map_err(|e| CError(R(L(e)))) - })?; + let chunk_contents = std::iter::from_fn(move || { + Python::attach(|py| + chunk_contents.bind(py).clone().next() + .map(|item| item.and_then( + |item| item.extract()).map_err(Into::into))) + }); + let stream = futures::stream::iter(chunk_contents); + self.object_store_client + .store_chunks(&prefix, ChunkType::Input, stream) + .await + .map_err(|e| CError(R(L(e)))) + })?; let chunk_count = chunk::ChunkIndex::from(chunk_count); tracing::debug!("Finished uploading to object store. {prefix} contains {chunk_count} chunks"); @@ -426,7 +430,7 @@ impl ProducerClient { /// /// Returns an error if polling or output streaming fails. #[allow(clippy::result_large_err, clippy::type_complexity)] - #[pyo3(signature = (submission_id))] + #[pyo3(signature = (submission_id, timeout=None))] pub fn blocking_stream_completed_submission_chunks( &self, py: Python<'_>, diff --git a/opsqueue/benches/chunks_select.rs b/opsqueue/benches/chunks_select.rs index a1ab7b96..950b4d19 100644 --- a/opsqueue/benches/chunks_select.rs +++ b/opsqueue/benches/chunks_select.rs @@ -191,6 +191,7 @@ async fn seed_or_extend( None, metadata, ChunkSize::default(), + false, &mut conn, ) .await diff --git a/opsqueue/opsqueue_example_database_schema.db b/opsqueue/opsqueue_example_database_schema.db index 7a4657366c44a59e8f911b594fae48f9f17d32e5..9fbf60230e2b27611292ff65a782b082ad0a653f 100644 GIT binary patch delta 602 zcmZozz}E19ZGyC*ECU0BG!Vmp{zM&PM%j%COZXM|xKtSU=kvSp-QmmS{mI+MtI4yI z$DjKucMjJtt_CiZ#fAc0F;RT%tkTAuIm|-2Vd5`u6eN}wXXd3F85o)A8d&NYnkpEY zSQ!~u8JMeiww~3L+9+Sqa<^4~_K%ECV_*JhLi3s3&fD(MapRnRCXN=?hG&{4=t(QLApW*0X!WNgs{JF2)eDL1pYI5R&F z<}yqX(a8-=l57yOCLiEkG}(cbZL$N4rh!57EVK zTt%ZE>yxo=5bGwa3HkT5a=848nfgy66hGA$XF$vb0&BbU)C=_5P_&3dxs03kmi*6#dF36L3544eJHw1gLS#e4g zD|%@~G^+~>6~ht-2ePx8C2G2*&Z~wgDyF!kl&Y$kVZ4bnc*5~$)W?xS0A6?oK61a> CHf{3& diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index be7be9ea..2b957688 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -313,12 +313,11 @@ pub mod db { output_content: Option>, mut conn: impl WriterConnection, ) -> Result<(), E> { - let chunks_moved = conn + let chunk_moved = conn .transaction(move |mut tx| { Box::pin(async move { - let chunks_moved = - complete_chunk_raw(chunk_id, output_content, &mut tx).await?; - if chunks_moved { + let chunk_moved = complete_chunk_raw(chunk_id, output_content, &mut tx).await?; + if chunk_moved { crate::common::submission::db::maybe_complete_submission( chunk_id.submission_id, &mut tx, @@ -332,12 +331,12 @@ pub mod db { ); } - Result::>::Ok(chunks_moved) + Result::>::Ok(chunk_moved) }) }) .await?; - if chunks_moved { + if chunk_moved { counter!(crate::prometheus::CHUNKS_COMPLETED_COUNTER).increment(1); } Ok(()) @@ -649,7 +648,7 @@ pub mod db { Ok(()) } - /// Move all chunks of a paused submission from `chunks_paused` back to `chunks`. + /// Move all chunks of a paused submission from `chunks_paused` to `chunks`. /// /// # Errors /// @@ -669,8 +668,8 @@ pub mod db { submission_id, submission_id, ) - .execute(conn.get_inner()) - .await?; + .execute(conn.get_inner()) + .await?; Ok(()) } @@ -699,8 +698,8 @@ pub mod db { submission_id, submission_id, ) - .execute(conn.get_inner()) - .await?; + .execute(conn.get_inner()) + .await?; counter!(crate::prometheus::CHUNKS_SKIPPED_COUNTER).increment(query_res.rows_affected()); Ok(()) diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index db69662e..fa40ea08 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -379,8 +379,8 @@ pub mod db { submission.otel_trace_carrier, submission.chunk_size.0, ) - .execute(conn.get_inner()) - .await?; + .execute(conn.get_inner()) + .await?; Ok(()) } @@ -437,11 +437,22 @@ pub mod db { &submission.strategic_metadata, &mut tx, ) - .await?; + .await?; super::chunk::db::insert_many_chunks(chunks, &mut tx).await?; + + // Empty submissions get special handling: we mark them as completed right away. + // See https://github.com/channable/opsqueue/issues/86 for rationale. + if chunks_total == 0 { + maybe_complete_submission(submission.id, tx).await + .map_err(|e| match e { + // Forward our database errors to the caller. + E::L(e) => e, + E::R(_) => panic!("Failed to find submission that was just inserted in this transaction"), + })?; + } Ok(()) } - .boxed() + .boxed() }) .await; @@ -503,8 +514,8 @@ pub mod db { submission.otel_trace_carrier, submission.chunk_size.0, ) - .execute(conn.get_inner()) - .await?; + .execute(conn.get_inner()) + .await?; Ok(()) } @@ -512,6 +523,9 @@ pub mod db { /// Unpause a paused submission. Atomically moves it back from `submissions_paused` /// to `submissions` and its chunks from `chunks_paused` to `chunks`. /// + /// If the submission has 0 chunks, the submission is marked as completed immediately + /// afterwards. + /// /// # Errors /// /// Returns [`DatabaseError`] if the transaction or any SQL query fails. @@ -614,27 +628,6 @@ pub mod db { insert_paused_submission(submission, chunks, &mut conn).await?; } else { insert_submission(submission, chunks, &mut conn).await?; - // Empty submissions get special handling: we mark them as completed right away. - // See https://github.com/channable/opsqueue/issues/86 for rationale. - if len == 0 { - match maybe_complete_submission(submission_id, conn).await { - // Forward our database errors to the caller. - Err(E::L(e)) => return Err(e), - // If the submission ID can't be found, that's too bad, but it's not our problem anymore I guess. - Err(E::R(_)) => { - tracing::warn!(%submission_id, "Presumed zero-length submission not found"); - } - // If everything went OK, this *could* still indicate a bug in producer code, so let's just log it. - // Our future selves might thank us. - Ok(true) => { - tracing::debug!(%submission_id, "Zero-length submission marked as completed"); - } - // This should never happen. If it does, better log it. - Ok(false) => { - tracing::warn!(%submission_id, "Zero-length submission wasn't zero-length?!"); - } - } - } } Ok(submission_id) } @@ -807,8 +800,11 @@ pub mod db { mut conn: impl Connection, ) -> Result, DatabaseError> { // NOTE: The order is important here; a concurrent writer could move a submission - // from InProgress to Completed/Failed in-between the queries. - // TODO: Rewrite the queries here into a single query using `UNION ALL`. + // from Paused to InProgress/Cancelled in-between the queries. + // from InProgress to Completed/Failed/Cancelled in-between the queries. + // TODO(opsqueue#183): Instead of relying on our attention to detail. Let's rewrite the + // queries here to use a transaction and/or merge them into a single query using + // `UNION ALL`. let paused_row_opt = submission_status_paused_query(id) .fetch_optional(conn.get_inner()) @@ -1130,13 +1126,18 @@ pub mod db { if submission.chunks_done == submission.chunks_total { complete_submission_raw(id, &mut tx).await?; + if submission.chunks_total == 0 { + // There are no chunks, this might indicate a bug in producer code, so + // let's just log it. Our future selves might thank us. + tracing::debug!(%submission.id, "Zero-length submission marked as completed"); + } Ok(true) } else { Ok(false) } }) }) - .await + .await } /// Cancel a submission if it is still cancellable. @@ -1244,8 +1245,8 @@ pub mod db { id, id, ) - .execute(conn.get_inner()) - .await?; + .execute(conn.get_inner()) + .await?; if res.rows_affected() == 0 { Err(E::R(SubmissionNotFound(id))) } else { @@ -1276,8 +1277,8 @@ pub mod db { id, id, ) - .execute(conn.get_inner()) - .await?; + .execute(conn.get_inner()) + .await?; if res.rows_affected() == 0 { Err(E::R(SubmissionNotFound(id))) } else { @@ -1347,8 +1348,8 @@ pub mod db { id, id, ) - .fetch_one(conn.get_inner()) - .await?; + .fetch_one(conn.get_inner()) + .await?; counter!(crate::prometheus::SUBMISSIONS_FAILED_COUNTER).increment(1); histogram!(crate::prometheus::SUBMISSIONS_DURATION_FAIL_HISTOGRAM).record( crate::prometheus::time_delta_as_f64(Utc::now() - id.timestamp()), @@ -1508,8 +1509,8 @@ pub mod db { );", older_than ) - .execute(tx.get_inner()) - .await?; + .execute(tx.get_inner()) + .await?; query!( "DELETE FROM submissions_metadata WHERE submission_id IN ( @@ -1517,8 +1518,8 @@ pub mod db { );", older_than ) - .execute(tx.get_inner()) - .await?; + .execute(tx.get_inner()) + .await?; query!( "DELETE FROM submissions_metadata WHERE submission_id IN ( @@ -1526,41 +1527,41 @@ pub mod db { );", older_than ) - .execute(tx.get_inner()) - .await?; + .execute(tx.get_inner()) + .await?; // Clean up old submissions: let n_submissions_completed = query!( "DELETE FROM submissions_completed WHERE completed_at < julianday($1);", older_than ) - .execute(tx.get_inner()) - .await?.rows_affected(); + .execute(tx.get_inner()) + .await?.rows_affected(); let n_submissions_failed = query!( "DELETE FROM submissions_failed WHERE failed_at < julianday($1);", older_than ) - .execute(tx.get_inner()) - .await?.rows_affected(); + .execute(tx.get_inner()) + .await?.rows_affected(); let n_submissions_cancelled = query!( "DELETE FROM submissions_cancelled WHERE cancelled_at < julianday($1);", older_than ) - .execute(tx.get_inner()) - .await?.rows_affected(); + .execute(tx.get_inner()) + .await?.rows_affected(); let n_chunks_completed = query!( "DELETE FROM chunks_completed WHERE completed_at < julianday($1);", older_than ) - .execute(tx.get_inner()) - .await?.rows_affected(); + .execute(tx.get_inner()) + .await?.rows_affected(); let n_chunks_failed = query!( "DELETE FROM chunks_failed WHERE failed_at < julianday($1);", older_than ) - .execute(tx.get_inner()) - .await?.rows_affected(); + .execute(tx.get_inner()) + .await?.rows_affected(); tracing::info!("Deleted {n_submissions_completed} completed submissions (with {n_chunks_completed} chunks completed)"); tracing::info!("Deleted {n_submissions_failed} failed submissions (with {n_chunks_failed} chunks failed)"); @@ -1568,7 +1569,7 @@ pub mod db { Ok(()) }) }) - .await + .await } pub async fn periodically_cleanup_old(db: &WriterPool, max_age: Duration) { @@ -1629,6 +1630,20 @@ pub mod test { .join("\n") } + fn assert_no_materialization(explained: &str) { + assert!( + !explained.contains("MATERIALIZE"), + "Query should contain no materialization, but it did." + ); + } + + fn assert_no_temporary_b_trees(explained: &str) { + assert!( + !explained.contains("TEMP B-TREE"), + "Query should contain no temporary B-tree construction, but it did." + ); + } + #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_query_plan_lookup_by_strategic_metadata(db: sqlx::SqlitePool) { let strategic_metadata: StrategicMetadataMap = @@ -1655,6 +1670,8 @@ pub mod test { ? "); let explained = explain_query_plan(qb.build_query_scalar::(), db).await; + assert_no_materialization(explained.as_str()); + assert_no_temporary_b_trees(explained.as_str()); insta::assert_snapshot!(explained, @" 8, 0, SEARCH s0 USING COVERING INDEX lookup_submission_by_metadata (metadata_key=? AND metadata_value=?) 16, 0, SEARCH submissions USING COVERING INDEX sqlite_autoindex_submissions_1 (id=?) @@ -1666,6 +1683,8 @@ pub mod test { pub async fn test_query_plan_submission_status_in_progress(db: sqlx::SqlitePool) { let query = submission_status_in_progress_query(SubmissionId::new()); let explained = explain_query_plan(query, db).await; + assert_no_materialization(explained.as_str()); + assert_no_temporary_b_trees(explained.as_str()); insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions USING INDEX sqlite_autoindex_submissions_1 (id=?) 17, 0, CORRELATED SCALAR SUBQUERY 1 @@ -1678,6 +1697,8 @@ pub mod test { let query = submission_status_completed_query(SubmissionId::new()); let explained = explain_query_plan(query, db).await; + assert_no_materialization(explained.as_str()); + assert_no_temporary_b_trees(explained.as_str()); insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions_completed USING INDEX sqlite_autoindex_submissions_completed_1 (id=?) 16, 0, CORRELATED SCALAR SUBQUERY 1 @@ -1689,6 +1710,8 @@ pub mod test { pub async fn test_query_plan_submission_status_failed(db: sqlx::SqlitePool) { let query = submission_status_failed_query(SubmissionId::new()); let explained = explain_query_plan(query, db).await; + assert_no_materialization(explained.as_str()); + assert_no_temporary_b_trees(explained.as_str()); insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions_failed USING INDEX sqlite_autoindex_submissions_failed_1 (id=?) 17, 0, CORRELATED SCALAR SUBQUERY 1 @@ -1700,6 +1723,8 @@ pub mod test { pub async fn test_query_plan_submission_status_cancelled(db: sqlx::SqlitePool) { let query = submission_status_cancelled_query(SubmissionId::new()); let explained = explain_query_plan(query, db).await; + assert_no_materialization(explained.as_str()); + assert_no_temporary_b_trees(explained.as_str()); insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions_cancelled USING INDEX sqlite_autoindex_submissions_cancelled_1 (id=?) 16, 0, CORRELATED SCALAR SUBQUERY 1 @@ -2095,6 +2120,8 @@ pub mod test { pub async fn test_query_plan_submission_status_paused(db: sqlx::SqlitePool) { let query = submission_status_paused_query(SubmissionId::new()); let explained = explain_query_plan(query, db).await; + assert_no_materialization(explained.as_str()); + assert_no_temporary_b_trees(explained.as_str()); insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions_paused USING INDEX sqlite_autoindex_submissions_paused_1 (id=?) 17, 0, CORRELATED SCALAR SUBQUERY 1 diff --git a/opsqueue/src/consumer/dispatcher/mod.rs b/opsqueue/src/consumer/dispatcher/mod.rs index ce3e9a23..a21d8466 100644 --- a/opsqueue/src/consumer/dispatcher/mod.rs +++ b/opsqueue/src/consumer/dispatcher/mod.rs @@ -379,6 +379,7 @@ mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut writer_conn, ) .await diff --git a/opsqueue/src/producer/client.rs b/opsqueue/src/producer/client.rs index 17a94691..f4ff7124 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -441,7 +441,6 @@ impl InternalProducerClientError { #[cfg(test)] #[cfg(feature = "server-logic")] mod tests { - use crate::{ common::{ StrategicMetadataMap, From adffc65b4ce7e62ef5556201d2230ef282903387 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 16 Sep 2026 10:30:55 +0200 Subject: [PATCH 09/11] Print an actionable message when `just semver` fails --- justfile | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/justfile b/justfile index 3bd50ba3..50ec5902 100644 --- a/justfile +++ b/justfile @@ -113,7 +113,10 @@ semver: set -euo pipefail export DATABASE_URL="sqlite://{{justfile_directory()}}/opsqueue/opsqueue_example_database_schema.db" # We select the latest git tag, not to be confused with the latest Cargo version. - cargo semver-checks --workspace --target x86_64-unknown-linux-gnu --baseline-rev "$(git tag -l --sort=-version:refname | head -1)" + cargo semver-checks --workspace --target x86_64-unknown-linux-gnu --baseline-rev "$(git tag -l --sort=-version:refname | head -1)" || { + echo "Semver checks failed. Please bump the version in Cargo.toml" + exit 1 + } # Rust static analysis [group('lint')] From 16dd14724e77df239820630b0ec0b80938bc3ecf Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 16 Sep 2026 13:50:04 +0200 Subject: [PATCH 10/11] Clarify test --- opsqueue/src/common/chunk.rs | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index 2b957688..a865cb8b 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -1027,13 +1027,15 @@ pub mod test { .await .expect("insertion failed"); - let res = retry_or_fail_chunk(chunk_id, "kapot".into(), &mut conn, 2).await; - assert_matches!(res, Ok(false)); + let max_retries = 2; - let res = retry_or_fail_chunk(chunk_id, "kapot".into(), &mut conn, 2).await; - assert_matches!(res, Ok(true)); + let res = retry_or_fail_chunk(chunk_id, "kapot".into(), &mut conn, max_retries).await; + assert_matches!(res, Ok(false)); // Retry limit not yet reached. - let res = retry_or_fail_chunk(chunk_id, "kapot".into(), &mut conn, 2).await; - assert_matches!(res, Ok(false)); + let res = retry_or_fail_chunk(chunk_id, "kapot".into(), &mut conn, max_retries).await; + assert_matches!(res, Ok(true)); // Retry limit reached, submission is now permanently failed. + + let res = retry_or_fail_chunk(chunk_id, "kapot".into(), &mut conn, max_retries).await; + assert_matches!(res, Ok(false)); // Submission was already failed, check that we ignore. } } From dca3abb1bed296e13eea34820512aadd1f3f5f2f Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 16 Sep 2026 14:00:07 +0200 Subject: [PATCH 11/11] Use an InitialSubmissionStatus enum instead of a paused boolean --- .../python/opsqueue/producer.py | 27 +++++------ libs/opsqueue_python/src/common.rs | 18 ++++++++ libs/opsqueue_python/src/lib.rs | 1 + libs/opsqueue_python/src/producer.rs | 26 ++++++----- libs/opsqueue_python/tests/test_roundtrip.py | 34 +++++++------- opsqueue/benches/chunks_select.rs | 3 +- opsqueue/src/common/chunk.rs | 4 +- opsqueue/src/common/submission.rs | 45 ++++++++++++------- opsqueue/src/consumer/client.rs | 6 +-- opsqueue/src/consumer/dispatcher/mod.rs | 3 +- opsqueue/src/consumer/strategy.rs | 3 +- opsqueue/src/producer/client.rs | 12 ++--- opsqueue/src/producer/common.rs | 5 ++- opsqueue/src/producer/server.rs | 8 ++-- 14 files changed, 117 insertions(+), 78 deletions(-) diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index 9bb72d6c..bdb6b593 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/libs/opsqueue_python/python/opsqueue/producer.py @@ -1,25 +1,24 @@ from __future__ import annotations -from collections.abc import Iterable, Iterator, AsyncIterator -from typing import Any, cast import itertools +from collections.abc import Iterable, Iterator, AsyncIterator +from typing import Any, cast from opentelemetry import trace - from opsqueue.common import ( SerializationFormat, encode_chunk, decode_chunk, DEFAULT_SERIALIZATION_FORMAT, ) -from . import opsqueue_internal -from . import tracing from opsqueue.exceptions import ( SubmissionFailedError, SubmissionNotCancellableError, SubmissionNotFoundError, TooManyMatchingSubmissionsError, ) +from . import opsqueue_internal +from . import tracing from .opsqueue_internal import ( # type: ignore[import-not-found] SubmissionId, SubmissionStatus, @@ -28,21 +27,23 @@ ChunkFailed, SubmissionNotCancellable, SubmissionPaused, + InitialSubmissionStatus, ) __all__ = [ + "ChunkFailed", + "InitialSubmissionStatus", "ProducerClient", - "SubmissionId", - "SubmissionStatus", "SubmissionCompleted", - "SubmissionFailedError", "SubmissionFailed", + "SubmissionFailedError", + "SubmissionId", "SubmissionNotCancellable", "SubmissionNotCancellableError", "SubmissionNotFoundError", "SubmissionPaused", + "SubmissionStatus", "TooManyMatchingSubmissionsError", - "ChunkFailed", ] @@ -150,7 +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, + initial_status: InitialSubmissionStatus = InitialSubmissionStatus.InProgress, ) -> SubmissionId: """ Inserts a submission into the queue, @@ -167,7 +168,7 @@ def insert_submission( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, - paused=paused, + initial_status=initial_status, ) def blocking_stream_completed_submission( @@ -267,7 +268,7 @@ def insert_submission_chunks( metadata: None | bytes = None, strategic_metadata: None | dict[str, int] = None, chunk_size: None | int = None, - paused: bool = False, + initial_status: InitialSubmissionStatus = InitialSubmissionStatus.InProgress, ) -> SubmissionId: """ Inserts an already-chunked submission into the queue, @@ -284,7 +285,7 @@ def insert_submission_chunks( strategic_metadata=strategic_metadata, chunk_size=chunk_size, otel_trace_carrier=otel_trace_carrier, - paused=paused, + initial_status=initial_status, ) def blocking_stream_completed_submission_chunks( diff --git a/libs/opsqueue_python/src/common.rs b/libs/opsqueue_python/src/common.rs index e2f6d378..855dc9f8 100644 --- a/libs/opsqueue_python/src/common.rs +++ b/libs/opsqueue_python/src/common.rs @@ -350,6 +350,24 @@ impl From for SubmissionCance } } +#[pyclass(from_py_object, eq, eq_int)] +#[derive(Default, Debug, Clone, PartialEq, Eq)] +pub enum InitialSubmissionStatus { + Paused, + #[default] + InProgress, +} + +impl From for opsqueue::common::submission::InitialSubmissionStatus { + fn from(value: InitialSubmissionStatus) -> Self { + use opsqueue::common::submission::InitialSubmissionStatus::{InProgress, Paused}; + match value { + InitialSubmissionStatus::Paused => Paused, + InitialSubmissionStatus::InProgress => InProgress, + } + } +} + #[pyclass(from_py_object, frozen, module = "opsqueue")] #[derive(Debug, Clone, PartialEq, Eq)] pub enum SubmissionStatus { diff --git a/libs/opsqueue_python/src/lib.rs b/libs/opsqueue_python/src/lib.rs index 27f6835c..868ae8d6 100644 --- a/libs/opsqueue_python/src/lib.rs +++ b/libs/opsqueue_python/src/lib.rs @@ -26,6 +26,7 @@ fn opsqueue_internal(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index e21a5f53..5aed787f 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -7,6 +7,14 @@ use pyo3::{ use std::time::TryFromFloatSecsError; use std::{future::IntoFuture, sync::Arc, time::Duration}; +use crate::{ + async_util, + common::{ + InitialSubmissionStatus, SubmissionId, SubmissionStatus, run_unless_interrupted, + start_runtime, + }, + errors::{self, CError, CPyResult, FatalPythonException}, +}; use futures::{StreamExt, TryStreamExt, stream::BoxStream}; use opsqueue::{ E, @@ -21,12 +29,6 @@ use opsqueue::{ use tokio::time::error::Elapsed; use ux::u63; -use crate::{ - async_util, - common::{SubmissionId, SubmissionStatus, run_unless_interrupted, start_runtime}, - errors::{self, CError, CPyResult, FatalPythonException}, -}; - create_exception!(opsqueue_internal, ProducerClientError, PyException); const SUBMISSION_POLLING_INTERVAL: Duration = Duration::from_secs(5); @@ -283,7 +285,7 @@ impl ProducerClient { /// /// Returns an error if submission insertion fails. #[allow(clippy::too_many_arguments)] - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), paused=false) + #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), initial_status=InitialSubmissionStatus::default()) )] pub fn insert_submission_direct( &self, @@ -293,7 +295,7 @@ impl ProducerClient { strategic_metadata: Option, chunk_size: Option, otel_trace_carrier: CarrierMap, - paused: bool, + initial_status: InitialSubmissionStatus, ) -> CPyResult> { py.detach(|| { let submission = opsqueue::producer::InsertSubmission { @@ -303,7 +305,7 @@ impl ProducerClient { }, metadata, strategic_metadata: strategic_metadata.unwrap_or_default(), - paused, + initial_status: initial_status.into(), }; self.block_unless_interrupted(async move { self.client @@ -321,7 +323,7 @@ impl ProducerClient { /// /// Returns an error if chunk upload or submission insertion fails. #[allow(clippy::type_complexity, clippy::too_many_arguments)] - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), paused=false) + #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), initial_status=InitialSubmissionStatus::default()) )] pub fn insert_submission_chunks( &self, @@ -331,7 +333,7 @@ impl ProducerClient { strategic_metadata: Option, chunk_size: Option, otel_trace_carrier: CarrierMap, - paused: bool, + initial_status: InitialSubmissionStatus, ) -> CPyResult< SubmissionId, E![ @@ -371,7 +373,7 @@ impl ProducerClient { }, metadata, strategic_metadata: strategic_metadata.unwrap_or_default(), - paused, + initial_status: initial_status.into(), }; self.client .insert_submission(&submission, &otel_trace_carrier) diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index 44a68edf..ecfd0064 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -2,7 +2,21 @@ # - use pytest's `--log-cli-level=info` (or `=debug`) argument to get more detailed logs from the producer/consumer clients # - use `RUST_LOG="opsqueue=info"` (or `opsqueue=debug` or `debug` for even more verbosity), together with to the pytest option `-s` AKA `--capture=no`, to debug the opsqueue binary itself. +import logging +import time from collections.abc import Iterator, Sequence + +import pytest +from conftest import ( + background_process, + multiple_background_processes, + OpsqueueProcess, + opsqueue_service, + StrategyDescription, + strategy_from_description, +) +from opsqueue.common import SerializationFormat +from opsqueue.consumer import ConsumerClient, Chunk from opsqueue.producer import ( SubmissionId, ProducerClient, @@ -15,20 +29,8 @@ SubmissionNotCancellable, SubmissionNotCancellableError, TooManyMatchingSubmissionsError, + InitialSubmissionStatus, ) -from opsqueue.consumer import ConsumerClient, Chunk -from opsqueue.common import SerializationFormat -from conftest import ( - background_process, - multiple_background_processes, - OpsqueueProcess, - opsqueue_service, - StrategyDescription, - strategy_from_description, -) -import logging -import time -import pytest SUBMISSION_COMPLETED_TIMEOUT = 10.0 @@ -762,7 +764,7 @@ def test_unpause_and_complete(opsqueue: OpsqueueProcess) -> None: url = "file:///tmp/opsqueue/test_unpause_and_complete" producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) submission_id = producer_client.insert_submission( - (1, 2, 3), chunk_size=1, paused=True + (1, 2, 3), chunk_size=1, initial_status=InitialSubmissionStatus.Paused ) assert isinstance( @@ -796,7 +798,7 @@ def test_unpause_not_found(opsqueue: OpsqueueProcess) -> None: url = "file:///tmp/opsqueue/test_unpause_not_found" producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) submission_id = producer_client.insert_submission( - (1, 2, 3), chunk_size=1, paused=False + (1, 2, 3), chunk_size=1, initial_status=InitialSubmissionStatus.InProgress ) assert isinstance( producer_client.get_submission_status(submission_id), @@ -811,7 +813,7 @@ def test_cancel_paused(opsqueue: OpsqueueProcess) -> None: url = "file:///tmp/opsqueue/test_cancel_paused" producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) submission_id = producer_client.insert_submission( - (1, 2, 3), chunk_size=1, paused=True + (1, 2, 3), chunk_size=1, initial_status=InitialSubmissionStatus.Paused ) assert isinstance( diff --git a/opsqueue/benches/chunks_select.rs b/opsqueue/benches/chunks_select.rs index 950b4d19..5396f419 100644 --- a/opsqueue/benches/chunks_select.rs +++ b/opsqueue/benches/chunks_select.rs @@ -8,6 +8,7 @@ /// Write result use opsqueue::common::StrategicMetadataMap; use opsqueue::common::chunk::{ChunkId, ChunkSize}; +use opsqueue::common::submission::InitialSubmissionStatus; use opsqueue::common::submission::db::insert_submission_from_chunks; use opsqueue::consumer::dispatcher::Dispatcher; use opsqueue::consumer::strategy::Strategy; @@ -191,7 +192,7 @@ async fn seed_or_extend( None, metadata, ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index a865cb8b..acf2a7d0 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -824,7 +824,7 @@ pub mod db { pub mod test { use crate::common::StrategicMetadataMap; use crate::common::submission::db::{insert_submission, insert_submission_raw}; - use crate::common::submission::{Submission, SubmissionStatus}; + use crate::common::submission::{InitialSubmissionStatus, Submission, SubmissionStatus}; use crate::db::{Connection as _, WriterPool}; use std::assert_matches; @@ -916,7 +916,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index fa40ea08..a304c30a 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -231,6 +231,13 @@ pub struct SubmissionPaused { pub otel_trace_carrier: String, } +#[derive(Debug, Clone, PartialEq, Eq, Default, serde::Serialize, serde::Deserialize)] +pub enum InitialSubmissionStatus { + Paused, + #[default] + InProgress, +} + #[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)] pub enum SubmissionStatus { InProgress(Submission), @@ -313,9 +320,9 @@ pub mod db { use sqlx::{Database, QueryBuilder, Sqlite, query, query_as, query_scalar}; use super::{ - Chunk, ChunkCount, ChunkIndex, DateTime, Duration, E, Metadata, Submission, - SubmissionCancelled, SubmissionCompleted, SubmissionFailed, SubmissionId, SubmissionStatus, - Utc, chunk, + Chunk, ChunkCount, ChunkIndex, DateTime, Duration, E, InitialSubmissionStatus, Metadata, + Submission, SubmissionCancelled, SubmissionCompleted, SubmissionFailed, SubmissionId, + SubmissionStatus, Utc, chunk, }; impl<'q> sqlx::Encode<'q, Sqlite> for SubmissionId { @@ -599,7 +606,7 @@ pub mod db { metadata: Option, strategic_metadata: StrategicMetadataMap, chunk_size: ChunkSize, - paused: bool, + initial_status: InitialSubmissionStatus, mut conn: impl WriterConnection, ) -> Result { let submission_id = SubmissionId::new(); @@ -624,11 +631,15 @@ pub mod db { }) .collect(); - if paused { - insert_paused_submission(submission, chunks, &mut conn).await?; - } else { - insert_submission(submission, chunks, &mut conn).await?; + match initial_status { + InitialSubmissionStatus::Paused => { + insert_paused_submission(submission, chunks, &mut conn).await?; + } + InitialSubmissionStatus::InProgress => { + insert_submission(submission, chunks, &mut conn).await?; + } } + Ok(submission_id) } @@ -1791,7 +1802,7 @@ pub mod test { None, strategic_metadata.clone(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await @@ -1867,7 +1878,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await @@ -1878,7 +1889,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await @@ -1889,7 +1900,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await @@ -1900,7 +1911,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await @@ -1933,7 +1944,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await @@ -1944,7 +1955,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await @@ -1955,7 +1966,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await @@ -2009,7 +2020,7 @@ pub mod test { StrategicMetadataMap::default(), // chunk size ChunkSize::default(), - false, + InitialSubmissionStatus::InProgress, &mut conn, ) .await diff --git a/opsqueue/src/consumer/client.rs b/opsqueue/src/consumer/client.rs index ac37557c..de049359 100644 --- a/opsqueue/src/consumer/client.rs +++ b/opsqueue/src/consumer/client.rs @@ -552,9 +552,9 @@ mod tests { use tokio::task::yield_now; use tokio_util::task::TaskTracker; - use crate::{common::StrategicMetadataMap, db}; - use super::*; + use crate::common::submission::InitialSubmissionStatus; + use crate::{common::StrategicMetadataMap, db}; #[sqlx::test(migrator = "crate::MIGRATOR")] pub async fn test_fetch_chunks(pool: sqlx::SqlitePool) { @@ -578,7 +578,7 @@ mod tests { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await diff --git a/opsqueue/src/consumer/dispatcher/mod.rs b/opsqueue/src/consumer/dispatcher/mod.rs index a21d8466..b352f7d8 100644 --- a/opsqueue/src/consumer/dispatcher/mod.rs +++ b/opsqueue/src/consumer/dispatcher/mod.rs @@ -362,6 +362,7 @@ mod test { use super::*; use crate::common::chunk::ChunkId; use crate::common::chunk::ChunkSize; + use crate::common::submission::InitialSubmissionStatus; use crate::db::DBPools; use tokio::sync::mpsc::unbounded_channel; use ux::u63; @@ -379,7 +380,7 @@ mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut writer_conn, ) .await diff --git a/opsqueue/src/consumer/strategy.rs b/opsqueue/src/consumer/strategy.rs index 8339be4f..3785e33b 100644 --- a/opsqueue/src/consumer/strategy.rs +++ b/opsqueue/src/consumer/strategy.rs @@ -829,6 +829,7 @@ pub mod test { "); } + use crate::common::submission::InitialSubmissionStatus; use crate::db::Connection; use futures::stream::TryStreamExt as _; @@ -849,7 +850,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await diff --git a/opsqueue/src/producer/client.rs b/opsqueue/src/producer/client.rs index f4ff7124..cdd57f6c 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -441,6 +441,8 @@ impl InternalProducerClientError { #[cfg(test)] #[cfg(feature = "server-logic")] mod tests { + use super::*; + use crate::common::submission::InitialSubmissionStatus; use crate::{ common::{ StrategicMetadataMap, @@ -451,8 +453,6 @@ mod tests { producer::common::ChunkContents, }; - use super::*; - async fn start_server_in_background(pool: &sqlx::SqlitePool, url: &str) { let db_pools = DBPools::from_test_pool(pool); // We spawn the separate server @@ -480,7 +480,7 @@ mod tests { None, StrategicMetadataMap::default(), ChunkSize::default(), - false, + InitialSubmissionStatus::default(), &mut conn, ) .await @@ -510,7 +510,7 @@ mod tests { metadata: None, strategic_metadata: StrategicMetadataMap::default(), chunk_size: None, - paused: false, + initial_status: InitialSubmissionStatus::default(), }; client .insert_submission(&submission, &std::collections::HashMap::default()) @@ -554,7 +554,7 @@ mod tests { metadata: None, strategic_metadata: StrategicMetadataMap::default(), chunk_size: None, - paused: false, + initial_status: InitialSubmissionStatus::default(), }; let submission_id = client .insert_submission(&submission, &std::collections::HashMap::default()) @@ -601,7 +601,7 @@ mod tests { metadata: None, strategic_metadata: StrategicMetadataMap::default(), chunk_size: None, - paused: true, + initial_status: InitialSubmissionStatus::Paused, }; let submission_id = client .insert_submission(&submission, &std::collections::HashMap::default()) diff --git a/opsqueue/src/producer/common.rs b/opsqueue/src/producer/common.rs index f98b4b86..2d24b817 100644 --- a/opsqueue/src/producer/common.rs +++ b/opsqueue/src/producer/common.rs @@ -1,5 +1,6 @@ use crate::common::StrategicMetadataMap; +use crate::common::submission::InitialSubmissionStatus; use crate::common::{chunk, submission::Metadata}; /// A producer's request to create a new submission. @@ -10,10 +11,10 @@ pub struct InsertSubmission { #[serde(default)] pub strategic_metadata: StrategicMetadataMap, pub chunk_size: Option, - /// When `true`, the submission is inserted in a paused state and will not + /// When `Paused`, the submission is inserted in a paused state and will not /// be dispatched to consumers until explicitly unpaused. #[serde(default)] - pub paused: bool, + pub initial_status: InitialSubmissionStatus, } /// 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 7e603233..2e51a3ad 100644 --- a/opsqueue/src/producer/server.rs +++ b/opsqueue/src/producer/server.rs @@ -1,7 +1,7 @@ use std::sync::Arc; use crate::common::errors::E::{L, R}; -use crate::common::submission::{self, SubmissionId}; +use crate::common::submission::{self, InitialSubmissionStatus, SubmissionId}; use crate::common::{MaxSubmissions, StrategicMetadataMap}; use crate::db::{self, DBPools}; use crate::tracing::anyhow_as_dyn_error; @@ -226,7 +226,7 @@ async fn insert_submission( request.metadata, request.strategic_metadata, request.chunk_size.unwrap_or_default(), - request.paused, + request.initial_status.clone(), &mut conn, ) .await?; @@ -236,8 +236,8 @@ async fn insert_submission( // this is the moment to perform an extra WAL checkpoint let _ = db::perform_explicit_wal_checkpoint(conn).await; - // Notify waiting consumers, but only for non-paused submissions. - if !request.paused { + // Notify waiting consumers, but only for submissions which start running immediately. + if request.initial_status == InitialSubmissionStatus::InProgress { state.notify_on_insert.notify_waiters(); }