fix counting of reallocations

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() {

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
     }