Skip to content

Commit 5f624c2

Browse files
ErigaraErigara
committed
fix: use T aligned pointer in TempFdArray
Signed-off-by: Shanin Roman <[email protected]>
1 parent 4939f73 commit 5f624c2

File tree

2 files changed

+108
-59
lines changed

2 files changed

+108
-59
lines changed

‎src/collector.rs

+63-22
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@ use std::convert::TryInto;
55
use std::fmt::Debug;
66
use std::hash::{Hash, Hasher};
77
use std::io::{Read, Seek, SeekFrom, Write};
8+
use std::mem::ManuallyDrop;
89

910
use crate::frames::UnresolvedFrames;
1011

@@ -148,6 +149,7 @@ pub struct TempFdArray<T: 'static> {
148149
file: NamedTempFile,
149150
buffer: Box<[T; BUFFER_LENGTH]>,
150151
buffer_index: usize,
152+
flush_n: usize,
151153
}
152154

153155
impl<T: Default + Debug> TempFdArray<T> {
@@ -162,6 +164,7 @@ impl<T: Default + Debug> TempFdArray<T> {
162164
file,
163165
buffer,
164166
buffer_index: 0,
167+
flush_n: 0,
165168
})
166169
}
167170
}
@@ -175,6 +178,7 @@ impl<T> TempFdArray<T> {
175178
BUFFER_LENGTH * std::mem::size_of::<T>(),
176179
)
177180
};
181+
self.flush_n += 1;
178182
self.file.write_all(buf)?;
179183

180184
Ok(())
@@ -191,24 +195,56 @@ impl<T> TempFdArray<T> {
191195
Ok(())
192196
}
193197

194-
fn try_iter(&self) -> std::io::Result<impl Iterator<Item = &T>> {
195-
let mut file_vec = Vec::new();
196-
let mut file = self.file.reopen()?;
197-
file.seek(SeekFrom::Start(0))?;
198-
file.read_to_end(&mut file_vec)?;
199-
file.seek(SeekFrom::End(0))?;
198+
fn try_iter<'lt>(
199+
&'lt self,
200+
file_buffer_container: &'lt mut Option<Box<[ManuallyDrop<T>]>>,
201+
) -> std::io::Result<impl Iterator<Item = &'lt T>> {
202+
let file_buffer = self.file_buffer()?;
203+
let file_buffer = file_buffer_container.insert(file_buffer);
200204

201205
Ok(TempFdArrayIterator {
202206
buffer: &self.buffer[0..self.buffer_index],
203-
file_vec,
207+
file_buffer,
204208
index: 0,
205209
})
206210
}
211+
212+
fn file_buffer(&self) -> std::io::Result<Box<[ManuallyDrop<T>]>> {
213+
if self.flush_n == 0 {
214+
return Ok(Vec::new().into_boxed_slice());
215+
}
216+
217+
let mut file = self.file.reopen()?;
218+
file.seek(SeekFrom::Start(0))?;
219+
let file_buffer = unsafe {
220+
// Get properly aligned pointer
221+
let len = BUFFER_LENGTH * self.flush_n;
222+
// Expect T to be non-ZST
223+
let layout = std::alloc::Layout::array::<ManuallyDrop<T>>(len).unwrap();
224+
let ptr = std::alloc::alloc(layout);
225+
if ptr.is_null() {
226+
std::alloc::handle_alloc_error(layout);
227+
}
228+
// Populate with bytes
229+
file.read_exact(std::slice::from_raw_parts_mut(
230+
ptr,
231+
len * std::mem::size_of::<T>(),
232+
))?;
233+
// Cast to proper type
234+
Box::from_raw(std::ptr::slice_from_raw_parts_mut(
235+
ptr.cast::<ManuallyDrop<T>>(),
236+
len,
237+
))
238+
};
239+
file.seek(SeekFrom::End(0))?;
240+
241+
Ok(file_buffer)
242+
}
207243
}
208244

209245
pub struct TempFdArrayIterator<'a, T> {
210246
pub buffer: &'a [T],
211-
pub file_vec: Vec<u8>,
247+
pub file_buffer: &'a [ManuallyDrop<T>],
212248
pub index: usize,
213249
}
214250

