tor-browser

frame_allocator.rs (16184B)

---

/* This Source Code Form is subject to the terms of the Mozilla Public
* License, v. 2.0. If a copy of the MPL was not distributed with this
* file, You can obtain one at http://mozilla.org/MPL/2.0/. */

///! A custom allocator for memory allocations that have the lifetime of a frame.
///!
///! See also `internal_types::FrameVec`.
///!

use allocator_api2::alloc::{Allocator, AllocError, Layout, Global};

use std::{cell::UnsafeCell, ptr::NonNull, sync::{atomic::{AtomicI32, Ordering}, Arc}};

use crate::{bump_allocator::{BumpAllocator, ChunkPool, Stats}, internal_types::{FrameId, FrameVec}};

/// A memory allocator for allocations that have the same lifetime as a built frame.
///
/// A custom allocator is used because:
/// - The frame is created on a thread and dropped on another thread, which causes
///   lock contention in jemalloc.
/// - Since all allocations have a very similar lifetime, we can implement much faster
///   allocation and deallocation with a specialized allocator than can be achieved
///   with a general purpose allocator.
///
/// If the allocator is created using `FrameAllocator::fallback()`, it is not
/// attached to a `FrameMemory` and simply falls back to the global allocator. This
/// should only be used to handle deserialization (for wrench replays) and tests.
///
/// # Safety
///
/// None of the safety restrictions below apply if the allocator is created using
/// `FrameAllocator::fallback`.
///
/// `FrameAllocator` can move between thread if and only if it does so along with
/// the `FrameMemory` it is associated to (if any). The opposite is also true: it
/// is safe to move `FrameMemory` between threads if and only if all live frame
/// allocators associated to it move along with it.
///
/// `FrameAllocator` must be dropped before the `FrameMemory` it is associated to.
///
/// In other words, `FrameAllocator` should only be used for containers that are
/// in the `Frame` data structure and not stored elsewhere. The `Frame` holds on
/// to its `FrameMemory`, allowing it all to be sent from the frame builder thread
/// to the renderer thread together.
///
/// Another way to think of it is that the frame is a large self-referential data
/// structure, holding on to its memory and a large number of containers that
/// point into the memory.
pub struct FrameAllocator {
   // If this pointer is null, fall back to the global allocator.
   inner: *mut FrameInnerAllocator,

   #[cfg(debug_assertions)]
   frame_id: Option<FrameId>,
}

impl FrameAllocator {
   /// Creates a `FrameAllocator` that defaults to the global allocator.
   ///
   /// Should only be used for testing purposes or desrialization in wrench replays.
   pub fn fallback() -> Self {
       FrameAllocator {
           inner: std::ptr::null_mut(),
           #[cfg(debug_assertions)]
           frame_id: None,
       }
   }

   /// Shorthand for creating a FrameVec.
   #[inline]
   pub fn new_vec<T>(self) -> FrameVec<T> {
       FrameVec::new_in(self)
   }

   /// Shorthand for creating a FrameVec.
   #[inline]
   pub fn new_vec_with_capacity<T>(self, cap: usize) -> FrameVec<T> {
       FrameVec::with_capacity_in(cap, self)
   }

   #[inline]
   fn allocate_impl(mem: *mut FrameInnerAllocator, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
       unsafe {
           (*mem).live_alloc_count.fetch_add(1, Ordering::Relaxed);
           (*mem).bump.allocate_item(layout)
       }
   }

   #[inline]
   unsafe fn deallocate_impl(mem: *mut FrameInnerAllocator, ptr: NonNull<u8>, layout: Layout) {
       (*mem).live_alloc_count.fetch_sub(1, Ordering::Relaxed);
       (*mem).bump.deallocate_item(ptr, layout)
   }

   #[inline]
   unsafe fn grow_impl(mem: *mut FrameInnerAllocator, ptr: NonNull<u8>, old_layout: Layout, new_layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
       (*mem).bump.grow_item(ptr, old_layout, new_layout)
   }

