Simplify the Allocator trait and ScratchSpaceMessageAllocator

Author: dwrensha

benchmark/benchmark.rs

@@ -19,7 +19,7 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 // THE SOFTWARE.
 
-use std::{mem, io};
+use std::{io};
 
 use capnp::{message, serialize, serialize_packed};
 use capnp::traits::Owned;
@@ -102,7 +102,7 @@ impl Serialize for Packed {
 trait Scratch<'a> {
     type Allocator: message::Allocator;
 
-    fn get_builders(&'a mut self) -> (message::Builder<Self::Allocator>, message::Builder<Self::Allocator>);
+    fn get_allocators(&'a mut self) -> (Self::Allocator, Self::Allocator);
 }
 
 const SCRATCH_SIZE: usize = 128 * 1024;
@@ -113,51 +113,43 @@ pub struct NoScratch;
 impl <'a> Scratch<'a> for NoScratch {
     type Allocator = message::HeapAllocator;
 
-    fn get_builders(&'a mut self) -> (message::Builder<Self::Allocator>, message::Builder<Self::Allocator>) {
-        (message::Builder::new_default(), message::Builder::new_default())
+    fn get_allocators(&'a mut self) -> (Self::Allocator, Self::Allocator) {
+        (message::HeapAllocator::new(), message::HeapAllocator::new())
     }
 }
 
 pub struct UseScratch {
-    _owned_space: ::std::vec::Vec<::std::vec::Vec<capnp::Word>>,
-    scratch_space: ::std::vec::Vec<message::ScratchSpace<'static>>,
+    buffer1: Vec<capnp::Word>,
+    buffer2: Vec<capnp::Word>,
 }
 
 impl UseScratch {
     pub fn new() -> UseScratch {
-        let mut owned = Vec::new();
-        let mut scratch = Vec::new();
-        for _ in 0..2 {
-            let mut words = capnp::Word::allocate_zeroed_vec(SCRATCH_SIZE);
-            scratch.push(message::ScratchSpace::new(
-                unsafe { mem::transmute(capnp::Word::words_to_bytes_mut(&mut words[..])) }));
-            owned.push(words);
-        }
         UseScratch {
-            _owned_space: owned,
-            scratch_space: scratch,
+            buffer1: capnp::Word::allocate_zeroed_vec(SCRATCH_SIZE),
+            buffer2: capnp::Word::allocate_zeroed_vec(SCRATCH_SIZE),
         }
     }
 }
 
 impl <'a> Scratch<'a> for UseScratch {
-    type Allocator = message::ScratchSpaceHeapAllocator<'a, 'a>;
-
-    fn get_builders(&'a mut self) -> (message::Builder<Self::Allocator>, message::Builder<Self::Allocator>) {
-        (message::Builder::new(message::ScratchSpaceHeapAllocator::new(
-            unsafe { mem::transmute(&mut self.scratch_space[0]) })),
-         message::Builder::new(message::ScratchSpaceHeapAllocator::new(
-             unsafe { mem::transmute(&mut self.scratch_space[1]) })))
+    type Allocator = message::ScratchSpaceHeapAllocator<'a>;
 
+    fn get_allocators(&'a mut self) -> (Self::Allocator, Self::Allocator) {
+        let UseScratch {ref mut buffer1, ref mut buffer2 } = self;
+        (message::ScratchSpaceHeapAllocator::new(capnp::Word::words_to_bytes_mut(buffer1)),
+         message::ScratchSpaceHeapAllocator::new(capnp::Word::words_to_bytes_mut(buffer2)))
     }
 }
 
 fn pass_by_object<S, T>(testcase: T, mut reuse: S, iters: u64) -> ::capnp::Result<()>
     where S: for<'a> Scratch<'a>, T: TestCase,
 {
     let mut rng = common::FastRand::new();
+    let (mut allocator_req, mut allocator_res) = reuse.get_allocators();
     for _ in 0..iters {
-        let (mut message_req, mut message_res) = reuse.get_builders();
+        let mut message_req = message::Builder::new(allocator_req);
+        let mut message_res = message::Builder::new(allocator_res);
 
         let expected = testcase.setup_request(
             &mut rng,
@@ -170,6 +162,9 @@ fn pass_by_object<S, T>(testcase: T, mut reuse: S, iters: u64) -> ::capnp::Resul
         testcase.check_response(
             message_res.get_root_as_reader()?,
             expected)?;
+
+        allocator_req = message_req.into_allocator();
+        allocator_res = message_res.into_allocator();
     }
     Ok(())
 }
@@ -180,8 +175,10 @@ fn pass_by_bytes<C, S, T>(testcase: T, mut reuse: S, compression: C, iters: u64)
     let mut request_bytes = vec![0u8; SCRATCH_SIZE * 8];
     let mut response_bytes = vec![0u8; SCRATCH_SIZE * 8];
     let mut rng = common::FastRand::new();
+    let (mut allocator_req, mut allocator_res) = reuse.get_allocators();
     for _ in 0..iters {
-        let (mut message_req, mut message_res) = reuse.get_builders();
+        let mut message_req = message::Builder::new(allocator_req);
+        let mut message_res = message::Builder::new(allocator_res);
 
         let expected = {
             let request = message_req.init_root();
@@ -217,6 +214,8 @@ fn pass_by_bytes<C, S, T>(testcase: T, mut reuse: S, compression: C, iters: u64)
 
         let response_reader = message_reader.get_root()?;
         testcase.check_response(response_reader, expected)?;
+        allocator_req = message_req.into_allocator();
+        allocator_res = message_res.into_allocator();
     }
     Ok(())
 }
@@ -227,9 +226,10 @@ fn server<C, S, T, R, W>(testcase: T, mut reuse: S, compression: C, iters: u64,
 {
     let mut out_buffered = io::BufWriter::new(&mut output);
     let mut in_buffered = io::BufReader::new(&mut input);
+    let (mut allocator_res, _) = reuse.get_allocators();
     for _ in 0..iters {
         use std::io::Write;
-        let (mut message_res, _) = reuse.get_builders();
+        let mut message_res = message::Builder::new(allocator_res);
 
         {
             let response = message_res.init_root();
@@ -242,6 +242,7 @@ fn server<C, S, T, R, W>(testcase: T, mut reuse: S, compression: C, iters: u64,
 
         compression.write_message(&mut out_buffered, &mut message_res)?;
         out_buffered.flush()?;
+        allocator_res = message_res.into_allocator();
     }
     Ok(())
 }
@@ -255,9 +256,10 @@ fn sync_client<C, S, T>(testcase: T, mut reuse: S, compression: C, iters: u64)
     let mut in_buffered = io::BufReader::new(&mut in_stream);
     let mut out_buffered = io::BufWriter::new(&mut out_stream);
     let mut rng = common::FastRand::new();
+    let (mut allocator_req, _) = reuse.get_allocators();
     for _ in 0..iters {
         use std::io::Write;
-        let (mut message_req, _) = reuse.get_builders();
+        let mut message_req = message::Builder::new(allocator_req);
 
         let expected = {
             let request = message_req.init_root();
@@ -271,6 +273,7 @@ fn sync_client<C, S, T>(testcase: T, mut reuse: S, compression: C, iters: u64)
             Default::default())?;
         let response_reader = message_reader.get_root()?;
         testcase.check_response(response_reader, expected)?;
+        allocator_req = message_req.into_allocator();
     }
     Ok(())
 }

capnp/src/lib.rs

@@ -27,6 +27,8 @@
 //! [capnpc-rust](https://github.com/capnproto/capnproto-rust/capnpc) crate.
 #![cfg_attr(feature = "rpc_try", feature(try_trait))]
 
+extern crate alloc;
+
 pub mod any_pointer;
 pub mod any_pointer_list;
 pub mod capability;

capnp/src/message.rs

@@ -262,32 +262,20 @@ impl <A, T> From<Builder<A>> for TypedReader<Builder<A>, T>
 /// An object that allocates memory for a Cap'n Proto message as it is being built.
 pub unsafe trait Allocator {
     /// Allocates zeroed memory for a new segment, returning a pointer to the start of the segment
-    /// and a u32 indicating the length of the segment (in words).
+    /// and a u32 indicating the length of the segment (in words). The allocated segment must be
+    /// at least `minimum_size` words long.
     ///
     /// UNSAFETY ALERT: Implementors must ensure all of the following:
     ///     1. the returned memory is initialized to all zeroes,
-    ///     2. the returned memory is valid for the remaining lifetime of the Allocator object,
+    ///     2. the returned memory is valid until deallocate_segment() is called on it,
     ///     3. the memory doesn't overlap with other allocated memory,
     ///     4. the memory is 8-byte aligned (or the "unaligned" feature is enabled for the capnp crate).
     fn allocate_segment(&mut self, minimum_size: u32) -> (*mut u8, u32);
 
-    /// Called before the Allocator is dropped, to allow the first segment to be re-zeroed
-    /// before possibly being reused in another allocator.
-    fn pre_drop(&mut self, _segment0_currently_allocated: u32) {}
+    /// Indicates that a segment, previously allocated via allocate_segment(), is no longer in use.
+    fn deallocate_segment(&mut self, ptr: *mut u8, word_size: u32, words_used: u32);
 }
 
-/* TODO(version 0.9): update to a more user-friendly trait here?
-pub trait Allocator {
-    /// Allocates memory for a new segment.
-    fn allocate_segment(&mut self, minimum_size: u32) -> Result<()>;
-
-    fn get_segment<'a>(&'a self, id: u32) -> &'a [u8];
-    fn get_segment_mut<'a>(&'a mut self, id: u32) -> &'a mut [u8];
-
-    fn pre_drop(&mut self, _segment0_currently_allocated: u32) {}
-}
-*/
-
 /// A container used to build a message.
 pub struct Builder<A> where A: Allocator {
     arena: BuilderArenaImpl<A>,
@@ -380,6 +368,10 @@ impl <A> Builder<A> where A: Allocator {
             nesting_limit: i32::max_value()
         })
     }
+
+    pub fn into_allocator(self) -> A {
+        self.arena.into_allocator()
+    }
 }
 
 impl <A> ReaderSegments for Builder<A> where A: Allocator {
@@ -392,11 +384,9 @@ impl <A> ReaderSegments for Builder<A> where A: Allocator {
     }
 }
 
-/// Standard segment allocator. Allocates each segment as a `Vec<Word>`, where `Word` has
-/// 8-byte alignment.
+/// Standard segment allocator. Allocates each segment via `alloc::alloc::alloc_zeroed()`.
 #[derive(Debug)]
 pub struct HeapAllocator {
-    owned_memory: Vec<Vec<crate::Word>>,
     next_size: u32,
     allocation_strategy: AllocationStrategy,
 }
@@ -412,8 +402,7 @@ pub const SUGGESTED_ALLOCATION_STRATEGY: AllocationStrategy = AllocationStrategy
 
 impl HeapAllocator {
     pub fn new() -> HeapAllocator {
-        HeapAllocator { owned_memory: Vec::new(),
-                        next_size: SUGGESTED_FIRST_SEGMENT_WORDS,
+        HeapAllocator { next_size: SUGGESTED_FIRST_SEGMENT_WORDS,
                         allocation_strategy: SUGGESTED_ALLOCATION_STRATEGY }
     }
 
@@ -431,16 +420,24 @@ impl HeapAllocator {
 unsafe impl Allocator for HeapAllocator {
     fn allocate_segment(&mut self, minimum_size: u32) -> (*mut u8, u32) {
         let size = ::std::cmp::max(minimum_size, self.next_size);
-        let mut new_words = crate::Word::allocate_zeroed_vec(size as usize);
-        let ptr = new_words.as_mut_ptr() as *mut u8;
-        self.owned_memory.push(new_words);
-
+        let ptr = unsafe {
+            alloc::alloc::alloc_zeroed(alloc::alloc::Layout::from_size_align(size as usize * BYTES_PER_WORD, 8).unwrap())
+        };
         match self.allocation_strategy {
             AllocationStrategy::GrowHeuristically => { self.next_size += size; }
             _ => { }
         }
         (ptr, size as u32)
     }
+
+    fn deallocate_segment(&mut self, ptr: *mut u8, word_size: u32, _words_used: u32) {
+        unsafe {
+            alloc::alloc::dealloc(ptr,
+                                  alloc::alloc::Layout::from_size_align(word_size as usize * BYTES_PER_WORD, 8).unwrap());
+        }
+        // TODO move this next_size stuff out of what the Allocator trait should be resposible for.
+        self.next_size = SUGGESTED_FIRST_SEGMENT_WORDS;
+    }
 }
 
 impl Builder<HeapAllocator> {
@@ -449,57 +446,59 @@ impl Builder<HeapAllocator> {
     }
 }
 
-#[derive(Debug)]
-pub struct ScratchSpace<'a> {
-    slice: &'a mut [u8],
-    in_use: bool,
-}
-
-impl <'a> ScratchSpace<'a> {
-    pub fn new(slice: &'a mut [u8]) -> ScratchSpace<'a> {
-        ScratchSpace { slice: slice, in_use: false }
-    }
-}
-
-pub struct ScratchSpaceHeapAllocator<'a, 'b: 'a> {
-    scratch_space: &'a mut ScratchSpace<'b>,
+/// An Allocator whose first segment is a backed by a user-provided buffer.
+pub struct ScratchSpaceHeapAllocator<'a> {
+    scratch_space: &'a mut [u8],
+    scratch_space_allocated: bool,
     allocator: HeapAllocator,
 }
 
-impl <'a, 'b: 'a> ScratchSpaceHeapAllocator<'a, 'b> {
-    pub fn new(scratch_space: &'a mut ScratchSpace<'b>) -> ScratchSpaceHeapAllocator<'a, 'b> {
+impl <'a> ScratchSpaceHeapAllocator<'a> {
+    /// Writes zeroes into the entire buffer and constructs a new allocator from it.
+    ///
+    /// If you want to reuse the same buffer and to minimize the cost of zeroing for each message,
+    /// you can call `message::Builder::into_allocator()` to recover the allocator from
+    /// the previous message and then pass it into the new message.
+    pub fn new(scratch_space: &'a mut [u8]) -> ScratchSpaceHeapAllocator<'a> {
+        // We need to ensure that the buffer is zeroed.
+        for b in &mut scratch_space[..] {
+            *b = 0;
+        }
         ScratchSpaceHeapAllocator { scratch_space: scratch_space,
+                                    scratch_space_allocated: false,
                                     allocator: HeapAllocator::new()}
     }
 
-    pub fn second_segment_words(self, value: u32) -> ScratchSpaceHeapAllocator<'a, 'b> {
-        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
-                                    allocator: self.allocator.first_segment_words(value) }
+    pub fn second_segment_words(self, value: u32) -> ScratchSpaceHeapAllocator<'a> {
+        ScratchSpaceHeapAllocator { allocator: self.allocator.first_segment_words(value), ..self }
 
     }
 
-    pub fn allocation_strategy(self, value: AllocationStrategy) -> ScratchSpaceHeapAllocator<'a, 'b> {
-        ScratchSpaceHeapAllocator { scratch_space: self.scratch_space,
-                                    allocator: self.allocator.allocation_strategy(value) }
+    pub fn allocation_strategy(self, value: AllocationStrategy) -> ScratchSpaceHeapAllocator<'a> {
+        ScratchSpaceHeapAllocator { allocator: self.allocator.allocation_strategy(value), ..self }
     }
-
 }
 
-unsafe impl <'a, 'b: 'a> Allocator for ScratchSpaceHeapAllocator<'a, 'b> {
+unsafe impl <'a> Allocator for ScratchSpaceHeapAllocator<'a> {
     fn allocate_segment(&mut self, minimum_size: u32) -> (*mut u8, u32) {
-        if !self.scratch_space.in_use {
-            self.scratch_space.in_use = true;
-            (self.scratch_space.slice.as_mut_ptr(), (self.scratch_space.slice.len() / BYTES_PER_WORD) as u32)
+        if (minimum_size as usize) < (self.scratch_space.len() / BYTES_PER_WORD) && !self.scratch_space_allocated {
+            self.scratch_space_allocated = true;
+            (self.scratch_space.as_mut_ptr(), (self.scratch_space.len() / BYTES_PER_WORD) as u32)
         } else {
             self.allocator.allocate_segment(minimum_size)
         }
     }
 
-    fn pre_drop(&mut self, segment0_currently_allocated: u32) {
-        let ptr = self.scratch_space.slice.as_mut_ptr();
-        unsafe {
-            ::std::ptr::write_bytes(ptr, 0u8, (segment0_currently_allocated as usize) * BYTES_PER_WORD);
+    fn deallocate_segment(&mut self, ptr: *mut u8, word_size: u32, words_used: u32) {
+        if ptr == self.scratch_space.as_mut_ptr() {
+            // Rezero the slice to allow reuse of the allocator. We only need to write
+            // words that we know might contain nonzero values.
+            unsafe {
+                ::std::ptr::write_bytes(ptr, 0u8, (words_used as usize) * BYTES_PER_WORD);
+            }
+            self.scratch_space_allocated = false;
+        } else {
+            self.allocator.deallocate_segment(ptr, word_size, words_used);
         }
-        self.scratch_space.in_use = false;
     }
 }