Add exponential backoff and log suppression for scheduler disconnection

Author: phillipleblanc

Cargo.lock

@@ -68,6 +68,7 @@ dashmap = { version = "6.1" }
 async-trait = { version = "0.1" }
 serde = { version = "1.0" }
 tokio-stream = { version = "0.1" }
+backoff = { version = "0.4" }
 url = { version = "2.5" }
 
 # cargo build --profile release-lto

Cargo.toml

@@ -830,15 +830,8 @@ pub mod execute_query_params {
     pub enum Query {
         #[prost(bytes, tag = "1")]
         LogicalPlan(::prost::alloc::vec::Vec<u8>),
-
         #[prost(bytes, tag = "6")]
         SubstraitPlan(::prost::alloc::vec::Vec<u8>),
-
-        /// I'd suggest to remove this, if SQL needed use `flight-sql`
-        #[deprecated]
-        #[prost(string, tag = "2")]
-        Sql(::prost::alloc::string::String),
-
     }
 }
 #[derive(Clone, PartialEq, ::prost::Message)]

ballista/core/src/serde/generated/ballista.rs

@@ -35,7 +35,7 @@ use std::sync::Arc;
 use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use std::{fs::File, pin::Pin};
 use tonic::codegen::StdError;
-use tonic::transport::{Channel, Endpoint, Error, Server};
+use tonic::transport::{Channel, Error, Server};
 
 /// Configuration for gRPC client connections.
 ///
@@ -196,21 +196,11 @@ pub async fn collect_stream(
     Ok(batches)
 }
 
-
-
 /// Creates a gRPC client connection with the specified configuration.
-
-
 pub async fn create_grpc_client_connection<D>(
     dst: D,
     config: &GrpcClientConfig,
 ) -> std::result::Result<Channel, Error>
-
-
-pub fn create_grpc_client_endpoint<D>(dst: D) -> std::result::Result<Endpoint, Error>
-
-
-
 where
     D: std::convert::TryInto<tonic::transport::Endpoint>,
     D::Error: Into<StdError>,
@@ -231,6 +221,26 @@ where
     endpoint.connect().await
 }
 
+/// Creates a gRPC client endpoint (returns Endpoint without connecting).
+/// Used for TLS/API key customization before establishing connection.
+pub fn create_grpc_client_endpoint<D>(dst: D) -> std::result::Result<tonic::transport::Endpoint, Error>
+where
+    D: std::convert::TryInto<tonic::transport::Endpoint>,
+    D::Error: Into<StdError>,
+{
+    let endpoint = tonic::transport::Endpoint::new(dst)?
+        .connect_timeout(Duration::from_secs(20))
+        .timeout(Duration::from_secs(20))
+        // Disable Nagle's Algorithm since we don't want packets to wait
+        .tcp_nodelay(true)
+        .tcp_keepalive(Option::Some(Duration::from_secs(3600)))
+        .http2_keep_alive_interval(Duration::from_secs(300))
+        .keep_alive_timeout(Duration::from_secs(20))
+        .keep_alive_while_idle(true);
+
+    Ok(endpoint)
+}
+
 /// Creates a gRPC server builder with the specified configuration.
 pub fn create_grpc_server(config: &GrpcServerConfig) -> Server {
     Server::builder()
@@ -270,24 +280,3 @@ pub fn get_time_before(interval_seconds: u64) -> u64 {
         .unwrap_or_else(|| Duration::from_secs(0))
         .as_secs()
 }
-
-/// Create a gRPC client endpoint with configurable settings
-///
-/// This is used by Spice extensions for customizing gRPC endpoint creation
-pub fn create_grpc_client_endpoint<D>(dst: D) -> std::result::Result<Endpoint, Error>
-where
-    D: std::convert::TryInto<tonic::transport::Endpoint>,
-    D::Error: Into<StdError>,
-{
-    let endpoint = tonic::transport::Endpoint::new(dst)?
-        .connect_timeout(Duration::from_secs(20))
-        .timeout(Duration::from_secs(20))
-        // Disable Nagle's Algorithm since we don't want packets to wait
-        .tcp_nodelay(true)
-        .tcp_keepalive(Option::Some(Duration::from_secs(3600)))
-        .http2_keep_alive_interval(Duration::from_secs(300))
-        .keep_alive_timeout(Duration::from_secs(20))
-        .keep_alive_while_idle(true);
-
-    Ok(endpoint)
-}

ballista/core/src/utils.rs

@@ -40,7 +40,11 @@ default = ["build-binary", "mimalloc"]
 arrow = { workspace = true }
 arrow-flight = { workspace = true }
 async-trait = { workspace = true }
+
 ballista-core = { path = "../core", version = "51.0.0" }
+
+backoff = { workspace = true }
+
 clap = { workspace = true, optional = true }
 dashmap = { workspace = true }
 datafusion = { workspace = true }