[3/3] allow customization of executor grpc endpoint

Author: mach-kernel

ballista/executor/src/config.rs

@@ -147,6 +147,7 @@ impl TryFrom<Config> for ExecutorProcessConfig {
             override_logical_codec: None,
             override_physical_codec: None,
             override_arrow_flight_service: None,
+            override_create_grpc_client_endpoint: None,
         })
     }
 }

ballista/executor/src/executor_process.rs

@@ -56,6 +56,7 @@ use ballista_core::utils::{
     get_time_before,
 };
 use ballista_core::{ConfigProducer, RuntimeProducer, BALLISTA_VERSION};
+use tonic::transport::{Endpoint, Error as TonicTransportError};
 
 use crate::execution_engine::ExecutionEngine;
 use crate::executor::{Executor, TasksDrainedFuture};
@@ -104,6 +105,9 @@ pub struct ExecutorProcessConfig {
     pub override_physical_codec: Option<Arc<dyn PhysicalExtensionCodec>>,
     /// [ArrowFlightServerProvider] implementation override option
     pub override_arrow_flight_service: Option<Arc<ArrowFlightServerProvider>>,
+    /// Override function for customizing gRPC client endpoints before they are used
+    pub override_create_grpc_client_endpoint:
+        Option<Arc<dyn Fn(Endpoint) -> Result<Endpoint, TonicTransportError> + Send + Sync>>,
 }
 
 impl ExecutorProcessConfig {
@@ -147,6 +151,7 @@ impl Default for ExecutorProcessConfig {
             override_logical_codec: None,
             override_physical_codec: None,
             override_arrow_flight_service: None,
+            override_create_grpc_client_endpoint: None,
         }
     }
 }
@@ -246,12 +251,22 @@ pub async fn start_executor_process(
 
     let connect_timeout = opt.scheduler_connect_timeout_seconds as u64;
     let connection = if connect_timeout == 0 {
-        create_grpc_client_endpoint(scheduler_url)
+        let mut endpoint = create_grpc_client_endpoint(scheduler_url)
             .map_err(|_| {
                 BallistaError::GrpcConnectionError(
                     "Could not create endpoint to scheduler".to_string(),
                 )
-            })?
+            })?;
+
+        if let Some(ref override_fn) = opt.override_create_grpc_client_endpoint {
+            endpoint = override_fn(endpoint).map_err(|_| {
+                BallistaError::GrpcConnectionError(
+                    "Failed to apply endpoint override".to_string(),
+                )
+            })?;
+        }
+
+        endpoint
             .connect()
             .await
             .map_err(|_| {
@@ -269,16 +284,31 @@ pub async fn start_executor_process(
             && Instant::now().elapsed().as_secs() - start_time < connect_timeout
         {
             match create_grpc_client_endpoint(scheduler_url.clone()) {
-                Ok(endpoint) => match endpoint.connect().await {
-                    Ok(connection) => {
-                        info!("Connected to scheduler at {scheduler_url}");
-                        x = Some(connection);
+                Ok(mut endpoint) => {
+                    if let Some(ref override_fn) = opt.override_create_grpc_client_endpoint {
+                        match override_fn(endpoint) {
+                            Ok(overridden_endpoint) => endpoint = overridden_endpoint,
+                            Err(e) => {
+                                warn!(
+                                    "Failed to apply endpoint override to scheduler at {scheduler_url} ({e}); retrying ..."
+                                );
+                                tokio::time::sleep(time::Duration::from_millis(500)).await;
+                                continue;
+                            }
+                        }
                     }
-                    Err(e) => {
-                        warn!(
-                            "Failed to connect to scheduler at {scheduler_url} ({e}); retrying ..."
-                        );
-                        tokio::time::sleep(time::Duration::from_millis(500)).await;
+
+                    match endpoint.connect().await {
+                        Ok(connection) => {
+                            info!("Connected to scheduler at {scheduler_url}");
+                            x = Some(connection);
+                        }
+                        Err(e) => {
+                            warn!(
+                                "Failed to connect to scheduler at {scheduler_url} ({e}); retrying ..."
+                            );
+                            tokio::time::sleep(time::Duration::from_millis(500)).await;
+                        }
                     }
                 },
                 Err(e) => {

ballista/executor/src/executor_server.rs

@@ -24,7 +24,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use tokio::sync::mpsc;
 
 use log::{debug, error, info, warn};
-use tonic::transport::Channel;
+use tonic::transport::{Channel, Endpoint, Error as TonicTransportError};
 use tonic::{Request, Response, Status};
 
 use ballista_core::error::BallistaError;
@@ -97,6 +97,7 @@ pub async fn startup<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
         codec,
         config.grpc_max_encoding_message_size as usize,
         config.grpc_max_decoding_message_size as usize,
+        config.override_create_grpc_client_endpoint.clone(),
     );
 
     // 1. Start executor grpc service
@@ -183,6 +184,8 @@ pub struct ExecutorServer<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPl
     schedulers: SchedulerClients,
     grpc_max_encoding_message_size: usize,
     grpc_max_decoding_message_size: usize,
+    override_create_grpc_client_endpoint:
+        Option<Arc<dyn Fn(Endpoint) -> Result<Endpoint, TonicTransportError> + Send + Sync>>,
 }
 
 #[derive(Clone)]
@@ -209,6 +212,9 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
         codec: BallistaCodec<T, U>,
         grpc_max_encoding_message_size: usize,
         grpc_max_decoding_message_size: usize,
+        override_create_grpc_client_endpoint: Option<
+            Arc<dyn Fn(Endpoint) -> Result<Endpoint, TonicTransportError> + Send + Sync>,
+        >,
     ) -> Self {
         Self {
             _start_time: SystemTime::now()
@@ -222,6 +228,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
             schedulers: Default::default(),
             grpc_max_encoding_message_size,
             grpc_max_decoding_message_size,
+            override_create_grpc_client_endpoint,
         }
     }
 
@@ -235,9 +242,13 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
             Ok(scheduler)
         } else {
             let scheduler_url = format!("http://{scheduler_id}");
-            let connection = create_grpc_client_endpoint(scheduler_url)?
-                .connect()
-                .await?;
+            let mut endpoint = create_grpc_client_endpoint(scheduler_url)?;
+
+            if let Some(ref override_fn) = self.override_create_grpc_client_endpoint {
+                endpoint = override_fn(endpoint)?;
+            }
+
+            let connection = endpoint.connect().await?;
             let scheduler = SchedulerGrpcClient::new(connection)
                 .max_encoding_message_size(self.grpc_max_encoding_message_size)
                 .max_decoding_message_size(self.grpc_max_decoding_message_size);