diff --git a/extra/experimentation/consumer.py b/extra/experimentation/consumer.py index 45e5188d..72634112 100755 --- a/extra/experimentation/consumer.py +++ b/extra/experimentation/consumer.py @@ -74,9 +74,7 @@ def main(): # HACK: It's ugly that we need to modify the URL here. Does the google.cloud library have no option to download a gs:// URL directly? blob_name = ( - chunk["submission_directory"].replace( - "gs://channable-opsqueue-experimentation/", "" - ) + chunk["submission_directory"].replace("gs://channable-opsqueue-experimentation/", "") + "/" + str(chunk["chunk_id"]) ) diff --git a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer.py b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer.py index 792ea8ac..ebcf5a52 100644 --- a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer.py +++ b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer.py @@ -4,9 +4,7 @@ def run_a_consumer(id: int) -> None: - logging.basicConfig( - format=f"Consumer {id} - %(levelname)s: %(message)s", level=logging.INFO - ) + logging.basicConfig(format=f"Consumer {id} - %(levelname)s: %(message)s", level=logging.INFO) def my_operation(data: int) -> int: # time.sleep(0.01) @@ -19,8 +17,7 @@ def my_operation(data: int) -> int: def main() -> None: n_consumers = 16 processes = [ - multiprocessing.Process(target=run_a_consumer, args=(id,)) - for id in range(n_consumers) + multiprocessing.Process(target=run_a_consumer, args=(id,)) for id in range(n_consumers) ] for p in processes: p.start() diff --git a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer_preferdistinct.py b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer_preferdistinct.py index bc35f7c3..0967955b 100644 --- a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer_preferdistinct.py +++ b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_consumer_preferdistinct.py @@ -9,9 +9,7 @@ def run_a_consumer(id: int) -> None: - logging.basicConfig( - format=f"Consumer {id} - %(levelname)s: %(message)s", level=logging.INFO - ) + logging.basicConfig(format=f"Consumer {id} - %(levelname)s: %(message)s", level=logging.INFO) def my_operation(data: int) -> int: # time.sleep(0.01) @@ -25,8 +23,7 @@ def my_operation(data: int) -> int: def main() -> None: n_consumers = 4 processes = [ - multiprocessing.Process(target=run_a_consumer, args=(id,)) - for id in range(n_consumers) + multiprocessing.Process(target=run_a_consumer, args=(id,)) for id in range(n_consumers) ] for p in processes: p.start() diff --git a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_producer_many.py b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_producer_many.py index a91698f7..3215a7f4 100644 --- a/libs/opsqueue_python/examples/multiprocessing/multiprocessing_producer_many.py +++ b/libs/opsqueue_python/examples/multiprocessing/multiprocessing_producer_many.py @@ -2,9 +2,7 @@ import multiprocessing from opsqueue.producer import ProducerClient -logging.basicConfig( - format="%(levelname)s - %(asctime)s: %(message)s", level=logging.INFO -) +logging.basicConfig(format="%(levelname)s - %(asctime)s: %(message)s", level=logging.INFO) def run_a_producer(id: int) -> None: @@ -22,8 +20,7 @@ def main() -> None: f"Starting {n_producers} producers... When `multiprocessing_consumer_preferdistinct.py` is used, expecting all submissions to finish roughly at the same time" ) processes = [ - multiprocessing.Process(target=run_a_producer, args=(id,)) - for id in range(n_producers) + multiprocessing.Process(target=run_a_producer, args=(id,)) for id in range(n_producers) ] for p in processes: p.start() diff --git a/libs/opsqueue_python/examples/tracing/tracing_consumer.py b/libs/opsqueue_python/examples/tracing/tracing_consumer.py index a319117f..81ecd12a 100644 --- a/libs/opsqueue_python/examples/tracing/tracing_consumer.py +++ b/libs/opsqueue_python/examples/tracing/tracing_consumer.py @@ -17,9 +17,7 @@ def set_up_global_tracer() -> None: """ This is usually called once per app, at startup time. """ - resource = Resource( - attributes={SERVICE_NAME: "tracing_with_opsqueue_example_consumer"} - ) + resource = Resource(attributes={SERVICE_NAME: "tracing_with_opsqueue_example_consumer"}) provider = TracerProvider(resource=resource) processor = BatchSpanProcessor(OTLPSpanExporter()) provider.add_span_processor(processor) diff --git a/libs/opsqueue_python/examples/tracing/tracing_producer.py b/libs/opsqueue_python/examples/tracing/tracing_producer.py index 24f0f136..cf481878 100644 --- a/libs/opsqueue_python/examples/tracing/tracing_producer.py +++ b/libs/opsqueue_python/examples/tracing/tracing_producer.py @@ -24,9 +24,7 @@ def set_up_global_tracer() -> None: """ This is usually called once per app, at startup time. """ - resource = Resource( - attributes={SERVICE_NAME: "tracing_with_opsqueue_example_producer"} - ) + resource = Resource(attributes={SERVICE_NAME: "tracing_with_opsqueue_example_producer"}) provider = TracerProvider(resource=resource) processor = BatchSpanProcessor(OTLPSpanExporter()) provider.add_span_processor(processor) @@ -39,9 +37,7 @@ def do_something() -> None: with added_baggage(baggage={"app_mode": "preview"}): # print(opentelemetry.baggage.get_all()) - client = ProducerClient( - "localhost:3999", "file:///tmp/opsqueue/tracing_example" - ) + client = ProducerClient("localhost:3999", "file:///tmp/opsqueue/tracing_example") input_iter = range(0, 100) output_iter = client.run_submission(input_iter, chunk_size=10) diff --git a/libs/opsqueue_python/python/opsqueue/common.py b/libs/opsqueue_python/python/opsqueue/common.py index 80abf9aa..974a910c 100644 --- a/libs/opsqueue_python/python/opsqueue/common.py +++ b/libs/opsqueue_python/python/opsqueue/common.py @@ -29,15 +29,11 @@ def loads(cls, data: bytes) -> Any: DEFAULT_SERIALIZATION_FORMAT: SerializationFormat = cbor2 -def encode_chunk( - chunk: Sequence[Any], serialization_format: SerializationFormat -) -> bytes: +def encode_chunk(chunk: Sequence[Any], serialization_format: SerializationFormat) -> bytes: return serialization_format.dumps(chunk) -def decode_chunk( - chunk: bytes, serialization_format: SerializationFormat -) -> Sequence[Any]: +def decode_chunk(chunk: bytes, serialization_format: SerializationFormat) -> Sequence[Any]: res = serialization_format.loads(chunk) assert isinstance(res, Sequence), ( f"Decoding a chunk should always return a sequence, got unexpected type {type(res)}" diff --git a/libs/opsqueue_python/python/opsqueue/consumer.py b/libs/opsqueue_python/python/opsqueue/consumer.py index a1311d57..207a5734 100644 --- a/libs/opsqueue_python/python/opsqueue/consumer.py +++ b/libs/opsqueue_python/python/opsqueue/consumer.py @@ -85,9 +85,7 @@ def run_each_chunk( ) -> None: def raw_chunk_callback(chunk: Chunk) -> bytes: ctx = _trace_context_from_chunk(chunk) - with opentelemetry.trace.get_tracer( - "opsqueue.consumer" - ).start_as_current_span( + with opentelemetry.trace.get_tracer("opsqueue.consumer").start_as_current_span( "run_chunk", context=ctx, kind=opentelemetry.trace.SpanKind.CONSUMER ) as span: span.set_attribute("submission_id", chunk.submission_id.id) @@ -96,9 +94,7 @@ def raw_chunk_callback(chunk: Chunk) -> bytes: safe_v = typing.cast(opentelemetry.util.types.AttributeValue, v) span.set_attribute(k, safe_v) - chunk_contents = common.decode_chunk( - chunk.input_content, serialization_format - ) + chunk_contents = common.decode_chunk(chunk.input_content, serialization_format) chunk_result_contents = chunk_callback(chunk_contents, chunk) return common.encode_chunk(chunk_result_contents, serialization_format) @@ -121,9 +117,7 @@ def run_each_chunk_raw( """ self.inner.run_per_chunk(strategy, chunk_callback) - def reserve_chunks( - self, max: int = 1, strategy: Strategy = Strategy.Newest - ) -> list[Chunk]: + def reserve_chunks(self, max: int = 1, strategy: Strategy = Strategy.Newest) -> list[Chunk]: """ Low-level function to manually reserve one or more chunks for processing. @@ -155,9 +149,7 @@ def complete_chunk( Raises: - InternalConsumerClientError if there is a low-level internal error """ - self.inner.complete_chunk( - submission_id, submission_prefix, chunk_index, output_content - ) + self.inner.complete_chunk(submission_id, submission_prefix, chunk_index, output_content) def fail_chunk( self, @@ -183,6 +175,4 @@ def fail_chunk( def _trace_context_from_chunk(chunk: Chunk) -> opentelemetry.context.Context: - return tracing.carrier_to_opentelemetry_tracecontext( - chunk.submission_otel_trace_carrier - ) + return tracing.carrier_to_opentelemetry_tracecontext(chunk.submission_otel_trace_carrier) 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/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index d72f0cd2..c35cc3d2 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,9 +339,37 @@ def cancel_submission(self, submission_id: SubmissionId) -> None: """ self.inner.cancel_submission(submission_id) - def get_submission_status( - self, submission_id: SubmissionId - ) -> SubmissionStatus | None: + 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: """ Retrieve the status (in progress, completed, or failed) of a specific submission. diff --git a/libs/opsqueue_python/src/common.rs b/libs/opsqueue_python/src/common.rs index f24305b4..872f9d16 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(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 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/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..e99fab5a 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.detach(|| { + 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.detach(|| { + 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/conftest.py b/libs/opsqueue_python/tests/conftest.py index 0052bca1..c8ed23ca 100644 --- a/libs/opsqueue_python/tests/conftest.py +++ b/libs/opsqueue_python/tests/conftest.py @@ -161,10 +161,7 @@ def multiple_background_processes( function: Callable[[int], None], count: int ) -> Generator[list[multiprocess.Process], None, None]: with ExitStack() as stack: - yield [ - stack.enter_context(background_process(function, args=(p,))) - for p in range(count) - ] + yield [stack.enter_context(background_process(function, args=(p,))) for p in range(count)] type StrategyDescription = str | tuple[str, str, StrategyDescription] diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index 23d9a149..4b98c95f 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. @@ -84,9 +85,7 @@ def test_complete_then_fail_chunks( chunk_size = 10 n_chunks = n_ops // chunk_size n_consumers = 3 - assert n_chunks // n_consumers > 1, ( - "Need more than one chunk per consumer for this test" - ) + assert n_chunks // n_consumers > 1, "Need more than one chunk per consumer for this test" def run_consumer(_consumer_id: int) -> None: consumer_client = ConsumerClient( @@ -145,9 +144,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 +300,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 +329,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 +563,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: @@ -594,16 +585,10 @@ def test_lookup_submission_ids_by_strategic_metadata(opsqueue: OpsqueueProcess) [1], chunk_size=1, strategic_metadata={"foo": 1, "bar": 2, "moo": 3} ) # Inserting some similar data to that above, which shouldn't get matched. - producer_client.insert_submission( - [1], chunk_size=1, strategic_metadata={"foo": 2, "bar": 1} - ) + producer_client.insert_submission([1], chunk_size=1, strategic_metadata={"foo": 2, "bar": 1}) - def test_lookup( - strategic_metadata: dict[str, int], expected_ids: list[int] - ) -> None: - found_ids = producer_client.lookup_submission_ids_by_strategic_metadata( - strategic_metadata - ) + def test_lookup(strategic_metadata: dict[str, int], expected_ids: list[int]) -> None: + found_ids = producer_client.lookup_submission_ids_by_strategic_metadata(strategic_metadata) assert isinstance(found_ids, list) assert all(map(lambda x: isinstance(x, SubmissionId), found_ids)) assert found_ids == expected_ids @@ -641,19 +626,14 @@ def test_lookup_too_many_submission_ids_by_strategic_metadata() -> None: max_ = 2 # We didn't request the OpsQueueProcess as a parameter so an instance isn't # started, instead we start one here with custom args. - with opsqueue_service( - command_args=["--max-submissions-returned", str(max_)] - ) as opsqueue: + with opsqueue_service(command_args=["--max-submissions-returned", str(max_)]) as opsqueue: url = "file:///tmp/opsqueue/test_lookup_too_many_matching_submissions" producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) inserted: list[SubmissionId] = [] strategic_metadata = {"k": 1} for _ in range(max_ + 1): - assert ( - inserted - == producer_client.lookup_submission_ids_by_strategic_metadata( - strategic_metadata - ) + assert inserted == producer_client.lookup_submission_ids_by_strategic_metadata( + strategic_metadata ) inserted.append( producer_client.insert_submission( @@ -662,8 +642,145 @@ def test_lookup_too_many_submission_ids_by_strategic_metadata() -> None: ) with pytest.raises(TooManyMatchingSubmissionsError) as exc: assert len(inserted) == max_ + 1 - producer_client.lookup_submission_ids_by_strategic_metadata( - strategic_metadata - ) + producer_client.lookup_submission_ids_by_strategic_metadata(strategic_metadata) 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 ae49303b..9780cd30 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 740fc616..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,26 +323,43 @@ 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. // // (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))] @@ -366,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 @@ -375,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) + } } }) }) @@ -529,6 +552,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 +704,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 +820,7 @@ pub mod test { None, StrategicMetadataMap::default(), ChunkSize::default(), + false, &mut conn, ) .await 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); 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,