Make tests timeout properly, in preparation for showing failure logs properly

Author: SemMulder

libs/opsqueue_python/python/opsqueue/producer.py

@@ -98,6 +98,7 @@ def run_submission(
         serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT,
         metadata: None | bytes = None,
         strategic_metadata: None | dict[str, int] = None,
+        timeout: float | None = None,
     ) -> Iterator[Any]:
         """
         Inserts a submission into the queue, and blocks until it is completed.
@@ -118,6 +119,7 @@ def run_submission(
                 metadata=metadata,
                 strategic_metadata=strategic_metadata,
                 chunk_size=chunk_size,
+                timeout=timeout,
             )
             return _unchunk_iterator(results_iter, serialization_format)
 
@@ -171,6 +173,7 @@ def blocking_stream_completed_submission(
         submission_id: SubmissionId,
         *,
         serialization_format: SerializationFormat = DEFAULT_SERIALIZATION_FORMAT,
+        timeout: float | None = None,
     ) -> Iterator[Any]:
         """
         Blocks until the submission is completed.
@@ -183,7 +186,7 @@ def blocking_stream_completed_submission(
           (after retrying a consumer kept failing on one of the chunks)
         """
         return _unchunk_iterator(
-            self.blocking_stream_completed_submission_chunks(submission_id),
+            self.blocking_stream_completed_submission_chunks(submission_id, timeout),
             serialization_format,
         )
 
@@ -213,6 +216,7 @@ def run_submission_chunks(
         metadata: None | bytes = None,
         strategic_metadata: None | dict[str, int] = None,
         chunk_size: None | int = None,
+        timeout: float | None = None,
     ) -> Iterator[bytes]:
         """
         Inserts an already-chunked submission into the queue, and blocks until it is completed.
@@ -231,7 +235,7 @@ def run_submission_chunks(
             strategic_metadata=strategic_metadata,
             chunk_size=chunk_size,
         )
-        return self.blocking_stream_completed_submission_chunks(submission_id)
+        return self.blocking_stream_completed_submission_chunks(submission_id, timeout)
 
     async def async_run_submission_chunks(
         self,
@@ -280,7 +284,9 @@ def insert_submission_chunks(
         )
 
     def blocking_stream_completed_submission_chunks(
-        self, submission_id: SubmissionId
+        self,
+        submission_id: SubmissionId,
+        timeout: float | None = None,
     ) -> Iterator[bytes]:
         """
         Blocks until the submission is completed, and returns an iterator that lazily
@@ -291,7 +297,7 @@ def blocking_stream_completed_submission_chunks(
         - `SubmissionFailedError` if the submission failed permanently
           (after retrying a consumer kept failing on one of the chunks)
         """
-        return self.inner.blocking_stream_completed_submission_chunks(submission_id)  # type: ignore[no-any-return]
+        return self.inner.blocking_stream_completed_submission_chunks(submission_id, timeout)  # type: ignore[no-any-return]
 
     async def async_stream_completed_submission_chunks(
         self, submission_id: SubmissionId

libs/opsqueue_python/src/errors.rs

@@ -6,7 +6,7 @@ use opsqueue::common::errors::{
     E, IncorrectUsage, SubmissionNotCancellable, SubmissionNotFound, TooManyMatchingSubmissions,
     UnexpectedOpsqueueConsumerServerResponse,
 };
-use pyo3::exceptions::PyBaseException;
+use pyo3::exceptions::{PyBaseException, PyTimeoutError};
 use pyo3::{Bound, PyErr, Python, import_exception};
 
 use crate::common;
@@ -182,6 +182,12 @@ impl From<CError<UnexpectedOpsqueueConsumerServerResponse>> for PyErr {
     }
 }
 
+impl From<CError<tokio::time::error::Elapsed>> for PyErr {
+    fn from(_value: CError<tokio::time::error::Elapsed>) -> Self {
+        PyTimeoutError::new_err("timeout was reached")
+    }
+}
+
 impl<T> From<PyErr> for CError<E<FatalPythonException, T>> {
     fn from(value: PyErr) -> Self {
         CError(E::L(FatalPythonException(value)))

libs/opsqueue_python/src/producer.rs

@@ -18,6 +18,7 @@ use opsqueue::{
     producer::client::{Client as ActualClient, InternalProducerClientError},
     tracing::CarrierMap,
 };
+use tokio::time::error::Elapsed;
 use ux::u63;
 
 use crate::{
@@ -390,52 +391,6 @@ impl ProducerClient {
         })
     }
 
-    #[pyo3(signature = (chunk_contents, metadata=None, strategic_metadata=None, chunk_size=None, otel_trace_carrier=CarrierMap::default()))]
-    #[allow(clippy::result_large_err, clippy::type_complexity)]
-    pub fn run_submission_chunks(
-        &self,
-        py: Python<'_>,
-        chunk_contents: Py<PyIterator>,
-        metadata: Option<submission::Metadata>,
-        strategic_metadata: Option<StrategicMetadataMap>,
-        chunk_size: Option<i64>,
-        otel_trace_carrier: CarrierMap,
-    ) -> CPyResult<
-        PyChunksIter,
-        E![
-            FatalPythonException,
-            errors::SubmissionFailed,
-            ChunksStorageError,
-            InternalProducerClientError,
-        ],
-    > {
-        let submission_id = self
-            .insert_submission_chunks(
-                py,
-                chunk_contents,
-                metadata,
-                strategic_metadata,
-                chunk_size,
-                otel_trace_carrier,
-            )
-            .map_err(|CError(e)| {
-                CError(match e {
-                    L(e) => L(e),
-                    R(e) => R(R(e)),
-                })
-            })?;
-        let res = self
-            .blocking_stream_completed_submission_chunks(py, submission_id)
-            .map_err(|CError(e)| {
-                CError(match e {
-                    L(e) => L(e),
-                    R(L(e)) => R(L(e)),
-                    R(R(e)) => R(R(R(e))),
-                })
-            })?;
-        Ok(res)
-    }
-
     /// Blocks (and short-polls) until the submission is completed.
     ///
     /// We start with a small short-polling interval
@@ -447,17 +402,34 @@ impl ProducerClient {
         &self,
         py: Python<'_>,
         submission_id: SubmissionId,
+        timeout: Option<f64>,
     ) -> CPyResult<
         PyChunksIter,
         E![
             FatalPythonException,
+            Elapsed,
             errors::SubmissionFailed,
             InternalProducerClientError
         ],
     > {
         py.detach(|| {
             self.block_unless_interrupted(async move {
-                self.stream_completed_submission_chunks(submission_id).await
+                let fut = self.stream_completed_submission_chunks(submission_id);
+                match timeout {
+                    Some(duration) => tokio::time::timeout(Duration::from_secs_f64(duration), fut)
+                        .await
+                        .map_err(|err| CError(R(L(err))))
+                        .and_then(|err| {
+                            err.map_err(|err| match err.0 {
+                                L(err) => CError(L(err)),
+                                R(err) => CError(R(R(err))),
+                            })
+                        }),
+                    None => fut.await.map_err(|err| match err.0 {
+                        L(err) => CError(L(err)),
+                        R(err) => CError(R(R(err))),
+                    }),
+                }
             })
         })
     }

libs/opsqueue_python/tests/test_roundtrip.py

@@ -29,9 +29,11 @@
 )
 import logging
 import multiprocess  # type: ignore[import-untyped]
-
+import time
 import pytest
 
+SUBMISSION_COMPLETED_TIMEOUT = 10.0
+
 
 def increment(data: int) -> int:
     return data + 1
@@ -57,7 +59,10 @@ def run_consumer() -> None:
         input_iter = range(0, 100)
 
         output_iter: Iterator[int] = producer_client.run_submission(
-            input_iter, chunk_size=20, strategic_metadata={"id": 42}
+            input_iter,
+            chunk_size=20,
+            strategic_metadata={"id": 42},
+            timeout=SUBMISSION_COMPLETED_TIMEOUT,
         )
         res = sum(output_iter)
 
@@ -127,6 +132,7 @@ def run_consumer(_consumer_id: int) -> None:
             input_iter,
             chunk_size=chunk_size,
             strategic_metadata={"id": 42, "second_id": 69},
+            timeout=SUBMISSION_COMPLETED_TIMEOUT,
         )
         res = sum(output_iter)
 
@@ -144,7 +150,11 @@ 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,
+        timeout=SUBMISSION_COMPLETED_TIMEOUT,
+    )
     res = sum(output_iter)
     assert res == 0
 
