use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tokio::sync::{mpsc, Notify};
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),
}
}
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 {
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();
}
}
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
}
}
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);
tokio::spawn(async move {
while let Some(msg) = source.recv().await {
producer.push(msg);
}
producer.close();
});
DropOldestReceiver::new(shared)
}