easyio/
lib.rs

1use std::io::{self, Read, Write};
2use std::sync::atomic::{AtomicUsize, Ordering};
3use std::sync::{Arc};
4
5pub mod conv;
6
7/// read_full reads from r until buffer is full, EOF is met, or an io:Error occured.
8///
9/// On ok return:
10/// If return size == buffer.len() the read is successful and there may be more data available from r.
11/// If return size < buffer.len(), EOF is met before buffer is filled.
12pub fn read_full(buffer: &mut [u8], r: &mut dyn Read) -> Result<usize, io::Error> {
13    let mut len_read: usize = 0;
14    loop {
15        match r.read(&mut buffer[len_read..]) {
16            Ok(size) => {
17                len_read += size;
18                if size == 0 || buffer.len() == len_read {
19                    return Ok(len_read);
20                }
21            }
22            Err(e) => return Err(e),
23        }
24    }
25}
26
27/// BlackHole implements io::Write trait.
28///
29/// Writes to BlackHole always succeeds.
30pub struct BlackHole {}
31impl Write for BlackHole {
32    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
33        Ok(buf.len())
34    }
35    fn flush(&mut self) -> io::Result<()> {
36        Ok(())
37    }
38}
39
40struct MeteringReaderHandle<'a> {
41    underlying_reader: &'a mut dyn Read,
42    counter: Arc<AtomicUsize>,
43}
44
45/// MeteringReader wraps around a reader and atomically accumulates the total count of bytes written to it.
46///
47/// Use as_reader() to obtain a io::Reader handle to it.
48pub struct MeteringReader<'a> {
49    inner: MeteringReaderHandle<'a>,
50    counter: Arc<AtomicUsize>,
51}
52
53impl MeteringReader<'_> {
54    pub fn new(r: &mut dyn Read) -> MeteringReader {
55        let counter = Arc::new(AtomicUsize::new(0));
56        MeteringReader{
57            inner : MeteringReaderHandle{
58                underlying_reader: r,
59                counter: Arc::clone(&counter),
60            },
61            counter: Arc::clone(&counter),
62        }
63    }
64
65    pub fn as_reader(&mut self) -> &mut dyn Read {
66        &mut self.inner
67    }
68
69    pub fn get_counter(&self) -> usize {
70        self.counter.load(Ordering::Relaxed)
71    }
72}
73
74impl Read for MeteringReaderHandle<'_> {
75    fn read(&mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
76        match self.underlying_reader.read(buf) {
77            Ok(size) => {
78                self.counter.fetch_add(size, Ordering::Relaxed);
79                Ok(size)
80            },
81            Err(e) => {
82                Err(e)
83            }
84        }
85    }
86}
87
88#[cfg(test)]
89mod tests {
90    use crate as lib;
91    use std::io::{self, Read};
92
93    pub struct SlowReader<'a> {
94        underlying_reader: &'a mut dyn (Read),
95    }
96
97    impl Read for SlowReader<'_> {
98        fn read(& mut self, buf: &mut [u8]) -> Result<usize, io::Error> {
99            // return two byte at a tim e
100            if buf.len() >= 2 {
101                self.underlying_reader.read(&mut buf[..2])
102            } else {
103                self.underlying_reader.read(&mut buf[..])
104            }
105        }
106    }
107
108    mod test_read_full {
109
110        #[test]
111        fn test_read_full_slow() {
112            let mut underlying_data: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7];
113            let mut reader = super::SlowReader {
114                underlying_reader: &mut underlying_data,
115            };
116            let mut buf = vec![0u8; 4];
117            let res = super::lib::read_full(&mut buf[..4], &mut reader);
118            assert_eq!(res.unwrap(), 4usize);
119            assert_eq!(buf[..4], [0, 1, 2, 3]);
120            let res = super::lib::read_full(&mut buf[..3], &mut reader);
121            assert_eq!(res.unwrap(), 3usize);
122            assert_eq!(buf[..3], [4, 5, 6]);
123            let res = super::lib::read_full(&mut buf[..3], &mut reader);
124            assert_eq!(res.unwrap(), 1usize);
125            assert_eq!(buf[0], 7);
126        }
127
128        #[test]
129        fn test_read_full_once() {
130            let mut underlying_data: &[u8] = &[0, 1, 2, 3, 4, 5, 6, 7];
131            let mut buf = vec![0u8; 9];
132            let res = super::lib::read_full(&mut buf[..], &mut underlying_data);
133            assert_eq!(res.unwrap(), 8usize);
134        }
135    }
136
137    mod test_metering_reader {
138        use std::io;
139        use crate::{MeteringReader, BlackHole};
140        use std::sync::{Arc};
141
142        #[test]
143        fn test_metering_reader_update() {
144            let mut input = "123456".as_bytes();
145            let input_len = input.len();
146            let mut meter = MeteringReader::new(&mut input);
147            let mut meter_reader = meter.as_reader();
148            io::copy(&mut meter_reader, &mut BlackHole{}).unwrap();
149            let result = meter.get_counter();
150            assert_eq!(input_len, result);
151        }
152
153        #[test]
154        fn test_metering_drop_counter_when_meter_is_dropped() {
155            let counter_ref;
156            {
157                let mut input = "123456".as_bytes();
158                let input_len = input.len();
159                let mut meter = MeteringReader::new(&mut input);
160                counter_ref = Arc::downgrade(&meter.counter);
161                let mut meter_reader = meter.as_reader();
162                io::copy(&mut meter_reader, &mut BlackHole{}).unwrap();
163                let result = meter.get_counter();
164                assert_eq!(input_len, result);
165            }
166            assert_eq!(counter_ref.upgrade().is_none(), true);
167        }
168    }
169}
easyio/lib.rs

easyio/
lib.rs
