use std::cmp;
use std::io;
use std::io::{Read, Write};
use super::{TReadTransport, TReadTransportFactory, TWriteTransport, TWriteTransportFactory};
const READ_CAPACITY: usize = 4096;
const WRITE_CAPACITY: usize = 4096;
#[derive(Debug)]
pub struct TBufferedReadTransport<C>
where
C: Read,
{
buf: Box<[u8]>,
pos: usize,
cap: usize,
chan: C,
}
impl<C> TBufferedReadTransport<C>
where
C: Read,
{
pub fn new(channel: C) -> TBufferedReadTransport<C> {
TBufferedReadTransport::with_capacity(READ_CAPACITY, channel)
}
pub fn with_capacity(read_capacity: usize, channel: C) -> TBufferedReadTransport<C> {
TBufferedReadTransport {
buf: vec![0; read_capacity].into_boxed_slice(),
pos: 0,
cap: 0,
chan: channel,
}
}
fn get_bytes(&mut self) -> io::Result<&[u8]> {
if self.cap - self.pos == 0 {
self.pos = 0;
self.cap = self.chan.read(&mut self.buf)?;
}
Ok(&self.buf[self.pos..self.cap])
}
fn consume(&mut self, consumed: usize) {
self.pos = cmp::min(self.cap, self.pos + consumed);
}
}
impl<C> Read for TBufferedReadTransport<C>
where
C: Read,
{
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
let mut bytes_read = 0;
loop {
let nread = {
let avail_bytes = self.get_bytes()?;
let avail_space = buf.len() - bytes_read;
let nread = cmp::min(avail_space, avail_bytes.len());
buf[bytes_read..(bytes_read + nread)].copy_from_slice(&avail_bytes[..nread]);
nread
};
self.consume(nread);
bytes_read += nread;
if bytes_read == buf.len() || nread == 0 {
break;
}
}
Ok(bytes_read)
}
}
#[derive(Default)]
pub struct TBufferedReadTransportFactory;
impl TBufferedReadTransportFactory {
pub fn new() -> TBufferedReadTransportFactory {
TBufferedReadTransportFactory {}
}
}
impl TReadTransportFactory for TBufferedReadTransportFactory {
fn create(&self, channel: Box<dyn Read + Send>) -> Box<dyn TReadTransport + Send> {
Box::new(TBufferedReadTransport::new(channel))
}
}
#[derive(Debug)]
pub struct TBufferedWriteTransport<C>
where
C: Write,
{
buf: Vec<u8>,
cap: usize,
channel: C,
}
impl<C> TBufferedWriteTransport<C>
where
C: Write,
{
pub fn new(channel: C) -> TBufferedWriteTransport<C> {
TBufferedWriteTransport::with_capacity(WRITE_CAPACITY, channel)
}
pub fn with_capacity(write_capacity: usize, channel: C) -> TBufferedWriteTransport<C> {
assert!(
write_capacity > 0,
"write buffer size must be a positive integer"
);
TBufferedWriteTransport {
buf: Vec::with_capacity(write_capacity),
cap: write_capacity,
channel: channel,
}
}
}
impl<C> Write for TBufferedWriteTransport<C>
where
C: Write,
{
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if !buf.is_empty() {
let mut avail_bytes;
loop {
avail_bytes = cmp::min(buf.len(), self.cap - self.buf.len());
if avail_bytes == 0 {
self.flush()?;
} else {
break;
}
}
let avail_bytes = avail_bytes;
self.buf.extend_from_slice(&buf[..avail_bytes]);
assert!(self.buf.len() <= self.cap, "copy overflowed buffer");
Ok(avail_bytes)
} else {
Ok(0)
}
}
fn flush(&mut self) -> io::Result<()> {
self.channel.write_all(&self.buf)?;
self.channel.flush()?;
self.buf.clear();
Ok(())
}
}
#[derive(Default)]
pub struct TBufferedWriteTransportFactory;
impl TBufferedWriteTransportFactory {
pub fn new() -> TBufferedWriteTransportFactory {
TBufferedWriteTransportFactory {}
}
}
impl TWriteTransportFactory for TBufferedWriteTransportFactory {
fn create(&self, channel: Box<dyn Write + Send>) -> Box<dyn TWriteTransport + Send> {
Box::new(TBufferedWriteTransport::new(channel))
}
}
#[cfg(test)]
mod tests {
use std::io::{Read, Write};
use super::*;
use transport::TBufferChannel;
#[test]
fn must_return_zero_if_read_buffer_is_empty() {
let mem = TBufferChannel::with_capacity(10, 0);
let mut t = TBufferedReadTransport::with_capacity(10, mem);
let mut b = vec![0; 10];
let read_result = t.read(&mut b);
assert_eq!(read_result.unwrap(), 0);
}
#[test]
fn must_return_zero_if_caller_reads_into_zero_capacity_buffer() {
let mem = TBufferChannel::with_capacity(10, 0);
let mut t = TBufferedReadTransport::with_capacity(10, mem);
let read_result = t.read(&mut []);
assert_eq!(read_result.unwrap(), 0);
}
#[test]
fn must_return_zero_if_nothing_more_can_be_read() {
let mem = TBufferChannel::with_capacity(4, 0);
let mut t = TBufferedReadTransport::with_capacity(4, mem);
t.chan.set_readable_bytes(&[0, 1, 2, 3]);
let mut buf = vec![0u8; 4];
let read_result = t.read(&mut buf);
assert_eq!(read_result.unwrap(), 4);
assert_eq!(&buf, &[0, 1, 2, 3]);
let buf_again = vec![0u8; 4];
let read_result = t.read(&mut buf);
assert_eq!(read_result.unwrap(), 0);
assert_eq!(&buf_again, &[0, 0, 0, 0])
}
#[test]
fn must_fill_user_buffer_with_only_as_many_bytes_as_available() {
let mem = TBufferChannel::with_capacity(4, 0);
let mut t = TBufferedReadTransport::with_capacity(4, mem);
t.chan.set_readable_bytes(&[0, 1, 2, 3]);
let mut buf = vec![0u8; 8];
let read_result = t.read(&mut buf);
assert_eq!(read_result.unwrap(), 4);
assert_eq!(&buf[..4], &[0, 1, 2, 3]);
let read_result = t.read(&mut buf[4..]);
assert_eq!(read_result.unwrap(), 0);
assert_eq!(&buf, &[0, 1, 2, 3, 0, 0, 0, 0])
}
#[test]
fn must_read_successfully() {
let mem = TBufferChannel::with_capacity(10, 0);
let mut t = TBufferedReadTransport::with_capacity(2, mem);
let mut readable_bytes = [0u8; 10];
for i in 0..10 {
readable_bytes[i] = i as u8;
}
t.chan.set_readable_bytes(&readable_bytes);
let mut buf = [0u8; 8];
let read_result = t.read(&mut buf);
assert_eq!(read_result.unwrap(), 8);
assert_eq!(&buf, &[0, 1, 2, 3, 4, 5, 6, 7]);
for i in 0..8 {
buf[i] = 0;
}
let read_result = t.read(&mut buf);
assert_eq!(read_result.unwrap(), 2);
assert_eq!(&buf[0..2], &[8, 9]);
assert_eq!(&buf[2..], &[0, 0, 0, 0, 0, 0]);
let read_result = t.read(&mut buf);
assert_eq!(read_result.unwrap(), 0);
assert_eq!(&buf[0..2], &[8, 9]);
assert_eq!(&buf[2..], &[0, 0, 0, 0, 0, 0]);
}
#[test]
fn must_return_error_when_nothing_can_be_written_to_underlying_channel() {
let mem = TBufferChannel::with_capacity(0, 0);
let mut t = TBufferedWriteTransport::with_capacity(1, mem);
let b = vec![0; 10];
let r = t.write(&b);
assert_eq!(r.unwrap(), 1);
let r = t.write(&b[1..]);
assert!(r.is_err());
}
#[test]
fn must_return_zero_if_caller_calls_write_with_empty_buffer() {
let mem = TBufferChannel::with_capacity(0, 10);
let mut t = TBufferedWriteTransport::with_capacity(10, mem);
let r = t.write(&[]);
let expected: [u8; 0] = [];
assert_eq!(r.unwrap(), 0);
assert_eq_transport_written_bytes!(t, expected);
}
#[test]
fn must_auto_flush_if_write_buffer_full() {
let mem = TBufferChannel::with_capacity(0, 8);
let mut t = TBufferedWriteTransport::with_capacity(4, mem);
let b0 = [0x00, 0x01, 0x02, 0x03];
let b1 = [0x04, 0x05, 0x06, 0x07];
let r = t.write(&b0);
assert_eq!(r.unwrap(), 4);
let r = t.write(&b1);
assert_eq!(r.unwrap(), 4);
assert_eq_transport_num_written_bytes!(t, 4);
assert_eq_transport_written_bytes!(t, b0);
t.channel.empty_write_buffer();
assert!(t.flush().is_ok());
assert_eq_transport_written_bytes!(t, b1);
}
#[test]
fn must_write_to_inner_transport_on_flush() {
let mem = TBufferChannel::with_capacity(10, 10);
let mut t = TBufferedWriteTransport::new(mem);
let b: [u8; 5] = [0, 1, 2, 3, 4];
assert_eq!(t.write(&b).unwrap(), 5);
assert_eq_transport_num_written_bytes!(t, 0);
assert!(t.flush().is_ok());
assert_eq_transport_written_bytes!(t, b);
}
#[test]
fn must_write_successfully_after_flush() {
let mem = TBufferChannel::with_capacity(0, 5);
let mut t = TBufferedWriteTransport::with_capacity(5, mem);
let b: [u8; 5] = [0, 1, 2, 3, 4];
assert_eq!(t.write(&b).unwrap(), 5);
assert!(t.flush().is_ok());
assert_eq_transport_written_bytes!(t, b);
t.channel.empty_write_buffer();
assert_eq!(t.write(&b).unwrap(), 5);
assert!(t.flush().is_ok());
assert_eq_transport_written_bytes!(t, b);
}
}