use alloc::vec::Vec;
pub trait BatchSink {
type Error;
fn write_batch(&mut self, bytes: &[u8]) -> Result<(), Self::Error>;
}
#[derive(Debug, Default)]
pub struct BatchBuffer {
bytes: Vec<u8>,
}
impl BatchBuffer {
#[must_use]
pub fn new() -> Self {
Self { bytes: Vec::new() }
}
#[must_use]
pub fn with_capacity(capacity: usize) -> Self {
Self {
bytes: Vec::with_capacity(capacity),
}
}
pub fn push_str(&mut self, text: &str) {
self.bytes.extend_from_slice(text.as_bytes());
}
pub fn push_bytes(&mut self, bytes: &[u8]) {
self.bytes.extend_from_slice(bytes);
}
#[must_use]
pub fn is_empty(&self) -> bool {
self.bytes.is_empty()
}
#[must_use]
pub fn len(&self) -> usize {
self.bytes.len()
}
#[must_use]
pub fn capacity(&self) -> usize {
self.bytes.capacity()
}
#[must_use]
pub fn as_bytes(&self) -> &[u8] {
&self.bytes
}
pub fn clear(&mut self) {
self.bytes.clear();
}
pub fn flush_to<S: BatchSink>(&mut self, sink: &mut S) -> Result<(), S::Error> {
if self.bytes.is_empty() {
return Ok(());
}
sink.write_batch(&self.bytes)?;
self.bytes.clear();
Ok(())
}
}
impl core::fmt::Write for BatchBuffer {
fn write_str(&mut self, text: &str) -> core::fmt::Result {
self.push_str(text);
Ok(())
}
}
#[cfg(feature = "std")]
pub struct WriteSink<W>(pub W);
#[cfg(feature = "std")]
impl<W: std::io::Write> BatchSink for WriteSink<W> {
type Error = std::io::Error;
fn write_batch(&mut self, bytes: &[u8]) -> Result<(), std::io::Error> {
self.0.write_all(bytes)
}
}
#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use alloc::vec::Vec;
use core::convert::Infallible;
use core::fmt::Write as _;
#[derive(Default)]
struct RecordingSink {
writes: usize,
bytes: Vec<u8>,
}
impl BatchSink for RecordingSink {
type Error = Infallible;
fn write_batch(&mut self, bytes: &[u8]) -> Result<(), Infallible> {
self.writes += 1;
self.bytes.extend_from_slice(bytes);
Ok(())
}
}
#[test]
fn many_appends_flush_in_one_write() {
let mut buffer = BatchBuffer::new();
buffer.push_str("alpha\n");
buffer.push_str("beta\n");
buffer.push_str("gamma\n");
let mut sink = RecordingSink::default();
buffer.flush_to(&mut sink).unwrap();
assert_eq!(sink.writes, 1, "three appends flush in a single write");
assert_eq!(&sink.bytes, b"alpha\nbeta\ngamma\n");
assert!(buffer.is_empty(), "buffer cleared after flush");
}
#[test]
fn empty_flush_does_not_write() {
let mut buffer = BatchBuffer::new();
let mut sink = RecordingSink::default();
buffer.flush_to(&mut sink).unwrap();
assert_eq!(sink.writes, 0, "no write for an empty batch");
}
#[test]
fn buffer_reuses_after_flush() {
let mut buffer = BatchBuffer::with_capacity(64);
let mut sink = RecordingSink::default();
buffer.push_str("first\n");
buffer.flush_to(&mut sink).unwrap();
buffer.push_str("second\n");
buffer.flush_to(&mut sink).unwrap();
assert_eq!(sink.writes, 2, "one write per batch");
assert_eq!(&sink.bytes, b"first\nsecond\n");
}
#[test]
fn push_bytes_carries_raw_bytes() {
let mut buffer = BatchBuffer::new();
buffer.push_bytes(&[0xDE, 0xAD, 0xBE, 0xEF]);
let mut sink = RecordingSink::default();
buffer.flush_to(&mut sink).unwrap();
assert_eq!(&sink.bytes, &[0xDE, 0xAD, 0xBE, 0xEF]);
}
#[test]
fn with_capacity_preallocates() {
let buffer = BatchBuffer::with_capacity(8192);
assert!(buffer.capacity() >= 8192, "reserves the requested bytes");
assert!(buffer.is_empty(), "reserved but holds nothing yet");
}
#[cfg(feature = "std")]
#[test]
fn write_sink_adapts_io_write() {
let mut out: Vec<u8> = Vec::new();
let mut buffer = BatchBuffer::new();
buffer.push_str("alpha\n");
buffer.push_str("beta\n");
buffer.flush_to(&mut WriteSink(&mut out)).unwrap();
assert_eq!(&out, b"alpha\nbeta\n");
}
#[test]
fn fmt_write_appends_formatted_output() {
let mut buffer = BatchBuffer::new();
let seq = 12_345;
let peer = "10.0.0.1:443";
writeln!(buffer, "seq={seq} peer={peer}").unwrap();
let mut sink = RecordingSink::default();
buffer.flush_to(&mut sink).unwrap();
assert_eq!(&sink.bytes, b"seq=12345 peer=10.0.0.1:443\n");
}
}