From e5c812fbcbad552e01597a87db82bb0db65781e9 Mon Sep 17 00:00:00 2001 From: Matin Aniss <76515905+MatinAniss@users.noreply.github.com> Date: Thu, 26 Jun 2025 05:29:13 +1000 Subject: [PATCH] gpui: Dynamic element arena (#32079) Implements a chunking strategy for the element arena that allows it to grow dynamically based on allocations, it is initialised with a single chunk of a total size of 1 mebibyte. On allocation of data with a size greater than the remaining space of the current chunk a new chunk is created. This reduces the memory allocation from the static 32 mebibytes, this especially helps GPUI applications that don't need such a large element arena and even Zed in most cases. This also prevents the panic when allocations ever exceed the element arena. Release Notes: - N/A --------- Co-authored-by: Michael Sloan --- crates/gpui/src/arena.rs | 143 ++++++++++++++++++++++++++++---------- crates/gpui/src/window.rs | 9 +-- 2 files changed, 109 insertions(+), 43 deletions(-) diff --git a/crates/gpui/src/arena.rs b/crates/gpui/src/arena.rs index f30f4b6480..2448746a88 100644 --- a/crates/gpui/src/arena.rs +++ b/crates/gpui/src/arena.rs @@ -1,5 +1,5 @@ use std::{ - alloc, + alloc::{self, handle_alloc_error}, cell::Cell, ops::{Deref, DerefMut}, ptr, @@ -20,43 +20,98 @@ impl Drop for ArenaElement { } } -pub struct Arena { +struct Chunk { start: *mut u8, end: *mut u8, offset: *mut u8, - elements: Vec, - valid: Rc>, } -impl Arena { - pub fn new(size_in_bytes: usize) -> Self { +impl Drop for Chunk { + fn drop(&mut self) { unsafe { - let layout = alloc::Layout::from_size_align(size_in_bytes, 1).unwrap(); + let chunk_size = self.end.offset_from_unsigned(self.start); + // this never fails as it succeeded during allocation + let layout = alloc::Layout::from_size_align(chunk_size, 1).unwrap(); + alloc::dealloc(self.start, layout); + } + } +} + +impl Chunk { + fn new(chunk_size: usize) -> Self { + unsafe { + // this only fails if chunk_size is unreasonably huge + let layout = alloc::Layout::from_size_align(chunk_size, 1).unwrap(); let start = alloc::alloc(layout); - let end = start.add(size_in_bytes); + if start.is_null() { + handle_alloc_error(layout); + } + let end = start.add(chunk_size); Self { start, end, offset: start, - elements: Vec::new(), - valid: Rc::new(Cell::new(true)), } } } - pub fn len(&self) -> usize { - self.offset as usize - self.start as usize + fn allocate(&mut self, layout: alloc::Layout) -> Option<*mut u8> { + unsafe { + let aligned = self.offset.add(self.offset.align_offset(layout.align())); + let next = aligned.add(layout.size()); + + if next <= self.end { + self.offset = next; + Some(aligned) + } else { + None + } + } + } + + fn reset(&mut self) { + self.offset = self.start; + } +} + +pub struct Arena { + chunks: Vec, + elements: Vec, + valid: Rc>, + current_chunk_index: usize, + chunk_size: usize, +} + +impl Drop for Arena { + fn drop(&mut self) { + self.clear(); + } +} + +impl Arena { + pub fn new(chunk_size: usize) -> Self { + assert!(chunk_size > 0); + Self { + chunks: vec![Chunk::new(chunk_size)], + elements: Vec::new(), + valid: Rc::new(Cell::new(true)), + current_chunk_index: 0, + chunk_size, + } } pub fn capacity(&self) -> usize { - self.end as usize - self.start as usize + self.chunks.len() * self.chunk_size } pub fn clear(&mut self) { self.valid.set(false); self.valid = Rc::new(Cell::new(true)); self.elements.clear(); - self.offset = self.start; + for chunk_index in 0..=self.current_chunk_index { + self.chunks[chunk_index].reset(); + } + self.current_chunk_index = 0; } #[inline(always)] @@ -79,33 +134,45 @@ impl Arena { unsafe { let layout = alloc::Layout::new::(); - let offset = self.offset.add(self.offset.align_offset(layout.align())); - let next_offset = offset.add(layout.size()); - assert!(next_offset <= self.end, "not enough space in Arena"); - - let result = ArenaBox { - ptr: offset.cast(), - valid: self.valid.clone(), + let mut current_chunk = &mut self.chunks[self.current_chunk_index]; + let ptr = if let Some(ptr) = current_chunk.allocate(layout) { + ptr + } else { + self.current_chunk_index += 1; + if self.current_chunk_index >= self.chunks.len() { + self.chunks.push(Chunk::new(self.chunk_size)); + assert_eq!(self.current_chunk_index, self.chunks.len() - 1); + log::info!( + "increased element arena capacity to {}kb", + self.capacity() / 1024, + ); + } + current_chunk = &mut self.chunks[self.current_chunk_index]; + if let Some(ptr) = current_chunk.allocate(layout) { + ptr + } else { + panic!( + "Arena chunk_size of {} is too small to allocate {} bytes", + self.chunk_size, + layout.size() + ); + } }; - inner_writer(result.ptr, f); + inner_writer(ptr.cast(), f); self.elements.push(ArenaElement { - value: offset, + value: ptr, drop: drop::, }); - self.offset = next_offset; - result + ArenaBox { + ptr: ptr.cast(), + valid: self.valid.clone(), + } } } } -impl Drop for Arena { - fn drop(&mut self) { - self.clear(); - } -} - pub struct ArenaBox { ptr: *mut T, valid: Rc>, @@ -215,13 +282,17 @@ mod tests { } #[test] - #[should_panic(expected = "not enough space in Arena")] - fn test_arena_overflow() { - let mut arena = Arena::new(16); + fn test_arena_grow() { + let mut arena = Arena::new(8); arena.alloc(|| 1u64); arena.alloc(|| 2u64); - // This should panic. - arena.alloc(|| 3u64); + + assert_eq!(arena.capacity(), 16); + + arena.alloc(|| 3u32); + arena.alloc(|| 4u32); + + assert_eq!(arena.capacity(), 24); } #[test] diff --git a/crates/gpui/src/window.rs b/crates/gpui/src/window.rs index 0e3f5763da..be3b753d6a 100644 --- a/crates/gpui/src/window.rs +++ b/crates/gpui/src/window.rs @@ -206,8 +206,7 @@ slotmap::new_key_type! { } thread_local! { - /// 8MB wasn't quite enough... - pub(crate) static ELEMENT_ARENA: RefCell = RefCell::new(Arena::new(32 * 1024 * 1024)); + pub(crate) static ELEMENT_ARENA: RefCell = RefCell::new(Arena::new(1024 * 1024)); } /// Returned when the element arena has been used and so must be cleared before the next draw. @@ -218,12 +217,8 @@ impl ArenaClearNeeded { /// Clear the element arena. pub fn clear(self) { ELEMENT_ARENA.with_borrow_mut(|element_arena| { - let percentage = (element_arena.len() as f32 / element_arena.capacity() as f32) * 100.; - if percentage >= 80. { - log::warn!("elevated element arena occupation: {}.", percentage); - } element_arena.clear(); - }) + }); } }