diff --git a/crates/turbo-tasks-malloc/src/counter.rs b/crates/turbo-tasks-malloc/src/counter.rs index e0765f51101d1..b3cc5b07219f9 100644 --- a/crates/turbo-tasks-malloc/src/counter.rs +++ b/crates/turbo-tasks-malloc/src/counter.rs @@ -49,6 +49,35 @@ impl ThreadLocalCounter { } } + fn update(&mut self, old_size: usize, new_size: usize) { + self.allocation_counters.deallocations += old_size; + self.allocation_counters.deallocation_count += 1; + self.allocation_counters.allocations += new_size; + self.allocation_counters.allocation_count += 1; + match old_size.cmp(&new_size) { + std::cmp::Ordering::Equal => {} + std::cmp::Ordering::Less => { + let size = new_size - old_size; + if self.buffer >= size { + self.buffer -= size; + } else { + let offset = size - self.buffer + TARGET_BUFFER; + self.buffer = TARGET_BUFFER; + ALLOCATED.fetch_add(offset, Ordering::Relaxed); + } + } + std::cmp::Ordering::Greater => { + let size = old_size - new_size; + self.buffer += size; + if self.buffer > MAX_BUFFER { + let offset = self.buffer - TARGET_BUFFER; + self.buffer = TARGET_BUFFER; + ALLOCATED.fetch_sub(offset, Ordering::Relaxed); + } + } + } + } + fn unload(&mut self) { if self.buffer > 0 { ALLOCATED.fetch_sub(self.buffer, Ordering::Relaxed); @@ -93,6 +122,11 @@ pub fn remove(size: usize) { with_local_counter(|local| local.remove(size)); } +/// Adds some `size` to the global counter in a thread-local buffered way. +pub fn update(old_size: usize, new_size: usize) { + with_local_counter(|local| local.update(old_size, new_size)); +} + /// Flushes the thread-local buffer to the global counter. This should be called /// e. g. when a thread is stopped or goes to sleep for a long time. pub fn flush() { diff --git a/crates/turbo-tasks-malloc/src/lib.rs b/crates/turbo-tasks-malloc/src/lib.rs index 5fa6cc4c7d7cf..f00648694bcac 100644 --- a/crates/turbo-tasks-malloc/src/lib.rs +++ b/crates/turbo-tasks-malloc/src/lib.rs @@ -5,7 +5,7 @@ use std::{ marker::PhantomData, }; -use self::counter::{add, flush, get, remove}; +use self::counter::{add, flush, get, remove, update}; #[derive(Default, Clone, Debug)] pub struct AllocationInfo { @@ -97,11 +97,7 @@ unsafe impl GlobalAlloc for TurboMalloc { let ret = mimalloc::MiMalloc.realloc(ptr, layout, new_size); if !ret.is_null() { let old_size = layout.size(); - if old_size < new_size { - add(new_size - old_size); - } else { - remove(old_size - new_size); - } + update(old_size, new_size); } ret } @@ -137,11 +133,7 @@ unsafe impl GlobalAlloc for TurboMalloc { let ret = std::alloc::System.realloc(ptr, layout, new_size); if !ret.is_null() { let old_size = layout.size(); - if old_size < new_size { - add(new_size - old_size); - } else { - remove(old_size - new_size); - } + update(old_size, new_size); } ret }