[2/3] allow customization of scheduler grpc endpoint

Author: mach-kernel

ballista/scheduler/src/config.rs

@@ -32,6 +32,7 @@ use datafusion_proto::logical_plan::LogicalExtensionCodec;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use std::fmt::Display;
 use std::sync::Arc;
+use tonic::transport::{Endpoint, Error as TonicTransportError};
 
 /// Configuration of the application
 #[cfg(feature = "build-binary")]
@@ -215,6 +216,9 @@ pub struct SchedulerConfig {
     pub override_logical_codec: Option<Arc<dyn LogicalExtensionCodec>>,
     /// [PhysicalExtensionCodec] override option
     pub override_physical_codec: Option<Arc<dyn PhysicalExtensionCodec>>,
+    /// 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 Default for SchedulerConfig {
@@ -241,6 +245,7 @@ impl Default for SchedulerConfig {
             override_session_builder: None,
             override_logical_codec: None,
             override_physical_codec: None,
+            override_create_grpc_client_endpoint: None,
         }
     }
 }
@@ -343,6 +348,14 @@ impl SchedulerConfig {
         self.override_session_builder = Some(override_session_builder);
         self
     }
+
+    pub fn with_override_create_grpc_client_endpoint(
+        mut self,
+        override_fn: Arc<dyn Fn(Endpoint) -> Result<Endpoint, TonicTransportError> + Send + Sync>,
+    ) -> Self {
+        self.override_create_grpc_client_endpoint = Some(override_fn);
+        self
+    }
 }
 
 /// Policy of distributing tasks to available executor slots
@@ -450,6 +463,7 @@ impl TryFrom<Config> for SchedulerConfig {
             override_logical_codec: None,
             override_physical_codec: None,
             override_session_builder: None,
+            override_create_grpc_client_endpoint: None,
         };
 
         Ok(config)

ballista/scheduler/src/state/executor_manager.rs

@@ -427,7 +427,13 @@ impl ExecutorManager {
                 "http://{}:{}",
                 executor_metadata.host, executor_metadata.grpc_port
             );
-            let connection = create_grpc_client_endpoint(executor_url)?.connect().await?;
+            let mut endpoint = create_grpc_client_endpoint(executor_url)?;
+
+            if let Some(ref override_fn) = self.config.override_create_grpc_client_endpoint {
+                endpoint = override_fn(endpoint)?;
+            }
+
+            let connection = endpoint.connect().await?;
             let client = ExecutorGrpcClient::new(connection);
 
             {