dtrace_hdl.c

/* Ruby-DTrace
 * (c) 2007 Chris Andrews <chris@nodnol.org>
 */

#include "dtrace_api.h"

RUBY_EXTERN VALUE eDTraceException;

RUBY_EXTERN VALUE cDTrace;
RUBY_EXTERN VALUE cDTraceProbeDesc;
RUBY_EXTERN VALUE cDTraceProgram;
RUBY_EXTERN VALUE cDTraceRecDesc;
RUBY_EXTERN VALUE cDTraceProbeData;
RUBY_EXTERN VALUE cDTraceBufData;
RUBY_EXTERN VALUE cDTraceProcess;
RUBY_EXTERN VALUE cDTraceDropData;
RUBY_EXTERN VALUE cDTraceErrData;

static void dtrace_hdl_free(void *arg)
{
  dtrace_handle_t *handle = (dtrace_handle_t *)arg;
  VALUE proc;

  if (handle->hdl != NULL) {
    if (handle->procs != Qnil) {
      while ((proc = rb_ary_pop(handle->procs)) != Qnil) {
        dtrace_process_release(proc);
      }
    }
    dtrace_close(handle->hdl);
  }
  free(handle);
}

static void dtrace_hdl_mark(void *arg)
{
  dtrace_handle_t *handle = (dtrace_handle_t *)arg;

  if (handle) {
    rb_gc_mark(handle->probe);
    rb_gc_mark(handle->rec);
    rb_gc_mark(handle->buf);
    rb_gc_mark(handle->err);
    rb_gc_mark(handle->drop);
    rb_gc_mark(handle->procs);
  }
}

VALUE dtrace_hdl_alloc(VALUE klass)
{
  dtrace_hdl_t *hdl;
  dtrace_handle_t *handle;
  int err;
  VALUE obj;

  hdl = dtrace_open(DTRACE_VERSION, 0, &err);

  if (hdl) {
    /*
     * Leopard's DTrace requires symbol resolution to be
     * switched on explicitly
     */
#ifdef __APPLE__
    (void) dtrace_setopt(hdl, "stacksymbols", "enabled");
#endif

    /* always request flowindent information */
    (void) dtrace_setopt(hdl, "flowindent", 0);

    handle = ALLOC(dtrace_handle_t);
    if (!handle) {
      rb_raise(eDTraceException, "alloc failed");
      return Qnil;
    }

    handle->hdl   = hdl;
    handle->probe = Qnil;
    handle->rec   = Qnil;
    handle->buf   = Qnil;
    handle->err   = Qnil;
    handle->drop  = Qnil;
    handle->procs = Qnil;

    obj = Data_Wrap_Struct(klass, dtrace_hdl_mark, dtrace_hdl_free, handle);
    return obj;
  }
  else {
    rb_raise(eDTraceException, "unable to open dtrace: %s (not root?)", strerror(err));
    return Qnil;
  }
}

VALUE dtrace_hdl_close(VALUE self)
{
  dtrace_handle_t *handle;

  Data_Get_Struct(self, dtrace_handle_t, handle);
  dtrace_close(handle->hdl);
  handle->hdl = NULL;

  return Qnil;
}

static
int _dtrace_next_probe(dtrace_hdl_t *hdl, const dtrace_probedesc_t *pdp, void *arg)
{
  VALUE probe;

  probe = Data_Wrap_Struct(cDTraceProbeDesc, 0, NULL, (dtrace_probedesc_t *)pdp);

  rb_yield(probe);
  return 0;
}

/*
 * Yields each probe found on the system.
 * (equivalent to dtrace -l)
 *
 * Each probe is represented by a DTraceProbe object
 */
VALUE dtrace_each_probe_all(VALUE self)
{
  dtrace_handle_t *handle;

  Data_Get_Struct(self, dtrace_handle_t, handle);
  (void) dtrace_probe_iter(handle->hdl, NULL, _dtrace_next_probe, NULL);

  return self;
}

/*
 * Yields each probe found on the system, matching against a
 * partial name.
 * (equivalent to dtrace -l -n 'probe:::spec')
 *
 * Each probe is represented by a DTraceProbe object
 */
