A simple C dtrace aggregation consumer

This is a little bit of gluing from other sources, but it’s a simple piece of code to consume the content of an aggregation from dtrace using the `libdtrace` library. I’ve tested it on Mac OSX 10.9 and it works, the usual caveats that it wont work all the time, and you need privileged access to get it to run:

#include <assert.h>
#include <dtrace.h>
#include <stdio.h>

// This is the program.
static const char *progstr =
"syscall::recvfrom:return \
{ @agg[execname,pid] = sum(arg0); }";

// This is the aggregation walk function
static int
aggfun(const dtrace_aggdata_t *data, void *gndn __attribute__((unused)))
{
    dtrace_aggdesc_t *aggdesc = data->dtada_desc;
    dtrace_recdesc_t *name_rec, *pid_rec, *sum_rec;
    char *name;
    int32_t *ppid;
    int64_t *psum;
    static const dtrace_aggdata_t *count;

    if (count == NULL) {
        count = data;
        return (DTRACE_AGGWALK_NEXT);
    }

    // Our agrgegation has 4 records (id, execname, pid, sum)
    assert(aggdesc->dtagd_nrecs == 4);

    name_rec = &aggdesc->dtagd_rec[1];
    pid_rec = &aggdesc->dtagd_rec[2];
    sum_rec = &aggdesc->dtagd_rec[3];

    name = data->dtada_data + name_rec->dtrd_offset;
    assert(pid_rec->dtrd_size == sizeof(pid_t));
    ppid = (int32_t *)(data->dtada_data + pid_rec->dtrd_offset);
    assert(sum_rec->dtrd_size == sizeof(int64_t));
    psum = (int64_t *)(data->dtada_data + sum_rec->dtrd_offset);

    printf("%1$-30s %2$-20d %3$-20ld\n", name, *ppid, (long)*psum);
    return (DTRACE_AGGWALK_NEXT);
}

// set the option, otherwise print an error & return -1
int
set_opt(dtrace_hdl_t *dtp, const char *opt, const char *value)
{
    if (-1 == dtrace_setopt(dtp, opt, value)) {
        fprintf(stderr, "Failed to set '%1$s' to '%2$s'.\n", opt, value);
        return (-1);
    }
    return (0);
}

// set all the options, otherwise return an error
int
set_opts(dtrace_hdl_t *dtp)
{
    return (set_opt(dtp, "strsize", "4096")
        | set_opt(dtp, "bufsize", "1m")
        | set_opt(dtp, "aggsize", "1m")
        | set_opt(dtp, "aggrate", "2msec")
        | set_opt(dtp, "arch", "x86_64"));
}

int
main(int argc __attribute__((unused)), char **argv __attribute__((unused)))
{
    int err;
    dtrace_proginfo_t info;
    dtrace_hdl_t *dtp;
    dtrace_prog_t *prog;

    dtp = dtrace_open(DTRACE_VERSION, DTRACE_O_LP64, &err);

    if (dtp == 0) {
        perror("dtrace_open");
        return (1);
    }
    if (-1 == set_opts(dtp))
        return (1);

    prog = dtrace_program_strcompile(dtp, progstr, DTRACE_PROBESPEC_NAME, 0, 0, NULL);
    if (prog == 0) {
        printf("dtrace_program_compile failed\n");
        return (1);
    }
    if (-1 == dtrace_program_exec(dtp, prog, &info)) {
        printf("Failed to dtrace exec.\n");
        return (1);
    }
    if (-1 == dtrace_go(dtp)) {
        fprintf(stderr, "Failed to dtrace_go.\n");
        return (1);
    }

    while(1) {
        int status = dtrace_status(dtp);
        if (status == DTRACE_STATUS_OKAY) {
            dtrace_aggregate_snap(dtp);
            dtrace_aggregate_walk(dtp, aggfun, 0);
        } else if (status != DTRACE_STATUS_NONE) {
            break;
        }
        dtrace_sleep(dtp);
    }

    dtrace_stop(dtp);
    dtrace_close(dtp);
    return (0);
}