vortex_file/
counting.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::io;
5use std::sync::Arc;
6use std::sync::atomic::AtomicU64;
7use std::sync::atomic::Ordering;
8
9use vortex_io::IoBuf;
10use vortex_io::VortexWrite;
11
12/// A wrapper around a [`VortexWrite`] that counts the number of bytes written.
13pub struct CountingVortexWrite<W> {
14    inner: W,
15    bytes_written: Arc<AtomicU64>,
16}
17
18impl<W: VortexWrite> CountingVortexWrite<W> {
19    /// Wrap a writer with a new byte counter.
20    pub fn new(inner: W) -> Self {
21        Self {
22            inner,
23            bytes_written: Default::default(),
24        }
25    }
26
27    /// Returns the shared byte counter updated by this writer.
28    pub fn counter(&self) -> Arc<AtomicU64> {
29        Arc::clone(&self.bytes_written)
30    }
31}
32
33impl<W: VortexWrite + Unpin> VortexWrite for CountingVortexWrite<W> {
34    async fn write_all<B: IoBuf>(&mut self, buffer: B) -> io::Result<B> {
35        let buf_len = buffer.as_slice().len() as u64;
36        let result = self.inner.write_all(buffer).await;
37        if result.is_ok() {
38            self.bytes_written.fetch_add(buf_len, Ordering::Relaxed);
39        }
40        result
41    }
42
43    fn flush(&mut self) -> impl Future<Output = io::Result<()>> {
44        self.inner.flush()
45    }
46
47    fn shutdown(&mut self) -> impl Future<Output = io::Result<()>> {
48        self.inner.shutdown()
49    }
50}
vortex_file/counting.rs

vortex_file/
counting.rs
