use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use crate::buf::RingBuffer;
use crate::err::{Error, Result};
pub struct Receiver<T, const N: usize> {
buf: Arc<RingBuffer<T, N>>,
closed: Arc<AtomicBool>,
}
impl<T, const N: usize> Receiver<T, N> {
#[allow(dead_code)]
pub(crate) fn new(buf: Arc<RingBuffer<T, N>>) -> Self {
Self {
buf,
closed: Arc::new(AtomicBool::new(false)),
}
}
pub(crate) fn new_with_closed(buf: Arc<RingBuffer<T, N>>, closed: Arc<AtomicBool>) -> Self {
Self { buf, closed }
}
#[allow(dead_code)]
pub(crate) fn closed_flag(&self) -> Arc<AtomicBool> {
self.closed.clone()
}
#[inline]
pub fn recv(&self) -> Result<T> {
if self.closed.load(Ordering::Acquire) && self.buf.is_empty() {
return Err(Error::ChannelClosed);
}
self.buf.recv()
}
#[inline]
pub fn try_recv(&self) -> Result<T> {
self.recv()
}
pub fn recv_spin(&self) -> Result<T> {
loop {
if self.closed.load(Ordering::Acquire) && self.buf.is_empty() {
return Err(Error::ChannelClosed);
}
if let Ok(v) = self.buf.recv() {
return Ok(v);
}
std::hint::spin_loop();
}
}
#[inline]
pub fn is_empty(&self) -> bool {
self.buf.is_empty()
}
#[inline]
pub fn len(&self) -> usize {
self.buf.len()
}
#[inline]
pub fn capacity(&self) -> usize {
N
}
#[inline]
pub fn is_closed(&self) -> bool {
self.closed.load(Ordering::Acquire)
}
pub fn close(&self) {
self.closed.store(true, Ordering::Release);
}
pub fn recv_batch(&self, buf: &mut [T], max: usize) -> usize {
let count = max.min(buf.len());
let mut received = 0;
for item in buf.iter_mut().take(count) {
match self.buf.recv() {
Ok(value) => {
*item = value;
received += 1;
}
Err(_) => break,
}
}
received
}
pub fn recv_timeout(&self, timeout: std::time::Duration) -> Result<T> {
let start = std::time::Instant::now();
loop {
if self.closed.load(Ordering::Acquire) && self.buf.is_empty() {
return Err(Error::ChannelClosed);
}
match self.buf.recv() {
Ok(value) => return Ok(value),
Err(e) => {
if start.elapsed() >= timeout {
return Err(e);
}
std::hint::spin_loop();
}
}
}
}
pub fn drain(&self) -> DrainIter<'_, T, N> {
DrainIter { receiver: self }
}
}
pub struct DrainIter<'a, T, const N: usize> {
receiver: &'a Receiver<T, N>,
}
impl<T, const N: usize> Iterator for DrainIter<'_, T, N> {
type Item = T;
fn next(&mut self) -> Option<Self::Item> {
self.receiver.buf.recv().ok()
}
}
impl<T, const N: usize> Clone for Receiver<T, N> {
fn clone(&self) -> Self {
Self {
buf: self.buf.clone(),
closed: self.closed.clone(),
}
}
}
unsafe impl<T: Send, const N: usize> Send for Receiver<T, N> {}
unsafe impl<T: Send, const N: usize> Sync for Receiver<T, N> {}