diff --git a/crates/arroyo-controller/src/lib.rs b/crates/arroyo-controller/src/lib.rs index 11f100d8..09d1995f 100644 --- a/crates/arroyo-controller/src/lib.rs +++ b/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) { diff --git a/crates/arroyo-controller/src/states/mod.rs b/crates/arroyo-controller/src/states/mod.rs index 3d5199a4..d4ce12a3 100644 --- a/crates/arroyo-controller/src/states/mod.rs +++ b/crates/arroyo-controller/src/states/mod.rs @@ -1207,6 +1207,10 @@ impl StateMachine { } } + pub(crate) fn sender(&self) -> Option> { + self.tx.clone() + } + pub fn done(&self) -> bool { if let Some(tx) = &self.tx { tx.is_closed() diff --git a/crates/arroyo-rpc/src/config.rs b/crates/arroyo-rpc/src/config.rs index 7238b677..35dfe0ac 100644 --- a/crates/arroyo-rpc/src/config.rs +++ b/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, +} + #[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() { diff --git a/crates/arroyo-rpc/src/lib.rs b/crates/arroyo-rpc/src/lib.rs index 7acc7bb8..3ecbef43 100644 --- a/crates/arroyo-rpc/src/lib.rs +++ b/crates/arroyo-rpc/src/lib.rs @@ -1004,7 +1004,7 @@ pub async fn grpc_channel_builder( target_tls: &Option, ) -> Result { 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) } }