VALUE dtrace_each_probe_match(VALUE self, VALUE provider, VALUE mod, VALUE func, VALUE name)
{
  dtrace_handle_t *handle;

  dtrace_probedesc_t desc;
  desc.dtpd_id = 0;
  strcpy(desc.dtpd_provider, RSTRING_PTR(provider));
  strcpy(desc.dtpd_mod,      RSTRING_PTR(mod));
  strcpy(desc.dtpd_func,     RSTRING_PTR(func));
  strcpy(desc.dtpd_name,     RSTRING_PTR(name));

  Data_Get_Struct(self, dtrace_handle_t, handle);
  (void) dtrace_probe_iter(handle->hdl, &desc, _dtrace_next_probe, NULL);

  return self;
}

static int
_dtrace_next_stmt(dtrace_hdl_t *hdl, dtrace_prog_t *program,
		  dtrace_stmtdesc_t *stp, dtrace_ecbdesc_t **last)
{
  dtrace_ecbdesc_t *edp = stp->dtsd_ecbdesc;

  if (edp == *last)
    return 0;

  if (dtrace_probe_iter(hdl, &edp->dted_probe, _dtrace_next_probe, NULL) != 0) {
    rb_raise(eDTraceException, "failed to match %s:%s:%s:%s: %s\n",
	     edp->dted_probe.dtpd_provider, edp->dted_probe.dtpd_mod,
	     edp->dted_probe.dtpd_func, edp->dted_probe.dtpd_name,
	     dtrace_errmsg(hdl, dtrace_errno(hdl)));

  }

  *last = edp;
  return 0;
}

/*
 * Yields each probe enabled by the given D program.
 * (equivalent to dtrace -n -s program.d)
 */
VALUE dtrace_each_probe_prog(VALUE self, VALUE program)
{
  dtrace_handle_t *handle;
  dtrace_prog_t *prog;
  dtrace_ecbdesc_t *last = NULL;

  Data_Get_Struct(self, dtrace_handle_t, handle);
  Data_Get_Struct(program, dtrace_prog_t, prog);

  (void) dtrace_stmt_iter(handle->hdl, prog, (dtrace_stmt_f *)_dtrace_next_stmt, &last);
  return Qnil;
}

/*
 * Compile a D program.
 *
 * Arguments:
 * * The program text to compile
 * * (Optionally) any arguments required by the program
 *
 * Raises a DTraceException if the program cannot be compiled.
 */
VALUE dtrace_strcompile(int argc, VALUE *argv, VALUE self)
{
  dtrace_handle_t *handle;
  dtrace_prog_t *program;
  VALUE dtrace_program;

  VALUE dtrace_text;
  int dtrace_argc;
  VALUE dtrace_argv_array;

  char **dtrace_argv;
  int i;

  rb_scan_args(argc, argv, "1*", &dtrace_text, &dtrace_argv_array);

  dtrace_argc = rb_ary_len(dtrace_argv_array);
  dtrace_argv = ALLOC_N(char *, dtrace_argc + 1);
  if (!dtrace_argv) {
    rb_raise(eDTraceException, "alloc failed");
    return Qnil;
  }

  for (i = 0; i < dtrace_argc; i++) {
    dtrace_argv[i + 1] = strdup(RSTRING_PTR(rb_ary_entry(dtrace_argv_array, i)));
  }

  dtrace_argv[0] = "ruby";
  dtrace_argc++;

  Data_Get_Struct(self, dtrace_handle_t, handle);
  program = dtrace_program_strcompile(
    handle->hdl, RSTRING_PTR(dtrace_text),
    DTRACE_PROBESPEC_NAME, DTRACE_C_PSPEC,
    dtrace_argc, dtrace_argv
  );

  if (!program) {
    rb_raise(eDTraceException, "%s", dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl)));
    return Qnil;
  }
  else {
    dtrace_program = Data_Wrap_Struct(cDTraceProgram, 0, NULL, program);
    rb_iv_set(dtrace_program, "@handle", self);
    return dtrace_program;
  }
}

/*
 * Start tracing. Must be called once a program has been successfully
 * compiled and executed.
 *
 * Raises a DTraceException on any error.
 */
