Merge pull request #3681 from albinsuresh/fix/3632/track-all-entity-retained-messages-v2

fix: clear all retained entity messages on entity deregistration

Author: albinsuresh

crates/core/tedge_agent/src/agent.rs

@@ -58,6 +58,7 @@ use tedge_log_manager::LogManagerConfig;
 use tedge_log_manager::LogManagerOptions;
 use tedge_mqtt_ext::MqttActorBuilder;
 use tedge_mqtt_ext::MqttConfig;
+use tedge_mqtt_ext::MqttDynamicConnector;
 use tedge_mqtt_ext::TopicFilter;
 use tedge_script_ext::ScriptActor;
 use tedge_signal_ext::SignalActor;
@@ -271,7 +272,7 @@ impl Agent {
         let mut restart_actor_builder = RestartManagerBuilder::new(self.config.restart_config);
 
         // Mqtt actor
-        let mut mqtt_actor_builder = MqttActorBuilder::new(self.config.mqtt_config);
+        let mut mqtt_actor_builder = MqttActorBuilder::new(self.config.mqtt_config.clone());
 
         // Software update actor
         let mut software_update_builder = SoftwareManagerBuilder::new(self.config.sw_update_config);
@@ -387,6 +388,7 @@ impl Agent {
             let entity_store_server = EntityStoreServer::new(
                 entity_store,
                 mqtt_schema.clone(),
+                Box::new(MqttDynamicConnector::new(self.config.mqtt_config)),
                 &mut mqtt_actor_builder,
                 self.config.entity_auto_register,
             );

crates/core/tedge_agent/src/entity_manager/server.rs

@@ -12,13 +12,17 @@ use tedge_api::entity_store::EntityTwinMessage;
 use tedge_api::entity_store::EntityUpdateMessage;
 use tedge_api::entity_store::ListFilters;
 use tedge_api::mqtt_topics::Channel;
+use tedge_api::mqtt_topics::ChannelFilter;
+use tedge_api::mqtt_topics::EntityFilter;
 use tedge_api::mqtt_topics::EntityTopicId;
 use tedge_api::mqtt_topics::MqttSchema;
 use tedge_api::pending_entity_store::RegisteredEntityData;
 use tedge_api::EntityStore;
+use tedge_mqtt_ext::MqttConnector;
 use tedge_mqtt_ext::MqttMessage;
-use tedge_mqtt_ext::QoS;
 use tedge_mqtt_ext::TopicFilter;
+use tokio::time::timeout;
+use tokio::time::Duration;
 use tracing::error;
 
 #[derive(Debug)]
@@ -52,6 +56,7 @@ pub enum EntityStoreResponse {
 pub struct EntityStoreServer {
     entity_store: EntityStore,
     mqtt_schema: MqttSchema,
+    mqtt_connector: Box<dyn MqttConnector>,
     mqtt_publisher: LoggingSender<MqttMessage>,
     entity_auto_register: bool,
 }
@@ -60,6 +65,7 @@ impl EntityStoreServer {
     pub fn new(
         entity_store: EntityStore,
         mqtt_schema: MqttSchema,
+        mqtt_connector: Box<dyn MqttConnector>,
         mqtt_actor: &mut impl MessageSink<MqttMessage>,
         entity_auto_register: bool,
     ) -> Self {
@@ -68,6 +74,7 @@ impl EntityStoreServer {
         Self {
             entity_store,
             mqtt_schema,
+            mqtt_connector,
             mqtt_publisher,
             entity_auto_register,
         }
@@ -306,24 +313,56 @@ impl EntityStoreServer {
 
     async fn deregister_entity(&mut self, topic_id: &EntityTopicId) -> Vec<EntityMetadata> {
         let deleted = self.entity_store.deregister_entity(topic_id);
-        for entity in deleted.iter().rev() {
-            for twin_key in entity.twin_data.keys() {
-                let topic = self.mqtt_schema.topic_for(
-                    &entity.topic_id,
-                    &Channel::EntityTwinData {
-                        fragment_key: twin_key.to_string(),
-                    },
-                );
-                let clear_twin_msg = MqttMessage::new(&topic, "").with_retain();
-
-                self.publish_message(clear_twin_msg).await;
+
+        let mut topics = TopicFilter::empty();
+        for entity in deleted.iter() {
+            for channel_filter in [
+                ChannelFilter::MeasurementMetadata,
+                ChannelFilter::EventMetadata,
+                ChannelFilter::AlarmMetadata,
+                ChannelFilter::Alarm,
+                ChannelFilter::EntityTwinData,
+                ChannelFilter::AnyCommand,
+                ChannelFilter::AnyCommandMetadata,
+            ] {
+                let topic = self
+                    .mqtt_schema
+                    .topics(EntityFilter::Entity(&entity.topic_id), channel_filter);
+                topics.add_all(topic);
+            }
+        }
+
+        // A single connection to retrieve all retained metadata messages for all deleted entities
+        match self.mqtt_connector.connect(topics).await {
+            Ok(mut connection) => {
+                while let Ok(Some(message)) =
+                    timeout(Duration::from_secs(1), connection.next_message()).await
+                {
+                    if message.retain && !message.payload_bytes().is_empty() {
+                        let clear_msg = MqttMessage::new(&message.topic, "").with_retain();
+                        if let Err(err) = self.mqtt_publisher.send(clear_msg).await {
+                            error!(
+                                    "Failed to clear retained message on topic {} while de-registering {} due to: {err}",
+                                    topic_id,
+                                    message.topic
+                                );
+                        }
+                    }
+                }
+
+                connection.disconnect().await;
+            }
+            Err(err) => {
+                error!("Failed to create MQTT connection for clearing entity data: {err}");
             }
+        }
+
+        // Clear the entity metadata of all deleted entities bottom up
+        for entity in deleted.iter().rev() {
             let topic = self
                 .mqtt_schema
                 .topic_for(&entity.topic_id, &Channel::EntityMetadata);
-            let clear_entity_msg = MqttMessage::new(&topic, "")
-                .with_retain()
-                .with_qos(QoS::AtLeastOnce);
+            let clear_entity_msg = MqttMessage::new(&topic, "").with_retain();
 
             self.publish_message(clear_entity_msg).await;
         }