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/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index 82a877e..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", ] @@ -96,6 +98,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 +119,7 @@ def run_submission( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + timeout=timeout, ) return _unchunk_iterator(results_iter, serialization_format) @@ -146,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, @@ -162,6 +167,7 @@ def insert_submission( metadata=metadata, strategic_metadata=strategic_metadata, chunk_size=chunk_size, + paused=paused, ) def blocking_stream_completed_submission( @@ -169,6 +175,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 +188,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 +218,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 +237,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, @@ -259,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, @@ -275,10 +284,13 @@ 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( - 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 +301,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 @@ -326,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. @@ -337,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/errors.rs b/libs/opsqueue_python/src/errors.rs index 76b4172..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; +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()) @@ -201,6 +182,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/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 e3167a9..1c40ad2 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::{ @@ -158,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, @@ -246,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<'_>, @@ -254,6 +285,7 @@ impl ProducerClient { metadata: Option, chunk_size: Option, otel_trace_carrier: CarrierMap, + paused: bool, ) -> CPyResult> { let strategic_metadata = std::collections::HashMap::default(); @@ -265,6 +297,7 @@ impl ProducerClient { }, metadata, strategic_metadata, + paused, }; self.block_unless_interrupted(async move { self.client @@ -276,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 @@ -291,6 +324,7 @@ impl ProducerClient { strategic_metadata: Option, chunk_size: Option, otel_trace_carrier: CarrierMap, + paused: bool, ) -> CPyResult< SubmissionId, E![ @@ -330,6 +364,7 @@ impl ProducerClient { }, metadata, strategic_metadata: strategic_metadata.unwrap_or_default(), + paused, }; self.client .insert_submission(&submission, &otel_trace_carrier) @@ -376,57 +411,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 +426,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..6279ccc 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -27,8 +27,11 @@ strategy_from_description, ) import logging +import time import pytest +SUBMISSION_COMPLETED_TIMEOUT = 10.0 + def increment(data: int) -> int: return data + 1 @@ -56,7 +59,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 +134,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 +153,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 +191,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 +237,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 +281,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 +323,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 +400,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 +447,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 +538,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 +574,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 +609,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 +703,90 @@ 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, + ) + + +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 ae49303..9780cd3 100644 Binary files a/opsqueue/opsqueue_example_database_schema.db and b/opsqueue/opsqueue_example_database_schema.db differ diff --git a/opsqueue/src/common/chunk.rs b/opsqueue/src/common/chunk.rs index c27c17e..2526f5a 100644 --- a/opsqueue/src/common/chunk.rs +++ b/opsqueue/src/common/chunk.rs @@ -225,14 +225,13 @@ impl Chunk { #[cfg(feature = "server-logic")] pub mod db { use super::{ - Chunk, ChunkCompleted, ChunkFailed, ChunkId, ChunkIndex, ChunkSize, DateTime, SubmissionId, - Utc, u63, + 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( @@ -301,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(()) @@ -335,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) @@ -354,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 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. // // (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. @@ -395,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 @@ -404,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) } }) @@ -583,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 @@ -620,13 +725,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 +742,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 +759,33 @@ 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")) + } + + /// 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. @@ -692,11 +823,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 +883,9 @@ pub mod test { .await .expect("complete chunk failed"); - assert_eq!(count_chunks(&mut conn).await.unwrap(), u63::new(0)); - assert_eq!( - count_chunks_completed(&mut conn).await.unwrap(), - u63::new(1) - ); - assert_eq!(count_chunks_failed(&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(), 1); + assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), 0); } #[sqlx::test(migrator = "crate::MIGRATOR")] @@ -770,12 +898,13 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .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 +953,8 @@ pub mod test { .await .expect("Succeed chunk failed"); - assert_eq!(count_chunks(&mut conn).await.unwrap(), u63::new(0)); - assert_eq!( - count_chunks_completed(&mut conn).await.unwrap(), - u63::new(0) - ); - assert_eq!(count_chunks_failed(&mut conn).await.unwrap(), u63::new(1)); + assert_eq!(count_chunks(&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/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); diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index 9cc469a..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,13 +304,13 @@ pub mod db { DatabaseError, E, SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions, }, + submission::SubmissionPaused, }, db::{Connection, True, WriterConnection, WriterPool}, }; 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, @@ -429,9 +450,123 @@ pub mod db { res } + #[tracing::instrument(skip(chunks, conn))] + pub(crate) async fn insert_paused_submission( + submission: Submission, + chunks: Vec, + mut conn: impl WriterConnection, + ) -> Result<(), DatabaseError> { + use axum_prometheus::metrics::counter; + use futures::FutureExt as _; + + let chunks_total = submission.chunks_total.into(); + tracing::debug!("Inserting paused submission {}", submission.id); + + let res = conn + .transaction(move |mut tx| { + async move { + insert_paused_submission_raw(&submission, &mut tx).await?; + insert_submission_metadata_raw( + &submission, + &submission.strategic_metadata, + &mut tx, + ) + .await?; + super::chunk::db::insert_many_paused_chunks(&chunks, &mut tx).await?; + Ok(()) + } + .boxed() + }) + .await; + + counter!(crate::prometheus::SUBMISSIONS_PAUSED_COUNTER).increment(1); + counter!(crate::prometheus::SUBMISSIONS_TOTAL_COUNTER).increment(1); + counter!(crate::prometheus::CHUNKS_TOTAL_COUNTER).increment(chunks_total); + res + } + + #[tracing::instrument(skip(conn))] + async fn insert_paused_submission_raw( + submission: &Submission, + mut conn: impl WriterConnection, + ) -> Result<(), DatabaseError> { + sqlx::query!( + " + INSERT INTO submissions_paused (id, prefix, chunks_total, chunks_done, metadata, otel_trace_carrier, chunk_size) + VALUES ($1, $2, $3, $4, $5, $6, $7) + ", + submission.id, + submission.prefix, + submission.chunks_total, + submission.chunks_done, + submission.metadata, + submission.otel_trace_carrier, + submission.chunk_size.0, + ) + .execute(conn.get_inner()) + .await?; + + Ok(()) + } + + /// Unpause a paused submission. Atomically moves it back from `submissions_paused` + /// to `submissions` and its chunks from `chunks_paused` to `chunks`. + /// + /// # Errors + /// + /// Returns [`DatabaseError`] if the transaction or any SQL query fails. + /// + /// Returns [`SubmissionNotFound`] if the submission is not currently paused. + #[tracing::instrument(skip(conn))] + pub async fn unpause_submission( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + conn.transaction(move |mut tx| { + Box::pin(async move { + unpause_submission_raw(id, &mut tx).await?; + super::chunk::db::restore_paused_chunks(id, &mut tx).await?; + Ok(()) + }) + }) + .await + } + + #[tracing::instrument(skip(conn))] + pub(super) async fn unpause_submission_raw( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + let row = query!( + " + INSERT INTO submissions + (id, chunks_total, chunks_done, prefix, metadata, otel_trace_carrier, chunk_size) + SELECT id, chunks_total, chunks_done, prefix, metadata, otel_trace_carrier, chunk_size + FROM submissions_paused WHERE id = $1; + + DELETE FROM submissions_paused WHERE id = $2 RETURNING *; + ", + id, + id, + ) + .fetch_optional(conn.get_inner()) + .await?; + if row.is_none() { + Err(E::R(SubmissionNotFound(id))) + } else { + counter!(crate::prometheus::SUBMISSIONS_UNPAUSED_COUNTER).increment(1); + Ok(()) + } + } + /// Creates a new submission with the given chunks and inserts it into the database. /// - /// If the number of chunks is 0, the submission is marked as completed immediately afterwards. + /// If `paused` is false and the number of chunks is 0, the submission is marked + /// as completed immediately afterwards. + /// + /// If `paused` is true, the submission is inserted directly into `submissions_paused` + /// (and its chunks into `chunks_paused`), so it won't be picked up by consumers + /// until explicitly unpaused. Zero-chunk paused submissions stay paused. /// /// # Panics /// @@ -447,6 +582,7 @@ pub mod db { metadata: Option, strategic_metadata: StrategicMetadataMap, chunk_size: ChunkSize, + paused: bool, mut conn: impl WriterConnection, ) -> Result { let submission_id = SubmissionId::new(); @@ -462,7 +598,7 @@ pub mod db { strategic_metadata, otel_trace_carrier, }; - let iter = chunks_contents + let chunks: Vec = chunks_contents .into_iter() .enumerate() .map(move |(chunk_index, uri)| { @@ -470,25 +606,30 @@ pub mod db { Chunk::new(submission_id, chunk_index.try_into().unwrap(), uri) }) .collect(); - insert_submission(submission, iter, &mut conn).await?; - // Empty submissions get special handling: we mark them as completed right away. - // See https://github.com/channable/opsqueue/issues/86 for rationale. - if len == 0 { - match maybe_complete_submission(submission_id, conn).await { - // Forward our database errors to the caller. - Err(E::L(e)) => return Err(e), - // If the submission ID can't be found, that's too bad, but it's not our problem anymore i guess. - Err(E::R(_)) => { - tracing::warn!(%submission_id, "Presumed zero-length submission not found"); - } - // If everything went OK, this *could* still indicate a bug in producer code, so let's just log it. - // Our future selves might thank us. - Ok(true) => { - tracing::debug!(%submission_id, "Zero-length submission marked as completed"); - } - // This should never happen. If it does, better log it. - Ok(false) => { - tracing::warn!(%submission_id, "Zero-length submission wasn't zero-length?!"); + + if paused { + insert_paused_submission(submission, chunks, &mut conn).await?; + } else { + insert_submission(submission, chunks, &mut conn).await?; + // Empty submissions get special handling: we mark them as completed right away. + // See https://github.com/channable/opsqueue/issues/86 for rationale. + if len == 0 { + match maybe_complete_submission(submission_id, conn).await { + // Forward our database errors to the caller. + Err(E::L(e)) => return Err(e), + // If the submission ID can't be found, that's too bad, but it's not our problem anymore i guess. + Err(E::R(_)) => { + tracing::warn!(%submission_id, "Presumed zero-length submission not found"); + } + // If everything went OK, this *could* still indicate a bug in producer code, so let's just log it. + // Our future selves might thank us. + Ok(true) => { + tracing::debug!(%submission_id, "Zero-length submission marked as completed"); + } + // This should never happen. If it does, better log it. + Ok(false) => { + tracing::warn!(%submission_id, "Zero-length submission wasn't zero-length?!"); + } } } } @@ -579,12 +720,15 @@ pub mod db { r#" SELECT id AS "id: SubmissionId" FROM submissions WHERE prefix = $1 UNION ALL - SELECT id AS "id: SubmissionId" FROM submissions_completed WHERE prefix = $2 + SELECT id AS "id: SubmissionId" FROM submissions_paused WHERE prefix = $2 + UNION ALL + SELECT id AS "id: SubmissionId" FROM submissions_completed WHERE prefix = $3 UNION ALL - SELECT id AS "id: SubmissionId" FROM submissions_failed WHERE prefix = $3 + SELECT id AS "id: SubmissionId" FROM submissions_failed WHERE prefix = $4 "#, prefix, prefix, + prefix, prefix ) .fetch_optional(conn.get_inner()) @@ -808,6 +952,40 @@ pub mod db { return Ok(Some(SubmissionStatus::Cancelled(cancelled_submission))); } + let paused_row_opt = query!( + r#" + SELECT + id AS "id: SubmissionId" + , prefix + , chunks_total AS "chunks_total: ChunkCount" + , chunks_done AS "chunks_done: ChunkCount" + , chunk_size AS "chunk_size!: ChunkSize" + , metadata + , ( SELECT json_group_object(metadata_key, metadata_value) + FROM submissions_metadata + WHERE submission_id = submissions_paused.id + ) AS "strategic_metadata!: sqlx::types::Json" + , otel_trace_carrier + FROM submissions_paused WHERE id = $1 + "#, + id + ) + .fetch_optional(conn.get_inner()) + .await?; + if let Some(row) = paused_row_opt { + let paused_submission = SubmissionPaused { + id: row.id, + prefix: row.prefix, + chunks_total: row.chunks_total, + chunks_done: row.chunks_done, + chunk_size: row.chunk_size, + metadata: row.metadata, + strategic_metadata: row.strategic_metadata.0, + otel_trace_carrier: row.otel_trace_carrier, + }; + return Ok(Some(SubmissionStatus::Paused(paused_submission))); + } + Ok(None) } @@ -877,6 +1055,15 @@ pub mod db { Ok(Some(SubmissionStatus::Cancelled(submission))) => { Err(E::R(E::R(SubmissionNotCancellable::Cancelled(submission)))) } + Ok(Some(SubmissionStatus::Paused(_))) => { + // Paused submissions are cancellable. + cancel_paused_submission_notx(id, &mut tx).await.map_err( + |e| match e { + E::L(db_err) => E::L(db_err), + E::R(not_found) => E::R(E::L(not_found)), + }, + ) + } Err(db_err) => Err(E::L(db_err)), } } @@ -900,6 +1087,22 @@ pub mod db { Ok(()) } + /// Do not call directly! Must be called inside a transaction. + /// + /// # Errors + /// + /// Returns [`DatabaseError`] if any SQL query fails. + /// + /// Returns [`SubmissionNotFound`] if the submission is not found in `submissions_paused`. + pub async fn cancel_paused_submission_notx( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + cancel_paused_submission_raw(id, &mut conn).await?; + super::chunk::db::skip_remaining_paused_chunks(id, conn).await?; + Ok(()) + } + #[tracing::instrument(skip(conn))] pub(super) async fn cancel_submission_raw( id: SubmissionId, @@ -932,6 +1135,38 @@ pub mod db { } } + #[tracing::instrument(skip(conn))] + pub(super) async fn cancel_paused_submission_raw( + id: SubmissionId, + mut conn: impl WriterConnection, + ) -> Result<(), E> { + let now = chrono::prelude::Utc::now(); + + let submission_opt = query!( + " + INSERT INTO submissions_cancelled + (id, chunks_total, prefix, metadata, cancelled_at, chunks_done) + SELECT id, chunks_total, prefix, metadata, julianday($1), chunks_done FROM submissions_paused WHERE id = $2; + + DELETE FROM submissions_paused WHERE id = $3 RETURNING *; + ", + now, + id, + id, + ) + .fetch_optional(conn.get_inner()) + .await?; + if submission_opt.is_none() { + Err(E::R(SubmissionNotFound(id))) + } else { + counter!(crate::prometheus::SUBMISSIONS_CANCELLED_COUNTER).increment(1); + histogram!(crate::prometheus::SUBMISSIONS_DURATION_CANCEL_HISTOGRAM).record( + crate::prometheus::time_delta_as_f64(Utc::now() - id.timestamp()), + ); + Ok(()) + } + } + #[tracing::instrument(skip(conn))] /// Do not call directly! MUST be called inside a transaction. pub(super) async fn complete_submission_raw( @@ -1047,12 +1282,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 +1299,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 +1316,50 @@ 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")) + } + + /// 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, @@ -1189,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; @@ -1196,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::*; @@ -1380,7 +1664,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 +1676,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")] @@ -1434,6 +1718,7 @@ pub mod test { None, strategic_metadata.clone(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1465,15 +1750,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 +1777,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")] @@ -1521,6 +1794,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1531,6 +1805,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1541,6 +1816,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1551,6 +1827,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1583,6 +1860,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1593,6 +1871,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1603,6 +1882,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -1625,18 +1905,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 @@ -1662,20 +1936,15 @@ pub mod test { StrategicMetadataMap::default(), // chunk size ChunkSize::default(), + false, &mut conn, ) .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. @@ -1773,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 13de24a..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. @@ -380,7 +423,6 @@ impl InternalProducerClientError { #[cfg(test)] #[cfg(feature = "server-logic")] mod tests { - use ux::u63; use crate::{ common::{ @@ -421,6 +463,7 @@ mod tests { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await @@ -441,7 +484,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 { @@ -450,6 +493,7 @@ mod tests { metadata: None, strategic_metadata: StrategicMetadataMap::default(), chunk_size: None, + paused: false, }; client .insert_submission(&submission, &std::collections::HashMap::default()) @@ -459,7 +503,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 +521,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")] @@ -493,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()) @@ -507,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) => { @@ -517,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 b0d76a1..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)) } @@ -218,7 +248,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 +256,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..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, @@ -210,9 +222,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?;