VALUE dtrace_hdl_go(VALUE self)
{
  dtrace_handle_t *handle;

  Data_Get_Struct(self, dtrace_handle_t, handle);
  if (dtrace_go(handle->hdl) < 0)
    rb_raise(eDTraceException, "%s", dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl)));

  return Qnil;
}

/*
 * Returns the status of the DTrace handle.
 *
 * Status values are defined as:
 *
 * * 0 - none
 * * 1 - ok
 * * 4 - stopped
 */
VALUE dtrace_hdl_status(VALUE self)
{
  dtrace_handle_t *handle;
  int status;

  Data_Get_Struct(self, dtrace_handle_t, handle);
  if ((status = dtrace_status(handle->hdl)) < 0)
    rb_raise(eDTraceException, "%s", dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl)));

  return INT2FIX(status);
}

/*
 * Set an option on the DTrace handle.
 *
 * Options which may be set:
 *
 * * aggsize
 * * bufsize
 */
VALUE dtrace_hdl_setopt(VALUE self, VALUE key, VALUE value)
{
  dtrace_handle_t *handle;
  int ret;

  Data_Get_Struct(self, dtrace_handle_t, handle);

  if (NIL_P(value)) {
    ret = dtrace_setopt(handle->hdl, RSTRING_PTR(key), 0);
  }
  else {
    ret = dtrace_setopt(handle->hdl, RSTRING_PTR(key), RSTRING_PTR(value));
  }

  if (ret < 0)
    rb_raise(eDTraceException, "%s", dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl)));

  return Qnil;
}

/* Stop tracing.
 *
 * Must be called after go has been called to start tracing.
 */
VALUE dtrace_hdl_stop(VALUE self)
{
  dtrace_handle_t *handle;

  Data_Get_Struct(self, dtrace_handle_t, handle);
  if (dtrace_stop(handle->hdl) < 0)
    rb_raise(eDTraceException, "%s",
             dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl)));

  return Qnil;
}

/*
 * Return the most recent DTrace error.
 */
VALUE dtrace_hdl_error(VALUE self)
{
  dtrace_handle_t *handle;
  const char *error_string;

  Data_Get_Struct(self, dtrace_handle_t, handle);
  error_string = dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl));
  return rb_str_new2(error_string);
}

/*
 * Sleep until we need to wake up to honour D options controlling
 * consumption rates.
 */
VALUE dtrace_hdl_sleep(VALUE self)
{
  dtrace_handle_t *handle;

  Data_Get_Struct(self, dtrace_handle_t, handle);
  dtrace_sleep(handle->hdl);
  return Qnil;
}

static int _probe_consumer(const dtrace_probedata_t *data, void *arg)
{
  VALUE proc;
  dtrace_work_handlers_t handlers;
  VALUE probedata;

  handlers = *(dtrace_work_handlers_t *) arg;
  proc = handlers.probe;

  if (!NIL_P(proc)) {
    probedata = Data_Wrap_Struct(cDTraceProbeData, 0, NULL,
                                 (dtrace_probedata_t *)data);

    rb_iv_set(probedata, "@handle", handlers.handle);
    rb_funcall(proc, rb_intern("call"), 1, probedata);
  }

  return (DTRACE_CONSUME_THIS);
}

static int _rec_consumer(const dtrace_probedata_t *data, const dtrace_recdesc_t *rec, void *arg)
{
  VALUE proc;
  dtrace_work_handlers_t handlers;
  VALUE recdesc;
  VALUE probedata;

  dtrace_actkind_t act;

  handlers = *(dtrace_work_handlers_t *) arg;
  proc = handlers.rec;
  if (!NIL_P(proc)) {
    if (rec) {
      recdesc = Data_Wrap_Struct(cDTraceRecDesc, 0, NULL, (dtrace_recdesc_t *)rec);
      rb_iv_set(recdesc, "@handle", handlers.handle);
      rb_funcall(proc, rb_intern("call"), 1, recdesc);
    }
    else {
      rb_funcall(proc, rb_intern("call"), 1, Qnil);
      return (DTRACE_CONSUME_NEXT);
    }
  }

  if (rec) {
    act = rec->dtrd_action;
    if (act == DTRACEACT_EXIT)
      return (DTRACE_CONSUME_NEXT);
  }

  return (DTRACE_CONSUME_THIS);
}

