Merge pull request #88 from michaelwoerister/opt-file-sink

Optimize FileSerializationSink by using parking_lot::Mutex and avoiding heap allocations in write_atomic.

Author: wesleywiser

analyzeme/benches/serialization_bench.rs

@@ -8,17 +8,35 @@ use measureme::{FileSerializationSink, MmapSerializationSink};
 #[bench]
 fn bench_file_serialization_sink(bencher: &mut test::Bencher) {
     bencher.iter(|| {
-        testing_common::run_end_to_end_serialization_test::<FileSerializationSink>(
-            "file_serialization_sink_test",
+        testing_common::run_serialization_bench::<FileSerializationSink>(
+            "file_serialization_sink_test", 500_000, 1
         );
     });
 }
 
 #[bench]
 fn bench_mmap_serialization_sink(bencher: &mut test::Bencher) {
     bencher.iter(|| {
-        testing_common::run_end_to_end_serialization_test::<MmapSerializationSink>(
-            "mmap_serialization_sink_test",
+        testing_common::run_serialization_bench::<MmapSerializationSink>(
+            "mmap_serialization_sink_test", 500_000, 1
+        );
+    });
+}
+
+#[bench]
+fn bench_file_serialization_sink_8_threads(bencher: &mut test::Bencher) {
+    bencher.iter(|| {
+        testing_common::run_serialization_bench::<FileSerializationSink>(
+            "file_serialization_sink_test", 50_000, 8
+        );
+    });
+}
+
+#[bench]
+fn bench_mmap_serialization_sink_8_threads(bencher: &mut test::Bencher) {
+    bencher.iter(|| {
+        testing_common::run_serialization_bench::<MmapSerializationSink>(
+            "mmap_serialization_sink_test", 50_000, 8
         );
     });
 }

analyzeme/src/testing_common.rs

@@ -1,5 +1,5 @@
 use crate::timestamp::Timestamp;
-use crate::{Event, LightweightEvent, ProfilingData};
+use crate::{Event, ProfilingData};
 use measureme::{Profiler, SerializationSink, StringId};
 use rustc_hash::FxHashMap;
 use std::borrow::Cow;
@@ -19,12 +19,16 @@ fn mk_filestem(file_name_stem: &str) -> PathBuf {
 }
 
 // Generate some profiling data. This is the part that would run in rustc.
-fn generate_profiling_data<S: SerializationSink>(filestem: &Path) -> Vec<Event<'static>> {
+fn generate_profiling_data<S: SerializationSink>(
+    filestem: &Path,
+    num_stacks: usize,
+    num_threads: usize,
+) -> Vec<Event<'static>> {
     let profiler = Arc::new(Profiler::<S>::new(Path::new(filestem)).unwrap());
 
     let event_id_reserved = StringId::reserved(42);
 
-    let event_ids = &[
+    let event_ids = vec![
         (
             profiler.alloc_string("Generic"),
             profiler.alloc_string("SomeGenericActivity"),
@@ -39,20 +43,33 @@ fn generate_profiling_data<S: SerializationSink>(filestem: &Path) -> Vec<Event<'
     event_ids_as_str.insert(event_ids[1].0, "Query");
     event_ids_as_str.insert(event_ids[1].1, "SomeQuery");
 
-    let mut expected_events = Vec::new();
+    let threads: Vec<_> = (0.. num_threads).map(|thread_id| {
+        let event_ids = event_ids.clone();
+        let profiler = profiler.clone();
+        let event_ids_as_str = event_ids_as_str.clone();
 
-    for i in 0..10_000 {
-        // Allocate some invocation stacks
+        std::thread::spawn(move || {
+            let mut expected_events = Vec::new();
 
-        pseudo_invocation(
-            &profiler,
-            i,
-            4,
-            event_ids,
-            &event_ids_as_str,
-            &mut expected_events,
-        );
-    }
+            for i in 0..num_stacks {
+                // Allocate some invocation stacks
+
+                pseudo_invocation(
+                    &profiler,
+                    i,
+                    thread_id as u32,
+                    4,
+                    &event_ids[..],
+                    &event_ids_as_str,
+                    &mut expected_events,
+                );
+            }
+
+            expected_events
+        })
+    }).collect();
+
+    let expected_events: Vec<_> = threads.into_iter().flat_map(|t| t.join().unwrap()).collect();
 
     // An example of allocating the string contents of an event id that has
     // already been used
@@ -67,53 +84,83 @@ fn process_profiling_data(filestem: &Path, expected_events: &[Event<'static>]) {
     let profiling_data = ProfilingData::new(filestem).unwrap();
 
     check_profiling_data(
-        &mut profiling_data.iter(),
+        &mut profiling_data.iter().map(|e| e.to_event()),
         &mut expected_events.iter().cloned(),
         expected_events.len(),
     );
     check_profiling_data(
-        &mut profiling_data.iter().rev(),
+        &mut profiling_data.iter().rev().map(|e| e.to_event()),
         &mut expected_events.iter().rev().cloned(),
         expected_events.len(),
     );
 }
 
 fn check_profiling_data(
-    actual_lightweight_events: &mut dyn Iterator<Item = LightweightEvent<'_>>,
+    actual_events: &mut dyn Iterator<Item = Event<'_>>,
     expected_events: &mut dyn Iterator<Item = Event<'_>>,
     num_expected_events: usize,
 ) {
     let mut count = 0;
 
+    // This assertion makes sure that the ExactSizeIterator impl works as expected.
     assert_eq!(
         (num_expected_events, Some(num_expected_events)),
-        actual_lightweight_events.size_hint()
+        actual_events.size_hint()
     );
 
-    for (actual_lightweight_event, expected_event) in actual_lightweight_events.zip(expected_events) {
-        let actual_event = actual_lightweight_event.to_event();
-        assert_eq!(actual_event.event_kind, expected_event.event_kind);
-        assert_eq!(actual_event.label, expected_event.label);
-        assert_eq!(actual_event.additional_data, expected_event.additional_data);
-        assert_eq!(
-            actual_event.timestamp.is_instant(),
-            expected_event.timestamp.is_instant()
-        );
-
-        count += 1;
+    let actual_events_per_thread = collect_events_per_thread(actual_events);
+    let expected_events_per_thread = collect_events_per_thread(expected_events);
+
+    let thread_ids: Vec<_> = actual_events_per_thread.keys().collect();
+    assert_eq!(thread_ids, expected_events_per_thread.keys().collect::<Vec<_>>());
+
+    for thread_id in thread_ids {
+        let actual_events = &actual_events_per_thread[thread_id];
+        let expected_events = &expected_events_per_thread[thread_id];
+
+        assert_eq!(actual_events.len(), expected_events.len());
+
+        for (actual_event, expected_event) in actual_events.iter().zip(expected_events.iter()) {
+            assert_eq!(actual_event.event_kind, expected_event.event_kind);
+            assert_eq!(actual_event.label, expected_event.label);
+            assert_eq!(actual_event.additional_data, expected_event.additional_data);
+            assert_eq!(
+                actual_event.timestamp.is_instant(),
+                expected_event.timestamp.is_instant()
+            );
+
+            count += 1;
+        }
     }
+
     assert_eq!(count, num_expected_events);
 }
 
-pub fn run_end_to_end_serialization_test<S: SerializationSink>(file_name_stem: &str) {
+fn collect_events_per_thread<'a>(events: &mut dyn Iterator<Item = Event<'a>>) -> FxHashMap<u32, Vec<Event<'a>>> {
+    let mut per_thread: FxHashMap<_, _> = Default::default();
+
+    for event in events {
+        per_thread.entry(event.thread_id).or_insert(Vec::new()).push(event);
+    }
+
+    per_thread
+}
+
+pub fn run_serialization_bench<S: SerializationSink>(file_name_stem: &str, num_events: usize, num_threads: usize) {
     let filestem = mk_filestem(file_name_stem);
-    let expected_events = generate_profiling_data::<S>(&filestem);
+    generate_profiling_data::<S>(&filestem, num_events, num_threads);
+}
+
+pub fn run_end_to_end_serialization_test<S: SerializationSink>(file_name_stem: &str, num_threads: usize) {
+    let filestem = mk_filestem(file_name_stem);
+    let expected_events = generate_profiling_data::<S>(&filestem, 10_000, num_threads);
     process_profiling_data(&filestem, &expected_events);
 }
 
 fn pseudo_invocation<S: SerializationSink>(
     profiler: &Profiler<S>,
     random: usize,
+    thread_id: u32,
     recursions_left: usize,
     event_ids: &[(StringId, StringId)],
     event_ids_as_str: &FxHashMap<StringId, &'static str>,
@@ -123,15 +170,14 @@ fn pseudo_invocation<S: SerializationSink>(
         return;
     }
 
-    let thread_id = (random % 3) as u32;
-
     let (event_kind, event_id) = event_ids[random % event_ids.len()];
 
     let _prof_guard = profiler.start_recording_interval_event(event_kind, event_id, thread_id);
 
     pseudo_invocation(
         profiler,
         random,
+        thread_id,
         recursions_left - 1,
         event_ids,
         event_ids_as_str,

analyzeme/tests/serialization.rs

@@ -2,11 +2,21 @@ use analyzeme::testing_common::run_end_to_end_serialization_test;
 use measureme::{FileSerializationSink, MmapSerializationSink};
 
 #[test]
-fn test_file_serialization_sink() {
-    run_end_to_end_serialization_test::<FileSerializationSink>("file_serialization_sink_test");
+fn test_file_serialization_sink_1_thread() {
+    run_end_to_end_serialization_test::<FileSerializationSink>("file_serialization_sink_test_1_thread", 1);
 }
 
 #[test]
-fn test_mmap_serialization_sink() {
-    run_end_to_end_serialization_test::<MmapSerializationSink>("mmap_serialization_sink_test");
+fn test_file_serialization_sink_8_threads() {
+    run_end_to_end_serialization_test::<FileSerializationSink>("file_serialization_sink_test_8_threads", 8);
+}
+
+#[test]
+fn test_mmap_serialization_sink_1_thread() {
+    run_end_to_end_serialization_test::<MmapSerializationSink>("mmap_serialization_sink_test_1_thread", 1);
+}
+
+#[test]
+fn test_mmap_serialization_sink_8_threads() {
+    run_end_to_end_serialization_test::<MmapSerializationSink>("mmap_serialization_sink_test_8_threads", 8);
 }

measureme/Cargo.toml

@@ -15,6 +15,7 @@ travis-ci = { repository = "rust-lang/measureme" }
 [dependencies]
 byteorder = "1.2.7"
 rustc-hash = "1.0.1"
+parking_lot = "0.9"
 
 [target.'cfg(not(target_arch="wasm32"))'.dependencies]
 memmap = "0.6.0"

measureme/src/file_serialization_sink.rs

@@ -1,12 +1,19 @@
 use crate::serialization::{Addr, SerializationSink};
 use std::error::Error;
 use std::fs;
-use std::io::{BufWriter, Write};
+use std::io::{Write};
 use std::path::Path;
-use std::sync::Mutex;
+use parking_lot::Mutex;
 
 pub struct FileSerializationSink {
-    data: Mutex<(BufWriter<fs::File>, u32)>,
+    data: Mutex<Inner>,
+}
+
+struct Inner {
+    file: fs::File,
+    buffer: Vec<u8>,
+    buf_pos: usize,
+    addr: u32,
 }
 
 impl SerializationSink for FileSerializationSink {
@@ -16,7 +23,12 @@ impl SerializationSink for FileSerializationSink {
         let file = fs::File::create(path)?;
 
         Ok(FileSerializationSink {
-            data: Mutex::new((BufWriter::new(file), 0)),
+            data: Mutex::new(Inner {
+                file,
+                buffer: vec![0; 1024*512],
+                buf_pos: 0,
+                addr: 0
+            }),
         })
     }
 
@@ -25,17 +37,58 @@ impl SerializationSink for FileSerializationSink {
     where
         W: FnOnce(&mut [u8]),
     {
-        let mut buffer = vec![0; num_bytes];
-        write(buffer.as_mut_slice());
+        let mut data = self.data.lock();
+        let Inner {
+            ref mut file,
+            ref mut buffer,
+            ref mut buf_pos,
+            ref mut addr
+        } = *data;
 
-        let mut data = self.data.lock().expect("couldn't acquire lock");
-        let curr_addr = data.1;
-        let file = &mut data.0;
+        let curr_addr = *addr;
+        *addr += num_bytes as u32;
 
-        file.write_all(&buffer).expect("failed to write buffer");
+        let buf_start = *buf_pos;
+        let buf_end = buf_start + num_bytes;
 
-        data.1 += num_bytes as u32;
+        if buf_end <= buffer.len() {
+            // We have enough space in the buffer, just write the data to it.
+            write(&mut buffer[buf_start .. buf_end]);
+            *buf_pos = buf_end;
+        } else {
+            // We don't have enough space in the buffer, so flush to disk
+            file.write_all(&buffer[..buf_start]).unwrap();
+
+            if num_bytes <= buffer.len() {
+                // There's enough space in the buffer, after flushing
+                write(&mut buffer[0 .. num_bytes]);
+                *buf_pos = num_bytes;
+            } else {
+                // Even after flushing the buffer there isn't enough space, so
+                // fall back to dynamic allocation
+                let mut temp_buffer = vec![0; num_bytes];
+                write(&mut temp_buffer[..]);
+                file.write_all(&temp_buffer[..]).unwrap();
+                *buf_pos = 0;
+            }
+        }
 
         Addr(curr_addr)
     }
 }
+
+impl Drop for FileSerializationSink {
+    fn drop(&mut self) {
+        let mut data = self.data.lock();
+        let Inner {
+            ref mut file,
+            ref mut buffer,
+            ref mut buf_pos,
+            addr: _,
+        } = *data;
+
+        if *buf_pos > 0 {
+            file.write_all(&buffer[..*buf_pos]).unwrap();
+        }
+    }
+}