use std::sync::Mutex;
use crate::Error;
use crate::codec::DRAIN_FRAME_OVERHEAD;
struct QueueInner {
wire_data: Vec<u8>,
frame_count: u32,
read_pos: usize,
bytes_used: usize,
max_bytes: usize,
}
impl QueueInner {
fn new(max_bytes: usize) -> Self {
Self {
wire_data: Vec::new(),
frame_count: 0,
read_pos: 0,
bytes_used: 0,
max_bytes,
}
}
#[inline]
fn frame_cost(frame: &[u8]) -> usize {
DRAIN_FRAME_OVERHEAD + frame.len()
}
}
pub struct Queue {
inner: Mutex<QueueInner>,
}
impl Queue {
pub fn new(max_bytes: usize) -> Self {
Self {
inner: Mutex::new(QueueInner::new(max_bytes)),
}
}
pub fn unbounded() -> Self {
Self::new(0)
}
pub fn push(&self, frame: &[u8]) -> Result<(), Error> {
if frame.len() > u32::MAX as usize
|| DRAIN_FRAME_OVERHEAD.checked_add(frame.len()).is_none()
{
return Err(Error::PayloadTooLarge(frame.len()));
}
let cost = QueueInner::frame_cost(frame);
let mut inner = crate::lock_or_recover(&self.inner);
if inner.max_bytes > 0 && inner.bytes_used + cost > inner.max_bytes {
return Err(Error::ChannelFull);
}
if inner.frame_count == u32::MAX {
return Err(Error::ChannelFull);
}
inner
.wire_data
.extend_from_slice(&(frame.len() as u32).to_le_bytes());
inner.wire_data.extend_from_slice(frame);
inner.frame_count += 1;
inner.bytes_used = inner.bytes_used.saturating_add(cost);
Ok(())
}
#[must_use]
pub fn try_pop(&self) -> Option<Vec<u8>> {
let mut inner = crate::lock_or_recover(&self.inner);
if inner.frame_count == 0 {
return None;
}
let len_bytes: [u8; 4] = inner.wire_data[inner.read_pos..inner.read_pos + 4]
.try_into()
.unwrap();
let payload_len = u32::from_le_bytes(len_bytes) as usize;
let payload_start = inner.read_pos + 4;
let frame = inner.wire_data[payload_start..payload_start + payload_len].to_vec();
let cost = DRAIN_FRAME_OVERHEAD + payload_len;
inner.read_pos += cost;
inner.frame_count -= 1;
inner.bytes_used -= cost;
if inner.frame_count == 0 {
inner.wire_data.clear();
inner.read_pos = 0;
} else if inner.read_pos > inner.wire_data.len() / 2 {
let rp = inner.read_pos;
inner.wire_data.copy_within(rp.., 0);
let new_len = inner.wire_data.len() - rp;
inner.wire_data.truncate(new_len);
inner.read_pos = 0;
}
Some(frame)
}
#[must_use]
pub fn drain_all(&self) -> Vec<u8> {
let (wire_data, read_pos, frame_count) = {
let mut inner = crate::lock_or_recover(&self.inner);
if inner.frame_count == 0 {
return Vec::new();
}
let wire_data = std::mem::take(&mut inner.wire_data);
let read_pos = inner.read_pos;
let frame_count = inner.frame_count;
inner.read_pos = 0;
inner.frame_count = 0;
inner.bytes_used = 0;
(wire_data, read_pos, frame_count)
};
let live_data = &wire_data[read_pos..];
let output_size = 4 + live_data.len();
let mut buf = Vec::with_capacity(output_size);
buf.extend_from_slice(&frame_count.to_le_bytes());
buf.extend_from_slice(live_data);
buf
}
#[must_use]
pub fn frame_count(&self) -> usize {
crate::lock_or_recover(&self.inner).frame_count as usize
}
#[must_use]
pub fn bytes_used(&self) -> usize {
crate::lock_or_recover(&self.inner).bytes_used
}
#[must_use]
pub fn max_bytes(&self) -> usize {
crate::lock_or_recover(&self.inner).max_bytes
}
pub fn clear(&self) {
let mut inner = crate::lock_or_recover(&self.inner);
inner.wire_data.clear();
inner.frame_count = 0;
inner.read_pos = 0;
inner.bytes_used = 0;
}
}
impl std::fmt::Debug for Queue {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let inner = crate::lock_or_recover(&self.inner);
f.debug_struct("Queue")
.field("frame_count", &inner.frame_count)
.field("bytes_used", &inner.bytes_used)
.field("max_bytes", &inner.max_bytes)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn push_and_pop() {
let q = Queue::new(1024);
q.push(b"alpha").unwrap();
q.push(b"beta").unwrap();
q.push(b"gamma").unwrap();
assert_eq!(q.frame_count(), 3);
assert_eq!(q.try_pop().unwrap(), b"alpha");
assert_eq!(q.try_pop().unwrap(), b"beta");
assert_eq!(q.try_pop().unwrap(), b"gamma");
assert!(q.try_pop().is_none());
}
#[test]
fn push_within_limit() {
let q = Queue::new(16);
q.push(b"aaaa").unwrap(); q.push(b"bbbb").unwrap(); assert_eq!(q.frame_count(), 2);
assert_eq!(q.bytes_used(), 16);
}
#[test]
fn push_exceeds_limit() {
let q = Queue::new(16);
q.push(b"aaaa").unwrap(); q.push(b"bbbb").unwrap();
let err = q.push(b"cccc").unwrap_err();
assert!(matches!(err, Error::ChannelFull));
assert_eq!(err.to_string(), "channel full: byte limit reached");
assert_eq!(q.frame_count(), 2);
assert_eq!(q.try_pop().unwrap(), b"aaaa");
assert_eq!(q.try_pop().unwrap(), b"bbbb");
}
#[test]
fn drain_all_format() {
let q = Queue::new(1024);
q.push(b"hello").unwrap();
q.push(b"world").unwrap();
let blob = q.drain_all();
let count = u32::from_le_bytes(blob[0..4].try_into().unwrap());
assert_eq!(count, 2);
let len1 = u32::from_le_bytes(blob[4..8].try_into().unwrap()) as usize;
assert_eq!(len1, 5);
assert_eq!(&blob[8..8 + len1], b"hello");
let offset2 = 8 + len1;
let len2 = u32::from_le_bytes(blob[offset2..offset2 + 4].try_into().unwrap()) as usize;
assert_eq!(len2, 5);
assert_eq!(&blob[offset2 + 4..offset2 + 4 + len2], b"world");
assert_eq!(q.frame_count(), 0);
assert_eq!(q.bytes_used(), 0);
}
#[test]
fn drain_frees_capacity() {
let q = Queue::new(16);
q.push(b"aaaa").unwrap(); q.push(b"bbbb").unwrap();
assert!(q.push(b"cccc").is_err());
let blob = q.drain_all();
assert!(!blob.is_empty());
assert_eq!(q.bytes_used(), 0);
q.push(b"dddd").unwrap();
q.push(b"eeee").unwrap();
assert_eq!(q.frame_count(), 2);
}
#[test]
fn unbounded_mode() {
let q = Queue::unbounded();
assert_eq!(q.max_bytes(), 0);
for i in 0u32..10_000 {
q.push(&i.to_le_bytes()).unwrap();
}
assert_eq!(q.frame_count(), 10_000);
}
#[test]
fn frame_count_and_bytes() {
let q = Queue::new(1024);
assert_eq!(q.frame_count(), 0);
assert_eq!(q.bytes_used(), 0);
assert_eq!(q.max_bytes(), 1024);
q.push(b"abc").unwrap(); assert_eq!(q.frame_count(), 1);
assert_eq!(q.bytes_used(), 7);
q.push(b"de").unwrap(); assert_eq!(q.frame_count(), 2);
assert_eq!(q.bytes_used(), 13);
let _ = q.try_pop();
assert_eq!(q.frame_count(), 1);
assert_eq!(q.bytes_used(), 6);
}
#[test]
fn clear() {
let q = Queue::new(1024);
q.push(b"one").unwrap();
q.push(b"two").unwrap();
q.push(b"three").unwrap();
assert_eq!(q.frame_count(), 3);
q.clear();
assert_eq!(q.frame_count(), 0);
assert_eq!(q.bytes_used(), 0);
assert!(q.try_pop().is_none());
}
#[test]
fn concurrent_push_pop() {
use std::sync::Arc;
let q = Arc::new(Queue::unbounded());
let q_producer = Arc::clone(&q);
let q_consumer = Arc::clone(&q);
let producer = std::thread::spawn(move || {
for i in 0u32..1000 {
q_producer.push(&i.to_le_bytes()).unwrap();
}
});
let consumer = std::thread::spawn(move || {
let mut popped = 0usize;
loop {
if q_consumer.try_pop().is_some() {
popped += 1;
}
if popped >= 1000 {
break;
}
std::thread::yield_now();
}
popped
});
producer.join().unwrap();
let consumer_popped = consumer.join().unwrap();
let remaining = q.frame_count();
assert_eq!(consumer_popped + remaining, 1000);
}
#[test]
fn empty_drain() {
let q = Queue::new(1024);
let blob = q.drain_all();
assert!(blob.is_empty());
}
#[test]
fn drain_then_push() {
let q = Queue::new(1024);
q.push(b"first").unwrap();
let blob = q.drain_all();
assert!(!blob.is_empty());
q.push(b"second").unwrap();
assert_eq!(q.frame_count(), 1);
assert_eq!(q.try_pop().unwrap(), b"second");
}
}