use std::io::{self, Write};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[non_exhaustive]
pub enum FramingContract {
PerRecord,
WholePayload,
}
pub trait RecordFramer: Send {
fn push(&mut self, bytes: &[u8]) -> io::Result<()>;
fn finish(&mut self) -> io::Result<()>;
fn pop(&mut self) -> Option<Vec<u8>>;
fn decoded_bytes(&self) -> u64;
}
pub struct FramerWriter {
framer: Box<dyn RecordFramer>,
}
impl FramerWriter {
#[must_use]
pub fn new(framer: Box<dyn RecordFramer>) -> FramerWriter {
FramerWriter { framer }
}
#[must_use]
pub fn framer(&self) -> &dyn RecordFramer {
&*self.framer
}
pub fn framer_mut(&mut self) -> &mut dyn RecordFramer {
&mut *self.framer
}
#[must_use]
pub fn into_inner(self) -> Box<dyn RecordFramer> {
self.framer
}
}
impl std::fmt::Debug for FramerWriter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("FramerWriter").finish_non_exhaustive()
}
}
impl Write for FramerWriter {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
self.framer.push(buf)?;
Ok(buf.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::VecDeque;
#[derive(Default)]
struct WholeFramer {
buf: Vec<u8>,
ready: VecDeque<Vec<u8>>,
decoded: u64,
}
impl RecordFramer for WholeFramer {
fn push(&mut self, bytes: &[u8]) -> io::Result<()> {
self.decoded += bytes.len() as u64;
self.buf.extend_from_slice(bytes);
Ok(())
}
fn finish(&mut self) -> io::Result<()> {
if !self.buf.is_empty() {
self.ready.push_back(std::mem::take(&mut self.buf));
}
Ok(())
}
fn pop(&mut self) -> Option<Vec<u8>> {
self.ready.pop_front()
}
fn decoded_bytes(&self) -> u64 {
self.decoded
}
}
#[test]
fn record_framer_is_object_safe() {
let mut framer: Box<dyn RecordFramer> = Box::new(WholeFramer::default());
framer.push(b"x").unwrap();
framer.finish().unwrap();
assert_eq!(framer.pop(), Some(b"x".to_vec()));
}
#[test]
fn framer_writer_forwards_writes_to_push_and_counts_bytes() {
let mut w = FramerWriter::new(Box::new(WholeFramer::default()));
w.write_all(b"a\nb\n").unwrap();
assert_eq!(
w.framer().decoded_bytes(),
4,
"every written byte reaches push"
);
w.framer_mut().finish().unwrap();
let mut out = Vec::new();
while let Some(r) = w.framer_mut().pop() {
out.push(r);
}
assert_eq!(out, vec![b"a\nb\n".to_vec()]);
}
#[test]
fn framer_writer_into_inner_reclaims_the_framer() {
let w = FramerWriter::new(Box::new(WholeFramer::default()));
let framer = w.into_inner();
assert_eq!(framer.decoded_bytes(), 0);
}
}