Support for custom allocators/ref counting #269
Description
Allocation and reference counting of the buffers could be abstracted behind a generic type parameter. This can enable:
- Lightweight non-atomic reference counting for single-threaded use (raised previously as unsync / non-thread safe version #200).
- Fast task-bound, non-global allocators.
- Use in
no_stdconfigurations with a custom allocator.
Here's a sketch of a possible abstraction API:
use core::alloc::AllocErr;
use core::ptr::NonNull;
#[cfg(feature = "system_alloc")]
use std::alloc::System as SystemAlloc;
/// Represents results of an allocation: a handle that represents
/// the allocated block and is used to keep track of
/// the references, and the buffer's total capacity (which may be larger
/// than requested).
pub struct BufAllocation<H>(pub H, pub usize);
/// Handle to a shared buffer descriptor.
///
/// A type implementing this trait is an opaque handle that provides
/// low-level access to a shared data buffer and is used by higher-level
/// container implementations to keep track of references to the buffer.
/// The implementing type should also implement `Clone` by producing another
/// handle referring to the same buffer.
///
/// Neither the handle, nor a shared descriptor structure that it possibly
/// points to, own buffer's data. If the handle is dropped without the
/// `release` method called on it first, the implementation will leak the data,
/// the reference, or both.
///
/// The implementation type can exhibit specific thread safety
/// characteristics with regard to `Send` and `Sync`.
/// For global thread-safe allocators, it can be modeled as an `Arc`
/// pointing to a structure with the data pointer and the buffer's
/// allocated size.
///
/// It is up to the implementation whether to allocate the descriptor
/// and the data block in a single contiguous allocation or separately.
/// For arena-like allocators that never free individual allocations,
/// `BufHandle` could contain just the data pointer and a pointer to the
/// shared arena state.
///
pub trait BufHandle: Sized {
fn ptr(&self) -> NonNull<u8>;
/// Release the reference on the buffer represented by `self`
/// and deallocate buffer data if the reference was unique.
/// The typical place for calling this method is the `Drop`
/// implementation of a type containing the handle referenced by `self`.
///
/// # Safety
///
/// The only safe way to use this function is to call it exactly once
/// per the lifetime of the handle, without calling any other methods
/// afterwards before the handle is dropped. The caller also has to
/// ensure that no use of the allocated buffer data occurs after the
/// last reference to it is released.
///
unsafe fn release(&mut self);
unsafe fn make_unique(
&mut self,
data_ptr: *const u8,
data_len: usize,
new_capacity: usize,
) -> Result<BufAllocation<Self>, AllocErr>;
}
pub trait AllocBuf {
type Handle: BufHandle;
fn alloc_buf(
&mut self,
capacity: usize,
) -> Result<BufAllocation<Self::Handle>, AllocErr>;
}
#[cfg(feature = "system_alloc")]
impl AllocBuf for SystemAlloc {
...
}Bytes and BytesMut could then be parameterized with a BufHandle implementation in a backward-compatible way:
pub struct BytesMut<H = ArcBuf<SystemAlloc>> {
inner: InnerMut<H>,
}
impl BytesMut<ArcBuf<SystemAlloc>> {
pub fn with_capacity(capacity: usize) -> Self {
BytesMut::with_capacity_in(SystemAlloc, capacity)
}
}
impl<H: BufHandle> BytesMut<H> {
pub fn with_capacity_in<A>(mut alloc: A, capacity: usize) -> Self
where
A: AllocBuf<Handle = H>,
{
let BufAllocation(buf, cap) =
alloc.alloc_buf(capacity).unwrap_or_else(|_| {
alloc::handle_alloc_error(
Layout::from_size_align(capacity, 1).unwrap(),
)
});
let ptr = buf.ptr();
BytesMut {
inner: InnerMut::new(buf, ptr, 0, cap)
}
}
}
#[derive(Clone)]
pub struct Bytes<H = ArcBuf<SystemAlloc>> {
buf: H,
ptr: *const u8,
len: usize,
}Thread safety shall then depend on the choice of the allocator:
unsafe impl<H: Send> Send for Bytes<H> {}
unsafe impl<H: Sync> Sync for Bytes<H> {}