@@ -219,16 +255,11 @@ impl<'a, T> Iterator for TempFdArrayIterator<'a, T> {
219255
if self.index < self.buffer.len() {
220256
self.index += 1;
221257
Some(&self.buffer[self.index - 1])
258+
} else if self.index - self.buffer.len() < self.file_buffer.len() {
259+
self.index += 1;
260+
Some(&self.file_buffer[self.index - self.buffer.len() - 1])
222261
} else {
223-
let length = self.file_vec.len() / std::mem::size_of::<T>();
224-
let ts =
225-
unsafe { std::slice::from_raw_parts(self.file_vec.as_ptr() as *const T, length) };
226-
if self.index - self.buffer.len() < ts.len() {
227-
self.index += 1;
228-
Some(&ts[self.index - self.buffer.len() - 1])
229-
} else {
230-
None
231-
}
262+
None
232263
}
233264
}
234265
}
@@ -256,8 +287,14 @@ impl<T: Hash + Eq + 'static> Collector<T> {
256287
Ok(())
257288
}
258289

259-
pub fn try_iter(&self) -> std::io::Result<impl Iterator<Item = &Entry<T>>> {
260-
Ok(self.map.iter().chain(self.temp_array.try_iter()?))
290+
pub fn try_iter<'lt>(
291+
&'lt self,
292+
file_buffer_store: &'lt mut Option<Box<[ManuallyDrop<Entry<T>>]>>,
293+
) -> std::io::Result<impl Iterator<Item = &'lt Entry<T>>> {
294+
Ok(self
295+
.map
296+
.iter()
297+
.chain(self.temp_array.try_iter(file_buffer_store)?))
261298
}
262299
}
263300

@@ -343,9 +380,13 @@ mod tests {
343380
}
344381
}
345382

346-
collector.try_iter().unwrap().for_each(|entry| {
347-
test_utils::add_map(&mut real_map, entry);
348-
});
383+
let mut file_buffer_store = None;
384+
collector
385+
.try_iter(&mut file_buffer_store)
386+
.unwrap()
387+
.for_each(|entry| {
388+
test_utils::add_map(&mut real_map, entry);
389+
});
349390

350391
for item in 0..(1 << 12) * 4 {
351392
let count = (item % 4) as isize;

‎src/report.rs

+45-37
Original file line numberDiff line numberDiff line change
@@ -69,25 +69,29 @@ impl<'a> ReportBuilder<'a> {
6969
Err(Error::CreatingError)
7070
}
7171
Ok(profiler) => {
72-
profiler.data.try_iter()?.for_each(|entry| {
73-
let count = entry.count;
74-
if count > 0 {
75-
let key = &entry.item;
76-
match hash_map.get_mut(key) {
77-
Some(value) => {
78-
*value += count;
79-
}
80-
None => {
81-
match hash_map.insert(key.clone(), count) {
82-
None => {}
83-
Some(_) => {
84-
unreachable!();
85-
}
86-
};
72+
let mut file_buffer_store = None;
73+
profiler
74+
.data
75+
.try_iter(&mut file_buffer_store)?
76+
.for_each(|entry| {
77+
let count = entry.count;
78+
if count > 0 {
79+
let key = &entry.item;
80+
match hash_map.get_mut(key) {
81+
Some(value) => {
82+
*value += count;
83+
}
84+
None => {
85+
match hash_map.insert(key.clone(), count) {
86+
None => {}
87+
Some(_) => {
88+
unreachable!();
89+
}
90+
};
91+
}
8792
}
8893
}
89-
}
90-
});
94+
});
9195

9296
Ok(UnresolvedReport {
9397
data: hash_map,
@@ -107,29 +111,33 @@ impl<'a> ReportBuilder<'a> {
107111
Err(Error::CreatingError)
108112
}
109113
Ok(profiler) => {
110-
profiler.data.try_iter()?.for_each(|entry| {
111-
let count = entry.count;
112-
if count > 0 {
113-
let mut key = Frames::from(entry.item.clone());
114-
if let Some(processor) = &self.frames_post_processor {
115-
processor(&mut key);
116-
}
117-
118-
match hash_map.get_mut(&key) {
119-
Some(value) => {
120-
*value += count;
114+
let mut file_buffer_store = None;
115+
profiler
116+
.data
117+
.try_iter(&mut file_buffer_store)?
118+
.for_each(|entry| {
119+
let count = entry.count;
120+
if count > 0 {
121+
let mut key = Frames::from(entry.item.clone());
122+
if let Some(processor) = &self.frames_post_processor {
123+
processor(&mut key);
121124
}
122-
None => {
123-
match hash_map.insert(key, count) {
124-
None => {}
125-
Some(_) => {
126-
unreachable!();
127-
}
128-
};
125+
126+
match hash_map.get_mut(&key) {
127+
Some(value) => {
128+
*value += count;
129+
}
130+
None => {
131+
match hash_map.insert(key, count) {
132+
None => {}
133+
Some(_) => {
134+
unreachable!();
135+
}
136+
};
137+
}
129138
}
130139
}
131-
}
132-
});
140+
});
133141

134142
Ok(Report {
135143
data: hash_map,

0 commit comments

Comments
 (0)