allow customized schedulergrpcclient

Author: mach-kernel

ballista/core/src/serde/mod.rs

@@ -50,7 +50,6 @@ use crate::remote_catalog::remote_table_provider::RemoteTableProvider;
 use crate::serde::protobuf::ballista_physical_plan_node::PhysicalPlanType;
 use crate::serde::scheduler::PartitionLocation;
 use datafusion::catalog::TableProvider;
-use datafusion::logical_expr::UserDefinedLogicalNode;
 pub use generated::ballista as protobuf;
 use prost::Message;
 use std::fmt::Debug;

ballista/executor/src/execution_loop.rs

@@ -45,14 +45,20 @@ use std::time::{SystemTime, UNIX_EPOCH};
 use std::{sync::Arc, time::Duration};
 use tokio::sync::oneshot::Sender as OneShotSender;
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
-use tonic::transport::Channel;
 
-pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
-    mut scheduler: SchedulerGrpcClient<Channel>,
+pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan, C>(
+    mut scheduler: SchedulerGrpcClient<C>,
     executor: Arc<Executor>,
     codec: BallistaCodec<T, U>,
     readiness: Option<OneShotSender<String>>,
-) -> Result<(), BallistaError> {
+) -> Result<(), BallistaError>
+where
+    C: tonic::client::GrpcService<tonic::body::Body> + Send + 'static,
+    C::Error: Into<Box<dyn Error + Send + Sync>>,
+    C::ResponseBody: tonic::codegen::Body<Data = tonic::codegen::Bytes> + Send + 'static,
+    <C::ResponseBody as tonic::codegen::Body>::Error:
+        Into<Box<dyn Error + Send + Sync>> + Send,
+{
     let executor_specification: ExecutorSpecification = executor
         .metadata
         .specification