1cbyc commented on code in PR #2608:
URL:
https://github.com/apache/datafusion-ballista/pull/2608#discussion_r4432019106
##########
ballista/core/src/client.rs:
##########
@@ -720,7 +684,73 @@ mod tests {
use futures::{StreamExt, TryStreamExt};
use prost::bytes::Bytes;
- use crate::client::BlockDataStream;
+ use crate::client::{BallistaClient, BlockDataStream};
+
+ #[tokio::test]
+ async fn retry_io_attempts_once_when_retries_are_zero() {
+ let mut client = BallistaClient::new_for_test("localhost", 50050);
+ client.io_retries_times = 0;
+ client.io_retry_wait_time_ms = 0;
+ let attempts = Arc::new(AtomicUsize::new(0));
+ let observed = Arc::clone(&attempts);
+
+ let error = client
+ .retry_io(move |_| {
+ observed.fetch_add(1, Ordering::Relaxed);
+ async { Err::<(), _>(tonic::Status::unavailable("offline")) }
+ })
+ .await
+ .unwrap_err();
+
+ assert_eq!(error.code(), tonic::Code::Unavailable);
+ assert_eq!(attempts.load(Ordering::Relaxed), 1);
+ }
+
+ #[tokio::test]
+ async fn retry_io_retries_unknown_errors() {
+ let mut client = BallistaClient::new_for_test("localhost", 50050);
+ client.io_retries_times = 3;
+ client.io_retry_wait_time_ms = 0;
+ let attempts = Arc::new(AtomicUsize::new(0));
+ let observed = Arc::clone(&attempts);
+
+ let result = client
+ .retry_io(move |_| {
+ let attempt = observed.fetch_add(1, Ordering::Relaxed);
+ async move {
+ if attempt < 2 {
+ Err(tonic::Status::unknown("transient"))
+ } else {
+ Ok("done")
+ }
+ }
+ })
+ .await
+ .unwrap();
+
+ assert_eq!(result, "done");
+ assert_eq!(attempts.load(Ordering::Relaxed), 3);
+ }
+
+ #[tokio::test]
+ async fn retry_io_does_not_retry_other_statuses() {
+ let mut client = BallistaClient::new_for_test("localhost", 50050);
+ client.io_retries_times = 3;
+ client.io_retry_wait_time_ms = 0;
+ let attempts = Arc::new(AtomicUsize::new(0));
+ let observed = Arc::clone(&attempts);
+
+ let error = client
+ .retry_io(move |_| {
+ observed.fetch_add(1, Ordering::Relaxed);
+ async { Err::<(), _>(tonic::Status::invalid_argument("bad")) }
+ })
+ .await
+ .unwrap_err();
+
+ assert_eq!(error.code(), tonic::Code::InvalidArgument);
Review Comment:
Done. The retry-budget tests now assert the returned messages (`bad-1` and
`bad-3`), not only the status code.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
