mqtt_channel/

lib.rs

//! Provide channel-based subscription to a MQTT broker.

#![warn(missing_docs)]
#![deny(warnings)]

use futures::{SinkExt, StreamExt};
use std::{collections::HashMap, sync::Arc};

// Re-export some rumqttc::4v5 useful types
pub use rumqttc::v5::mqttbytes::v5 as mqttbytes;
pub use rumqttc::v5::mqttbytes::QoS;

/// Result type for mqtt-channel
pub type Result<T> = std::result::Result<T, Error>;

//  ____  _        _   _     _   _
// / ___|| |_ __ _| |_(_)___| |_(_) ___ ___
// \___ \| __/ _` | __| / __| __| |/ __/ __|
//  ___) | || (_| | |_| \__ \ |_| | (__\__ \
// |____/ \__\__,_|\__|_|___/\__|_|\___|___/
//

/// Statistics of the MQTT Client
#[cfg(feature = "statistics")]
#[derive(Default, Clone)]
pub struct Statistics
{
  /// number of received messages
  pub received_messages: u64,
  /// number of duplicated received messages
  pub received_duplicated_messages: u64,
  /// number of published messages
  pub published_messages: u64,
  /// Number of connection acknowledgement received [CONNACK](https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html#_Toc3901074)
  pub connection_ack: u64,
  /// Number of subscription acknowledgement received
  pub subscription_ack: u64,
  /// Number of publish acknowledgement received [PUBACK](https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html#_Toc3901121)
  pub publish_ack: u64,
  /// Number of publish released received [PUBREL](https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html#_Toc3901141)
  pub publish_release: u64,
  /// Number of publish received received [PUBREC](https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html#_Toc3901131)
  pub publish_received: u64,
  /// Number of publish complete received [PUBCOMP](https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html#_Toc3901151)
  pub publish_complete: u64,
  /// Number of ping response released received [PINGRESP](https://docs.oasis-open.org/mqtt/mqtt/v5.0/os/mqtt-v5.0-os.html#_Toc3901200)
  pub ping_response: u64,
}

