From 264a81589492a997f2bad68ea8d77e41d8f376e1 Mon Sep 17 00:00:00 2001
From: Thomas Orozco <torozco@fb.com>
Date: Wed, 8 Jan 2020 03:57:56 -0800
Subject: [PATCH] mononoke: allocation_tracing: support realloc + handle errors

Summary: This makes the implementation a little more comprehensive.

Reviewed By: mitrandir77

Differential Revision: D19308863

fbshipit-source-id: b265ea645e390a9638eeb66301f69c2211b0cfa4
---
 common/allocation_tracing/src/core.rs | 68 +++++++++++++++++++++++++--
 1 file changed, 63 insertions(+), 5 deletions(-)

diff --git a/common/allocation_tracing/src/core.rs b/common/allocation_tracing/src/core.rs
index de1b3592..26d7ab81 100644
--- a/common/allocation_tracing/src/core.rs
+++ b/common/allocation_tracing/src/core.rs
@@ -27,19 +27,50 @@ struct TracingAllocator;
 // if it did allocate, any allocation with TracingAllocator would hang forever.
 unsafe impl GlobalAlloc for TracingAllocator {
     unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        ALLOCATION_STATS.with(|cell| {
-            (*cell.get()).allocated += layout.size() as u64;
-        });
+        let ret = System.alloc(layout);
+
+        if !ret.is_null() {
+            ALLOCATION_STATS.with(|cell| {
+                (*cell.get()).allocated += layout.size() as u64;
+            });
+        }
+
+        ret
+    }
+
+    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+        let ret = System.alloc_zeroed(layout);
+
+        if !ret.is_null() {
+            ALLOCATION_STATS.with(|cell| {
+                (*cell.get()).allocated += layout.size() as u64;
+            });
+        }
+
+        ret
+    }
+
+    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+        let ret = System.realloc(ptr, layout, new_size);
+
+        if !ret.is_null() {
+            ALLOCATION_STATS.with(|cell| {
+                (*cell.get()).freed += layout.size() as u64;
+                (*cell.get()).allocated += new_size as u64;
+            });
+        }
 
-        System.alloc(layout)
+        ret
     }
 
     unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+        let ret = System.dealloc(ptr, layout);
+
         ALLOCATION_STATS.with(|cell| {
             (*cell.get()).freed += layout.size() as u64;
         });
 
-        System.dealloc(ptr, layout)
+        ret
     }
 }
 
@@ -179,4 +210,31 @@ mod test {
         run_test::<TestStruct5>()?;
         Ok(())
     }
+
+    #[test]
+    fn test_realloc() -> Result<(), Error> {
+        fn run_test<T: Default>() -> Result<(), Error> {
+            let (_, stats) = trace_allocations(|| {
+                let mut v = Vec::<T>::with_capacity(2);
+                v.push(T::default());
+                v.push(T::default());
+                v.push(T::default());
+                v.push(T::default());
+                v.push(T::default());
+                v.pop();
+                v.shrink_to_fit();
+                v
+            });
+
+            assert_eq!(stats.delta()?, (mem::size_of::<T>() * 4) as i64);
+            Ok(())
+        }
+
+        run_test::<TestStruct1>()?;
+        run_test::<TestStruct2>()?;
+        run_test::<TestStruct3>()?;
+        run_test::<TestStruct4>()?;
+        run_test::<TestStruct5>()?;
+        Ok(())
+    }
 }