mecha10-core 0.6.2

//! Bounded buffer that evicts the oldest queued item when full.
//!
//! Backs [`super::Receiver::with_drop_oldest`].

use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::{mpsc, Notify};

/// Shared state backing a [`DropOldestReceiver`]/producer pair.
///
/// Unlike `tokio::sync::mpsc`, this buffer allows the *producer* side to evict
/// the oldest queued item when the buffer is full, which is what a true
/// "keep only the N most recent messages" strategy requires. `tokio::sync::mpsc`
/// only allows the consumer to pop, so it cannot implement this policy.
pub(super) struct DropOldestShared<T> {
    queue: Mutex<VecDeque<T>>,
    capacity: usize,
    notify: Notify,
    closed: AtomicBool,
}

impl<T> DropOldestShared<T> {
    pub(super) fn new(capacity: usize) -> Self {
        Self {
            queue: Mutex::new(VecDeque::with_capacity(capacity)),
            capacity,
            notify: Notify::new(),
            closed: AtomicBool::new(false),
        }
    }

    /// Push a message, evicting the oldest one first if the buffer is full.
    pub(super) fn push(&self, msg: T) {
        let mut queue = self.queue.lock().unwrap_or_else(|poisoned| poisoned.into_inner());
        if queue.len() >= self.capacity {
            // Drop the oldest message to make room for the newest one.
            queue.pop_front();
        }
        queue.push_back(msg);
        drop(queue);
        self.notify.notify_one();
    }

    pub(super) fn close(&self) {
        self.closed.store(true, Ordering::Release);
        self.notify.notify_one();
    }
}

/// Consumer handle for a drop-oldest bounded buffer.
///
/// See [`super::Receiver::with_drop_oldest`] for the public API this backs.
pub(crate) struct DropOldestReceiver<T> {
    shared: Arc<DropOldestShared<T>>,
}

impl<T> DropOldestReceiver<T> {
    pub(super) fn new(shared: Arc<DropOldestShared<T>>) -> Self {
        Self { shared }
    }

    pub(crate) async fn recv(&mut self) -> Option<T> {
        loop {
            {
                let mut queue = self
                    .shared
                    .queue
                    .lock()
                    .unwrap_or_else(|poisoned| poisoned.into_inner());
                if let Some(msg) = queue.pop_front() {
                    return Some(msg);
                }
                if self.shared.closed.load(Ordering::Acquire) {
                    return None;
                }
            }
            self.shared.notify.notified().await;
        }
    }

    pub(crate) fn try_recv(&mut self) -> std::result::Result<T, mpsc::error::TryRecvError> {
        let mut queue = self
            .shared
            .queue
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner());
        if let Some(msg) = queue.pop_front() {
            Ok(msg)
        } else if self.shared.closed.load(Ordering::Acquire) {
            Err(mpsc::error::TryRecvError::Disconnected)
        } else {
            Err(mpsc::error::TryRecvError::Empty)
        }
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.shared
            .queue
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .is_empty()
    }

    pub(crate) fn len(&self) -> usize {
        self.shared
            .queue
            .lock()
            .unwrap_or_else(|poisoned| poisoned.into_inner())
            .len()
    }

    pub(crate) fn capacity(&self) -> usize {
        self.shared.capacity
    }
}

/// Spawn a task that drains `source` into a fresh drop-oldest buffer of the
/// given `capacity`, returning the consumer handle.
///
/// This is the shared scaffolding behind [`super::Receiver::with_drop_oldest`].
pub(super) fn spawn_drop_oldest<T>(mut source: mpsc::UnboundedReceiver<T>, capacity: usize) -> DropOldestReceiver<T>
where
    T: Send + 'static,
{
    let shared = Arc::new(DropOldestShared::new(capacity));
    let producer = Arc::clone(&shared);

    // Spawn task to forward messages, evicting the oldest buffered
    // message whenever the buffer is full so the newest message is
    // always kept (true "drop oldest" semantics).
    tokio::spawn(async move {
        while let Some(msg) = source.recv().await {
            producer.push(msg);
        }
        // Source closed; mark the buffer closed so consumers see `None`
        // once it's drained instead of waiting forever.
        producer.close();
    });

    DropOldestReceiver::new(shared)
}