/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 | | } |