1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89
use std::{self, cmp, thread, time}; use {Result, Read, Write}; pub trait MockRead { #[inline] fn limited(self, limit: usize) -> Limited<Self> where Self: Sized { Limited { limit: limit, inner: self } } #[inline] fn throttled(self, rate: usize) -> Throttled<Self> where Self: Sized { Throttled { rate: rate, inner: self } } } impl<R: Read> MockRead for R {} pub trait MockWrite { #[inline] fn limited(self, limit: usize) -> Limited<Self> where Self: Sized { Limited { limit: limit, inner: self } } } impl<W: Write> MockWrite for W {} pub struct Limited<T: ?Sized> { limit: usize, inner: T, } impl<T> Limited<T> { #[inline] pub fn limit(&self) -> usize { self.limit } } impl<T: Read> Read for Limited<T> { fn read(&mut self, buf: &mut [u8]) -> Result<usize> { let max = cmp::min(buf.len(), self.limit); self.inner.read(&mut buf[..max]) } } impl<T: Read> std::io::Read for Limited<T> { fn read(&mut self, buf: &mut [u8]) -> Result<usize> { Read::read(self, buf) } } impl<T: Write> Write for Limited<T> { fn write(&mut self, buf: &[u8]) -> Result<usize> { let amt = cmp::min(self.limit, buf.len()); self.inner.write(&buf[..amt]) } #[inline] fn flush(&mut self) -> Result<()> { self.inner.flush() } } impl<T: Write> std::io::Write for Limited<T> { fn write(&mut self, buf: &[u8]) -> Result<usize> { Write::write(self, buf) } #[inline] fn flush(&mut self) -> Result<()> { Write::flush(self) } } pub struct Throttled<T: ?Sized> { pub rate: usize, pub inner: T, } impl<R: Read> Read for Throttled<R> { #[inline] fn read(&mut self, buf: &mut [u8]) -> Result<usize> { let amt = try!(self.inner.read(buf)); let nanos = 1000000000u64 * amt as u64 / self.rate as u64; thread::sleep(time::Duration::new(nanos / 1000000000u64, (nanos % 1000000000u64) as u32)); Ok(amt) } } impl<R: Read> std::io::Read for Throttled<R> { #[inline] fn read(&mut self, buf: &mut [u8]) -> Result<usize> { Read::read(self, buf) } }