ruststream-rumqttc 0.7.1

//! [`MqttSubscriber`]: a stream of deliveries fed by the connection task.

use std::num::NonZeroUsize;
use std::sync::Arc;
use std::time::Duration;

use futures::Stream;
use ruststream::{BatchSubscriber, BufferedSubscriber, Subscriber};
use tokio::sync::mpsc;

use crate::broker::Link;
use crate::conn::Shared;
use crate::error::MqttError;
use crate::message::MqttMessage;

/// How long a partial batch waits for more deliveries after its first one.
///
/// The crate's own choice, not the mount site's: MQTT delivers one PUBLISH packet at a time over
/// a network event loop, so a batch that is not yet full waits a network-shaped moment for the
/// next packet rather than an in-process one. The batch size is the registration's, and arrives
/// per subscription as the argument of [`BatchSubscriber::batches`].
///
/// The in-process mode batches with the same deadline, so a batch handler behaves the same under
/// the harness as it does on a server.
pub(crate) const BATCH_MAX_WAIT: Duration = Duration::from_millis(20);

/// What dropping a subscription does to its filter on the server.
#[derive(Clone, Copy, Debug)]
pub(crate) enum OnDrop {
    /// Opened by hand, so the drop is the caller's decision to stop reading while the connection
    /// lives on: the filter leaves the server once no other subscription on the connection shares
    /// it.
    Unsubscribe,
    /// Opened for the service's app, whose runtime drops its subscriptions only while the service
    /// stops, right before it shuts the connection down. The filter stays on the server: a
    /// persistent session keeps receiving for it while the service is away, and a clean one ends
    /// with the `DISCONNECT` anyway.
    Keep,
}

/// The wire half of a subscription: whatever the connection demultiplexed onto this filter's
/// channel, one delivery at a time, which is all the protocol offers.
///
/// The in-process mode demultiplexes onto the same channel through the same registry, so this is
/// the one subscriber of both transports; only the link its `Drop` unsubscribes through differs.
struct WireSubscriber {
    id: u64,
    shared: Arc<Shared>,
    link: Link,
    on_drop: OnDrop,
    rx: mpsc::UnboundedReceiver<Result<MqttMessage, MqttError>>,
}

impl std::fmt::Debug for WireSubscriber {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("WireSubscriber").finish_non_exhaustive()
    }
}

impl Drop for WireSubscriber {
    fn drop(&mut self) {
        match self.on_drop {
            OnDrop::Unsubscribe => self.shared.release(self.id, &self.link),
            OnDrop::Keep => self.shared.leave(self.id),
        }
    }
}

impl Subscriber for WireSubscriber {
    type Message = MqttMessage;
    type Error = MqttError;

    fn stream(&mut self) -> impl Stream<Item = Result<MqttMessage, MqttError>> + Send + '_ {
        // Poll the channel in place rather than wrapping it in an owning stream, so `stream`
        // can be called again after the returned stream is dropped (the runtime and the
        // conformance helpers re-enter it per call).
        futures::stream::poll_fn(move |cx| self.rx.poll_recv(cx))
    }
}

/// A subscription to one MQTT topic filter; yields [`MqttMessage`]s one at a time, or in batches.
///
/// Delivery back-pressure is the protocol's receive-maximum: the broker bounds unacked
/// `QoS` 1/2 deliveries, so unsettled messages cap what sits in this subscriber's queue
/// (`QoS` 0 has no such bound by design).
///
/// A subscription the app opened stays in the session when the service stops: the broker keeps
/// its filter, and a persistent session queues what arrives for it until the service starts
/// again. Dropping one opened by hand ([`subscribe_topic`], [`subscribe_filter`]) unsubscribes
/// the filter once no other subscriber on this connection shares it.
///
/// [`subscribe_topic`]: crate::ConnectedMqttBroker::subscribe_topic
/// [`subscribe_filter`]: crate::ConnectedMqttBroker::subscribe_filter
///
/// MQTT has no batch fetch - a PUBLISH packet carries one message - so the batches a `&[T]`
/// handler consumes are assembled on the client by the framework's
/// [`BufferedSubscriber`]: a batch closes when it holds the size the mount site named, or once
/// the crate's own 20 ms deadline has elapsed after its first delivery. Nothing at the mount
/// site says which of the two happened, which is the point.
pub struct MqttSubscriber {
    filter: String,
    buffered: BufferedSubscriber<WireSubscriber>,
}

impl std::fmt::Debug for MqttSubscriber {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("MqttSubscriber")
            .field("filter", &self.filter)
            .finish_non_exhaustive()
    }
}

impl MqttSubscriber {
    pub(crate) fn new(
        filter: String,
        id: u64,
        shared: Arc<Shared>,
        link: Link,
        on_drop: OnDrop,
        rx: mpsc::UnboundedReceiver<Result<MqttMessage, MqttError>>,
    ) -> Self {
        Self {
            filter,
            buffered: BufferedSubscriber::new(WireSubscriber {
                id,
                shared,
                link,
                on_drop,
                rx,
            })
            .max_wait(BATCH_MAX_WAIT),
        }
    }

    /// The plain topic filter this subscription matches.
    #[must_use]
    pub fn filter(&self) -> &str {
        &self.filter
    }
}

impl Subscriber for MqttSubscriber {
    type Message = MqttMessage;
    type Error = MqttError;

    fn stream(&mut self) -> impl Stream<Item = Result<MqttMessage, MqttError>> + Send + '_ {
        self.buffered.stream()
    }
}

impl BatchSubscriber for MqttSubscriber {
    type Batch = Vec<MqttMessage>;

    /// # Cancel safety
    ///
    /// Cancel-safe between polls, like [`Subscriber::stream`]: the batch being filled lives
    /// inside the returned stream and survives a cancelled poll. Dropping the whole stream
    /// abandons that batch's deliveries unacknowledged, and MQTT redelivers them when a
    /// persistent session resumes.
    fn batches(
        &mut self,
        size: NonZeroUsize,
    ) -> impl Stream<Item = Result<Self::Batch, MqttError>> + Send + '_ {
        self.buffered.batches(size)
    }
}