NoahKusaba commented on code in PR #2608:
URL:
https://github.com/apache/datafusion-ballista/pull/2608#discussion_r4355635620
##########
ballista/core/src/client.rs:
##########
@@ -62,6 +64,45 @@ pub struct BallistaClient {
}
impl BallistaClient {
+ fn encode_action(action: &Action) -> BResult<Vec<u8>> {
+ let serialized_action: protobuf::Action =
action.to_owned().try_into()?;
+ let mut buf = Vec::with_capacity(serialized_action.encoded_len());
+ serialized_action
+ .encode(&mut buf)
+ .map_err(|e| BallistaError::GrpcActionError(format!("{e:?}")))?;
+ Ok(buf)
+ }
+
+ async fn retry_io<T, F, Fut>(&self, mut op: F) -> Result<T, Status>
Review Comment:
There's already a generic `with_retry(retries, wait_ms, should_retry, f)` in
`shuffle_reader.rs:945` with the same shape. Could we make it generic over the
error type, move it to `utils.rs`, and call it here with `|s: &Status| s.code()
== Code::Unknown`? Then the crate has one retry helper, not two. One catch:
`with_retry` counts *retries* and the client counts *total attempts*, so the
call would pass `self.io_retries_times.saturating_sub(1)`. Happy to leave this
as a follow-up if you'd rather keep this PR small.
##########
ballista/core/src/client.rs:
##########
@@ -62,6 +64,45 @@ pub struct BallistaClient {
}
impl BallistaClient {
+ fn encode_action(action: &Action) -> BResult<Vec<u8>> {
Review Comment:
Nit: encoding into a `Vec` can't fail, so the `map_err` here is dead code.
`encode_to_vec()` is already used in `serde/mod.rs`. Returning `Bytes` also
makes the per-attempt `buf.clone()` a refcount bump instead of a copy:
```rust
let proto: protobuf::Action = action.clone().try_into()?;
Ok(proto.encode_to_vec().into())
```
##########
ballista/core/src/client.rs:
##########
@@ -62,6 +64,45 @@ pub struct BallistaClient {
}
impl BallistaClient {
+ fn encode_action(action: &Action) -> BResult<Vec<u8>> {
+ let serialized_action: protobuf::Action =
action.to_owned().try_into()?;
+ let mut buf = Vec::with_capacity(serialized_action.encoded_len());
+ serialized_action
+ .encode(&mut buf)
+ .map_err(|e| BallistaError::GrpcActionError(format!("{e:?}")))?;
+ Ok(buf)
+ }
+
+ async fn retry_io<T, F, Fut>(&self, mut op: F) -> Result<T, Status>
+ where
+ F: FnMut(FlightServiceClient<Channel>) -> Fut,
+ Fut: Future<Output = Result<T, Status>>,
+ {
+ let attempts = self.io_retries_times.max(1);
Review Comment:
Could you add a short comment here explaining how 0 can show up? Besides an
explicit `0`, the config entry is `UInt16` but `utils.rs:99` casts it with `as
u8`, so 256 truncates to 0 (and 300 becomes 44). Something like `//
io_retries_times counts total attempts; 0 can come from config or from u16→u8
truncation, so always make at least one attempt.` That way the next reader
knows why the clamp is there.
##########
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:
Nit: could the tests compare the message too (e.g.
`assert_eq!(error.message(), "bad")`)? Checking only the code would still pass
if a different error with the same code came back.
##########
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() {
Review Comment:
Could we add a test where the op always returns `Unknown`, with
`io_retries_times = 3`? It should run exactly 3 times and return the last
error. `attempt + 1 < attempts` is where an off-by-one would hide, and right
now nothing covers it.
##########
ballista/core/src/client.rs:
##########
@@ -62,6 +64,45 @@ pub struct BallistaClient {
}
impl BallistaClient {
+ fn encode_action(action: &Action) -> BResult<Vec<u8>> {
+ let serialized_action: protobuf::Action =
action.to_owned().try_into()?;
+ let mut buf = Vec::with_capacity(serialized_action.encoded_len());
+ serialized_action
+ .encode(&mut buf)
+ .map_err(|e| BallistaError::GrpcActionError(format!("{e:?}")))?;
+ Ok(buf)
+ }
+
+ async fn retry_io<T, F, Fut>(&self, mut op: F) -> Result<T, Status>
+ where
+ F: FnMut(FlightServiceClient<Channel>) -> Fut,
+ Fut: Future<Output = Result<T, Status>>,
+ {
+ let attempts = self.io_retries_times.max(1);
+ let mut attempt = 0;
+
+ loop {
+ if attempt > 0 {
Review Comment:
Nit: if you keep `retry_io`, the sleep could move into the retry arm. That
drops the `if attempt > 0` block and merges the two return arms into one:
```rust
let mut attempt = 1;
loop {
match op(self.flight_client.clone()).await {
Err(e) if e.code() == Code::Unknown && attempt < attempts => {
warn!("Remote shuffle read failed: {e}; retry
{attempt}/{attempts} in {} ms", self.io_retry_wait_time_ms);
tokio::time::sleep(Duration::from_millis(self.io_retry_wait_time_ms)).await;
attempt += 1;
}
result => return result,
}
}
```
It also lets the warning log the error that caused the retry, which it
doesn't today.
##########
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")) }
Review Comment:
`Unavailable` is never retried, so this test passes no matter what
`io_retries_times` is. And since the code uses `loop`, it would also pass
without `.max(1)`. Could this use `Status::unknown(...)` so it actually checks
that zero means one attempt?
--
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]
