diff --git a/.vscode/opsqueue.code-workspace b/.vscode/opsqueue.code-workspace new file mode 100644 index 00000000..dc0b7ae7 --- /dev/null +++ b/.vscode/opsqueue.code-workspace @@ -0,0 +1,13 @@ +{ + "folders": [ + { + "path": "..", + "name": "opsqueue", + }, + { + "path": "../libs/opsqueue_python", + "name": "opsqueue_python", + } + ], + "settings": {} +} diff --git a/Cargo.lock b/Cargo.lock index 25de1869..51e92265 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1995,7 +1995,7 @@ dependencies = [ [[package]] name = "opsqueue" -version = "0.41.0" +version = "0.42.0" dependencies = [ "anyhow", "arc-swap", @@ -2048,7 +2048,7 @@ dependencies = [ [[package]] name = "opsqueue_python" -version = "0.41.0" +version = "0.42.0" dependencies = [ "anyhow", "chrono", diff --git a/Cargo.toml b/Cargo.toml index a650f0b6..0021f7d9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,7 +14,7 @@ hakari-package = "workspace-hack" [workspace.package] edition = "2024" -version = "0.41.0" +version = "0.42.0" [workspace.dependencies] anyhow = { version = "1.0.102", default-features = false } diff --git a/libs/opsqueue_python/python/opsqueue/producer.py b/libs/opsqueue_python/python/opsqueue/producer.py index bdb6b593..fbe21dbe 100644 --- a/libs/opsqueue_python/python/opsqueue/producer.py +++ b/libs/opsqueue_python/python/opsqueue/producer.py @@ -28,6 +28,7 @@ SubmissionNotCancellable, SubmissionPaused, InitialSubmissionStatus, + Strategy, ) __all__ = [ @@ -240,6 +241,26 @@ def run_submission_chunks( ) return self.blocking_stream_completed_submission_chunks(submission_id, timeout) + def stream_submission_chunks( + self, + submission_id: SubmissionId, + strategy: Strategy, + ) -> Iterator[bytes]: + """Stream chunks progressively; strategy must be Oldest for this submission.""" + return self.inner.stream_submission_chunks( # type: ignore[no-any-return] + submission_id, strategy + ) + + async def async_stream_submission_chunks( + self, + submission_id: SubmissionId, + strategy: Strategy, + ) -> AsyncIterator[bytes]: + """Stream chunks progressively; strategy must be Oldest for this submission.""" + return await self.inner.async_stream_submission_chunks( # type: ignore[no-any-return] + submission_id, strategy + ) + async def async_run_submission_chunks( self, chunk_contents: Iterable[bytes], diff --git a/libs/opsqueue_python/src/errors.rs b/libs/opsqueue_python/src/errors.rs index 043ac11f..4681d1c7 100644 --- a/libs/opsqueue_python/src/errors.rs +++ b/libs/opsqueue_python/src/errors.rs @@ -151,6 +151,7 @@ impl From> for PyErr { } } +#[derive(Debug)] pub struct SubmissionFailed( pub crate::common::SubmissionFailed, pub crate::common::ChunkFailed, diff --git a/libs/opsqueue_python/src/producer.rs b/libs/opsqueue_python/src/producer.rs index 5aed787f..0157e553 100644 --- a/libs/opsqueue_python/src/producer.rs +++ b/libs/opsqueue_python/src/producer.rs @@ -1,6 +1,6 @@ use pyo3::{ create_exception, - exceptions::{PyException, PyStopAsyncIteration}, + exceptions::{PyException, PyStopAsyncIteration, PyValueError}, prelude::*, types::PyIterator, }; @@ -10,7 +10,7 @@ use std::{future::IntoFuture, sync::Arc, time::Duration}; use crate::{ async_util, common::{ - InitialSubmissionStatus, SubmissionId, SubmissionStatus, run_unless_interrupted, + InitialSubmissionStatus, Strategy, SubmissionId, SubmissionStatus, run_unless_interrupted, start_runtime, }, errors::{self, CError, CPyResult, FatalPythonException}, @@ -21,6 +21,7 @@ use opsqueue::{ common::errors::E::{self, L, R}, common::errors::{SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions}, common::{StrategicMetadataMap, chunk, submission}, + consumer::strategy::Strategy as ConsumerStrategy, object_store::{ChunkRetrievalError, ChunkType, ChunksStorageError, NewObjectStoreClientError}, producer::ChunkContents, producer::client::{Client as ActualClient, InternalProducerClientError}, @@ -32,6 +33,7 @@ use ux::u63; create_exception!(opsqueue_internal, ProducerClientError, PyException); const SUBMISSION_POLLING_INTERVAL: Duration = Duration::from_secs(5); +const INITIAL_SUBMISSION_POLLING_INTERVAL: Duration = Duration::from_millis(10); // NOTE: ProducerClient is reasonably cheap to clone, as most of its fields are behind Arcs. #[pyclass(from_py_object, module = "opsqueue")] @@ -421,6 +423,133 @@ impl ProducerClient { }) } + /// Stream output chunks as soon as each consumer has completed them. + /// + /// `strategy` must be `Oldest` or `PreferDistinct` with an underlying strategy + /// that eventually resolves to `Oldest`. Consumers processing this submission + /// must also reserve chunks using the same strategy. + /// + /// # Errors + /// + /// Returns `ValueError` if the strategy does not resolve to `Oldest`. + pub fn stream_submission_chunks( + &self, + submission_id: SubmissionId, + strategy: &Strategy, + ) -> PyResult { + self.streaming_submission_chunks(submission_id, strategy) + } + + #[allow(clippy::too_many_lines)] + fn streaming_submission_chunks( + &self, + submission_id: SubmissionId, + strategy: &Strategy, + ) -> PyResult { + let internal_strategy = ConsumerStrategy::from(strategy); + let mut meta_keys = internal_strategy.meta_keys(); + meta_keys.by_ref().for_each(drop); + if !matches!(meta_keys.take(), ConsumerStrategy::Oldest) { + return Err(PyValueError::new_err( + "streaming submission chunks requires Strategy.Oldest or Strategy.PreferDistinct ending in Strategy.Oldest; consumers must use the same strategy", + )); + } + + let client = self.client.clone(); + let object_store_client = self.object_store_client.clone(); + let stream = futures::stream::unfold( + StreamingChunkState { + client, + object_store_client, + submission_id, + index: u63::new(0), + prefix: None, + ready_until: u63::new(0), + finished: false, + interval: INITIAL_SUBMISSION_POLLING_INTERVAL, + }, + |mut state| async move { + loop { + if state.index < state.ready_until { + let prefix = state + .prefix + .as_deref() + .expect("ready submissions have an object-store prefix"); + let result = state + .object_store_client + .retrieve_chunk(prefix, state.index.into(), ChunkType::Output) + .await + .map_err(StreamingChunkError::Retrieval); + state.index = state.index + u63::new(1); + state.interval = INITIAL_SUBMISSION_POLLING_INTERVAL; + return Some((result, state)); + } + + if state.finished { + return None; + } + + let status = match state + .client + .get_submission(state.submission_id.into()) + .await + { + Ok(Some(status)) => status, + Ok(None) => { + return Some(( + Err(StreamingChunkError::SubmissionNotFound(SubmissionNotFound( + state.submission_id.into(), + ))), + state, + )); + } + Err(error) => { + return Some((Err(StreamingChunkError::Internal(error)), state)); + } + }; + + match status { + submission::SubmissionStatus::InProgress(submission) => { + state.prefix = state.prefix.take().or(submission.prefix); + state.ready_until = submission.chunks_ready.into(); + } + submission::SubmissionStatus::Completed(submission) => { + state.prefix = state.prefix.take().or(submission.prefix); + state.ready_until = submission.chunks_total.into(); + state.finished = true; + } + submission::SubmissionStatus::Failed(submission, chunk) => { + let failure = + crate::common::ChunkFailed::from_internal(chunk, &submission); + state.finished = true; + return Some(( + Err(StreamingChunkError::Failed(Box::new( + crate::errors::SubmissionFailed(submission.into(), failure), + ))), + state, + )); + } + submission::SubmissionStatus::Paused(_) => {} + submission::SubmissionStatus::Cancelled(_) => { + return Some((Err(StreamingChunkError::Cancelled), state)); + } + } + + if state.index < state.ready_until || state.finished { + continue; + } + tokio::time::sleep(state.interval).await; + if state.interval < SUBMISSION_POLLING_INTERVAL { + state.interval = (state.interval * 2).min(SUBMISSION_POLLING_INTERVAL); + } + } + }, + ) + .map(|item| item.map_err(CError)) + .boxed(); + Ok(PyChunksIter::from_stream(self, stream)) + } + /// Blocks (and short-polls) until the submission is completed. /// /// We start with a small short-polling interval @@ -474,6 +603,28 @@ impl ProducerClient { }) } + /// Return an awaitable that resolves immediately to an async iterator of output chunks. + /// + /// The iterator polls submission progress and yields each output chunk as soon as it is ready. + /// + /// # Errors + /// + /// Returns a Python error if creating the awaitable fails. + pub fn async_stream_submission_chunks<'p>( + &self, + py: Python<'p>, + submission_id: SubmissionId, + strategy: &Strategy, + ) -> PyResult> { + let me = self.clone(); + let stream = me.streaming_submission_chunks(submission_id, strategy)?; + let _tokio_active_runtime_guard = me.runtime.enter(); + async_util::future_into_py( + py, + async_util::async_detach(Box::pin(async move { Ok(PyChunksAsyncIter::from(stream)) })), + ) + } + /// Return an awaitable that resolves to an async iterator of output chunks. /// /// # Errors @@ -583,7 +734,53 @@ impl ProducerClient { } } -pub type ChunksStream = BoxStream<'static, CPyResult, ChunkRetrievalError>>; +struct StreamingChunkState { + client: ActualClient, + object_store_client: opsqueue::object_store::ObjectStoreClient, + submission_id: SubmissionId, + index: u63, + prefix: Option, + ready_until: u63, + finished: bool, + interval: Duration, +} + +#[derive(Debug)] +enum StreamingChunkError { + Retrieval(ChunkRetrievalError), + Internal(InternalProducerClientError), + Failed(Box), + SubmissionNotFound(SubmissionNotFound), + Cancelled, +} + +impl std::fmt::Display for StreamingChunkError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Retrieval(error) => error.fmt(f), + Self::Internal(error) => error.fmt(f), + Self::Failed(_) => write!(f, "Submission failed"), + Self::SubmissionNotFound(error) => error.fmt(f), + Self::Cancelled => write!(f, "Submission cancelled"), + } + } +} + +impl std::error::Error for StreamingChunkError {} + +impl From> for PyErr { + fn from(value: CError) -> Self { + match value.0 { + StreamingChunkError::Retrieval(error) => CError(error).into(), + StreamingChunkError::Internal(error) => CError(error).into(), + StreamingChunkError::Failed(error) => CError(*error).into(), + StreamingChunkError::SubmissionNotFound(error) => CError(error).into(), + error @ StreamingChunkError::Cancelled => PyException::new_err(error.to_string()), + } + } +} + +type ChunksStream = BoxStream<'static, CPyResult, StreamingChunkError>>; #[pyclass(module = "opsqueue")] pub struct PyChunksIter { @@ -592,16 +789,20 @@ pub struct PyChunksIter { } impl PyChunksIter { + fn from_stream(client: &ProducerClient, stream: ChunksStream) -> Self { + Self { + stream: Arc::new(tokio::sync::Mutex::new(stream)), + runtime: client.runtime.clone(), + } + } + pub(crate) fn new(client: &ProducerClient, prefix: String, chunks_total: u63) -> Self { let stream = client .object_store_client .retrieve_chunks(prefix, chunks_total, ChunkType::Output) - .map_err(CError) + .map_err(|error| CError(StreamingChunkError::Retrieval(error))) .boxed(); - Self { - stream: Arc::new(tokio::sync::Mutex::new(stream)), - runtime: client.runtime.clone(), - } + Self::from_stream(client, stream) } } @@ -611,7 +812,7 @@ impl PyChunksIter { slf } - fn __next__(&self, py: Python<'_>) -> Option, ChunkRetrievalError>> { + fn __next__(&self, py: Python<'_>) -> Option, StreamingChunkError>> { // The only time we need the GIL is when turning the result back. // By unlocking here, we reduce the chance of deadlocks. py.detach(move || { diff --git a/libs/opsqueue_python/tests/conftest.py b/libs/opsqueue_python/tests/conftest.py index e8c0f7ed..5c49e57b 100644 --- a/libs/opsqueue_python/tests/conftest.py +++ b/libs/opsqueue_python/tests/conftest.py @@ -366,6 +366,21 @@ def multiple_background_processes( ) +def _ends_with_oldest(strategy: StrategyDescription) -> bool: + match strategy: + case "Oldest": + return True + case ("PreferDistinct", _, underlying): + return _ends_with_oldest(underlying) + case _: + return False + + +oldest_strategies: tuple[StrategyDescription, ...] = tuple( + strategy for strategy in any_strategies if _ends_with_oldest(strategy) +) + + @pytest.fixture( scope="function", ids=lambda s: f"Strategy.{strategy_from_description(s)}", @@ -388,6 +403,17 @@ def any_consumer_strategy( yield request.param +@pytest.fixture( + scope="function", + ids=lambda s: f"Strategy.{strategy_from_description(s)}", + params=oldest_strategies, +) +def oldest_consumer_strategy( + request: pytest.FixtureRequest, +) -> Generator[StrategyDescription, None, None]: + yield request.param + + @pytest.fixture(scope="function", params=[json_as_bytes, cbor2, pickle]) def serialization_format( request: pytest.FixtureRequest, diff --git a/libs/opsqueue_python/tests/test_roundtrip.py b/libs/opsqueue_python/tests/test_roundtrip.py index ecfd0064..d96fb525 100644 --- a/libs/opsqueue_python/tests/test_roundtrip.py +++ b/libs/opsqueue_python/tests/test_roundtrip.py @@ -2,34 +2,40 @@ # - use pytest's `--log-cli-level=info` (or `=debug`) argument to get more detailed logs from the producer/consumer clients # - use `RUST_LOG="opsqueue=info"` (or `opsqueue=debug` or `debug` for even more verbosity), together with to the pytest option `-s` AKA `--capture=no`, to debug the opsqueue binary itself. +import asyncio import logging import time from collections.abc import Iterator, Sequence import pytest from conftest import ( + OpsqueueProcess, + StrategyDescription, background_process, multiple_background_processes, - OpsqueueProcess, opsqueue_service, - StrategyDescription, strategy_from_description, ) from opsqueue.common import SerializationFormat -from opsqueue.consumer import ConsumerClient, Chunk +from opsqueue.consumer import ( # type: ignore[attr-defined] + Chunk, + ConsumerClient, + Strategy, + opsqueue_internal, +) from opsqueue.producer import ( - SubmissionId, + ChunkFailed, + InitialSubmissionStatus, ProducerClient, SubmissionCompleted, SubmissionFailed, - ChunkFailed, - SubmissionStatus, SubmissionFailedError, - SubmissionNotFoundError, + SubmissionId, SubmissionNotCancellable, SubmissionNotCancellableError, + SubmissionNotFoundError, + SubmissionStatus, TooManyMatchingSubmissionsError, - InitialSubmissionStatus, ) SUBMISSION_COMPLETED_TIMEOUT = 10.0 @@ -58,7 +64,7 @@ def run_consumer() -> None: consumer_client.run_each_op(increment, strategy=strategy) with background_process(run_consumer) as _consumer: - input_iter = range(0, 100) + input_iter = range(100) output_iter: Iterator[int] = producer_client.run_submission( input_iter, @@ -130,7 +136,7 @@ def run_consumer(_consumer_id: int) -> None: ) with multiple_background_processes(run_consumer, n_consumers) as _consumers: - input_iter = range(0, n_ops) + input_iter = range(n_ops) output_iter: Iterator[int] = producer_client.run_submission( input_iter, @@ -140,7 +146,7 @@ def run_consumer(_consumer_id: int) -> None: ) res = sum(output_iter) - assert res == sum(range(0, n_ops)) + assert res == sum(range(n_ops)) def test_empty_submission(opsqueue: OpsqueueProcess) -> None: @@ -190,7 +196,7 @@ def run_consumer() -> None: ) with background_process(run_consumer) as _consumer: - input_iter = range(0, 100) + input_iter = range(100) output_iter: Iterator[int] = producer_client.run_submission( input_iter, @@ -236,7 +242,7 @@ def broken_increment(input: int) -> float: with background_process(run_consumer) as consumer: logging.error(f"Opsqueue: {opsqueue}") logging.error(f"Consumer: {consumer}") - input_iter = range(0, 100) + input_iter = range(100) with pytest.raises(SubmissionFailedError) as exc_info: producer_client.run_submission( @@ -280,7 +286,7 @@ def increment_list(ints: Sequence[int], _chunk: Chunk) -> Sequence[int]: consumer_client.run_each_chunk(increment_list, strategy=strategy) with background_process(run_consumer) as _consumer: - input_iter = map(lambda i: cbor2.dumps([i, i, i]), range(0, 10)) + input_iter = map(lambda i: cbor2.dumps([i, i, i]), range(10)) output_iter: Iterator[list[int]] = map( lambda c: cbor2.loads(c), producer_client.run_submission_chunks( @@ -323,7 +329,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) + input_iter = range(1000) output_iter: Iterator[int] = producer_client.run_submission( input_iter, chunk_size=100, @@ -356,7 +362,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) + input_iter = range(top) output_iter = await producer_client.async_run_submission( input_iter, chunk_size=1000 ) @@ -749,13 +755,12 @@ def process_op(x: int) -> int: 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, - ) + with background_process(run_consumer) as _consumer, pytest.raises(TimeoutError): + producer_client.run_submission( + [1], + chunk_size=1, + timeout=0.1, + ) def test_unpause_and_complete(opsqueue: OpsqueueProcess) -> None: @@ -824,3 +829,239 @@ def test_cancel_paused(opsqueue: OpsqueueProcess) -> None: assert isinstance( producer_client.get_submission_status(submission_id), SubmissionStatus.Cancelled ) + + +def test_streams_completed_chunks_before_submission_finishes( + opsqueue: OpsqueueProcess, +) -> None: + url = "file:///tmp/opsqueue/test_streaming_results" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission_chunks( + [b"[1]", b"[2]"], chunk_size=1 + ) + + def complete_chunks(_submission_id_value: int) -> None: + consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url) + chunks = sorted( + consumer_client.reserve_chunks( + max=2, + strategy=Strategy.Oldest(), + ), + key=lambda chunk: chunk.chunk_index, + ) + consumer_client.complete_chunk( + chunks[0].submission_id, + chunks[0].submission_prefix, + chunks[0].chunk_index, + chunks[0].input_content, + ) + time.sleep(0.25) + consumer_client.complete_chunk( + chunks[1].submission_id, + chunks[1].submission_prefix, + chunks[1].chunk_index, + chunks[1].input_content, + ) + + with background_process( + complete_chunks, + args=(submission_id.id,), + ): + results = producer_client.stream_submission_chunks( + submission_id, Strategy.Oldest() + ) + assert next(results) == b"[1]" + assert isinstance( + producer_client.get_submission_status(submission_id), + SubmissionStatus.InProgress, + ) + assert next(results) == b"[2]" + + +def test_async_streams_completed_chunks_before_submission_finishes( + opsqueue: OpsqueueProcess, + oldest_consumer_strategy: StrategyDescription, +) -> None: + url = "file:///tmp/opsqueue/test_async_streaming_results" + strategy = strategy_from_description(oldest_consumer_strategy) + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission_chunks( + [b"[1]", b"[2]"], chunk_size=1 + ) + + def complete_chunks(_submission_id_value: int) -> None: + consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url) + chunks = sorted( + consumer_client.reserve_chunks( + max=2, + strategy=strategy_from_description(oldest_consumer_strategy), + ), + key=lambda chunk: chunk.chunk_index, + ) + consumer_client.complete_chunk( + chunks[1].submission_id, + chunks[1].submission_prefix, + chunks[1].chunk_index, + chunks[1].input_content, + ) + time.sleep(0.25) + consumer_client.complete_chunk( + chunks[0].submission_id, + chunks[0].submission_prefix, + chunks[0].chunk_index, + chunks[0].input_content, + ) + + async def collect() -> list[bytes]: + results = await producer_client.async_stream_submission_chunks( + submission_id, strategy + ) + return [chunk async for chunk in results] + + with background_process( + complete_chunks, + args=(submission_id.id,), + ): + assert asyncio.run(collect()) == [b"[1]", b"[2]"] + + +def test_streams_chunks_in_order_when_consumers_complete_out_of_order( + opsqueue: OpsqueueProcess, + oldest_consumer_strategy: StrategyDescription, +) -> None: + url = "file:///tmp/opsqueue/test_streaming_out_of_order_consumers" + strategy = strategy_from_description(oldest_consumer_strategy) + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission_chunks( + [b"[1]", b"[2]"], chunk_size=1 + ) + first_consumer = ConsumerClient(f"localhost:{opsqueue.port}", url) + second_consumer = ConsumerClient(f"localhost:{opsqueue.port}", url) + + [first_chunk] = first_consumer.reserve_chunks(max=1, strategy=strategy) + [second_chunk] = second_consumer.reserve_chunks(max=1, strategy=strategy) + assert (first_chunk.chunk_index.id, second_chunk.chunk_index.id) == (0, 1) + + second_consumer.complete_chunk( + second_chunk.submission_id, + second_chunk.submission_prefix, + second_chunk.chunk_index, + second_chunk.input_content, + ) + + opsqueue_address = f"localhost:{opsqueue.port}" + + def complete_first_chunk( + _submission_id_value: int, + chunk_submission_id: int, + submission_prefix: str, + chunk_index: int, + input_content: bytes, + ) -> None: + time.sleep(0.25) + consumer_client = ConsumerClient(opsqueue_address, url) + consumer_client.complete_chunk( + SubmissionId(chunk_submission_id), + submission_prefix, + opsqueue_internal.ChunkIndex(chunk_index), + input_content, + ) + + with background_process( + complete_first_chunk, + args=( + submission_id.id, + first_chunk.submission_id.id, + first_chunk.submission_prefix, + first_chunk.chunk_index.id, + first_chunk.input_content, + ), + ): + results = producer_client.stream_submission_chunks(submission_id, strategy) + assert next(results) == b"[1]" + assert next(results) == b"[2]" + + +def test_stream_submission_chunks_fails_if_submission_failed_before_read() -> None: + url = "file:///tmp/opsqueue/test_streaming_failed_before_read" + with opsqueue_service(command_args=("--max-chunk-retries", "1")) as opsqueue: + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission_chunks( + [b"[1]", b"[2]"], chunk_size=1 + ) + results = producer_client.stream_submission_chunks( + submission_id, Strategy.Oldest() + ) + consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url) + first_chunk, second_chunk = sorted( + consumer_client.reserve_chunks(max=2, strategy=Strategy.Oldest()), + key=lambda chunk: chunk.chunk_index, + ) + consumer_client.complete_chunk( + first_chunk.submission_id, + first_chunk.submission_prefix, + first_chunk.chunk_index, + first_chunk.input_content, + ) + consumer_client.fail_chunk( + second_chunk.submission_id, + second_chunk.submission_prefix, + second_chunk.chunk_index, + "Simulated failure", + ) + with pytest.raises(SubmissionFailedError): + producer_client.blocking_stream_completed_submission_chunks( + submission_id, timeout=SUBMISSION_COMPLETED_TIMEOUT + ) + with pytest.raises(SubmissionFailedError): + next(results) + with pytest.raises(StopIteration): + next(results) + + +def test_stream_submission_chunks_not_found(opsqueue: OpsqueueProcess) -> None: + url = "file:///tmp/opsqueue/test_streaming_submission_not_found" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = SubmissionId(0) + + results = producer_client.stream_submission_chunks(submission_id, Strategy.Oldest()) + with pytest.raises(SubmissionNotFoundError) as exc_info: + next(results) + assert exc_info.value.submission_id == submission_id.id + + async def read_missing_chunk() -> None: + results = await producer_client.async_stream_submission_chunks( + submission_id, Strategy.Oldest() + ) + await results.__anext__() + + with pytest.raises(SubmissionNotFoundError) as exc_info: + asyncio.run(read_missing_chunk()) + assert exc_info.value.submission_id == submission_id.id + + +def test_stream_submission_chunks_requires_oldest_strategy( + opsqueue: OpsqueueProcess, +) -> None: + url = "file:///tmp/opsqueue/test_streaming_requires_oldest" + producer_client = ProducerClient(f"localhost:{opsqueue.port}", url) + submission_id = producer_client.insert_submission_chunks([b"[1]"], chunk_size=1) + + with pytest.raises(ValueError, match="requires Strategy.Oldest"): + producer_client.stream_submission_chunks(submission_id, Strategy.Random()) + + with pytest.raises(ValueError, match="requires Strategy.Oldest"): + producer_client.stream_submission_chunks( + submission_id, + Strategy.PreferDistinct( + meta_key="company_id", underlying=Strategy.Newest() + ), + ) + + async def collect_with_newest() -> None: + await producer_client.async_stream_submission_chunks( + submission_id, Strategy.Newest() + ) + + with pytest.raises(ValueError, match="requires Strategy.Oldest"): + asyncio.run(collect_with_newest()) diff --git a/opsqueue/src/common/submission.rs b/opsqueue/src/common/submission.rs index a304c30a..5efbf2e2 100644 --- a/opsqueue/src/common/submission.rs +++ b/opsqueue/src/common/submission.rs @@ -156,6 +156,8 @@ pub struct Submission { pub prefix: Option, pub chunks_total: ChunkCount, pub chunks_done: ChunkCount, + #[serde(default = "ChunkCount::zero")] + pub chunks_ready: ChunkCount, pub chunk_size: ChunkSize, pub metadata: Option, #[serde(default)] @@ -262,6 +264,7 @@ impl Submission { prefix: None, chunks_total: ChunkCount::zero(), chunks_done: ChunkCount::zero(), + chunks_ready: ChunkCount::zero(), chunk_size: ChunkSize::default(), metadata: None, strategic_metadata: StrategicMetadataMap::default(), @@ -283,6 +286,7 @@ impl Submission { prefix: None, chunks_total: len, chunks_done: ChunkCount::zero(), + chunks_ready: ChunkCount::zero(), chunk_size, metadata, strategic_metadata: StrategicMetadataMap::default(), @@ -617,6 +621,7 @@ pub mod db { prefix, chunks_total: len, chunks_done: ChunkCount::zero(), + chunks_ready: ChunkCount::zero(), chunk_size, metadata, strategic_metadata, @@ -659,6 +664,7 @@ pub mod db { , prefix , chunks_total AS "chunks_total: ChunkCount" , chunks_done AS "chunks_done: ChunkCount" + , COALESCE((SELECT MIN(chunk_index) FROM chunks WHERE submission_id = submissions.id), chunks_total) AS "chunks_ready: ChunkCount" , chunk_size AS "chunk_size!: ChunkSize" , metadata , ( SELECT json_group_object(metadata_key, metadata_value) @@ -679,6 +685,7 @@ pub mod db { prefix: row.prefix, chunks_total: row.chunks_total, chunks_done: row.chunks_done, + chunks_ready: row.chunks_ready, chunk_size: row.chunk_size, metadata: row.metadata, strategic_metadata: row.strategic_metadata.0, @@ -843,6 +850,7 @@ pub mod db { prefix: row.prefix, chunks_total: row.chunks_total, chunks_done: row.chunks_done, + chunks_ready: row.chunks_ready, chunk_size: row.chunk_size, metadata: row.metadata, strategic_metadata: row.strategic_metadata.0, @@ -916,6 +924,7 @@ pub mod db { prefix: Option, chunks_total: ChunkCount, chunks_done: ChunkCount, + chunks_ready: ChunkCount, chunk_size: ChunkSize, metadata: Option, strategic_metadata: sqlx::types::Json, @@ -939,6 +948,7 @@ pub mod db { , prefix , chunks_total AS "chunks_total: ChunkCount" , chunks_done AS "chunks_done: ChunkCount" + , COALESCE((SELECT MIN(chunk_index) FROM chunks WHERE submission_id = submissions.id), chunks_total) AS "chunks_ready: ChunkCount" , chunk_size AS "chunk_size!: ChunkSize" , metadata , ( SELECT json_group_object(metadata_key, metadata_value) @@ -1698,8 +1708,10 @@ pub mod test { assert_no_temporary_b_trees(explained.as_str()); insta::assert_snapshot!(explained, @" 3, 0, SEARCH submissions USING INDEX sqlite_autoindex_submissions_1 (id=?) - 17, 0, CORRELATED SCALAR SUBQUERY 1 - 22, 17, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) + 15, 0, CORRELATED SCALAR SUBQUERY 1 + 20, 15, SEARCH chunks USING PRIMARY KEY (submission_id=?) + 40, 0, CORRELATED SCALAR SUBQUERY 2 + 45, 40, SEARCH submissions_metadata USING PRIMARY KEY (submission_id=?) "); } diff --git a/opsqueue/src/consumer/strategy.rs b/opsqueue/src/consumer/strategy.rs index 3785e33b..9d0f954a 100644 --- a/opsqueue/src/consumer/strategy.rs +++ b/opsqueue/src/consumer/strategy.rs @@ -98,7 +98,7 @@ impl Strategy { "opsqueue_is_reserved(chunks.submission_id, chunks.chunk_index) = FALSE"; match self { Oldest => qb.push(format!( - "SELECT * FROM chunks WHERE {ffi_is_not_reserved} ORDER BY submission_id ASC" + "SELECT * FROM chunks WHERE {ffi_is_not_reserved} ORDER BY submission_id ASC, chunk_index ASC" )), Newest => qb.push(format!( "SELECT * FROM chunks WHERE {ffi_is_not_reserved} ORDER BY submission_id DESC" @@ -379,7 +379,8 @@ pub mod test { WHERE opsqueue_is_reserved(chunks.submission_id, chunks.chunk_index) = FALSE ORDER BY - submission_id ASC + submission_id ASC, + chunk_index ASC "); let explained = explain(qb, &mut conn).await;