rmqtt 0.23.0

//! A bounded asynchronous channel with rate-limiting and customizable backpressure policies.
//!
//! This module provides a mechanism for sending items into a bounded queue with an associated
//! receiver that consumes items at a controlled rate. It supports custom policies for handling
//! queue overflow and hook functions to track queue events.
//!
//! # Components
//!
//! - [`Sender<T>`]: Allows sending items into a bounded queue asynchronously. Supports custom
//!   overflow policies via [`PolicyFn`] and can notify via an MPSC signal channel.
//!
//! - [`Receiver<'a, T>`]: A `Stream` that yields elements from the queue at a rate-limited pace,
//!   using the [`governor`] crate.
//!
//! - [`Policy`]: Defines behavior when the queue is full:
//!   - `Current`: Discard the current value being sent.
//!   - `Early`: Discard the earliest value in the queue to make space.
//!
//! - [`Queue<T>`]: A thread-safe bounded queue using `SegQueue`, with optional hooks on push/pop.
//!
//! - [`Limiter`]: Creates rate-limited channels using `RateLimiter`, managing burst and refill
//!   periods.
//!
//! # Features
//!
//! - Thread-safe and cloneable senders with policy injection.
//! - Rate-limiting receivers with configurable burst size and interval.
//! - Event hooks for push/pop operations in the queue.
//! - Graceful handling of overflows based on selected policy.
//!
//!
//! # Crates Used
//!
//! - `governor`: For rate limiting
//! - `crossbeam`: For the lock-free `SegQueue`
//! - `futures`: For async MPSC and stream support

use std::collections::vec_deque::VecDeque;
use std::num::NonZeroU32;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use std::time::Duration;

use anyhow::Result;
// use crossbeam::queue::SegQueue;
use futures::{channel::mpsc, SinkExt, Stream};
use governor::{
    clock::DefaultClock,
    middleware::NoOpMiddleware,
    prelude::StreamRateLimitExt,
    state::{InMemoryState, NotKeyed},
    Quota, RateLimiter, RatelimitedStream,
};
use parking_lot::Mutex;

type DirectLimiter = RateLimiter<NotKeyed, InMemoryState, DefaultClock>;

pub type Receiver<'a, T> =
    RatelimitedStream<'a, ReceiverStream<T>, InMemoryState, DefaultClock, NoOpMiddleware>;

/// Defines the overflow policy when the bounded queue is full.
///
/// Selects which item to discard to make room for new entries.
pub enum Policy {
    /// Discard the current (incoming) value.
    Current,
    /// Discard the earliest (front) value in the queue.
    Early,
}

/// Trait for queue overflow policy functions.
///
/// A `PolicyFn` receives a reference to the item being sent
/// and returns the [`Policy`] to apply when the queue is full.
pub trait PolicyFn<P>: 'static + Sync + Send + Fn(&P) -> Policy {}

impl<T, P> PolicyFn<P> for T where T: 'static + Sync + Send + Clone + Fn(&P) -> Policy {}

/// Trait for queue event notification callbacks.
///
/// Invoked on push and pop operations to track queue activity.
/// Used for statistics and monitoring integration.
pub trait OnEventFn: 'static + Sync + Send + Fn() {}
impl<T> OnEventFn for T where T: 'static + Sync + Send + Clone + Fn() {}

/// A bounded, asynchronous sender with configurable overflow policy.
///
/// Wraps a [`Queue<T>`] and an MPSC signal channel. When an item
/// is successfully enqueued, a signal is sent on the MPSC channel
/// to wake the receiver task.
///
/// # Overflow Behavior
///
/// When the queue is full, the behavior depends on the configured
/// [`PolicyFn`]:
/// - `Policy::Current` — The new item is rejected and returned as `Err`.
/// - `Policy::Early` — The oldest item is evicted to make room.
pub struct Sender<T> {
    tx: mpsc::Sender<()>,
    queue: Arc<Queue<T>>,
    policy_fn: Arc<dyn PolicyFn<T>>,
}

impl<T> Sender<T> {
    /// Close the sender.
    ///
    /// Signals the receiver that no more items will be sent.
    #[inline]
    pub async fn close(&mut self) -> Result<()> {
        self.tx.close().await?;
        Ok(())
    }

    /// Number of items currently in the queue.
    #[inline]
    pub fn len(&self) -> usize {
        self.queue.len()
    }

    /// Whether the queue is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Set a custom overflow policy function for this sender.
    #[inline]
    pub fn policy<F>(mut self, f: F) -> Self
    where
        F: PolicyFn<T>,
    {
        self.policy_fn = Arc::new(f);
        self
    }

    /// Send a value into the bounded queue.
    ///
    /// If the queue is full, the configured policy determines
    /// which item is discarded. Returns `Ok(())` on success,
    /// or `Err(T)` with the rejected item.
    #[inline]
    pub async fn send(&self, v: T) -> Result<(), T> {
        if let Err(v) = self.queue.push(v) {
            match (self.policy_fn)(&v) {
                Policy::Current => return Err(v),
                Policy::Early => {
                    let removed = self.queue.pop();
                    if let Err(v) = self.queue.push(v) {
                        log::warn!("queue is full, queue len is {}", self.queue.len());
                        return Err(v);
                    }
                    if let Some(removed) = removed {
                        return Err(removed);
                    } else {
                        return Ok(());
                    }
                }
            }
        } else if let Err(e) = self.tx.clone().try_send(()) {
            log::warn!("channel is full, {e}");
        }
        Ok(())
    }

    /// Pop and return the oldest item from the queue, if any.
    #[inline]
    pub fn pop(&self) -> Option<T> {
        self.queue.pop()
    }
}

