From 64f0fe40f8d182c3085b2370938023064cc94181 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 22 Jul 2026 15:02:33 +0200 Subject: [PATCH 1/7] 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 c27c17e..a56515d 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -226,7 +226,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}; @@ -620,13 +619,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. @@ -634,13 +636,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. @@ -648,13 +653,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. @@ -692,11 +700,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")] @@ -752,12 +760,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")] @@ -775,7 +783,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 { @@ -824,11 +832,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 991744d..a742b38 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, @@ -1048,12 +1047,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. @@ -1061,12 +1064,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. @@ -1074,12 +1081,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, @@ -1381,7 +1392,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())], @@ -1393,7 +1404,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")] @@ -1466,15 +1477,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")] @@ -1499,15 +1504,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")] @@ -1626,18 +1625,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 @@ -1668,15 +1661,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 4f1dc92..1632050 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 74b8707..1fa4d67 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 8bc7f02..d4b308a 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 8780d964703e1a1fa6b3ff93ed2b9fe941e33737 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 22 Jul 2026 12:55:48 +0200 Subject: [PATCH 2/7] 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 82a877e..dc807cc 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 45f0f7d..110dded 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 e3167a9..1b4071a 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::{ @@ -376,57 +377,6 @@ impl ProducerClient { }) } - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default()))] - #[allow(clippy::result_large_err, clippy::type_complexity)] - /// Submit chunks and then stream the completed output chunks. - /// - /// # Errors - /// - /// Returns an error if upload, submission creation, or streaming fails. - 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 @@ -442,17 +392,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 23d9a14..6c86e2c 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 @@ -667,3 +702,25 @@ def test_lookup_too_many_submission_ids_by_strategic_metadata() -> None: ) assert exc.type is TooManyMatchingSubmissionsError assert exc.value.max_submissions == max_ + + +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 77c8b11c26f6f029e3f0515f91d43616c5186fa2 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 23 Jul 2026 16:16:43 +0200 Subject: [PATCH 3/7] 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 c946f33..542aa7a 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 110dded..b0de8e7 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 a56515d..ba60b74 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -225,13 +225,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( @@ -300,25 +300,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(()) @@ -334,9 +327,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) @@ -353,26 +346,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. @@ -394,7 +403,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 @@ -403,23 +412,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 6527810..6504ecd 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 4d59aa2d49b32da468cfd734fb7749da62fe529c Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 23 Jul 2026 16:32:15 +0200 Subject: [PATCH 4/7] 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/opsqueue_example_database_schema.db | Bin 102400 -> 106496 bytes opsqueue/src/common/chunk.rs | 115 ++++- opsqueue/src/common/submission.rs | 447 +++++++++++++++++- 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 + 16 files changed, 896 insertions(+), 39 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 dc807cc..9bb72d6 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 b0fbc0b..abe6b53 100644 --- a/libs/opsqueue_python/src/common.rs +++ b/libs/opsqueue_python/src/common.rs @@ -364,12 +364,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 { @@ -386,6 +389,9 @@ impl From for SubmissionStatus { Cancelled(s) => SubmissionStatus::Cancelled { submission: s.into(), }, + Paused(s) => SubmissionStatus::Paused { + submission: s.into(), + }, } } } @@ -510,6 +516,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 b5f804f..27f6835 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 1b4071a..1c40ad2 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -159,6 +159,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, @@ -247,7 +277,7 @@ impl ProducerClient { /// # Errors /// /// Returns an error if submission insertion fails. - #[pyo3(signature = (chunk_contents, metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default()))] + #[pyo3(signature = (chunk_contents, metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), paused=false))] pub fn insert_submission_direct( &self, py: Python<'_>, @@ -255,6 +285,7 @@ impl ProducerClient { metadata: Option, chunk_size: Option, otel_trace_carrier: CarrierMap, + paused: bool, ) -> CPyResult> { let strategic_metadata = std::collections::HashMap::default(); @@ -266,6 +297,7 @@ impl ProducerClient { }, metadata, strategic_metadata, + paused, }; self.block_unless_interrupted(async move { self.client @@ -277,8 +309,8 @@ impl ProducerClient { }) } - #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default()))] - #[allow(clippy::type_complexity)] + #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default(), paused=false))] + #[allow(clippy::type_complexity, clippy::too_many_arguments)] /// Insert submission chunks via object storage and enqueue the submission. /// /// # Errors @@ -292,6 +324,7 @@ impl ProducerClient { strategic_metadata: Option, chunk_size: Option, otel_trace_carrier: CarrierMap, + paused: bool, ) -> CPyResult< SubmissionId, E![ @@ -331,6 +364,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 6c86e2c..6279ccc 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 @@ -724,3 +725,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 0000000..eed0c27 --- /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 0000000..2a3a60f --- /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/opsqueue_example_database_schema.db b/opsqueue/opsqueue_example_database_schema.db index 83e09941a8e01e276a01805f785de1774aa7e3b4..3d3cdb8976d9c200f9abdf79f3fce5f15aded4b1 100644 GIT binary patch delta 598 zcmZozz}9epZGyC*Gy?;J6cEFJ?nE79M(K?SOZeq@x#lwPFW`6OyUUly`-``qSBqyC zPXPBd?i{XPTyr-o3RH8&M)R_>N*iDUq$2l5dDpi|k_+mC?;Lm1ob~gw z{GTu@MxOUNn}11hvC8rC7cuaE!|;{v-QM>E@A)@2hXe%mgk z$k@op!f}F`d;5GPMs{YGCJ|Y7admaZ=Jb-pq@2{`jMBX9;&_msQd3YkQqDoHjv=lJ zA&yQyt_mnp)6c6jYEO?8VYHggC(5WXy;O$Ljk%FgX8Kz#MhO;>S<@$IGm5f9OckHL zUYk)8EIj!?tE4+bR6(PlC^ap!LPsGpMYG9Xie22$kg-J*?5N_>q}`R!n9DFl zL?$;dNwPuAntXtF(PRf!w#g1GnhNUGwdy)t3P1pMb9`}TRjPugUx=%_YY>;Fv-IS> sJQ~Unees5v=86In2?r>lY`>tvXu!V6Kw*)Cz$6Fm?c52BU-f}(0J{CZ4gdfE delta 446 zcmZoTz}B#UZGyC*Bm)Bj2#W(TGZ1S{)G=n1+?cS0UxtV47X$wSepkM`e0jXTc>8&^ zcy{pwa9`uj;rg}NP~jHW=3i1=tTH_OwG8~<_+Rrs;J?a$ihn=<7XFp|^Y|z8cktJ4 z7Bq|R*crJtGaCGtU*sS#F@T+sX*)*( H;}3lRZ=!@V diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index ba60b74..2526f5a 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -601,6 +601,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 @@ -684,6 +771,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, @@ -780,10 +884,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); } @@ -797,6 +898,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -852,10 +954,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 a742b38..82fa39a 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,12 +720,15 @@ pub mod db { r#" SELECT id AS "id: SubmissionId" FROM submissions WHERE prefix = $1 UNION ALL - SELECT id AS "id: SubmissionId" FROM submissions_completed WHERE prefix = $2 + SELECT id AS "id: SubmissionId" FROM submissions_paused WHERE prefix = $2 UNION ALL - SELECT id AS "id: SubmissionId" FROM submissions_failed WHERE prefix = $3 + SELECT id AS "id: SubmissionId" FROM submissions_completed WHERE prefix = $3 + UNION ALL + SELECT id AS "id: SubmissionId" FROM submissions_failed WHERE prefix = $4 "#, prefix, prefix, + prefix, prefix ) .fetch_optional(conn.get_inner()) @@ -808,6 +952,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) } @@ -877,6 +1055,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)), } } @@ -900,6 +1087,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, @@ -932,6 +1135,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( @@ -1093,6 +1328,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. /// @@ -1206,8 +1475,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::*; @@ -1446,6 +1717,7 @@ pub mod test { None, strategic_metadata.clone(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1521,6 +1793,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1531,6 +1804,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1541,6 +1815,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1551,6 +1826,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1583,6 +1859,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1593,6 +1870,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1603,6 +1881,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1656,6 +1935,7 @@ pub mod test { StrategicMetadataMap::default(), // chunk size ChunkSize::default(), + false, &mut conn, ) .await @@ -1761,4 +2041,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 94abc81..ac37557 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 a725d13..5ac8ae1 100644 --- a/opsqueue/src/consumer/strategy.rs +++ b/opsqueue/src/consumer/strategy.rs @@ -781,6 +781,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 1632050..09e4ab5 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 41f3c8e..f98b4b8 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 1fa4d67..7e60323 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 d4b308a..29f63ba 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 ad14db989fee9e93fd41e16540e2088663715c8c Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 5 Aug 2026 12:06:01 +0200 Subject: [PATCH 5/7] Address feedback --- libs/opsqueue_python/src/producer.rs | 2 +- opsqueue/src/common/chunk.rs | 4 +-- opsqueue/src/common/submission.rs | 41 +++++++++++++--------------- opsqueue/src/producer/client.rs | 2 +- opsqueue/src/prometheus.rs | 2 +- 5 files changed, 24 insertions(+), 27 deletions(-) diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index 1c40ad2..fcf2b63 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -159,7 +159,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 2526f5a..9fa4e88 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -673,7 +673,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; ", @@ -705,7 +705,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 82fa39a..40746d6 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"); } @@ -720,16 +720,13 @@ 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 = $2 + SELECT id AS "id: SubmissionId" FROM submissions_paused WHERE prefix = $1 UNION ALL - SELECT id AS "id: SubmissionId" FROM submissions_completed WHERE prefix = $3 + SELECT id AS "id: SubmissionId" FROM submissions_completed WHERE prefix = $1 UNION ALL - SELECT id AS "id: SubmissionId" FROM submissions_failed WHERE prefix = $4 + SELECT id AS "id: SubmissionId" FROM submissions_failed WHERE prefix = $1 "#, prefix, - prefix, - prefix, - prefix ) .fetch_optional(conn.get_inner()) .await?; @@ -1078,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> { @@ -1094,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> { @@ -1110,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); @@ -1142,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); @@ -1260,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 09e4ab5..17a9469 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 29f63ba..99e7643 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 22c5d1edaa6d8df40f9f2bd46c9aa3ec6175af91 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 5 Aug 2026 13:06:38 +0200 Subject: [PATCH 6/7] 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 40746d6..fab2682 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 3ee2454b371d0f4121d6f35d13c4ab39f985c031 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 5 Aug 2026 15:26:08 +0200 Subject: [PATCH 7/7] 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 b0de8e7..043ac11 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 fcf2b63..ec64acc 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::{ @@ -431,6 +431,7 @@ impl ProducerClient { PyChunksIter, E![ FatalPythonException, + TryFromFloatSecsError, Elapsed, errors::SubmissionFailed, InternalProducerClientError @@ -440,18 +441,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 6279ccc..e0c9568 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -728,7 +728,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) @@ -751,7 +751,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 9fa4e88..a2239a0 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -301,19 +301,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(()) } @@ -327,7 +341,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!( " @@ -336,8 +350,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, @@ -346,9 +359,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. @@ -373,15 +387,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. @@ -438,7 +445,6 @@ pub mod db { completed, failed, or cancelled before. Ignoring.", chunk_id ); - Ok::<_, sqlx::Error>(false) } }) @@ -806,9 +812,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::*; @@ -932,6 +939,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); @@ -957,4 +993,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 fab2682..bc7f3b3 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);