Add Hook for ActiveSupport::Notifications

examples/rails.rb

@@ -110,15 +110,7 @@ def show
 # warm up
 make_request.call
 
-Vernier.trace(out: "rails.json") do |collector|
-  ActiveSupport::Notifications.monotonic_subscribe do |name, start, finish, id, payload|
-    collector.add_marker(
-      name:,
-      start: (start * 1_000_000_000).to_i,
-      finish: (finish * 1_000_000_000).to_i,
-    )
-  end
-
+Vernier.trace(out: "rails.json", hooks: [:rails]) do |collector|
   1000.times do
     make_request.call
   end

ext/vernier/vernier.cc

@@ -1620,7 +1620,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_funcall(obj, rb_intern("initialize"), 1, mode);
+    rb_funcall(obj, rb_intern("initialize"), 2, mode, options);
     return obj;
 }
 

lib/vernier.rb

@@ -3,15 +3,16 @@
 require_relative "vernier/version"
 require_relative "vernier/collector"
 require_relative "vernier/result"
+require_relative "vernier/hooks"
 require_relative "vernier/vernier"
 require_relative "vernier/output/firefox"
 require_relative "vernier/output/top"
 
 module Vernier
   class Error < StandardError; end
 
-  def self.trace(mode: :wall, out: nil, interval: nil)
-    collector = Vernier::Collector.new(mode, { interval: })
+  def self.trace(mode: :wall, out: nil, interval: nil, hooks: nil)
+    collector = Vernier::Collector.new(mode, { interval:, hooks: })
     collector.start
 
     result = nil

lib/vernier/collector.rb

@@ -4,9 +4,28 @@
 
 module Vernier
   class Collector
-    def initialize(mode)
+    def initialize(mode, options={})
       @mode = mode
       @markers = []
+      @hooks = []
+
+      if options[:hooks]
+        Array(options[:hooks]).each do |hook|
+          add_hook(hook)
+        end
+      end
+      @hooks.each do |hook|
+        hook.enable
+      end
+    end
+
+    private def add_hook(hook)
+      case hook
+      when :rails, :activesupport
+        @hooks << Vernier::Hooks::ActiveSupport.new(self)
+      else
+        warn "Unknown hook: #{hook}"
+      end
     end
 
     ##
@@ -47,6 +66,8 @@ def record_interval(category, name = category)
     def stop
       result = finish
 
+      result.hooks = @hooks
+
       end_time = Process.clock_gettime(Process::CLOCK_REALTIME, :nanosecond)
       result.pid = Process.pid
       result.end_time = end_time

lib/vernier/hooks.rb

@@ -0,0 +1,7 @@
+# frozen_string_literal: true
+
+module Vernier
+  module Hooks
+    autoload :ActiveSupport, "vernier/hooks/active_support"
+  end
+end

lib/vernier/hooks/active_support.rb

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module Vernier
+  module Hooks
+    class ActiveSupport
+      def initialize(collector)
+        @collector = collector
+      end
+
+      def enable
+        require "active_support"
+        @subscription = ::ActiveSupport::Notifications.monotonic_subscribe do |name, start, finish, id, payload|
+          data = { type: name }
+          data.update(payload)
+          @collector.add_marker(
+            name: name,
+            start: (start * 1_000_000_000.0).to_i,
+            finish: (finish * 1_000_000_000.0).to_i,
+            data: data
+          )
+        end
+      end
+
+      def disable
+        ActiveSupport::Notifications.unsubscribe(@subscription)
+        @subscription = nil
+      end
+
+      FIREFOX_MARKER_SCHEMA = Ractor.make_shareable([
+        {
+          name: "sql.active_record",
+          display: [ "marker-chart", "marker-table" ],
+          data: [
+            {
+              key: "sql",
+              format: "string"
+            },
+            {
+              key: "name",
+              format: "string"
+            },
+            {
+              key: "type_casted_binds",
+              label: "binds",
+              format: "string"
+            }
+          ]
+        },
+        {
+          name: "instantiation.active_record",
+          display: [ "marker-chart", "marker-table" ],
+          data: [
+            {
+              key: "record_count",
+              format: "integer"
+            },
+            {
+              key: "class_name",
+              format: "string"
+            }
+          ]
+        },
+        {
+          name: "process_action.action_controller",
+          display: [ "marker-chart", "marker-table" ],
+          data: [
+            {
+              key: "controller",
+              format: "string"
+            },
+            {
+              key: "action",
+              format: "string"
+            },
+            {
+              key: "status",
+              format: "integer"
+            },
+            {
+              key: "path",
+              format: "string"
+            },
+            {
+              key: "method",
+              format: "string"
+            }
+          ]
+        }
+      ])
+
+      def firefox_marker_schema
+        FIREFOX_MARKER_SCHEMA
+      end
+    end
+  end
+end

lib/vernier/output/firefox.rb

@@ -141,6 +141,12 @@ def data
       end
 
       def marker_schema
+        hook_additions = profile.hooks.flat_map do |hook|
+          if hook.respond_to?(:firefox_marker_schema)
+            hook.firefox_marker_schema
+          end
+        end.compact
+
         [
           {
             name: "GVL_THREAD_RESUMED",
@@ -151,7 +157,8 @@ def marker_schema
                 value: "The thread has acquired the GVL and is executing"
               }
             ]
-          }
+          },
+          *hook_additions
         ]
       end
 

lib/vernier/result.rb

@@ -3,6 +3,8 @@ class Result
     attr_reader :stack_table, :frame_table, :func_table
     attr_reader :markers
 
+    attr_accessor :hooks
+
     attr_accessor :pid, :end_time
     attr_accessor :threads
     attr_accessor :meta