use std::io::Read;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
pub struct ReadCounting<T: Read> {
reader: T,
count: u64,
}
impl<T: Read> ReadCounting<T> {
pub fn new(reader: T) -> Self {
ReadCounting { reader, count: 0 }
}
pub fn get_read_count(&self) -> u64 {
self.count
}
}
impl<T: Read> Read for ReadCounting<T> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let read = self.reader.read(buf)?;
self.count += read as u64;
Ok(read)
}
}
pub struct SharedReadCounting<T: Read> {
reader: T,
counter: Arc<AtomicU64>,
}
impl<T: Read> SharedReadCounting<T> {
pub fn new(reader: T) -> Self {
SharedReadCounting {
reader,
counter: Arc::new(AtomicU64::new(0)),
}
}
pub fn get_counter(&self) -> SharedROCounter {
SharedROCounter {
counter: self.counter.clone(),
}
}
}
impl<T: Read> Read for SharedReadCounting<T> {
fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
let read = self.reader.read(buf)?;
self.counter.fetch_add(read as u64, Ordering::Relaxed);
Ok(read)
}
}
#[derive(Debug, Clone)]
pub struct SharedROCounter {
counter: Arc<AtomicU64>,
}
impl SharedROCounter {
pub fn new(counter: Arc<AtomicU64>) -> Self {
SharedROCounter { counter }
}
pub fn get_count(&self) -> u64 {
self.counter.load(Ordering::Relaxed)
}
}