/// Stream adapter that polls the underlying MPSC receiver and
/// pops items from the bounded queue.
pub struct ReceiverStream<T> {
    rx: mpsc::Receiver<()>,
    queue: Arc<Queue<T>>,
}

impl<T> Stream for ReceiverStream<T> {
    type Item = Option<T>;
    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
        let result: Option<_> = futures::ready!(Pin::new(&mut self.rx).poll_next(cx));
        Poll::Ready(match result {
            Some(_) => Some(self.queue.pop()),
            None => None,
        })
    }
}

/// A rate limiter backed by the `governor` crate.
///
/// Creates bounded channels with token-bucket rate limiting
/// for controlled consumer throughput.
///
/// # Rate Calculation
///
/// The replenish period per token is computed as:
/// `replenish_n_per / burst`, ensuring an even distribution
/// of capacity across the burst window.
pub struct Limiter {
    l: DirectLimiter,
}

impl Limiter {
    #[inline]
    pub fn new(burst: NonZeroU32, replenish_n_per: Duration) -> Self {
        let period = replenish_n_per.as_nanos() as u64 / burst.get() as u64;
        let period = if period > 0 { Duration::from_nanos(period) } else { Duration::from_nanos(1) };
        log::debug!("burst: {burst:?}, {replenish_n_per:?}, {period:?}");
        let q = if let Some(q) = Quota::with_period(period) { q.allow_burst(burst) } else { unreachable!() };
        let l = RateLimiter::direct(q);
        Self { l }
    }

    #[inline]
    pub fn channel<T>(&self, queue: Arc<Queue<T>>) -> (Sender<T>, Receiver<'_, T>) {
        let (tx, rx) = mpsc::channel::<()>((queue.capacity() as f64 * 1.5) as usize);
        let s = ReceiverStream { rx, queue: queue.clone() }.ratelimit_stream(&self.l);
        (0..queue.len()).for_each(|_| {
            if let Err(e) = tx.clone().try_send(()) {
                //send offline message
                log::warn!("channel is full, {e}");
            }
        });
        (Sender { tx, queue, policy_fn: Arc::new(|_v: &T| -> Policy { Policy::Current }) }, s)
    }
}

/// A thread-safe bounded FIFO queue with optional push/pop event hooks.
///
/// Uses [`parking_lot::Mutex`] for interior mutability and a
/// [`VecDeque`] for O(1) amortized push/pop operations.
///
/// # Capacity Enforcement
///
/// When the queue reaches its capacity, [`push`](Queue::push) returns
/// `Err(v)` and the caller must decide how to handle the overflow
/// (typically via the enclosing [`Sender`]'s policy).
pub struct Queue<T> {
    cap: usize,
    inner: Mutex<VecDeque<T>>,
    on_push_fn: Option<Arc<dyn OnEventFn>>,
    on_pop_fn: Option<Arc<dyn OnEventFn>>,
}

impl<T> Drop for Queue<T> {
    #[inline]
    fn drop(&mut self) {
        log::debug!("Queue Drop ... len: {}", self.len());
    }
}

impl<T> Queue<T> {
    #[inline]
    pub fn new(cap: usize) -> Self {
        Self { cap, inner: Mutex::new(VecDeque::new()), on_push_fn: None, on_pop_fn: None }
    }

    #[inline]
    pub fn on_push<F>(&mut self, f: F)
    where
        F: OnEventFn,
    {
        self.on_push_fn = Some(Arc::new(f));
    }

    #[inline]
    pub fn on_pop<F>(&mut self, f: F)
    where
        F: OnEventFn,
    {
        self.on_pop_fn = Some(Arc::new(f));
    }

    #[inline]
    pub fn push(&self, v: T) -> Result<(), T> {
        let mut inner = self.inner.lock();
        if inner.len() >= self.cap {
            return Err(v);
        }
        if let Some(f) = self.on_push_fn.as_ref() {
            f();
        }
        inner.push_back(v);
        Ok(())
    }

    #[inline]
    pub fn pop(&self) -> Option<T> {
        let mut inner = self.inner.lock();
        if let Some(v) = inner.pop_front() {
            if let Some(f) = self.on_pop_fn.as_ref() {
                f();
            }
            Some(v)
        } else {
            None
        }
    }

    #[inline]
    pub fn capacity(&self) -> usize {
        self.cap
    }

    #[inline]
    pub fn len(&self) -> usize {
        self.inner.lock().len()
    }

    #[inline]
    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }
}

mod test {
    #[tokio::main]
    #[test]
    async fn channel() {
        use super::Policy;
        use super::Queue;
        use chrono::Local;
        use futures::StreamExt;
        use std::num::NonZeroU32;
        use std::sync::Arc;
        use std::time::Duration;
        let limiter = super::Limiter::new(NonZeroU32::new(1).unwrap(), Duration::from_millis(0));
        println!("limiter: {:?}", limiter.l);

        let limiter = super::Limiter::new(NonZeroU32::new(1).unwrap(), Duration::from_millis(100));
        let (tx, mut rx) = limiter.channel::<u64>(Arc::new(Queue::new(100)));

        let tx = tx.policy(|_v: &u64| -> Policy { Policy::Early });

        tokio::spawn(async move {
            for i in 0..20 {
                if let Err(e) = tx.send(i).await {
                    log::warn!("[send] {} => removed: {:?}  queue len: {}", i, e, tx.len());
                } else {
                    //log::info!("[send] {} => Ok   {}", i, tx.len());
                }
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        });

        while let Some(v) = rx.next().await {
            println!("{} queue recv: {:?}", Local::now().format("%Y-%m-%d %H:%M:%S%.3f %z"), v);
        }
    }
}