@@ -179,7 +189,10 @@ def run_consumer() -> None:
         input_iter = range(0, 100)
 
         output_iter: Iterator[int] = producer_client.run_submission(
-            input_iter, chunk_size=20, serialization_format=serialization_format
+            input_iter,
+            chunk_size=20,
+            serialization_format=serialization_format,
+            timeout=SUBMISSION_COMPLETED_TIMEOUT,
         )
         res = sum(output_iter)
 
@@ -222,7 +235,11 @@ def broken_increment(input: int) -> float:
         input_iter = range(0, 100)
 
         with pytest.raises(SubmissionFailedError) as exc_info:
-            producer_client.run_submission(input_iter, chunk_size=20)
+            producer_client.run_submission(
+                input_iter,
+                chunk_size=20,
+                timeout=SUBMISSION_COMPLETED_TIMEOUT,
+            )
 
         # We expect the intended attributes to be there:
         assert isinstance(exc_info.value.failure, str)
@@ -262,7 +279,10 @@ def increment_list(ints: Sequence[int], _chunk: Chunk) -> Sequence[int]:
         input_iter = map(lambda i: cbor2.dumps([i, i, i]), range(0, 10))
         output_iter: Iterator[list[int]] = map(
             lambda c: cbor2.loads(c),
-            producer_client.run_submission_chunks(input_iter),
+            producer_client.run_submission_chunks(
+                input_iter,
+                timeout=SUBMISSION_COMPLETED_TIMEOUT,
+            ),
         )
         import itertools
 
