From dc25c7564e0830d138a14137e64d7ec980223009 Mon Sep 17 00:00:00 2001
From: Carl Lerche <me@carllerche.com>
Date: Wed, 20 Jul 2016 13:27:41 +0200
Subject: [PATCH] Refactor heap allocation
---
Cargo.toml | 3 ++
src/alloc/heap.rs | 108 ++++++++++++++++++++++++++++++++++++----------
src/alloc/mod.rs | 103 ++++++++++++++++++-------------------------
src/buf/byte.rs | 34 ++++-----------
src/buf/ring.rs | 31 +------------
src/lib.rs | 2 +
src/str/seq.rs | 6 +--
7 files changed, 142 insertions(+), 145 deletions(-)
diff --git a/Cargo.toml b/Cargo.toml
index 3f9f385..5f5c745 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -18,6 +18,9 @@ exclude = [
"test/**/*"
]
+[dependencies]
+stable-heap = "0.1.0"
+
[dev-dependencies]
rand = "0.3.5"
diff --git a/src/alloc/heap.rs b/src/alloc/heap.rs
index 0de08eb..5508248 100644
--- a/src/alloc/heap.rs
+++ b/src/alloc/heap.rs
@@ -1,11 +1,17 @@
-use alloc::{Allocator, Mem, MemRef};
-use std::{mem, ptr, usize};
-use std::ops::DerefMut;
+use alloc::{Mem, MemRef};
+use stable_heap as heap;
+use std::{mem, ptr, isize, usize};
+use std::sync::atomic::{self, AtomicUsize, Ordering};
const MAX_ALLOC_SIZE: usize = usize::MAX;
+const MAX_REFCOUNT: usize = (isize::MAX) as usize;
pub struct Heap;
+struct Allocation {
+ refs: AtomicUsize,
+}
+
impl Heap {
pub fn allocate(&self, len: usize) -> MemRef {
// Make sure that the allocation is within the permitted range
@@ -13,40 +19,96 @@ impl Heap {
return MemRef::none();
}
- let alloc_len = len +
- mem::size_of::<Mem>() +
- mem::size_of::<Vec<u8>>();
-
unsafe {
- let mut vec: Vec<u8> = Vec::with_capacity(alloc_len);
- vec.set_len(alloc_len);
+ let mut ptr = heap::allocate(alloc_len(len), align());
+ let mut off = 0;
- let ptr = vec.deref_mut().as_mut_ptr();
+ ptr::write(ptr as *mut Allocation, Allocation::new());
- ptr::write(ptr as *mut Vec<u8>, vec);
+ off += mem::size_of::<Allocation>();
+ ptr::write(ptr.offset(off as isize) as *mut &Mem, &*(ptr as *const Allocation));
- let ptr = ptr.offset(mem::size_of::<Vec<u8>>() as isize);
- ptr::write(ptr as *mut Mem, Mem::new(len, mem::transmute(self as &Allocator)));
+ off += mem::size_of::<&Mem>();
+ ptr::write(ptr.offset(off as isize) as *mut usize, len);
- // Return the info
- MemRef::new(ptr as *mut Mem)
+ ptr = ptr.offset(mem::size_of::<Allocation>() as isize);
+
+ MemRef::new(ptr)
}
}
- pub fn deallocate(&self, mem: *mut Mem) {
+ fn deallocate(ptr: *mut u8) {
unsafe {
- let ptr = mem as *mut u8;
- let _ = ptr::read(ptr.offset(-(mem::size_of::<Vec<u8>>() as isize)) as *const Vec<u8>);
+ let off = mem::size_of::<Allocation>() + mem::size_of::<&Mem>();
+ let len = ptr::read(ptr.offset(off as isize) as *const usize);
+
+ heap::deallocate(ptr, alloc_len(len), align());
}
}
}
-impl Allocator for Heap {
- fn allocate(&self, len: usize) -> MemRef {
- Heap::allocate(self, len)
+impl Allocation {
+ fn new() -> Allocation {
+ Allocation {
+ refs: AtomicUsize::new(1),
+ }
}
+}
+
+impl Mem for Allocation {
+ fn ref_inc(&self) {
+ // Using a relaxed ordering is alright here, as knowledge of the
+ // original reference prevents other threads from erroneously deleting
+ // the object.
+ //
+ // As explained in the [Boost documentation][1], Increasing the
+ // reference counter can always be done with memory_order_relaxed: New
+ // references to an object can only be formed from an existing
+ // reference, and passing an existing reference from one thread to
+ // another must already provide any required synchronization.
+ //
+ // [1]: (www.boost.org/doc/libs/1_55_0/doc/html/atomic/usage_examples.html)
+ let old_size = self.refs.fetch_add(1, Ordering::Relaxed);
- fn deallocate(&self, mem: *mut Mem) {
- Heap::deallocate(self, mem)
+ // However we need to guard against massive refcounts in case someone
+ // is `mem::forget`ing Arcs. If we don't do this the count can overflow
+ // and users will use-after free. We racily saturate to `isize::MAX` on
+ // the assumption that there aren't ~2 billion threads incrementing
+ // the reference count at once. This branch will never be taken in
+ // any realistic program.
+ //
+ // We abort because such a program is incredibly degenerate, and we
+ // don't care to support it.
+ if old_size > MAX_REFCOUNT {
+ panic!("too many refs");
+ }
}
+
+ fn ref_dec(&self) {
+ if self.refs.fetch_sub(1, Ordering::Release) != 1 {
+ return;
+ }
+
+ atomic::fence(Ordering::Acquire);
+ Heap::deallocate(self as *const Allocation as *const u8 as *mut u8);
+ }
+}
+
+#[inline]
+fn alloc_len(bytes_len: usize) -> usize {
+ let len = bytes_len +
+ mem::size_of::<Allocation>() +
+ mem::size_of::<&Mem>() +
+ mem::size_of::<usize>();
+
+ if len & (align() - 1) == 0 {
+ len
+ } else {
+ (len & !align()) + align()
+ }
+}
+
+#[inline]
+fn align() -> usize {
+ mem::size_of::<usize>()
}
diff --git a/src/alloc/mod.rs b/src/alloc/mod.rs
index 3d8f04d..b597283 100644
--- a/src/alloc/mod.rs
+++ b/src/alloc/mod.rs
@@ -1,43 +1,35 @@
mod heap;
-pub use self::heap::{Heap};
+pub use self::heap::Heap;
+pub use self::pool::Pool;
use std::{mem, ptr};
-use std::sync::atomic::{AtomicUsize, Ordering};
pub fn heap(len: usize) -> MemRef {
Heap.allocate(len)
}
-/// Allocates memory to be used by Bufs or Bytes. Allows allocating memory
-/// using alternate stratgies than the default Rust heap allocator. Also does
-/// not require that allocations are continuous in memory.
-///
-/// For example, an alternate allocator could use a slab of 4kb chunks of
-/// memory and return as many chunks as needed to satisfy the length
-/// requirement.
-pub trait Allocator: Sync + Send {
+pub trait Mem: Send + Sync {
+ /// Increment the ref count
+ fn ref_inc(&self);
- /// Allocate memory. May or may not be contiguous.
- fn allocate(&self, len: usize) -> MemRef;
-
- /// Deallocate a chunk of memory
- fn deallocate(&self, mem: *mut Mem);
+ /// Decrement the ref count
+ fn ref_dec(&self);
}
pub struct MemRef {
+ // Pointer to the memory
+ // Layout:
+ // - &Mem
+ // - usize (len)
+ // - u8... bytes
ptr: *mut u8,
}
impl MemRef {
- pub fn new(mem: *mut Mem) -> MemRef {
- let ptr = mem as *mut u8;
-
- unsafe {
- MemRef {
- ptr: ptr.offset(mem::size_of::<Mem>() as isize),
- }
- }
+ #[inline]
+ pub unsafe fn new(ptr: *mut u8) -> MemRef {
+ MemRef { ptr: ptr }
}
#[inline]
@@ -51,14 +43,16 @@ impl MemRef {
}
#[inline]
- pub fn ptr(&self) -> *mut u8 {
- self.ptr
+ pub fn len(&self) -> usize {
+ unsafe { *self.len_ptr() }
}
+ #[inline]
pub fn bytes(&self) -> &[u8] {
use std::slice;
+
unsafe {
- slice::from_raw_parts(self.ptr(), self.mem().len)
+ slice::from_raw_parts(self.bytes_ptr(), self.len())
}
}
@@ -66,66 +60,51 @@ impl MemRef {
pub fn bytes_mut(&mut self) -> &mut [u8] {
use std::slice;
unsafe {
- slice::from_raw_parts_mut(self.ptr(), self.mem().len)
+ slice::from_raw_parts_mut(self.bytes_ptr(), self.len())
}
}
#[inline]
- fn mem_ptr(&self) -> *mut Mem {
+ fn mem(&self) -> &Mem {
unsafe {
- self.ptr.offset(-(mem::size_of::<Mem>() as isize)) as *mut Mem
+ *(self.ptr as *const &Mem)
}
}
#[inline]
- fn mem(&self) -> &Mem {
- unsafe {
- mem::transmute(self.mem_ptr())
- }
+ unsafe fn len_ptr(&self) -> *mut usize {
+ let off = mem::size_of::<&Mem>();
+ self.ptr.offset(off as isize) as *mut usize
+ }
+
+ #[inline]
+ unsafe fn bytes_ptr(&self) -> *mut u8 {
+ let off = mem::size_of::<&Mem>() + mem::size_of::<usize>();
+ self.ptr.offset(off as isize)
}
}
impl Clone for MemRef {
#[inline]
fn clone(&self) -> MemRef {
- self.mem().refs.fetch_add(1, Ordering::Relaxed);
+ if self.is_none() {
+ return MemRef::none();
+ }
+
+ self.mem().ref_inc();
MemRef { ptr: self.ptr }
}
}
impl Drop for MemRef {
fn drop(&mut self) {
- // Guard against the ref having already been dropped
- if self.ptr.is_null() { return; }
-
- // Decrement the ref count
- if 1 == self.mem().refs.fetch_sub(1, Ordering::Relaxed) {
- // Last ref dropped, free the memory
- unsafe {
- let alloc: &Allocator = mem::transmute(self.mem().allocator);
- alloc.deallocate(self.mem_ptr());
- }
+ if self.is_none() {
+ return;
}
+
+ self.mem().ref_dec();
}
}
unsafe impl Send for MemRef { }
unsafe impl Sync for MemRef { }
-
-/// Memory allocated by an Allocator must be prefixed with Mem
-pub struct Mem {
- // TODO: It should be possible to reduce the size of this struct
- allocator: *const Allocator,
- refs: AtomicUsize,
- len: usize,
-}
-
-impl Mem {
- pub fn new(len: usize, allocator: *const Allocator) -> Mem {
- Mem {
- allocator: allocator,
- refs: AtomicUsize::new(1),
- len: len,
- }
- }
-}
diff --git a/src/buf/byte.rs b/src/buf/byte.rs
index e536ab2..661abf9 100644
--- a/src/buf/byte.rs
+++ b/src/buf/byte.rs
@@ -1,5 +1,5 @@
use {alloc, Buf, Bytes, MutBuf, SeqByteStr, MAX_CAPACITY};
-use std::{cmp, fmt, ptr};
+use std::{cmp, fmt};
/*
*
@@ -101,13 +101,9 @@ impl ByteBuf {
pub fn read_slice(&mut self, dst: &mut [u8]) -> usize {
let len = cmp::min(dst.len(), self.remaining());
let cnt = len as u32;
+ let pos = self.pos as usize;
- unsafe {
- ptr::copy_nonoverlapping(
- self.mem.ptr().offset(self.pos as isize),
- dst.as_mut_ptr(),
- len);
- }
+ dst[0..len].copy_from_slice(&self.mem.bytes()[pos..pos+len]);
self.pos += cnt;
len
@@ -295,27 +291,15 @@ impl MutByteBuf {
#[inline]
pub fn write_slice(&mut self, src: &[u8]) -> usize {
- let cnt = src.len() as u32;
- let rem = self.buf.remaining_u32();
+ let cnt = cmp::min(src.len(), self.buf.remaining());
+ let pos = self.buf.pos as usize;
- if rem < cnt {
- self.write_ptr(src.as_ptr(), rem)
- } else {
- self.write_ptr(src.as_ptr(), cnt)
- }
- }
+ self.buf.mem.bytes_mut()[pos..pos+cnt]
+ .copy_from_slice(&src[0..cnt]);
- #[inline]
- fn write_ptr(&mut self, src: *const u8, len: u32) -> usize {
- unsafe {
- ptr::copy_nonoverlapping(
- src,
- self.buf.mem.ptr().offset(self.buf.pos as isize),
- len as usize);
+ self.buf.pos += cnt as u32;
- self.buf.pos += len;
- len as usize
- }
+ cnt
}
pub fn bytes<'a>(&'a self) -> &'a [u8] {
diff --git a/src/buf/ring.rs b/src/buf/ring.rs
index 75ab492..891eb12 100644
--- a/src/buf/ring.rs
+++ b/src/buf/ring.rs
@@ -1,5 +1,5 @@
use {alloc, Buf, MutBuf};
-use std::{cmp, fmt, ptr};
+use std::{cmp, fmt};
enum Mark {
NoMark,
@@ -137,35 +137,6 @@ impl RingBuf {
}
}
-impl Clone for RingBuf {
- fn clone(&self) -> RingBuf {
- use std::cmp;
-
- let mut ret = RingBuf::new(self.cap);
-
- ret.pos = self.pos;
- ret.len = self.len;
-
- unsafe {
- let to = self.pos + self.len;
-
- if to > self.cap {
- ptr::copy(self.ptr.ptr() as *const u8, ret.ptr.ptr(), to % self.cap);
- }
-
- ptr::copy(
- self.ptr.ptr().offset(self.pos as isize) as *const u8,
- ret.ptr.ptr().offset(self.pos as isize),
- cmp::min(self.len, self.cap - self.pos));
- }
-
- ret
- }
-
- // TODO: an improved version of clone_from is possible that potentially
- // re-uses the buffer
-}
-
impl fmt::Debug for RingBuf {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
write!(fmt, "RingBuf[.. {}]", self.len)
diff --git a/src/lib.rs b/src/lib.rs
index d5eca9b..0d15e78 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,6 +1,8 @@
#![crate_name = "bytes"]
#![deny(warnings)]
+extern crate stable_heap;
+
pub mod alloc;
pub mod buf;
pub mod str;
diff --git a/src/str/seq.rs b/src/str/seq.rs
index 2b67e6b..044f507 100644
--- a/src/str/seq.rs
+++ b/src/str/seq.rs
@@ -81,11 +81,7 @@ impl ops::Index<usize> for SeqByteStr {
fn index(&self, index: usize) -> &u8 {
assert!(index < self.len());
-
- unsafe {
- &*self.mem.ptr()
- .offset(index as isize + self.pos as isize)
- }
+ self.mem.bytes().index(index + self.pos as usize)
}
}
--
GitLab