use std::sync::Mutex;
use crate::codec::DRAIN_FRAME_OVERHEAD;
const DEFAULT_CAPACITY: usize = 64 * 1024;
struct Inner {
wire_data: Vec<u8>,
frame_count: u32,
read_pos: usize,
bytes_used: usize,
capacity: usize,
}
impl Inner {
fn new(capacity: usize) -> Self {
Self {
wire_data: Vec::new(),
frame_count: 0,
read_pos: 0,
bytes_used: 0,
capacity,
}
}
#[inline]
fn frame_cost(frame: &[u8]) -> usize {
DRAIN_FRAME_OVERHEAD + frame.len()
}
fn drop_oldest(&mut self) -> bool {
if self.frame_count == 0 {
return false;
}
let len_bytes: [u8; 4] = self.wire_data[self.read_pos..self.read_pos + 4]
.try_into()
.unwrap();
let payload_len = u32::from_le_bytes(len_bytes) as usize;
let cost = DRAIN_FRAME_OVERHEAD + payload_len;
self.read_pos += cost;
self.frame_count -= 1;
self.bytes_used -= cost;
true
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum PushOutcome {
Accepted(usize),
TooLarge,
}
pub struct RingBuffer {
inner: Mutex<Inner>,
}
impl RingBuffer {
pub fn new(capacity: usize) -> Self {
assert!(
capacity > DRAIN_FRAME_OVERHEAD,
"capacity must be at least {} bytes (DRAIN_FRAME_OVERHEAD + 1)",
DRAIN_FRAME_OVERHEAD + 1,
);
Self {
inner: Mutex::new(Inner::new(capacity)),
}
}
pub fn with_default_capacity() -> Self {
Self::new(DEFAULT_CAPACITY)
}
pub fn push(&self, frame: &[u8]) -> usize {
match self.push_checked(frame) {
PushOutcome::Accepted(n) => n,
PushOutcome::TooLarge => 0,
}
}
#[must_use]
pub fn push_checked(&self, frame: &[u8]) -> PushOutcome {
if frame.len() > u32::MAX as usize
|| DRAIN_FRAME_OVERHEAD.checked_add(frame.len()).is_none()
{
return PushOutcome::TooLarge;
}
let cost = Inner::frame_cost(frame);
let mut inner = crate::lock_or_recover(&self.inner);
if cost > inner.capacity {
return PushOutcome::TooLarge;
}
let mut dropped = 0usize;
while inner.bytes_used + cost > inner.capacity {
if !inner.drop_oldest() {
break;
}
dropped += 1;
}
if inner.read_pos > 0 && 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;
}
if inner.frame_count == u32::MAX {
return PushOutcome::TooLarge;
}
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 += cost;
PushOutcome::Accepted(dropped)
}
#[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 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;
}
Some(frame)
}
#[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 capacity(&self) -> usize {
crate::lock_or_recover(&self.inner).capacity
}
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 RingBuffer {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let inner = crate::lock_or_recover(&self.inner);
f.debug_struct("RingBuffer")
.field("frame_count", &inner.frame_count)
.field("bytes_used", &inner.bytes_used)
.field("capacity", &inner.capacity)
.finish()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn push_and_pop() {
let rb = RingBuffer::new(1024);
let _ = rb.push(b"alpha");
let _ = rb.push(b"beta");
let _ = rb.push(b"gamma");
assert_eq!(rb.frame_count(), 3);
assert_eq!(rb.try_pop().unwrap(), b"alpha");
assert_eq!(rb.try_pop().unwrap(), b"beta");
assert_eq!(rb.try_pop().unwrap(), b"gamma");
assert!(rb.try_pop().is_none());
}
#[test]
fn drain_all_format() {
let rb = RingBuffer::new(1024);
let _ = rb.push(b"hello");
let _ = rb.push(b"world");
let blob = rb.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!(rb.frame_count(), 0);
assert_eq!(rb.bytes_used(), 0);
}
#[test]
fn overflow_drops_oldest() {
let rb = RingBuffer::new(16);
let dropped = rb.push(b"aaaa"); assert_eq!(dropped, 0);
let dropped = rb.push(b"bbbb"); assert_eq!(dropped, 0);
let dropped = rb.push(b"cccc"); assert_eq!(dropped, 1);
assert_eq!(rb.frame_count(), 2);
assert_eq!(rb.try_pop().unwrap(), b"bbbb");
assert_eq!(rb.try_pop().unwrap(), b"cccc");
}
#[test]
fn empty_drain() {
let rb = RingBuffer::new(1024);
let blob = rb.drain_all();
assert!(blob.is_empty());
}
#[test]
fn frame_count_and_bytes() {
let rb = RingBuffer::new(1024);
assert_eq!(rb.frame_count(), 0);
assert_eq!(rb.bytes_used(), 0);
assert_eq!(rb.capacity(), 1024);
let _ = rb.push(b"abc"); assert_eq!(rb.frame_count(), 1);
assert_eq!(rb.bytes_used(), 7);
let _ = rb.push(b"de"); assert_eq!(rb.frame_count(), 2);
assert_eq!(rb.bytes_used(), 13);
let _ = rb.try_pop();
assert_eq!(rb.frame_count(), 1);
assert_eq!(rb.bytes_used(), 6);
}
#[test]
fn clear() {
let rb = RingBuffer::new(1024);
let _ = rb.push(b"one");
let _ = rb.push(b"two");
let _ = rb.push(b"three");
assert_eq!(rb.frame_count(), 3);
rb.clear();
assert_eq!(rb.frame_count(), 0);
assert_eq!(rb.bytes_used(), 0);
assert!(rb.try_pop().is_none());
}
#[tokio::test]
async fn concurrent_push_pop() {
use std::sync::Arc;
let rb = Arc::new(RingBuffer::new(64 * 1024));
let rb_producer = Arc::clone(&rb);
let rb_consumer = Arc::clone(&rb);
let producer = tokio::spawn(async move {
for i in 0u32..1000 {
let _ = rb_producer.push(&i.to_le_bytes());
}
});
let consumer = tokio::spawn(async move {
let mut popped = 0usize;
loop {
if let Some(_frame) = rb_consumer.try_pop() {
popped += 1;
} else {
tokio::task::yield_now().await;
}
if popped >= 1000 {
break;
}
}
popped
});
producer.await.unwrap();
let consumer_popped = consumer.await.unwrap();
let remaining = rb.frame_count();
assert_eq!(consumer_popped + remaining, 1000);
}
#[test]
fn single_large_frame() {
let rb = RingBuffer::new(32);
let _ = rb.push(b"ok"); let dropped = rb.push(&[0xFFu8; 100]); assert_eq!(dropped, 0);
assert_eq!(rb.frame_count(), 1);
assert_eq!(rb.try_pop().unwrap(), b"ok");
}
#[test]
fn drain_then_push() {
let rb = RingBuffer::new(1024);
let _ = rb.push(b"first");
let blob = rb.drain_all();
assert!(!blob.is_empty());
let _ = rb.push(b"second");
assert_eq!(rb.frame_count(), 1);
assert_eq!(rb.try_pop().unwrap(), b"second");
}
#[test]
fn overflow_cascade() {
let rb = RingBuffer::new(8);
let _ = rb.push(b"aaaa"); assert_eq!(rb.frame_count(), 1);
let dropped = rb.push(&[0u8; 6]);
assert_eq!(dropped, 0);
assert_eq!(rb.frame_count(), 1);
assert_eq!(rb.try_pop().unwrap(), b"aaaa");
}
#[test]
#[should_panic(expected = "capacity must be at least 5 bytes")]
fn tiny_capacity_panics() {
RingBuffer::new(4); }
#[test]
fn with_default_capacity() {
let rb = RingBuffer::with_default_capacity();
assert_eq!(rb.capacity(), 64 * 1024);
}
}