Support for custom allocators/ref counting

[mzabaluev]

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_std configurations 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> {}

[mzabaluev]

I have made the PoC API simpler and more flexible.

[DoumanAsh]

Custom ref counting would most likely require another type parameter.
But atomic ref counting should not be an issue in most cases

To be honest I think non-global allocator might over-complicated for 99% of users.

But if we want it, I think it might be better to wait for Alloc trait in std https://doc.rust-lang.org/std/alloc/trait.Alloc.html

We shouldn't come up with own allocator traits and instead use something common.
Obviously it requires us to wait for it's stabilization

I think having global allocator support as first step would be good start and we already have PR for this

[seanmonstar]

Something allowing this was added in #298.