Merge pull request #2631 from didier-wenzek/fix/tedge-agent-restart-o…

…n-self-update

Fix tedge agent restart on self update

crates/core/tedge_actors/src/message_boxes.rs

@@ -178,6 +178,12 @@ impl<Input: Debug> LoggingReceiver<Input> {
     pub fn into_split(self) -> (mpsc::Receiver<Input>, mpsc::Receiver<RuntimeRequest>) {
         (self.receiver.input_receiver, self.receiver.signal_receiver)
     }
+
+    /// Close the input so no new messages can be sent to this receiver
+    pub fn close_input(&mut self) {
+        self.receiver.input_receiver.close();
+        self.receiver.signal_receiver.close();
+    }
 }
 
 #[async_trait]

crates/core/tedge_agent/src/software_manager/actor.rs

@@ -14,6 +14,7 @@ use serde::Deserialize;
 use serde::Serialize;
 use std::path::PathBuf;
 use std::process::Command;
+use std::time::Duration;
 use tedge_actors::fan_in_message_type;
 use tedge_actors::Actor;
 use tedge_actors::LoggingReceiver;
@@ -107,6 +108,8 @@ impl Actor for SoftwareManagerActor {
                 _ = self.handle_request(request, &mut plugins, &operation_logs) => {
                     if let Err(SoftwareManagerError::NotRunningLatestVersion) = Self::detect_self_update() {
                         error!("Tedge-agent is no more running the latest-version => a restart is required");
+                        // Make sure the operation status is properly reported before the restart
+                        tokio::time::sleep(Duration::from_secs(5)).await;
                         return Err(RuntimeError::ActorError(Box::new(SoftwareManagerError::NotRunningLatestVersion)));
                     }
                 }

crates/extensions/tedge_mqtt_ext/src/lib.rs

@@ -134,12 +134,18 @@ impl FromPeers {
         &mut self,
         outgoing_mqtt: &mut mpsc::UnboundedSender<MqttMessage>,
     ) -> Result<(), RuntimeError> {
-        loop {
-            match self.try_recv().await {
-                Ok(Some(message)) => outgoing_mqtt.send(message).await.map_err(Box::new)?,
-                Ok(None) | Err(RuntimeRequest::Shutdown) => break Ok(()),
-            }
+        while let Ok(Some(message)) = self.try_recv().await {
+            outgoing_mqtt.send(message).await.map_err(Box::new)?;
         }
+
+        // On shutdown, first close input so no new messages can be pushed
+        self.input_receiver.close_input();
+
+        // Then, publish all the messages awaiting to be sent over MQTT
+        while let Some(message) = self.recv().await {
+            outgoing_mqtt.send(message).await.map_err(Box::new)?;
+        }
+        Ok(())
     }
 }