   #[inline]
   unsafe fn shrink_impl(mem: *mut FrameInnerAllocator, ptr: NonNull<u8>, old_layout: Layout, new_layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
       (*mem).bump.shrink_item(ptr, old_layout, new_layout)
   }

   #[cold]
   #[inline(never)]
   fn allocate_fallback(layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
       Global.allocate(layout)
   }

   #[cold]
   #[inline(never)]
   fn deallocate_fallback(ptr: NonNull<u8>, layout: Layout) {
       unsafe { Global.deallocate(ptr, layout) }
   }

   #[cold]
   #[inline(never)]
   fn grow_fallback(ptr: NonNull<u8>, old_layout: Layout, new_layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
       unsafe { Global.grow(ptr, old_layout, new_layout) }
   }

   #[cfg(not(debug_assertions))]
   fn check_frame_id(&self) {}

   #[cfg(debug_assertions)]
   fn check_frame_id(&self) {
       if self.inner.is_null() {
           return;
       }
       unsafe {
           assert_eq!(self.frame_id, (*self.inner).frame_id);
       }
   }
}

impl Clone for FrameAllocator {
   fn clone(&self) -> Self {
       unsafe {
           if let Some(inner) = self.inner.as_mut() {
               // When cloning a `FrameAllocator`, we have to decrement the
               // counter of dropped references in the inner allocator to
               // balance the fact that an extra `FrameAllocator` will be
               // dropped (that hasn't been accounted in `FrameMemory`).
               inner.references_dropped.fetch_sub(1, Ordering::Relaxed);
           }
       }

       FrameAllocator {
           inner: self.inner,
           #[cfg(debug_assertions)]
           frame_id: self.frame_id,
       }
   }
}

impl Drop for FrameAllocator {
   fn drop(&mut self) {
       unsafe {
           if let Some(inner) = self.inner.as_mut() {
               inner.references_dropped.fetch_add(1, Ordering::Relaxed);
           }
       }
   }
}

unsafe impl Send for FrameAllocator {}

unsafe impl Allocator for FrameAllocator {
   #[inline(never)]
   fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
       if self.inner.is_null() {
           return FrameAllocator::allocate_fallback(layout);
       }

       self.check_frame_id();

       FrameAllocator::allocate_impl(self.inner, layout)
   }

   #[inline(never)]
   unsafe fn deallocate(&self, ptr: NonNull<u8>, layout: Layout) {
       if self.inner.is_null() {
           return FrameAllocator::deallocate_fallback(ptr, layout);
       }

       self.check_frame_id();

       FrameAllocator::deallocate_impl(self.inner, ptr, layout)
   }

   #[inline(never)]
   unsafe fn grow(
       &self,
       ptr: NonNull<u8>,
       old_layout: Layout,
       new_layout: Layout
   ) -> Result<NonNull<[u8]>, AllocError> {
       if self.inner.is_null() {
           return FrameAllocator::grow_fallback(ptr, old_layout, new_layout);
       }

       self.check_frame_id();

       FrameAllocator::grow_impl(self.inner, ptr, old_layout, new_layout)
   }

   #[inline(never)]
   unsafe fn shrink(
       &self,
       ptr: NonNull<u8>,
       old_layout: Layout,
       new_layout: Layout
   ) -> Result<NonNull<[u8]>, AllocError> {
       if self.inner.is_null() {
           return FrameAllocator::grow_fallback(ptr, old_layout, new_layout);
       }

       self.check_frame_id();

       FrameAllocator::shrink_impl(self.inner, ptr, old_layout, new_layout)
   }
}

#[cfg(feature = "capture")]
impl serde::Serialize for FrameAllocator {
   fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
   where S: serde::Serializer
   {
       ().serialize(serializer)
   }
}

#[cfg(feature = "replay")]
impl<'de> serde::Deserialize<'de> for FrameAllocator {
   fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
   where
       D: serde::Deserializer<'de>,
   {
       let _ = <() as serde::Deserialize>::deserialize(deserializer)?;
       Ok(FrameAllocator::fallback())
   }
}

