cubecl_environment/bytes/
file_reads.rs

1//! The bytes file-backed [`Bytes`](super::Bytes) read from their files, for a
2//! caller that shows how far a lazy load has got.
3
4use std::{
5    path::{Path, PathBuf},
6    sync::{
7        Arc, Mutex, Weak,
8        atomic::{AtomicU64, AtomicUsize, Ordering},
9    },
10    vec::Vec,
11};
12
13/// Every counter alive, which each read of a file is counted on.
14static REGISTRY: CounterRegistry = CounterRegistry {
15    counters: Mutex::new(Vec::new()),
16    alive: AtomicUsize::new(0),
17};
18
19/// The bytes read from one file through file-backed
20/// [`Bytes`](super::Bytes) while this counter is alive.
21///
22/// A file-backed allocation reads its file only when it is first accessed —
23/// most often when a runtime copies it to a device, on the runtime's own
24/// thread — so whoever made the allocation does not see the read happen. A
25/// counter is how they learn how far it has got: create one for the file
26/// before its bytes are read, then read [`read`](Self::read) from any thread.
27///
28/// The file is matched by the path [`Bytes::from_file`](super::Bytes::from_file)
29/// was given, compared as given and never canonicalized. Two counters of one
30/// file each count every read of it.
31pub struct FileReadCounter {
32    counted: Arc<Counted>,
33}
34
35impl FileReadCounter {
36    /// Start counting the reads of `path`, from zero.
37    pub fn new(path: impl Into<PathBuf>) -> Self {
38        let counted = Arc::new(Counted {
39            path: path.into(),
40            read: AtomicU64::new(0),
41        });
42        REGISTRY.register(&counted);
43        Self { counted }
44    }
45
46    /// The bytes read from the file since this counter was created.
47    pub fn read(&self) -> u64 {
48        self.counted.read.load(Ordering::Relaxed)
49    }
50}
51
52impl Drop for FileReadCounter {
53    fn drop(&mut self) {
54        REGISTRY.unregister();
55    }
56}
57
58/// Count `size` bytes read from `path` on every counter of it.
59pub(crate) fn count(path: &Path, size: u64) {
60    REGISTRY.count(path, size);
61}
62
63/// One file's count, shared by its counter and the registry.
64struct Counted {
65    path: PathBuf,
66    read: AtomicU64,
67}
68
69/// The counters alive, and how many there are: a read with none alive
70/// counts nothing, and skips the lock every other read of a file would take.
71///
72/// The orderings are `Relaxed` throughout: the lock orders the counters'
73/// list, and a count is read on its own. A read racing a counter's creation
74/// may go uncounted, which a caller avoids by creating the counter before
75/// the reads it counts.
76struct CounterRegistry {
77    counters: Mutex<Vec<Weak<Counted>>>,
78    alive: AtomicUsize,
79}
80
81impl CounterRegistry {
82    fn register(&self, counted: &Arc<Counted>) {
83        let mut counters = self.lock();
84        counters.retain(|counter| counter.strong_count() > 0);
85        counters.push(Arc::downgrade(counted));
86        self.alive.fetch_add(1, Ordering::Relaxed);
87    }
88
89    fn unregister(&self) {
90        self.alive.fetch_sub(1, Ordering::Relaxed);
91    }
92
93    fn count(&self, path: &Path, size: u64) {
94        if self.alive.load(Ordering::Relaxed) == 0 {
95            return;
96        }
97        self.lock()
98            .iter()
99            .filter_map(Weak::upgrade)
100            .filter(|counted| counted.path == path)
101            .for_each(|counted| {
102                counted.read.fetch_add(size, Ordering::Relaxed);
103            });
104    }
105
106    fn lock(&self) -> std::sync::MutexGuard<'_, Vec<Weak<Counted>>> {
107        self.counters
108            .lock()
109            .unwrap_or_else(|poison| poison.into_inner())
110    }
111}
112
113#[cfg(test)]
114#[cfg(not(miri))]
115mod tests {
116    use super::super::{Bytes, Reader};
117    use super::*;
118    use std::io::Write;
119
120    /// A read is counted once, when the file is first accessed, whether it
121    /// materializes the bytes or copies them out, and a file nobody counts
122    /// is read as before.
123    #[test]
124    fn reads_are_counted_when_the_file_is_accessed() {
125        let dir = tempfile::tempdir().unwrap();
126        let path = dir.path().join("weights");
127        std::fs::File::create(&path)
128            .unwrap()
129            .write_all(&[7u8; 64])
130            .unwrap();
131
132        let counter = FileReadCounter::new(&path);
133        let bytes = Bytes::from_file(&path, 64, 0);
134        let view = bytes.view(0, 16).unwrap();
135        assert_eq!(counter.read(), 0);
136
137        let mut copied = Bytes::from_bytes_vec(std::vec![0u8; 16]);
138        view.copy_into(&mut copied);
139        assert_eq!(counter.read(), 16);
140
141        bytes.read(Reader::new()).unwrap();
142        bytes.read(Reader::new()).unwrap();
143        assert_eq!(counter.read(), 16 + 64);
144
145        let other = dir.path().join("other");
146        std::fs::File::create(&other)
147            .unwrap()
148            .write_all(&[1])
149            .unwrap();
150        Bytes::from_file(&other, 1, 0).read(Reader::new()).unwrap();
151        assert_eq!(counter.read(), 16 + 64);
152    }
153}
cubecl_environment/bytes/file_reads.rs

cubecl_environment/bytes/
file_reads.rs