static int _buf_consumer(const dtrace_bufdata_t *bufdata, void *arg)
{
  VALUE proc;
  VALUE dtracebufdata;

  proc = (VALUE)arg;

  if (!NIL_P(proc)) {
    dtracebufdata = Data_Wrap_Struct(cDTraceBufData, 0, NULL, (dtrace_bufdata_t *)bufdata);
    rb_funcall(proc, rb_intern("call"), 1, dtracebufdata);
  }

  return (DTRACE_HANDLE_OK);
}

/*
 * Process any data waiting from the D program.
 *
 * Takes a Proc to which DTraceProbeData objects will be yielded, and
 * an optional second Proc to which DTraceRecDesc objects will be
 * yielded.
 *
 */
VALUE dtrace_hdl_work(int argc, VALUE *argv, VALUE self)
{
  dtrace_handle_t *handle;
  dtrace_workstatus_t status;
  dtrace_work_handlers_t handlers;
  VALUE probe_consumer;
  VALUE rec_consumer;

  Data_Get_Struct(self, dtrace_handle_t, handle);

  /* handle args - probe_consumer_proc is mandatory, rec_consumer_proc
     is optional */
  rb_scan_args(argc, argv, "11", &probe_consumer, &rec_consumer);

  /* to mark during GC */
  handle->probe = probe_consumer;
  if (!NIL_P(rec_consumer))
    handle->rec = rec_consumer;

  /* fill out the handlers struct */
  handlers.probe  = probe_consumer;
  handlers.rec    = rec_consumer;
  handlers.handle = self;

  FILE *devnull = fopen("/dev/null", "w");
  status = dtrace_work(handle->hdl, devnull, _probe_consumer, _rec_consumer, &handlers);
  fclose(devnull);

  if (status < 0)
    rb_raise(eDTraceException, "%s", dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl)));

  return INT2FIX(status);
}

/*
 * Set up the buffered output handler for this handle.
 */
VALUE dtrace_hdl_buf_consumer(VALUE self, VALUE buf_consumer)
{
  dtrace_handle_t *handle;
  Data_Get_Struct(self, dtrace_handle_t, handle);

  /* to mark during GC */
  handle->buf = buf_consumer;

  /* attach the buffered output handler */
  if (dtrace_handle_buffered(handle->hdl, &_buf_consumer, (void *)buf_consumer) == -1) {
    rb_raise(eDTraceException, "failed to establish buffered handler: %s",
	     (dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl))));
  }

  return Qnil;
}

static int _drop_consumer(const dtrace_dropdata_t *dropdata, void *arg)
{
  VALUE proc;
  VALUE dtracedropdata;

  proc = (VALUE)arg;

  if (!NIL_P(proc)) {
    dtracedropdata = Data_Wrap_Struct(cDTraceDropData, 0, NULL, (dtrace_dropdata_t *)dropdata);
    rb_funcall(proc, rb_intern("call"), 1, dtracedropdata);
  }

  return (DTRACE_HANDLE_OK);
}

/*
 * Set up the drop-record handler for this handle. Takes a block,
 * which will be called with any drop records returned by DTrace,
 * represented by DTraceDropData objects.
 */
VALUE dtrace_hdl_drop_consumer(VALUE self, VALUE drop_consumer)
{
  dtrace_handle_t *handle;
  Data_Get_Struct(self, dtrace_handle_t, handle);

  /* to mark during GC */
  handle->drop = drop_consumer;

  /* attach the drop-record handler */
  if (dtrace_handle_drop(handle->hdl, &_drop_consumer, (void *)drop_consumer) == -1) {
    rb_raise(eDTraceException, "failed to establish drop-record handler");
  }

  return Qnil;
}

static int _err_consumer(const dtrace_errdata_t *errdata, void *arg)
{
  VALUE proc;
  VALUE dtraceerrdata;

  proc = (VALUE)arg;

  /* guard against bad invocations where arg is not what we provided... */
  if (TYPE(proc) == T_DATA) {
    dtraceerrdata = Data_Wrap_Struct(cDTraceErrData, 0, NULL,
                                     (dtrace_errdata_t *)errdata);
    rb_funcall(proc, rb_intern("call"), 1, dtraceerrdata);
  }
  else {
    /* arg looked bad, throw an exception */
    rb_raise(eDTraceException,
             "bad argument to _err_consumer: %p -> type 0x%x\n", arg, TYPE(proc));
  }

  return (DTRACE_HANDLE_OK);
}

