memprof 0.0.1 → 0.1.0

data/.gitignore

@@ -0,0 +1,4 @@
+*.o
+*.so
+*.bundle
+Makefile

data/README

@@ -2,17 +2,50 @@ memprof (c) Joe Damato  @joedamato   http://timetobleed.com
 
 What is memprof?
 ================
+Memprof is a memory profiler for Ruby that requires no patches to the Ruby VM.
+It can help you find Ruby level memory leaks in your application.
 
-Memprof hopes to become a memory profiler for Ruby that will work without
-modifying your Ruby binary. You just require the gem and off you go.
+How to use
+==========
 
+require 'memprof'
+Memprof.start
 
-When will it be done?
-=====================
+# ruby code
 
-No idea, but soon I hope.
+Memprof.stop
+Memprof.dump
 
-Why release this then?
-======================
+That will monitor allocations and frees that happen between "start" and "stop"
+and will output information to standard error.
 
-This gem illustrates an ugly hack that I thought might interest other people.
+Memprof.dump also takes an (optional) file name to write the output to a file.
+
+Supported systems
+=================
+Currently supporting:
+
+  Linux:
+    x86_64 builds of Ruby Enterprise Edition 1.8.6/1.8.7
+    x86_64 builds of MRI Ruby if built with --disable-shared
+
+Experimental support:
+
+  Snow Leopard:
+    x86_64 builds of MRI (both enable-shared and disable-shared)
+
+Coming soon:
+
+  Official support for Snow Leopard.
+
+  Linux:
+    Tracking object allocationns in C extensions.
+    x86_64 builds of MRI Ruby with --enable-shared
+
+  i386/i686 support for all the above.
+
+CREDITS
+=======
+Jake Douglas for the Mach O/snow leopard support.
+
+Aman Gupta for various bug fixes and other cleanup.

data/ext/bin_api.h

@@ -0,0 +1,83 @@
+#if !defined(BIN_API__)
+#define BIN_API__
+#include <stddef.h>
+#include <stdint.h>
+
+/* generic file format stuff */
+extern void *text_segment;
+extern unsigned long text_segment_len;
+extern size_t pagesize;
+
+/*
+ *  trampoline specific stuff
+ */
+extern struct tramp_tbl_entry *tramp_table;
+extern size_t tramp_size;
+
+/*
+ *  inline trampoline specific stuff
+ */
+extern size_t inline_tramp_size;
+extern struct inline_tramp_tbl_entry *inline_tramp_table;
+
+/* trampoline types */
+struct tramp_inline {
+  unsigned char jmp[1];
+  uint32_t displacement;
+  unsigned char pad[2];
+} __attribute__((__packed__));
+
+struct tramp_tbl_entry {
+  unsigned char rbx_save[1];
+  unsigned char mov[2];
+  void *addr;
+  unsigned char callq[2];
+  unsigned char rbx_restore[1];
+  unsigned char ret[1];
+} __attribute__((__packed__));
+
+struct inline_tramp_tbl_entry {
+  unsigned char rex[1];
+  unsigned char mov[1];
+  unsigned char src_reg[1];
+  uint32_t mov_displacement;
+
+  struct {
+    unsigned char push_rdi[1];
+    unsigned char mov_rdi[3];
+    uint32_t rdi_source_displacement;
+    unsigned char push_rbx[1];
+    unsigned char push_rbp[1];
+    unsigned char save_rsp[3];
+    unsigned char align_rsp[4];
+    unsigned char mov[2];
+    void *addr;
+    unsigned char callq[2];
+    unsigned char leave[1];
+    unsigned char rbx_restore[1];
+    unsigned char rdi_restore[1];
+  } __attribute__((__packed__)) frame;
+
+  unsigned char jmp[1];
+  uint32_t jmp_displacement;
+} __attribute__((__packed__));
+
+void
+update_callqs(int entry, void *trampee_addr);
+
+/*
+ *  EXPORTED API.
+ */
+void
+bin_init();
+
+void *
+bin_find_symbol(char *sym, size_t *size);
+
+void
+bin_update_image(int entry, void *trampee_addr);
+
+void *
+bin_allocate_page();
+
+#endif

data/ext/elf.c

