use bytes::Bytes;
use std::{
alloc::{Layout, alloc, alloc_zeroed, dealloc, handle_alloc_error},
mem::ManuallyDrop,
ptr::{self, NonNull},
sync::atomic::Ordering,
};
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
use loom::sync::atomic::{AtomicPtr, AtomicUsize, fence};
} else {
use std::sync::atomic::{AtomicPtr, AtomicUsize, fence};
}
}
const MAX_REFCOUNT: usize = isize::MAX as usize;
const TAG: usize = 1;
const _: () = assert!(align_of::<Record>() >= 2);
#[derive(Clone, Copy, Debug)]
struct Format {
align_bits: u32,
tail_bits: u32,
size_bits: u32,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct Fields {
size: usize,
align: usize,
tail: usize,
}
impl Format {
#[inline(always)]
const fn bits(self) -> u32 {
1 + self.align_bits + self.tail_bits + self.size_bits
}
#[inline(always)]
const fn tail_shift(self) -> u32 {
1 + self.align_bits
}
#[inline(always)]
const fn size_shift(self) -> u32 {
self.tail_shift() + self.tail_bits
}
#[inline(always)]
const fn tail_mask(self) -> u64 {
((1 << self.tail_bits) - 1) << self.tail_shift()
}
#[inline(always)]
const fn tail_field(self, tail: usize) -> Option<u64> {
let tail = tail as u64;
if tail >> self.tail_bits != 0 {
return None;
}
Some(tail << self.tail_shift())
}
#[inline(always)]
const fn pack(self, layout: Layout, tail: usize) -> Option<u64> {
let size = layout.size() as u64;
let log2_align = layout.align().trailing_zeros() as u64;
if size == 0 || size >> self.size_bits != 0 || log2_align >> self.align_bits != 0 {
return None;
}
match self.tail_field(tail) {
Some(tail) => Some((size << self.size_shift()) | tail | (log2_align << 1) | TAG as u64),
None => None,
}
}
#[inline(always)]
const fn with_tail(self, word: u64, tail: usize) -> Option<u64> {
match self.tail_field(tail) {
Some(tail) => Some((word & !self.tail_mask()) | tail),
None => None,
}
}
#[inline(always)]
const fn unpack(self, word: u64) -> Fields {
let log2_align = (word >> 1) & ((1 << self.align_bits) - 1);
Fields {
size: (word >> self.size_shift()) as usize,
align: 1 << log2_align,
tail: ((word & self.tail_mask()) >> self.tail_shift()) as usize,
}
}
}
const FORMAT_64: Format = Format {
align_bits: 6,
tail_bits: 17,
size_bits: 40,
};
const FORMAT_32: Format = Format {
align_bits: 0,
tail_bits: 7,
size_bits: 24,
};
const _: () = assert!(FORMAT_64.bits() == 64 && FORMAT_32.bits() == 32);
cfg_if::cfg_if! {
if #[cfg(target_pointer_width = "64")] {
const FORMAT: Format = FORMAT_64;
} else if #[cfg(target_pointer_width = "32")] {
const FORMAT: Format = FORMAT_32;
} else {
const FORMAT: Format = Format {
align_bits: 0,
tail_bits: 0,
size_bits: 0,
};
}
}
const _: () = assert!(FORMAT.bits() <= usize::BITS);
#[derive(Clone, Copy, Debug)]
pub(crate) struct Descriptor(usize);
impl Descriptor {
#[inline(always)]
const fn new(layout: Layout, tail: usize) -> Option<Self> {
match FORMAT.pack(layout, tail) {
Some(word) => Some(Self(word as usize)),
None => None,
}
}
#[inline(always)]
const fn fields(self) -> Fields {
FORMAT.unpack(self.0 as u64)
}
#[inline(always)]
const fn layout(self) -> Layout {
let Fields { size, align, .. } = self.fields();
unsafe { Layout::from_size_align_unchecked(size, align) }
}
#[inline(always)]
pub(crate) const fn tail(self) -> usize {
self.fields().tail
}
#[inline(always)]
const fn with_tail(self, tail: usize) -> Option<Self> {
match FORMAT.with_tail(self.0 as u64, tail) {
Some(word) => Some(Self(word as usize)),
None => None,
}
}
#[inline(always)]
pub(crate) const fn without_tail(self) -> Self {
Self(self.0 & !(FORMAT.tail_mask() as usize))
}
#[inline(always)]
const unsafe fn base(self, ptr: NonNull<u8>, extent: usize) -> NonNull<u8> {
let Fields { size, tail, .. } = self.fields();
let offset = size.wrapping_sub(tail).wrapping_sub(extent);
unsafe { ptr.sub(offset) }
}
#[inline(always)]
unsafe fn into_record_inline(self, ptr: NonNull<u8>, extent: usize, refs: usize) -> Owner {
let base = unsafe { self.base(ptr, extent) };
Owner::with_refs(
Storage::Native {
base,
layout: self.layout(),
},
refs,
)
}
#[cold]
#[inline(never)]
unsafe fn into_record(self, ptr: NonNull<u8>, extent: usize, refs: usize) -> Owner {
unsafe { self.into_record_inline(ptr, extent, refs) }
}
#[inline(always)]
unsafe fn release(self, ptr: NonNull<u8>, extent: usize) {
unsafe { dealloc(self.base(ptr, extent).as_ptr(), self.layout()) };
}
}
impl From<Descriptor> for Owner {
#[inline(always)]
fn from(descriptor: Descriptor) -> Self {
Self(ptr::without_provenance_mut(descriptor.0))
}
}
#[cold]
#[inline(never)]
fn native_record(base: NonNull<u8>, layout: Layout) -> Owner {
Owner::new(Storage::Native { base, layout })
}
#[repr(C)]
struct BytesOwner {
owner: Owner,
ptr: NonNull<u8>,
len: usize,
}
unsafe impl Send for BytesOwner {}
impl AsRef<[u8]> for BytesOwner {
#[inline]
fn as_ref(&self) -> &[u8] {
unsafe { std::slice::from_raw_parts(self.ptr.as_ptr(), self.len) }
}
}
impl Drop for BytesOwner {
#[inline]
fn drop(&mut self) {
unsafe { self.owner.release_view(self.ptr, self.len) };
}
}
pub(crate) struct OwnerCell(AtomicPtr<Record>);
impl OwnerCell {
#[allow(clippy::missing_const_for_fn)]
#[inline(always)]
pub(crate) fn new(word: Owner) -> Self {
Self(AtomicPtr::new(word.0))
}
#[inline(always)]
pub(crate) fn empty() -> Self {
Self::new(Owner::empty())
}
#[inline(always)]
pub(crate) fn load(&self) -> Owner {
Owner(self.0.load(Ordering::Acquire))
}
#[inline(always)]
pub(crate) fn get(&mut self) -> Owner {
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
Owner(self.0.with_mut(|word| *word))
} else {
Owner(*self.0.get_mut())
}
}
}
#[inline(always)]
fn set(&mut self, word: Owner) {
cfg_if::cfg_if! {
if #[cfg(feature = "loom")] {
self.0.with_mut(|slot| *slot = word.0);
} else {
*self.0.get_mut() = word.0;
}
}
}
#[inline(always)]
pub(crate) unsafe fn promote(
&self,
observed: Descriptor,
ptr: NonNull<u8>,
len: usize,
) -> Owner {
let candidate = unsafe { observed.into_record_inline(ptr, len, 2) };
match self.0.compare_exchange(
Owner::from(observed).0,
candidate.0,
Ordering::Release,
Ordering::Acquire,
) {
Ok(_) => candidate,
Err(winner) => {
unsafe { candidate.discard() };
let winner = Owner(winner);
assert!(
winner.descriptor().is_none(),
"a promoted owner never returns to a descriptor"
);
unsafe { winner.clone_shared() };
winner
}
}
}
#[inline(always)]
pub(crate) unsafe fn retain_mut(&mut self, ptr: NonNull<u8>, len: usize) -> Owner {
let owner = self.get();
if let Some(descriptor) = owner.descriptor() {
let record = unsafe { descriptor.into_record(ptr, len, 2) };
self.set(record);
return record;
}
unsafe { owner.clone_shared() };
owner
}
#[cfg(all(test, feature = "loom"))]
pub(crate) fn holds_record_relaxed(&self) -> bool {
let owner = Owner(self.0.load(Ordering::Relaxed));
owner.descriptor().is_none() && !owner.is_empty()
}
#[cfg(all(test, feature = "loom"))]
pub(crate) unsafe fn refcount_relaxed(&self) -> usize {
let owner = Owner(self.0.load(Ordering::Relaxed));
assert!(owner.descriptor().is_none(), "owner must be a record");
unsafe { owner.refcount_relaxed() }
}
}
#[derive(Clone, Copy, Debug)]
#[repr(transparent)]
pub(crate) struct Owner(*mut Record);
#[repr(C)]
pub(crate) struct Record {
refs: AtomicUsize,
storage: Storage,
#[cfg(feature = "loom")]
witness: Witness,
}
#[cfg(feature = "loom")]
struct Witness {
fields: loom::cell::UnsafeCell<()>,
_leak: loom::alloc::Track<()>,
}
pub(crate) enum Storage {
Native { base: NonNull<u8>, layout: Layout },
External(Bytes),
}
impl Owner {
#[inline(always)]
pub(crate) const fn empty() -> Self {
Self(ptr::null_mut())
}
#[inline(always)]
pub(crate) const fn is_empty(self) -> bool {
self.0.is_null()
}
#[inline(always)]
pub(crate) fn descriptor(self) -> Option<Descriptor> {
let word = self.0.addr();
(word & TAG != 0).then_some(Descriptor(word))
}
#[inline(always)]
fn native(base: NonNull<u8>, layout: Layout) -> Self {
Descriptor::new(layout, 0).map_or_else(|| native_record(base, layout), Self::from)
}
#[inline]
pub(crate) fn allocate(capacity: usize, alignment: usize, zeroed: bool) -> (NonNull<u8>, Self) {
assert!(capacity > 0, "capacity must be greater than zero");
let layout = Layout::from_size_align(capacity, alignment)
.expect("alignment must be a power of two and the size must fit a layout");
let ptr = if zeroed {
unsafe { alloc_zeroed(layout) }
} else {
unsafe { alloc(layout) }
};
let base = NonNull::new(ptr).unwrap_or_else(|| handle_alloc_error(layout));
(base, Self::native(base, layout))
}
#[inline]
pub(crate) fn adopt_vec(vec: Vec<u8>) -> (NonNull<u8>, usize, usize, Self) {
let mut vec = ManuallyDrop::new(vec);
let len = vec.len();
let cap = vec.capacity();
assert!(cap > 0, "capacity must be greater than zero");
let layout = Layout::array::<u8>(cap).expect("vec capacity must fit a layout");
let base = unsafe { NonNull::new_unchecked(vec.as_mut_ptr()) };
(base, len, cap, Self::native(base, layout))
}
#[inline(always)]
pub(crate) unsafe fn frozen(self, ptr: NonNull<u8>, len: usize, cap: usize) -> Self {
let Some(descriptor) = self.descriptor() else {
return self;
};
descriptor.with_tail(cap.wrapping_sub(len)).map_or_else(
|| unsafe { descriptor.into_record(ptr, cap, 1) },
Self::from,
)
}
#[inline(always)]
pub(crate) unsafe fn release_view(self, ptr: NonNull<u8>, len: usize) {
if unsafe { self.drop_view() } {
unsafe { self.release_unique(ptr, len) };
}
}
#[inline(always)]
pub(crate) unsafe fn drop_view(self) -> bool {
if self.descriptor().is_some() {
return true;
}
unsafe { self.drop_unless_unique() }
}
#[inline]
pub(crate) unsafe fn into_bytes(self, ptr: NonNull<u8>, len: usize) -> Bytes {
Bytes::from_owner(BytesOwner {
owner: self,
ptr,
len,
})
}
#[inline(never)]
pub(crate) unsafe fn release_unique(self, ptr: NonNull<u8>, extent: usize) {
match self.descriptor() {
Some(descriptor) => unsafe { descriptor.release(ptr, extent) },
None => unsafe { self.release_record() },
}
}
#[inline]
pub(crate) fn new(storage: Storage) -> Self {
Self::with_refs(storage, 1)
}
#[inline]
fn with_refs(storage: Storage, refs: usize) -> Self {
let record = Box::into_raw(Box::new(Record {
refs: AtomicUsize::new(refs),
storage,
#[cfg(feature = "loom")]
witness: Witness {
fields: loom::cell::UnsafeCell::new(()),
_leak: loom::alloc::Track::new(()),
},
}));
Self(record)
}
#[inline(always)]
unsafe fn record<'a>(self) -> &'a Record {
let record = unsafe { &*self.0 };
#[cfg(feature = "loom")]
record.witness.fields.with(|_| ());
record
}
#[inline(always)]
unsafe fn refs<'a>(self) -> &'a AtomicUsize {
&unsafe { self.record() }.refs
}
#[inline(always)]
pub(crate) unsafe fn storage<'a>(self) -> Option<&'a Storage> {
if self.is_empty() {
return None;
}
Some(&unsafe { self.record() }.storage)
}
#[inline(always)]
pub(crate) unsafe fn clone_shared(self) {
if self.is_empty() {
return;
}
let old = unsafe { self.refs() }.fetch_add(1, Ordering::Relaxed);
if old > MAX_REFCOUNT {
std::process::abort();
}
}
#[cfg(test)]
pub(crate) unsafe fn drop_shared(self) {
if unsafe { self.drop_unless_unique() } {
unsafe { self.release_record() };
}
}
#[inline(always)]
unsafe fn drop_unless_unique(self) -> bool {
if self.is_empty() {
return false;
}
let refs = unsafe { self.refs() };
if refs.load(Ordering::Acquire) == 1 {
return true;
}
let old = refs.fetch_sub(1, Ordering::Release);
if old == 1 {
unsafe { self.drop_shared_race_final() };
}
false
}
#[cold]
#[inline(never)]
unsafe fn drop_shared_race_final(self) {
fence(Ordering::Acquire);
unsafe { self.release_record() };
}
#[inline(never)]
unsafe fn release_record(self) {
let Some(record_ptr) = NonNull::new(self.0) else {
return;
};
let record = unsafe { record_ptr.as_ref() };
#[cfg(feature = "loom")]
record.witness.fields.with_mut(|_| ());
assert!(
record.refs.load(Ordering::Relaxed) <= 1,
"a record must not be released while shared"
);
match record.storage {
Storage::Native { base, layout } => {
unsafe { dealloc(base.as_ptr(), layout) };
#[cfg(feature = "loom")]
unsafe {
ptr::drop_in_place(&raw mut (*record_ptr.as_ptr()).witness)
};
}
Storage::External(_) => unsafe { record_ptr.drop_in_place() },
}
unsafe { dealloc(record_ptr.as_ptr().cast(), Layout::new::<Record>()) };
}
unsafe fn discard(self) {
let record = unsafe { NonNull::new_unchecked(self.0) };
#[cfg(feature = "loom")]
unsafe {
ptr::drop_in_place(&raw mut (*record.as_ptr()).witness)
};
unsafe { dealloc(record.as_ptr().cast(), Layout::new::<Record>()) };
}
#[inline(always)]
pub(crate) unsafe fn is_unique(self) -> bool {
unsafe { self.refs() }.load(Ordering::Acquire) == 1
}
#[cfg(all(test, not(feature = "loom")))]
pub(crate) unsafe fn refcount(self) -> Option<usize> {
if self.is_empty() {
return None;
}
Some(unsafe { self.refs() }.load(Ordering::Acquire))
}
#[cfg(all(test, feature = "loom"))]
pub(crate) unsafe fn refcount_relaxed(self) -> usize {
unsafe { self.refs() }.load(Ordering::Relaxed)
}
}
#[cfg(all(test, not(feature = "loom")))]
#[commonware_macros::test_group("miri")]
mod tests {
use super::*;
use std::{
panic::{self, AssertUnwindSafe},
sync::Arc,
};
const PAGE_ALIGNMENT: usize = 4096;
struct Tracker(Arc<AtomicUsize>);
impl Drop for Tracker {
fn drop(&mut self) {
self.0.fetch_add(1, Ordering::SeqCst);
}
}
impl AsRef<[u8]> for Tracker {
fn as_ref(&self) -> &[u8] {
&[1, 2, 3]
}
}
fn layout(size: usize, align: usize) -> Layout {
Layout::from_size_align(size, align).unwrap()
}
fn round_trip(format: Format, size: usize, align: usize, tail: usize) -> Fields {
let word = format
.pack(layout(size, align), tail)
.expect("fields must fit");
assert_eq!(word & 1, 1);
assert!(format.bits() == 64 || word >> format.bits() == 0);
format.unpack(word)
}
fn descriptor(word: Owner) -> Descriptor {
word.descriptor().expect("owner must be a descriptor")
}
fn record(word: Owner) -> Owner {
assert!(
word.descriptor().is_none() && !word.is_empty(),
"owner must be a record"
);
word
}
fn record_addr(owner: Owner) -> usize {
owner.0.addr()
}
fn heap(owner: Owner) -> (NonNull<u8>, Layout) {
match unsafe { owner.storage() } {
Some(&Storage::Native { base, layout }) => (base, layout),
_ => panic!("owner must be native"),
}
}
fn is_external(owner: Owner) -> bool {
matches!(unsafe { owner.storage() }, Some(Storage::External(_)))
}
fn native_record(capacity: usize, alignment: usize) -> (NonNull<u8>, Owner) {
let (data, word) = Owner::allocate(capacity, alignment, false);
let owner = unsafe { descriptor(word).into_record(data, capacity, 1) };
(data, owner)
}
fn frozen_view(
capacity: usize,
offset: usize,
len: usize,
) -> (NonNull<u8>, NonNull<u8>, Owner) {
let (base, word) = Owner::allocate(capacity, 8, false);
let ptr = unsafe { base.add(offset) };
let word = unsafe { word.frozen(ptr, len, capacity - offset) };
(base, ptr, word)
}
#[test]
fn test_format_64_round_trips() {
let max_size = (1 << 40) - 1;
let max_tail = (1 << 17) - 1;
for (size, align, tail) in [
(1, 1, 0),
(10, 1, 9),
(4096, PAGE_ALIGNMENT, max_tail),
(max_size, 1, 0),
(1, 1, max_tail),
(max_size, 1 << 30, max_tail),
(1, 1 << 62, 0),
] {
let fields = round_trip(FORMAT_64, size, align, tail);
assert_eq!(fields, Fields { size, align, tail });
}
}
#[test]
fn test_format_64_rejects_fields_that_do_not_fit() {
assert_eq!(FORMAT_64.pack(layout(1 << 40, 1), 0), None);
assert_eq!(FORMAT_64.pack(layout(64, 1), 1 << 17), None);
assert_eq!(FORMAT_64.pack(layout(0, 1), 0), None);
}
#[test]
fn test_format_32_round_trips() {
let max_size = (1 << 24) - 1;
for (size, tail) in [
(1, 0),
(1, 127),
(255, 127),
(256, 0),
(max_size, 0),
(max_size, 127),
] {
let fields = round_trip(FORMAT_32, size, 1, tail);
assert_eq!(
fields,
Fields {
size,
align: 1,
tail
}
);
}
}
#[test]
fn test_format_32_rejects_fields_that_do_not_fit() {
assert_eq!(FORMAT_32.pack(layout(1 << 24, 1), 0), None);
assert_eq!(FORMAT_32.pack(layout(64, 2), 0), None);
assert_eq!(FORMAT_32.pack(layout(64, 1), 128), None);
assert_eq!(FORMAT_32.pack(layout(0, 1), 0), None);
}
#[test]
fn test_format_tail_replacement_keeps_layout() {
for (format, size, align, max_tail) in [
(FORMAT_64, 1000, 64, (1 << 17) - 1),
(FORMAT_32, 1000, 1, 127),
] {
let word = format.pack(layout(size, align), 0).unwrap();
let replaced = format.with_tail(word, max_tail).unwrap();
assert_eq!(
format.unpack(replaced),
Fields {
size,
align,
tail: max_tail
}
);
let cleared = format.with_tail(replaced, 0).unwrap();
assert_eq!(cleared, word);
assert_eq!(format.with_tail(word, max_tail + 1), None);
}
}
#[test]
fn test_format_without_bits_describes_nothing() {
let none = Format {
align_bits: 0,
tail_bits: 0,
size_bits: 0,
};
assert_eq!(none.pack(layout(1, 1), 0), None);
}
#[test]
fn test_heap_allocation_is_exact_and_aligned() {
let page = PAGE_ALIGNMENT;
let (data, word) = Owner::allocate(4096, page, false);
assert!((data.as_ptr() as usize).is_multiple_of(page));
let descriptor = descriptor(word);
assert_eq!(descriptor.layout(), layout(4096, page));
assert_eq!(descriptor.tail(), 0);
assert_eq!(unsafe { descriptor.base(data, 4096) }, data);
unsafe { word.release_unique(data, 4096) };
}
#[test]
fn test_heap_zeroed_only_exposes_usable_region() {
let (data, word) = Owner::allocate(64, PAGE_ALIGNMENT, true);
let bytes = unsafe { std::slice::from_raw_parts(data.as_ptr(), 64) };
assert_eq!(bytes, &[0u8; 64]);
unsafe { word.release_unique(data, 64) };
}
#[test]
fn test_heap_capacity_is_not_rounded() {
let (data, word) = Owner::allocate(10, 1, false);
assert_eq!(descriptor(word).layout().size(), 10);
unsafe { word.release_unique(data, 10) };
}
#[test]
#[should_panic(expected = "alignment must be a power of two and the size must fit a layout")]
fn test_allocation_rejects_isize_max_overflow() {
let _ = Owner::allocate(isize::MAX as usize, 64, false);
}
#[test]
#[should_panic(expected = "alignment must be a power of two and the size must fit a layout")]
fn test_allocation_rejects_non_power_of_two_alignment() {
let _ = Owner::allocate(64, 3, false);
}
#[test]
fn test_descriptor_recovers_base_from_any_view() {
let (base, ptr, word) = frozen_view(64, 8, 32);
let frozen = descriptor(word);
assert_eq!(frozen.layout(), layout(64, 8));
assert_eq!(frozen.tail(), 24);
unsafe {
assert_eq!(frozen.base(ptr, 32), base);
assert_eq!(frozen.base(ptr.add(32), 0), base);
assert_eq!(frozen.without_tail().base(ptr, 56), base);
assert_eq!(frozen.without_tail().base(base.add(64), 0), base);
}
assert_eq!(frozen.without_tail().layout(), frozen.layout());
unsafe { word.release_unique(ptr.add(32), 0) };
}
#[test]
fn test_frozen_tail_overflow_allocates_record() {
let capacity = (1 << FORMAT.tail_bits) + 1;
let (base, ptr, word) = frozen_view(capacity, 0, 1);
let owner = record(word);
assert_eq!(ptr, base);
assert_eq!(heap(owner), (base, layout(capacity, 8)));
assert_eq!(unsafe { owner.refcount() }, Some(1));
unsafe { word.release_unique(ptr, 1) };
}
#[test]
fn test_vec_adoption_keeps_allocation_and_capacity() {
let mut vec = Vec::with_capacity(256);
vec.extend_from_slice(&[1u8, 2, 3, 4]);
let base_addr = vec.as_ptr() as usize;
let vec_cap = vec.capacity();
let (ptr, len, cap, word) = Owner::adopt_vec(vec);
assert_eq!(ptr.as_ptr() as usize, base_addr);
assert_eq!(len, 4);
assert_eq!(cap, vec_cap);
assert_eq!(
descriptor(word).layout(),
Layout::array::<u8>(vec_cap).unwrap()
);
let payload = unsafe { std::slice::from_raw_parts(ptr.as_ptr(), len) };
assert_eq!(payload, &[1, 2, 3, 4]);
unsafe { word.release_unique(ptr, cap) };
}
#[test]
fn test_vec_adoption_exact_size_is_native() {
let vec = vec![5u8, 6, 7];
assert_eq!(vec.len(), vec.capacity());
let base_addr = vec.as_ptr() as usize;
let (ptr, len, cap, word) = Owner::adopt_vec(vec);
assert_eq!(ptr.as_ptr() as usize, base_addr);
assert_eq!(len, 3);
assert_eq!(cap, 3);
assert_eq!(descriptor(word).layout().size(), 3);
let payload = unsafe { std::slice::from_raw_parts(ptr.as_ptr(), len) };
assert_eq!(payload, &[5, 6, 7]);
unsafe { word.release_unique(ptr, cap) };
}
#[test]
fn test_adopt_vec_with_capacity_but_no_length() {
let vec: Vec<u8> = Vec::with_capacity(64);
let base_addr = vec.as_ptr() as usize;
let (ptr, len, cap, word) = Owner::adopt_vec(vec);
assert_eq!(ptr.as_ptr() as usize, base_addr);
assert_eq!(len, 0);
assert_eq!(cap, 64);
unsafe { word.release_unique(ptr, cap) };
}
#[test]
#[should_panic(expected = "capacity must be greater than zero")]
fn test_adopt_vec_rejects_unallocated_vec() {
let _ = Owner::adopt_vec(Vec::new());
}
#[test]
fn test_external_owner_refcount() {
let owner = Owner::new(Storage::External(Bytes::from_static(b"abc")));
assert!(is_external(owner));
assert_eq!(unsafe { owner.refcount() }, Some(1));
unsafe { owner.clone_shared() };
assert_eq!(unsafe { owner.refcount() }, Some(2));
unsafe { owner.drop_shared() };
assert_eq!(unsafe { owner.refcount() }, Some(1));
unsafe { owner.drop_shared() };
}
#[test]
fn test_external_owner_keeps_inner_bytes_alive() {
let payload = Bytes::from(vec![7u8; 32]);
let inner_ptr = payload.as_ptr();
let owner = Owner::new(Storage::External(payload));
let Some(Storage::External(inner)) = (unsafe { owner.storage() }) else {
panic!("owner must be external");
};
assert_eq!(inner.as_ptr(), inner_ptr);
assert_eq!(inner.as_ref(), &[7u8; 32]);
unsafe { owner.drop_shared() };
}
#[test]
fn test_empty_owner_has_no_record() {
let owner = Owner::empty();
assert!(owner.is_empty());
assert!(unsafe { owner.storage() }.is_none());
assert_eq!(unsafe { owner.refcount() }, None);
unsafe { owner.release_record() };
unsafe { owner.drop_shared() };
unsafe { Owner::empty().release_unique(NonNull::dangling(), 0) };
}
#[test]
fn test_heap_owner_shared_clone_and_drop() {
let (_, owner) = native_record(64, 1);
unsafe { owner.clone_shared() };
assert_eq!(unsafe { owner.refcount() }, Some(2));
assert!(!unsafe { owner.is_unique() });
unsafe { owner.drop_shared() };
assert!(unsafe { owner.is_unique() });
unsafe { owner.drop_shared() };
}
#[test]
fn test_shared_release_panics_before_releasing_storage() {
let released = Arc::new(AtomicUsize::new(0));
let bytes = Bytes::from_owner(Tracker(released.clone()));
let owner = Owner::new(Storage::External(bytes));
unsafe { owner.refs() }.store(2, Ordering::Relaxed);
let result = panic::catch_unwind(AssertUnwindSafe(|| {
unsafe { owner.release_record() }
}));
assert!(result.is_err());
assert_eq!(released.load(Ordering::SeqCst), 0);
unsafe { owner.refs() }.store(1, Ordering::Relaxed);
unsafe { owner.drop_shared() };
assert_eq!(released.load(Ordering::SeqCst), 1);
}
#[test]
fn test_shared_promotion_publishes_two_references() {
let (base, ptr, word) = frozen_view(64, 8, 40);
let cell = OwnerCell::new(word);
let observed = descriptor(cell.load());
let view = unsafe { cell.promote(observed, ptr, 40) };
let installed = record(cell.load());
assert_eq!(record_addr(installed), record_addr(view));
assert_eq!(heap(view), (base, layout(64, 8)));
assert_eq!(unsafe { view.refcount() }, Some(2));
unsafe { view.drop_shared() };
let mut cell = cell;
unsafe { cell.get().release_view(ptr, 40) };
}
#[test]
fn test_losing_promotion_retains_winner() {
let (_, ptr, word) = frozen_view(64, 0, 64);
let cell = OwnerCell::new(word);
let observed = descriptor(cell.load());
let winner = unsafe { cell.promote(observed, ptr, 64) };
let loser = unsafe { cell.promote(observed, ptr, 64) };
assert_eq!(record_addr(loser), record_addr(winner));
assert_eq!(unsafe { winner.refcount() }, Some(3));
unsafe { winner.drop_shared() };
unsafe { loser.drop_shared() };
let mut cell = cell;
unsafe { cell.get().release_view(ptr, 64) };
}
#[test]
fn test_exclusive_promotion_stores_record() {
let (base, ptr, word) = frozen_view(64, 4, 20);
let mut cell = OwnerCell::new(word);
let first = unsafe { cell.retain_mut(ptr, 20) };
assert_eq!(record_addr(record(cell.get())), record_addr(first));
assert_eq!(heap(first), (base, layout(64, 8)));
assert_eq!(unsafe { first.refcount() }, Some(2));
let second = unsafe { cell.retain_mut(ptr, 20) };
assert_eq!(record_addr(second), record_addr(first));
assert_eq!(unsafe { first.refcount() }, Some(3));
unsafe { first.drop_shared() };
unsafe { second.drop_shared() };
unsafe { cell.get().release_view(ptr, 20) };
}
}
#[cfg(all(test, feature = "loom"))]
mod loom_tests {
use super::{
super::{IoBuf, IoBufMut},
*,
};
use bytes::BufMut as _;
use loom::{
cell::UnsafeCell,
sync::{Arc, atomic::AtomicUsize},
thread,
};
struct Tracker(Arc<AtomicUsize>);
impl Drop for Tracker {
fn drop(&mut self) {
self.0.fetch_add(1, Ordering::SeqCst);
}
}
impl AsRef<[u8]> for Tracker {
fn as_ref(&self) -> &[u8] {
&[1, 2, 3]
}
}
struct PayloadCell(UnsafeCell<usize>);
unsafe impl Send for PayloadCell {}
unsafe impl Sync for PayloadCell {}
struct TrackedPayload {
payload: Arc<PayloadCell>,
released: Arc<AtomicUsize>,
}
impl Drop for TrackedPayload {
fn drop(&mut self) {
self.payload.0.with(|cell| {
assert_eq!(unsafe { *cell }, 1);
});
self.released.fetch_add(1, Ordering::SeqCst);
}
}
impl AsRef<[u8]> for TrackedPayload {
fn as_ref(&self) -> &[u8] {
&[1, 2, 3]
}
}
fn unshared() -> IoBuf {
let mut buf = IoBufMut::with_capacity(12);
buf.put_bytes(7, 8);
buf.freeze()
}
#[test]
fn shared_clone_drop_releases_exactly_once() {
loom::model(|| {
let released = Arc::new(AtomicUsize::new(0));
let bytes = Bytes::from_owner(Tracker(released.clone()));
let owner = Owner::new(Storage::External(bytes));
unsafe { owner.clone_shared() };
unsafe { owner.clone_shared() };
let t1 = thread::spawn(move || {
unsafe { owner.drop_shared() };
});
let t2 = thread::spawn(move || {
unsafe { owner.drop_shared() };
});
unsafe { owner.drop_shared() };
t1.join().unwrap();
t2.join().unwrap();
assert_eq!(released.load(Ordering::SeqCst), 1);
});
}
#[test]
fn clone_races_concurrent_drop() {
loom::model(|| {
let released = Arc::new(AtomicUsize::new(0));
let bytes = Bytes::from_owner(Tracker(released.clone()));
let owner = Owner::new(Storage::External(bytes));
unsafe { owner.clone_shared() };
let t1 = thread::spawn(move || {
unsafe { owner.clone_shared() };
unsafe { owner.drop_shared() };
unsafe { owner.drop_shared() };
});
unsafe { owner.drop_shared() };
t1.join().unwrap();
assert_eq!(released.load(Ordering::SeqCst), 1);
});
}
#[test]
fn is_unique_races_final_drop() {
loom::model(|| {
let released = Arc::new(AtomicUsize::new(0));
let bytes = Bytes::from_owner(Tracker(released.clone()));
let owner = Owner::new(Storage::External(bytes));
let payload = Arc::new(UnsafeCell::new(0usize));
unsafe { owner.clone_shared() };
let t1 = thread::spawn({
let payload = payload.clone();
move || {
payload.with_mut(|cell| unsafe { *cell = 1 });
unsafe { owner.drop_shared() };
}
});
loop {
if unsafe { owner.is_unique() } {
payload.with(|cell| assert_eq!(unsafe { *cell }, 1));
assert_eq!(released.load(Ordering::SeqCst), 0);
break;
}
thread::yield_now();
}
unsafe { owner.drop_shared() };
t1.join().unwrap();
assert_eq!(released.load(Ordering::SeqCst), 1);
});
}
#[test]
fn payload_writes_happen_before_race_final_release() {
loom::model(|| {
let released = Arc::new(AtomicUsize::new(0));
let payload = Arc::new(PayloadCell(UnsafeCell::new(0)));
let bytes = Bytes::from_owner(TrackedPayload {
payload: payload.clone(),
released: released.clone(),
});
let owner = Owner::new(Storage::External(bytes));
unsafe { owner.clone_shared() };
let t1 = thread::spawn({
move || {
payload.0.with_mut(|cell| unsafe { *cell = 1 });
unsafe { owner.drop_shared() };
}
});
unsafe { owner.drop_shared() };
t1.join().unwrap();
assert_eq!(released.load(Ordering::SeqCst), 1);
});
}
#[test]
fn payload_writes_happen_before_fast_path_release() {
loom::model(|| {
let released = Arc::new(AtomicUsize::new(0));
let payload = Arc::new(PayloadCell(UnsafeCell::new(0)));
let bytes = Bytes::from_owner(TrackedPayload {
payload: payload.clone(),
released: released.clone(),
});
let owner = Owner::new(Storage::External(bytes));
unsafe { owner.clone_shared() };
let t1 = thread::spawn({
move || {
payload.0.with_mut(|cell| unsafe { *cell = 1 });
unsafe { owner.drop_shared() };
}
});
while unsafe { owner.refcount_relaxed() } != 1 {
thread::yield_now();
}
unsafe { owner.drop_shared() };
t1.join().unwrap();
assert_eq!(released.load(Ordering::SeqCst), 1);
});
}
#[test]
fn promotion_publishes_record_to_observer() {
loom::model(|| {
let buf = Arc::new(unshared());
let promoter = thread::spawn({
let buf = buf.clone();
move || drop(IoBuf::clone(&buf))
});
while !buf.owner_cell().holds_record_relaxed() {
thread::yield_now();
}
let view = IoBuf::clone(&buf);
assert_eq!(view, [7u8; 8]);
drop(view);
promoter.join().unwrap();
let buf = Arc::try_unwrap(buf).unwrap();
let recovered = buf.try_into_mut().expect("the views are gone");
assert_eq!(recovered.capacity(), 12);
});
}
#[test]
fn losing_promotion_retains_winner() {
loom::model(|| {
let buf = Arc::new(unshared());
let a = thread::spawn({
let buf = buf.clone();
move || drop(IoBuf::clone(&buf))
});
let b = thread::spawn({
let buf = buf.clone();
move || {
let view = buf.slice(1..7);
let again = view.clone();
assert_eq!(again, [7u8; 6]);
drop(view);
drop(again);
}
});
a.join().unwrap();
b.join().unwrap();
let buf = Arc::try_unwrap(buf).unwrap();
assert!(buf.owner_cell().holds_record_relaxed());
let recovered = buf.try_into_mut().expect("the views are gone");
assert_eq!(recovered.capacity(), 12);
});
}
#[test]
fn promoted_try_into_mut_observes_prior_accesses() {
loom::model(|| {
let buf = unshared();
let view = buf.clone();
let payload = Arc::new(PayloadCell(UnsafeCell::new(0)));
let t1 = thread::spawn({
let payload = payload.clone();
move || {
payload.0.with_mut(|cell| unsafe { *cell = 1 });
drop(view);
}
});
while unsafe { buf.owner_cell().refcount_relaxed() } != 1 {
thread::yield_now();
}
let recovered = buf.try_into_mut().expect("the view is gone");
payload.0.with(|cell| assert_eq!(unsafe { *cell }, 1));
drop(recovered);
t1.join().unwrap();
});
}
#[test]
fn promoted_fast_path_release_observes_prior_accesses() {
loom::model(|| {
let buf = unshared();
let view = buf.slice(2..6);
let t1 = thread::spawn(move || drop(view));
while unsafe { buf.owner_cell().refcount_relaxed() } != 1 {
thread::yield_now();
}
drop(buf);
t1.join().unwrap();
});
}
#[test]
fn promoted_concurrent_drops_release_once() {
loom::model(|| {
let buf = unshared();
let view = buf.clone();
let t1 = thread::spawn(move || drop(view));
drop(buf);
t1.join().unwrap();
});
}
#[test]
fn exclusive_promotion_shares_record_across_threads() {
loom::model(|| {
let mut buf = unshared();
let prefix = buf.split_to(3);
let t1 = thread::spawn(move || {
assert_eq!(prefix, [7u8; 3]);
drop(prefix);
});
let rest = buf.clone();
assert_eq!(rest, [7u8; 5]);
drop(buf);
drop(rest);
t1.join().unwrap();
});
}
}