/// The default impl is required for Deserialize to work in FrameVec.
/// It's fine to fallback to the global allocator when replaying wrench
/// recording but we don't want to accidentally use `FrameAllocator::default()`
/// in regular webrender usage, so we only implement it when the replay
/// feature is enabled.
#[cfg(feature = "replay")]
impl Default for FrameAllocator {
   fn default() -> Self {
       Self::fallback()
   }
}

/// The backing storage for `FrameAllocator`
///
/// This object is meant to be stored in the built frame and must not be dropped or
/// recycled before all allocations have been deallocated and all `FrameAllocators`
/// have been dropped. In other words, drop or recycle this after dropping the rest
/// of the built frame.
pub struct FrameMemory {
   // Box would be nice but it is not adequate for this purpose because
   // it is "no-alias". So we do it the hard way and manage this pointer
   // manually.

   /// Safety: The pointed `FrameInnerAllocator` must not move or be deallocated
   /// while there are live `FrameAllocator`s pointing to it. This is ensured
   /// by respecting that the `FrameMemory` is dropped last and by the
   /// `FrameInnerAllocator` not being exposed to the outside world.
   /// It is also checked at runtime via the reference count.
   allocator: Option<NonNull<FrameInnerAllocator>>,
   /// The number of `FrameAllocator`s created during the current frame. This is
   /// used to compare aganst the inner allocator's dropped references counter
   /// to check that references have all been dropped before freeing or recycling
   /// the memory.
   references_created: UnsafeCell<i32>,
}

impl FrameMemory {
   /// Creates a fallback FrameMemory that uses the global allocator.
   ///
   /// This should only be used for testing purposes and to handle the
   /// deserialization of webrender recordings.
   #[allow(unused)]
   pub fn fallback() -> Self {
       FrameMemory {
           allocator: None,
           references_created: UnsafeCell::new(0)
       }
   }

   /// # Panics
   ///
   /// A `FrameMemory` must not be dropped until all of the associated
   /// `FrameAllocators` as well as their allocations have been dropped,
   /// otherwise the `FrameMemory::drop` will panic.
   pub fn new(pool: Arc<ChunkPool>, _frame_id: FrameId) -> Self {
       let layout = Layout::from_size_align(
           std::mem::size_of::<FrameInnerAllocator>(),
           std::mem::align_of::<FrameInnerAllocator>(),
       ).unwrap();

       let uninit_u8 = Global.allocate(layout).unwrap();

       unsafe {
           let allocator: NonNull<FrameInnerAllocator> = uninit_u8.cast();
           allocator.as_ptr().write(FrameInnerAllocator {
               bump: BumpAllocator::new(pool),

               live_alloc_count: AtomicI32::new(0),
               references_dropped: AtomicI32::new(0),
               #[cfg(debug_assertions)]
               frame_id: Some(_frame_id),
           });

           FrameMemory {
               allocator: Some(allocator),
               references_created: UnsafeCell::new(0),
           }
       }
   }

   /// Create a `FrameAllocator` for the current frame.
   pub fn allocator(&self) -> FrameAllocator {
       if let Some(alloc) = &self.allocator {
           unsafe { *self.references_created.get() += 1 };

           return FrameAllocator {
               inner: alloc.as_ptr(),
               #[cfg(debug_assertions)]
               frame_id: unsafe { alloc.as_ref().frame_id },
           };
       }

       FrameAllocator::fallback()
   }

   /// Shorthand for creating a FrameVec.
   #[inline]
   pub fn new_vec<T>(&self) -> FrameVec<T> {
       FrameVec::new_in(self.allocator())
   }

   /// Shorthand for creating a FrameVec.
   #[inline]
   pub fn new_vec_with_capacity<T>(&self, cap: usize) -> FrameVec<T> {
       FrameVec::with_capacity_in(cap, self.allocator())
   }

