Assert no memory leaks by stubbing malloc/free in the test suite

cli/src/tests/allocations.rs

@@ -0,0 +1,104 @@
+#![cfg(test)]
+#![allow(dead_code)]
+
+use spin::Mutex;
+use std::collections::HashMap;
+use std::os::raw::{c_ulong, c_void};
+
+#[derive(Debug, PartialEq, Eq, Hash)]
+struct Allocation(*const c_void);
+unsafe impl Send for Allocation {}
+unsafe impl Sync for Allocation {}
+
+#[derive(Default)]
+struct AllocationRecorder {
+    enabled: bool,
+    allocation_count: u64,
+    outstanding_allocations: HashMap<Allocation, u64>,
+}
+
+lazy_static! {
+    static ref RECORDER: Mutex<AllocationRecorder> = Mutex::new(AllocationRecorder::default());
+}
+
+extern "C" {
+    fn malloc(size: c_ulong) -> *mut c_void;
+    fn calloc(count: c_ulong, size: c_ulong) -> *mut c_void;
+    fn realloc(ptr: *mut c_void, size: c_ulong) -> *mut c_void;
+    fn free(ptr: *mut c_void);
+}
+
+pub fn start_recording() {
+    let mut recorder = RECORDER.lock();
+    recorder.enabled = true;
+    recorder.allocation_count = 0;
+    recorder.outstanding_allocations.clear();
+}
+
+pub fn stop_recording() {
+    let mut recorder = RECORDER.lock();
+    recorder.enabled = false;
+
+    if !recorder.outstanding_allocations.is_empty() {
+        panic!(
+            "Leaked allocation indices: {:?}",
+            recorder
+                .outstanding_allocations
+                .iter()
+                .map(|e| e.1)
+                .collect::<Vec<_>>()
+        );
+    }
+}
+
+fn record_alloc(ptr: *mut c_void) {
+    let mut recorder = RECORDER.lock();
+    if recorder.enabled {
+        let count = recorder.allocation_count;
+        recorder.allocation_count += 1;
+        recorder
+            .outstanding_allocations
+            .insert(Allocation(ptr), count);
+    }
+}
+
+fn record_dealloc(ptr: *mut c_void) {
+    let mut recorder = RECORDER.lock();
+    if recorder.enabled {
+        recorder.outstanding_allocations.remove(&Allocation(ptr));
+    }
+}
+
+#[no_mangle]
+extern "C" fn ts_record_malloc(size: c_ulong) -> *const c_void {
+    let result = unsafe { malloc(size) };
+    record_alloc(result);
+    result
+}
+
+#[no_mangle]
+extern "C" fn ts_record_calloc(count: c_ulong, size: c_ulong) -> *const c_void {
+    let result = unsafe { calloc(count, size) };
+    record_alloc(result);
+    result
+}
+
+#[no_mangle]
+extern "C" fn ts_record_realloc(ptr: *mut c_void, size: c_ulong) -> *const c_void {
+    record_dealloc(ptr);
+    let result = unsafe { realloc(ptr, size) };
+    record_alloc(result);
+    result
+}
+
+#[no_mangle]
+extern "C" fn ts_record_free(ptr: *mut c_void) {
+    record_dealloc(ptr);
+    unsafe { free(ptr) };
+}
+
+#[no_mangle]
+extern "C" fn ts_record_allocations_toggle() {
+    let mut recorder = RECORDER.lock();
+    recorder.enabled = !recorder.enabled;
+}