fix(ballista): recover from transient shuffle-fetch connection failures (#36)

* fix(error): surface buried FetchFailed from wrapped DataFusionError

A BallistaError::FetchFailed bubbled through the DataFusion stream chain
on a shared shuffle stream typically arrives as

  DataFusionError::Shared(Arc(DataFusionError::ArrowError(
    ArrowError::ExternalError(BallistaError::FetchFailed(...)))))

The existing From<DataFusionError> for BallistaError only unwrapped
ArrowError, so the inner FetchFailed stayed buried as
BallistaError::DataFusionError(...). FailedTask::from then fell through
its pattern match to the catch-all arm and emitted
FailedReason::ExecutionError (non-retryable), bypassing the schedulers
FetchPartitionError recovery path that reruns the offending map stage.

Drill through Shared, Context, Diagnostic, Collection, and External so
the FetchFailed surfaces and the scheduler can recover. For Shared with
multiple strong references where Arc::try_unwrap fails, fall back to a
borrowed walk that reconstructs FetchFailed from its cloneable fields.

* fix(client): retry transient shuffle-fetch connects and apply gRPC transport defaults

BallistaClient::try_new previously called endpoint.connect() once with no
client-side connect_timeout, tcp_keepalive, or http2 keepalive (a None
GrpcClientConfig was passed to create_grpc_client_endpoint), and a
single transport-level failure propagated straight to FetchFailed
without any retry. The IO_RETRIES_TIMES loop in execute_do_get /
execute_do_action only applies *after* the channel has been
established, so connect-time failures got no second chance — a single
TCP RST during a shuffle connection storm against the same executor
was enough to fail a task.

Match the scheduler flight proxy: pass a real GrpcClientConfig::default
to create_grpc_client_endpoint, and wrap endpoint.connect() in the same
3-attempt / 3s-backoff retry budget used by the post-connect retry
loops. After exhausting the budget the error message reports the
attempt count so it stays distinguishable from a permanent failure.

Author: phillipleblanc

ballista/core/src/client.rs

@@ -47,7 +47,7 @@ use datafusion::error::Result;
 use crate::extension::BallistaConfigGrpcEndpoint;
 use crate::serde::protobuf;
 
-use crate::utils::create_grpc_client_endpoint;
+use crate::utils::{GrpcClientConfig, create_grpc_client_endpoint};
 
 use datafusion::physical_plan::{RecordBatchStream, SendableRecordBatchStream};
 use futures::{Stream, StreamExt};
@@ -80,7 +80,14 @@ impl BallistaClient {
         let addr = format!("{scheme}://{host}:{port}");
         debug!("BallistaClient connecting to {addr}");
 
-        let mut endpoint = create_grpc_client_endpoint(addr.clone(), None)
+        // Apply the same transport defaults (connect_timeout, tcp_keepalive,
+        // http2 keepalive) the scheduler's flight proxy uses for its own
+        // executor connections. Without this the shuffle client gets a bare
+        // tonic Endpoint with no timeouts and no keepalive — a single bad peer
+        // can hang the fetch indefinitely, and a half-open connection isn't
+        // detected.
+        let grpc_config = GrpcClientConfig::default();
+        let mut endpoint = create_grpc_client_endpoint(addr.clone(), Some(&grpc_config))
             .map_err(|e| {
                 BallistaError::GrpcConnectionError(format!(
                     "Error creating endpoint to Ballista scheduler or executor at {addr}: {e:?}"
@@ -97,11 +104,42 @@ impl BallistaClient {
                 })?;
         }
 
-        let connection = endpoint.connect().await.map_err(|e| {
-            BallistaError::GrpcConnectionError(format!(
-                "Error connecting to Ballista scheduler or executor at {addr}: {e:?}"
-            ))
-        })?;
+        // Retry transient connect failures with the same budget the do_get /
+        // do_action paths below use. A single TCP RST during connection setup
+        // (typically when many shuffle clients dial the same executor at once)
+        // would otherwise propagate straight to FetchFailed without any retry,
+        // since the IO_RETRIES_TIMES loop in execute_do_get only applies after
+        // the channel is established.
+        let mut last_err: Option<tonic::transport::Error> = None;
+        let mut connection = None;
+        for attempt in 0..IO_RETRIES_TIMES {
+            if attempt > 0 {
+                warn!(
+                    "Connection to {addr} failed (attempt {attempt}/{IO_RETRIES_TIMES}), sleeping {IO_RETRY_WAIT_TIME_MS} ms before retry"
+                );
+                tokio::time::sleep(std::time::Duration::from_millis(
+                    IO_RETRY_WAIT_TIME_MS,
+                ))
+                .await;
+            }
+            match endpoint.connect().await {
+                Ok(c) => {
+                    connection = Some(c);
+                    break;
+                }
+                Err(e) => last_err = Some(e),
+            }
+        }
+        let connection = match connection {
+            Some(c) => c,
+            None => {
+                return Err(BallistaError::GrpcConnectionError(format!(
+                    "Error connecting to Ballista scheduler or executor at {addr} after {IO_RETRIES_TIMES} attempts: {:?}",
+                    last_err
+                        .expect("at least one attempt failed when connection is None")
+                )));
+            }
+        };
 
         let flight_client = FlightServiceClient::new(connection)
             .max_decoding_message_size(max_message_size)

ballista/core/src/error.rs

@@ -21,6 +21,7 @@ use std::{
     error::Error,
     fmt::{Display, Formatter},
     io, result,
+    sync::Arc,
 };
 
 use crate::serde::protobuf::failed_task::FailedReason;
@@ -113,13 +114,88 @@ impl From<parser::ParserError> for BallistaError {
 
 impl From<DataFusionError> for BallistaError {
     fn from(e: DataFusionError) -> Self {
+        // BallistaError::FetchFailed must reach the top level so FailedTask::from
+        // routes it to FailedReason::FetchPartitionError (the scheduler's
+        // map-stage rerun path). When a stream is shared between consumers,
+        // DataFusion wraps the underlying error as
+        //   DataFusionError::Shared(Arc(DataFusionError::ArrowError(
+        //     ArrowError::ExternalError(BallistaError::FetchFailed(...)))))
+        // and similar wrappers appear for Context, Diagnostic, Collection, and
+        // External. Drill through them so the inner BallistaError surfaces.
         match e {
             DataFusionError::ArrowError(e, _) => Self::from(*e),
+            DataFusionError::External(e) => match e.downcast::<BallistaError>() {
+                Ok(b) => *b,
+                Err(e) => match e.downcast::<DataFusionError>() {
+                    Ok(d) => Self::from(*d),
+                    Err(e) => BallistaError::DataFusionError(Box::new(
+                        DataFusionError::External(e),
+                    )),
+                },
+            },
+            DataFusionError::Context(_, inner) => Self::from(*inner),
+            DataFusionError::Diagnostic(_, inner) => Self::from(*inner),
+            DataFusionError::Collection(mut errs) if !errs.is_empty() => {
+                Self::from(errs.swap_remove(0))
+            }
+            DataFusionError::Shared(arc) => match Arc::try_unwrap(arc) {
+                Ok(inner) => Self::from(inner),
+                Err(arc) => find_fetch_failed(arc.as_ref()).unwrap_or_else(|| {
+                    BallistaError::DataFusionError(Box::new(DataFusionError::Shared(arc)))
+                }),
+            },
             _ => BallistaError::DataFusionError(Box::new(e)),
         }
     }
 }
 
+/// Walk a borrowed [`DataFusionError`] chain looking for a buried
+/// [`BallistaError::FetchFailed`], reconstructing it from its cloneable fields.
+///
+/// Used when ownership of the inner error cannot be taken — e.g. for
+/// [`DataFusionError::Shared`] with multiple strong references — so the
+/// classification path in `From<BallistaError> for FailedTask` still sees the
+/// FetchFailed variant.
+fn find_fetch_failed(e: &DataFusionError) -> Option<BallistaError> {
+    match e {
+        DataFusionError::ArrowError(arrow, _) => find_fetch_failed_in_arrow(arrow),
+        DataFusionError::External(err) => find_fetch_failed_in_dyn(err.as_ref()),
+        DataFusionError::Context(_, inner) | DataFusionError::Diagnostic(_, inner) => {
+            find_fetch_failed(inner)
+        }
+        DataFusionError::Collection(errs) => errs.iter().find_map(find_fetch_failed),
+        DataFusionError::Shared(arc) => find_fetch_failed(arc.as_ref()),
+        _ => None,
+    }
+}
+
+fn find_fetch_failed_in_arrow(e: &ArrowError) -> Option<BallistaError> {
+    match e {
+        ArrowError::ExternalError(err) => find_fetch_failed_in_dyn(err.as_ref()),
+        _ => None,
+    }
+}
+
+fn find_fetch_failed_in_dyn(err: &(dyn Error + 'static)) -> Option<BallistaError> {
+    if let Some(BallistaError::FetchFailed(executor_id, map_stage, map_partition, desc)) =
+        err.downcast_ref::<BallistaError>()
+    {
+        return Some(BallistaError::FetchFailed(
+            executor_id.clone(),
+            *map_stage,
+            *map_partition,
+            desc.clone(),
+        ));
+    }
+    if let Some(df) = err.downcast_ref::<DataFusionError>() {
+        return find_fetch_failed(df);
+    }
+    if let Some(arrow) = err.downcast_ref::<ArrowError>() {
+        return find_fetch_failed_in_arrow(arrow);
+    }
+    None
+}
+
 impl From<io::Error> for BallistaError {
     fn from(e: io::Error) -> Self {
         BallistaError::IoError(e)
@@ -256,3 +332,90 @@ impl From<BallistaError> for FailedTask {
 }
 
 impl Error for BallistaError {}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn fetch_failed() -> BallistaError {
+        BallistaError::FetchFailed(
+            "executor-1:50052".to_string(),
+            1,
+            74,
+            "Error connecting to Ballista scheduler or executor".to_string(),
+        )
+    }
+
+    fn assert_routes_to_fetch_partition_error(err: BallistaError) {
+        let failed: FailedTask = err.into();
+        match failed.failed_reason {
+            Some(FailedReason::FetchPartitionError(fp)) => {
+                assert_eq!(fp.executor_id, "executor-1:50052");
+                assert_eq!(fp.map_stage_id, 1);
+                assert_eq!(fp.map_partition_id, 74);
+            }
+            other => panic!(
+                "expected FetchPartitionError, got {other:?} (error: {})",
+                failed.error
+            ),
+        }
+    }
+
+    #[test]
+    fn fetch_failed_through_arrow_external_routes_to_fetch_partition_error() {
+        let arrow = ArrowError::ExternalError(Box::new(fetch_failed()));
+        let df = DataFusionError::ArrowError(Box::new(arrow), Some(String::new()));
+        assert_routes_to_fetch_partition_error(BallistaError::from(df));
+    }
+
+    #[test]
+    fn fetch_failed_through_shared_routes_to_fetch_partition_error() {
+        // Reproduces the production wrapping from a shared shuffle stream:
+        //   DataFusionError::Shared(Arc(ArrowError(ExternalError(FetchFailed))))
+        let arrow = ArrowError::ExternalError(Box::new(fetch_failed()));
+        let inner = DataFusionError::ArrowError(Box::new(arrow), Some(String::new()));
+        let shared = DataFusionError::Shared(Arc::new(inner));
+        assert_routes_to_fetch_partition_error(BallistaError::from(shared));
+    }
+
+    #[test]
+    fn fetch_failed_through_shared_with_extra_refs_still_routes() {
+        // When Arc::try_unwrap fails (extra strong refs alive), the borrow-walk
+        // fallback must still surface the inner FetchFailed.
+        let arrow = ArrowError::ExternalError(Box::new(fetch_failed()));
+        let inner = DataFusionError::ArrowError(Box::new(arrow), Some(String::new()));
+        let arc = Arc::new(inner);
+        let extra_ref = Arc::clone(&arc);
+        let shared = DataFusionError::Shared(arc);
+        assert_routes_to_fetch_partition_error(BallistaError::from(shared));
+        drop(extra_ref);
+    }
+
+    #[test]
+    fn fetch_failed_through_context_routes_to_fetch_partition_error() {
+        let arrow = ArrowError::ExternalError(Box::new(fetch_failed()));
+        let inner = DataFusionError::ArrowError(Box::new(arrow), Some(String::new()));
+        let ctx =
+            DataFusionError::Context("reading shuffle".to_string(), Box::new(inner));
+        assert_routes_to_fetch_partition_error(BallistaError::from(ctx));
+    }
+
+    #[test]
+    fn fetch_failed_through_external_routes_to_fetch_partition_error() {
+        let external = DataFusionError::External(Box::new(fetch_failed()));
+        assert_routes_to_fetch_partition_error(BallistaError::from(external));
+    }
+
+    #[test]
+    fn unrelated_shared_error_is_not_misclassified() {
+        let inner = DataFusionError::Plan("planning failed".to_string());
+        let shared = DataFusionError::Shared(Arc::new(inner));
+        let bal = BallistaError::from(shared);
+        let failed: FailedTask = bal.into();
+        assert!(
+            matches!(failed.failed_reason, Some(FailedReason::ExecutionError(_))),
+            "expected ExecutionError, got {:?}",
+            failed.failed_reason
+        );
+    }
+}