refactor(config): make rpc client option configurable (#20993) (#21002)

Author: zwang28

src/batch/src/execution/local_exchange.rs

@@ -79,6 +79,7 @@ mod tests {
     use std::sync::Arc;
     use std::time::Duration;
 
+    use risingwave_common::config::RpcClientConfig;
     use risingwave_pb::batch_plan::{TaskId, TaskOutputId};
     use risingwave_pb::data::DataChunk;
     use risingwave_pb::task_service::exchange_service_server::{
@@ -154,7 +155,9 @@ mod tests {
         sleep(Duration::from_secs(1)).await;
         assert!(server_run.load(Ordering::SeqCst));
 
-        let client = ComputeClient::new(addr.into()).await.unwrap();
+        let client = ComputeClient::new(addr.into(), &RpcClientConfig::default())
+            .await
+            .unwrap();
         let task_output_id = TaskOutputId {
             task_id: Some(TaskId::default()),
             ..Default::default()

src/common/src/config.rs

@@ -543,6 +543,21 @@ pub struct MetaDeveloperConfig {
     /// Max number of epoch-to-version inserted into meta store per INSERT, during time travel metadata writing.
     #[serde(default = "default::developer::hummock_time_travel_epoch_version_insert_batch_size")]
     pub hummock_time_travel_epoch_version_insert_batch_size: usize,
+
+    #[serde(default)]
+    pub compute_client_config: RpcClientConfig,
+
+    #[serde(default)]
+    pub stream_client_config: RpcClientConfig,
+
+    #[serde(default)]
+    pub frontend_client_config: RpcClientConfig,
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize, DefaultFromSerde, ConfigDoc)]
+pub struct RpcClientConfig {
+    #[serde(default = "default::developer::rpc_client_connect_timeout_secs")]
+    pub connect_timeout_secs: u64,
 }
 
 /// The section `[server]` in `risingwave.toml`.
@@ -1176,6 +1191,9 @@ pub struct StreamingDeveloperConfig {
     /// When true, all jdbc sinks with connector='jdbc' and jdbc.url="jdbc:postgresql://..."
     /// will be switched from jdbc postgresql sinks to rust native (connector='postgres') sinks.
     pub switch_jdbc_pg_to_native: bool,
+
+    #[serde(default)]
+    pub compute_client_config: RpcClientConfig,
 }
 
 /// The subsections `[batch.developer]`.
@@ -1206,6 +1224,9 @@ pub struct BatchDeveloperConfig {
     /// If not specified, the value of `server.connection_pool_size` will be used.
     #[serde(default = "default::developer::batch_exchange_connection_pool_size")]
     exchange_connection_pool_size: Option<u16>,
+
+    #[serde(default)]
+    pub compute_client_config: RpcClientConfig,
 }
 
 macro_rules! define_system_config {
@@ -2160,6 +2181,9 @@ pub mod default {
         pub fn switch_jdbc_pg_to_native() -> bool {
             false
         }
+        pub fn rpc_client_connect_timeout_secs() -> u64 {
+            5
+        }
     }
 
     pub use crate::system_param::default as system;

src/compute/src/server.rs

@@ -349,6 +349,7 @@ pub async fn compute_node_serve(
     // Initialize batch environment.
     let batch_client_pool = Arc::new(ComputeClientPool::new(
         config.batch_exchange_connection_pool_size(),
+        config.batch.developer.compute_client_config.clone(),
     ));
     let batch_env = BatchEnvironment::new(
         batch_mgr.clone(),
@@ -368,6 +369,7 @@ pub async fn compute_node_serve(
     // Initialize the streaming environment.
     let stream_client_pool = Arc::new(ComputeClientPool::new(
         config.streaming_exchange_connection_pool_size(),
+        config.streaming.developer.compute_client_config.clone(),
     ));
     let stream_env = StreamEnvironment::new(
         advertise_addr.clone(),

src/config/example.toml

@@ -105,6 +105,15 @@ meta_hummock_time_travel_sst_info_insert_batch_size = 100
 meta_time_travel_vacuum_interval_sec = 30
 meta_hummock_time_travel_epoch_version_insert_batch_size = 1000
 
+[meta.developer.meta_compute_client_config]
+connect_timeout_secs = 5
+
+[meta.developer.meta_stream_client_config]
+connect_timeout_secs = 5
+
+[meta.developer.meta_frontend_client_config]
+connect_timeout_secs = 5
+
 [meta.meta_store_config]
 max_connections = 10
 min_connections = 1
@@ -127,6 +136,9 @@ batch_receiver_channel_size = 1000
 batch_root_stage_channel_size = 100
 batch_chunk_size = 1024
 
+[batch.developer.batch_compute_client_config]
+connect_timeout_secs = 5
+
 [streaming]
 in_flight_barrier_nums = 10000
 async_stack_trace = "ReleaseVerbose"
@@ -161,6 +173,9 @@ stream_enable_auto_schema_change = true
 stream_enable_shared_source = true
 stream_switch_jdbc_pg_to_native = false
 
+[streaming.developer.stream_compute_client_config]
+connect_timeout_secs = 5
+
 [storage]
 share_buffers_sync_parallelism = 1
 share_buffer_compaction_worker_threads_number = 4

src/ctl/src/cmd_impl/compute.rs

@@ -12,13 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use risingwave_common::config::{BatchConfig, StreamingConfig};
+use risingwave_common::config::{BatchConfig, RpcClientConfig, StreamingConfig};
 use risingwave_common::util::addr::HostAddr;
 use risingwave_rpc_client::ComputeClient;
 
 pub async fn show_config(host: &str) -> anyhow::Result<()> {
     let listen_addr = HostAddr::try_from(host)?;
-    let client = ComputeClient::new(listen_addr).await?;
+    let client = ComputeClient::new(listen_addr, &RpcClientConfig::default()).await?;
     let config_response = client.show_config().await?;
     let batch_config: BatchConfig = serde_json::from_str(&config_response.batch_config)?;
     let stream_config: StreamingConfig = serde_json::from_str(&config_response.stream_config)?;

src/ctl/src/cmd_impl/hummock/resize_cache.rs

@@ -15,6 +15,7 @@
 use std::process::exit;
 
 use futures::future::try_join_all;
+use risingwave_common::config::RpcClientConfig;
 use risingwave_pb::compute::ResizeCacheRequest;
 use risingwave_pb::meta::GetClusterInfoResponse;
 use risingwave_rpc_client::ComputeClient;
@@ -45,7 +46,7 @@ pub async fn resize_cache(
 
     let futures = worker_nodes.iter().map(|worker| async {
         let addr = worker.get_host().expect("worker host must be set");
-        let client = ComputeClient::new(addr.into())
+        let client = ComputeClient::new(addr.into(), &RpcClientConfig::default())
             .await
             .unwrap_or_else(|_| panic!("Cannot open client to compute node {addr:?}"));
         client

src/frontend/src/session.rs

@@ -320,6 +320,7 @@ impl FrontendEnv {
 
         let compute_client_pool = Arc::new(ComputeClientPool::new(
             config.batch_exchange_connection_pool_size(),
+            config.batch.developer.compute_client_config.clone(),
         ));
         let query_manager = QueryManager::new(
             worker_node_manager.clone(),

src/meta/node/src/lib.rs

@@ -476,6 +476,9 @@ pub fn start(
                     .developer
                     .actor_cnt_per_worker_parallelism_soft_limit,
                 license_key_path: opts.license_key_path,
+                compute_client_config: config.meta.developer.compute_client_config.clone(),
+                stream_client_config: config.meta.developer.stream_client_config.clone(),
+                frontend_client_config: config.meta.developer.frontend_client_config.clone(),
             },
             config.system.into_init_system_params(),
             Default::default(),

src/meta/node/src/server.rs

@@ -402,7 +402,7 @@ pub async fn start_service_as_election_leader(
             prometheus_client,
             prometheus_selector,
             metadata_manager: metadata_manager.clone(),
-            compute_clients: ComputeClientPool::new(1), // typically no need for plural clients
+            compute_clients: ComputeClientPool::new(1, env.opts.compute_client_config.clone()), /* typically no need for plural clients */
             diagnose_command,
             trace_state,
         };

src/meta/src/manager/env.rs

@@ -17,7 +17,9 @@ use std::path::PathBuf;
 use std::sync::Arc;
 
 use anyhow::Context;
-use risingwave_common::config::{CompactionConfig, DefaultParallelism, ObjectStoreConfig};
+use risingwave_common::config::{
+    CompactionConfig, DefaultParallelism, ObjectStoreConfig, RpcClientConfig,
+};
 use risingwave_common::session_config::SessionConfig;
 use risingwave_common::system_param::reader::SystemParamsReader;
 use risingwave_common::{bail, system_param};
@@ -254,6 +256,10 @@ pub struct MetaOpts {
     pub actor_cnt_per_worker_parallelism_soft_limit: usize,
 
     pub license_key_path: Option<PathBuf>,
+
+    pub compute_client_config: RpcClientConfig,
+    pub stream_client_config: RpcClientConfig,
+    pub frontend_client_config: RpcClientConfig,
 }
 
 impl MetaOpts {
@@ -331,6 +337,9 @@ impl MetaOpts {
             table_stat_throuput_window_seconds_for_merge: 240,
             periodic_scheduling_compaction_group_merge_interval_sec: 60 * 10,
             license_key_path: None,
+            compute_client_config: RpcClientConfig::default(),
+            stream_client_config: RpcClientConfig::default(),
+            frontend_client_config: RpcClientConfig::default(),
         }
     }
 }
@@ -343,8 +352,12 @@ impl MetaSrvEnv {
         meta_store_impl: SqlMetaStore,
     ) -> MetaResult<Self> {
         let idle_manager = Arc::new(IdleManager::new(opts.max_idle_ms));
-        let stream_client_pool = Arc::new(StreamClientPool::new(1)); // typically no need for plural clients
-        let frontend_client_pool = Arc::new(FrontendClientPool::new(1));
+        let stream_client_pool =
+            Arc::new(StreamClientPool::new(1, opts.stream_client_config.clone())); // typically no need for plural clients
+        let frontend_client_pool = Arc::new(FrontendClientPool::new(
+            1,
+            opts.frontend_client_config.clone(),
+        ));
         let event_log_manager = Arc::new(start_event_log_manager(
             opts.event_log_enabled,
             opts.event_log_channel_max_size,

src/rpc_client/src/compute_client.rs

@@ -18,7 +18,7 @@ use std::time::Duration;
 use async_trait::async_trait;
 use futures::StreamExt;
 use risingwave_common::catalog::DatabaseId;
-use risingwave_common::config::{MAX_CONNECTION_WINDOW_SIZE, STREAM_WINDOW_SIZE};
+use risingwave_common::config::{RpcClientConfig, MAX_CONNECTION_WINDOW_SIZE, STREAM_WINDOW_SIZE};
 use risingwave_common::monitor::{EndpointExt, TcpConfig};
 use risingwave_common::util::addr::HostAddr;
 use risingwave_common::util::tracing::TracingContext;
@@ -67,11 +67,11 @@ pub struct ComputeClient {
 }
 
 impl ComputeClient {
-    pub async fn new(addr: HostAddr) -> Result<Self> {
+    pub async fn new(addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
         let channel = Endpoint::from_shared(format!("http://{}", &addr))?
             .initial_connection_window_size(MAX_CONNECTION_WINDOW_SIZE)
             .initial_stream_window_size(STREAM_WINDOW_SIZE)
-            .connect_timeout(Duration::from_secs(5))
+            .connect_timeout(Duration::from_secs(opts.connect_timeout_secs))
             .monitored_connect(
                 "grpc-compute-client",
                 TcpConfig {
@@ -295,8 +295,8 @@ impl ComputeClient {
 
 #[async_trait]
 impl RpcClient for ComputeClient {
-    async fn new_client(host_addr: HostAddr) -> Result<Self> {
-        Self::new(host_addr).await
+    async fn new_client(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
+        Self::new(host_addr, opts).await
     }
 }
 

src/rpc_client/src/frontend_client.rs

@@ -16,7 +16,7 @@ use std::sync::Arc;
 use std::time::Duration;
 
 use async_trait::async_trait;
-use risingwave_common::config::MAX_CONNECTION_WINDOW_SIZE;
+use risingwave_common::config::{RpcClientConfig, MAX_CONNECTION_WINDOW_SIZE};
 use risingwave_common::monitor::{EndpointExt, TcpConfig};
 use risingwave_common::util::addr::HostAddr;
 use risingwave_pb::frontend_service::frontend_service_client::FrontendServiceClient;
@@ -37,10 +37,10 @@ const DEFAULT_RETRY_MAX_ATTEMPTS: usize = 10;
 struct FrontendClient(FrontendServiceClient<Channel>);
 
 impl FrontendClient {
-    async fn new(host_addr: HostAddr) -> Result<Self> {
+    async fn new(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
         let channel = Endpoint::from_shared(format!("http://{}", &host_addr))?
             .initial_connection_window_size(MAX_CONNECTION_WINDOW_SIZE)
-            .connect_timeout(Duration::from_secs(5))
+            .connect_timeout(Duration::from_secs(opts.connect_timeout_secs))
             .monitored_connect(
                 "grpc-frontend-client",
                 TcpConfig {
@@ -63,8 +63,8 @@ pub type FrontendClientPoolRef = Arc<FrontendClientPool>;
 
 #[async_trait]
 impl RpcClient for FrontendRetryClient {
-    async fn new_client(host_addr: HostAddr) -> Result<Self> {
-        Self::new(host_addr).await
+    async fn new_client(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
+        Self::new(host_addr, opts).await
     }
 }
 
@@ -74,8 +74,8 @@ pub struct FrontendRetryClient {
 }
 
 impl FrontendRetryClient {
-    async fn new(host_addr: HostAddr) -> Result<Self> {
-        let client = FrontendClient::new(host_addr).await?;
+    async fn new(host_addr: HostAddr, opts: &RpcClientConfig) -> Result<Self> {
+        let client = FrontendClient::new(host_addr, opts).await?;
         Ok(Self { client })
     }