1cbyc commented on code in PR #2608:
URL:
https://github.com/apache/datafusion-ballista/pull/2608#discussion_r4432018358
##########
ballista/core/src/client.rs:
##########
@@ -62,6 +64,45 @@ pub struct BallistaClient {
}
impl BallistaClient {
+ fn encode_action(action: &Action) -> BResult<Vec<u8>> {
Review Comment:
Thanks. I kept the existing `encode_protobuf` path unchanged so this PR
stays focused on consolidating retry behavior. The `Bytes` conversion is a
useful allocation cleanup, but separate from this fix.
##########
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:
Done. `with_retry` now lives in `utils.rs`, is generic over the error type,
and is shared by the client, shuffle, and pipelined paths. The client converts
total attempts to retries through `retry_budget()` with `saturating_sub(1)`.
##########
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:
Done. `retry_budget()` now documents both ways zero can appear: an explicit
zero and the current `u16` to `u8` config conversion.
##########
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:
Done. The shared helper sleeps only in the retry arm and logs each retry at
`warn` with the triggering error.
##########
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:
Done. The zero-budget test now uses `Status::unknown`, asserts exactly one
call, and checks the returned `bad-1` message.
##########
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:
Done. The exhaustion test configures three attempts, asserts exactly three
calls, and verifies that the final error is returned.
--
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]
