Coverage Report

Created: 2025-02-25 06:39

/rust/registry/src/index.crates.io-6f17d22bba15001f/pprof-0.14.0/src/collector.rs
Line
Count
Source (jump to first uncovered line)
1
// Copyright 2019 TiKV Project Authors. Licensed under Apache-2.0.
2
3
use std::collections::hash_map::DefaultHasher;
4
use std::convert::TryInto;
5
use std::fmt::Debug;
6
use std::hash::{Hash, Hasher};
7
use std::io::{Read, Seek, SeekFrom, Write};
8
9
use crate::frames::UnresolvedFrames;
10
11
use aligned_vec::AVec;
12
use tempfile::NamedTempFile;
13
14
pub const BUCKETS: usize = 1 << 12;
15
pub const BUCKETS_ASSOCIATIVITY: usize = 4;
16
pub const BUFFER_LENGTH: usize = (1 << 18) / std::mem::size_of::<Entry<UnresolvedFrames>>();
17
18
#[derive(Debug)]
19
pub struct Entry<T> {
20
    pub item: T,
21
    pub count: isize,
22
}
23
24
impl<T: Default> Default for Entry<T> {
25
0
    fn default() -> Self {
26
0
        Entry {
27
0
            item: Default::default(),
28
0
            count: 0,
29
0
        }
30
0
    }
31
}
32
33
#[derive(Debug)]
34
pub struct Bucket<T: 'static> {
35
    pub length: usize,
36
    entries: Box<[Entry<T>; BUCKETS_ASSOCIATIVITY]>,
37
}
38
39
impl<T: Eq + Default> Default for Bucket<T> {
40
0
    fn default() -> Bucket<T> {
41
0
        let entries = Box::default();
42
0
43
0
        Self { length: 0, entries }
44
0
    }
45
}
46
47
impl<T: Eq> Bucket<T> {
48
0
    pub fn add(&mut self, key: T, count: isize) -> Option<Entry<T>> {
49
0
        let mut done = false;
50
0
        self.entries[0..self.length].iter_mut().for_each(|ele| {
51
0
            if ele.item == key {
52
0
                ele.count += count;
53
0
                done = true;
54
0
            }
55
0
        });
56
0
57
0
        if done {
58
0
            None
59
0
        } else if self.length < BUCKETS_ASSOCIATIVITY {
60
0
            let ele = &mut self.entries[self.length];
61
0
            ele.item = key;
62
0
            ele.count = count;
63
0
64
0
            self.length += 1;
65
0
            None
66
        } else {
67
0
            let mut min_index = 0;
68
0
            let mut min_count = self.entries[0].count;
69
0
            for index in 0..self.length {
70
0
                let count = self.entries[index].count;
71
0
                if count < min_count {
72
0
                    min_index = index;
73
0
                    min_count = count;
74
0
                }
75
            }
76
77
0
            let mut new_entry = Entry { item: key, count };
78
0
            std::mem::swap(&mut self.entries[min_index], &mut new_entry);
79
0
            Some(new_entry)
80
        }
81
0
    }
82
83
0
    pub fn iter(&self) -> BucketIterator<T> {
84
0
        BucketIterator::<T> {
85
0
            related_bucket: self,
86
0
            index: 0,
87
0
        }
88
0
    }
89
}
90
91
pub struct BucketIterator<'a, T: 'static> {
92
    related_bucket: &'a Bucket<T>,
93
    index: usize,
94
}
95
96
impl<'a, T> Iterator for BucketIterator<'a, T> {
97
    type Item = &'a Entry<T>;
98
99
0
    fn next(&mut self) -> Option<Self::Item> {
100
0
        if self.index < self.related_bucket.length {
101
0
            self.index += 1;
102
0
            Some(&self.related_bucket.entries[self.index - 1])
103
        } else {
104
0
            None
105
        }
106
0
    }
107
}
108
109
pub struct HashCounter<T: Hash + Eq + 'static> {
110
    buckets: Box<[Bucket<T>; BUCKETS]>,
111
}
112
113
impl<T: Hash + Eq + Default + Debug> Default for HashCounter<T> {
114
0
    fn default() -> Self {
115
0
        let mut v: Vec<Bucket<T>> = Vec::with_capacity(BUCKETS);
116
0
        v.resize_with(BUCKETS, Default::default);
117
0
        let buckets = v.into_boxed_slice().try_into().unwrap();
118
0
119
0
        Self { buckets }
120
0
    }
121
}
122
123
impl<T: Hash + Eq> HashCounter<T> {
124
0
    fn hash(key: &T) -> u64 {
125
0
        let mut s = DefaultHasher::new();
126
0
        key.hash(&mut s);
127
0
        s.finish()
128
0
    }
129
130
0
    pub fn add(&mut self, key: T, count: isize) -> Option<Entry<T>> {
131
0
        let hash_value = Self::hash(&key);
132
0
        let bucket = &mut self.buckets[(hash_value % BUCKETS as u64) as usize];
133
0
134
0
        bucket.add(key, count)
135
0
    }
136
137
0
    pub fn iter(&self) -> impl Iterator<Item = &Entry<T>> {
138
0
        let mut iter: Box<dyn Iterator<Item = &Entry<T>>> =
139
0
            Box::new(self.buckets[0].iter().chain(std::iter::empty()));
140
0
        for bucket in self.buckets[1..].iter() {
141
0
            iter = Box::new(iter.chain(bucket.iter()));
142
0
        }
143
144
0
        iter
145
0
    }
146
}
147
148
pub struct TempFdArray<T: 'static> {
149
    file: NamedTempFile,
