vernier 1.8.1 → 1.9.0

data/ext/vernier/vernier.cc

@@ -19,6 +19,7 @@
 #include "timestamp.hh"
 #include "periodic_thread.hh"
 #include "signal_safe_semaphore.hh"
+#include "stack_table.hh"
 
 #include "ruby/ruby.h"
 #include "ruby/encoding.h"
@@ -43,18 +44,13 @@
   RUBY_EVENT_FIBER_SWITCH | \
   RUBY_EVENT_THREAD_END
 
-#define sym(name) ID2SYM(rb_intern_const(name))
-
-// HACK: This isn't public, but the objspace ext uses it
-extern "C" size_t rb_obj_memsize_of(VALUE);
-
 using namespace std;
 
 VALUE rb_mVernier;
 static VALUE rb_cVernierResult;
 static VALUE rb_mVernierMarkerType;
 static VALUE rb_cVernierCollector;
-static VALUE rb_cStackTable;
+static VALUE rb_cTimeCollector;
 
 static VALUE sym_state, sym_gc_by, sym_fiber_id;
 
@@ -74,126 +70,6 @@ static const char *gvl_event_name(rb_event_flag_t event) {
     return "no-event";
 }
 
-struct FuncInfo {
-    static const char *label_cstr(VALUE frame) {
-        VALUE label = rb_profile_frame_full_label(frame);
-        // Currently (2025-03-22, Ruby 3.4.2) this occurs when an iseq method
-        // entry is replaced with a refinement
-        if (NIL_P(label)) return "(nil)";
-        return StringValueCStr(label);
-    }
-
-    static const char *file_cstr(VALUE frame) {
-        VALUE file = rb_profile_frame_absolute_path(frame);
-        if (NIL_P(file))
-            file = rb_profile_frame_path(frame);
-        if (NIL_P(file)) {
-            return "(nil)";
-        } else {
-            return StringValueCStr(file);
-        }
-    }
-
-    static int first_lineno_int(VALUE frame) {
-        VALUE first_lineno = rb_profile_frame_first_lineno(frame);
-        return NIL_P(first_lineno) ? 0 : FIX2INT(first_lineno);
-    }
-
-    FuncInfo(VALUE frame) :
-        label(label_cstr(frame)),
-        file(file_cstr(frame)),
-        first_lineno(first_lineno_int(frame)) { }
-
-    std::string label;
-    std::string file;
-    int first_lineno;
-};
-
-bool operator==(const FuncInfo& lhs, const FuncInfo& rhs) noexcept {
-    return
-        lhs.label == rhs.label &&
-        lhs.file == rhs.file &&
-        lhs.first_lineno == rhs.first_lineno;
-}
-
-struct Frame {
-    VALUE frame;
-    int line;
-};
-
-bool operator==(const Frame& lhs, const Frame& rhs) noexcept {
-    return lhs.frame == rhs.frame && lhs.line == rhs.line;
-}
-
-bool operator!=(const Frame& lhs, const Frame& rhs) noexcept {
-    return !(lhs == rhs);
-}
-
-namespace std {
-    template<>
-    struct hash<Frame>
-    {
-        std::size_t operator()(Frame const& s) const noexcept
-        {
-            return s.frame ^ s.line;
-        }
-    };
-}
-
-class RawSample {
-    public:
-
-    constexpr static int MAX_LEN = 2048;
-
-    private:
-
-    VALUE frames[MAX_LEN];
-    int lines[MAX_LEN];
-    int len;
-    int offset;
-    bool gc;
-
-    public:
-
-    RawSample() : len(0), gc(false), offset(0) { }
-
-    int size() const {
-        return len - offset;
-    }
-
-    Frame frame(int i) const {
-        int idx = len - i - 1;
-        if (idx < 0) throw std::out_of_range("VERNIER BUG: index out of range");
-        const Frame frame = {frames[idx], lines[idx]};
-        return frame;
-    }
-
-    void sample(int offset = 0) {
-        clear();
-
-        if (!ruby_native_thread_p()) {
-            return;
-        }
-
-        if (rb_during_gc()) {
-          gc = true;
-        } else {
-          len = rb_profile_frames(0, MAX_LEN, frames, lines);
-          this->offset = std::min(offset, len);
-        }
-    }
-
-    void clear() {
-        len = 0;
-        offset = 0;
-        gc = false;
-    }
-
-    bool empty() const {
-        return len <= offset;
-    }
-};
-
 // Based very loosely on the design of Gecko's SigHandlerCoordinator
 // This is used for communication between the profiler thread and the signal
 // handlers in the observed thread.
