Merge pull request #3661 from jarhodes314/refactor/dynamic-mqtt-subs

refactor: support dynamic MQTT subscriptions in `mqtt-channel`/`tedge-mqtt-ext`

Author: jarhodes314

Cargo.lock

@@ -4,6 +4,7 @@ use crate::MqttError;
 use crate::MqttMessage;
 use crate::PubChannel;
 use crate::SubChannel;
+use crate::TopicFilter;
 use futures::channel::mpsc;
 use futures::channel::oneshot;
 use futures::SinkExt;
@@ -16,10 +17,12 @@ use rumqttc::Event;
 use rumqttc::EventLoop;
 use rumqttc::Outgoing;
 use rumqttc::Packet;
+use rumqttc::SubscribeFilter;
 use std::collections::HashSet;
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
+use std::sync::Mutex;
 use std::time::Duration;
 use tokio::sync::OwnedSemaphorePermit;
 use tokio::sync::Semaphore;
@@ -38,6 +41,76 @@ pub struct Connection {
 
     /// A channel to notify that all the published messages have been actually published.
     pub pub_done: oneshot::Receiver<()>,
+
+    pub subscriptions: SubscriberHandle,
+}
+
+#[derive(Clone)]
+/// A client for changing the subscribed topics
+pub struct SubscriberHandle {
+    client: AsyncClient,
+    pub(crate) subscriptions: Arc<Mutex<TopicFilter>>,
+}
+
+impl SubscriberHandle {
+    pub fn new(client: AsyncClient, subscriptions: Arc<Mutex<TopicFilter>>) -> Self {
+        Self {
+            client,
+            subscriptions,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+pub trait SubscriberOps {
+    async fn subscribe_many(
+        &self,
+        topics: impl IntoIterator<Item = String> + Send,
+    ) -> Result<(), MqttError>;
+    async fn unsubscribe_many(
+        &self,
+        topics: impl IntoIterator<Item = String> + Send,
+    ) -> Result<(), MqttError>;
+}
+
+#[async_trait::async_trait]
+impl SubscriberOps for SubscriberHandle {
+    async fn subscribe_many(
+        &self,
+        topics: impl IntoIterator<Item = String> + Send,
+    ) -> Result<(), MqttError> {
+        let topics = topics.into_iter().collect::<Vec<_>>();
+        {
+            let mut subs = self.subscriptions.lock().unwrap();
+            for topic in &topics {
+                subs.add(topic)?;
+            }
+        }
+        self.client
+            .subscribe_many(topics.into_iter().map(|path| SubscribeFilter {
+                path,
+                qos: rumqttc::QoS::AtLeastOnce,
+            }))
+            .await?;
+        Ok(())
+    }
+
+    async fn unsubscribe_many(
+        &self,
+        topics: impl IntoIterator<Item = String> + Send,
+    ) -> Result<(), MqttError> {
+        let topics = topics.into_iter().collect::<Vec<_>>();
+        {
+            let mut subs = self.subscriptions.lock().unwrap();
+            for topic in &topics {
+                subs.remove(topic);
+            }
+        }
+        for topic in topics {
+            self.client.unsubscribe(topic).await?;
+        }
+        Ok(())
+    }
 }
 
 impl Connection {
@@ -91,9 +164,15 @@ impl Connection {
         let (published_sender, published_receiver) = mpsc::unbounded();
         let (error_sender, error_receiver) = mpsc::unbounded();
         let (pub_done_sender, pub_done_receiver) = oneshot::channel();
+        let subscriptions = Arc::new(Mutex::new(config.subscriptions.clone()));
 
-        let (mqtt_client, event_loop) =
-            Connection::open(config, received_sender.clone(), error_sender.clone()).await?;
+        let (mqtt_client, event_loop) = Connection::open(
+            config,
+            received_sender.clone(),
+            error_sender.clone(),
+            subscriptions.clone(),
+        )
+        .await?;
         let permits = Arc::new(Semaphore::new(1));
         let permit = permits.clone().acquire_owned().await.unwrap();
         let pub_count = Arc::new(AtomicUsize::new(0));
@@ -106,9 +185,10 @@ impl Connection {
             pub_done_sender,
             permits,
             pub_count.clone(),
+            subscriptions.clone(),
         ));
         tokio::spawn(Connection::sender_loop(
-            mqtt_client,
+            mqtt_client.clone(),
             published_receiver,
             error_sender,
             config.last_will_message.clone(),
@@ -121,6 +201,7 @@ impl Connection {
             published: published_sender,
             errors: error_receiver,
             pub_done: pub_done_receiver,
+            subscriptions: SubscriberHandle::new(mqtt_client, subscriptions),
         })
     }
 
@@ -133,6 +214,7 @@ impl Connection {
         config: &Config,
         mut message_sender: mpsc::UnboundedSender<MqttMessage>,
         mut error_sender: mpsc::UnboundedSender<MqttError>,
+        subscriptions: Arc<Mutex<TopicFilter>>,
     ) -> Result<(AsyncClient, EventLoop), MqttError> {
         const INSECURE_MQTT_PORT: u16 = 1883;
         const SECURE_MQTT_PORT: u16 = 8883;
@@ -160,7 +242,7 @@ impl Connection {
                     };
                     info!(target: "MQTT", "Connection established");
 
-                    let subscriptions = config.subscriptions.filters();
+                    let subscriptions = subscriptions.lock().unwrap().filters();
 
                     // Need check here otherwise it will hang waiting for a SubAck, and none will come when there is no subscription.
                     if subscriptions.is_empty() {
@@ -217,6 +299,7 @@ impl Connection {
         done: oneshot::Sender<()>,
         permits: Arc<Semaphore>,
         pub_count: Arc<AtomicUsize>,
+        subscriptions: Arc<Mutex<TopicFilter>>,
     ) -> Result<(), MqttError> {
         let mut triggered_disconnect = false;
         let mut disconnect_permit = None;
@@ -289,7 +372,7 @@ impl Connection {
                             // If session_name is not provided or if the broker session persistence
                             // is not enabled or working, then re-subscribe
 
-                            let subscriptions = config.subscriptions.filters();
+                            let subscriptions = subscriptions.lock().unwrap().filters();
                             // Need check here otherwise it will hang waiting for a SubAck, and none will come when there is no subscription.
                             if subscriptions.is_empty() {
                                 break;

crates/common/mqtt_channel/src/connection.rs

@@ -474,3 +474,97 @@ async fn test_max_packet_size_validation() -> Result<(), anyhow::Error> {
 
     Ok(())
 }
+
+#[tokio::test]
+async fn dynamic_subscriptions() {
+    // Given an MQTT broker
+    let broker = mqtt_tests::test_mqtt_broker();
+    let mqtt_config = Config::default().with_port(broker.port);
+
+    // A client subscribes to a topic on connect
+    let topic = uniquify!("a/test/topic");
+    let topic2 = uniquify!("a/test/topic/2");
+
+    // Publish a retain message before any client connects
+    broker
+        .publish_with_opts(topic2, "msg 0", QoS::AtLeastOnce, true)
+        .await
+        .unwrap();
+
+    let mqtt_config = mqtt_config.with_session_name(uniquify!("test_client"));
+    let mut con = Connection::new(&mqtt_config).await.unwrap();
+
+    assert_eq!(
+        *con.subscriptions.subscriptions.lock().unwrap(),
+        TopicFilter::empty()
+    );
+    con.subscriptions
+        .subscribe_many([topic.to_owned()])
+        .await
+        .unwrap();
+
+    // Assert we have added the newly subscribed topic to the list we need to
+    // resubscribe to if the connection drops
+    assert_eq!(
+        *con.subscriptions.subscriptions.lock().unwrap(),
+        TopicFilter::new_unchecked(topic)
+    );
+
+    broker
+        .publish_with_opts(topic, "msg 1", QoS::AtLeastOnce, true)
+        .await
+        .unwrap();
+
+    // Assert just against the payload since the retain flag may be true or
+    // false depending on the ordering of the publish/subscribe messages (retain
+    // flag is only set on an incoming message if the message was published
+    // before we subscribe)
+    assert_payload_received(&mut con, "msg 1").await;
+
+    // Unsubscribe from one topic and subscribe to another
+    con.subscriptions
+        .unsubscribe_many([topic.to_owned()])
+        .await
+        .unwrap();
+    con.subscriptions
+        .subscribe_many([topic2.to_owned()])
+        .await
+        .unwrap();
+
+    // Check we've updated the subscription list with both those changes
+    assert_eq!(
+        *con.subscriptions.subscriptions.lock().unwrap(),
+        TopicFilter::new_unchecked(topic2)
+    );
+
+    // We expect now to receive the retained message that has been published before the connection created
+    // on the topic2 which the connection only subscribed to now.
+    assert_payload_received(&mut con, "msg 0").await;
+
+    // Wait for the new subscription to be enacted
+    broker
+        .publish_with_opts(topic2, "msg 2", QoS::AtLeastOnce, true)
+        .await
+        .unwrap();
+
+    assert_payload_received(&mut con, "msg 2").await;
+
+    // At this point, we are unsubscribed from topic, and therefore we shouldn't receive the messages
+    broker
+        .publish_with_opts(topic, "msg 3", QoS::AtLeastOnce, true)
+        .await
+        .unwrap();
+    broker
+        .publish_with_opts(topic2, "msg 4", QoS::AtLeastOnce, true)
+        .await
+        .unwrap();
+
+    assert_payload_received(&mut con, "msg 4").await;
+}
+
+async fn assert_payload_received(con: &mut Connection, payload: &'static str) {
+    match next_message(&mut con.received).await {
+        MaybeMessage::Next(msg) => assert_eq!(msg.payload.as_str().unwrap(), payload),
+        not_msg => panic!("Expected message to be received, got {not_msg:?}"),
+    }
+}

crates/common/mqtt_channel/src/tests.rs

@@ -169,7 +169,7 @@ impl TopicFilter {
     }
 
     /// The list of `SubscribeFilter` expected by `mqttc`
-    pub(crate) fn filters(&self) -> Vec<SubscribeFilter> {
+    pub fn filters(&self) -> Vec<SubscribeFilter> {
         let qos = self.qos;
         self.patterns
             .iter()
@@ -183,6 +183,14 @@ impl TopicFilter {
     pub fn patterns(&self) -> &Vec<String> {
         &self.patterns
     }
+
+    pub fn remove(&mut self, topic: &str) -> Option<String> {
+        if let Some((index, _)) = self.patterns.iter().enumerate().find(|(_, p)| *p == topic) {
+            Some(self.patterns.swap_remove(index))
+        } else {
+            None
+        }
+    }
 }
 
 impl TryInto<Topic> for &str {

crates/common/mqtt_channel/src/topics.rs

@@ -27,6 +27,7 @@ tracing = { workspace = true }
 [dev-dependencies]
 futures = { workspace = true }
 mqtt_tests = { path = "../../tests/mqtt_tests" }
+proptest = { workspace = true }
 
 [lints]
 workspace = true

crates/extensions/tedge_mqtt_ext/Cargo.toml

@@ -0,0 +1,26 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc c2af50542642d11276b3414120dba5d319b7e79bc5c4b593b0f0e3e551bb216e # shrinks to subscriptions = ["/#", "+/+"]
+cc 35ff8c20ac897723b2f45738d88b5c3fdfb3e5f1dbdbfd3a170a70913ebc368d # shrinks to subscriptions = ["+/a/a", "/#"]
+cc c305db36e62ac105b2011d441fca4e88224159738d637aae88679fdeb967de3d # shrinks to subscriptions = ["a/+/a"], unsubscriptions = ["a"]
+cc 3faad8e1ecbe0502bf89b9dfed357f774e1912712f66a63943df74a6310f93f1 # shrinks to subscriptions = ["+"], unsubscriptions = ["+/#"]
+cc 8509be48aeb210d202da60447c6fbe0b99d491808274aefe12a2a0e3828426cd # shrinks to subscriptions = [], unsubscriptions = ["+/+/a/+", "a/#"]
+cc e055cb4e6e9c214f73ee66f2080422dbf5ba5fffca89008a304081ceb942e0c2 # shrinks to subscriptions = [], unsubscriptions = ["+/+", "/#"]
+cc 0951d4466b6ca4c30105fee3c2d5fcc047f0133be17b69ee9b12a6ba7358f48a # shrinks to subscriptions = [], unsubscriptions = ["c/+", "c/#"]
+cc 77a623802e2f52b94bf7fcf10040927a15037ada1641a3209594985fadf76d14 # shrinks to subscriptions = ["a/+/a"], unsubscriptions = ["+/a/+"]
+cc ab19cb84992e7141ae210779904cd516408022c7bc85937942030ebec26ebef1 # shrinks to subscriptions = ["+/+/a/+", "a/a/+/a"]
+cc 5637b36f5f378ae9bcc4ed1d442f93368bcd7cc03cfdb052cc98ab942c4aa29d # shrinks to subscriptions = ["c/+/+", "c/a/+", "+/#"]
+cc 5255ea6f0768fd462c21e020972c4221a202345c0b898b502012f557d1542b76 # shrinks to subscriptions = ["/#", "+/a/+"]
+cc 9638071b32ed1148858f0b0a693cac06b2857b38a8181a56154429b181a61fc1 # shrinks to subscriptions = ["+", "", "/#"]
+cc 09fc525d653ff597e455644305a5bcf0a9143b5b50c5e0b668ab69b72a8094f4 # shrinks to subscriptions = ["/#", "", "+"]
+cc 72dbc7091f98279af02f2b825dcd9440526b9d475a0a33222e2a69389581af6e # shrinks to subscriptions = ["c/a", "+/+", "c/+"]
+cc e6fe7dfd3c0771d88dc0e1d8da448d9ec3e5402059423622b36d2dc6020bd607 # shrinks to subscriptions = ["b/a", "+/#", "b/+"]
+cc 7fcaadecf97f9f29c2b119f46c53992ff8b687dee107eec3e6bb775d4cf467e3 # shrinks to subscriptions = ["+/+/a", "+/+/#", "+/#"]
+cc 6d35905cd0547e000e5a89edfe24818f4fa2e2947ff8a40ad19f7cc89f75bb0d # shrinks to subscriptions = ["+/a/#", "b/a/+", "b/+/+"]
+cc 3368368c515c9bbb99afc5c7f7079f3930cd56049c3e164cb9be5ea916e3d183 # shrinks to subscriptions = ["c/+", "c/c", "+/c"]
+cc 83857d88a8be534674f2326f109b667f7076aecb6b8c8bcc0b3d282959fb3a1e # shrinks to subscriptions = ["a/+", "+/+", "+/+/#", "+/+"]
+cc 321d11e2a9be7ffbd83775f65e48763715103d79f1cf6d83938199b25fc66f99 # shrinks to subscriptions = ["+/c/a/+", "b/c/a/c", "b/+/+/c"]