150
    buffer: Box<[T; BUFFER_LENGTH]>,
151
    buffer_index: usize,
152
    flush_n: usize,
153
}
154
155
impl<T: Default + Debug> TempFdArray<T> {
156
0
    fn new() -> std::io::Result<TempFdArray<T>> {
157
0
        let file = NamedTempFile::new()?;
158
159
0
        let mut v: Vec<T> = Vec::with_capacity(BUFFER_LENGTH);
160
0
        v.resize_with(BUFFER_LENGTH, Default::default);
161
0
        let buffer = v.into_boxed_slice().try_into().unwrap();
162
0
163
0
        Ok(Self {
164
0
            file,
165
0
            buffer,
166
0
            buffer_index: 0,
167
0
            flush_n: 0,
168
0
        })
169
0
    }
170
}
171
172
impl<T> TempFdArray<T> {
173
0
    fn flush_buffer(&mut self) -> std::io::Result<()> {
174
0
        self.buffer_index = 0;
175
0
        let buf = unsafe {
176
0
            std::slice::from_raw_parts(
177
0
                self.buffer.as_ptr() as *const u8,
178
0
                BUFFER_LENGTH * std::mem::size_of::<T>(),
179
0
            )
180
0
        };
181
0
        self.flush_n += 1;
182
0
        self.file.write_all(buf)?;
183
184
0
        Ok(())
185
0
    }
186
187
0
    fn push(&mut self, entry: T) -> std::io::Result<()> {
188
0
        if self.buffer_index >= BUFFER_LENGTH {
189
0
            self.flush_buffer()?;
190
0
        }
191
192
0
        self.buffer[self.buffer_index] = entry;
193
0
        self.buffer_index += 1;
194
0
195
0
        Ok(())
196
0
    }
197
198
0
    fn try_iter(&self) -> std::io::Result<impl Iterator<Item = &T>> {
199
0
        let size = BUFFER_LENGTH * self.flush_n * std::mem::size_of::<T>();
200
0
201
0
        let mut file_vec = AVec::with_capacity(std::mem::align_of::<T>(), size);
202
0
        let mut file = self.file.reopen()?;
203
204
0
        unsafe {
205
0
            // it's safe as the capacity is initialized to `size`, and it'll be filled with `size` bytes
206
0
            file_vec.set_len(size);
207
0
        }
208
0
        file.read_exact(&mut file_vec[0..size])?;
209
0
        file.seek(SeekFrom::End(0))?;
210
211
0
        Ok(TempFdArrayIterator {
212
0
            buffer: &self.buffer[0..self.buffer_index],
213
0
            file_vec,
214
0
            index: 0,
215
0
        })
216
0
    }
217
}
218
219
pub struct TempFdArrayIterator<'a, T> {
220
    pub buffer: &'a [T],
221
    pub file_vec: AVec<u8>,
222
    pub index: usize,
223
}
224
225
impl<'a, T> Iterator for TempFdArrayIterator<'a, T> {
226
    type Item = &'a T;
227
228
0
    fn next(&mut self) -> Option<Self::Item> {
229
0
        if self.index < self.buffer.len() {
230
0
            self.index += 1;
231
0
            Some(&self.buffer[self.index - 1])
232
        } else {
233
0
            let length = self.file_vec.len() / std::mem::size_of::<T>();
234
0
            let ts =
235
0
                unsafe { std::slice::from_raw_parts(self.file_vec.as_ptr() as *const T, length) };
236
0
            if self.index - self.buffer.len() < ts.len() {
237
0
                self.index += 1;
238
0
                Some(&ts[self.index - self.buffer.len() - 1])
239
            } else {
240
0
                None
241
            }
242
        }
243
0
    }
244
}
245
246
pub struct Collector<T: Hash + Eq + 'static> {
247
    map: HashCounter<T>,
248
    temp_array: TempFdArray<Entry<T>>,
249
}
250
251
impl<T: Hash + Eq + Default + Debug + 'static> Collector<T> {
252
0
    pub fn new() -> std::io::Result<Self> {
253
0
        Ok(Self {
254
0
            map: HashCounter::<T>::default(),
255
0
            temp_array: TempFdArray::<Entry<T>>::new()?,
256
        })
257
0
    }
