use bytes::Bytes;
use std::collections::VecDeque;
pub(crate) struct EgressBuffer {
chunks: VecDeque<(Bytes, usize)>, write_offset: usize,
total_bytes: usize,
message_count: usize,
pub(crate) peak_messages: usize,
pub(crate) peak_bytes: usize,
}
impl EgressBuffer {
pub(crate) fn new() -> Self {
Self {
chunks: VecDeque::new(),
write_offset: 0,
total_bytes: 0,
message_count: 0,
peak_messages: 0,
peak_bytes: 0,
}
}
pub(crate) fn push(&mut self, data: Bytes, msg_count: usize) {
if data.is_empty() {
return;
}
self.total_bytes += data.len();
self.message_count += msg_count;
self.chunks.push_back((data, msg_count));
self.peak_messages = self.peak_messages.max(self.message_count);
self.peak_bytes = self.peak_bytes.max(self.total_bytes);
}
pub(crate) fn peak_messages(&self) -> usize {
self.peak_messages
}
pub(crate) fn peak_bytes(&self) -> usize {
self.peak_bytes
}
pub(crate) fn push_priority(&mut self, data: Bytes) {
if data.is_empty() {
return;
}
self.total_bytes += data.len();
if self.write_offset > 0 {
self.chunks.insert(1, (data, 0));
} else {
self.chunks.push_front((data, 0));
}
}
pub(crate) fn current_slice(&self) -> Option<&[u8]> {
self.chunks.front().map(|(c, _)| &c[self.write_offset..])
}
pub(crate) fn fill_slices<'a>(&'a self, out: &mut [std::io::IoSlice<'a>]) -> usize {
if self.chunks.is_empty() || out.is_empty() {
return 0;
}
let mut iter = self.chunks.iter();
let mut count = 0;
if let Some((front, _)) = iter.next() {
out[count] = std::io::IoSlice::new(&front[self.write_offset..]);
count += 1;
}
for (chunk, _) in iter {
if count >= out.len() {
break;
}
out[count] = std::io::IoSlice::new(chunk);
count += 1;
}
count
}
pub(crate) fn advance(&mut self, mut n: usize) -> usize {
self.total_bytes = self.total_bytes.saturating_sub(n);
let mut popped_messages = 0;
while n > 0 {
if let Some((front, _)) = self.chunks.front() {
let remaining = front.len() - self.write_offset;
if n >= remaining {
n -= remaining;
let (_, mc) = self.chunks.pop_front().unwrap();
self.message_count = self.message_count.saturating_sub(mc);
popped_messages += mc; self.write_offset = 0;
} else {
self.write_offset += n;
break;
}
} else {
break;
}
}
popped_messages
}
pub(crate) fn pending_messages(&self) -> usize {
self.message_count
}
pub(crate) fn total_pending_bytes(&self) -> usize {
self.total_bytes
}
pub(crate) fn is_empty(&self) -> bool {
self.chunks.is_empty()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn partial_writes_advance_and_pop() {
let mut buf = EgressBuffer::new();
buf.push(Bytes::from_static(b"0123456789"), 1);
assert_eq!(buf.current_slice().unwrap(), b"0123456789");
buf.advance(6);
assert_eq!(buf.current_slice().unwrap(), b"6789");
assert_eq!(buf.total_pending_bytes(), 4);
buf.advance(4);
assert!(buf.is_empty());
assert_eq!(buf.total_pending_bytes(), 0);
assert_eq!(buf.pending_messages(), 0);
}
#[test]
fn total_pending_bytes_tracks_advances() {
let mut buf = EgressBuffer::new();
buf.push(Bytes::from_static(b"aaa"), 1);
buf.push(Bytes::from_static(b"bb"), 1);
buf.push(Bytes::from_static(b"c"), 1);
assert_eq!(buf.total_pending_bytes(), 6);
assert_eq!(buf.pending_messages(), 3);
buf.advance(3);
assert_eq!(buf.total_pending_bytes(), 3);
assert_eq!(buf.pending_messages(), 2);
buf.advance(2);
assert_eq!(buf.total_pending_bytes(), 1);
assert_eq!(buf.pending_messages(), 1);
buf.advance(1);
assert!(buf.is_empty());
assert_eq!(buf.pending_messages(), 0);
}
#[test]
fn push_priority_front_of_non_empty() {
let mut buf = EgressBuffer::new();
buf.push(Bytes::from_static(b"data"), 1);
buf.push_priority(Bytes::from_static(b"PONG"));
assert_eq!(buf.current_slice().unwrap(), b"PONG");
assert_eq!(buf.pending_messages(), 1); buf.advance(4);
assert_eq!(buf.current_slice().unwrap(), b"data");
buf.advance(4);
assert_eq!(buf.pending_messages(), 0);
}
#[test]
fn advance_spans_multiple_chunks() {
let mut buf = EgressBuffer::new();
buf.push(Bytes::from_static(b"0123456789"), 1); buf.push(Bytes::from_static(b"abcde"), 1); buf.push(Bytes::from_static(b"XY"), 1); assert_eq!(buf.pending_messages(), 3);
buf.advance(12);
assert_eq!(buf.current_slice().unwrap(), b"cde");
assert_eq!(buf.total_pending_bytes(), 5);
assert_eq!(buf.pending_messages(), 2);
buf.advance(5);
assert!(buf.is_empty());
assert_eq!(buf.pending_messages(), 0);
}
#[test]
fn push_priority_after_partial_write() {
let mut buf = EgressBuffer::new();
buf.push(Bytes::from_static(b"data"), 1);
buf.advance(2); buf.push_priority(Bytes::from_static(b"PONG"));
assert_eq!(buf.current_slice().unwrap(), b"ta");
assert_eq!(buf.pending_messages(), 1); buf.advance(2);
assert_eq!(buf.current_slice().unwrap(), b"PONG");
buf.advance(4);
assert!(buf.is_empty());
assert_eq!(buf.pending_messages(), 0);
}
}