@@ -0,0 +1,106 @@
+#if defined(HAVE_ELF)
+
+#include "bin_api.h"
+
+#include <stdio.h>
+#include <fcntl.h>
+#include <gelf.h>
+#include <link.h>
+#include <sysexits.h>
+#include <unistd.h>
+
+#include <sys/mman.h>
+
+static ElfW(Shdr) symtab_shdr;
+static Elf *elf = NULL;
+static Elf_Data *symtab_data = NULL;
+
+void *
+bin_allocate_page()
+{
+  return mmap(NULL, pagesize, PROT_WRITE|PROT_READ|PROT_EXEC, MAP_ANON|MAP_PRIVATE|MAP_32BIT, -1, 0);
+}
+
+void
+bin_update_image(int entry, void *trampee_addr)
+{
+  update_callqs(entry, trampee_addr);
+}
+
+void *
+bin_find_symbol(char *sym, size_t *size)
+{
+  char *name = NULL;
+
+  /*now print the symbols*/
+  ElfW(Sym) *esym = (ElfW(Sym)*) symtab_data->d_buf;
+  ElfW(Sym) *lastsym = (ElfW(Sym)*) ((char*) symtab_data->d_buf + symtab_data->d_size);
+  /* now loop through the symbol table and print it*/
+  for (; esym < lastsym; esym++){
+    if ((esym->st_value == 0) ||
+        (ELF32_ST_BIND(esym->st_info)== STB_WEAK) ||
+        (ELF32_ST_BIND(esym->st_info)== STB_NUM))
+      continue;
+    name = elf_strptr(elf, symtab_shdr.sh_link, (size_t)esym->st_name);
+    if (strcmp(name, sym) == 0) {
+      if (size) {
+        *size = esym->st_size;
+      }
+      return (void *)esym->st_value;
+    }
+  }
+  return NULL;
+}
+
+
+void
+bin_init()
+{
+  int fd;
+  ElfW(Shdr) shdr;
+  size_t shstrndx;
+  char *filename;
+  Elf_Scn *scn;
+
+  if (elf_version(EV_CURRENT) == EV_NONE)
+    errx(EX_SOFTWARE, "ELF library initialization failed: %s",
+        elf_errmsg(-1));
+
+  asprintf(&filename, "/proc/%ld/exe", (long)getpid());
+
+  if ((fd = open(filename, O_RDONLY, 0)) < 0)
+    err(EX_NOINPUT, "open \%s\" failed", filename);
+
+  if ((elf = elf_begin(fd, ELF_C_READ, NULL)) == NULL)
+    errx(EX_SOFTWARE, "elf_begin() failed: %s.",
+        elf_errmsg(-1));
+
+  if (elf_kind(elf) != ELF_K_ELF)
+    errx(EX_DATAERR, "%s is not an ELF object.", filename);
+
+  if (elf_getshstrndx(elf, &shstrndx) == 0)
+    errx(EX_SOFTWARE, "getshstrndx() failed: %s.",
+        elf_errmsg(-1));
+
+  scn = NULL;
+
+  while ((scn = elf_nextscn(elf, scn)) != NULL) {
+    if (gelf_getshdr(scn, &shdr) != &shdr)
+      errx(EX_SOFTWARE, "getshdr() failed: %s.",
+          elf_errmsg(-1));
+
+    if (shdr.sh_type == SHT_PROGBITS &&
+        (shdr.sh_flags == (SHF_ALLOC | SHF_EXECINSTR)) &&
+        strcmp(elf_strptr(elf, shstrndx, shdr.sh_name), ".text") == 0) {
+
+        text_segment = (void *)shdr.sh_addr;
+        text_segment_len = shdr.sh_size;
+    } else if (shdr.sh_type == SHT_SYMTAB) {
+        symtab_shdr = shdr;
+        if ((symtab_data = elf_getdata(scn,symtab_data)) == 0 || symtab_data->d_size == 0) {
+          return;
+        }
+    }
+  }
+}
+#endif

data/ext/extconf.rb

@@ -1,5 +1,10 @@
 require 'mkmf'
 
-if (have_library('elf', 'gelf_getshdr'))
-  create_makefile ('memprof')
+def add_define(name)
+  $defs.push("-D#{name}")
+end
+
+if (add_define("HAVE_ELF") if have_library('elf', 'gelf_getshdr')) ||
+  (add_define("HAVE_MACH") if have_header('mach-o/dyld.h'))
+  create_makefile('memprof')
 end

data/ext/mach.c