258
}
259
260
impl<T: Hash + Eq + 'static> Collector<T> {
261
0
    pub fn add(&mut self, key: T, count: isize) -> std::io::Result<()> {
262
0
        if let Some(evict) = self.map.add(key, count) {
263
0
            self.temp_array.push(evict)?;
264
0
        }
265
266
0
        Ok(())
267
0
    }
268
269
0
    pub fn try_iter(&self) -> std::io::Result<impl Iterator<Item = &Entry<T>>> {
270
0
        Ok(self.map.iter().chain(self.temp_array.try_iter()?))
271
0
    }
272
}
273
274
#[cfg(test)]
275
mod test_utils {
276
    use super::*;
277
    use std::collections::BTreeMap;
278
279
    pub fn add_map<T: std::cmp::Ord + Copy>(hashmap: &mut BTreeMap<T, isize>, entry: &Entry<T>) {
280
        match hashmap.get_mut(&entry.item) {
281
            None => {
282
                hashmap.insert(entry.item, entry.count);
283
            }
284
            Some(count) => *count += entry.count,
285
        }
286
    }
287
}
288
289
#[cfg(test)]
290
mod tests {
291
    use super::*;
292
    use std::collections::BTreeMap;
293
294
    #[test]
295
    fn stack_hash_counter() {
296
        let mut stack_hash_counter = HashCounter::<usize>::default();
297
        stack_hash_counter.add(0, 1);
298
        stack_hash_counter.add(1, 1);
299
        stack_hash_counter.add(1, 1);
300
301
        stack_hash_counter.iter().for_each(|item| {
302
            if item.item == 0 {
303
                assert_eq!(item.count, 1);
304
            } else if item.item == 1 {
305
                assert_eq!(item.count, 2);
306
            } else {
307
                unreachable!();
308
            }
309
        });
310
    }
311
312
    #[test]
313
    fn evict_test() {
314
        let mut stack_hash_counter = HashCounter::<usize>::default();
315
        let mut real_map = BTreeMap::new();
316
317
        for item in 0..(1 << 10) * 4 {
318
            for _ in 0..(item % 4) {
319
                match stack_hash_counter.add(item, 1) {
320
                    None => {}
321
                    Some(evict) => {
322
                        test_utils::add_map(&mut real_map, &evict);
323
                    }
324
                }
325
            }
326
        }
327
328
        stack_hash_counter.iter().for_each(|entry| {
329
            test_utils::add_map(&mut real_map, entry);
330
        });
331
332
        for item in 0..(1 << 10) * 4 {
333
            let count = (item % 4) as isize;
334
            match real_map.get(&item) {
335
                Some(item) => {
336
                    assert_eq!(*item, count);
337
                }
338
                None => {
339
                    assert_eq!(count, 0);
340
                }
341
            }
342
        }
343
    }
344
345
    #[test]
346
    fn collector_test() {
347
        let mut collector = Collector::new().unwrap();
348
        let mut real_map = BTreeMap::new();
349
350
        for item in 0..(1 << 12) * 4 {
351
            for _ in 0..(item % 4) {
352
                collector.add(item, 1).unwrap();
353
            }
354
        }
355
356
        collector.try_iter().unwrap().for_each(|entry| {
357
            test_utils::add_map(&mut real_map, entry);
358
        });
359
360
        for item in 0..(1 << 12) * 4 {
361
            let count = (item % 4) as isize;
362
            match real_map.get(&item) {
363
                Some(value) => {
364
                    assert_eq!(count, *value);
365
                }
366
                None => {
367
                    assert_eq!(count, 0);
368
                }
369
            }
370
        }
371
    }
372
373
    #[derive(Debug, Hash, Eq, PartialEq, PartialOrd, Ord, Default, Clone, Copy)]
374
    struct AlignTest {
375
        a: u16,
376
        b: u32,
377
        c: u64,
378
        d: u64,
379
    }
380
381
    // collector_align_test uses a bigger item to test the alignment of the collector
382
    #[test]
383
    fn collector_align_test() {
384
        let mut collector = Collector::new().unwrap();
385
        let mut real_map = BTreeMap::new();
386
387
        for item in 0..(1 << 12) * 4 {
388
            for _ in 0..(item % 4) {
389
                collector
390
                    .add(
391
                        AlignTest {
392
                            a: item as u16,
393
                            b: item as u32,
394
                            c: item as u64,
395
                            d: item as u64,
396
                        },
397
                        1,
398
                    )
399
                    .unwrap();
400
            }
401
        }
402
403
        collector.try_iter().unwrap().for_each(|entry| {
404
            test_utils::add_map(&mut real_map, entry);
405
        });
406
407
        for item in 0..(1 << 12) * 4 {
408
            let count = (item % 4) as isize;
409
            let align_item = AlignTest {
410
                a: item as u16,
411
                b: item as u32,
412
                c: item as u64,
413
                d: item as u64,
414
            };
415
            match real_map.get(&align_item) {
416
                Some(value) => {
417
                    assert_eq!(count, *value);
418
                }
419
                None => {
420
                    assert_eq!(count, 0);
421
                }
422
            }
423
        }
424
    }
425
}