@@ -227,395 +103,6 @@ struct LiveSample {
     }
 };
 
-template <typename K>
-class IndexMap {
-    public:
-        std::unordered_map<K, int> to_idx;
-        std::vector<K> list;
-
-        const K& operator[](int i) const noexcept {
-            return list[i];
-        }
-
-        size_t size() const noexcept {
-            return list.size();
-        }
-
-        int index(const K key) {
-            auto it = to_idx.find(key);
-            if (it == to_idx.end()) {
-                int idx = list.size();
-                list.push_back(key);
-
-                auto result = to_idx.insert({key, idx});
-                it = result.first;
-            }
-
-            return it->second;
-        }
-
-        void clear() {
-            list.clear();
-            to_idx.clear();
-        }
-};
-
-struct StackTable {
-    private:
-
-    struct FrameWithInfo {
-        Frame frame;
-        FuncInfo info;
-    };
-
-    IndexMap<Frame> frame_map;
-
-    IndexMap<VALUE> func_map;
-    std::vector<FuncInfo> func_info_list;
-
-    struct StackNode {
-        std::unordered_map<Frame, int> children;
-        Frame frame;
-        int parent;
-        int index;
-
-        StackNode(Frame frame, int index, int parent) : frame(frame), index(index), parent(parent) {}
-
-        // root
-        StackNode() : frame(Frame{0, 0}), index(-1), parent(-1) {}
-    };
-
-    // This mutex guards the StackNodes only. The rest of the maps and vectors
-    // should be guarded by the GVL
-    std::mutex stack_mutex;
-
-    StackNode root_stack_node;
-    vector<StackNode> stack_node_list;
-    int stack_node_list_finalized_idx = 0;
-
-    StackNode *next_stack_node(StackNode *node, Frame frame) {
-        auto search = node->children.find(frame);
-        if (search == node->children.end()) {
-            // insert a new node
-            int next_node_idx = stack_node_list.size();
-            node->children[frame] = next_node_idx;
-            stack_node_list.emplace_back(
-                    frame,
-                    next_node_idx,
-                    node->index
-                    );
-            return &stack_node_list[next_node_idx];
-        } else {
-            int node_idx = search->second;
-            return &stack_node_list[node_idx];
-        }
-    }
-
-    public:
-
-    int stack_index(const RawSample &stack) {
-        if (stack.empty()) {
-            throw std::runtime_error("VERNIER BUG: empty stack");
-        }
-
-        const std::lock_guard<std::mutex> lock(stack_mutex);
-
-        StackNode *node = &root_stack_node;
-        for (int i = 0; i < stack.size(); i++) {
-            Frame frame = stack.frame(i);
-            node = next_stack_node(node, frame);
-        }
-        return node->index;
-    }
-
-    int stack_parent(int stack_idx) {
-        const std::lock_guard<std::mutex> lock(stack_mutex);
-        if (stack_idx < 0 || stack_idx >= stack_node_list.size()) {
-            return -1;
-        } else {
-            return stack_node_list[stack_idx].parent;
-        }
-    }
-
-    int stack_frame(int stack_idx) {
-        const std::lock_guard<std::mutex> lock(stack_mutex);
-        if (stack_idx < 0 || stack_idx >= stack_node_list.size()) {
-            return -1;
-        } else {
-            return frame_map.index(stack_node_list[stack_idx].frame);
-        }
-    }
-
-    // Converts Frames from stacks other tables. "Symbolicates" the frames
-    // which allocates.
-    void finalize() {
-        {
-            const std::lock_guard<std::mutex> lock(stack_mutex);
-            for (int i = stack_node_list_finalized_idx; i < stack_node_list.size(); i++) {
-                const auto &stack_node = stack_node_list[i];
-                frame_map.index(stack_node.frame);
-                func_map.index(stack_node.frame.frame);
-                stack_node_list_finalized_idx = i;
-            }
-        }
-
-        for (int i = func_info_list.size(); i < func_map.size(); i++) {
-            const auto &func = func_map[i];
-            // must not hold a mutex here
-            func_info_list.push_back(FuncInfo(func));
-        }
-    }
-
-    void mark_frames() {
-        const std::lock_guard<std::mutex> lock(stack_mutex);
-
-        for (auto stack_node: stack_node_list) {
-            rb_gc_mark(stack_node.frame.frame);
-        }
-    }
-
-    StackNode *convert_stack(StackTable &other, int original_idx) {
-        if (original_idx < 0) {
-            return &root_stack_node;
-        }
-
-        StackNode &original_node = other.stack_node_list[original_idx];
-        StackNode *parent_node = convert_stack(other, original_node.parent);
-        StackNode *node = next_stack_node(parent_node, original_node.frame);
-
-        return node;
-    }
-
-    static VALUE stack_table_convert(VALUE self, VALUE other, VALUE original_stack);
-
-    static VALUE stack_table_stack_count(VALUE self);
-    static VALUE stack_table_frame_count(VALUE self);
-    static VALUE stack_table_func_count(VALUE self);
-
-    static VALUE stack_table_frame_line_no(VALUE self, VALUE idxval);
-    static VALUE stack_table_frame_func_idx(VALUE self, VALUE idxval);
-    static VALUE stack_table_func_name(VALUE self, VALUE idxval);
-    static VALUE stack_table_func_filename(VALUE self, VALUE idxval);
-    static VALUE stack_table_func_first_lineno(VALUE self, VALUE idxval);
-
-    friend class SampleTranslator;
-};
-
-static void
-stack_table_mark(void *data) {
-    StackTable *stack_table = static_cast<StackTable *>(data);
-    stack_table->mark_frames();
-}
-
-static void
-stack_table_free(void *data) {
-    StackTable *stack_table = static_cast<StackTable *>(data);
-    delete stack_table;
-}
-
-static const rb_data_type_t rb_stack_table_type = {
-    .wrap_struct_name = "vernier/stack_table",
-    .function = {
-        //.dmemsize = rb_collector_memsize,
-        .dmark = stack_table_mark,
-        .dfree = stack_table_free,
-    },
-};
-
-static VALUE
-stack_table_new(VALUE self) {
-    StackTable *stack_table = new StackTable();
-    VALUE obj = TypedData_Wrap_Struct(self, &rb_stack_table_type, stack_table);
-    return obj;
-}
-
-static StackTable *get_stack_table(VALUE obj) {
-    StackTable *stack_table;
-    TypedData_Get_Struct(obj, StackTable, &rb_stack_table_type, stack_table);
-    return stack_table;
-}
-
-static VALUE
-stack_table_current_stack(int argc, VALUE *argv, VALUE self) {
-    int offset;
-    VALUE offset_v;
-
-    rb_scan_args(argc, argv, "01", &offset_v);
-    if (argc > 0) {
-        offset = NUM2INT(offset_v) + 1;
-    } else {
-        offset = 1;
-    }
-
-    StackTable *stack_table = get_stack_table(self);
-    RawSample stack;
-    stack.sample(offset);
-    int stack_index = stack_table->stack_index(stack);
-    return INT2NUM(stack_index);
-}
-
-static VALUE
-stack_table_stack_parent_idx(VALUE self, VALUE idxval) {
-    StackTable *stack_table = get_stack_table(self);
-    int idx = NUM2INT(idxval);
-    int parent_idx = stack_table->stack_parent(idx);
-    if (parent_idx < 0) {
-        return Qnil;
-    } else {
-        return INT2NUM(parent_idx);
-    }
-}
-
-static VALUE
-stack_table_stack_frame_idx(VALUE self, VALUE idxval) {
-    StackTable *stack_table = get_stack_table(self);
-    //stack_table->finalize();
-    int idx = NUM2INT(idxval);
-    int frame_idx = stack_table->stack_frame(idx);
-    return frame_idx < 0 ? Qnil : INT2NUM(frame_idx);
-}
-
-VALUE
-StackTable::stack_table_stack_count(VALUE self) {
-    StackTable *stack_table = get_stack_table(self);
-    int count;
-    {
-        const std::lock_guard<std::mutex> lock(stack_table->stack_mutex);
-        count = stack_table->stack_node_list.size();
-    }
-    return INT2NUM(count);
-}
-
-VALUE
-StackTable::stack_table_convert(VALUE self, VALUE original_tableval, VALUE original_idxval) {
-    StackTable *stack_table = get_stack_table(self);
-    StackTable *original_table = get_stack_table(original_tableval);
-    int original_idx = NUM2INT(original_idxval);
-
-    int original_size;
-    {
-        const std::lock_guard<std::mutex> lock(original_table->stack_mutex);
-        original_size = original_table->stack_node_list.size();
-    }
-
-    if (original_idx >= original_size || original_idx < 0) {
-        rb_raise(rb_eRangeError, "index out of range");
-    }
-
-    int result_idx;
-    {
-        const std::lock_guard<std::mutex> lock1(stack_table->stack_mutex);
-        const std::lock_guard<std::mutex> lock2(original_table->stack_mutex);
-        StackNode *node = stack_table->convert_stack(*original_table, original_idx);
-        result_idx = node->index;
-    }
-    return INT2NUM(result_idx);
-}
-
-VALUE
-StackTable::stack_table_frame_count(VALUE self) {
-    StackTable *stack_table = get_stack_table(self);
-    stack_table->finalize();
-    int count = stack_table->frame_map.size();
-    return INT2NUM(count);
-}
-
-VALUE
-StackTable::stack_table_func_count(VALUE self) {
-    StackTable *stack_table = get_stack_table(self);
-    stack_table->finalize();
-    int count = stack_table->func_map.size();
-    return INT2NUM(count);
-}
-
-VALUE
-StackTable::stack_table_frame_line_no(VALUE self, VALUE idxval) {
-    StackTable *stack_table = get_stack_table(self);
-    stack_table->finalize();
-    int idx = NUM2INT(idxval);
-    if (idx < 0 || idx >= stack_table->frame_map.size()) {
-        return Qnil;
-    } else {
-        const auto &frame = stack_table->frame_map[idx];
-        return INT2NUM(frame.line);
-    }
-}
-
-VALUE
-StackTable::stack_table_frame_func_idx(VALUE self, VALUE idxval) {
-    StackTable *stack_table = get_stack_table(self);
-    stack_table->finalize();
-    int idx = NUM2INT(idxval);
-    if (idx < 0 || idx >= stack_table->frame_map.size()) {
-        return Qnil;
-    } else {
-        const auto &frame = stack_table->frame_map[idx];
-        int func_idx = stack_table->func_map.index(frame.frame);
-        return INT2NUM(func_idx);
-    }
-}
-
-VALUE
-StackTable::stack_table_func_name(VALUE self, VALUE idxval) {
-    StackTable *stack_table = get_stack_table(self);
-    stack_table->finalize();
-    int idx = NUM2INT(idxval);
-    auto &table = stack_table->func_info_list;
-    if (idx < 0 || idx >= table.size()) {
-        return Qnil;
-    } else {
-        const auto &func_info = table[idx];
-        const std::string &label = func_info.label;
-
-        // Ruby constants are in an arbitrary (ASCII compatible) encoding and
-        // method names are in an arbitrary (ASCII compatible) encoding. These
-        // can be mixed in the same program.
-        //
-        // However, by this point we've lost the chain of what the correct
-        // encoding should be. Oops!
-        //
-        // Instead we'll just guess at UTF-8 which should satisfy most. It won't
-        // necessarily be valid but that can be scrubbed on the Ruby side.
-        //
-        // In the future we might keep class and method name separate for
-        // longer, preserve encodings, and defer formatting to the Ruby side.
-        return rb_enc_interned_str(label.c_str(), label.length(), rb_utf8_encoding());
-    }
-}
-
-VALUE
-StackTable::stack_table_func_filename(VALUE self, VALUE idxval) {
-    StackTable *stack_table = get_stack_table(self);
-    stack_table->finalize();
-    int idx = NUM2INT(idxval);
-    auto &table = stack_table->func_info_list;
-    if (idx < 0 || idx >= table.size()) {
-        return Qnil;
-    } else {
-        const auto &func_info = table[idx];
-        const std::string &filename = func_info.file;
-
-        // Technically filesystems are binary and then Ruby interprets that as
-        // default_external encoding. But to keep things simple for now we are
-        // going to assume UTF-8.
-        return rb_enc_interned_str(filename.c_str(), filename.length(), rb_utf8_encoding());
-    }
-}
-
-VALUE
-StackTable::stack_table_func_first_lineno(VALUE self, VALUE idxval) {
-    StackTable *stack_table = get_stack_table(self);
-    stack_table->finalize();
-    int idx = NUM2INT(idxval);
-    auto &table = stack_table->func_info_list;
-    if (idx < 0 || idx >= table.size()) {
-        return Qnil;
-    } else {
-        const auto &func_info = table[idx];
-        return INT2NUM(func_info.first_lineno);
-    }
-}
-
 class SampleTranslator {
     public:
         int last_stack_index;
@@ -956,8 +443,7 @@ class Thread {
         unique_ptr<MarkerTable> markers;
 
         // FIXME: don't use pthread at start
-        Thread(State state, pthread_t pthread_id, VALUE ruby_thread) : pthread_id(pthread_id), ruby_thread(ruby_thread), state(state), stack_on_suspend_idx(-1) {
-            ruby_thread_id = rb_obj_id(ruby_thread);
+        Thread(State state, pthread_t pthread_id, VALUE ruby_thread, VALUE ruby_thread_id) : pthread_id(pthread_id), ruby_thread(ruby_thread), state(state), stack_on_suspend_idx(-1), ruby_thread_id(ruby_thread_id) {
             //ruby_thread_id = ULL2NUM(ruby_thread);
             native_tid = get_native_thread_id();
             started_at = state_changed_at = TimeStamp::Now();
@@ -989,7 +475,7 @@ class Thread {
             markers->record(Marker::Type::MARKER_FIBER_SWITCH, stack_idx, { .fiber_data = { .fiber_id = fiber_id } });
         }
 
-        void set_state(State new_state) {
+        void set_state(State new_state, bool current = true) {
             if (state == Thread::State::STOPPED) {
                 return;
             }
@@ -1117,9 +603,43 @@ class ThreadTable {
         }
 
     private:
-        void set_state(Thread::State new_state, VALUE th) {
+        Thread* find_thread(VALUE th) {
+            // Assumes lock is already held by caller
+            for (auto &threadptr : list) {
+                if (thread_equal(th, threadptr->ruby_thread)) {
+                    return threadptr.get();
+                }
+            }
+            return nullptr;
+        }
+
+        Thread& find_or_create_thread(VALUE th, Thread::State initial_state) {
+            {
+                const std::lock_guard<std::mutex> lock(mutex);
+                Thread* existing = find_thread(th);
+                if (existing) {
+                    return *existing;
+                }
+            }
+
+            // Thread not found, compute ID outside the lock
+            VALUE ruby_thread_id = rb_obj_id(th);
+
             const std::lock_guard<std::mutex> lock(mutex);
+            // Check again in case another thread created it while we were unlocked
+            Thread* existing = find_thread(th);
+            if (existing) {
+                return *existing;
+            }
+
+            //fprintf(stderr, "NEW THREAD: th: %p, state: %i\n", th, initial_state);
+            auto new_thread = std::make_unique<Thread>(initial_state, pthread_self(), th, ruby_thread_id);
+            Thread& thread_ref = *new_thread;
+            list.push_back(std::move(new_thread));
+            return thread_ref;
+        }
 
+        void set_state(Thread::State new_state, VALUE th) {
             //cerr << "set state=" << new_state << " thread=" << gettid() << endl;
 
             pid_t native_tid = get_native_thread_id();
@@ -1127,35 +647,28 @@ class ThreadTable {
 
             //fprintf(stderr, "th %p (tid: %i) from %s to %s\n", (void *)th, native_tid, gvl_event_name(state), gvl_event_name(new_state));
 
-            for (auto &threadptr : list) {
-                auto &thread = *threadptr;
-                if (thread_equal(th, thread.ruby_thread)) {
-                    if (new_state == Thread::State::SUSPENDED || new_state == Thread::State::READY && (thread.state != Thread::State::SUSPENDED)) {
+            Thread& thread = find_or_create_thread(th, new_state);
 
-                        RawSample sample;
-                        sample.sample();
+            const std::lock_guard<std::mutex> lock(mutex);
 
-                        thread.stack_on_suspend_idx = thread.translator.translate(frame_list, sample);
-                        //cerr << gettid() << " suspended! Stack size:" << thread.stack_on_suspend.size() << endl;
-                    }
+            if (new_state == Thread::State::SUSPENDED || new_state == Thread::State::READY && (thread.state != Thread::State::SUSPENDED)) {
 
-                    thread.set_state(new_state);
+                RawSample sample;
+                sample.sample();
 
-                    if (thread.state == Thread::State::RUNNING) {
-                        thread.pthread_id = pthread_self();
-                        thread.native_tid = get_native_thread_id();
-                    } else {
-                        thread.pthread_id = 0;
-                        thread.native_tid = 0;
-                    }
+                thread.stack_on_suspend_idx = thread.translator.translate(frame_list, sample);
+                //cerr << gettid() << " suspended! Stack size:" << thread.stack_on_suspend.size() << endl;
+            }
 
+            thread.set_state(new_state);
 
-                    return;
-                }
+            if (thread.state == Thread::State::RUNNING) {
+                thread.pthread_id = pthread_self();
+                thread.native_tid = get_native_thread_id();
+            } else {
+                thread.pthread_id = 0;
+                thread.native_tid = 0;
             }
-
-            //fprintf(stderr, "NEW THREAD: th: %p, state: %i\n", th, new_state);
-            list.push_back(std::make_unique<Thread>(new_state, pthread_self(), th));
         }
 
         bool thread_equal(VALUE a, VALUE b) {
@@ -1232,209 +745,6 @@ class BaseCollector {
     };
 };
 
-class CustomCollector : public BaseCollector {
-    SampleList samples;
-
-    void sample() {
-        RawSample sample;
-        sample.sample();
-        int stack_index = stack_table->stack_index(sample);
-
-        samples.record_sample(stack_index, TimeStamp::Now(), CATEGORY_NORMAL);
-    }
-
-    VALUE stop() {
-        BaseCollector::stop();
-
-        stack_table->finalize();
-
-        VALUE result = build_collector_result();
-
-        reset();
-
-        return result;
-    }
-
-    VALUE build_collector_result() {
-        VALUE result = BaseCollector::build_collector_result();
-
-        VALUE threads = rb_hash_new();
-        rb_ivar_set(result, rb_intern("@threads"), threads);
-
-        VALUE thread_hash = rb_hash_new();
-        samples.write_result(thread_hash);
-
-        rb_hash_aset(threads, ULL2NUM(0), thread_hash);
-        rb_hash_aset(thread_hash, sym("tid"), ULL2NUM(0));
-        rb_hash_aset(thread_hash, sym("name"), rb_str_new_cstr("custom"));
-        rb_hash_aset(thread_hash, sym("started_at"), ULL2NUM(started_at.nanoseconds()));
-
-        return result;
-    }
-
-    public:
-
-    CustomCollector(VALUE stack_table) : BaseCollector(stack_table) { }
-};
-
-class RetainedCollector : public BaseCollector {
-    void reset() {
-        object_frames.clear();
-        object_list.clear();
-
-        BaseCollector::reset();
-    }
-
-    void record(VALUE obj) {
-        RawSample sample;
-        sample.sample();
-        if (sample.empty()) {
-            // During thread allocation we allocate one object without a frame
-            // (as of Ruby 3.3)
-            // Ideally we'd allow empty samples to be represented
-            return;
-        }
-        int stack_index = stack_table->stack_index(sample);
-
-        object_list.push_back(obj);
-        object_frames.emplace(obj, stack_index);
-    }
-
-    std::unordered_map<VALUE, int> object_frames;
-    std::vector<VALUE> object_list;
-
-    VALUE tp_newobj = Qnil;
-    VALUE tp_freeobj = Qnil;
-
-    static void newobj_i(VALUE tpval, void *data) {
-        RetainedCollector *collector = static_cast<RetainedCollector *>(data);
-        rb_trace_arg_t *tparg = rb_tracearg_from_tracepoint(tpval);
-        VALUE obj = rb_tracearg_object(tparg);
-
-        collector->record(obj);
-    }
-
-    static void freeobj_i(VALUE tpval, void *data) {
-        RetainedCollector *collector = static_cast<RetainedCollector *>(data);
-        rb_trace_arg_t *tparg = rb_tracearg_from_tracepoint(tpval);
-        VALUE obj = rb_tracearg_object(tparg);
-
-        collector->object_frames.erase(obj);
-    }
-
-    public:
-
-    RetainedCollector(VALUE stack_table) : BaseCollector(stack_table) { }
-
-    bool start() {
-        if (!BaseCollector::start()) {
-            return false;
-        }
-
-        tp_newobj = rb_tracepoint_new(0, RUBY_INTERNAL_EVENT_NEWOBJ, newobj_i, this);
-        tp_freeobj = rb_tracepoint_new(0, RUBY_INTERNAL_EVENT_FREEOBJ, freeobj_i, this);
-
-        rb_tracepoint_enable(tp_newobj);
-        rb_tracepoint_enable(tp_freeobj);
-
-        return true;
-    }
-
-    VALUE stop() {
-        BaseCollector::stop();
-
-        // GC before we start turning stacks into strings
-        rb_gc();
-
-        // Stop tracking any more new objects, but we'll continue tracking free'd
-        // objects as we may be able to free some as we remove our own references
-        // to stack frames.
-        rb_tracepoint_disable(tp_newobj);
-        tp_newobj = Qnil;
-
-        stack_table->finalize();
-
-        // We should have collected info for all our frames, so no need to continue
-        // marking them
-        // FIXME: previously here we cleared the list of frames so we would stop
-        // marking them. Maybe now we should set a flag so that we stop marking them
-
-        // GC again
-        rb_gc();
-
-        rb_tracepoint_disable(tp_freeobj);
-        tp_freeobj = Qnil;
-
-        VALUE result = build_collector_result();
-
-        reset();
-
-        return result;
-    }
-
-    VALUE build_collector_result() {
-        RetainedCollector *collector = this;
-        StackTable &frame_list = *collector->stack_table;
-
-        VALUE result = BaseCollector::build_collector_result();
-
-        VALUE threads = rb_hash_new();
-        rb_ivar_set(result, rb_intern("@threads"), threads);
-        VALUE thread_hash = rb_hash_new();
-        rb_hash_aset(threads, ULL2NUM(0), thread_hash);
-
-        rb_hash_aset(thread_hash, sym("tid"), ULL2NUM(0));
-        VALUE samples = rb_ary_new();
-        rb_hash_aset(thread_hash, sym("samples"), samples);
-        VALUE weights = rb_ary_new();
-        rb_hash_aset(thread_hash, sym("weights"), weights);
-
-        rb_hash_aset(thread_hash, sym("name"), rb_str_new_cstr("retained memory"));
-        rb_hash_aset(thread_hash, sym("started_at"), ULL2NUM(collector->started_at.nanoseconds()));
-
-        for (auto& obj: collector->object_list) {
-            const auto search = collector->object_frames.find(obj);
-            if (search != collector->object_frames.end()) {
-                int stack_index = search->second;
-
-                rb_ary_push(samples, INT2NUM(stack_index));
-                rb_ary_push(weights, INT2NUM(rb_obj_memsize_of(obj)));
-            }
-        }
-
-        return result;
-    }
-
-    void mark() {
-        // We don't mark the objects, but we MUST mark the frames, otherwise they
-        // can be garbage collected.
-        // When we stop collection we will stringify the remaining frames, and then
-        // clear them from the set, allowing them to be removed from out output.
-        stack_table->mark_frames();
-        rb_gc_mark(stack_table_value);
-
-        rb_gc_mark(tp_newobj);
-        rb_gc_mark(tp_freeobj);
-    }
-
-    void compact() {
-        RetainedCollector *collector = this;
-        for (auto& obj: collector->object_list) {
-            VALUE reloc_obj = rb_gc_location(obj);
-            
-            const auto search = collector->object_frames.find(obj);
-            if (search != collector->object_frames.end()) {
-                int stack_index = search->second;
-                
-                collector->object_frames.erase(search);
-                collector->object_frames.emplace(reloc_obj, stack_index);
-            }
-            
-            obj = reloc_obj;
-        }
-    }
-};
-
 class GlobalSignalHandler {
     static LiveSample *live_sample;
 
@@ -1782,11 +1092,11 @@ class TimeCollector : public BaseCollector {
     VALUE build_collector_result() {
         VALUE result = BaseCollector::build_collector_result();
 
-        rb_ivar_set(result, rb_intern("@gc_markers"), this->gc_markers.to_array());
-
         VALUE threads = rb_hash_new();
         rb_ivar_set(result, rb_intern("@threads"), threads);
 
+        rb_ivar_set(result, rb_intern("@gc_markers"), this->gc_markers.to_array());
+
         for (const auto& thread: this->threads.list) {
             VALUE hash = rb_hash_new();
             thread->samples.write_result(hash);
@@ -1840,10 +1150,9 @@ collector_compact(void *data) {
 static const rb_data_type_t rb_collector_type = {
     .wrap_struct_name = "vernier/collector",
     .function = {
-        //.dmemsize = rb_collector_memsize,
         .dmark = collector_mark,
         .dfree = collector_free,
-        .dsize = NULL,
+        .dsize = nullptr,
         .dcompact = collector_compact,
     },
 };
@@ -1873,30 +1182,11 @@ collector_stop(VALUE self) {
     return result;
 }
 
-static VALUE
-collector_sample(VALUE self) {
-    auto *collector = get_collector(self);
-
-    collector->sample();
-    return Qtrue;
-}
-
-static VALUE
-collector_stack_table(VALUE self) {
-    auto *collector = get_collector(self);
-
-    return collector->stack_table_value;
-}
-
 static VALUE collector_new(VALUE self, VALUE mode, VALUE options) {
     BaseCollector *collector;
 
-    VALUE stack_table = stack_table_new(rb_cStackTable);
-    if (mode == sym("retained")) {
-        collector = new RetainedCollector(stack_table);
-    } else if (mode == sym("custom")) {
-        collector = new CustomCollector(stack_table);
-    } else if (mode == sym("wall")) {
+    VALUE stack_table = StackTable::stack_table_new();
+    if (mode == sym("wall")) {
         VALUE intervalv = rb_hash_aref(options, sym("interval"));
         TimeStamp interval;
         if (NIL_P(intervalv)) {
@@ -1920,6 +1210,7 @@ static VALUE collector_new(VALUE self, VALUE mode, VALUE options) {
         rb_raise(rb_eArgError, "invalid mode");
     }
     VALUE obj = TypedData_Wrap_Struct(self, &rb_collector_type, collector);
+    rb_ivar_set(obj, rb_intern("@stack_table"), stack_table);
     rb_funcall(obj, rb_intern("initialize"), 2, mode, options);
     return obj;
 }
@@ -1956,31 +1247,16 @@ Init_vernier(void)
   rb_mVernierMarkerType = rb_define_module_under(rb_mVernierMarker, "Type");
 
   rb_cVernierCollector = rb_define_class_under(rb_mVernier, "Collector", rb_cObject);
-  rb_undef_alloc_func(rb_cVernierCollector);
-  rb_define_singleton_method(rb_cVernierCollector, "_new", collector_new, 2);
-  rb_define_method(rb_cVernierCollector, "start", collector_start, 0);
-  rb_define_method(rb_cVernierCollector, "sample", collector_sample, 0);
-  rb_define_method(rb_cVernierCollector, "stack_table", collector_stack_table, 0);
-  rb_define_private_method(rb_cVernierCollector, "finish",  collector_stop, 0);
-
-  rb_cStackTable = rb_define_class_under(rb_mVernier, "StackTable", rb_cObject);
-  rb_undef_alloc_func(rb_cStackTable);
-  rb_define_singleton_method(rb_cStackTable, "new", stack_table_new, 0);
-  rb_define_method(rb_cStackTable, "current_stack", stack_table_current_stack, -1);
-  rb_define_method(rb_cStackTable, "convert", StackTable::stack_table_convert, 2);
-  rb_define_method(rb_cStackTable, "stack_parent_idx", stack_table_stack_parent_idx, 1);
-  rb_define_method(rb_cStackTable, "stack_frame_idx", stack_table_stack_frame_idx, 1);
-  rb_define_method(rb_cStackTable, "frame_line_no", StackTable::stack_table_frame_line_no, 1);
-  rb_define_method(rb_cStackTable, "frame_func_idx", StackTable::stack_table_frame_func_idx, 1);
-  rb_define_method(rb_cStackTable, "func_name", StackTable::stack_table_func_name, 1);
-  rb_define_method(rb_cStackTable, "func_filename", StackTable::stack_table_func_filename, 1);
-  rb_define_method(rb_cStackTable, "func_first_lineno", StackTable::stack_table_func_first_lineno, 1);
-  rb_define_method(rb_cStackTable, "stack_count", StackTable::stack_table_stack_count, 0);
-  rb_define_method(rb_cStackTable, "frame_count", StackTable::stack_table_frame_count, 0);
-  rb_define_method(rb_cStackTable, "func_count", StackTable::stack_table_func_count, 0);
+  rb_cTimeCollector = rb_define_class_under(rb_cVernierCollector, "TimeCollector", rb_cVernierCollector);
+  rb_undef_alloc_func(rb_cTimeCollector);
+  rb_define_singleton_method(rb_cTimeCollector, "new", collector_new, 2);
+  rb_define_method(rb_cTimeCollector, "start", collector_start, 0);
+  rb_define_private_method(rb_cTimeCollector, "finish",  collector_stop, 0);
 
   Init_consts(rb_mVernierMarkerPhase);
   Init_memory();
+  Init_stack_table();
+  Init_heap_tracker();
 
   //static VALUE gc_hook = Data_Wrap_Struct(rb_cObject, collector_mark, NULL, &_collector);
   //rb_global_variable(&gc_hook);