ext/liballocations/liballocations.c

#include "liballocations.h"

VALUE allocation_tracer;
VALUE free_tracer;
VALUE state_const;

ID id_enabled;

#define IS_SINGLETON(klass) \
    RB_TYPE_P(klass, T_CLASS) && FL_TEST(klass, FL_SINGLETON)

/**
 * Called whenever a new Ruby object is allocated.
 */
void newobj_callback(VALUE tracepoint, void* data) {
    rb_trace_arg_t *trace_arg = rb_tracearg_from_tracepoint(tracepoint);

    st_data_t count = 0;

    AllocationState *state = allocation_state_get_struct(state_const);

    VALUE klass = rb_tracearg_defined_class(trace_arg);

    /* These aren't actually allocated so there's no point in tracking them. */
    if ( klass == Qtrue || klass == Qfalse || klass == Qnil ) {
        return;
    }

    // We don't care about sigleton classes since only one of them exists at a
    // time. The logic here is stolen from MRI's implementation of
    // Class#singleton_class? as MRI sadly provides no public C function for
    // this method.
    if ( IS_SINGLETON(klass) ) return;

    st_lookup(state->object_counts, (st_data_t) klass, &count);
    st_insert(state->object_counts, (st_data_t) klass, count + 1);
}

/**
 * Called whenever a Ruby object is about to be released.
 *
 * Important: any Ruby allocations in this function will cause CRuby to
 * segfault.
 */
void freeobj_callback(VALUE tracepoint, void* data) {
    st_data_t count;

    rb_trace_arg_t *trace_arg = rb_tracearg_from_tracepoint(tracepoint);
    AllocationState *state = allocation_state_get_struct(state_const);

    VALUE klass = rb_tracearg_defined_class(trace_arg);

    if ( IS_SINGLETON(klass) ) return;

    if ( st_lookup(state->object_counts, (st_data_t) klass, &count) ) {
        if ( count > 0 && (count - 1) > 0 ) {
            st_insert(state->object_counts, (st_data_t) klass, count - 1);
        }
        /* Remove the entry if the count is now 0 */
        else {
            st_delete(state->object_counts, (st_data_t*) &klass, NULL);
        }
    }
}

/**
 * Copies every value in an st_table to a given Ruby Hash.
 */
static int each_count(st_data_t key, st_data_t value, st_data_t hash_ptr) {
    VALUE vkey = (VALUE) key;

    rb_hash_aset((VALUE) hash_ptr, vkey, INT2NUM(value));

    return ST_CONTINUE;
}

/**
 * Returns a Hash containing the current allocation statistics.
 *
 * The returned Hash contains its own copy of the statistics, any further object
 * allocations/frees will not modify said Hash.
 *
 * This method ignores singleton classes.
 *
 * call-seq:
 *     Allocations.to_hash -> Hash
 */
VALUE allocations_to_hash(VALUE self) {
    AllocationState *state = allocation_state_get_struct(state_const);

    st_table *local_counts;
    VALUE hash;

    if ( !state->object_counts ) {
        return rb_hash_new();
    }

    local_counts = allocation_state_copy_table(state);
    hash = rb_hash_new();

    st_foreach(local_counts, each_count, (st_data_t) hash);

    st_free_table(local_counts);

    return hash;
}

/**
 * Starts the counting of object allocations.
 *
 * call-seq:
 *     Allocations.start -> nil
 */
VALUE allocations_start(VALUE self) {
    AllocationState *state = allocation_state_get_struct(state_const);

    if ( rb_ivar_get(self, id_enabled) == Qtrue ) {
        return Qnil;
    }

    allocation_state_allocate_counts(state);

    rb_ivar_set(self, id_enabled, Qtrue);

    rb_tracepoint_enable(allocation_tracer);
    rb_tracepoint_enable(free_tracer);

    return Qnil;
}

/**
 * Stops the counting of object allocations and clears the current statistics.
 *
 * call-seq:
 *     Allocations.stop -> nil
 */
VALUE allocations_stop(VALUE self) {
    AllocationState *state = allocation_state_get_struct(state_const);

    if ( rb_ivar_get(self, id_enabled) != Qtrue ) {
        return Qnil;
    }

    rb_tracepoint_disable(allocation_tracer);
    rb_tracepoint_disable(free_tracer);

    allocation_state_reset_counts(state);

    rb_ivar_set(self, id_enabled, Qfalse);

    return Qnil;
}

/**
 * Returns true if tracking allocations has been enabled, false otherwise.
 *
 * call-seq:
 *     Allocations.enabled? -> true/false
 */
VALUE allocations_enabled_p(VALUE self) {
    VALUE enabled = Qfalse;

    if ( rb_ivar_get(self, id_enabled) == Qtrue ) {
        enabled = Qtrue;
    }

    return enabled;
}

void Init_liballocations() {
    VALUE mAllocations = rb_define_module_under(rb_cObject, "Allocations");

    allocation_tracer = rb_tracepoint_new(Qnil, RUBY_INTERNAL_EVENT_NEWOBJ,
                                          newobj_callback, NULL);

    free_tracer = rb_tracepoint_new(Qnil, RUBY_INTERNAL_EVENT_FREEOBJ,
                                    freeobj_callback, NULL);

    id_enabled = rb_intern("enabled");

    rb_define_singleton_method(mAllocations, "to_hash", allocations_to_hash, 0);
    rb_define_singleton_method(mAllocations, "start", allocations_start, 0);
    rb_define_singleton_method(mAllocations, "stop", allocations_stop, 0);
    rb_define_singleton_method(mAllocations, "enabled?", allocations_enabled_p, 0);

    rb_define_const(mAllocations, "ALLOCATION_TRACER", allocation_tracer);
    rb_define_const(mAllocations, "FREE_TRACER", free_tracer);

    Init_allocations_state();

    state_const = rb_const_get(mAllocations, rb_intern("STATE"));
}