From 2c15037fcf53051ac219c9fd92edc9c691494771 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Wed, 15 Jul 2026 15:07:40 +0200 Subject: [PATCH 1/2] Implement pausing for submissions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Introduce `submissions_paused` and `chunks_paused` tables (alongside the existing `submissions_{completed,failed,cancelled}` and `chunks_{completed,failed}` tables). Pausing a submission atomically moves it from `submissions` → `submissions_paused` and its remaining chunks from `chunks` → `chunks_paused`. Because paused chunks are no longer in the `chunks` table, the consumer dispatcher naturally skips them without any changes to the dispatch query. Unpausing reverses the move and notifies waiting consumers. Paused submissions can also be inserted directly in the paused state (via the new `paused: bool` field on `InsertSubmission`); in that case `notify_on_insert` is intentionally not fired. Paused submissions are cancellable; `cancel_submission` now handles the case where the submission is found in `submissions_paused`. --- .../python/opsqueue/producer.py | 34 +- libs/opsqueue_python/src/common.rs | 42 ++ libs/opsqueue_python/src/lib.rs | 1 + libs/opsqueue_python/src/producer.rs | 55 ++ libs/opsqueue_python/tests/test_roundtrip.py | 162 ++++- .../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 | 111 ++++ opsqueue/src/common/submission.rs | 583 +++++++++++++++++- opsqueue/src/consumer/client.rs | 1 + opsqueue/src/consumer/strategy.rs | 1 + opsqueue/src/producer/client.rs | 85 ++- opsqueue/src/producer/common.rs | 4 + opsqueue/src/producer/server.rs | 54 +- opsqueue/src/prometheus.rs | 12 + 16 files changed, 1126 insertions(+), 43 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 d72f0cd2..91e45a3d 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", ] @@ -326,7 +328,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. @@ -337,6 +339,36 @@ def cancel_submission(self, submission_id: SubmissionId) -> None: """ self.inner.cancel_submission(submission_id) + def pause_submission(self, submission_id: SubmissionId) -> None: + """ + Pause a specific submission that is currently in progress. + + Once paused, the submission's chunks will not be dispatched to consumers + until the submission is unpaused. + + Returns None if the submission was successfully paused. + + Raises: + - `SubmissionNotFoundError` if the submission is not currently in + the in-progress state (it may be already paused, completed, failed, + cancelled, or simply not found). + - `InternalProducerClientError` if there is a low-level internal error. + """ + self.inner.pause_submission(submission_id) + + def unpause_submission(self, submission_id: SubmissionId) -> None: + """ + Unpause a specific submission that is currently paused, + making it available to consumers again. + + 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 f24305b4..9948547d 100644 --- a/libs/opsqueue_python/src/common.rs +++ b/libs/opsqueue_python/src/common.rs @@ -352,6 +352,9 @@ pub enum SubmissionStatus { Cancelled { submission: SubmissionCancelled, }, + Paused { + submission: SubmissionPaused, + }, } impl From for SubmissionStatus { @@ -372,6 +375,9 @@ impl From for SubmissionStatus { Cancelled(s) => SubmissionStatus::Cancelled { submission: s.into(), }, + Paused(s) => SubmissionStatus::Paused { + submission: s.into(), + }, } } } @@ -496,6 +502,42 @@ pub struct SubmissionCancelled { pub cancelled_at: DateTime, } +#[pyclass(frozen, get_all, module = "opsqueue")] +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SubmissionPaused { + pub id: SubmissionId, + pub chunks_total: u64, + pub chunks_done: u64, + pub metadata: Option, + pub strategic_metadata: StrategicMetadataMap, +} + +impl From for SubmissionPaused { + fn from(value: opsqueue::common::submission::SubmissionPaused) -> Self { + Self { + id: value.id.into(), + chunks_total: value.chunks_total.into(), + chunks_done: value.chunks_done.into(), + metadata: value.metadata, + strategic_metadata: value.strategic_metadata, + } + } +} + +#[pymethods] +impl SubmissionPaused { + fn __repr__(&self) -> String { + format!( + "SubmissionPaused(id={0}, chunks_total={1}, chunks_done={2}, metadata={3:?}, strategic_metadata={4:?})", + self.id.__repr__(), + self.chunks_total, + self.chunks_done, + self.metadata, + self.strategic_metadata + ) + } +} + /// Submission could not be cancelled because it was already completed, failed /// or cancelled. #[pyclass(from_py_object, frozen, module = "opsqueue")] diff --git a/libs/opsqueue_python/src/lib.rs b/libs/opsqueue_python/src/lib.rs index b5f804f5..27f6835c 100644 --- a/libs/opsqueue_python/src/lib.rs +++ b/libs/opsqueue_python/src/lib.rs @@ -24,6 +24,7 @@ fn opsqueue_internal(m: &Bound<'_, PyModule>) -> PyResult<()> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_class::()?; m.add_class::()?; m.add_class::()?; diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index 03dabc39..f4272fd7 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -144,6 +144,59 @@ impl ProducerClient { }) } + /// Pause an in-progress submission. + /// + /// Will return an error if the submission is not currently in progress + /// (e.g. already paused, completed, failed, or not found). + #[allow(clippy::result_large_err, clippy::type_complexity)] + pub fn pause_submission( + &self, + py: Python<'_>, + id: SubmissionId, + ) -> CPyResult< + (), + E![ + FatalPythonException, + SubmissionNotFound, + InternalProducerClientError + ], + > { + py.allow_threads(|| { + self.block_unless_interrupted(async { + self.producer_client + .pause_submission(id.into()) + .await + .map_err(|e| CError(R(e))) + }) + }) + } + + /// Unpause a paused submission, making it available to consumers again. + /// + /// Will return an error if the submission is not currently paused. + #[allow(clippy::result_large_err, clippy::type_complexity)] + pub fn unpause_submission( + &self, + py: Python<'_>, + id: SubmissionId, + ) -> CPyResult< + (), + E![ + FatalPythonException, + SubmissionNotFound, + InternalProducerClientError + ], + > { + py.allow_threads(|| { + self.block_unless_interrupted(async { + self.producer_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, @@ -234,6 +287,7 @@ impl ProducerClient { }, metadata, strategic_metadata, + paused: false, }; self.block_unless_interrupted(async move { self.producer_client @@ -294,6 +348,7 @@ impl ProducerClient { }, metadata, strategic_metadata: strategic_metadata.unwrap_or_default(), + paused: false, }; self.producer_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 23d9a149..2d772f4b 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -14,6 +14,7 @@ SubmissionNotFoundError, SubmissionNotCancellable, SubmissionNotCancellableError, + SubmissionPaused, TooManyMatchingSubmissionsError, ) from opsqueue.consumer import ConsumerClient, Chunk @@ -27,6 +28,8 @@ strategy_from_description, ) import logging +import multiprocess # type: ignore[import-untyped] + import pytest @@ -34,9 +37,7 @@ def increment(data: int) -> int: return data + 1 -def test_roundtrip( - opsqueue: OpsqueueProcess, any_consumer_strategy: StrategyDescription -) -> None: +def test_roundtrip(opsqueue: OpsqueueProcess, any_consumer_strategy: StrategyDescription) -> None: """ A most basic test that round-trips all three components. If this fails, something is very wrong. @@ -145,9 +146,7 @@ def test_empty_submission(opsqueue: OpsqueueProcess) -> None: ) input_iter: list[int] = [] - output_iter: Iterator[int] = producer_client.run_submission( - input_iter, chunk_size=20 - ) + output_iter: Iterator[int] = producer_client.run_submission(input_iter, chunk_size=20) res = sum(output_iter) assert res == 0 @@ -303,9 +302,7 @@ def run_consumer(consumer_id: int) -> None: n_consumers = 16 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 - ) + output_iter: Iterator[int] = producer_client.run_submission(input_iter, chunk_size=100) res = sum(output_iter) assert res == sum(range(1, 1001)) @@ -334,9 +331,7 @@ def increment_list(ints: Sequence[int], _chunk: Chunk) -> Sequence[int]: async def run_one_submission(top: int) -> int: logging.debug(f"Running submission {top}") input_iter = range(0, top) - output_iter = await producer_client.async_run_submission( - input_iter, chunk_size=1000 - ) + output_iter = await producer_client.async_run_submission(input_iter, chunk_size=1000) logging.debug(f"Submission for {top} done!") res = 0 @@ -570,9 +565,7 @@ def consume(x: int) -> int | None: raise ValueError(f"Couldn't consume {x}") return x - consumer_client.run_each_op( - consume, strategy=strategy_from_description("Oldest") - ) + consumer_client.run_each_op(consume, strategy=strategy_from_description("Oldest")) with background_process(run_consumer): with pytest.raises(SubmissionFailedError) as exc_info: @@ -667,3 +660,142 @@ def test_lookup_too_many_submission_ids_by_strategic_metadata() -> None: ) assert exc.type is TooManyMatchingSubmissionsError assert exc.value.max_submissions == max_ + + +def test_pause_in_progress(opsqueue: OpsqueueProcess) -> None: + """Pausing an in-progress submission succeeds and its status becomes Paused.""" + url = "file:///tmp/opsqueue/test_pause_in_progress" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission((1, 2, 3), chunk_size=1) + # Sanity-check: submission is in progress before pausing. + assert isinstance( + producer_client.get_submission_status(submission_id), + SubmissionStatus.InProgress, + ) + + producer_client.pause_submission(submission_id) + + status = producer_client.get_submission_status(submission_id) + assert isinstance(status, SubmissionStatus.Paused) + assert isinstance(status.submission, SubmissionPaused) + + +def test_pause_not_found(opsqueue: OpsqueueProcess) -> None: + """Pausing a non-existent submission raises SubmissionNotFoundError.""" + url = "file:///tmp/opsqueue/test_pause_not_found" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + with pytest.raises(SubmissionNotFoundError) as exc_info: + producer_client.pause_submission(SubmissionId(0)) + assert exc_info.value.submission_id == 0 + + +def test_unpause_and_complete( + opsqueue: OpsqueueProcess, any_consumer_strategy: StrategyDescription +) -> 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) + + producer_client.pause_submission(submission_id) + 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) + strategy = strategy_from_description(any_consumer_strategy) + consumer_client.run_each_op(lambda x: x, strategy=strategy) + + 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_pause_after_completing_one_chunk( + opsqueue: OpsqueueProcess, any_consumer_strategy: StrategyDescription +) -> None: + url = "file:///tmp/opsqueue/test_pause_after_completing_one_chunk" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission((1, 2, 3), chunk_size=1) + + # Set by the consumer when it enters the callback for the second op. + at_second_chunk_op = multiprocess.Event() + # Set by the main process to release the consumer after pausing/unpausing. + can_continue = multiprocess.Event() + + def run_consumer() -> None: + consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url) + chunk_count = 0 + + def process_op(x: int) -> int: + nonlocal chunk_count + chunk_count += 1 + if chunk_count == 2: + at_second_chunk_op.set() + can_continue.wait() + return x + + strategy = strategy_from_description(any_consumer_strategy) + consumer_client.run_each_op(process_op, strategy=strategy) + + with background_process(run_consumer): + assert at_second_chunk_op.wait(timeout=10), "Timed out waiting for first chunk to complete" + + producer_client.pause_submission(submission_id) + status = producer_client.get_submission_status(submission_id) + assert isinstance(status, SubmissionStatus.Paused) + assert isinstance(status.submission, SubmissionPaused) + assert status.submission.chunks_done == 1 + + producer_client.unpause_submission(submission_id) + assert isinstance( + producer_client.get_submission_status(submission_id), + SubmissionStatus.InProgress, + ) + + # Release the consumer. + can_continue.set() + + 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) + 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) + + producer_client.pause_submission(submission_id) + 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 00000000..eed0c27f --- /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 00000000..2a3a60f4 --- /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 740fc616..b8c80ad3 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -340,6 +340,9 @@ pub mod db { // 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 // ) @@ -529,6 +532,104 @@ pub mod db { Ok(()) } + #[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 remaining (non-completed, non-failed) chunks of a submission + /// from `chunks` to `chunks_paused`. + #[tracing::instrument(skip(conn))] + pub async fn move_chunks_to_paused( + submission_id: SubmissionId, + mut conn: impl WriterConnection, + ) -> sqlx::Result<()> { + sqlx::query!( + " + INSERT INTO chunks_paused (submission_id, chunk_index, input_content, retries) + SELECT submission_id, chunk_index, input_content, retries FROM chunks WHERE submission_id = $1; + + DELETE FROM chunks WHERE submission_id = $2; + ", + submission_id, + submission_id, + ) + .execute(conn.get_inner()) + .await?; + Ok(()) + } + + /// Move all chunks of a paused submission from `chunks_paused` back to `chunks`. + #[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`. + #[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(()) + } + #[tracing::instrument(skip(conn))] pub async fn skip_remaining_chunks( submission_id: SubmissionId, @@ -583,6 +684,15 @@ pub mod db { Ok(count) } + #[tracing::instrument(skip(db))] + pub async fn count_chunks_paused(mut db: impl Connection) -> sqlx::Result { + let count = sqlx::query!("SELECT COUNT(1) as count FROM chunks_paused;") + .fetch_one(db.get_inner()) + .await?; + let count = u63::new(count.count as u64); + Ok(count) + } + /// Looks up the number of operations in the backlog. /// /// An estimation that returns a slightly too high number, @@ -690,6 +800,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 1ede03ec..8c973452 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -206,12 +206,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 { @@ -406,9 +427,166 @@ pub mod db { res } - /// Creates a new submission with the given chunks and inserts it into the database. + #[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_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(()) + } + + /// Pause an in-progress submission. Atomically moves it from `submissions` + /// to `submissions_paused` and its remaining chunks from `chunks` to `chunks_paused`. + /// + /// Returns `SubmissionNotFound` if the submission is not currently in progress + /// (it may be completed, failed, cancelled, or already paused). + #[tracing::instrument(skip(conn))] + pub async fn pause_submission( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + conn.transaction(move |mut tx| { + Box::pin(async move { + pause_submission_raw(id, &mut tx).await?; + super::chunk::db::move_chunks_to_paused(id, &mut tx).await?; + Ok(()) + }) + }) + .await + } + + #[tracing::instrument(skip(conn))] + pub(super) async fn pause_submission_raw( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + let row = query!( + " + INSERT INTO submissions_paused + (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 WHERE id = $1; + + DELETE FROM submissions WHERE id = $2 RETURNING *; + ", + id, + id, + ) + .fetch_optional(conn.get_inner()) + .await?; + match row { + None => Err(E::R(SubmissionNotFound(id))), + Some(_) => { + counter!(crate::prometheus::SUBMISSIONS_PAUSED_COUNTER).increment(1); + Ok(()) + } + } + } + + /// Unpause a paused submission. Atomically moves it back from `submissions_paused` + /// to `submissions` and its chunks from `chunks_paused` to `chunks`. + /// + /// 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?; + match row { + None => Err(E::R(SubmissionNotFound(id))), + Some(_) => { + counter!(crate::prometheus::SUBMISSIONS_UNPAUSED_COUNTER).increment(1); + Ok(()) + } + } + } + + /// + /// If `paused` is false and the number of chunks is 0, the submission is marked + /// as completed immediately afterwards. /// - /// If 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. #[tracing::instrument(skip(metadata, chunks_contents, conn))] pub async fn insert_submission_from_chunks( prefix: Option, @@ -416,6 +594,7 @@ pub mod db { metadata: Option, strategic_metadata: StrategicMetadataMap, chunk_size: ChunkSize, + paused: bool, mut conn: impl WriterConnection, ) -> Result { let submission_id = SubmissionId::new(); @@ -431,7 +610,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)| { @@ -439,25 +618,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?!") + } } } } @@ -534,12 +718,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()) @@ -751,6 +938,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) } @@ -805,6 +1026,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)), } } @@ -824,6 +1054,16 @@ pub mod db { Ok(()) } + /// Do not call directly! Must be called inside a transaction. + 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, @@ -857,6 +1097,39 @@ 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?; + match submission_opt { + None => Err(E::R(SubmissionNotFound(id))), + Some(_) => { + 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( @@ -982,6 +1255,22 @@ pub mod db { Ok(count.count as usize) } + #[tracing::instrument(skip(db))] + pub async fn count_submissions_paused(mut db: impl Connection) -> sqlx::Result { + let count = sqlx::query!("SELECT COUNT(1) as count FROM submissions_paused;") + .fetch_one(db.get_inner()) + .await?; + Ok(count.count as usize) + } + + #[tracing::instrument(skip(db))] + pub async fn count_submissions_cancelled(mut db: impl Connection) -> sqlx::Result { + let count = sqlx::query!("SELECT COUNT(1) as count FROM submissions_cancelled;") + .fetch_one(db.get_inner()) + .await?; + Ok(count.count as usize) + } + /// Transactionally removes all completed/failed submissions, /// including all their chunks and associated strategic metadata. /// @@ -1095,6 +1384,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::*; @@ -1333,6 +1623,7 @@ pub mod test { None, strategic_metadata.clone(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1408,6 +1699,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1418,6 +1710,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1428,6 +1721,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1438,6 +1732,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1470,6 +1765,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1480,6 +1776,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1490,6 +1787,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1539,6 +1837,7 @@ pub mod test { StrategicMetadataMap::default(), // chunk size ChunkSize::default(), + false, &mut conn, ) .await @@ -1644,4 +1943,240 @@ pub mod test { let deserialized: SubmissionCancelled = serde_json::from_value(json).unwrap(); assert_eq!(deserialized, cancelled); } + + #[sqlx::test] + 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] + pub async fn test_pause_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_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + pause_submission(submission.id, &mut conn).await.unwrap(); + + assert_matches!(count_submissions(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_paused(&mut conn).await, Ok(1)); + assert_matches!(count_submissions_completed(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_failed(&mut conn).await, Ok(0)); + assert_matches!(count_chunks(&mut conn).await, Ok(n) if n == u63::new(0)); + assert_matches!(count_chunks_paused(&mut conn).await, Ok(n) if n == u63::new(3)); + } + + #[sqlx::test] + 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_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + pause_submission(submission.id, &mut conn).await.unwrap(); + assert_matches!(count_submissions(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_paused(&mut conn).await, Ok(1)); + assert_matches!(count_chunks(&mut conn).await, Ok(n) if n == u63::new(0)); + assert_matches!(count_chunks_paused(&mut conn).await, Ok(n) if n == u63::new(3)); + + unpause_submission(submission.id, &mut conn).await.unwrap(); + assert_matches!(count_submissions(&mut conn).await, Ok(1)); + assert_matches!(count_submissions_paused(&mut conn).await, Ok(0)); + assert_matches!(count_chunks(&mut conn).await, Ok(n) if n == u63::new(3)); + assert_matches!(count_chunks_paused(&mut conn).await, Ok(n) if n == u63::new(0)); + } + + #[sqlx::test] + pub async fn test_pause_and_unpause_after_completing_one_chunk(db: sqlx::SqlitePool) { + let db = WriterPool::new(db); + let mut conn = db.writer_conn().await.unwrap(); + let submission_id = insert_submission_from_chunks( + None, + vec![Some("foo".into()), Some("bar".into()), Some("baz".into())], + None, + StrategicMetadataMap::default(), + ChunkSize::default(), + false, + &mut conn, + ) + .await + .expect("insertion failed"); + + // Complete chunk index 0, leaving chunks 1 and 2 pending. + conn.transaction(move |mut tx| { + Box::pin(async move { + crate::common::chunk::db::complete_chunk_raw( + (submission_id, u63::new(0).into()).into(), + None, + &mut tx, + ) + .await + }) + }) + .await + .expect("complete chunk failed"); + + assert_matches!(count_submissions(&mut conn).await, Ok(1)); + assert_matches!(count_chunks(&mut conn).await, Ok(n) if n == u63::new(2)); + + pause_submission(submission_id, &mut conn).await.unwrap(); + + assert_matches!(count_submissions(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_paused(&mut conn).await, Ok(1)); + assert_matches!(count_chunks(&mut conn).await, Ok(n) if n == u63::new(0)); + assert_matches!(count_chunks_paused(&mut conn).await, Ok(n) if n == u63::new(2)); + let status = submission_status(submission_id, &mut conn) + .await + .unwrap() + .unwrap(); + assert_matches!(status, SubmissionStatus::Paused(ref s) if s.chunks_done == 1); + + unpause_submission(submission_id, &mut conn).await.unwrap(); + + assert_matches!(count_submissions(&mut conn).await, Ok(1)); + assert_matches!(count_submissions_paused(&mut conn).await, Ok(0)); + assert_matches!(count_chunks(&mut conn).await, Ok(n) if n == u63::new(2)); + assert_matches!(count_chunks_paused(&mut conn).await, Ok(n) if n == u63::new(0)); + let status = submission_status(submission_id, &mut conn) + .await + .unwrap() + .unwrap(); + assert_matches!(status, SubmissionStatus::InProgress(ref s) if s.chunks_done == 1); + } + + #[sqlx::test] + 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_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + pause_submission(submission.id, &mut conn).await.unwrap(); + assert_matches!(count_submissions_paused(&mut conn).await, Ok(1)); + assert_matches!(count_chunks_paused(&mut conn).await, Ok(n) if n == u63::new(3)); + + cancel_submission(submission.id, &mut conn).await.unwrap(); + + assert_matches!(count_submissions(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_paused(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_completed(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_failed(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_cancelled(&mut conn).await, Ok(1)); + assert_matches!(count_chunks(&mut conn).await, Ok(n) if n == u63::new(0)); + assert_matches!(count_chunks_failed(&mut conn).await, Ok(n) if n == u63::new(3)); + assert_matches!(count_chunks_paused(&mut conn).await, Ok(n) if n == u63::new(0)); + } + + #[sqlx::test] + 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_submission(submission.clone(), chunks, &mut conn) + .await + .expect("insertion failed"); + + pause_submission(submission.id, &mut conn).await.unwrap(); + + let status = submission_status(submission.id, &mut conn) + .await + .unwrap() + .unwrap(); + assert_matches!(status, SubmissionStatus::Paused(_)); + } + + #[sqlx::test] + /// Test that a submission inserted directly in the paused state stays paused + /// (unlike empty non-paused submissions which are auto-completed). + pub async fn insert_paused_submission_stays_paused(db: sqlx::SqlitePool) { + let db = WriterPool::new(db); + let mut conn = db.writer_conn().await.unwrap(); + insert_submission_from_chunks( + None, + vec![Some("foo".into()), Some("bar".into())], + None, + StrategicMetadataMap::default(), + ChunkSize::default(), + true, + &mut conn, + ) + .await + .expect("insertion failed"); + + assert_matches!(count_submissions(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_paused(&mut conn).await, Ok(1)); + assert_matches!(count_submissions_completed(&mut conn).await, Ok(0)); + } + + #[sqlx::test] + /// 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(); + insert_submission_from_chunks( + None, + vec![], + None, + StrategicMetadataMap::default(), + ChunkSize::default(), + true, + &mut conn, + ) + .await + .expect("insertion failed"); + + assert_matches!(count_submissions(&mut conn).await, Ok(0)); + assert_matches!(count_submissions_paused(&mut conn).await, Ok(1)); + assert_matches!(count_submissions_completed(&mut conn).await, Ok(0)); + } } diff --git a/opsqueue/src/consumer/client.rs b/opsqueue/src/consumer/client.rs index 3b2c4aee..970ad869 100644 --- a/opsqueue/src/consumer/client.rs +++ b/opsqueue/src/consumer/client.rs @@ -521,6 +521,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 573475d1..a4a220ba 100644 --- a/opsqueue/src/consumer/strategy.rs +++ b/opsqueue/src/consumer/strategy.rs @@ -743,6 +743,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 b7c9983c..9eab98df 100644 --- a/opsqueue/src/producer/client.rs +++ b/opsqueue/src/producer/client.rs @@ -174,6 +174,85 @@ impl Client { .await } + /// Pause an in-progress submission. + /// + /// Returns an error if the submission is not currently in progress + /// (e.g. it is already paused, completed, failed, or cancelled). + pub async fn pause_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/pause/{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 + } + + /// Unpause a paused submission, making it available to consumers again. + /// + /// Returns an error if the submission is not currently paused. + 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. @@ -379,6 +458,7 @@ mod tests { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -408,6 +488,7 @@ mod tests { metadata: None, strategic_metadata: Default::default(), chunk_size: None, + paused: false, }; client .insert_submission(&submission, &Default::default()) @@ -451,6 +532,7 @@ mod tests { metadata: None, strategic_metadata: Default::default(), chunk_size: None, + paused: false, }; let submission_id = client .insert_submission(&submission, &Default::default()) @@ -465,7 +547,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) => { diff --git a/opsqueue/src/producer/common.rs b/opsqueue/src/producer/common.rs index 41f3c8e9..f98b4b86 100644 --- a/opsqueue/src/producer/common.rs +++ b/opsqueue/src/producer/common.rs @@ -10,6 +10,10 @@ pub struct InsertSubmission { #[serde(default)] pub strategic_metadata: StrategicMetadataMap, pub chunk_size: Option, + /// When `true`, the submission is inserted in a paused state and will not + /// be dispatched to consumers until explicitly unpaused. + #[serde(default)] + pub paused: bool, } /// Either embedded chunk contents or a reference to object storage. diff --git a/opsqueue/src/producer/server.rs b/opsqueue/src/producer/server.rs index 2788d071..dd9c798d 100644 --- a/opsqueue/src/producer/server.rs +++ b/opsqueue/src/producer/server.rs @@ -60,6 +60,11 @@ impl ServerState { "/submissions/cancel/{submission_id}", post(cancel_submission), ) + .route("/submissions/pause/{submission_id}", post(pause_submission)) + .route( + "/submissions/unpause/{submission_id}", + post(unpause_submission), + ) .route( "/submissions/count_completed", get(submissions_count_completed), @@ -128,6 +133,48 @@ async fn cancel_submission( } } +/// 200 if the submission was successfully paused. +/// 404 if the submission could not be found in the in-progress state. +/// 500 if a DatabaseError occurred. +async fn pause_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::pause_submission(submission_id, &mut conn).await { + Ok(_) => 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()), + } +} + +/// 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, @@ -189,6 +236,7 @@ async fn insert_submission( request.metadata, request.strategic_metadata, request.chunk_size.unwrap_or_default(), + request.paused, &mut conn, ) .await?; @@ -198,8 +246,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 d0c3dda5..eb42144e 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"; @@ -65,6 +67,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 71ace4eab16666feadd10a4909a8fb332ba3f498 Mon Sep 17 00:00:00 2001 From: Sem Mulder Date: Fri, 17 Jul 2026 16:20:42 +0200 Subject: [PATCH 2/2] Make complete_chunk and fail_chunk not error when processing previously completed, failed, or paused chunks --- .../python/opsqueue/exceptions.py | 9 -- libs/opsqueue_python/src/errors.rs | 23 +--- opsqueue/src/common/chunk.rs | 130 ++++++++++-------- opsqueue/src/common/errors.rs | 6 +- 4 files changed, 78 insertions(+), 90 deletions(-) diff --git a/libs/opsqueue_python/python/opsqueue/exceptions.py b/libs/opsqueue_python/python/opsqueue/exceptions.py index 615dc958..3f05fa59 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 760c2d7e..212f91ba 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; 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 b8c80ad3..77f5c2e4 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -218,11 +218,11 @@ impl Chunk { #[cfg(feature = "server-logic")] pub mod db { use super::*; - 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( @@ -281,25 +281,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(()) @@ -311,9 +304,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) @@ -330,15 +323,20 @@ 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 or pausing 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, and for paused submissions we will retry the chunk when it + // becomes unpaused. // - // 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. + // By only updating `chunks_done` when we actually moved a chunk, we ensure that 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. @@ -346,13 +344,22 @@ pub mod db { // (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, cancelled, or paused before. Ignoring.", + chunk_id + ); + } + + Ok(()) } #[tracing::instrument(skip(conn))] @@ -369,7 +376,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 @@ -378,24 +385,37 @@ 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) - } 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) + match retries { + Some(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) + } + } + None => { + tracing::warn!( + "Could not fail chunk {:?} because it was either: \ + completed, failed, cancelled, or paused before. Ignoring.", + chunk_id + ); + + Ok::<_, sqlx::Error>(false) + } } }) }) diff --git a/opsqueue/src/common/errors.rs b/opsqueue/src/common/errors.rs index efc82dba..bad01a2b 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);