object_counter.rs

//! A simple utility to count the number of allocated objects per type:
//! declared a field of type `ObjectCounter<T>` inside your type `T`,
//! and then `ObjectCounter::<T>::count()` will return the number of
//! allocated objects.
//!
//! This could be used for example a quick and dirty way to count
//! domain-specific objects e. g. sockets, connections, buffers,
//! requests etc.

#![feature(alloc_static)]

use std::sync::atomic::*;
use std::marker;

/// Place this inside your type, and it will report the number of allocations.
pub struct ObjectCounter<T: 'static> {
    marker: marker::PhantomData<T>,
}

impl<T: 'static> ObjectCounter<T> {
    /// Return the number of allocated `<T>` objects.
    pub fn count() -> usize {
        unsafe {
            let counter = &mut *::std::alloc::alloc_static::<ObjectCounter<T>, AtomicUsize>();
            counter.load(Ordering::SeqCst)
        }
    }
}

impl<T: 'static> Default for ObjectCounter<T> {
    /// Increment a per-type object counter by 1.
    fn default() -> ObjectCounter<T> {
        unsafe {
            let counter = &mut *::std::alloc::alloc_static::<ObjectCounter<T>, AtomicUsize>();
            counter.fetch_add(1, Ordering::SeqCst);
        }
        ObjectCounter {
            marker: marker::PhantomData,
        }
    }
}

impl<T: 'static> Drop for ObjectCounter<T> {
    /// Decrement a per-type object counter by 1.
    fn drop(&mut self) {
        unsafe {
            let counter = &mut *::std::alloc::alloc_static::<ObjectCounter<T>, AtomicUsize>();
            counter.fetch_sub(1, Ordering::SeqCst);
        }
    }
}


#[derive(Default)]
struct MySampleData {
    content: String,
    counter: ObjectCounter<MySampleData>,
}

#[derive(Default)]
struct MyAnotherData {
    content: u32,
    counter: ObjectCounter<MyAnotherData>,
}

fn main() {
    let _x = Box::new(MySampleData::default());
    let _y = Box::new(MySampleData::default());
    let _z = Box::new(MyAnotherData::default());

    // Prints `2 1`
    println!("{} {}",
        ObjectCounter::<MySampleData>::count(),
        ObjectCounter::<MyAnotherData>::count());

    assert_eq!(2, ObjectCounter::<MySampleData>::count());
    assert_eq!(1, ObjectCounter::<MyAnotherData>::count());
}