From c7fe7fed2bdeecfdffbd2bbdd8ffd6e7e091ee47 Mon Sep 17 00:00:00 2001 From: Otavio Salvador Date: Fri, 28 Aug 2026 14:41:45 -0300 Subject: [PATCH 1/4] fix: signal the waker without blocking The waker carries a single bit of information: the main loop should stop waiting on the current transition and step the machine. The loop treats it that way already, discharging it with `try_recv` on every iteration, but the handlers filled it with a blocking `send().await` on a one-slot channel. That made a full slot mean "wait until the loop drains me", which the loop cannot do while a handler is running. Two accepted requests reaching a handler between two transitions were enough to park the agent for good: the first took the slot, the second waited on a loop that would not run again until the handler it was blocking returned. Sending without waiting says what was meant. An occupied slot already carries the message, so there is nothing to wait for. --- updatehub/src/states/machine/mod.rs | 18 ++++++++++++++---- 1 file changed, 14 insertions(+), 4 deletions(-) diff --git a/updatehub/src/states/machine/mod.rs b/updatehub/src/states/machine/mod.rs index 7e597bdf..93ce423f 100644 --- a/updatehub/src/states/machine/mod.rs +++ b/updatehub/src/states/machine/mod.rs @@ -128,7 +128,7 @@ pub(super) trait CommunicationState: StateChangeImpl { ProbeResponse::NoUpdate => { info!("no update is current available for this device"); - context.waker.sender.send(()).await?; + context.wake(); // Store timestamp of last polling context.runtime_settings.set_last_polling(Utc::now())?; @@ -137,7 +137,7 @@ pub(super) trait CommunicationState: StateChangeImpl { ProbeResponse::Update(package, sign) => { info!("update received: {} ({})", package.version(), package.package_uid()); - context.waker.sender.send(()).await?; + context.wake(); // Store timestamp of last polling context.runtime_settings.set_last_polling(Utc::now())?; @@ -173,7 +173,7 @@ pub(super) trait CommunicationState: StateChangeImpl { // Starting logging a new scope of operation since we are // starting to handle a user request crate::logger::start_memory_logging(); - context.waker.sender.send(()).await?; + context.wake(); Ok(( address::StateResponse::RequestAccepted(name), @@ -195,7 +195,7 @@ pub(super) trait CommunicationState: StateChangeImpl { // Starting logging a new scope of operation since we are // starting to handle a user request crate::logger::start_memory_logging(); - context.waker.sender.send(()).await?; + context.wake(); Ok(( address::StateResponse::RequestAccepted(name), @@ -222,6 +222,16 @@ impl Context { } } + /// Asks the main loop to stop waiting on the current transition. + /// + /// The waker holds a single pending wake up, so an occupied slot already + /// carries the message. Signalling must never block: a handler runs to + /// completion before the loop that drains the waker gets to run again, so + /// waiting for room here would deadlock the agent. + pub(super) fn wake(&self) { + let _ = self.waker.sender.try_send(()); + } + pub(super) fn server_address(&self) -> &str { self.runtime_settings .custom_server_address() From 68fd1c77fdad0fc8491a390cf2dd801dc21e6fff Mon Sep 17 00:00:00 2001 From: Otavio Salvador Date: Fri, 28 Aug 2026 14:43:30 -0300 Subject: [PATCH 2/4] fix: answer requests that are still being handled when a wait ends While a state waits on its transition, the loop raced the sleep and the waker against a future that both received a request and handled it. Whenever the sleep or the wake won, `select` dropped that future wherever it happened to be, and a request suspended on a slow server went down with it: its reply channel was dropped unsent, and the caller was left holding a receiver nobody would ever send on. Racing the receive alone settles it. Dropping a pending receive loses nothing, so the request is only taken off the channel once it is going to be answered, and the handling runs where nothing can cancel it. The wait itself now outlives the requests answered against it, so a query no longer restarts a poll interval, and a request that redirects the machine ends the wait it was made against instead of relying on a wake up to do it. Two smaller repairs come along, both of which cost a request its answer: - An accepted request kept its state change even when the reply could not be delivered. Tolerating a client that hung up is right; throwing away the local install it just asked for is not. - `Addr` treated a closed reply channel as unreachable and panicked the caller's task. A dead state machine is a runtime condition, and `TransitionError::CommunicationFailed` already describes it. The remaining `unreachable!` covers a mismatched response variant, which really cannot happen. The tests hold a probe at the server and wake the machine while the answer is outstanding, which is what a poll timer firing mid-request looks like in production. --- updatehub/src/cloud_mock.rs | 26 ++++++ updatehub/src/states/machine/address.rs | 38 ++++++-- updatehub/src/states/machine/mod.rs | 113 ++++++++++++++---------- updatehub/src/states/machine/tests.rs | 109 +++++++++++++++++++++++ 4 files changed, 233 insertions(+), 53 deletions(-) create mode 100644 updatehub/src/states/machine/tests.rs diff --git a/updatehub/src/cloud_mock.rs b/updatehub/src/cloud_mock.rs index 9787b723..37a4e89e 100644 --- a/updatehub/src/cloud_mock.rs +++ b/updatehub/src/cloud_mock.rs @@ -13,6 +13,12 @@ std::thread_local! { static OBJECT_DATA: RefCell>> = const { RefCell::new(Option::None) }; } +type ProbeGate = (async_channel::Receiver<()>, async_channel::Sender<()>); + +std::thread_local! { + static PROBE_GATE: RefCell> = const { RefCell::new(None) }; +} + pub(crate) enum FakeResponse { NoUpdate, HasUpdate, @@ -28,6 +34,18 @@ pub(crate) fn setup_fake_response(res: FakeResponse) { RESPONSE_CONFIG.with(|conf| conf.replace_with(move |&mut _| res)); } +/// Holds every probe until it is released, standing in for a server that has +/// not answered yet. +/// +/// Returns the release side and a receiver that reports each probe reaching the +/// server, so a test can act while a request is provably half-answered. +pub(crate) fn gate_probes() -> (async_channel::Sender<()>, async_channel::Receiver<()>) { + let (release, gate) = async_channel::unbounded(); + let (reached, arrived) = async_channel::unbounded(); + PROBE_GATE.with(|conf| conf.borrow_mut().replace((gate, reached))); + (release, arrived) +} + pub(crate) fn set_download_data(data: Vec) { OBJECT_DATA.with(|conf| conf.borrow_mut().replace(data)); } @@ -42,6 +60,14 @@ impl<'a> Client<'a> { _num_retries: usize, _firmware: api::FirmwareMetadata<'_>, ) -> Result { + // Cloned out before awaiting, so the `RefCell` is never borrowed across + // a suspension point. + let gate = PROBE_GATE.with(|conf| conf.borrow().clone()); + if let Some((gate, reached)) = gate { + let _ = reached.send(()).await; + let _ = gate.recv().await; + } + RESPONSE_CONFIG.with(|conf| match std::ops::Deref::deref(&conf.borrow()) { FakeResponse::NoUpdate => Ok(api::ProbeResponse::NoUpdate), FakeResponse::ExtraPoll => Ok(api::ProbeResponse::ExtraPoll(10)), diff --git a/updatehub/src/states/machine/address.rs b/updatehub/src/states/machine/address.rs index 9d63fb68..bf1bee1a 100644 --- a/updatehub/src/states/machine/address.rs +++ b/updatehub/src/states/machine/address.rs @@ -49,9 +49,18 @@ pub(crate) enum StateResponse { InvalidState(String), } +// Either half of a request's channel going away means the state machine is no +// longer running. Report that rather than taking the caller's task down with a +// panic. impl From> for crate::states::TransitionError { - fn from(err: async_channel::SendError) -> Self { - unreachable!("Unexpected sending error for {:?}", err) + fn from(_: async_channel::SendError) -> Self { + crate::states::TransitionError::CommunicationFailed + } +} + +impl From for crate::states::TransitionError { + fn from(_: async_channel::RecvError) -> Self { + crate::states::TransitionError::CommunicationFailed } } @@ -62,7 +71,10 @@ impl Addr { match recv.recv().await { Ok(Ok(Response::Info(resp))) => Ok(*resp), Ok(Err(e)) => Err(e), - res => unreachable!("Unexpected response: {:?}", res), + Err(async_channel::RecvError) => { + Err(crate::states::TransitionError::CommunicationFailed) + } + Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), } } @@ -75,7 +87,10 @@ impl Addr { match recv.recv().await { Ok(Ok(Response::Probe(resp))) => Ok(resp), Ok(Err(e)) => Err(e), - res => unreachable!("Unexpected response: {:?}", res), + Err(async_channel::RecvError) => { + Err(crate::states::TransitionError::CommunicationFailed) + } + Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), } } @@ -85,7 +100,10 @@ impl Addr { match recv.recv().await { Ok(Ok(Response::AbortDownload(resp))) => Ok(resp), Ok(Err(e)) => Err(e), - res => unreachable!("Unexpected response: {:?}", res), + Err(async_channel::RecvError) => { + Err(crate::states::TransitionError::CommunicationFailed) + } + Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), } } @@ -99,7 +117,10 @@ impl Addr { match recv.recv().await { Ok(Ok(Response::LocalInstall(resp))) => Ok(resp), Ok(Err(e)) => Err(e), - res => unreachable!("Unexpected response: {:?}", res), + Err(async_channel::RecvError) => { + Err(crate::states::TransitionError::CommunicationFailed) + } + Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), } } @@ -110,7 +131,10 @@ impl Addr { match recv.recv().await { Ok(Ok(Response::RemoteInstall(resp))) => Ok(resp), Ok(Err(e)) => Err(e), - res => unreachable!("Unexpected response: {:?}", res), + Err(async_channel::RecvError) => { + Err(crate::states::TransitionError::CommunicationFailed) + } + Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), } } } diff --git a/updatehub/src/states/machine/mod.rs b/updatehub/src/states/machine/mod.rs index 93ce423f..8a09cd68 100644 --- a/updatehub/src/states/machine/mod.rs +++ b/updatehub/src/states/machine/mod.rs @@ -4,6 +4,9 @@ mod address; +#[cfg(test)] +mod tests; + use super::{ DirectDownload, EntryPoint, Metadata, PrepareLocalInstall, Result, RuntimeSettings, Settings, State, StateChangeImpl, Validation, @@ -83,17 +86,22 @@ pub(super) trait CommunicationState: StateChangeImpl { .map(|(res, st)| (address::Response::RemoteInstall(res), st)), }; - match res { - Ok((response, state)) => { - responder.send(Ok(response)).await.ok()?; - state - } + let (reply, state) = match res { + Ok((response, state)) => (Ok(response), state), Err(e) => { error!("Request failed with: {}", e); - responder.send(Err(e)).await.ok()?; - None + (Err(e), None) } + }; + + // A client that hung up before we answered is tolerated, but the work + // already done on its behalf is not thrown away: an accepted request + // still moves the machine to the state it asked for. + if responder.send(reply).await.is_err() { + trace!("client gave up before the response was delivered"); } + + state } async fn handle_probe( @@ -261,10 +269,13 @@ impl StateMachine { } pub(super) async fn start(mut self) { + let waker = self.context.waker.receiver.clone(); + let communication = self.context.communication.receiver.clone(); + loop { // Since the loop is already currently running, we can // discharges any wake message received. - let _ = self.context.waker.receiver.try_recv(); + let _ = waker.try_recv(); self.consume_pending_communication().await; @@ -275,39 +286,48 @@ impl StateMachine { .unwrap_or_else(|e| (State::from(e), StepTransition::Immediate)); self.state = state; - match transition { - StepTransition::Immediate => {} + let delay = match transition { + StepTransition::Immediate => continue, StepTransition::Delayed(t) => { trace!("delaying transition for: {} seconds", t.num_seconds()); - let waker = self.context.waker.receiver.clone(); - - let sleep_fut = tokio::time::sleep(t.to_std().unwrap_or_default()); - let waker_fut = async { - let _ = waker.recv().await; - }; - let comm_fut = self.await_communication(); - - futures_util::pin_mut!(sleep_fut); - futures_util::pin_mut!(waker_fut); - futures_util::pin_mut!(comm_fut); - - let _ = futures_util::future::select( - futures_util::future::select(sleep_fut, waker_fut), - comm_fut, - ) - .await; + Some(t.to_std().unwrap_or_default()) } StepTransition::Never => { trace!("stopping transition until awoken"); - let waker_recv = self.context.waker.receiver.clone(); - let recv_fut = waker_recv.recv(); - let comm_fut = async { - self.await_communication().await; - std::result::Result::<_, async_channel::RecvError>::Ok(()) - }; - - futures_util::pin_mut!(recv_fut, comm_fut); - let _ = futures_util::future::select(recv_fut, comm_fut).await; + None + } + }; + + // Built once and polled across every request answered below, so + // answering one does not restart the wait the current state asked + // for. `Never` is the same wait with a deadline that never fires. + let mut elapsed = std::pin::pin!(async { + match delay { + Some(t) => tokio::time::sleep(t).await, + None => std::future::pending::<()>().await, + } + }); + + loop { + let (msg, responder) = tokio::select! { + // Polled in the order the previous `select` chain used. + biased; + () = elapsed.as_mut() => break, + _ = waker.recv() => break, + received = communication.recv() => match received { + Ok(request) => request, + Err(_) => break, + }, + }; + + // Handled out here on purpose. Only the receive takes part in + // the race above: dropping a pending receive loses nothing, + // whereas dropping a request already being handled would leave + // its caller waiting on a reply channel nobody holds any more. + if self.handle_request(msg, responder).await { + // The request moved the machine, so the wait it was made + // against no longer describes what to do next. + break; } } } @@ -315,21 +335,22 @@ impl StateMachine { async fn consume_pending_communication(&mut self) { while let Ok((msg, responder)) = self.context.communication.receiver.try_recv() { - if let Some(new_state) = - self.state.handle_communication(msg, responder, &mut self.context).await - { - self.state = new_state; - } + let _ = self.handle_request(msg, responder).await; } } - async fn await_communication(&mut self) { - while let Ok((msg, responder)) = self.context.communication.receiver.recv().await { - if let Some(new_state) = - self.state.handle_communication(msg, responder, &mut self.context).await - { + /// Answers one request, reporting whether it moved the machine. + async fn handle_request( + &mut self, + msg: Message, + responder: async_channel::Sender>, + ) -> bool { + match self.state.handle_communication(msg, responder, &mut self.context).await { + Some(new_state) => { self.state = new_state; + true } + None => false, } } } diff --git a/updatehub/src/states/machine/tests.rs b/updatehub/src/states/machine/tests.rs new file mode 100644 index 00000000..a9a657a6 --- /dev/null +++ b/updatehub/src/states/machine/tests.rs @@ -0,0 +1,109 @@ +// Copyright (C) 2026 O.S. Systems Sofware LTDA +// +// SPDX-License-Identifier: Apache-2.0 + +use super::*; +use crate::states::{Park, Poll}; +use chrono::Utc; + +/// Wakes the machine while a probe is still waiting on the server, then lets +/// the server answer. +/// +/// A handler suspended on a slow server is the case that matters: the wait the +/// machine was racing resolves while a request is only half answered, and the +/// answer still has to reach the client. +async fn probe_interrupted_by_a_wake(state: State, setup: crate::tests::TestEnvironment) { + let context = setup.gen_context(); + let waker = context.waker.sender.clone(); + let machine = StateMachine { state, context }; + let addr = machine.address(); + + tokio::task::LocalSet::new() + .run_until(async move { + tokio::task::spawn_local(machine.start()); + + // A completed round trip proves the machine reached the transition + // under test instead of still walking the loop. + addr.request_info().await.unwrap(); + + let (release, arrived) = crate::cloud_mock::gate_probes(); + let (response, ()) = futures_util::future::join(addr.request_probe(None), async { + arrived.recv().await.unwrap(); + waker.send(()).await.unwrap(); + release.send(()).await.unwrap(); + }) + .await; + + assert!( + matches!(response, Ok(ProbeResponse::Unavailable)), + "the probe was answered with {response:?}" + ); + }) + .await; +} + +#[tokio::test] +async fn a_slow_probe_is_answered_while_parked() { + let setup = crate::tests::TestEnvironment::build().disable_polling().finish(); + probe_interrupted_by_a_wake(State::Park(Park {}), setup).await; +} + +#[tokio::test] +async fn a_slow_probe_is_answered_while_delaying() { + let mut setup = crate::tests::TestEnvironment::build().finish(); + // Polling just happened, so `Poll` delays for the whole interval instead of + // probing right away. + setup.runtime_settings.data.polling.last = Utc::now(); + probe_interrupted_by_a_wake(State::Poll(Poll {}), setup).await; +} + +/// Answering one request must not stall the next. Each accepted request signals +/// the waker, whose single slot the loop only drains between transitions. +#[tokio::test] +async fn queued_requests_are_all_answered() { + let setup = crate::tests::TestEnvironment::build().disable_polling().finish(); + let machine = StateMachine { state: State::Park(Park {}), context: setup.gen_context() }; + let addr = machine.address(); + + tokio::task::LocalSet::new() + .run_until(async move { + tokio::task::spawn_local(machine.start()); + addr.request_info().await.unwrap(); + + let (first, second, third) = futures_util::future::join3( + addr.request_probe(None), + addr.request_probe(None), + addr.request_probe(None), + ) + .await; + + for (which, response) in [("first", first), ("second", second), ("third", third)] { + assert!( + matches!(response, Ok(ProbeResponse::Unavailable)), + "{which} probe was answered with {response:?}" + ); + } + }) + .await; +} + +#[tokio::test] +async fn state_transition_survives_an_abandoned_request() { + let setup = crate::tests::TestEnvironment::build().finish(); + let mut context = setup.gen_context(); + + let (responder, reply) = async_channel::bounded(1); + // The client gave up before the agent got to answer. + drop(reply); + + let new_state = State::Park(Park {}) + .handle_communication( + Message::LocalInstall(std::path::PathBuf::from("/tmp/update.uhupkg")), + responder, + &mut context, + ) + .await; + + let machine = new_state.expect("the accepted local install was dropped"); + assert_state!(machine, PrepareLocalInstall); +} From 212d58c389ffe0730602d01fd9775c4c857f56af Mon Sep 17 00:00:00 2001 From: Otavio Salvador Date: Fri, 28 Aug 2026 14:44:25 -0300 Subject: [PATCH 3/4] fix: share one request-handling race across the download states `Download` and `DirectDownload` each hand-rolled the shape the main loop just had fixed: race the download against a future that receives a request and handles it, and let `select` drop whichever loses. A handler suspended when the download finished went the same way its counterpart in the loop did, taking the caller's reply with it. Neither state can reach that today. Both are non-preemptive, so every handler they see answers without suspending, and the drop has nothing to interrupt. That makes this a repair of the shape rather than of an observed failure, which is also why one preemptive state away it would have become a live bug in a place nobody was looking. `handle_communication_while` now owns the race for all three sites. It goes a step further than the loop: the download keeps being polled while a request is answered, so a query no longer stalls a transfer, and when the download wins the race `select` hands back the unfinished handler to be driven to completion rather than dropped. A request that redirects the machine outranks work that is no longer wanted, which matches the reply the client was already given. Both states shrink to building their work future and handing it over, and `start_download` takes the update package by reference, dropping two of its three clones. --- updatehub/src/states/direct_download.rs | 26 ++----------- updatehub/src/states/download.rs | 52 +++++-------------------- updatehub/src/states/machine/mod.rs | 52 ++++++++++++++++++++++++- updatehub/src/states/machine/tests.rs | 34 ++++++++++++++++ 4 files changed, 98 insertions(+), 66 deletions(-) diff --git a/updatehub/src/states/direct_download.rs b/updatehub/src/states/direct_download.rs index 61dbac82..7eda95bf 100644 --- a/updatehub/src/states/direct_download.rs +++ b/updatehub/src/states/direct_download.rs @@ -29,10 +29,9 @@ impl StateChangeImpl for DirectDownload { async fn handle(self, context: &mut Context) -> Result<(State, machine::StepTransition)> { info!("fetching update package directly from url: {:?}", self.url); - let communication_receiver = &context.communication.receiver.clone(); let context = Mutex::new(context); - let download_future = async { + let download = async { let download_dir = context.lock().await.settings.update.download_dir.clone(); tokio::fs::create_dir_all(&download_dir) .await @@ -46,27 +45,8 @@ impl StateChangeImpl for DirectDownload { Ok(State::PrepareLocalInstall(PrepareLocalInstall { update_file })) }; - let message_handle_future = async { - while let Ok((msg, responder)) = communication_receiver.recv().await { - if let Some(new_state) = - self.handle_communication(msg, responder, *context.lock().await).await - { - return Ok(new_state); - } - } + let state = self.handle_communication_while(&context, download).await?; - Err(super::TransitionError::CommunicationFailed) - }; - - futures_util::pin_mut!(download_future); - futures_util::pin_mut!(message_handle_future); - - Ok(( - futures_util::future::select(download_future, message_handle_future) - .await - .factor_first() - .0?, - machine::StepTransition::Immediate, - )) + Ok((state, machine::StepTransition::Immediate)) } } diff --git a/updatehub/src/states/download.rs b/updatehub/src/states/download.rs index a8950490..c7648488 100644 --- a/updatehub/src/states/download.rs +++ b/updatehub/src/states/download.rs @@ -23,7 +23,7 @@ pub(super) struct Download { impl Download { async fn start_download( - update_package: UpdatePackage, + update_package: &UpdatePackage, context: &Mutex<&mut Context>, ) -> Result<()> { let installation_set = @@ -129,52 +129,20 @@ impl StateChangeImpl for Download { } async fn handle(self, context: &mut Context) -> Result<(State, machine::StepTransition)> { - let communication_receiver = &context.communication.receiver.clone(); let context = Mutex::new(context); - let update_package = self.update_package.clone(); - let download_future = async { - Download::start_download(update_package.clone(), &context).await?; - Result::Ok(None) + let download = async { + Download::start_download(&self.update_package, &context).await?; + Ok(State::Validation(Validation { + package: self.update_package.clone(), + sign: self.sign.clone(), + require_download: false, + })) }; - let message_handle_future = async { - while let Ok((msg, responder)) = communication_receiver.recv().await { - if let Some(new_state) = - self.handle_communication(msg, responder, *context.lock().await).await - { - return Ok(Some(new_state)); - } - } - Ok(None) - }; + let state = self.handle_communication_while(&context, download).await?; - // Clone update package and object_context so self can be freely held by - // message_handle_future - let update_package = self.update_package.clone(); - let sign = self.sign.clone(); - - // download_future dones't need to be pinned as it doesn't borrow context - futures_util::pin_mut!(download_future); - futures_util::pin_mut!(message_handle_future); - - if let Some(new_state) = - futures_util::future::select(download_future, message_handle_future) - .await - .factor_first() - .0? - { - return Ok((new_state, machine::StepTransition::Immediate)); - } - - Ok(( - State::Validation(Validation { - package: update_package, - sign, - require_download: false, - }), - machine::StepTransition::Immediate, - )) + Ok((state, machine::StepTransition::Immediate)) } } diff --git a/updatehub/src/states/machine/mod.rs b/updatehub/src/states/machine/mod.rs index 8a09cd68..6973e230 100644 --- a/updatehub/src/states/machine/mod.rs +++ b/updatehub/src/states/machine/mod.rs @@ -11,8 +11,10 @@ use super::{ DirectDownload, EntryPoint, Metadata, PrepareLocalInstall, Result, RuntimeSettings, Settings, State, StateChangeImpl, Validation, }; +use async_lock::Mutex; +use futures_util::future::{Either, select}; use slog_scope::{error, info, trace}; -use std::path::PathBuf; +use std::{future::Future, path::PathBuf}; pub(crate) use address::{ AbortDownloadResponse, Addr, Message, ProbeResponse, Response, StateResponse, @@ -104,6 +106,54 @@ pub(super) trait CommunicationState: StateChangeImpl { state } + /// Runs `work` while answering requests, returning whichever settles first: + /// the state `work` produced, or the state a request moved the machine to. + /// + /// A request already being handled is never dropped, whichever way the + /// races below go: dropping one leaves its caller waiting on a reply + /// channel nobody holds any more, and the caller answers that with an + /// error the request never recovers from. + /// + /// `work` must not hold the context lock across an await. Answering a + /// request takes that lock, and `work` is not polled while the lock is + /// being acquired, so a guard held over a suspension point deadlocks both. + async fn handle_communication_while( + &self, + context: &Mutex<&mut Context>, + work: impl Future>, + ) -> Result { + let communication = context.lock().await.communication.receiver.clone(); + futures_util::pin_mut!(work); + + loop { + let received = communication.recv(); + futures_util::pin_mut!(received); + + let (msg, responder) = match select(work.as_mut(), received).await { + Either::Left((done, _)) => return done, + Either::Right((Ok(request), _)) => request, + // Nobody can ask us anything any more, so just finish the work. + Either::Right((Err(_), _)) => return work.as_mut().await, + }; + + let mut locked = context.lock().await; + let handling = self.handle_communication(msg, responder, *locked); + futures_util::pin_mut!(handling); + + // `work` keeps being polled while the request is answered, but the + // answering itself is never what gets dropped: when `work` wins + // this race it hands back the unfinished `handling`, which is then + // driven to completion rather than discarded. + match select(work.as_mut(), handling).await { + // The request still gets its answer, and a request that + // redirects the machine outranks work we no longer want. + Either::Left((done, handling)) => return handling.await.map_or(done, Ok), + Either::Right((Some(new_state), _)) => return Ok(new_state), + Either::Right((None, _)) => {} + } + } + } + async fn handle_probe( &self, context: &mut Context, diff --git a/updatehub/src/states/machine/tests.rs b/updatehub/src/states/machine/tests.rs index a9a657a6..e157eaec 100644 --- a/updatehub/src/states/machine/tests.rs +++ b/updatehub/src/states/machine/tests.rs @@ -107,3 +107,37 @@ async fn state_transition_survives_an_abandoned_request() { let machine = new_state.expect("the accepted local install was dropped"); assert_state!(machine, PrepareLocalInstall); } + +/// The download states race a request against work that can finish first. +#[tokio::test] +async fn a_half_answered_request_survives_the_work_finishing() { + let setup = crate::tests::TestEnvironment::build().finish(); + let mut context = setup.gen_context(); + + let (responder, reply) = async_channel::bounded(1); + context.communication.sender.send((Message::Probe(None), responder)).await.unwrap(); + + let (release, arrived) = crate::cloud_mock::gate_probes(); + let (finish_work, work_finished) = async_channel::bounded::<()>(1); + let context = async_lock::Mutex::new(&mut context); + let work = async { + work_finished.recv().await.unwrap(); + Ok(State::Park(Park {})) + }; + + let (state, ()) = futures_util::future::join( + State::Park(Park {}).handle_communication_while(&context, work), + async { + // The request is waiting on the server, so finish the work while it + // is still only half answered. + arrived.recv().await.unwrap(); + finish_work.send(()).await.unwrap(); + release.send(()).await.unwrap(); + }, + ) + .await; + + assert!(reply.try_recv().is_ok(), "the request was dropped half-answered"); + let machine = state.unwrap(); + assert_state!(machine, EntryPoint); +} From 94e3bff633b232194b77e08dc0cee20f05c3616a Mon Sep 17 00:00:00 2001 From: Otavio Salvador Date: Fri, 28 Aug 2026 14:45:18 -0300 Subject: [PATCH 4/4] refactor: collapse the repeated request plumbing in Addr Every `request_*` method opened its own reply channel, sent the message, awaited the answer, and then matched four arms to unwrap it. Five copies, which is how the response matching had drifted apart in the first place, and it meant each new request type had to remember the closed-channel arm or reintroduce a panic. A `From` impl beside the existing `From` one lets `?` carry the closed channel, and a private `request` holds the plumbing. What is left in each method is the part that actually differs: the message it sends and the response variant it expects. --- updatehub/src/states/machine/address.rs | 67 ++++++++----------------- 1 file changed, 22 insertions(+), 45 deletions(-) diff --git a/updatehub/src/states/machine/address.rs b/updatehub/src/states/machine/address.rs index bf1bee1a..747ae96e 100644 --- a/updatehub/src/states/machine/address.rs +++ b/updatehub/src/states/machine/address.rs @@ -65,16 +65,17 @@ impl From for crate::states::TransitionError { } impl Addr { - pub(crate) async fn request_info(&self) -> super::Result { + /// Hands `msg` to the state machine and waits for the answer to come back. + async fn request(&self, msg: Message) -> super::Result { let (sndr, recv) = async_channel::bounded(1); - self.message.send((Message::Info, sndr)).await?; - match recv.recv().await { - Ok(Ok(Response::Info(resp))) => Ok(*resp), - Ok(Err(e)) => Err(e), - Err(async_channel::RecvError) => { - Err(crate::states::TransitionError::CommunicationFailed) - } - Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), + self.message.send((msg, sndr)).await?; + recv.recv().await? + } + + pub(crate) async fn request_info(&self) -> super::Result { + match self.request(Message::Info).await? { + Response::Info(resp) => Ok(*resp), + res => unreachable!("Unexpected response: {res:?}"), } } @@ -82,28 +83,16 @@ impl Addr { &self, custom_server: Option, ) -> super::Result { - let (sndr, recv) = async_channel::bounded(1); - self.message.send((Message::Probe(custom_server), sndr)).await?; - match recv.recv().await { - Ok(Ok(Response::Probe(resp))) => Ok(resp), - Ok(Err(e)) => Err(e), - Err(async_channel::RecvError) => { - Err(crate::states::TransitionError::CommunicationFailed) - } - Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), + match self.request(Message::Probe(custom_server)).await? { + Response::Probe(resp) => Ok(resp), + res => unreachable!("Unexpected response: {res:?}"), } } pub(crate) async fn request_abort_download(&self) -> super::Result { - let (sndr, recv) = async_channel::bounded(1); - self.message.send((Message::AbortDownload, sndr)).await?; - match recv.recv().await { - Ok(Ok(Response::AbortDownload(resp))) => Ok(resp), - Ok(Err(e)) => Err(e), - Err(async_channel::RecvError) => { - Err(crate::states::TransitionError::CommunicationFailed) - } - Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), + match self.request(Message::AbortDownload).await? { + Response::AbortDownload(resp) => Ok(resp), + res => unreachable!("Unexpected response: {res:?}"), } } @@ -112,29 +101,17 @@ impl Addr { path: PathBuf, ) -> super::Result { trace!("Local install requested"); - let (sndr, recv) = async_channel::bounded(1); - self.message.send((Message::LocalInstall(path), sndr)).await?; - match recv.recv().await { - Ok(Ok(Response::LocalInstall(resp))) => Ok(resp), - Ok(Err(e)) => Err(e), - Err(async_channel::RecvError) => { - Err(crate::states::TransitionError::CommunicationFailed) - } - Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), + match self.request(Message::LocalInstall(path)).await? { + Response::LocalInstall(resp) => Ok(resp), + res => unreachable!("Unexpected response: {res:?}"), } } pub(crate) async fn request_remote_install(&self, url: String) -> super::Result { trace!("Remote install requested"); - let (sndr, recv) = async_channel::bounded(1); - self.message.send((Message::RemoteInstall(url), sndr)).await?; - match recv.recv().await { - Ok(Ok(Response::RemoteInstall(resp))) => Ok(resp), - Ok(Err(e)) => Err(e), - Err(async_channel::RecvError) => { - Err(crate::states::TransitionError::CommunicationFailed) - } - Ok(Ok(res)) => unreachable!("Unexpected response: {res:?}"), + match self.request(Message::RemoteInstall(url)).await? { + Response::RemoteInstall(resp) => Ok(resp), + res => unreachable!("Unexpected response: {res:?}"), } } }