@@ -0,0 +1,108 @@
+#if defined(HAVE_MACH)
+
+#include "bin_api.h"
+
+#include <limits.h>
+#include <sysexits.h>
+#include <sys/mman.h>
+
+#include <mach-o/dyld.h>
+#include <mach-o/getsect.h>
+#include <mach-o/loader.h>
+#include <mach-o/ldsyms.h>
+
+static void
+set_text_segment(struct mach_header *header, const char *sectname)
+{
+  text_segment = getsectdatafromheader_64((struct mach_header_64*)header, "__TEXT", sectname, (uint64_t*)&text_segment_len);
+  if (!text_segment)
+    errx(EX_UNAVAILABLE, "Failed to locate the %s section", sectname);
+}
+
+static void
+update_dyld_stubs(int entry, void *trampee_addr)
+{
+  char *byte = text_segment;
+  size_t count = 0;
+
+  for(; count < text_segment_len; count++) {
+    if (*byte == '\xff') {
+      int off = *(int *)(byte+2);
+      if (trampee_addr == (void*)(*(long long*)(byte + 6 + off))) {
+        *(long long*)(byte + 6 + off) = tramp_table[entry].addr;
+      }
+    }
+    byte++;
+  }
+}
+
+void *
+bin_allocate_page()
+{
+  void *ret = NULL;
+  size_t i = 0;
+
+  for (i = pagesize; i < INT_MAX - pagesize; i += pagesize) {
+    ret = mmap((void*)(NULL + i), 2*pagesize, PROT_WRITE|PROT_READ|PROT_EXEC,
+               MAP_ANON|MAP_PRIVATE, -1, 0);
+
+    if (tramp_table != MAP_FAILED)
+      return ret;
+  }
+  return NULL;
+}
+
+void
+bin_update_image(int entry, void *trampee_addr)
+{
+  // Modify any callsites residing inside the text segment of the executable itself
+  set_text_segment((struct mach_header*)&_mh_execute_header, "__text");
+  update_callqs(entry, trampee_addr);
+
+  // Modify all dyld stubs in shared libraries that have been loaded
+  int i, j, k;
+  int header_count = _dyld_image_count();
+
+  // Go through all the mach objects that are loaded into this process
+  for (i=0; i < header_count; i++) {
+    const struct mach_header *current_hdr = _dyld_get_image_header(i);
+    int lc_count = current_hdr->ncmds;
+
+    // this as a char* because we need to step it forward by an arbitrary number of bytes
+    const char *lc = ((const char*) current_hdr) + sizeof(struct mach_header_64);
+
+    // Check all the load commands in the object to see if they are segment commands
+    for (j = 0; j < lc_count; j++) {
+      if (((struct load_command*)lc)->cmd == LC_SEGMENT_64) {
+        const struct segment_command_64 *seg = (const struct segment_command_64 *) lc;
+        const struct section_64 * sect = (const struct section_64*)(lc + sizeof(struct segment_command_64));
+        int section_count = (seg->cmdsize - sizeof(struct segment_command_64)) / sizeof(struct section_64);
+
+        // Search the segment for a section containing dyld stub functions
+        for (k=0; k < section_count; k++) {
+          if (strncmp(sect->sectname, "__symbol_stub", 13) == 0) {
+            set_text_segment((struct mach_header*)current_hdr, sect->sectname);
+            text_segment += _dyld_get_image_vmaddr_slide(i);
+            update_dyld_stubs(entry, trampee_addr);
+          }
+          sect++;
+        }
+      }
+      lc += ((struct load_command*)lc)->cmdsize;
+    }
+  }
+}
+
+void  *
+bin_find_symbol(char *sym,  size_t *size) {
+  void *ptr = NULL;
+  _dyld_lookup_and_bind((const char*)sym, &ptr, NULL);
+  return ptr;
+}
+
+void
+bin_init()
+{
+  /* mach-o is so cool it needs no initialization */
+}
+#endif

data/ext/memprof.c

@@ -1,78 +1,289 @@
 #define _GNU_SOURCE
 #include <err.h>
 #include <fcntl.h>
-#include <gelf.h>
+#include <stddef.h>
 #include <stdio.h>
 #include <stdint.h>
 #include <stdlib.h>
 #include <unistd.h>
-#include <link.h>
 #include <sysexits.h>
 #include <sys/mman.h>
+#include <err.h>
 
+#include <st.h>
 #include <ruby.h>
 #include <intern.h>
+#include <node.h>
 
-struct tramp_tbl_entry {
-  unsigned char mov[2];
-  long long addr;
-  unsigned char callq[2];
-  unsigned char ret;
-  unsigned char pad[3];
-} __attribute__((__packed__));;
+#include "bin_api.h"
 
+size_t pagesize;
+void *text_segment = NULL;
+unsigned long text_segment_len = 0;
 
