From 7e3de50977ec7afb666c23300ef7dda41e09feaa Mon Sep 17 00:00:00 2001 From: aa51513 Date: Wed, 29 Apr 2026 23:56:09 +0800 Subject: [PATCH] feat(rumqttc): return packet ID from publish methods All publish methods (sync/async, v4/v5) now return Result instead of Result<(), ClientError>. For QoS 0 publishes, 0 is returned. For QoS 1/2, the assigned packet ID is returned immediately. Implementation uses a shared Arc counter between the client and MqttState. The client pre-assigns packet IDs before sending through the channel, avoiding the need for async oneshot channels (which would deadlock the sync Client). - Add pkid_counter field to MqttState, EventLoop, and AsyncClient - MqttState: remove Clone derive, implement Debug manually - EventLoop: create and expose shared pkid_counter - AsyncClient: assign pkid via atomic fetch_add before sending - Update from_senders/from_sender signatures to accept pkid_counter - Update tests for new API and pkid wrapping behavior --- rumqttc/src/client.rs | 64 ++++++++++++++++++++++-------- rumqttc/src/eventloop.rs | 13 +++++- rumqttc/src/state.rs | 54 +++++++++++++++++-------- rumqttc/src/v5/client.rs | 79 ++++++++++++++++++++++++++----------- rumqttc/src/v5/eventloop.rs | 13 +++++- rumqttc/src/v5/state.rs | 64 +++++++++++++++++++----------- 6 files changed, 205 insertions(+), 82 deletions(-) diff --git a/rumqttc/src/client.rs b/rumqttc/src/client.rs index cb58cf82d..e318be9d4 100644 --- a/rumqttc/src/client.rs +++ b/rumqttc/src/client.rs @@ -1,5 +1,7 @@ //! This module offers a high level synchronous and asynchronous abstraction to //! async eventloop. +use std::sync::atomic::{AtomicU16, Ordering}; +use std::sync::Arc; use std::time::Duration; use crate::mqttbytes::{v4::*, QoS}; @@ -42,6 +44,8 @@ impl From> for ClientError { #[derive(Clone, Debug)] pub struct AsyncClient { request_tx: Sender, + pkid_counter: Arc, + max_inflight: u16, } impl AsyncClient { @@ -49,10 +53,16 @@ impl AsyncClient { /// /// `cap` specifies the capacity of the bounded async channel. pub fn new(options: MqttOptions, cap: usize) -> (AsyncClient, EventLoop) { + let max_inflight = options.inflight; let eventloop = EventLoop::new(options, cap); let request_tx = eventloop.requests_tx.clone(); + let pkid_counter = eventloop.pkid_counter(); - let client = AsyncClient { request_tx }; + let client = AsyncClient { + request_tx, + pkid_counter, + max_inflight, + }; (client, eventloop) } @@ -61,8 +71,24 @@ impl AsyncClient { /// /// This is mostly useful for creating a test instance where you can /// listen on the corresponding receiver. - pub fn from_senders(request_tx: Sender) -> AsyncClient { - AsyncClient { request_tx } + pub fn from_senders(request_tx: Sender, pkid_counter: Arc, max_inflight: u16) -> AsyncClient { + AsyncClient { + request_tx, + pkid_counter, + max_inflight, + } + } + + /// Assigns a packet id for QoS 1/2 publishes using the shared atomic counter. + fn assign_pkid(&self, publish: &mut Publish) -> u16 { + if publish.qos != QoS::AtMostOnce { + let raw = self.pkid_counter.fetch_add(1, Ordering::Relaxed); + let pkid = (raw % self.max_inflight) + 1; + publish.pkid = pkid; + pkid + } else { + 0 + } } /// Sends a MQTT Publish to the `EventLoop`. @@ -72,7 +98,7 @@ impl AsyncClient { qos: QoS, retain: bool, payload: V, - ) -> Result<(), ClientError> + ) -> Result where S: Into, V: Into>, @@ -80,12 +106,13 @@ impl AsyncClient { let topic = topic.into(); let mut publish = Publish::new(&topic, qos, payload); publish.retain = retain; + let pkid = self.assign_pkid(&mut publish); let publish = Request::Publish(publish); if !valid_topic(&topic) { return Err(ClientError::Request(publish)); } self.request_tx.send_async(publish).await?; - Ok(()) + Ok(pkid) } /// Attempts to send a MQTT Publish to the `EventLoop`. @@ -95,7 +122,7 @@ impl AsyncClient { qos: QoS, retain: bool, payload: V, - ) -> Result<(), ClientError> + ) -> Result where S: Into, V: Into>, @@ -103,12 +130,13 @@ impl AsyncClient { let topic = topic.into(); let mut publish = Publish::new(&topic, qos, payload); publish.retain = retain; + let pkid = self.assign_pkid(&mut publish); let publish = Request::Publish(publish); if !valid_topic(&topic) { return Err(ClientError::TryRequest(publish)); } self.request_tx.try_send(publish)?; - Ok(()) + Ok(pkid) } /// Sends a MQTT PubAck to the `EventLoop`. Only needed in if `manual_acks` flag is set. @@ -137,15 +165,16 @@ impl AsyncClient { qos: QoS, retain: bool, payload: Bytes, - ) -> Result<(), ClientError> + ) -> Result where S: Into, { let mut publish = Publish::from_bytes(topic, qos, payload); publish.retain = retain; + let pkid = self.assign_pkid(&mut publish); let publish = Request::Publish(publish); self.request_tx.send_async(publish).await?; - Ok(()) + Ok(pkid) } /// Sends a MQTT Subscribe to the `EventLoop` @@ -272,9 +301,9 @@ impl Client { /// /// This is mostly useful for creating a test instance where you can /// listen on the corresponding receiver. - pub fn from_sender(request_tx: Sender) -> Client { + pub fn from_sender(request_tx: Sender, pkid_counter: Arc, max_inflight: u16) -> Client { Client { - client: AsyncClient::from_senders(request_tx), + client: AsyncClient::from_senders(request_tx, pkid_counter, max_inflight), } } @@ -285,7 +314,7 @@ impl Client { qos: QoS, retain: bool, payload: V, - ) -> Result<(), ClientError> + ) -> Result where S: Into, V: Into>, @@ -293,12 +322,13 @@ impl Client { let topic = topic.into(); let mut publish = Publish::new(&topic, qos, payload); publish.retain = retain; + let pkid = self.client.assign_pkid(&mut publish); let publish = Request::Publish(publish); if !valid_topic(&topic) { return Err(ClientError::Request(publish)); } self.client.request_tx.send(publish)?; - Ok(()) + Ok(pkid) } pub fn try_publish( @@ -307,13 +337,12 @@ impl Client { qos: QoS, retain: bool, payload: V, - ) -> Result<(), ClientError> + ) -> Result where S: Into, V: Into>, { - self.client.try_publish(topic, qos, retain, payload)?; - Ok(()) + self.client.try_publish(topic, qos, retain, payload) } /// Sends a MQTT PubAck to the `EventLoop`. Only needed in if `manual_acks` flag is set. @@ -540,7 +569,8 @@ mod test { #[test] fn should_be_able_to_build_test_client_from_channel() { let (tx, rx) = flume::bounded(1); - let client = Client::from_sender(tx); + let pkid_counter = Arc::new(AtomicU16::new(0)); + let client = Client::from_sender(tx, pkid_counter, 100); client .publish("hello/world", QoS::ExactlyOnce, false, "good bye") .expect("Should be able to publish"); diff --git a/rumqttc/src/eventloop.rs b/rumqttc/src/eventloop.rs index c2d46f905..789c52220 100644 --- a/rumqttc/src/eventloop.rs +++ b/rumqttc/src/eventloop.rs @@ -13,6 +13,8 @@ use std::collections::VecDeque; use std::io; use std::net::SocketAddr; use std::pin::Pin; +use std::sync::atomic::AtomicU16; +use std::sync::Arc; use std::time::Duration; #[cfg(unix)] @@ -85,6 +87,8 @@ pub struct EventLoop { /// Keep alive time keepalive_timeout: Option>>, pub network_options: NetworkOptions, + /// Shared atomic counter for packet id assignment + pub(crate) pkid_counter: Arc, } /// Events which can be yielded by the event loop @@ -104,19 +108,26 @@ impl EventLoop { let pending = VecDeque::new(); let max_inflight = mqtt_options.inflight; let manual_acks = mqtt_options.manual_acks; + let pkid_counter = Arc::new(AtomicU16::new(0)); EventLoop { mqtt_options, - state: MqttState::new(max_inflight, manual_acks), + state: MqttState::new(max_inflight, manual_acks, pkid_counter.clone()), requests_tx, requests_rx, pending, network: None, keepalive_timeout: None, network_options: NetworkOptions::new(), + pkid_counter, } } + /// Returns a clone of the shared packet id counter for use by the client + pub fn pkid_counter(&self) -> Arc { + self.pkid_counter.clone() + } + /// Last session might contain packets which aren't acked. MQTT says these packets should be /// republished in the next session. Move pending messages from state to eventloop, drops the /// underlying network connection and clears the keepalive timeout if any. diff --git a/rumqttc/src/state.rs b/rumqttc/src/state.rs index f7cb34841..57e44d13b 100644 --- a/rumqttc/src/state.rs +++ b/rumqttc/src/state.rs @@ -4,6 +4,9 @@ use crate::mqttbytes::v4::*; use crate::mqttbytes::{self, *}; use fixedbitset::FixedBitSet; use std::collections::VecDeque; +use std::fmt::{self, Debug, Formatter}; +use std::sync::atomic::{AtomicU16, Ordering}; +use std::sync::Arc; use std::{io, time::Instant}; /// Errors during state handling @@ -40,7 +43,6 @@ pub enum StateError { // This is done for 2 reasons // Bad acks or out of order acks aren't O(n) causing cpu spikes // Any missing acks from the broker are detected during the next recycled use of packet ids -#[derive(Debug, Clone)] pub struct MqttState { /// Status of last ping pub await_pingresp: bool, @@ -72,19 +74,40 @@ pub struct MqttState { pub events: VecDeque, /// Indicates if acknowledgements should be send immediately pub manual_acks: bool, + /// Shared atomic counter for packet id assignment + pub(crate) pkid_counter: Arc, +} + +impl Debug for MqttState { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_struct("MqttState") + .field("await_pingresp", &self.await_pingresp) + .field("collision_ping_count", &self.collision_ping_count) + .field("last_pkid", &self.last_pkid) + .field("last_puback", &self.last_puback) + .field("inflight", &self.inflight) + .field("max_inflight", &self.max_inflight) + .field("outgoing_pub", &self.outgoing_pub) + .field("outgoing_rel", &self.outgoing_rel) + .field("collision", &self.collision) + .field("events", &self.events) + .field("manual_acks", &self.manual_acks) + .finish() + } } impl MqttState { /// Creates new mqtt state. Same state should be used during a /// connection for persistent sessions while new state should /// instantiated for clean sessions - pub fn new(max_inflight: u16, manual_acks: bool) -> Self { + pub fn new(max_inflight: u16, manual_acks: bool, pkid_counter: Arc) -> Self { + let last_pkid = pkid_counter.load(Ordering::Relaxed); MqttState { await_pingresp: false, collision_ping_count: 0, last_incoming: Instant::now(), last_outgoing: Instant::now(), - last_pkid: 0, + last_pkid, last_puback: 0, inflight: 0, max_inflight, @@ -96,6 +119,7 @@ impl MqttState { // TODO: Optimize these sizes later events: VecDeque::with_capacity(100), manual_acks, + pkid_counter, } } @@ -312,6 +336,9 @@ impl MqttState { if publish.qos != QoS::AtMostOnce { if publish.pkid == 0 { publish.pkid = self.next_pkid(); + } else { + // Client pre-assigned pkid via shared atomic counter + self.last_pkid = publish.pkid; } let pkid = publish.pkid; @@ -484,19 +511,10 @@ impl MqttState { /// Packet ids are incremented till maximum set inflight messages and reset to 1 after that. /// fn next_pkid(&mut self) -> u16 { - let next_pkid = self.last_pkid + 1; - - // When next packet id is at the edge of inflight queue, - // set await flag. This instructs eventloop to stop - // processing requests until all the inflight publishes - // are acked - if next_pkid == self.max_inflight { - self.last_pkid = 0; - return next_pkid; - } - - self.last_pkid = next_pkid; - next_pkid + let raw = self.pkid_counter.fetch_add(1, Ordering::Relaxed); + let pkid = (raw % self.max_inflight) + 1; + self.last_pkid = pkid; + pkid } } @@ -506,6 +524,8 @@ mod test { use crate::mqttbytes::v4::*; use crate::mqttbytes::*; use crate::{Event, Incoming, Outgoing, Request}; + use std::sync::atomic::AtomicU16; + use std::sync::Arc; fn build_outgoing_publish(qos: QoS) -> Publish { let topic = "hello/world".to_owned(); @@ -527,7 +547,7 @@ mod test { } fn build_mqttstate() -> MqttState { - MqttState::new(100, false) + MqttState::new(100, false, Arc::new(AtomicU16::new(0))) } #[test] diff --git a/rumqttc/src/v5/client.rs b/rumqttc/src/v5/client.rs index 7a86333f2..4772c3399 100644 --- a/rumqttc/src/v5/client.rs +++ b/rumqttc/src/v5/client.rs @@ -1,5 +1,7 @@ //! This module offers a high level synchronous and asynchronous abstraction to //! async eventloop. +use std::sync::atomic::{AtomicU16, Ordering}; +use std::sync::Arc; use std::time::Duration; use super::mqttbytes::v5::{ @@ -47,6 +49,8 @@ impl From> for ClientError { #[derive(Clone, Debug)] pub struct AsyncClient { request_tx: Sender, + pkid_counter: Arc, + max_inflight: u16, } impl AsyncClient { @@ -54,10 +58,16 @@ impl AsyncClient { /// /// `cap` specifies the capacity of the bounded async channel. pub fn new(options: MqttOptions, cap: usize) -> (AsyncClient, EventLoop) { + let max_inflight = options.outgoing_inflight_upper_limit.unwrap_or(u16::MAX); let eventloop = EventLoop::new(options, cap); let request_tx = eventloop.requests_tx.clone(); + let pkid_counter = eventloop.pkid_counter(); - let client = AsyncClient { request_tx }; + let client = AsyncClient { + request_tx, + pkid_counter, + max_inflight, + }; (client, eventloop) } @@ -66,8 +76,24 @@ impl AsyncClient { /// /// This is mostly useful for creating a test instance where you can /// listen on the corresponding receiver. - pub fn from_senders(request_tx: Sender) -> AsyncClient { - AsyncClient { request_tx } + pub fn from_senders(request_tx: Sender, pkid_counter: Arc, max_inflight: u16) -> AsyncClient { + AsyncClient { + request_tx, + pkid_counter, + max_inflight, + } + } + + /// Assigns a packet id for QoS 1/2 publishes using the shared atomic counter. + fn assign_pkid(&self, publish: &mut Publish) -> u16 { + if publish.qos != QoS::AtMostOnce { + let raw = self.pkid_counter.fetch_add(1, Ordering::Relaxed); + let pkid = (raw % self.max_inflight) + 1; + publish.pkid = pkid; + pkid + } else { + 0 + } } /// Sends a MQTT Publish to the `EventLoop`. @@ -78,7 +104,7 @@ impl AsyncClient { retain: bool, payload: P, properties: Option, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -86,12 +112,13 @@ impl AsyncClient { let topic = topic.into(); let mut publish = Publish::new(&topic, qos, payload, properties); publish.retain = retain; + let pkid = self.assign_pkid(&mut publish); let publish = Request::Publish(publish); if !valid_topic(&topic) { return Err(ClientError::Request(publish)); } self.request_tx.send_async(publish).await?; - Ok(()) + Ok(pkid) } pub async fn publish_with_properties( @@ -101,7 +128,7 @@ impl AsyncClient { retain: bool, payload: P, properties: PublishProperties, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -116,7 +143,7 @@ impl AsyncClient { qos: QoS, retain: bool, payload: P, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -132,7 +159,7 @@ impl AsyncClient { retain: bool, payload: P, properties: Option, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -140,12 +167,13 @@ impl AsyncClient { let topic = topic.into(); let mut publish = Publish::new(&topic, qos, payload, properties); publish.retain = retain; + let pkid = self.assign_pkid(&mut publish); let publish = Request::Publish(publish); if !valid_topic(&topic) { return Err(ClientError::TryRequest(publish)); } self.request_tx.try_send(publish)?; - Ok(()) + Ok(pkid) } pub fn try_publish_with_properties( @@ -155,7 +183,7 @@ impl AsyncClient { retain: bool, payload: P, properties: PublishProperties, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -169,7 +197,7 @@ impl AsyncClient { qos: QoS, retain: bool, payload: P, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -204,19 +232,20 @@ impl AsyncClient { retain: bool, payload: Bytes, properties: Option, - ) -> Result<(), ClientError> + ) -> Result where S: Into, { let topic = topic.into(); let mut publish = Publish::new(&topic, qos, payload, properties); publish.retain = retain; + let pkid = self.assign_pkid(&mut publish); let publish = Request::Publish(publish); if !valid_topic(&topic) { return Err(ClientError::TryRequest(publish)); } self.request_tx.send_async(publish).await?; - Ok(()) + Ok(pkid) } pub async fn publish_bytes_with_properties( @@ -226,7 +255,7 @@ impl AsyncClient { retain: bool, payload: Bytes, properties: PublishProperties, - ) -> Result<(), ClientError> + ) -> Result where S: Into, { @@ -240,7 +269,7 @@ impl AsyncClient { qos: QoS, retain: bool, payload: Bytes, - ) -> Result<(), ClientError> + ) -> Result where S: Into, { @@ -489,9 +518,9 @@ impl Client { /// /// This is mostly useful for creating a test instance where you can /// listen on the corresponding receiver. - pub fn from_sender(request_tx: Sender) -> Client { + pub fn from_sender(request_tx: Sender, pkid_counter: Arc, max_inflight: u16) -> Client { Client { - client: AsyncClient::from_senders(request_tx), + client: AsyncClient::from_senders(request_tx, pkid_counter, max_inflight), } } @@ -503,7 +532,7 @@ impl Client { retain: bool, payload: P, properties: Option, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -511,12 +540,13 @@ impl Client { let topic = topic.into(); let mut publish = Publish::new(&topic, qos, payload, properties); publish.retain = retain; + let pkid = self.client.assign_pkid(&mut publish); let publish = Request::Publish(publish); if !valid_topic(&topic) { return Err(ClientError::Request(publish)); } self.client.request_tx.send(publish)?; - Ok(()) + Ok(pkid) } pub fn publish_with_properties( @@ -526,7 +556,7 @@ impl Client { retain: bool, payload: P, properties: PublishProperties, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -540,7 +570,7 @@ impl Client { qos: QoS, retain: bool, payload: P, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -555,7 +585,7 @@ impl Client { retain: bool, payload: P, properties: PublishProperties, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -570,7 +600,7 @@ impl Client { qos: QoS, retain: bool, payload: P, - ) -> Result<(), ClientError> + ) -> Result where S: Into, P: Into, @@ -890,7 +920,8 @@ mod test { #[test] fn should_be_able_to_build_test_client_from_channel() { let (tx, rx) = flume::bounded(1); - let client = Client::from_sender(tx); + let pkid_counter = Arc::new(AtomicU16::new(0)); + let client = Client::from_sender(tx, pkid_counter, u16::MAX); client .publish("hello/world", QoS::ExactlyOnce, false, "good bye") .expect("Should be able to publish"); diff --git a/rumqttc/src/v5/eventloop.rs b/rumqttc/src/v5/eventloop.rs index 8fe6f22a5..9564c04a3 100644 --- a/rumqttc/src/v5/eventloop.rs +++ b/rumqttc/src/v5/eventloop.rs @@ -11,6 +11,8 @@ use tokio::time::{self, error::Elapsed, Instant, Sleep}; use std::collections::VecDeque; use std::io; use std::pin::Pin; +use std::sync::atomic::AtomicU16; +use std::sync::Arc; use std::time::Duration; use super::mqttbytes::v5::ConnectReturnCode; @@ -82,6 +84,8 @@ pub struct EventLoop { network: Option, /// Keep alive time keepalive_timeout: Option>>, + /// Shared atomic counter for packet id assignment + pub(crate) pkid_counter: Arc, } /// Events which can be yielded by the event loop @@ -101,18 +105,25 @@ impl EventLoop { let pending = VecDeque::new(); let inflight_limit = options.outgoing_inflight_upper_limit.unwrap_or(u16::MAX); let manual_acks = options.manual_acks; + let pkid_counter = Arc::new(AtomicU16::new(0)); EventLoop { options, - state: MqttState::new(inflight_limit, manual_acks), + state: MqttState::new(inflight_limit, manual_acks, pkid_counter.clone()), requests_tx, requests_rx, pending, network: None, keepalive_timeout: None, + pkid_counter, } } + /// Returns a clone of the shared packet id counter for use by the client + pub fn pkid_counter(&self) -> Arc { + self.pkid_counter.clone() + } + /// Last session might contain packets which aren't acked. MQTT says these packets should be /// republished in the next session. Move pending messages from state to eventloop, drops the /// underlying network connection and clears the keepalive timeout if any. diff --git a/rumqttc/src/v5/state.rs b/rumqttc/src/v5/state.rs index 9a7485f3b..42d19f996 100644 --- a/rumqttc/src/v5/state.rs +++ b/rumqttc/src/v5/state.rs @@ -10,6 +10,9 @@ use super::{Event, Incoming, Outgoing, Request}; use bytes::Bytes; use fixedbitset::FixedBitSet; use std::collections::{HashMap, VecDeque}; +use std::fmt::{self, Debug, Formatter}; +use std::sync::atomic::{AtomicU16, Ordering}; +use std::sync::Arc; use std::{io, time::Instant}; /// Errors during state handling @@ -74,7 +77,6 @@ impl From for StateError { // This is done for 2 reasons // Bad acks or out of order acks aren't O(n) causing cpu spikes // Any missing acks from the broker are detected during the next recycled use of packet ids -#[derive(Debug, Clone)] pub struct MqttState { /// Status of last ping pub await_pingresp: bool, @@ -110,19 +112,40 @@ pub struct MqttState { pub(crate) max_outgoing_inflight: u16, /// Upper limit on the maximum number of allowed inflight QoS1 & QoS2 requests max_outgoing_inflight_upper_limit: u16, + /// Shared atomic counter for packet id assignment + pub(crate) pkid_counter: Arc, +} + +impl Debug for MqttState { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + f.debug_struct("MqttState") + .field("await_pingresp", &self.await_pingresp) + .field("collision_ping_count", &self.collision_ping_count) + .field("last_pkid", &self.last_pkid) + .field("inflight", &self.inflight) + .field("max_outgoing_inflight", &self.max_outgoing_inflight) + .field("outgoing_pub", &self.outgoing_pub) + .field("outgoing_rel", &self.outgoing_rel) + .field("collision", &self.collision) + .field("events", &self.events) + .field("manual_acks", &self.manual_acks) + .field("broker_topic_alias_max", &self.broker_topic_alias_max) + .finish() + } } impl MqttState { /// Creates new mqtt state. Same state should be used during a /// connection for persistent sessions while new state should /// instantiated for clean sessions - pub fn new(max_inflight: u16, manual_acks: bool) -> Self { + pub fn new(max_inflight: u16, manual_acks: bool, pkid_counter: Arc) -> Self { + let last_pkid = pkid_counter.load(Ordering::Relaxed); MqttState { await_pingresp: false, collision_ping_count: 0, last_incoming: Instant::now(), last_outgoing: Instant::now(), - last_pkid: 0, + last_pkid, inflight: 0, // index 0 is wasted as 0 is not a valid packet id outgoing_pub: vec![None; max_inflight as usize + 1], @@ -137,6 +160,7 @@ impl MqttState { broker_topic_alias_max: 0, max_outgoing_inflight: max_inflight, max_outgoing_inflight_upper_limit: max_inflight, + pkid_counter, } } @@ -475,6 +499,9 @@ impl MqttState { if publish.qos != QoS::AtMostOnce { if publish.pkid == 0 { publish.pkid = self.next_pkid(); + } else { + // Client pre-assigned pkid via shared atomic counter + self.last_pkid = publish.pkid; } let pkid = publish.pkid; @@ -666,19 +693,10 @@ impl MqttState { /// Packet ids are incremented till maximum set inflight messages and reset to 1 after that. /// fn next_pkid(&mut self) -> u16 { - let next_pkid = self.last_pkid + 1; - - // When next packet id is at the edge of inflight queue, - // set await flag. This instructs eventloop to stop - // processing requests until all the inflight publishes - // are acked - if next_pkid == self.max_outgoing_inflight { - self.last_pkid = 0; - return next_pkid; - } - - self.last_pkid = next_pkid; - next_pkid + let raw = self.pkid_counter.fetch_add(1, Ordering::Relaxed); + let pkid = (raw % self.max_outgoing_inflight) + 1; + self.last_pkid = pkid; + pkid } } @@ -688,6 +706,8 @@ mod test { use super::mqttbytes::*; use super::{Event, Incoming, Outgoing, Request}; use super::{MqttState, StateError}; + use std::sync::atomic::AtomicU16; + use std::sync::Arc; fn build_outgoing_publish(qos: QoS) -> Publish { let topic = "hello/world".to_owned(); @@ -709,7 +729,7 @@ mod test { } fn build_mqttstate() -> MqttState { - MqttState::new(u16::MAX, false) + MqttState::new(u16::MAX, false, Arc::new(AtomicU16::new(0))) } #[test] @@ -770,7 +790,7 @@ mod test { #[test] fn outgoing_publish_with_max_inflight_is_ok() { - let mut mqtt = MqttState::new(2, false); + let mut mqtt = MqttState::new(2, false, Arc::new(AtomicU16::new(0))); // QoS2 publish let publish = build_outgoing_publish(QoS::ExactlyOnce); @@ -779,12 +799,12 @@ mod test { assert_eq!(mqtt.last_pkid, 1); assert_eq!(mqtt.inflight, 1); - // Packet id should be set back down to 0, since we hit the limit + // Second publish gets pkid 2 (max_inflight = 2) mqtt.outgoing_publish(publish.clone()).unwrap(); - assert_eq!(mqtt.last_pkid, 0); + assert_eq!(mqtt.last_pkid, 2); assert_eq!(mqtt.inflight, 2); - // This should cause a collition + // This should cause a collision (pkid wraps back to 1) mqtt.outgoing_publish(publish.clone()).unwrap(); assert_eq!(mqtt.last_pkid, 1); assert_eq!(mqtt.inflight, 2); @@ -796,7 +816,7 @@ mod test { // Now there should be space in the outgoing queue mqtt.outgoing_publish(publish.clone()).unwrap(); - assert_eq!(mqtt.last_pkid, 0); + assert_eq!(mqtt.last_pkid, 2); assert_eq!(mqtt.inflight, 2); }