@@ -300,7 +320,11 @@ 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,
+            timeout=SUBMISSION_COMPLETED_TIMEOUT,
+        )
         res = sum(output_iter)
 
         assert res == sum(range(1, 1001))
@@ -372,7 +396,10 @@ def run_consumer() -> None:
 
     with background_process(run_consumer):
         # Wait for the submission to complete.
-        producer_client.blocking_stream_completed_submission(submission_id)
+        producer_client.blocking_stream_completed_submission(
+            submission_id,
+            timeout=SUBMISSION_COMPLETED_TIMEOUT,
+        )
         submission = producer_client.get_submission_status(submission_id)
         assert submission is not None
         assert isinstance(submission.submission, SubmissionCompleted)
@@ -416,7 +443,10 @@ def assert_submission_failed_has_metadata(x: SubmissionFailed) -> None:
 
         with pytest.raises(SubmissionFailedError) as exc_info:
             # Wait for the submission to fail.
-            producer_client.blocking_stream_completed_submission(submission_id)
+            producer_client.blocking_stream_completed_submission(
+                submission_id,
+                timeout=SUBMISSION_COMPLETED_TIMEOUT,
+            )
         assert_submission_failed_has_metadata(exc_info.value.submission)
 
         submission = producer_client.get_submission_status(submission_id)
@@ -504,7 +534,10 @@ def run_consumer() -> None:
 
     with background_process(run_consumer):
         # Wait for the submission to complete.
-        producer_client.blocking_stream_completed_submission(submission_id)
+        producer_client.blocking_stream_completed_submission(
+            submission_id,
+            timeout=SUBMISSION_COMPLETED_TIMEOUT,
+        )
         submission = producer_client.get_submission_status(submission_id)
         assert submission is not None
         assert isinstance(submission.submission, SubmissionCompleted)
@@ -537,7 +570,10 @@ def consume(x: int) -> None:
 
     with background_process(run_consumer):
         with pytest.raises(SubmissionFailedError):
-            producer_client.blocking_stream_completed_submission(submission_id)
+            producer_client.blocking_stream_completed_submission(
+                submission_id,
+                timeout=SUBMISSION_COMPLETED_TIMEOUT,
+            )
         # Cancelling the failed submission should fail.
         with pytest.raises(SubmissionNotCancellableError) as exc_info:
             producer_client.cancel_submission(submission_id)
@@ -567,7 +603,10 @@ def consume(x: int) -> int | None:
 
     with background_process(run_consumer):
         with pytest.raises(SubmissionFailedError) as exc_info:
-            producer_client.blocking_stream_completed_submission(submission_id)
+            producer_client.blocking_stream_completed_submission(
+                submission_id,
+                timeout=SUBMISSION_COMPLETED_TIMEOUT,
+            )
         assert exc_info.value.submission.chunks_done == len(chunks) - 1
 
 
@@ -695,7 +734,10 @@ def run_consumer() -> None:
         consumer_client.run_each_op(lambda x: x)
 
     with background_process(run_consumer):
-        producer_client.blocking_stream_completed_submission(submission_id)
+        producer_client.blocking_stream_completed_submission(
+            submission_id,
+            timeout=SUBMISSION_COMPLETED_TIMEOUT,
+        )
         assert isinstance(
             producer_client.get_submission_status(submission_id),
             SubmissionStatus.Completed,
@@ -744,7 +786,10 @@ def process_op(x: int) -> int:
         # Release the consumer.
         can_continue.set()
 
-        producer_client.blocking_stream_completed_submission(submission_id)
+        producer_client.blocking_stream_completed_submission(
+            submission_id,
+            timeout=SUBMISSION_COMPLETED_TIMEOUT,
+        )
         assert isinstance(
             producer_client.get_submission_status(submission_id),
             SubmissionStatus.Completed,
@@ -778,3 +823,25 @@ def test_cancel_paused(opsqueue: OpsqueueProcess) -> None:
     assert isinstance(
         producer_client.get_submission_status(submission_id), SubmissionStatus.Cancelled
     )
+
+
+def test_run_submission_timeout(opsqueue: OpsqueueProcess) -> None:
+    url = "file:///tmp/opsqueue/test_run_submission_timeout"
+    producer_client = ProducerClient(f"localhost:{opsqueue.port}", url)
+
+    def run_consumer() -> None:
+        consumer_client = ConsumerClient(f"localhost:{opsqueue.port}", url)
+
+        def process_op(x: int) -> int:
+            time.sleep(2.0)
+            return x
+
+        consumer_client.run_each_op(process_op)
+
+    with background_process(run_consumer) as _consumer:
+        with pytest.raises(TimeoutError):
+            producer_client.run_submission(
+                [1],
+                chunk_size=1,
+                timeout=0.1,
+            )