Merge pull request #97 from michaelwoerister/write_bytes_atomic

SerializationSink: Add write_bytes_atomic method.

Author: wesleywiser

measureme/src/file_serialization_sink.rs

@@ -1,9 +1,9 @@
 use crate::serialization::{Addr, SerializationSink};
+use parking_lot::Mutex;
 use std::error::Error;
 use std::fs;
-use std::io::{Write};
+use std::io::Write;
 use std::path::Path;
-use parking_lot::Mutex;
 
 pub struct FileSerializationSink {
     data: Mutex<Inner>,
@@ -25,9 +25,9 @@ impl SerializationSink for FileSerializationSink {
         Ok(FileSerializationSink {
             data: Mutex::new(Inner {
                 file,
-                buffer: vec![0; 1024*512],
+                buffer: vec![0; 1024 * 512],
                 buf_pos: 0,
-                addr: 0
+                addr: 0,
             }),
         })
     }
@@ -42,7 +42,7 @@ impl SerializationSink for FileSerializationSink {
             ref mut file,
             ref mut buffer,
             ref mut buf_pos,
-            ref mut addr
+            ref mut addr,
         } = *data;
 
         let curr_addr = *addr;
@@ -53,15 +53,15 @@ impl SerializationSink for FileSerializationSink {
 
         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]);
+            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]);
+                write(&mut buffer[0..num_bytes]);
                 *buf_pos = num_bytes;
             } else {
                 // Even after flushing the buffer there isn't enough space, so
@@ -75,6 +75,36 @@ impl SerializationSink for FileSerializationSink {
 
         Addr(curr_addr)
     }
+
+    fn write_bytes_atomic(&self, bytes: &[u8]) -> Addr {
+        if bytes.len() < 128 {
+            // For "small" pieces of data, use the regular implementation so we
+            // don't repeatedly flush an almost empty buffer to disk.
+            return self.write_atomic(bytes.len(), |sink| sink.copy_from_slice(bytes));
+        }
+
+        let mut data = self.data.lock();
+        let Inner {
+            ref mut file,
+            ref mut buffer,
+            ref mut buf_pos,
+            ref mut addr,
+        } = *data;
+
+        let curr_addr = *addr;
+        *addr += bytes.len() as u32;
+
+        if *buf_pos > 0 {
+            // There's something in the buffer, flush it to disk
+            file.write_all(&buffer[..*buf_pos]).unwrap();
+            *buf_pos = 0;
+        }
+
+        // Now write the whole input to disk, skipping the write buffer
+        file.write_all(bytes).unwrap();
+
+        Addr(curr_addr)
+    }
 }
 
 impl Drop for FileSerializationSink {

measureme/src/serialization.rs

@@ -1,6 +1,6 @@
+use parking_lot::Mutex;
 use std::error::Error;
 use std::path::Path;
-use parking_lot::Mutex;
 
 #[derive(Clone, Copy, Eq, PartialEq, Debug)]
 pub struct Addr(pub u32);
@@ -14,9 +14,23 @@ impl Addr {
 pub trait SerializationSink: Sized + Send + Sync + 'static {
     fn from_path(path: &Path) -> Result<Self, Box<dyn Error>>;
 
+    /// Atomically write `num_bytes` to the sink. The implementation must ensure
+    /// that concurrent invocations of `write_atomic` do not conflict with each
+    /// other.
+    ///
+    /// The `write` argument is a function that must fill the output buffer
+    /// passed to it. The output buffer is guaranteed to be exactly `num_bytes`
+    /// large.
     fn write_atomic<W>(&self, num_bytes: usize, write: W) -> Addr
     where
         W: FnOnce(&mut [u8]);
+
+    /// Same as write_atomic() but might be faster in cases where bytes to be
+    /// written are already present in a buffer (as opposed to when it is
+    /// benefical to directly serialize into the output buffer).
+    fn write_bytes_atomic(&self, bytes: &[u8]) -> Addr {
+        self.write_atomic(bytes.len(), |sink| sink.copy_from_slice(bytes))
+    }
 }
 
 /// A `SerializationSink` that writes to an internal `Vec<u8>` and can be