Extend FileSerializationSink to support in-memory storage and use this instead of ByteVecSink.

Author: michaelwoerister

analyzeme/src/profiling_data.rs

@@ -6,8 +6,9 @@ use measureme::file_header::{
     read_file_header, write_file_header, CURRENT_FILE_FORMAT_VERSION, FILE_HEADER_SIZE,
     FILE_MAGIC_EVENT_STREAM,
 };
-use measureme::ByteVecSink;
-use measureme::{EventId, ProfilerFiles, RawEvent, SerializationSink, StringTableBuilder};
+use measureme::{
+    EventId, FileSerializationSink, ProfilerFiles, RawEvent, SerializationSink, StringTableBuilder,
+};
 use serde::{Deserialize, Deserializer};
 use std::error::Error;
 use std::fs;
@@ -200,17 +201,17 @@ impl<'a> DoubleEndedIterator for ProfilerEventIterator<'a> {
 /// implementation might not be efficient, which why it should only be used for
 /// writing tests and other things that are not performance sensitive.
 pub struct ProfilingDataBuilder {
-    event_sink: ByteVecSink,
-    string_table_data_sink: Arc<ByteVecSink>,
-    string_table_index_sink: Arc<ByteVecSink>,
-    string_table: StringTableBuilder<ByteVecSink>,
+    event_sink: FileSerializationSink,
+    string_table_data_sink: Arc<FileSerializationSink>,
+    string_table_index_sink: Arc<FileSerializationSink>,
+    string_table: StringTableBuilder<FileSerializationSink>,
 }
 
 impl ProfilingDataBuilder {
     pub fn new() -> ProfilingDataBuilder {
-        let event_sink = ByteVecSink::new();
-        let string_table_data_sink = Arc::new(ByteVecSink::new());
-        let string_table_index_sink = Arc::new(ByteVecSink::new());
+        let event_sink = FileSerializationSink::new_in_memory();
+        let string_table_data_sink = Arc::new(FileSerializationSink::new_in_memory());
+        let string_table_index_sink = Arc::new(FileSerializationSink::new_in_memory());
 
         // The first thing in every file we generate must be the file header.
         write_file_header(&event_sink, FILE_MAGIC_EVENT_STREAM);

analyzeme/src/stringtable.rs

@@ -243,13 +243,13 @@ impl StringTable {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use measureme::{ByteVecSink, StringComponent, StringTableBuilder};
+    use measureme::{FileSerializationSink, StringComponent, StringTableBuilder};
     use std::sync::Arc;
 
     #[test]
     fn simple_strings() {
-        let data_sink = Arc::new(ByteVecSink::new());
-        let index_sink = Arc::new(ByteVecSink::new());
+        let data_sink = Arc::new(FileSerializationSink::new_in_memory());
+        let index_sink = Arc::new(FileSerializationSink::new_in_memory());
 
         let expected_strings = &[
             "abc",
@@ -289,8 +289,8 @@ mod tests {
 
     #[test]
     fn composite_string() {
-        let data_sink = Arc::new(ByteVecSink::new());
-        let index_sink = Arc::new(ByteVecSink::new());
+        let data_sink = Arc::new(FileSerializationSink::new_in_memory());
+        let index_sink = Arc::new(FileSerializationSink::new_in_memory());
 
         let expected_strings = &[
             "abc",                  // 0

measureme/src/file_header.rs

@@ -54,11 +54,11 @@ pub fn strip_file_header(data: &[u8]) -> &[u8] {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::serialization::ByteVecSink;
+    use crate::file_serialization_sink::FileSerializationSink;
 
     #[test]
     fn roundtrip() {
-        let data_sink = ByteVecSink::new();
+        let data_sink = FileSerializationSink::new_in_memory();
 
         write_file_header(&data_sink, FILE_MAGIC_EVENT_STREAM);
 
@@ -72,7 +72,7 @@ mod tests {
 
     #[test]
     fn invalid_magic() {
-        let data_sink = ByteVecSink::new();
+        let data_sink = FileSerializationSink::new_in_memory();
         write_file_header(&data_sink, FILE_MAGIC_STRINGTABLE_DATA);
         let mut data = data_sink.into_bytes();
 
@@ -83,7 +83,7 @@ mod tests {
 
     #[test]
     fn other_version() {
-        let data_sink = ByteVecSink::new();
+        let data_sink = FileSerializationSink::new_in_memory();
 
         write_file_header(&data_sink, FILE_MAGIC_STRINGTABLE_INDEX);
 

measureme/src/file_serialization_sink.rs

@@ -1,21 +1,76 @@
 use crate::serialization::{Addr, SerializationSink};
 use parking_lot::Mutex;
 use std::error::Error;
+use std::fmt::Debug;
 use std::fs;
 use std::io::Write;
 use std::path::Path;
 
+#[derive(Debug)]
 pub struct FileSerializationSink {
     data: Mutex<Inner>,
 }
 
+/// The `BackingStorage` is what the data gets written to.
+trait BackingStorage: Write + Send + Debug {
+    fn drain_bytes(&mut self) -> Vec<u8>;
+}
+
+impl BackingStorage for fs::File {
+    fn drain_bytes(&mut self) -> Vec<u8> {
+        unimplemented!()
+    }
+}
+
+impl BackingStorage for Vec<u8> {
+    fn drain_bytes(&mut self) -> Vec<u8> {
+        let mut bytes = Vec::new();
+        std::mem::swap(&mut bytes, self);
+        bytes
+    }
+}
+
+#[derive(Debug)]
 struct Inner {
-    file: fs::File,
+    file: Box<dyn BackingStorage>,
     buffer: Vec<u8>,
     buf_pos: usize,
     addr: u32,
 }
 
+impl FileSerializationSink {
+    pub fn new_in_memory() -> FileSerializationSink {
+        FileSerializationSink {
+            data: Mutex::new(Inner {
+                file: Box::new(Vec::new()),
+                buffer: vec![0; 1024 * 512],
+                buf_pos: 0,
+                addr: 0,
+            }),
+        }
+    }
+
+    /// Create a copy of all data written so far. This method meant to be used
+    /// for writing unit tests. It will panic if the underlying `BackingStorage`
+    /// does not implement `extract_bytes`.
+    pub fn into_bytes(self) -> Vec<u8> {
+        let mut data = self.data.lock();
+        let Inner {
+            ref mut file,
+            ref mut buffer,
+            ref mut buf_pos,
+            addr: _,
+        } = *data;
+
+        // We need to flush the buffer first.
+        file.write_all(&buffer[..*buf_pos]).unwrap();
+        *buf_pos = 0;
+
+        // Then we can create a copy of the data written so far.
+        file.drain_bytes()
+    }
+}
+
 impl SerializationSink for FileSerializationSink {
     fn from_path(path: &Path) -> Result<Self, Box<dyn Error + Send + Sync>> {
         fs::create_dir_all(path.parent().unwrap())?;
@@ -24,7 +79,7 @@ impl SerializationSink for FileSerializationSink {
 
         Ok(FileSerializationSink {
             data: Mutex::new(Inner {
-                file,
+                file: Box::new(file),
                 buffer: vec![0; 1024 * 512],
                 buf_pos: 0,
                 addr: 0,

measureme/src/lib.rs

@@ -36,10 +36,7 @@
 //! [`Profiler::start_recording_interval_event()`]: Profiler::start_recording_interval_event
 //! [`StringId`]: StringId
 #![allow(renamed_and_removed_lints)] // intra_doc_link_resolution_failure is renamed on nightly
-#![deny(
-    warnings,
-    intra_doc_link_resolution_failure,
-)]
+#![deny(warnings, intra_doc_link_resolution_failure)]
 
 pub mod event_id;
 pub mod file_header;
@@ -61,5 +58,5 @@ pub use crate::file_serialization_sink::FileSerializationSink;
 pub use crate::mmap_serialization_sink::MmapSerializationSink;
 pub use crate::profiler::{Profiler, ProfilerFiles, TimingGuard};
 pub use crate::raw_event::{RawEvent, MAX_INSTANT_TIMESTAMP, MAX_INTERVAL_TIMESTAMP};
-pub use crate::serialization::{Addr, ByteVecSink, SerializationSink};
+pub use crate::serialization::{Addr, SerializationSink};
 pub use crate::stringtable::{SerializableString, StringComponent, StringId, StringTableBuilder};

measureme/src/serialization.rs

@@ -1,4 +1,3 @@
-use parking_lot::Mutex;
 use std::error::Error;
 use std::path::Path;
 
@@ -32,49 +31,3 @@ pub trait SerializationSink: Sized + Send + Sync + 'static {
         self.write_atomic(bytes.len(), |sink| sink.copy_from_slice(bytes))
     }
 }
-
-/// A `SerializationSink` that writes to an internal `Vec<u8>` and can be
-/// converted into this raw `Vec<u8>`. This implementation is only meant to be
-/// used for testing and is not very efficient.
-pub struct ByteVecSink {
-    data: Mutex<Vec<u8>>,
-}
-
-impl ByteVecSink {
-    pub fn new() -> ByteVecSink {
-        ByteVecSink {
-            data: Mutex::new(Vec::new()),
-        }
-    }
-
-    pub fn into_bytes(self) -> Vec<u8> {
-        self.data.into_inner()
-    }
-}
-
-impl SerializationSink for ByteVecSink {
-    fn from_path(_path: &Path) -> Result<Self, Box<dyn Error + Send + Sync>> {
-        unimplemented!()
-    }
-
-    fn write_atomic<W>(&self, num_bytes: usize, write: W) -> Addr
-    where
-        W: FnOnce(&mut [u8]),
-    {
-        let mut data = self.data.lock();
-
-        let start = data.len();
-
-        data.resize(start + num_bytes, 0);
-
-        write(&mut data[start..]);
-
-        Addr(start as u32)
-    }
-}
-
-impl std::fmt::Debug for ByteVecSink {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "ByteVecSink")
-    }
-}