msr_core/io/
mod.rs

1//! I/O related utilities
2
3use std::io::{Result, Write};
4
5use crate::sync::{
6    atomic::{AtomicU64, Ordering},
7    Arc,
8};
9
10#[derive(Clone, Debug)]
11pub struct BytesWritten(Arc<AtomicU64>);
12
13impl BytesWritten {
14    #[must_use]
15    pub fn value(&self) -> u64 {
16        self.0.load(Ordering::Relaxed)
17    }
18}
19
20#[derive(Debug)]
21pub struct CountingWrite<W: Write> {
22    writer: W,
23    bytes_written: Arc<AtomicU64>,
24}
25
26impl<W: Write> CountingWrite<W> {
27    /// Wrap a writer and start counting
28    pub fn from_writer(writer: W) -> (Self, BytesWritten) {
29        let bytes_written = Arc::new(AtomicU64::new(0));
30        (
31            Self {
32                writer,
33                bytes_written: Arc::clone(&bytes_written),
34            },
35            BytesWritten(bytes_written),
36        )
37    }
38
39    /// Dismantle the wrapped writer and stop counting
40    pub fn into_value(self) -> W {
41        self.writer
42    }
43}
44
45impl<W: Write> Write for CountingWrite<W> {
46    fn write(&mut self, buf: &[u8]) -> Result<usize> {
47        let bytes_written = self.writer.write(buf)?;
48        // self has exclusive mutable access on the number of octets written, i.e.
49        // we can safely get-modify-set this value without race conditions here!
50        let mut sum_bytes_written = self.bytes_written.load(Ordering::Relaxed);
51        sum_bytes_written = sum_bytes_written.saturating_add(bytes_written as u64);
52        self.bytes_written
53            .store(sum_bytes_written, Ordering::Relaxed);
54        Ok(bytes_written)
55    }
56
57    fn flush(&mut self) -> Result<()> {
58        self.writer.flush()
59    }
60}
61
62#[cfg(test)]
63mod tests;
msr_core/io/mod.rs

msr_core/io/
mod.rs