//  _____
// | ____|_ __ _ __ ___  _ __
// |  _| | '__| '__/ _ \| '__|
// | |___| |  | | | (_) | |
// |_____|_|  |_|  \___/|_|
//
#[derive(thiserror::Error, Debug)]
#[allow(missing_docs)]
pub enum Error
{
  #[error("Already subscribed to topic {0}.")]
  AlreadySubscribed(String),
  #[error("Not subscribed to topic {0}.")]
  NotSubscribed(String),
  #[error("An internal error occured when sending a request {0}.")]
  RequestChannel(#[from] futures::channel::mpsc::SendError),
  #[error("Could not lock a mutex: {0}.")]
  MutexPoison(String),
  #[error("An error occured in the communication middleware {0}.")]
  Communication(#[from] rumqttc::v5::ClientError),
  #[error("MQTT option error {0}.")]
  OptionError(#[from] rumqttc::v5::OptionError),
  #[cfg(feature = "json")]
  #[error("Error during json serialization {0}")]
  JsonSerialisation(#[from] serde_json::Error),
  #[cfg(feature = "cbor")]
  #[error("Error during cbor serialization {0}")]
  CborSerialisation(#[from] ciborium::ser::Error<std::io::Error>),
}

impl<T> From<std::sync::PoisonError<std::sync::MutexGuard<'_, T>>> for Error
{
  fn from(value: std::sync::PoisonError<std::sync::MutexGuard<'_, T>>) -> Self
  {
    Error::MutexPoison(value.to_string())
  }
}

//  ___       _                        _ _____
// |_ _|_ __ | |_ ___ _ __ _ __   __ _| | ____|_ __ _ __ ___  _ __
//  | || '_ \| __/ _ \ '__| '_ \ / _` | |  _| | '__| '__/ _ \| '__|
//  | || | | | ||  __/ |  | | | | (_| | | |___| |  | | | (_) | |
// |___|_| |_|\__\___|_|  |_| |_|\__,_|_|_____|_|  |_|  \___/|_|

#[derive(thiserror::Error, Debug)]
enum InternalError
{
  #[error("No channel for topic {0}")]
  NoChannelFor(String),
  #[error("An error occured in the communication middleware {0}.")]
  Communication(#[from] rumqttc::v5::ConnectionError),
  #[error("Error when broadcasting message {0}")]
  AsyncBroadcast(String),
  #[error("UTF-8 error {0}")]
  Utf8Error(#[from] std::str::Utf8Error),
}

impl<T> From<async_broadcast::SendError<T>> for InternalError
{
  fn from(value: async_broadcast::SendError<T>) -> Self
  {
    Self::AsyncBroadcast(value.to_string())
  }
}

//  __  __
// |  \/  | ___  ___ ___  __ _  __ _  ___
// | |\/| |/ _ \/ __/ __|/ _` |/ _` |/ _ \
// | |  | |  __/\__ \__ \ (_| | (_| |  __/
// |_|  |_|\___||___/___/\__,_|\__, |\___|
//                             |___/

/// Message received from MQTT
#[derive(Clone)]
pub struct Message<T>
{
  /// Raw data
  pub payload: T,
  /// Properties of the message, check [rumqttc documentation](https://docs.rs/rumqttc/latest/rumqttc/v5/mqttbytes/v5/struct.PublishProperties.html) for more information.
  pub properties: Option<rumqttc::v5::mqttbytes::v5::PublishProperties>,
}

/// Raw message
pub type RawMessage = Message<bytes::Bytes>;

//  ____                            _
// |  _ \ ___  __ _ _   _  ___  ___| |_
// | |_) / _ \/ _` | | | |/ _ \/ __| __|
// |  _ <  __/ (_| | |_| |  __/\__ \ |_
// |_| \_\___|\__, |\__,_|\___||___/\__|
//               |_|

/// Request sent to the internal MQTT tasks
enum Request
{
  Subscribe
  {
    topic: String,
    sender: async_broadcast::Sender<RawMessage>,
    qos: QoS,
  },
  Publish
  {
    topic: String,
    qos: QoS,
    retain: bool,
    payload: bytes::Bytes,
    properties: Option<rumqttc::v5::mqttbytes::v5::PublishProperties>,
  },
}

//   ____ _ _            _   ___
//  / ___| (_) ___ _ __ | |_|_ _|_ __  _ __   ___ _ __
// | |   | | |/ _ \ '_ \| __|| || '_ \| '_ \ / _ \ '__|
// | |___| | |  __/ | | | |_ | || | | | | | |  __/ |
//  \____|_|_|\___|_| |_|\__|___|_| |_|_| |_|\___|_|

#[derive(Default)]
struct ClientInner
{
  senders: HashMap<String, async_broadcast::Sender<RawMessage>>,
  #[cfg(feature = "statistics")]
  statistics: Statistics,
  #[cfg(feature = "debug")]
  last_logged_statistics: Statistics,
}

//   ____ _ _            _   ____        _ _     _
//  / ___| (_) ___ _ __ | |_| __ ) _   _(_) | __| | ___ _ __
// | |   | | |/ _ \ '_ \| __|  _ \| | | | | |/ _` |/ _ \ '__|
// | |___| | |  __/ | | | |_| |_) | |_| | | | (_| |  __/ |
//  \____|_|_|\___|_| |_|\__|____/ \__,_|_|_|\__,_|\___|_|

/// Used to setup a new connection to a MQTT Broker
pub struct ClientBuilder
{
  mqttoptions: rumqttc::v5::MqttOptions,
  cap: usize,
}

impl ClientBuilder
{
  /// Set the capacity of the reception buffer
  pub fn set_capacity(mut self, cap: usize) -> Self
  {
    self.cap = cap;
    self
  }

  async fn handle_events(
    event: std::result::Result<rumqttc::v5::Event, rumqttc::v5::ConnectionError>,
    client_inner: &Arc<futures::lock::Mutex<ClientInner>>,
  ) -> std::result::Result<(), InternalError>
  {
    match event?
    {
      rumqttc::v5::Event::Incoming(packet) => match packet
      {
        rumqttc::v5::Incoming::Publish(pub_msg) =>
        {
          let topic_str = std::str::from_utf8(pub_msg.topic.as_ref())?;
          #[cfg(feature = "statistics")]
          let mut client_inner = client_inner.lock().await;
          #[cfg(feature = "statistics")]
          {
            client_inner.statistics.received_messages += 1;
            if pub_msg.dup
            {
              client_inner.statistics.received_duplicated_messages += 1;
            }
          }
          #[cfg(not(feature = "statistics"))]
          let client_inner = client_inner.lock().await;

          match client_inner.senders.get(topic_str)
          {
            Some(s) =>
            {
              s.broadcast(RawMessage {
                payload: pub_msg.payload,
                properties: pub_msg.properties,
              })
              .await?;
            }
            None => Err(InternalError::NoChannelFor(topic_str.to_string()))?,
          }
        }
        rumqttc::v5::Incoming::ConnAck(_) =>
        {
          #[cfg(feature = "statistics")]
          {
            client_inner.lock().await.statistics.connection_ack += 1;
          }
        }
        rumqttc::v5::Incoming::SubAck(_) =>
        {
          #[cfg(feature = "statistics")]
          {
            client_inner.lock().await.statistics.subscription_ack += 1;
          }
        }
        rumqttc::v5::Incoming::PubAck(_) =>
        {
          #[cfg(feature = "statistics")]
          {
            client_inner.lock().await.statistics.publish_ack += 1;
          }
        }
        rumqttc::v5::Incoming::PubRel(_) =>
        {
          #[cfg(feature = "statistics")]
          {
            client_inner.lock().await.statistics.publish_release += 1;
          }
        }
        rumqttc::v5::Incoming::PubRec(_) =>
        {
          #[cfg(feature = "statistics")]
          {
            client_inner.lock().await.statistics.publish_received += 1;
          }
        }
        rumqttc::v5::Incoming::PubComp(_) =>
        {
          #[cfg(feature = "statistics")]
          {
            client_inner.lock().await.statistics.publish_complete += 1;
          }
        }
        rumqttc::v5::Incoming::PingResp(_) =>
        {
          #[cfg(feature = "statistics")]
          {
            client_inner.lock().await.statistics.ping_response += 1;
          }
        }
        _ =>
        {
          log::error!("Incoming unhandled packet {packet:?}");
        }
      },
      rumqttc::v5::Event::Outgoing(_) =>
      { /* ignore outgoing event */ }
    }
    #[cfg(feature = "debug")]
    {
      let mut client_inner = client_inner.lock().await;

      if client_inner.statistics.received_messages
        > client_inner.last_logged_statistics.received_messages + 1000
      {
        client_inner.last_logged_statistics = client_inner.statistics.clone();
        log::info!(
          "RECVMSG {} (DUP {}) PUBMSG {} CONACK {} PUBACK {} PUBREL {} PUBREC {} PUBCOMP {} PINGRESP {}",
          client_inner.last_logged_statistics.received_messages,
          client_inner.last_logged_statistics.received_duplicated_messages,
          client_inner.last_logged_statistics.published_messages,
          client_inner.last_logged_statistics.connection_ack,
          client_inner.last_logged_statistics.publish_ack,
          client_inner.last_logged_statistics.publish_release,
          client_inner.last_logged_statistics.publish_received,
          client_inner.last_logged_statistics.publish_complete,
          client_inner.last_logged_statistics.ping_response
        );
      }
    }

    Ok(())
  }
  /// Start the connection
  pub fn start(self) -> (Client, impl std::future::Future<Output = ()>)
  {
    let (client, mut eventloop) = rumqttc::v5::AsyncClient::new(self.mqttoptions, self.cap);

    let (request_sender, mut request_receiver) = futures::channel::mpsc::channel(self.cap);

    let fut = async move {
      let client_inner: Arc<futures::lock::Mutex<ClientInner>> = Default::default();
      let fut_mqtt = {
        let client_inner = client_inner.clone();
        async move {
          loop
          {
            if let Err(e) = Self::handle_events(eventloop.poll().await, &client_inner).await
            {
              log::error!("An error occured when handling MQTT event: {e:?}");
            }
          }
        }
      };
      let fut_req = async move {
        loop
        {
          let r = match request_receiver.next().await
          {
            Some(Request::Subscribe { topic, sender, qos }) =>
            {
              client_inner
                .lock()
                .await
                .senders
                .insert(topic.clone(), sender);
              client.subscribe(topic, qos).await
            }
            Some(Request::Publish {
              topic,
              qos,
              retain,
              payload,
              properties,
            }) =>
            {
              #[cfg(feature = "statistics")]
              {
                client_inner.lock().await.statistics.published_messages += 1;
              }
              match properties
              {
                Some(properties) =>
                {
                  client
                    .publish_with_properties(topic, qos, retain, payload, properties)
                    .await
                }
                None => client.publish(topic, qos, retain, payload).await,
              }
            }
            None => Ok(()),
          };
          if let Err(e) = r
          {
            log::error!("An error occured when evaluating requests: {e:?}");
          }
        }
      };
      futures::join!(fut_mqtt, fut_req);
    };

    (
      Client {
        request_sender,
        subscriptions: Default::default(),
      },
      fut,
    )
  }
}

//   ____ _ _            _
//  / ___| (_) ___ _ __ | |_
// | |   | | |/ _ \ '_ \| __|
// | |___| | |  __/ | | | |_
//  \____|_|_|\___|_| |_|\__|

/// Handle to the connection
#[derive(Clone)]
pub struct Client
{
  subscriptions: std::sync::Arc<
    futures::lock::Mutex<HashMap<String, async_broadcast::InactiveReceiver<RawMessage>>>,
  >,
  request_sender: futures::channel::mpsc::Sender<Request>,
}

macro_rules! _create_raw_subscription {
  ($request_sender: expr, $subs: ident, $topic: ident, $qos: ident, $cap: ident) => {{
    let (sender, r) = async_broadcast::broadcast::<RawMessage>($cap);
    $subs.insert($topic.clone(), r.clone().deactivate());
    $request_sender
      .send(Request::Subscribe {
        topic: $topic,
        sender,
        qos: $qos,
      })
      .await?;
    Ok(r)
  }};
}

impl Client
{
  /// Create a new MQTT Client, with the given name to connect at the given and MQTT broker
  pub fn build(node_id: impl Into<String>, hostname: impl Into<String>, port: u16)
    -> ClientBuilder
  {
    let mut mqttoptions = rumqttc::v5::MqttOptions::new(node_id, hostname, port);
    mqttoptions.set_keep_alive(std::time::Duration::from_secs(5));

    ClientBuilder {
      mqttoptions,
      cap: 1000,
    }
  }
  /// Create a new MQTT Client from the given url of the form `mqtt://example.com:1883?client_id=123`
  pub fn build_from_url(url: impl Into<String>) -> Result<ClientBuilder>
  {
    let mut mqttoptions = rumqttc::v5::MqttOptions::parse_url(url)?;
    mqttoptions.set_keep_alive(std::time::Duration::from_secs(5));

    Ok(ClientBuilder {
      mqttoptions,
      cap: 1000,
    })
  }
  /// Publish a message on a topic
  pub async fn publish_raw(
    mut self,
    topic: impl Into<String>,
    qos: QoS,
    retain: bool,
    payload: impl Into<bytes::Bytes>,
    properties: Option<rumqttc::v5::mqttbytes::v5::PublishProperties>,
  ) -> Result<()>
  {
    self
      .request_sender
      .send(Request::Publish {
        topic: topic.into(),
        qos,
        retain,
        payload: payload.into(),
        properties,
      })
      .await?;
    Ok(())
  }
  /// Publish a message on a topic, serialized to JSON
  #[cfg(feature = "json")]
  pub fn publish_json<T: serde::Serialize>(
    self,
    topic: impl Into<String>,
    qos: QoS,
    retain: bool,
    payload: T,
    properties: Option<rumqttc::v5::mqttbytes::v5::PublishProperties>,
  ) -> Result<impl std::future::Future<Output = Result<()>>>
  {
    let properties = match properties
    {
      None => rumqttc::v5::mqttbytes::v5::PublishProperties {
        content_type: Some("application/json").map(str::to_string),
        ..Default::default()
      },
      Some(props) =>
      {
        let mut props = props.clone();
        props.content_type = Some("application/json".to_string());
        props
      }
    };
    let payload_raw = serde_json::to_string(&payload)?;

    Ok(self.publish_raw(topic, qos, retain, payload_raw, Some(properties)))
  }
  /// Publish a message on a topic, serialized to JSON
  #[cfg(feature = "cbor")]
  pub fn publish_cbor<T: serde::Serialize>(
    self,
    topic: impl Into<String>,
    qos: QoS,
    retain: bool,
    payload: T,
    properties: Option<rumqttc::v5::mqttbytes::v5::PublishProperties>,
  ) -> Result<impl std::future::Future<Output = Result<()>>>
  {
    let properties = match properties
    {
      None => rumqttc::v5::mqttbytes::v5::PublishProperties {
        content_type: Some("application/cbor").map(str::to_string),
        ..Default::default()
      },
      Some(props) =>
      {
        let mut props = props.clone();
        props.content_type = Some("application/cbor".to_string());
        props
      }
    };
    let mut payload_raw = Vec::<u8>::new();
    ciborium::into_writer(&payload, &mut payload_raw)?;

    Ok(self.publish_raw(topic, qos, retain, payload_raw, Some(properties)))
  }
  /// Create a raw subscription
  pub async fn create_raw_subscription(
    mut self,
    topic: impl Into<String>,
    qos: QoS,
    cap: usize,
  ) -> Result<async_broadcast::Receiver<RawMessage>>
  {
    let mut subs = self.subscriptions.lock().await;
    let topic = topic.into();
    match subs.get(&topic)
    {
      Some(_) => Err(Error::AlreadySubscribed(topic)),
      None => _create_raw_subscription!(self.request_sender, subs, topic, qos, cap),
    }
  }
  /// Get an existing subscription
  pub fn get_raw_subscription(
    &self,
    topic: impl Into<String>,
  ) -> Result<async_broadcast::Receiver<RawMessage>>
  {
    let topic = topic.into();
    futures::executor::block_on(self.subscriptions.lock())
      .get(&topic)
      .ok_or_else(|| Error::NotSubscribed(topic))
      .map(|x| x.activate_cloned())
  }
  /// Get an existing subscription, or create a new one with the given QoS
  /// This function is guaranteed to return, however if the subscription
  /// already exists, the QoS is ignored.
  pub async fn get_or_create_raw_subscription(
    mut self,
    topic: impl Into<String>,
    qos: QoS,
    cap: usize,
  ) -> Result<async_broadcast::Receiver<RawMessage>>
  {
    let mut subs = self.subscriptions.lock().await;
    let topic = topic.into();
    match subs.get(&topic)
    {
      Some(s) => Ok(s.activate_cloned()),
      None => _create_raw_subscription!(self.request_sender, subs, topic, qos, cap),
    }
  }
  /// Create a JSON subscription, aka, topic will be parsed using serde_json.
  #[cfg(feature = "json")]
  pub async fn get_or_create_json_subscription<T>(
    self,
    topic: impl Into<String>,
    qos: QoS,
    cap: usize,
  ) -> Result<impl futures::Stream<Item = Message<T>>>
  where
    for<'de> T: serde::Deserialize<'de>,
  {
    use bytes::Buf;
    let raw_s = self.get_or_create_raw_subscription(topic, qos, cap).await?;
    let map_s = raw_s.filter_map(|msg| async move {
      let payload_r = serde_json::from_slice::<T>(msg.payload.chunk());
      match payload_r
      {
        Ok(payload) => Some(Message::<T> {
          payload,
          properties: msg.properties,
        }),
        Err(e) =>
        {
          log::error!("Failed to parse JSON message: {e}");
          None
        }
      }
    });
    Ok(map_s)
  }
}

#[cfg(test)]
mod tests
{
  #[test]
  fn it_works()
  {
    assert_eq!(4, 4);
  }
}