Add gRPC connect timeouts, prevent global locking of job_state

crates/arroyo-controller/src/lib.rs

@@ -651,22 +651,28 @@ impl ControllerServer {
     }
 
     async fn send_to_job_queue(&self, job_id: &str, msg: JobMessage) -> Result<(), Status> {
-        let mut jobs = self.job_state.lock().await;
-
-        if let Some(sm) = jobs.get_mut(job_id) {
-            if let Err(e) = sm.send(msg).await {
-                Err(Status::failed_precondition(format!(
-                    "Cannot handle message for {job_id}: {e}"
-                )))
-            } else {
-                Ok(())
-            }
-        } else {
-            warn!(message = "Received message for unknown job id", job_id);
-            Err(Status::failed_precondition(format!(
-                "No job with id {job_id}"
-            )))
-        }
+        // Keep per-job backpressure from holding the global job map lock.
+        let tx = {
+            let jobs = self.job_state.lock().await;
+            let Some(sm) = jobs.get(job_id) else {
+                warn!(message = "Received message for unknown job id", job_id);
+                return Err(Status::failed_precondition(format!(
+                    "No job with id {job_id}"
+                )));
+            };
+
+            sm.sender().ok_or_else(|| {
+                Status::failed_precondition(format!(
+                    "Cannot handle message for {job_id}: State machine is inactive"
+                ))
+            })?
+        };
+
+        tx.send(msg).await.map_err(|_| {
+            Status::failed_precondition(format!(
+                "Cannot handle message for {job_id}: State machine is inactive"
+            ))
+        })
     }
 
     fn start_updater(&self, guard: ShutdownGuard) {

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

@@ -1207,6 +1207,10 @@ impl StateMachine {
         }
     }
 
+    pub(crate) fn sender(&self) -> Option<Sender<JobMessage>> {
+        self.tx.clone()
+    }
+
     pub fn done(&self) -> bool {
         if let Some(tx) = &self.tx {
             tx.is_closed()

crates/arroyo-rpc/src/config.rs

@@ -193,6 +193,10 @@ fn load_config(paths: &[PathBuf]) -> Figment {
 #[derive(Debug, Deserialize, Serialize, Clone)]
 #[serde(rename_all = "kebab-case", deny_unknown_fields)]
 pub struct Config {
+    /// gRPC client configuration
+    #[serde(default)]
+    pub grpc: GrpcConfig,
+
     /// API service configuration
     pub api: ApiConfig,
 
@@ -275,6 +279,13 @@ pub struct Config {
     pub disable_telemetry: bool,
 }
 
+#[derive(Debug, Default, Deserialize, Serialize, Clone)]
+#[serde(rename_all = "kebab-case", deny_unknown_fields)]
+pub struct GrpcConfig {
+    /// Maximum time to establish a gRPC connection
+    pub connect_timeout: Option<HumanReadableDuration>,
+}
+
 #[derive(Debug, Deserialize, Serialize, Clone, Default)]
 #[serde(rename_all = "kebab-case", deny_unknown_fields)]
 pub enum JobControllerMode {
@@ -1047,6 +1058,7 @@ impl TlsConfig {
 #[cfg(test)]
 mod tests {
     use crate::config::{Config, DatabaseType, Scheduler, SchemaName, SqliteConfig, load_config};
+    use std::time::Duration;
     use url::Url;
 
     #[test]
@@ -1073,6 +1085,33 @@ mod tests {
         }
     }
 
+    #[test]
+    #[allow(clippy::result_large_err)]
+    fn grpc_connect_timeout_is_unset_by_default() {
+        figment::Jail::expect_with(|_| {
+            let config: Config = load_config(&[]).extract().unwrap();
+
+            assert!(config.grpc.connect_timeout.is_none());
+            Ok(())
+        });
+    }
+
+    #[test]
+    #[allow(clippy::result_large_err)]
+    fn grpc_connect_timeout_can_be_overridden_with_environment() {
+        figment::Jail::expect_with(|jail| {
+            jail.set_env("ARROYO__GRPC__CONNECT_TIMEOUT", "3s");
+
+            let config: Config = load_config(&[]).extract().unwrap();
+
+            assert_eq!(
+                **config.grpc.connect_timeout.as_ref().unwrap(),
+                Duration::from_secs(3)
+            );
+            Ok(())
+        });
+    }
+
     #[test]
     #[allow(clippy::result_large_err)]
     fn test_config() {

crates/arroyo-rpc/src/lib.rs

@@ -1004,7 +1004,7 @@ pub async fn grpc_channel_builder(
     target_tls: &Option<TlsConfig>,
 ) -> 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,10 +1031,16 @@ 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())?
+    };
+
+    if let Some(connect_timeout) = config.grpc.connect_timeout.as_deref() {
+        Ok(endpoint.connect_timeout(*connect_timeout))
+    } else {
+        Ok(endpoint)
     }
 }