   /// Panics if there are still live allocations or `FrameAllocator`s.
   pub fn assert_memory_reusable(&self) {
       if let Some(ptr) = self.allocator {
           unsafe {
               // If this assert blows up, it means an allocation is still alive.
               assert_eq!(ptr.as_ref().live_alloc_count.load(Ordering::Acquire), 0);
               // If this assert blows up, it means one or several FrameAllocators
               // from the previous frame are still alive.
               let references_created = *self.references_created.get();
               assert_eq!(ptr.as_ref().references_dropped.load(Ordering::Acquire), references_created);
           }
       }
   }

   #[allow(unused)]
   pub fn get_stats(&self) -> Stats {
       unsafe {
           self.allocator.map(|ptr| (*ptr.as_ptr()).bump.get_stats()).unwrap_or_else(Stats::default)
       }
   }
}

impl Drop for FrameMemory {
   fn drop(&mut self) {
       self.assert_memory_reusable();

       let layout = Layout::new::<FrameInnerAllocator>();

       unsafe {
           if let Some(ptr) = &mut self.allocator {
               std::ptr::drop_in_place(ptr.as_ptr());
               Global.deallocate(ptr.cast(), layout);
           }
       }
   }
}

unsafe impl Send for FrameMemory {}

#[cfg(feature = "capture")]
impl serde::Serialize for FrameMemory {
   fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
   where S: serde::Serializer
   {
       ().serialize(serializer)
   }
}

#[cfg(feature = "replay")]
impl<'de> serde::Deserialize<'de> for FrameMemory {
   fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
   where
       D: serde::Deserializer<'de>,
   {
       let _ = <() as serde::Deserialize>::deserialize(deserializer)?;
       Ok(FrameMemory::fallback())
   }
}

struct FrameInnerAllocator {
   bump: BumpAllocator,

   // Strictly speaking the live allocation and reference count do not need to
   // be atomic if the allocator is used correctly (the thread that
   // allocates/deallocates is also the thread where the allocator is).
   // Since the point of keeping track of the number of live allocations is to
   // check that the allocator is indeed used correctly, we stay on the safe
   // side for now.

   live_alloc_count: AtomicI32,
   /// We count the number of references dropped here and compare it against the
   /// number of references created by the `AllocatorMemory` when we need to check
   /// that the memory can be safely reused or released.
   /// This looks and is very similar to a reference counting scheme (`Arc`). The
   /// main differences are that we don't want the reference count to drive the
   /// lifetime of the allocator (only to check when we require all references to
   /// have been dropped), and we do half as many the atomic operations since we only
   /// count drops and not creations.
   references_dropped: AtomicI32,
   #[cfg(debug_assertions)]
   frame_id: Option<FrameId>,
}

#[test]
fn frame_memory_simple() {
   use std::sync::mpsc::channel;

   let chunk_pool = Arc::new(ChunkPool::new());
   let memory = FrameMemory::new(chunk_pool, FrameId::first());

   let alloc = memory.allocator();
   let a2 = memory.allocator();
   let a3 = memory.allocator();
   let v1: FrameVec<u32> = memory.new_vec_with_capacity(10);
   let v2: FrameVec<u32> = memory.new_vec_with_capacity(256);
   let v3: FrameVec<u32> = memory.new_vec_with_capacity(1024 * 128);
   let v4: FrameVec<u32> = memory.new_vec_with_capacity(128);
   let mut v5 = alloc.clone().new_vec();
   for i in 0..256u32 {
       v5.push(i);
   }
   let v6 = v2.clone().clone().clone().clone();
   let mut frame = alloc.new_vec();
   frame.push(v1);
   frame.push(v2);
   frame.push(v3);
   frame.push(v4);
   frame.push(v5);
   frame.push(v6);
   let (tx, rx) = channel();
   tx.send(frame).unwrap();

   let handle = std::thread::spawn(move || {
       let mut frame = rx.recv().unwrap();
       frame.push(memory.new_vec_with_capacity(10));
       std::mem::drop(a3);
       std::mem::drop(a2);
       std::mem::drop(frame);

       memory.assert_memory_reusable();
   });

   handle.join().unwrap();
}