diff --git a/crates/arroyo-controller/src/lib.rs b/crates/arroyo-controller/src/lib.rs index e4b7cca4..fcaf404b 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 66fc495e..c02e6e55 100644 --- a/crates/arroyo-controller/src/states/mod.rs +++ b/crates/arroyo-controller/src/states/mod.rs @@ -1208,6 +1208,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()