-static void *text_segment = NULL;
-static unsigned long text_segment_len = 0;
+/*
+   trampoline specific stuff
+ */
+struct tramp_tbl_entry *tramp_table = NULL;
+size_t tramp_size = 0;
 
 /*
-  trampoline specific stuff
-*/
-static struct tramp_tbl_entry *tramp_table = NULL;
-static size_t tramp_size = 0;
+   inline trampoline specific stuff
+ */
+size_t inline_tramp_size = 0;
+struct inline_tramp_tbl_entry *inline_tramp_table = NULL;
 
 /*
-  ELF specific stuff
-*/
-static ElfW(Shdr) symtab_shdr;
-static Elf *elf = NULL;
-static Elf_Data *symtab_data = NULL;
+ * bleak_house stuff
+ */
+static int track_objs = 0;
+static st_table *objs = NULL;
 
+struct obj_track {
+  VALUE obj;
+  char *source;
+  int line;
+};
 
 static void
-error_tramp() {
+error_tramp()
+{
   printf("WARNING: NO TRAMPOLINE SET.\n");
   return;
 }
 
 static VALUE
-newobj_tramp() {
-  printf("source = %s, line = %d\n", ruby_sourcefile, ruby_sourceline);
-  return rb_newobj();
+newobj_tramp()
+{
+  VALUE ret = rb_newobj();
+  struct obj_track *tracker = NULL;
+
+  if (track_objs) {
+    tracker = malloc(sizeof(*tracker));
+
+    if (tracker) {
+      if (ruby_current_node && ruby_current_node->nd_file && *ruby_current_node->nd_file) {
+        tracker->source = strdup(ruby_current_node->nd_file);
+        tracker->line = nd_line(ruby_current_node);
+      } else if (ruby_sourcefile) {
+        tracker->source = strdup(ruby_sourcefile);
+        tracker->line = ruby_sourceline;
+      } else {
+        tracker->source = strdup("__null__");
+        tracker->line = 0;
+      }
+
+      tracker->obj = ret;
+      st_insert(objs, (st_data_t)ret, (st_data_t)tracker);
+    } else {
+      fprintf(stderr, "Warning, unable to allocate a tracker. You are running dangerously low on RAM!\n");
+    }
+  }
+
+  return ret;
+}
+
+static void
+freelist_tramp(unsigned long rval)
+{
+  struct obj_track *tracker = NULL;
+
+  if (track_objs) {
+    st_delete(objs, (st_data_t *) &rval, (st_data_t *) &tracker);
+    if (tracker) {
+      free(tracker->source);
+      free(tracker);
+    }
+  }
+}
+
+static int
+memprof_tabulate(st_data_t key, st_data_t record, st_data_t arg)
+{
+  st_table *table = (st_table *)arg;
+  struct obj_track *tracker = (struct obj_track *)record;
+  char *source_key = NULL;
+  unsigned long count = 0;
+  char *type = NULL;
+
+  switch (TYPE(tracker->obj)) {
+    case T_NONE:
+      type = "__none__"; break;
+    case T_BLKTAG:
+      type = "__blktag__"; break;
+    case T_UNDEF:
+      type = "__undef__"; break;
+    case T_VARMAP:
+      type = "__varmap__"; break;
+    case T_SCOPE:
+      type = "__scope__"; break;
+    case T_NODE:
+      type = "__node__"; break;
+    default:
+      if (RBASIC(tracker->obj)->klass) {
+        type = rb_obj_classname(tracker->obj);
+      } else {
+        type = "__unknown__";
+      }
+  }
+
+  asprintf(&source_key, "%s:%d:%s", tracker->source, tracker->line, type);
+  st_lookup(table, (st_data_t)source_key, (st_data_t *)&count);
+  if (st_insert(table, (st_data_t)source_key, ++count)) {
+    free(source_key);
+  }
+
+  free(tracker->source);
+  return ST_DELETE;
+}
+
+struct results {
+  char **entries;
+  unsigned long num_entries;
+};
+
+static int
+memprof_do_dump(st_data_t key, st_data_t record, st_data_t arg)
+{
+  struct results *res = (struct results *)arg;
+  unsigned long count = (unsigned long)record;
+  char *source = (char *)key;
+  
+  asprintf(&(res->entries[res->num_entries++]), "%7d %s", count, source);
+
+  free(source);
+  return ST_DELETE;
+}
+
+static VALUE
+memprof_start(VALUE self)
+{
+  if (track_objs == 1)
+    return Qfalse;
+
+  track_objs = 1;
+  return Qtrue;
+}
+
+static VALUE
+memprof_stop(VALUE self)
+{
+  if (track_objs == 0)
+    return Qfalse;
+
+  track_objs = 0;
+  return Qtrue;
+}
+
+static int
+memprof_strcmp(const void *obj1, const void *obj2)
+{
+  char *str1 = *(char **)obj1;
+  char *str2 = *(char **)obj2;
+  return strcmp(str2, str1);
+}
+
+static VALUE
+memprof_dump(int argc, VALUE *argv, VALUE self)
+{
+  st_table *tmp_table;
+  struct results res;
+  int i;
+  VALUE str;
+  FILE *out = NULL;
+
+  rb_scan_args(argc, argv, "01", &str);
+
+  if (RTEST(str)) {
+    out = fopen(StringValueCStr(str), "w");
+    if (!out)
+      rb_raise(rb_eArgError, "unable to open output file");
+  }
+
+  track_objs = 0;
+
+  tmp_table = st_init_strtable();
+  st_foreach(objs, memprof_tabulate, (st_data_t)tmp_table);
+
+  res.num_entries = 0;
+  res.entries = malloc(sizeof(char*) * tmp_table->num_entries);
+
+  st_foreach(tmp_table, memprof_do_dump, (st_data_t)&res);
+  st_free_table(tmp_table);
+
+  qsort(res.entries, res.num_entries, sizeof(char*), &memprof_strcmp);
+  for (i=0; i < res.num_entries; i++) {
+    fprintf(out ? out : stderr, "%s\n", res.entries[i]);
+    free(res.entries[i]);
+  }
+  free(res.entries);
+
+  track_objs = 1;
+  return Qnil;
 }
 
 static void
-create_tramp_table() {
-  int i = 0;
+create_tramp_table()
+{
+  int i, j = 0;
 
   struct tramp_tbl_entry ent = {
-    .mov = {'\x48', '\xbb'},
-    .addr = (long long)&error_tramp,
-    .callq = { '\xff', '\xd3' },
-    .ret = '\xc3',
-    .pad =  { '\x90', '\x90', '\x90'},
+    .rbx_save      = {'\x53'},                // push rbx
+    .mov           = {'\x48', '\xbb'},        // mov addr into rbx
+    .addr          = error_tramp,             // ^^^
+    .callq         = {'\xff', '\xd3'},        // callq rbx
+    .rbx_restore   = {'\x5b'},                // pop rbx
+    .ret           = {'\xc3'},                // ret
   };
 
-  tramp_table = mmap(NULL, 4096, PROT_WRITE|PROT_READ|PROT_EXEC, MAP_32BIT|MAP_ANONYMOUS|MAP_PRIVATE, -1, 0);
-  if (tramp_table != MAP_FAILED) {
-    for (; i < 4096/sizeof(struct tramp_tbl_entry); i ++ ) {
-      memcpy(tramp_table + i, &ent, sizeof(struct tramp_tbl_entry));
-    }
+  struct inline_tramp_tbl_entry inline_ent = {
+    .rex     = {'\x48'},
+    .mov     = {'\x89'},
+    .src_reg = {'\x05'},
+    .mov_displacement = 0,
+
+    .frame = {
+      .push_rdi = {'\x57'},
+      .mov_rdi = {'\x48', '\x8b', '\x3d'},
+      .rdi_source_displacement = 0,
+      .push_rbx = {'\x53'},
+      .push_rbp = {'\x55'},
+      .save_rsp = {'\x48', '\x89', '\xe5'},
+      .align_rsp = {'\x48', '\x83', '\xe4', '\xf0'},
+      .mov = {'\x48', '\xbb'},
+      .addr = error_tramp,
+      .callq = {'\xff', '\xd3'},
+      .leave = {'\xc9'},
+      .rbx_restore = {'\x5b'},
+      .rdi_restore = {'\x5f'},
+    },
+
+    .jmp  = {'\xe9'},
+    .jmp_displacement = 0,
+  };
+
+  if ((tramp_table = bin_allocate_page()) == MAP_FAILED) {
+    fprintf(stderr, "Failed to allocate memory for stage 1 trampoline.\n");
+    return;
+  }
+
+  if ((inline_tramp_table = bin_allocate_page()) == MAP_FAILED) {
+    fprintf(stderr, "Faied to allocate memory for the stage 1 inline trampoline.\n");
+    return;
+  }
+
+  for (j = 0; j < pagesize/sizeof(struct tramp_tbl_entry); j ++ ) {
+    memcpy(tramp_table + j, &ent, sizeof(struct tramp_tbl_entry));
+  }
+
+  for (j = 0; j < pagesize/sizeof(struct inline_tramp_tbl_entry); j++) {
+    memcpy(inline_tramp_table + j, &inline_ent, sizeof(struct inline_tramp_tbl_entry));
   }
 }
 
-static void
-update_image(int entry, void *trampee_addr) {
+void
+update_callqs(int entry, void *trampee_addr)
+{
   char *byte = text_segment;
   size_t count = 0;
   int fn_addr = 0;
@@ -92,93 +303,174 @@ update_image(int entry, void *trampee_addr) {
   }
 }
 
-static void *
-find_symbol(char *sym) {
-  char *name = NULL;
-
-  /*now print the symbols*/
-  ElfW(Sym) *esym = (ElfW(Sym)*) symtab_data->d_buf;
-  ElfW(Sym) *lastsym = (ElfW(Sym)*) ((char*) symtab_data->d_buf + symtab_data->d_size);
-  /* now loop through the symbol table and print it*/
-  for (; esym < lastsym; esym++){
-    if ((esym->st_value == 0) ||
-        (ELF32_ST_BIND(esym->st_info)== STB_WEAK) ||
-        (ELF32_ST_BIND(esym->st_info)== STB_NUM) ||
-        (ELF32_ST_TYPE(esym->st_info)!= STT_FUNC))
-      continue;
-    name = elf_strptr(elf, symtab_shdr.sh_link, (size_t)esym->st_name);
-    if (strcmp(name, sym) == 0) {
-       return (void *)esym->st_value;
+
+static void
+hook_freelist(int entry)
+{
+  long sizes[] = { 0, 0, 0 };
+  void *sym1 = bin_find_symbol("gc_sweep", &sizes[0]);
+
+  if (sym1 == NULL) {
+    /* this is MRI ... */
+    sym1 = bin_find_symbol("garbage_collect", &sizes[0]);
+  }
+
+  void *sym2 = bin_find_symbol("finalize_list", &sizes[1]);
+  void *sym3 = bin_find_symbol("rb_gc_force_recycle", &sizes[2]);
+  void *freelist_callers[] = { sym1, sym2, sym3 };
+  int max = 3;
+  size_t i = 0;
+  char *byte = freelist_callers[0];
+  void *freelist = bin_find_symbol("freelist", NULL);
+  uint32_t mov_target =  0;
+  void *aligned_addr = NULL;
+  size_t count = 0;
+
+  /* This is the stage 1 trampoline for hooking the inlined add_freelist
+   * function .
+   *
+   * NOTE: The original instruction mov %reg, freelist is 7 bytes wide,
+   * whereas jmpq $displacement is only 5 bytes wide. We *must* pad out
+   * the next two bytes. This will be important to remember below.
+   */
+  struct tramp_inline tramp = {
+    .jmp           = {'\xe9'},
+    .displacement  = 0,
+    .pad           = {'\x90', '\x90'},
+  };
+
+  struct inline_tramp_tbl_entry *inl_tramp_st2 = NULL;
+
+  for (;i < max;) {
+    /* make sure it is a mov instruction */
+    if (byte[1] == '\x89') {
+
+      /* Read the REX byte to make sure it is a mov that we care about */
+      if ((byte[0] == '\x48') ||
+          (byte[0] == '\x4c')) {
+
+        /* Grab the target of the mov. REMEMBER: in this case the target is 
+         * a 32bit displacment that gets added to RIP (where RIP is the adress of
+         * the next instruction).
+         */
+        mov_target = *(uint32_t *)(byte + 3);
+
+        /* Sanity check. Ensure that the displacement from freelist to the next
+         * instruction matches the mov_target. If so, we know this mov is
+         * updating freelist.
+         */
+        if ((freelist - (void *)(byte+7)) == mov_target) {
+          /* Before the stage 1 trampoline gets written, we need to generate
+           * the code for the stage 2 trampoline. Let's copy over the REX byte
+           * and the byte which mentions the source register into the stage 2
+           * trampoline.
+           */
+          inl_tramp_st2 = inline_tramp_table + entry;
+          inl_tramp_st2->rex[0] = byte[0];
+          inl_tramp_st2->src_reg[0] = byte[2];
+
+          /* Setup the stage 1 trampoline. Calculate the displacement to
+           * the stage 2 trampoline from the next instruction.
+           *
+           * REMEMBER!!!! The next instruction will be NOP after our stage 1
+           * trampoline is written. This is 5 bytes into the structure, even
+           * though the original instruction we overwrote was 7 bytes.
+           */
+          tramp.displacement = (uint32_t)((void *)(inl_tramp_st2) - (void *)(byte+5));
+
+          /* Figure out what page the stage 1 tramp is gonna be written to, mark
+           * it WRITE, write the trampoline in, and then remove WRITE permission.
+           */
+          aligned_addr = (void*)(((long)byte)&~(0xffff));
+          mprotect(aligned_addr, (((void *)byte) - aligned_addr) + 10, PROT_READ|PROT_WRITE|PROT_EXEC);
+          memcpy(byte, &tramp, sizeof(struct tramp_inline));
+          mprotect(aligned_addr, (((void *)byte) - aligned_addr) + 10, PROT_READ|PROT_EXEC);
+
+          /* Finish setting up the stage 2 trampoline. */
+
+          /* calculate the displacement to freelist from the next instruction.
+           *
+           * This is used to replicate the original instruction we overwrote.
+           */
+          inl_tramp_st2->mov_displacement = freelist - (void *)&(inl_tramp_st2->frame);
+
+          /* fill in the displacement to freelist from the next instruction.
+           *
+           * This is to arrange for the new value in freelist to be in %rdi, and as such
+           * be the first argument to the C handler. As per the amd64 ABI.
+           */
+          inl_tramp_st2->frame.rdi_source_displacement = freelist - (void *)&(inl_tramp_st2->frame.push_rbx);
+
+          /* jmp back to the instruction after stage 1 trampoline was inserted 
+           *
+           * This can be 5 or 7, it doesn't matter. If its 5, we'll hit our 2
+           * NOPS. If its 7, we'll land directly on the next instruction.
+           */
+          inl_tramp_st2->jmp_displacement = (uint32_t)((void *)(byte + 7) -
+                                                       (void *)(inline_tramp_table + entry + 1));
+
+          /* write the address of our C level trampoline in to the structure */
+          inl_tramp_st2->frame.addr = freelist_tramp;
+
+          /* track the new entry and new trampoline size */
+          entry++;
+          inline_tramp_size++;
+        }
+      }
     }
+
+    if (count >= sizes[i]) {
+        count = 0;
+        i ++;
+        byte = freelist_callers[i];
+    }
+    count++;
+    byte++;
   }
-  return NULL;
 }
 
 static void
-insert_tramp(char *trampee, void *tramp) {
-  void *trampee_addr = find_symbol(trampee);
+insert_tramp(char *trampee, void *tramp)
+{
+  void *trampee_addr = bin_find_symbol(trampee, NULL);
   int entry = tramp_size;
-  tramp_table[tramp_size].addr = (long long)tramp;
-  tramp_size++;
-  update_image(entry, trampee_addr);
-}
+  int inline_ent = inline_tramp_size;
 
-void Init_memprof()
-{
-  int fd;
-  ElfW(Shdr) shdr;
-  size_t shstrndx;
-  char *filename; 
-  Elf_Scn *scn;
-
-  if (elf_version(EV_CURRENT) == EV_NONE)
-    errx(EX_SOFTWARE, "ELF library initialization failed: %s",
-        elf_errmsg(-1));
-
-  asprintf(&filename, "/proc/%ld/exe", (long)getpid());
-
-  if ((fd = open(filename, O_RDONLY, 0)) < 0)
-    err(EX_NOINPUT, "open \%s\" failed", filename);
-
-  if ((elf = elf_begin(fd, ELF_C_READ, NULL)) == NULL)
-    errx(EX_SOFTWARE, "elf_begin() failed: %s.",
-        elf_errmsg(-1));
-
-  if (elf_kind(elf) != ELF_K_ELF)
-    errx(EX_DATAERR, "%s is not an ELF object.", filename);
-
-  if (elf_getshstrndx(elf, &shstrndx) == 0)
-    errx(EX_SOFTWARE, "getshstrndx() failed: %s.",
-        elf_errmsg(-1));
-
-  scn = NULL;
-
-  while ((scn = elf_nextscn(elf, scn)) != NULL) {
-    if (gelf_getshdr(scn, &shdr) != &shdr)
-      errx(EX_SOFTWARE, "getshdr() failed: %s.",
-          elf_errmsg(-1));
-
-    if (shdr.sh_type == SHT_PROGBITS &&
-        (shdr.sh_flags == (SHF_ALLOC | SHF_EXECINSTR)) &&
-        strcmp(elf_strptr(elf, shstrndx, shdr.sh_name), ".text") == 0) {
-      
-        text_segment = (void *)shdr.sh_addr;
-        text_segment_len = shdr.sh_size;
-    } else if (shdr.sh_type == SHT_SYMTAB) {
-        symtab_shdr = shdr;
-        if ((symtab_data = elf_getdata(scn,symtab_data)) == 0 || symtab_data->d_size == 0) {
-          return;
-        }
+  if (trampee_addr == NULL) {
+    if (strcmp("add_freelist", trampee) == 0) {
+      /* XXX super hack */
+      inline_tramp_table[inline_tramp_size].frame.addr = tramp;
+      inline_tramp_size++;
+      hook_freelist(inline_ent);
+    } else {
+      return;
     }
+  } else {
+    tramp_table[tramp_size].addr = tramp;
+    tramp_size++;
+    bin_update_image(entry, trampee_addr);
   }
+}
 
+void
+Init_memprof()
+{
+  VALUE memprof = rb_define_module("Memprof");
+  rb_define_singleton_method(memprof, "start", memprof_start, 0);
+  rb_define_singleton_method(memprof, "stop", memprof_stop, 0);
+  rb_define_singleton_method(memprof, "dump", memprof_dump, -1);
 
+  pagesize = getpagesize();
+  objs = st_init_numtable();
+  bin_init();
   create_tramp_table();
 
+#if defined(HAVE_MACH)
+  insert_tramp("_rb_newobj", newobj_tramp);
+#elif defined(HAVE_ELF)
   insert_tramp("rb_newobj", newobj_tramp);
-#if 0
-  (void) elf_end(e);
-  (void) close(fd);
+  insert_tramp("add_freelist", freelist_tramp);
 #endif
+
   return;
 }

data/memprof.gemspec

@@ -1,17 +1,22 @@
 spec = Gem::Specification.new do |s|
   s.name = 'memprof'
-  s.version = '0.0.1'
-  s.date = '2009-11-21'
-  s.summary = 'Ruby memory profiler gem'
+  s.version = '0.1.0'
+  s.date = '2009-12-10'
+  s.summary = 'Ruby Memory Profiler'
+  s.description = "Ruby memory profiler similar to bleak_house, but without patches to the Ruby VM"
   s.email = "ice799@gmail.com"
   s.homepage = "http://github.com/ice799/memprof"
-  s.description = "Ruby memory profiler gem"
   s.has_rdoc = false
   s.authors = ["Joe Damato"]
   s.extensions = "ext/extconf.rb"
-  s.require_paths << "ext"
-  s.files = ["README",
-             "memprof.gemspec",
-             "ext/memprof.c",
-             "ext/extconf.rb"]
+  s.files = %w[
+    .gitignore
+    README
+    ext/bin_api.h
+    ext/elf.c
+    ext/extconf.rb
+    ext/mach.c
+    ext/memprof.c
+    memprof.gemspec
+  ]
 end

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification 
 name: memprof
 version: !ruby/object:Gem::Version 
-  version: 0.0.1
+  version: 0.1.0
 platform: ruby
 authors: 
 - Joe Damato
@@ -9,11 +9,11 @@ autorequire:
 bindir: bin
 cert_chain: []
 
-date: 2009-11-21 00:00:00 -08:00
+date: 2009-12-10 00:00:00 -08:00
 default_executable: 
 dependencies: []
 
-description: Ruby memory profiler gem
+description: Ruby memory profiler similar to bleak_house, but without patches to the Ruby VM
 email: ice799@gmail.com
 executables: []
 
@@ -22,10 +22,14 @@ extensions:
 extra_rdoc_files: []
 
 files: 
+- .gitignore
 - README
-- memprof.gemspec
-- ext/memprof.c
+- ext/bin_api.h
+- ext/elf.c
 - ext/extconf.rb
+- ext/mach.c
+- ext/memprof.c
+- memprof.gemspec
 has_rdoc: true
 homepage: http://github.com/ice799/memprof
 licenses: []
@@ -35,7 +39,6 @@ rdoc_options: []
 
 require_paths: 
 - lib
-- ext
 required_ruby_version: !ruby/object:Gem::Requirement 
   requirements: 
   - - ">="
@@ -54,6 +57,6 @@ rubyforge_project:
 rubygems_version: 1.3.5
 signing_key: 
 specification_version: 3
-summary: Ruby memory profiler gem
+summary: Ruby Memory Profiler
 test_files: []