diff --git a/ballista/core/src/client.rs b/ballista/core/src/client.rs index c01d431ca..00292e334 100644 --- a/ballista/core/src/client.rs +++ b/ballista/core/src/client.rs @@ -17,17 +17,11 @@ //! Client API for sending requests to executors. -use std::collections::HashMap; -use std::sync::Arc; - -use std::{ - convert::{TryFrom, TryInto}, - task::{Context, Poll}, -}; - use crate::error::{BallistaError, Result as BResult}; +use crate::extension::BallistaConfigGrpcEndpoint; +use crate::serde::protobuf; use crate::serde::scheduler::{Action, PartitionId}; - +use crate::utils::create_grpc_client_endpoint; use arrow_flight; use arrow_flight::Ticket; use arrow_flight::utils::flight_data_to_arrow_batch; @@ -43,21 +37,23 @@ use datafusion::arrow::{ }; use datafusion::error::DataFusionError; use datafusion::error::Result; - -use crate::extension::BallistaConfigGrpcEndpoint; -use crate::serde::protobuf; - -use crate::utils::create_grpc_client_endpoint; - use datafusion::physical_plan::{RecordBatchStream, SendableRecordBatchStream}; use futures::{Stream, StreamExt}; use log::{debug, warn}; use prost::Message; +use std::collections::HashMap; +use std::sync::Arc; +use std::{ + convert::{TryFrom, TryInto}, + task::{Context, Poll}, +}; use tonic::{Code, Streaming}; /// Client for interacting with Ballista executors. #[derive(Clone)] pub struct BallistaClient { + host: String, + port: u16, flight_client: FlightServiceClient, } @@ -109,7 +105,11 @@ impl BallistaClient { debug!("BallistaClient connected OK: {flight_client:?}"); - Ok(Self { flight_client }) + Ok(Self { + flight_client, + host: host.to_string(), + port, + }) } /// Retrieves a partition from an executor. @@ -122,8 +122,6 @@ impl BallistaClient { executor_id: &str, partition_id: &PartitionId, path: &str, - host: &str, - port: u16, flight_transport: bool, ) -> BResult { let action = Action::FetchPartition { @@ -131,8 +129,8 @@ impl BallistaClient { stage_id: partition_id.stage_id, partition_id: partition_id.partition_id, path: path.to_owned(), - host: host.to_owned(), - port, + host: self.host.to_owned(), + port: self.port, }; let result = if flight_transport { diff --git a/ballista/core/src/config.rs b/ballista/core/src/config.rs index ca151558c..d74dacf7a 100644 --- a/ballista/core/src/config.rs +++ b/ballista/core/src/config.rs @@ -31,10 +31,8 @@ use datafusion::{ pub const BALLISTA_JOB_NAME: &str = "ballista.job.name"; /// Configuration key for standalone processing parallelism. pub const BALLISTA_STANDALONE_PARALLELISM: &str = "ballista.standalone.parallelism"; - /// Configuration key for disabling default cache extension node. pub const BALLISTA_CACHE_NOOP: &str = "ballista.cache.noop"; - /// Configuration key for maximum concurrent shuffle read requests. pub const BALLISTA_SHUFFLE_READER_MAX_REQUESTS: &str = "ballista.shuffle.max_concurrent_read_requests"; @@ -44,7 +42,6 @@ pub const BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ: &str = /// Configuration key to prefer Flight protocol for remote shuffle reads. pub const BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT: &str = "ballista.shuffle.remote_read_prefer_flight"; - /// max message size for gRPC clients pub const BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE: &str = "ballista.grpc_client_max_message_size"; @@ -82,6 +79,8 @@ pub const BALLISTA_SHUFFLE_SORT_BASED_BATCH_SIZE: &str = "ballista.shuffle.sort_based.batch_size"; /// Should client employ pull or push job tracking strategy pub const BALLISTA_CLIENT_PULL: &str = "ballista.client.pull"; +/// Should client use tls connection +pub const BALLISTA_CLIENT_USE_TLS: &str = "ballista.client.use_tls"; /// Result type for configuration parsing operations. pub type ParseResult = result::Result; @@ -162,6 +161,10 @@ static CONFIG_ENTRIES: LazyLock> = LazyLock::new(|| ConfigEntry::new(BALLISTA_CLIENT_PULL.to_string(), "Should client employ pull or push job tracking. In pull mode client will make a request to server in the loop, until job finishes. Pull mode is kept for legacy clients.".to_string(), DataType::Boolean, + Some(false.to_string())), + ConfigEntry::new(BALLISTA_CLIENT_USE_TLS.to_string(), + "Should connection between client, scheduler, and executors use TLS.".to_string(), + DataType::Boolean, Some(false.to_string())) ]; entries @@ -274,11 +277,6 @@ impl BallistaConfig { &self.settings } - /// Returns the maximum message size for gRPC clients in bytes. - pub fn default_grpc_client_max_message_size(&self) -> usize { - self.get_usize_setting(BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE) - } - /// Returns the standalone processing parallelism level. pub fn default_standalone_parallelism(&self) -> usize { self.get_usize_setting(BALLISTA_STANDALONE_PARALLELISM) @@ -290,25 +288,30 @@ impl BallistaConfig { } /// Returns the gRPC client connection timeout in seconds. - pub fn default_grpc_client_connect_timeout_seconds(&self) -> usize { + pub fn grpc_client_connect_timeout_seconds(&self) -> usize { self.get_usize_setting(BALLISTA_GRPC_CLIENT_CONNECT_TIMEOUT_SECONDS) } /// Returns the gRPC client request timeout in seconds. - pub fn default_grpc_client_timeout_seconds(&self) -> usize { + pub fn grpc_client_timeout_seconds(&self) -> usize { self.get_usize_setting(BALLISTA_GRPC_CLIENT_TIMEOUT_SECONDS) } /// Returns the TCP keep-alive interval for gRPC clients in seconds. - pub fn default_grpc_client_tcp_keepalive_seconds(&self) -> usize { + pub fn grpc_client_tcp_keepalive_seconds(&self) -> usize { self.get_usize_setting(BALLISTA_GRPC_CLIENT_TCP_KEEPALIVE_SECONDS) } /// Returns the HTTP/2 keep-alive interval for gRPC clients in seconds. - pub fn default_grpc_client_http2_keepalive_interval_seconds(&self) -> usize { + pub fn grpc_client_http2_keepalive_interval_seconds(&self) -> usize { self.get_usize_setting(BALLISTA_GRPC_CLIENT_HTTP2_KEEPALIVE_INTERVAL_SECONDS) } + /// Returns the maximum message size for gRPC clients in bytes. + pub fn grpc_client_max_message_size(&self) -> usize { + self.get_usize_setting(BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE) + } + /// Returns whether the default cache node extension is disabled. pub fn cache_noop(&self) -> bool { self.get_bool_setting(BALLISTA_CACHE_NOOP) @@ -373,6 +376,11 @@ impl BallistaConfig { self.get_bool_setting(BALLISTA_CLIENT_PULL) } + /// should client use TLS to communicate with ballista cluster + pub fn client_use_tls(&self) -> bool { + self.get_bool_setting(BALLISTA_CLIENT_USE_TLS) + } + fn get_usize_setting(&self, key: &str) -> usize { if let Some(v) = self.settings.get(key) { // infallible because we validate all configs in the constructor @@ -419,6 +427,23 @@ impl BallistaConfig { v.parse::().unwrap() } } + /// sets the configuration value where key starts with ballista + /// prefix. + pub fn set_with_prefix( + &mut self, + key: &str, + value: &str, + ) -> datafusion::error::Result<()> { + let entries = Self::valid_entries(); + //let k = format!("{}.{key}", BallistaConfig::PREFIX); + + if entries.contains_key(key) { + self.settings.insert(key.to_string(), value.to_string()); + Ok(()) + } else { + config_err!("configuration key `{}` does not exist", key) + } + } } impl datafusion::config::ExtensionOptions for BallistaConfig { @@ -539,7 +564,7 @@ mod tests { #[test] fn default_config() -> Result<()> { let config = BallistaConfig::default(); - assert_eq!(16777216, config.default_grpc_client_max_message_size()); + assert_eq!(16777216, config.grpc_client_max_message_size()); Ok(()) } } diff --git a/ballista/core/src/execution_plans/distributed_query.rs b/ballista/core/src/execution_plans/distributed_query.rs index 5f1ad258d..599c9f967 100644 --- a/ballista/core/src/execution_plans/distributed_query.rs +++ b/ballista/core/src/execution_plans/distributed_query.rs @@ -253,7 +253,7 @@ impl ExecutionPlan for DistributedQueryExec { self.scheduler_url.clone(), self.session_id.clone(), query, - self.config.default_grpc_client_max_message_size(), + self.config.grpc_client_max_message_size(), GrpcClientConfig::from(&self.config), Arc::new(self.metrics.clone()), partition, @@ -280,7 +280,7 @@ impl ExecutionPlan for DistributedQueryExec { execute_query_push( self.scheduler_url.clone(), query, - self.config.default_grpc_client_max_message_size(), + self.config.grpc_client_max_message_size(), GrpcClientConfig::from(&self.config), Arc::new(self.metrics.clone()), partition, @@ -701,8 +701,6 @@ async fn fetch_partition( let partition_id = location.partition_id.ok_or_else(|| { DataFusionError::Internal("Received empty partition id".to_owned()) })?; - let host = metadata.host.as_str(); - let port = metadata.port as u16; let (client_host, client_port) = get_client_host_port(&metadata, &scheduler_url, &flight_proxy)?; @@ -721,8 +719,6 @@ async fn fetch_partition( &metadata.id, &partition_id.into(), &location.path, - host, - port, flight_transport, ) .await diff --git a/ballista/core/src/execution_plans/shuffle_reader.rs b/ballista/core/src/execution_plans/shuffle_reader.rs index c3676a93b..c7d9e8e3f 100644 --- a/ballista/core/src/execution_plans/shuffle_reader.rs +++ b/ballista/core/src/execution_plans/shuffle_reader.rs @@ -15,49 +15,45 @@ // specific language governing permissions and limitations // under the License. -use async_trait::async_trait; -use datafusion::arrow::ipc::reader::StreamReader; -use datafusion::common::stats::Precision; -use datafusion::physical_plan::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; -use std::any::Any; -use std::collections::HashMap; -use std::fmt::Debug; -use std::fs::File; -use std::io::BufReader; -use std::pin::Pin; -use std::result; -use std::sync::Arc; -use std::task::{Context, Poll}; - use crate::client::BallistaClient; +use crate::config::BallistaConfig; +use crate::error::BallistaError; use crate::execution_plans::sort_shuffle::{ get_index_path, is_sort_shuffle_output, stream_sort_shuffle_partition, }; use crate::extension::{BallistaConfigGrpcEndpoint, SessionConfigExt}; use crate::serde::scheduler::{PartitionLocation, PartitionStats}; - use datafusion::arrow::datatypes::SchemaRef; use datafusion::arrow::error::ArrowError; +use datafusion::arrow::ipc::reader::StreamReader; use datafusion::arrow::record_batch::RecordBatch; use datafusion::common::runtime::SpawnedTask; - +use datafusion::common::stats::Precision; use datafusion::error::{DataFusionError, Result}; +use datafusion::execution::context::TaskContext; +use datafusion::physical_plan::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; use datafusion::physical_plan::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, }; +use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ ColumnStatistics, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; use futures::{Stream, StreamExt, TryStreamExt, ready}; - -use crate::error::BallistaError; -use datafusion::execution::context::TaskContext; -use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use itertools::Itertools; use log::{debug, error, trace}; use rand::prelude::SliceRandom; use rand::rng; +use std::any::Any; +use std::collections::HashMap; +use std::fmt::Debug; +use std::fs::File; +use std::io::BufReader; +use std::pin::Pin; +use std::result; +use std::sync::Arc; +use std::task::{Context, Poll}; use tokio::sync::{Semaphore, mpsc}; use tokio_stream::wrappers::ReceiverStream; @@ -162,17 +158,13 @@ impl ExecutionPlan for ShuffleReaderExec { debug!("ShuffleReaderExec::execute({task_id})"); let config = context.session_config(); - - let max_request_num = - config.ballista_shuffle_reader_maximum_concurrent_requests(); - let max_message_size = config.ballista_grpc_client_max_message_size(); - let force_remote_read = config.ballista_shuffle_reader_force_remote_read(); - let prefer_flight = config.ballista_shuffle_reader_remote_prefer_flight(); let batch_size = config.batch_size(); - let customize_endpoint = config.ballista_override_create_grpc_client_endpoint(); - let use_tls = config.ballista_use_tls(); + let customize_endpoint: Option> = + config.ballista_override_create_grpc_client_endpoint(); + + let config = config.ballista_config(); - if force_remote_read { + if config.shuffle_reader_force_remote_read() { debug!( "All shuffle partitions will be read as remote partitions! To disable this behavior set: `{}=false`", crate::config::BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ @@ -180,7 +172,9 @@ impl ExecutionPlan for ShuffleReaderExec { } log::debug!( - "ShuffleReaderExec::execute({task_id}) max_request_num: {max_request_num}, max_message_size: {max_message_size}" + "ShuffleReaderExec::execute({task_id}) max_request_num: {}, max_message_size: {}", + config.shuffle_reader_maximum_concurrent_requests(), + config.grpc_client_max_message_size() ); let mut partition_locations = HashMap::new(); for p in &self.partition[partition] { @@ -198,15 +192,8 @@ impl ExecutionPlan for ShuffleReaderExec { .collect(); // Shuffle partitions for evenly send fetching partition requests to avoid hot executors within multiple tasks partition_locations.shuffle(&mut rng()); - let response_receiver = send_fetch_partitions( - partition_locations, - max_request_num, - max_message_size, - force_remote_read, - prefer_flight, - customize_endpoint, - use_tls, - ); + let response_receiver = + send_fetch_partitions(partition_locations, &config, customize_endpoint); let input_stream = Box::pin(RecordBatchStreamAdapter::new( self.schema.clone(), @@ -405,19 +392,19 @@ fn local_remote_read_split( fn send_fetch_partitions( partition_locations: Vec, - max_request_num: usize, - max_message_size: usize, - force_remote_read: bool, - flight_transport: bool, + config: &BallistaConfig, customize_endpoint: Option>, - use_tls: bool, ) -> AbortableReceiverStream { + let max_request_num = config.shuffle_reader_maximum_concurrent_requests(); + let (response_sender, response_receiver) = mpsc::channel(max_request_num); let semaphore = Arc::new(Semaphore::new(max_request_num)); let mut spawned_tasks: Vec> = vec![]; - let (local_locations, remote_locations): (Vec<_>, Vec<_>) = - local_remote_read_split(partition_locations, force_remote_read); + let (local_locations, remote_locations): (Vec<_>, Vec<_>) = local_remote_read_split( + partition_locations, + config.shuffle_reader_force_remote_read(), + ); debug!( "local shuffle file counts:{}, remote shuffle file count:{}.", @@ -427,46 +414,44 @@ fn send_fetch_partitions( // keep local shuffle files reading in serial order for memory control. let response_sender_c = response_sender.clone(); - let customize_endpoint_c = customize_endpoint.clone(); - spawned_tasks.push(SpawnedTask::spawn(async move { - for p in local_locations { - let r = PartitionReaderEnum::Local - .fetch_partition( - &p, - max_message_size, - flight_transport, - customize_endpoint_c.clone(), - use_tls, - ) - .await; - if let Err(e) = response_sender_c.send(r).await { - error!("Fail to send response event to the channel due to {e}"); + + // + // fetching local partitions (read from file) + // + + spawned_tasks.push(SpawnedTask::spawn_blocking({ + move || { + for p in local_locations { + let r = fetch_partition_local(&p); + if let Err(e) = response_sender_c.blocking_send(r) { + error!("Fail to send response event to the channel due to {e}"); + } } } })); + // + // fetching remote partitions (uses grpc flight protocol) + // for p in remote_locations.into_iter() { let semaphore = semaphore.clone(); let response_sender = response_sender.clone(); - let customize_endpoint_c = customize_endpoint.clone(); - spawned_tasks.push(SpawnedTask::spawn(async move { - // Block if exceeds max request number. - let permit = semaphore.acquire_owned().await.unwrap(); - let r = PartitionReaderEnum::FlightRemote - .fetch_partition( - &p, - max_message_size, - flight_transport, - customize_endpoint_c, - use_tls, - ) - .await; - // Block if the channel buffer is full. - if let Err(e) = response_sender.send(r).await { - error!("Fail to send response event to the channel due to {e}"); + + spawned_tasks.push(SpawnedTask::spawn({ + // TODO: make BallistaConfig cheaper to clone + let config = config.clone(); + let customize_endpoint = customize_endpoint.clone(); + async move { + // Block if exceeds max request number. + let permit = semaphore.acquire_owned().await.unwrap(); + let r = fetch_partition_remote(&p, &config, customize_endpoint).await; + // Block if the channel buffer is full. + if let Err(e) = response_sender.send(r).await { + error!("Fail to send response event to the channel due to {e}"); + } + // Increase semaphore by dropping existing permits. + drop(permit); } - // Increase semaphore by dropping existing permits. - drop(permit); })); } @@ -477,103 +462,51 @@ fn check_is_local_location(location: &PartitionLocation) -> bool { std::path::Path::new(location.path.as_str()).exists() } -/// Partition reader Trait, different partition reader can have -#[async_trait] -trait PartitionReader: Send + Sync + Clone { - // Read partition data from PartitionLocation - async fn fetch_partition( - &self, - location: &PartitionLocation, - max_message_size: usize, - flight_transport: bool, - customize_endpoint: Option>, - use_tls: bool, - ) -> result::Result; -} - -#[derive(Clone)] -enum PartitionReaderEnum { - Local, - FlightRemote, - #[allow(dead_code)] - ObjectStoreRemote, -} +async fn new_ballista_client( + host: &str, + port: u16, + config: &BallistaConfig, + customize_endpoint: Option>, +) -> result::Result { + let max_message_size = config.grpc_client_max_message_size(); + let use_tls = config.client_use_tls(); -#[async_trait] -impl PartitionReader for PartitionReaderEnum { - // Notice return `BallistaError::FetchFailed` will let scheduler re-schedule the task. - async fn fetch_partition( - &self, - location: &PartitionLocation, - max_message_size: usize, - flight_transport: bool, - customize_endpoint: Option>, - use_tls: bool, - ) -> result::Result { - match self { - PartitionReaderEnum::FlightRemote => { - fetch_partition_remote( - location, - max_message_size, - flight_transport, - customize_endpoint, - use_tls, - ) - .await - } - PartitionReaderEnum::Local => fetch_partition_local(location).await, - PartitionReaderEnum::ObjectStoreRemote => { - fetch_partition_object_store(location).await - } - } - } + BallistaClient::try_new(host, port, max_message_size, use_tls, customize_endpoint) + .await } async fn fetch_partition_remote( location: &PartitionLocation, - max_message_size: usize, - flight_transport: bool, + config: &BallistaConfig, customize_endpoint: Option>, - use_tls: bool, ) -> result::Result { let metadata = &location.executor_meta; let partition_id = &location.partition_id; - // TODO for shuffle client connections, we should avoid creating new connections again and again. - // And we should also avoid to keep alive too many connections for long time. + let prefer_flight = config.shuffle_reader_remote_prefer_flight(); let host = metadata.host.as_str(); let port = metadata.port; - let mut ballista_client = BallistaClient::try_new( - host, - port, - max_message_size, - use_tls, - customize_endpoint, - ) - .await - .map_err(|error| match error { - // map grpc connection error to partition fetch error. - BallistaError::GrpcConnectionError(msg) => BallistaError::FetchFailed( - metadata.id.clone(), - partition_id.stage_id, - partition_id.partition_id, - msg, - ), - other => other, - })?; + + // TODO for shuffle client connections, we should avoid creating new connections again and again. + // And we should also avoid to keep alive too many connections for long time. + let mut ballista_client = new_ballista_client(host, port, config, customize_endpoint) + .await + .map_err(|error| match error { + // map grpc connection error to partition fetch error. + BallistaError::GrpcConnectionError(msg) => BallistaError::FetchFailed( + metadata.id.clone(), + partition_id.stage_id, + partition_id.partition_id, + msg, + ), + other => other, + })?; ballista_client - .fetch_partition( - &metadata.id, - partition_id, - &location.path, - host, - port, - flight_transport, - ) + .fetch_partition(&metadata.id, partition_id, &location.path, prefer_flight) .await } -async fn fetch_partition_local( +fn fetch_partition_local( location: &PartitionLocation, ) -> result::Result { let path = &location.path; @@ -637,14 +570,6 @@ fn fetch_partition_local_inner( Ok(reader) } -async fn fetch_partition_object_store( - _location: &PartitionLocation, -) -> result::Result { - Err(BallistaError::NotImplemented( - "Should not use ObjectStorePartitionReader".to_string(), - )) -} - struct CoalescedShuffleReaderStream { schema: SchemaRef, input: SendableRecordBatchStream, @@ -1120,16 +1045,15 @@ mod tests { partition_num, file_path.to_str().unwrap().to_string(), ); + let mut config = BallistaConfig::default(); - let response_receiver = send_fetch_partitions( - partition_locations, - max_request_num, - 4 * 1024 * 1024, - false, - true, - None, - false, - ); + config + .set_with_prefix( + crate::config::BALLISTA_SHUFFLE_READER_MAX_REQUESTS, + &max_request_num.to_string(), + ) + .unwrap(); + let response_receiver = send_fetch_partitions(partition_locations, &config, None); let stream = RecordBatchStreamAdapter::new( Arc::new(schema), diff --git a/ballista/core/src/extension.rs b/ballista/core/src/extension.rs index 2c777097c..7374978eb 100644 --- a/ballista/core/src/extension.rs +++ b/ballista/core/src/extension.rs @@ -16,7 +16,7 @@ // under the License. use crate::config::{ - BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE, BALLISTA_JOB_NAME, + BALLISTA_CLIENT_USE_TLS, BALLISTA_GRPC_CLIENT_MAX_MESSAGE_SIZE, BALLISTA_JOB_NAME, BALLISTA_SHUFFLE_READER_FORCE_REMOTE_READ, BALLISTA_SHUFFLE_READER_MAX_REQUESTS, BALLISTA_SHUFFLE_READER_REMOTE_PREFER_FLIGHT, BALLISTA_STANDALONE_PARALLELISM, BallistaConfig, @@ -392,10 +392,8 @@ impl SessionConfigExt for SessionConfig { self.options() .extensions .get::() - .map(|c| c.default_grpc_client_max_message_size()) - .unwrap_or_else(|| { - BallistaConfig::default().default_grpc_client_max_message_size() - }) + .map(|c| c.grpc_client_max_message_size()) + .unwrap_or_else(|| BallistaConfig::default().grpc_client_max_message_size()) } fn with_ballista_job_name(self, job_name: &str) -> Self { @@ -538,13 +536,20 @@ impl SessionConfigExt for SessionConfig { } fn with_ballista_use_tls(self, use_tls: bool) -> Self { - self.with_extension(Arc::new(BallistaUseTls(use_tls))) + if self.options().extensions.get::().is_some() { + self.set_bool(BALLISTA_CLIENT_USE_TLS, use_tls) + } else { + self.with_option_extension(BallistaConfig::default()) + .set_bool(BALLISTA_CLIENT_USE_TLS, use_tls) + } } fn ballista_use_tls(&self) -> bool { - self.get_extension::() - .map(|ext| ext.0) - .unwrap_or(false) + self.options() + .extensions + .get::() + .map(|c| c.client_use_tls()) + .unwrap_or_else(|| BallistaConfig::default().client_use_tls()) } } @@ -746,10 +751,6 @@ impl BallistaConfigGrpcEndpoint { } } -/// Wrapper for cluster-wide TLS configuration -#[derive(Clone, Copy)] -pub struct BallistaUseTls(pub bool); - #[derive(Debug)] struct BallistaCacheFactory; diff --git a/ballista/core/src/utils.rs b/ballista/core/src/utils.rs index 0d6a3d833..20aa8df0c 100644 --- a/ballista/core/src/utils.rs +++ b/ballista/core/src/utils.rs @@ -68,13 +68,11 @@ pub struct GrpcClientConfig { impl From<&BallistaConfig> for GrpcClientConfig { fn from(config: &BallistaConfig) -> Self { Self { - connect_timeout_seconds: config.default_grpc_client_connect_timeout_seconds() - as u64, - timeout_seconds: config.default_grpc_client_timeout_seconds() as u64, - tcp_keepalive_seconds: config.default_grpc_client_tcp_keepalive_seconds() - as u64, + connect_timeout_seconds: config.grpc_client_connect_timeout_seconds() as u64, + timeout_seconds: config.grpc_client_timeout_seconds() as u64, + tcp_keepalive_seconds: config.grpc_client_tcp_keepalive_seconds() as u64, http2_keepalive_interval_seconds: config - .default_grpc_client_http2_keepalive_interval_seconds() + .grpc_client_http2_keepalive_interval_seconds() as u64, } } @@ -312,19 +310,19 @@ mod tests { // Verify the conversion picks up the right values assert_eq!( grpc_config.connect_timeout_seconds, - ballista_config.default_grpc_client_connect_timeout_seconds() as u64 + ballista_config.grpc_client_connect_timeout_seconds() as u64 ); assert_eq!( grpc_config.timeout_seconds, - ballista_config.default_grpc_client_timeout_seconds() as u64 + ballista_config.grpc_client_timeout_seconds() as u64 ); assert_eq!( grpc_config.tcp_keepalive_seconds, - ballista_config.default_grpc_client_tcp_keepalive_seconds() as u64 + ballista_config.grpc_client_tcp_keepalive_seconds() as u64 ); assert_eq!( grpc_config.http2_keepalive_interval_seconds, - ballista_config.default_grpc_client_http2_keepalive_interval_seconds() as u64 + ballista_config.grpc_client_http2_keepalive_interval_seconds() as u64 ); } diff --git a/examples/examples/standalone-substrait.rs b/examples/examples/standalone-substrait.rs index 7e8b2036c..95caa8370 100644 --- a/examples/examples/standalone-substrait.rs +++ b/examples/examples/standalone-substrait.rs @@ -416,8 +416,6 @@ impl SubstraitSchedulerClient { &metadata.id, &partition_id.into(), &location.path, - host, - port, flight_transport, ) .await