From 2a4c6b16bda9327c3c266adcb61917746da7435f Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 22 Jul 2026 15:02:33 +0200 Subject: [PATCH 1/4] 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 9cc469a..c5b6f26 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -289,7 +289,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, @@ -1047,12 +1046,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. @@ -1060,12 +1063,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. @@ -1073,12 +1080,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, @@ -1380,7 +1391,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())], @@ -1392,7 +1403,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")] @@ -1465,15 +1476,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")] @@ -1498,15 +1503,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")] @@ -1625,18 +1624,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 @@ -1667,15 +1660,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 13de24a..60ec46a 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -380,7 +380,6 @@ impl InternalProducerClientError { #[cfg(test)] #[cfg(feature = "server-logic")] mod tests { - use ux::u63; use crate::{ common::{ @@ -441,7 +440,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 { @@ -459,7 +458,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()) @@ -477,7 +476,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 b0d76a1..1d59692 100644 --- a/opsqueue/src/producer/server.rs +++ b/opsqueue/src/producer/server.rs @@ -218,7 +218,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( @@ -226,5 +226,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 3f12719..2ecdc95 100644 --- a/opsqueue/src/prometheus.rs +++ b/opsqueue/src/prometheus.rs @@ -210,9 +210,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 a62e1431befa7016bf0a4d675445f07c6c12ba32 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 22 Jul 2026 12:55:48 +0200 Subject: [PATCH 2/4] 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 76b4172..960112a 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 4ca536d0735c9ee305a11185d84c42ba7d126e98 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 23 Jul 2026 16:16:43 +0200 Subject: [PATCH 3/4] 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 960112a..f4e2faa 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 efc82db..bad01a2 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}, }; @@ -31,10 +31,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 9fc1444bda3f5f1a7e632199758b6b92239b6cd1 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Thu, 23 Jul 2026 16:32:15 +0200 Subject: [PATCH 4/4] 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 +++ .../20260715143000_pausing.down.sql | 2 + .../migrations/20260715143000_pausing.up.sql | 22 + opsqueue/opsqueue_example_database_schema.db | Bin 102400 -> 118784 bytes opsqueue/src/common/chunk.rs | 115 ++++- opsqueue/src/common/submission.rs | 449 +++++++++++++++++- 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, 898 insertions(+), 39 deletions(-) create mode 100644 opsqueue/migrations/20260715143000_pausing.down.sql create mode 100644 opsqueue/migrations/20260715143000_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/20260715143000_pausing.down.sql b/opsqueue/migrations/20260715143000_pausing.down.sql new file mode 100644 index 0000000..eed0c27 --- /dev/null +++ b/opsqueue/migrations/20260715143000_pausing.down.sql @@ -0,0 +1,2 @@ +DROP TABLE chunks_paused; +DROP TABLE submissions_paused; diff --git a/opsqueue/migrations/20260715143000_pausing.up.sql b/opsqueue/migrations/20260715143000_pausing.up.sql new file mode 100644 index 0000000..2a3a60f --- /dev/null +++ b/opsqueue/migrations/20260715143000_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 ae49303b2d692551e8a9a018ee8a0af45eadaebf..9780cd30caa4a2c8c7662229313843c19e52f86a 100644 GIT binary patch delta 503 zcmZozz}B#UeS);01Oo$uED*zh`a~ULRtW~Zvd)buv-uhMngpf^FtVudGHzkuyTUh_ za|W+Bhc_<+;})LX%>UU|atEK) zj8aSieES%#azEr<4s_Smi4)f|IZIC7%cCJ(kXTxrni6kVT$+@dSzMf%pI00Y5tH67 zrO4RG$i$tueXbJYMGh98uMEW#C+0IZGD=NF3oMwWmjlFj`IL6J^wxUMj=rX3}IY!7grS$k?I@HW<@Cm?Bb>8<=3O zoP2Osm&0^fXMUXL>g^9^*x}+GR=H^PKFO1VyNHL0T_m*HRVPazZ zH~An}Emu<&C%d?$Bx9@pc1{sSUB>AZdW>$9ck>+I-XzP|%p}Vw#lRQG^OgG_cOKU( z_A<83EU`SP%#BQo8Ksy4CQkgZT}qL$k&%gkVf$Pq#)}-2wb_17PuF4W+q{XLo0*Yo PGo!(O`9%%_iyRC9|5Q5M 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 c5b6f26..a75e38e 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -212,12 +212,33 @@ pub struct SubmissionCancelled { pub cancelled_at: DateTime, } +/// A submission that has been paused. +/// +/// Once a submission is paused, it gets moved to the `submissions_paused` +/// table, and its old `submissions` record gets deleted. All remaining +/// (non-completed, non-failed) chunks are similarly moved to `chunks_paused`. +/// +/// A paused submission can be unpaused (resumed) or cancelled. +#[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 { @@ -283,6 +304,7 @@ pub mod db { DatabaseError, E, SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions, }, + submission::SubmissionPaused, }, db::{Connection, True, WriterConnection, WriterPool}, }; @@ -428,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 /// @@ -446,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(); @@ -461,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)| { @@ -469,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?!"); + } } } } @@ -578,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()) @@ -807,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) } @@ -876,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)), } } @@ -899,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, @@ -931,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( @@ -1092,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. /// @@ -1200,6 +1470,8 @@ pub mod db { #[cfg(test)] #[cfg(feature = "server-logic")] pub mod test { + use assert_matches::assert_matches; + use chrono::Utc; use chunk::ChunkSize; use itertools::Itertools; @@ -1207,6 +1479,7 @@ pub mod test { use sqlx::{Row, SqliteConnection}; 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::*; @@ -1445,6 +1718,7 @@ pub mod test { None, strategic_metadata.clone(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1520,6 +1794,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1530,6 +1805,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1540,6 +1816,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1550,6 +1827,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1582,6 +1860,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1592,6 +1871,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1602,6 +1882,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1655,6 +1936,7 @@ pub mod test { StrategicMetadataMap::default(), // chunk size ChunkSize::default(), + false, &mut conn, ) .await @@ -1760,4 +2042,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 594f736..adfddab 100644 --- a/opsqueue/src/consumer/client.rs +++ b/opsqueue/src/consumer/client.rs @@ -573,6 +573,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 89a7ff8..9c756c1 100644 --- a/opsqueue/src/consumer/strategy.rs +++ b/opsqueue/src/consumer/strategy.rs @@ -767,6 +767,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 60ec46a..206cce0 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -194,6 +194,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. @@ -420,6 +463,7 @@ mod tests { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -449,6 +493,7 @@ mod tests { metadata: None, strategic_metadata: StrategicMetadataMap::default(), chunk_size: None, + paused: false, }; client .insert_submission(&submission, &std::collections::HashMap::default()) @@ -492,6 +537,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()) @@ -506,7 +552,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) => { @@ -516,4 +563,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 1d59692..13a7765 100644 --- a/opsqueue/src/producer/server.rs +++ b/opsqueue/src/producer/server.rs @@ -66,6 +66,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), @@ -134,6 +138,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, @@ -195,6 +222,7 @@ async fn insert_submission( request.metadata, request.strategic_metadata, request.chunk_size.unwrap_or_default(), + request.paused, &mut conn, ) .await?; @@ -204,8 +232,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 2ecdc95..bb63d74 100644 --- a/opsqueue/src/prometheus.rs +++ b/opsqueue/src/prometheus.rs @@ -18,6 +18,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"; @@ -66,6 +68,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,