/*
 * Set up the err-record handler for this handle. Takes a block, which
 * will be called with any error records returned by DTrace,
 * represented by DTraceErrData records.
 */
VALUE dtrace_hdl_err_consumer(VALUE self, VALUE err_consumer)
{
  dtrace_handle_t *handle;
  void *arg;
  Data_Get_Struct(self, dtrace_handle_t, handle);

  if (dtrace_status(handle->hdl) != 0) {
    rb_raise(eDTraceException, "too late to add error handler");
    return Qnil;
  }

  /* to mark during GC */
  handle->err = err_consumer;

  /* attach the err-record handler */
  if (dtrace_handle_err(handle->hdl, &_err_consumer, (void *)err_consumer) == -1) {
    rb_raise(eDTraceException, "failed to establish err-record handler: %s",
	     dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl)));
  }

  return Qnil;
}

static void _push_proc(dtrace_handle_t *handle, VALUE proc)
{
  if (handle->procs == Qnil)
    handle->procs = rb_ary_new();

  rb_ary_push(handle->procs, proc);
}

/*
 * Start a process which will be traced. The pid of the started
 * process will be available in D as $target.
 *
 * Pass an array, where the first element is the full path to the
 * program to start, and subsequent elements are its arguments.
 *
 * Returns a DTraceProcess object which is used to start the process
 * once tracing is set up.
 */
VALUE dtrace_hdl_createprocess(VALUE self, VALUE rbargv)
{
  dtrace_handle_t *handle;
  struct ps_prochandle *P;
  char **argv;
  long len;
  int i;
  dtrace_process_t *process;
  VALUE rb_process;

  Data_Get_Struct(self, dtrace_handle_t, handle);

  Check_Type(rbargv, T_ARRAY);
  len = rb_ary_len(rbargv);

  argv = ALLOC_N(char *, len + 1);
  if (!argv) {
    rb_raise(eDTraceException, "alloc failed");
    return Qnil;
  }

  for (i = 0; i < len; i++)
    argv[i] = strdup(RSTRING_PTR(rb_ary_entry(rbargv, i)));
  argv[len] = NULL;

  P = dtrace_proc_create(handle->hdl, argv[0], argv);

  for (i = 0; i < len; i++)
    free(argv[i]);
  free(argv);

  if (P == NULL)
    rb_raise(eDTraceException, "%s",
             dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl)));

  process = ALLOC(dtrace_process_t);
  if (!process) {
    rb_raise(eDTraceException, "alloc failed");
    return Qnil;
  }

  process->handle = handle;
  process->proc   = P;

  rb_process = Data_Wrap_Struct(cDTraceProcess, 0, dtrace_process_free,
                                (dtrace_process_t *)process);

  _push_proc(handle, rb_process);
  return rb_process;
}

/*
 * Grab a currently-running process by pid.
 *
 * Returns a Rb_Process object which is used to start the process
 * once tracing is set up.
 */
VALUE dtrace_hdl_grabprocess(VALUE self, VALUE pid)
{
  dtrace_handle_t *handle;
  struct ps_prochandle *P;
  dtrace_process_t *process;
  VALUE rb_process;

  Data_Get_Struct(self, dtrace_handle_t, handle);

  P = dtrace_proc_grab(handle->hdl, FIX2INT(pid), 0);

  if (P == NULL) {
    rb_raise(eDTraceException, "%s",
             dtrace_errmsg(handle->hdl, dtrace_errno(handle->hdl)));
  }

  process = ALLOC(dtrace_process_t);
  if (!process) {
    rb_raise(eDTraceException, "alloc failed");
    return Qnil;
  }

  process->handle = handle;
  process->proc   = P;

  rb_process = Data_Wrap_Struct(cDTraceProcess, 0, dtrace_process_free,
                                (dtrace_process_t *)process);

  _push_proc(handle, rb_process);
  return rb_process;
}