Add gRPC connect timeout on LeaderManager::connect

Refactors the related code to plumb the value through to the grpc builder
Config is currently limited to the controller with a default of None.

This _should_ only affect the Controller -> Leader connection, which currently
does not have any supervisor to detect issues. Most (if not all) other connections
have some supervisor that will timeout and handle the error accordingly.

Author: cmackenzie1

crates/arroyo-api/src/jobs.rs

@@ -59,6 +59,7 @@ async fn fetch_from_leader<
         &config().api.tls,
         leader_context.worker_id,
         leader_context.rpc_address,
+        None,
     )
     .await
     .map_err(|_| ErrorResp {

crates/arroyo-api/src/lib.rs

@@ -124,12 +124,13 @@ pub async fn compiler_service() -> Result<CompilerGrpcClient<Channel>, ErrorResp
     let config = config();
     let endpoint = config.compiler_endpoint();
 
-    let channel = arroyo_rpc::connect_grpc("api", endpoint, &config.api.tls, &config.compiler.tls)
-        .await
-        .map_err(|e| {
-            error!("Failed to connect to compiler service: {}", e);
-            service_unavailable("compiler-service")
-        })?;
+    let channel =
+        arroyo_rpc::connect_grpc("api", endpoint, &config.api.tls, &config.compiler.tls, None)
+            .await
+            .map_err(|e| {
+                error!("Failed to connect to compiler service: {}", e);
+                service_unavailable("compiler-service")
+            })?;
 
     Ok(CompilerGrpcClient::new(channel))
 }

crates/arroyo-controller/src/job_controller/leader_manager.rs

@@ -33,13 +33,15 @@ impl LeaderManager {
         generation: u64,
         worker_id: WorkerId,
         address: String,
+        connect_timeout: Option<Duration>,
     ) -> anyhow::Result<Self> {
         let leader_client = retry!(
             job_status_client(
                 "controller",
                 &config().worker.tls,
                 worker_id,
-                address.clone()
+                address.clone(),
+                connect_timeout,
             )
             .await,
             5,

crates/arroyo-controller/src/schedulers/mod.rs

@@ -605,6 +605,7 @@ impl NodeScheduler {
             format!("http://{}", node.addr),
             &config().controller.tls,
             &config().node.tls,
+            None,
         )
         .await?;
         Ok(NodeGrpcClient::new(channel))

crates/arroyo-controller/src/states/mod.rs

@@ -927,7 +927,7 @@ pub(crate) async fn state_backoff(retries_attempted: usize, job_id: &str, pipeli
 
 #[allow(clippy::too_many_arguments)]
 async fn run_to_completion(
-    config: Arc<RwLock<(JobConfig, AppliedStatus)>>,
+    job_config_and_status: Arc<RwLock<(JobConfig, AppliedStatus)>>,
     pipeline_info: Arc<PipelineInfo>,
     mut program: LogicalProgram,
     mut status: JobStatus,
@@ -937,7 +937,7 @@ async fn run_to_completion(
     scheduler: Arc<dyn Scheduler>,
     metrics: Arc<tokio::sync::RwLock<HashMap<Arc<String>, JobMetrics>>>,
 ) {
-    let job_config = config.read().unwrap().0.clone();
+    let job_config = job_config_and_status.read().unwrap().0.clone();
 
     let leader_manager = if let Some(ctx) = &status.state_context.leader {
         LeaderManager::connect(
@@ -946,6 +946,7 @@ async fn run_to_completion(
             ctx.generation,
             ctx.worker_id,
             ctx.rpc_address.clone(),
+            config().controller.connect_timeout.as_deref().copied(),
         )
         .await
         .map(Some)
@@ -977,7 +978,7 @@ async fn run_to_completion(
     };
 
     loop {
-        config.write().unwrap().1 = AppliedStatus::Applied;
+        job_config_and_status.write().unwrap().1 = AppliedStatus::Applied;
         match execute_state(state, ctx).await {
             (Some(new_state), new_ctx) => {
                 state = new_state;
@@ -986,7 +987,7 @@ async fn run_to_completion(
             (None, _) => break,
         }
 
-        ctx.config = config.read().unwrap().0.clone();
+        ctx.config = job_config_and_status.read().unwrap().0.clone();
     }
 }
 

crates/arroyo-controller/src/states/scheduling.rs

@@ -160,6 +160,7 @@ async fn handle_worker_connect<'a>(
                         rpc_address.clone(),
                         &config().controller.tls,
                         &config().worker.tls,
+                        None,
                     )
                     .await
                     .unwrap()
@@ -935,6 +936,7 @@ impl State for Scheduling {
                             ctx.status.generation,
                             id,
                             addr,
+                            config().controller.connect_timeout.as_deref().copied(),
                         ).await
                             && let Ok(status) = leader_manager.poll_leader_status().await {
                                 match JobState::try_from(status.job_state) {
@@ -1031,6 +1033,7 @@ impl State for Scheduling {
                     ctx.status.generation,
                     id,
                     addr.clone(),
+                    config().controller.connect_timeout.as_deref().copied(),
                 )
                 .await
                 {

crates/arroyo-rpc/src/config.rs

@@ -401,6 +401,10 @@ pub struct ControllerConfig {
     /// Poll interval for leader status
     pub leader_poll_interval: HumanReadableDuration,
 
+    /// Timeout for connecting to gRPC services
+    #[serde(default)]
+    pub connect_timeout: Option<HumanReadableDuration>,
+
     /// Metric system configurations
     pub metrics: MetricsConfig,
 }

crates/arroyo-rpc/src/lib.rs

@@ -1002,9 +1002,10 @@ pub async fn grpc_channel_builder(
     endpoint: String,
     our_tls: &Option<TlsConfig>,
     target_tls: &Option<TlsConfig>,
+    connect_timeout: Option<Duration>,
 ) -> Result<Endpoint> {
     let config = config();
-    if let Some(target_tls) = config.get_tls_config(target_tls) {
+    let endpoint = if let Some(target_tls) = config.get_tls_config(target_tls) {
         let mut endpoint = Url::parse(&endpoint)?;
         endpoint
             .set_scheme("https")
@@ -1031,11 +1032,17 @@ pub async fn grpc_channel_builder(
             config_builder = config_builder.identity(Identity::from_pem(our_tls.cert, our_tls.key));
         }
 
-        Ok(b.tls_config(config_builder).context("configuring TLS")?)
+        b.tls_config(config_builder).context("configuring TLS")?
     } else {
         debug!("connecting to grpc endpoint {endpoint}");
-        Ok(Channel::from_shared(endpoint.to_string())?)
-    }
+        Channel::from_shared(endpoint.to_string())?
+    };
+
+    Ok(if let Some(connect_timeout) = connect_timeout {
+        endpoint.connect_timeout(connect_timeout)
+    } else {
+        endpoint
+    })
 }
 
 /// Connect to a gRPC service with optional TLS
@@ -1044,17 +1051,20 @@ pub async fn connect_grpc(
     endpoint: String,
     our_tls: &Option<TlsConfig>,
     tls: &Option<TlsConfig>,
+    connect_timeout: Option<Duration>,
 ) -> Result<Channel> {
-    Ok(grpc_channel_builder(our_name, endpoint, our_tls, tls)
-        .await?
-        .connect()
-        .await?)
+    Ok(
+        grpc_channel_builder(our_name, endpoint, our_tls, tls, connect_timeout)
+            .await?
+            .connect()
+            .await?,
+    )
 }
 
 /// Connect a raw gRPC channel to the controller endpoint.
 pub async fn connect_controller(our_name: &str, our_tls: &Option<TlsConfig>) -> Result<Channel> {
     let endpoint = config().controller_endpoint();
-    connect_grpc(our_name, endpoint, our_tls, &config().controller.tls).await
+    connect_grpc(our_name, endpoint, our_tls, &config().controller.tls, None).await
 }
 
 pub async fn controller_client(
@@ -1077,7 +1087,7 @@ pub async fn job_controller_client(
         &config().controller.tls
     };
 
-    let channel = connect_grpc(our_name, addr, our_tls, their_tls).await?;
+    let channel = connect_grpc(our_name, addr, our_tls, their_tls, None).await?;
     Ok(job_controller_grpc_client::JobControllerGrpcClient::new(
         channel,
     ))
@@ -1088,9 +1098,17 @@ pub async fn job_status_client(
     our_tls: &Option<TlsConfig>,
     worker_id: WorkerId,
     addr: String,
+    connect_timeout: Option<Duration>,
 ) -> Result<job_status_grpc_client::JobStatusGrpcClient<InterceptedService<Channel, InjectWorkerId>>>
 {
-    let channel = connect_grpc(our_name, addr, our_tls, &config().worker.tls).await?;
+    let channel = connect_grpc(
+        our_name,
+        addr,
+        our_tls,
+        &config().worker.tls,
+        connect_timeout,
+    )
+    .await?;
     Ok(
         job_status_grpc_client::JobStatusGrpcClient::with_interceptor(
             channel,

crates/arroyo-worker/src/job_controller/mod.rs

@@ -317,6 +317,7 @@ pub(crate) async fn connect_to_worker(id: WorkerId, addr: String) -> anyhow::Res
             addr.clone(),
             &config().worker.tls,
             &config().worker.tls,
+            None,
         )
         .await?
         .timeout(Duration::from_secs(15))