Add support for custom hook classes with Firefox profiler integration

Author: joshuay03

README.md

@@ -126,6 +126,71 @@ ruby -r vernier -e 'Vernier.trace_retained(out: "irb_profile.json") { require "i
 > [!NOTE]
 > Retained-memory flamegraphs must be interpreted a little differently than a typical profiling flamegraph. In a retained-memory flamegraph, the x-axis represents a proportion of memory in bytes,  _not time or samples_ The topmost boxes on the y-axis represent the retained objects, with their stacktrace below; their width represents the percentage of overall retained memory each object occupies.
 
+### Hooks
+
+Hooks automatically add markers to profiles based on application events:
+
+```ruby
+Vernier.profile(hooks: [:activesupport, MyHook]) do
+  # code to profile
+end
+```
+
+#### Built-in hooks
+
+- `:activesupport` (`:rails`) - ActiveSupport notifications
+- `:memory_usage` - Memory usage tracking
+
+#### Custom hooks
+
+Define a class with `initialize(collector)`, `enable`, and `disable` methods:
+
+```ruby
+class MyHook
+  def initialize(collector)
+    @collector = collector
+  end
+
+  def enable
+    # Set up event listeners
+  end
+
+  def disable
+    # Clean up
+  end
+end
+```
+
+For Firefox profiler integration, add `firefox_marker_schema` or `firefox_counters` methods (see [Firefox profiler docs](https://profiler.firefox.com/docs/#/) for format details):
+
+```ruby
+def firefox_marker_schema
+  [{
+    name: "my_event",
+    display: ["marker-chart", "marker-table"],
+    tooltipLabel: "{marker.data.name}",
+    data: [
+      { key: "name", format: "string" }
+    ]
+  }]
+end
+
+def firefox_counters
+  {
+    name: "my_counter",
+    category: "Custom",
+    description: "My custom counter",
+    samples: {
+      time: [0, 1000, 2000],
+      count: [10, 20, 30],
+      length: 3
+    }
+  }
+end
+```
+
+See [`examples/custom_hook.rb`](examples/custom_hook.rb) for a complete example.
+
 ### Options
 
 | Option                | Middleware Param              | Description                                                   | Default (Middleware Default) |

examples/custom_hook.rb

@@ -0,0 +1,37 @@
+#!/usr/bin/env ruby
+
+require "vernier"
+
+class CustomHook
+  def initialize(collector)
+    @collector = collector
+    @events = []
+  end
+
+  def enable
+    # Simulate subscribing to application events
+    @thread = Thread.new do
+      3.times do |i|
+        sleep 0.03
+        start_time = @collector.current_time
+        sleep 0.01 # Simulate work
+        @collector.add_marker(
+          name: "custom_event",
+          start: start_time,
+          finish: @collector.current_time,
+          data: { type: "custom_event", event_id: i + 1 }
+        )
+      end
+    end
+  end
+
+  def disable
+    @thread&.join
+  end
+end
+
+result = Vernier.profile(hooks: [CustomHook]) do
+  sleep 0.15
+end
+
+puts "Profile complete with custom events"

lib/vernier/collector.rb

@@ -31,13 +31,17 @@ def initialize(mode, options = {})
     end
 
     private def add_hook(hook)
-      case hook.to_sym
+      case (hook.to_sym if hook.respond_to?(:to_sym))
       when :rails, :activesupport
         @hooks << Vernier::Hooks::ActiveSupport.new(self)
       when :memory_usage
         @hooks << Vernier::Hooks::MemoryUsage.new(self)
       else
-        warn "unknown hook: #{hook.inspect}"
+        if hook.respond_to?(:new)
+          @hooks << hook.new(self)
+        else
+          warn "unknown hook: #{hook.inspect}"
+        end
       end
     end
 

test/output/test_firefox.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 require "test_helper"
-require "firefox_test_helpers"
 
 class TestOutputFirefox < Minitest::Test
   include FirefoxTestHelpers

test/test_custom_hooks.rb

@@ -0,0 +1,122 @@
+# frozen_string_literal: true
+
+require "test_helper"
+
+class TestCustomHooks < Minitest::Test
+  include FirefoxTestHelpers
+
+  class SimpleHook
+    def initialize(collector)
+      @collector = collector
+    end
+
+    def enable
+      @start_time = @collector.current_time
+    end
+
+    def disable
+      @collector.add_marker(
+        name: "Simple Hook",
+        start: @start_time,
+        finish: @collector.current_time,
+        data: { type: "simple_hook" }
+      )
+    end
+  end
+
+  class HookWithSchema
+    def initialize(collector)
+      @collector = collector
+    end
+
+    def enable
+      @collector.add_marker(
+        name: "custom_event",
+        start: @collector.current_time,
+        finish: @collector.current_time + 1000,
+        data: { type: "custom_event", event_name: "test" }
+      )
+    end
+
+    def disable
+    end
+
+    def firefox_marker_schema
+      [{
+        name: "custom_event",
+        display: ["marker-chart", "marker-table"],
+        data: [{ key: "event_name", format: "string" }]
+      }]
+    end
+  end
+
+  class HookWithCounters
+    def initialize(collector)
+      @collector = collector
+    end
+
+    def enable
+    end
+
+    def disable
+    end
+
+    def firefox_counters
+      {
+        name: "test_counter",
+        category: "Test",
+        description: "Test counter",
+        samples: {
+          time: [0, 1000],
+          count: [10, 20],
+          length: 2
+        }
+      }
+    end
+  end
+
+  def test_custom_hook_class
+    result = Vernier.profile(hooks: [SimpleHook]) do
+      sleep 0.01
+    end
+
+    output = Vernier::Output::Firefox.new(result).output
+    assert_valid_firefox_profile(output)
+  end
+
+  def test_custom_hook_with_schema
+    result = Vernier.profile(hooks: [HookWithSchema]) do
+      sleep 0.01
+    end
+
+    output = Vernier::Output::Firefox.new(result).output
+    assert_valid_firefox_profile(output)
+
+    data = JSON.parse(output)
+    schemas = data.dig("meta", "markerSchema") || []
+    assert schemas.any? { |s| s["name"] == "custom_event" }
+  end
+
+  def test_mixed_hooks
+    result = Vernier.profile(hooks: [:memory_usage, SimpleHook]) do
+      sleep 0.01
+    end
+
+    output = Vernier::Output::Firefox.new(result).output
+    assert_valid_firefox_profile(output)
+    assert_equal 2, result.hooks.size
+  end
+
+  def test_custom_hook_with_counters
+    result = Vernier.profile(hooks: [HookWithCounters]) do
+      sleep 0.01
+    end
+
+    output = Vernier::Output::Firefox.new(result).output
+    assert_valid_firefox_profile(output)
+
+    data = JSON.parse(output)
+    counters = data["counters"]
+    assert counters.any? { |c| c["name"] == "test_counter" }
+  end
+end

test/test_rack_middleware.rb

@@ -1,7 +1,6 @@
 # frozen_string_literal: true
 
 require "test_helper"
-require "firefox_test_helpers"
 
 require "rack"