Repository: quartiq/minimq
Pull request #144 (2023-08-23 11:06:09)

Refactoring publication to support zero-copy serialization
This PR is designed to allow an end user to provide a `Publication` message with a yet-to-be-serialized payload. This then allows them to serialize their payload _directly_ into the MQTT TX buffer, which eliminates the need for them serializing and then copying it over into the MQTT payload.

Example:
```rs
client.publish(Publication::new_deferred(|buf| settings.get(key, buf)));
```


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index bfd058c..c28bf80 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@ This document describes the changes to Minimq between releases.
  to a new `Config` structure that is supplied to the `Minimq::new()` constructor to simplify the
  client.
 
+
 ## Added
 * Support for subscribing at `QoS::ExactlyOnce`
 * Support for downgrading the `QoS` to the maximum permitted by the server
diff --git a/src/de/received_packet.rs b/src/de/received_packet.rs
index b7fa9ae..27817d3 100644
--- a/src/de/received_packet.rs
+++ b/src/de/received_packet.rs
@@ -14,7 +14,7 @@ use serde::Deserialize;
 #[derive(Debug)]
 pub enum ReceivedPacket<'a> {
     ConnAck(ConnAck<'a>),
-    Publish(Pub<'a>),
+    Publish(Pub<'a, &'a [u8]>),
     PubAck(PubAck<'a>),
     SubAck(SubAck<'a>),
     PubRel(PubRel<'a>),
@@ -83,7 +83,7 @@ impl<'de> serde::de::Visitor<'de> for ControlPacketVisitor {
 
                 let properties = seq.next_element()?.unwrap();
 
-                let publish = Pub {
+                let publish: Pub<'_, &[u8]> = Pub {
                     topic,
                     packet_id,
                     properties,
diff --git a/src/lib.rs b/src/lib.rs
index 9992204..9b00ab1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -85,7 +85,7 @@ mod will;
 pub use broker::Broker;
 pub use config::Config;
 pub use properties::Property;
-pub use publication::Publication;
+pub use publication::{DeferredPublication, Publication};
 pub use reason_codes::ReasonCode;
 pub use will::Will;
 
@@ -152,6 +152,29 @@ pub enum ProtocolError {
     Deserialization(crate::de::Error),
 }
 
+#[derive(Debug, PartialEq)]
+pub enum PubError<T, E> {
+    Error(Error<T>),
+    Serialization(E),
+}
+
+impl<T, E> From<crate::ser::PubError<E>> for PubError<T, E> {
+    fn from(e: crate::ser::PubError<E>) -> Self {
+        match e {
+            crate::ser::PubError::Other(e) => crate::PubError::Serialization(e),
+            crate::ser::PubError::Error(e) => crate::PubError::Error(crate::Error::Minimq(
+                crate::MinimqError::Protocol(ProtocolError::from(e)),
+            )),
+        }
+    }
+}
+
+impl<T, E> From<Error<T>> for PubError<T, E> {
+    fn from(e: Error<T>) -> Self {
+        Self::Error(e)
+    }
+}
+
 impl From<crate::ser::Error> for ProtocolError {
     fn from(err: crate::ser::Error) -> Self {
         ProtocolError::Serialization(err)
diff --git a/src/message_types.rs b/src/message_types.rs
index 04afa5a..bfc26ca 100644
--- a/src/message_types.rs
+++ b/src/message_types.rs
@@ -43,9 +43,8 @@ impl<'a> ControlPacket for ConnAck<'a> {
     const MESSAGE_TYPE: MessageType = MessageType::ConnAck;
 }
 
-impl<'a> ControlPacket for Pub<'a> {
-    const MESSAGE_TYPE: MessageType = MessageType::Publish;
-    fn fixed_header_flags(&self) -> u8 {
+impl<'a, P: crate::publication::ToPayload> Pub<'a, P> {
+    pub fn fixed_header_flags(&self) -> u8 {
         *0u8.set_bits(1..=2, self.qos as u8)
             .set_bit(0, self.retain == Retain::Retained)
     }
diff --git a/src/mqtt_client.rs b/src/mqtt_client.rs
index 4ac1389..77bcdc3 100644
--- a/src/mqtt_client.rs
+++ b/src/mqtt_client.rs
@@ -380,7 +380,10 @@ impl<'buf, TcpStack: TcpClientStack, Clock: embedded_time::Clock, Broker: crate:
     /// # Args
     /// * `publish` - The publication to generate. See [Publication] for a builder pattern to
     /// generate a message.
-    pub fn publish(&mut self, mut publish: Pub<'_>) -> Result<(), Error<TcpStack::Error>> {
+    pub fn publish<P: crate::publication::ToPayload>(
+        &mut self,
+        mut publish: Pub<'_, P>,
+    ) -> Result<(), crate::PubError<TcpStack::Error, P::Error>> {
         // If we are not yet connected to the broker, we can't transmit a message.
         if !self.is_connected() {
             return Ok(());
@@ -393,7 +396,7 @@ impl<'buf, TcpStack: TcpClientStack, Clock: embedded_time::Clock, Broker: crate:
         }
 
         if !self.can_publish(publish.qos) {
-            return Err(Error::NotReady);
+            return Err(crate::PubError::Error(Error::NotReady));
         }
 
         publish.dup = false;
@@ -406,13 +409,14 @@ impl<'buf, TcpStack: TcpClientStack, Clock: embedded_time::Clock, Broker: crate:
                 .replace(self.sm.context_mut().session_state.get_packet_identifier());
         }
 
-        let packet = self.network.send_packet(&publish)?;
+        let packet = self.network.send_pub(&publish)?;
 
         if let Some(id) = publish.packet_id {
             let context = self.sm.context_mut();
             context
                 .session_state
-                .handle_publish(publish.qos, id, packet)?;
+                .handle_publish(publish.qos, id, packet)
+                .map_err(|e| Error::Minimq(e.into()))?;
             context.send_quota = context.send_quota.checked_sub(1).unwrap();
         }
 
diff --git a/src/network_manager.rs b/src/network_manager.rs
index 0af804e..7a7bda1 100644
--- a/src/network_manager.rs
+++ b/src/network_manager.rs
@@ -5,7 +5,7 @@
 //! simple ownership semantics of reading and writing to the network stack. This allows the network
 //! stack to be used to transmit buffers that may be stored internally in other structs without
 //! violating Rust's borrow rules.
-use crate::message_types::ControlPacket;
+use crate::{message_types::ControlPacket, packets::Pub};
 use embedded_nal::{nb, SocketAddr, TcpClientStack, TcpError};
 use serde::Serialize;
 
@@ -146,6 +146,24 @@ where
         Ok(&self.tx_buffer[offset..][..len])
     }
 
+    pub fn send_pub<P: crate::publication::ToPayload>(
+        &mut self,
+        pub_packet: &Pub<'_, P>,
+    ) -> Result<&[u8], crate::PubError<TcpStack::Error, P::Error>> {
+        // If there's an unfinished write pending, it's invalid to try to write new data. The
+        // previous write must first be completed.
+        assert!(self.pending_write.is_none());
+
+        crate::info!("Sending: {:?}", pub_packet);
+        let (offset, packet) =
+            crate::ser::MqttSerializer::pub_to_buffer_meta(self.tx_buffer, pub_packet)?;
+
+        let len = packet.len();
+
+        self.commit_write(offset, len)?;
+        Ok(&self.tx_buffer[offset..][..len])
+    }
+
     /// Finish writing an MQTT control packet to the interface if one exists.
     pub fn finish_write(&mut self) -> Result<(), Error<TcpStack::Error>> {
         if let Some((head, tail)) = self.pending_write.take() {
diff --git a/src/packets.rs b/src/packets.rs
index 1d23c7d..9f2f838 100644
--- a/src/packets.rs
+++ b/src/packets.rs
@@ -73,8 +73,8 @@ pub struct ConnAck<'a> {
 }
 
 /// An MQTT PUBLISH packet, containing data to be sent or received.
-#[derive(Debug, Serialize)]
-pub struct Pub<'a> {
+#[derive(Serialize)]
+pub struct Pub<'a, P: crate::publication::ToPayload> {
     /// The topic that the message was received on.
     pub topic: Utf8String<'a>,
 
@@ -85,7 +85,8 @@ pub struct Pub<'a> {
     pub properties: Properties<'a>,
 
     /// The message to be transmitted.
-    pub payload: &'a [u8],
+    #[serde(skip)]
+    pub payload: P,
 
     /// Specifies whether or not the message should be retained on the broker.
     #[serde(skip)]
@@ -100,6 +101,20 @@ pub struct Pub<'a> {
     pub dup: bool,
 }
 
+impl<'a, P: crate::publication::ToPayload> core::fmt::Debug for Pub<'a, P> {
+    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
+        f.debug_struct("Pub")
+            .field("topic", &self.topic)
+            .field("packet_id", &self.packet_id)
+            .field("properties", &self.properties)
+            .field("retain", &self.retain)
+            .field("qos", &self.qos)
+            .field("dup", &self.dup)
+            .field("payload", &"<deferred>")
+            .finish()
+    }
+}
+
 /// An MQTT SUBSCRIBE control packet
 #[derive(Debug, Serialize)]
 pub struct Subscribe<'a> {
@@ -255,7 +270,7 @@ mod tests {
         };
 
         let mut buffer: [u8; 900] = [0; 900];
-        let message = MqttSerializer::to_buffer(&mut buffer, &publish).unwrap();
+        let message = MqttSerializer::pub_to_buffer(&mut buffer, &publish).unwrap();
 
         assert_eq!(message, good_publish);
     }
@@ -282,7 +297,7 @@ mod tests {
         };
 
         let mut buffer: [u8; 900] = [0; 900];
-        let message = MqttSerializer::to_buffer(&mut buffer, &publish).unwrap();
+        let message = MqttSerializer::pub_to_buffer(&mut buffer, &publish).unwrap();
 
         assert_eq!(message, good_publish);
     }
@@ -334,7 +349,7 @@ mod tests {
         };
 
         let mut buffer: [u8; 900] = [0; 900];
-        let message = MqttSerializer::to_buffer(&mut buffer, &publish).unwrap();
+        let message = MqttSerializer::pub_to_buffer(&mut buffer, &publish).unwrap();
 
         assert_eq!(message, good_publish);
     }
@@ -366,7 +381,7 @@ mod tests {
         };
 
         let mut buffer: [u8; 900] = [0; 900];
-        let message = MqttSerializer::to_buffer(&mut buffer, &publish).unwrap();
+        let message = MqttSerializer::pub_to_buffer(&mut buffer, &publish).unwrap();
 
         assert_eq!(message, good_publish);
     }
diff --git a/src/publication.rs b/src/publication.rs
index 7509609..ffadfe2 100644
--- a/src/publication.rs
+++ b/src/publication.rs
@@ -5,6 +5,55 @@ use crate::{
     ProtocolError, QoS, Retain,
 };
 
+pub trait ToPayload {
+    type Error;
+    fn serialize(&self, buffer: &mut [u8]) -> Result<usize, Self::Error>;
+}
+
+impl<'a> ToPayload for &'a [u8] {
+    type Error = ();
+
+    fn serialize(&self, buffer: &mut [u8]) -> Result<usize, Self::Error> {
+        if buffer.len() < self.len() {
+            return Err(());
+        }
+        buffer[..self.len()].copy_from_slice(self);
+        Ok(self.len())
+    }
+}
+
+impl<const N: usize> ToPayload for [u8; N] {
+    type Error = ();
+
+    fn serialize(&self, buffer: &mut [u8]) -> Result<usize, ()> {
+        (&self[..]).serialize(buffer)
+    }
+}
+impl<const N: usize> ToPayload for &[u8; N] {
+    type Error = ();
+
+    fn serialize(&self, buffer: &mut [u8]) -> Result<usize, ()> {
+        (&self[..]).serialize(buffer)
+    }
+}
+
+pub struct DeferredPublication<E, F: Fn(&mut [u8]) -> Result<usize, E>> {
+    func: F,
+}
+
+impl<E, F: Fn(&mut [u8]) -> Result<usize, E>> DeferredPublication<E, F> {
+    pub fn new<'a>(func: F) -> Publication<'a, Self> {
+        Publication::new(Self { func })
+    }
+}
+
+impl<E, F: Fn(&mut [u8]) -> Result<usize, E>> ToPayload for DeferredPublication<E, F> {
+    type Error = E;
+    fn serialize(&self, buffer: &mut [u8]) -> Result<usize, E> {
+        (self.func)(buffer)
+    }
+}
+
 /// Builder pattern for generating MQTT publications.
 ///
 /// # Note
@@ -17,17 +66,17 @@ use crate::{
 /// It is expected that the user provide a topic either by directly specifying a publication topic
 /// in [Publication::topic], or by parsing a topic from the [Property::ResponseTopic] property
 /// contained within received properties by using the [Publication::reply] API.
-pub struct Publication<'a> {
+pub struct Publication<'a, P: ToPayload> {
     topic: Option<&'a str>,
     properties: Properties<'a>,
     qos: QoS,
-    payload: &'a [u8],
+    payload: P,
     retain: Retain,
 }
 
-impl<'a> Publication<'a> {
+impl<'a, P: ToPayload> Publication<'a, P> {
     /// Construct a new publication with a payload.
-    pub fn new(payload: &'a [u8]) -> Self {
+    pub fn new(payload: P) -> Self {
         Self {
             payload,
             qos: QoS::AtMostOnce,
@@ -121,7 +170,7 @@ impl<'a> Publication<'a> {
     /// # Returns
     /// The message to be published if a publication topic was specified. If no publication topic
     /// was identified, an error is returned.
-    pub fn finish(self) -> Result<Pub<'a>, ProtocolError> {
+    pub fn finish(self) -> Result<Pub<'a, P>, ProtocolError> {
         Ok(Pub {
             topic: Utf8String(self.topic.ok_or(ProtocolError::NoTopic)?),
             properties: self.properties,
diff --git a/src/ser/mod.rs b/src/ser/mod.rs
index aaae9fe..82ee09b 100644
--- a/src/ser/mod.rs
+++ b/src/ser/mod.rs
@@ -26,8 +26,11 @@
 //! * Structs
 //!
 //! Other types are explicitly not implemented and there is no plan to implement them.
-use crate::message_types::{ControlPacket, MessageType};
 use crate::varint::VarintBuffer;
+use crate::{
+    message_types::{ControlPacket, MessageType},
+    packets::Pub,
+};
 use bit_field::BitField;
 use serde::Serialize;
 use varint_rs::VarintWriter;
@@ -46,6 +49,12 @@ pub enum Error {
     Custom,
 }
 
+#[derive(Debug, Copy, Clone, PartialEq)]
+pub enum PubError<E> {
+    Error(Error),
+    Other(E),
+}
+
 impl serde::ser::StdError for Error {}
 
 impl serde::ser::Error for Error {
@@ -109,6 +118,37 @@ impl<'a> MqttSerializer<'a> {
         Ok((offset, packet))
     }
 
+    pub fn pub_to_buffer_meta<P: crate::publication::ToPayload>(
+        buf: &'a mut [u8],
+        pub_packet: &Pub<'a, P>,
+    ) -> Result<(usize, &'a [u8]), PubError<P::Error>> {
+        let mut serializer = crate::ser::MqttSerializer::new(buf);
+        pub_packet
+            .serialize(&mut serializer)
+            .map_err(PubError::Error)?;
+
+        // Next, serialize the payload into the remaining buffer
+        let len = pub_packet
+            .payload
+            .serialize(serializer.remainder())
+            .map_err(PubError::Other)?;
+        serializer.commit(len).map_err(PubError::Error)?;
+
+        // Finally, finish the packet and send it.
+        let (offset, packet) = serializer
+            .finalize(MessageType::Publish, pub_packet.fixed_header_flags())
+            .map_err(PubError::Error)?;
+        Ok((offset, packet))
+    }
+
+    pub fn pub_to_buffer<P: crate::publication::ToPayload>(
+        buf: &'a mut [u8],
+        pub_packet: &Pub<'a, P>,
+    ) -> Result<&'a [u8], PubError<P::Error>> {
+        let (_, packet) = Self::pub_to_buffer_meta(buf, pub_packet)?;
+        Ok(packet)
+    }
+
     /// Encode an MQTT control packet into a buffer.
     ///
     /// # Args
@@ -179,6 +219,27 @@ impl<'a> MqttSerializer<'a> {
 
         Ok(())
     }
+
+    /// Get the remaining buffer to serialize into directly.
+    ///
+    /// # Note
+    /// You must call `commit` after serializing.
+    pub fn remainder(&mut self) -> &mut [u8] {
+        &mut self.buf[self.index..]
+    }
+
+    /// Commit previously-serialized data into the buffer.
+    ///
+    /// # Args
+    /// * `len` - The number of bytes serialized.
+    pub fn commit(&mut self, len: usize) -> Result<(), Error> {
+        if self.buf.len() < (self.index + len) {
+            return Err(Error::InsufficientMemory);
+        }
+
+        self.index += len;
+        Ok(())
+    }
 }
 
 impl<'a> serde::Serializer for &mut MqttSerializer<'a> {
diff --git a/src/types.rs b/src/types.rs
index 78d1ffc..98dacbe 100644
--- a/src/types.rs
+++ b/src/types.rs
@@ -31,6 +31,17 @@ pub struct PropertiesIter<'a> {
     index: usize,
 }
 
+impl<'a> PropertiesIter<'a> {
+    pub fn response_topic(&mut self) -> Option<&'a str> {
+        self.find_map(|prop| {
+            if let Ok(crate::Property::ResponseTopic(topic)) = prop {
+                Some(topic.0)
+            } else {
+                None
+            }
+        })
+    }
+}
 impl<'a> core::iter::Iterator for PropertiesIter<'a> {
     type Item = Result<Property<'a>, ProtocolError>;
 
diff --git a/src/varint.rs b/src/varint.rs
index 9e62c65..6be3db6 100644
--- a/src/varint.rs
+++ b/src/varint.rs
@@ -63,7 +63,7 @@ impl<'de> serde::de::Visitor<'de> for VarintVisitor {
     fn visit_seq<A: serde::de::SeqAccess<'de>>(self, seq: A) -> Result<Self::Value, A::Error> {
         let mut reader = VarintParser {
             seq,
-            _data: core::marker::PhantomData::default(),
+            _data: core::marker::PhantomData,
         };
         Ok(Varint(reader.read_u32_varint()?))
     }
diff --git a/src/will.rs b/src/will.rs
index b5a4650..09027dc 100644
--- a/src/will.rs
+++ b/src/will.rs
@@ -57,7 +57,7 @@ impl<'a> Will<'a> {
         })
     }
 
-    /// Serialize the will contents into a flattened, borrowed buffer.
+    /// Serialize the will contents into a flattened, owned buffer.
     pub(crate) fn serialize<'b>(
         &self,
         buf: &'b mut [u8],


Test patch:
diff --git a/tests/deferred_payload.rs b/tests/deferred_payload.rs
new file mode 100644
index 0000000..8d201d5
--- /dev/null
+++ b/tests/deferred_payload.rs
@@ -0,0 +1,40 @@
+use minimq::{DeferredPublication, Minimq, QoS};
+
+use embedded_nal::{self, IpAddr, Ipv4Addr};
+use std_embedded_time::StandardClock;
+
+#[test]
+fn main() -> std::io::Result<()> {
+    env_logger::init();
+
+    let mut rx_buffer = [0u8; 256];
+    let mut tx_buffer = [0u8; 256];
+    let mut session = [0u8; 256];
+    let stack = std_embedded_nal::Stack;
+    let localhost = IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1));
+    let mut mqtt: Minimq<'_, _, _, minimq::broker::IpBroker> = Minimq::new(
+        stack,
+        StandardClock::default(),
+        minimq::Config::new(localhost.into(), &mut rx_buffer, &mut tx_buffer)
+            .session_state(&mut session)
+            .keepalive_interval(60),
+    );
+
+    while !mqtt.client().is_connected() {
+        mqtt.poll(|_client, _topic, _payload, _properties| {})
+            .unwrap();
+    }
+
+    assert!(matches!(
+        mqtt.client().publish(
+            DeferredPublication::new(|_buf| { Err("Oops!") })
+                .topic("data")
+                .qos(QoS::ExactlyOnce)
+                .finish()
+                .unwrap(),
+        ),
+        Err(minimq::PubError::Serialization("Oops!"))
+    ));
+
+    Ok(())
+}