From 5f624c20fdd36f3396b78b954edbc7fc8b85773c Mon Sep 17 00:00:00 2001 From: Shanin Roman Date: Wed, 31 Jan 2024 14:48:26 +0300 Subject: [PATCH] fix: use T aligned pointer in TempFdArray Signed-off-by: Shanin Roman --- src/collector.rs | 85 +++++++++++++++++++++++++++++++++++------------- src/report.rs | 82 +++++++++++++++++++++++++--------------------- 2 files changed, 108 insertions(+), 59 deletions(-) diff --git a/src/collector.rs b/src/collector.rs index c05fb18b..0225465d 100644 --- a/src/collector.rs +++ b/src/collector.rs @@ -5,6 +5,7 @@ use std::convert::TryInto; use std::fmt::Debug; use std::hash::{Hash, Hasher}; use std::io::{Read, Seek, SeekFrom, Write}; +use std::mem::ManuallyDrop; use crate::frames::UnresolvedFrames; @@ -148,6 +149,7 @@ pub struct TempFdArray { file: NamedTempFile, buffer: Box<[T; BUFFER_LENGTH]>, buffer_index: usize, + flush_n: usize, } impl TempFdArray { @@ -162,6 +164,7 @@ impl TempFdArray { file, buffer, buffer_index: 0, + flush_n: 0, }) } } @@ -175,6 +178,7 @@ impl TempFdArray { BUFFER_LENGTH * std::mem::size_of::(), ) }; + self.flush_n += 1; self.file.write_all(buf)?; Ok(()) @@ -191,24 +195,56 @@ impl TempFdArray { Ok(()) } - fn try_iter(&self) -> std::io::Result> { - let mut file_vec = Vec::new(); - let mut file = self.file.reopen()?; - file.seek(SeekFrom::Start(0))?; - file.read_to_end(&mut file_vec)?; - file.seek(SeekFrom::End(0))?; + fn try_iter<'lt>( + &'lt self, + file_buffer_container: &'lt mut Option]>>, + ) -> std::io::Result> { + let file_buffer = self.file_buffer()?; + let file_buffer = file_buffer_container.insert(file_buffer); Ok(TempFdArrayIterator { buffer: &self.buffer[0..self.buffer_index], - file_vec, + file_buffer, index: 0, }) } + + fn file_buffer(&self) -> std::io::Result]>> { + if self.flush_n == 0 { + return Ok(Vec::new().into_boxed_slice()); + } + + let mut file = self.file.reopen()?; + file.seek(SeekFrom::Start(0))?; + let file_buffer = unsafe { + // Get properly aligned pointer + let len = BUFFER_LENGTH * self.flush_n; + // Expect T to be non-ZST + let layout = std::alloc::Layout::array::>(len).unwrap(); + let ptr = std::alloc::alloc(layout); + if ptr.is_null() { + std::alloc::handle_alloc_error(layout); + } + // Populate with bytes + file.read_exact(std::slice::from_raw_parts_mut( + ptr, + len * std::mem::size_of::(), + ))?; + // Cast to proper type + Box::from_raw(std::ptr::slice_from_raw_parts_mut( + ptr.cast::>(), + len, + )) + }; + file.seek(SeekFrom::End(0))?; + + Ok(file_buffer) + } } pub struct TempFdArrayIterator<'a, T> { pub buffer: &'a [T], - pub file_vec: Vec, + pub file_buffer: &'a [ManuallyDrop], pub index: usize, } @@ -219,16 +255,11 @@ impl<'a, T> Iterator for TempFdArrayIterator<'a, T> { if self.index < self.buffer.len() { self.index += 1; Some(&self.buffer[self.index - 1]) + } else if self.index - self.buffer.len() < self.file_buffer.len() { + self.index += 1; + Some(&self.file_buffer[self.index - self.buffer.len() - 1]) } else { - let length = self.file_vec.len() / std::mem::size_of::(); - let ts = - unsafe { std::slice::from_raw_parts(self.file_vec.as_ptr() as *const T, length) }; - if self.index - self.buffer.len() < ts.len() { - self.index += 1; - Some(&ts[self.index - self.buffer.len() - 1]) - } else { - None - } + None } } } @@ -256,8 +287,14 @@ impl Collector { Ok(()) } - pub fn try_iter(&self) -> std::io::Result>> { - Ok(self.map.iter().chain(self.temp_array.try_iter()?)) + pub fn try_iter<'lt>( + &'lt self, + file_buffer_store: &'lt mut Option>]>>, + ) -> std::io::Result>> { + Ok(self + .map + .iter() + .chain(self.temp_array.try_iter(file_buffer_store)?)) } } @@ -343,9 +380,13 @@ mod tests { } } - collector.try_iter().unwrap().for_each(|entry| { - test_utils::add_map(&mut real_map, entry); - }); + let mut file_buffer_store = None; + collector + .try_iter(&mut file_buffer_store) + .unwrap() + .for_each(|entry| { + test_utils::add_map(&mut real_map, entry); + }); for item in 0..(1 << 12) * 4 { let count = (item % 4) as isize; diff --git a/src/report.rs b/src/report.rs index 971cb8d6..3070ade7 100644 --- a/src/report.rs +++ b/src/report.rs @@ -69,25 +69,29 @@ impl<'a> ReportBuilder<'a> { Err(Error::CreatingError) } Ok(profiler) => { - profiler.data.try_iter()?.for_each(|entry| { - let count = entry.count; - if count > 0 { - let key = &entry.item; - match hash_map.get_mut(key) { - Some(value) => { - *value += count; - } - None => { - match hash_map.insert(key.clone(), count) { - None => {} - Some(_) => { - unreachable!(); - } - }; + let mut file_buffer_store = None; + profiler + .data + .try_iter(&mut file_buffer_store)? + .for_each(|entry| { + let count = entry.count; + if count > 0 { + let key = &entry.item; + match hash_map.get_mut(key) { + Some(value) => { + *value += count; + } + None => { + match hash_map.insert(key.clone(), count) { + None => {} + Some(_) => { + unreachable!(); + } + }; + } } } - } - }); + }); Ok(UnresolvedReport { data: hash_map, @@ -107,29 +111,33 @@ impl<'a> ReportBuilder<'a> { Err(Error::CreatingError) } Ok(profiler) => { - profiler.data.try_iter()?.for_each(|entry| { - let count = entry.count; - if count > 0 { - let mut key = Frames::from(entry.item.clone()); - if let Some(processor) = &self.frames_post_processor { - processor(&mut key); - } - - match hash_map.get_mut(&key) { - Some(value) => { - *value += count; + let mut file_buffer_store = None; + profiler + .data + .try_iter(&mut file_buffer_store)? + .for_each(|entry| { + let count = entry.count; + if count > 0 { + let mut key = Frames::from(entry.item.clone()); + if let Some(processor) = &self.frames_post_processor { + processor(&mut key); } - None => { - match hash_map.insert(key, count) { - None => {} - Some(_) => { - unreachable!(); - } - }; + + match hash_map.get_mut(&key) { + Some(value) => { + *value += count; + } + None => { + match hash_map.insert(key, count) { + None => {} + Some(_) => { + unreachable!(); + } + }; + } } } - } - }); + }); Ok(Report { data: hash_map,