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/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/address.rs b/updatehub/src/states/machine/address.rs index 9d63fb68..747ae96e 100644 --- a/updatehub/src/states/machine/address.rs +++ b/updatehub/src/states/machine/address.rs @@ -49,20 +49,33 @@ 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 } } 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), - 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:?}"), } } @@ -70,22 +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), - 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), - res => unreachable!("Unexpected response: {:?}", res), + match self.request(Message::AbortDownload).await? { + Response::AbortDownload(resp) => Ok(resp), + res => unreachable!("Unexpected response: {res:?}"), } } @@ -94,23 +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), - 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), - res => unreachable!("Unexpected response: {:?}", res), + match self.request(Message::RemoteInstall(url)).await? { + Response::RemoteInstall(resp) => Ok(resp), + res => unreachable!("Unexpected response: {res:?}"), } } } diff --git a/updatehub/src/states/machine/mod.rs b/updatehub/src/states/machine/mod.rs index 7e597bdf..6973e230 100644 --- a/updatehub/src/states/machine/mod.rs +++ b/updatehub/src/states/machine/mod.rs @@ -4,12 +4,17 @@ mod address; +#[cfg(test)] +mod tests; + 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, @@ -83,15 +88,68 @@ 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 + } + + /// 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, _)) => {} } } } @@ -128,7 +186,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 +195,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 +231,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 +253,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 +280,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() @@ -251,10 +319,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; @@ -265,39 +336,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; } } } @@ -305,21 +385,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..e157eaec --- /dev/null +++ b/updatehub/src/states/machine/tests.rs @@ -0,0 +1,143 @@ +// 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); +} + +/// 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); +}