Track specific object allocations

track_alloc.rs

use core::any::type_name;
use std::ops::{Deref, DerefMut};
use std::sync::Mutex;

use std::sync::atomic::{AtomicUsize, Ordering};

use std::collections::HashMap;

pub struct TrackingStats {
    pub instance_count: usize,
    pub instance_name: String,
}

lazy_static! {
    static ref TRACKING_STATS: Mutex<HashMap<String, AtomicUsize>> = Mutex::new(HashMap::new());
}

pub fn get_statistics() -> Vec<TrackingStats> {
    let stats = TRACKING_STATS.lock().unwrap();
    let mut result = vec![];

    for (key, value) in stats.iter() {
        result.push(TrackingStats {
            instance_count: value.load(Ordering::SeqCst),
            instance_name: key.clone(),
        });
    }

    result
}

pub struct Tracked<T> (pub T);

impl<'a, T> Tracked<T> {
    pub fn new(t: T) -> Self {
        let mut mutex = TRACKING_STATS.lock().unwrap();
        let type_name: String = type_name::<T>().into();
        let this_type_instance = mutex.get_mut(&type_name);

        let instance_id = if let Some(tracking_info) = this_type_instance {
            tracking_info.fetch_add(1, Ordering::SeqCst)
        } else {
            {
                mutex.insert(type_name.clone(), AtomicUsize::from(0));
            }

            let tracking_info = mutex.get(&type_name).unwrap();

            tracking_info.fetch_add(1, Ordering::SeqCst)
        };

        println!(
            "[Shared<{:?}>] Creating Resource {:?}",
            type_name, instance_id
        );

        Self(t)
    }
}

impl<T> Drop for Tracked<T> {
    fn drop(&mut self) {
        let mut mutex = TRACKING_STATS.lock().unwrap();
        let type_name: String = type_name::<T>().into();
        let this_type_instance = mutex.get_mut(&type_name).unwrap();

        this_type_instance.fetch_sub(1, Ordering::SeqCst);
    }
}

impl<'a, T> Deref for Tracked<T> {
    type Target = T;

    fn deref(&self) -> &T {
        &self.0
    }
}

impl<'a, T> DerefMut for Tracked<T> {
    fn deref_mut(&mut self) -> &mut T {
        &mut self.0
    }
}