rails_trace_viewer 0.1.0

data/lib/rails_trace_viewer/broadcaster.rb

@@ -0,0 +1,13 @@
+module RailsTraceViewer
+  class Broadcaster
+    def self.emit(trace_id, node)
+      ActionCable.server.broadcast(
+        "rails_trace_viewer",
+        {
+          trace_id: trace_id,
+          node: node
+        }
+      )
+    end
+  end
+end

data/lib/rails_trace_viewer/collector.rb

@@ -0,0 +1,151 @@
+module RailsTraceViewer
+  class Collector
+    TRACE_TREES = {}
+
+    @sweeper_started = false
+    @mutex = Mutex.new
+
+    class << self
+      def start_sweeper!
+        return if @sweeper_started
+
+        @sweeper_started = true
+
+        Thread.new do
+          Thread.current.name = "RailsTraceViewer-OrphanSweeper"
+          loop do
+            sleep 3
+            sweep_all_orphans
+          end
+        end
+      end
+
+      def start_trace(trace_id)
+        @mutex.synchronize do
+          TRACE_TREES[trace_id] ||= {
+            root: nil,
+            index: {},
+            orphans: {}
+          }
+        end
+      end
+
+      def add_node(trace_id, node)
+        after_broadcasts = []
+
+        @mutex.synchronize do
+          after_broadcasts = process_node(trace_id, node)
+        end
+
+        after_broadcasts.each(&:call)
+      end
+
+      def process_node(trace_id, node)
+        after = []
+        tree = TRACE_TREES[trace_id]
+
+        unless tree
+           tree = { root: nil, index: {}, orphans: {} }
+           TRACE_TREES[trace_id] = tree
+        end
+
+        id        = node[:id]
+        parent_id = node[:parent_id]
+
+        tree[:index][id] = node
+
+        if parent_id.nil?
+          tree[:root] ||= node
+          after << -> { Broadcaster.emit(trace_id, node) }
+          after.concat attach_waiting_children(tree, trace_id, id, node)
+
+        else
+          parent = tree[:index][parent_id]
+
+          if parent
+            parent[:children] ||= []
+            parent[:children] << node
+
+            after << -> { Broadcaster.emit(trace_id, node) }
+            after.concat attach_waiting_children(tree, trace_id, id, node)
+          elsif node[:type] == "job_perform"
+            after << -> { Broadcaster.emit(trace_id, node) }
+            
+            after.concat attach_waiting_children(tree, trace_id, id, node)
+          else
+            tree[:orphans][parent_id] ||= []
+            tree[:orphans][parent_id] << {
+              node: node,
+              time: Process.clock_gettime(Process::CLOCK_MONOTONIC)
+            }
+          end
+        end
+
+        after
+      end
+
+      def attach_waiting_children(tree, trace_id, parent_id, parent_node)
+        return [] unless tree[:orphans][parent_id]
+
+        after = []
+
+        tree[:orphans][parent_id].each do |entry|
+          child = entry[:node]
+          parent_node[:children] ||= []
+          parent_node[:children] << child
+
+          after << -> { Broadcaster.emit(trace_id, child) }
+
+          after.concat attach_waiting_children(tree, trace_id, child[:id], child)
+        end
+
+        tree[:orphans].delete(parent_id)
+        after
+      end
+
+      def sweep_all_orphans
+        after = []
+
+        @mutex.synchronize do
+          now = Process.clock_gettime(Process::CLOCK_MONOTONIC)
+
+          TRACE_TREES.each do |trace_id, tree|
+            next if tree[:orphans].empty?
+
+            tree[:orphans].each do |parent_id, entries|
+              entries.reject! do |entry|
+                if now - entry[:time] > 3
+                  orphan = entry[:node]
+
+                  after << -> { Broadcaster.emit(trace_id, orphan) }
+                  true
+                else
+                  false
+                end
+              end
+            end
+          end
+        end
+
+        after.each(&:call)
+      end
+
+      def finalize_trace(trace_id)
+        trace = nil
+
+        @mutex.synchronize do
+          tree = TRACE_TREES[trace_id]
+          return unless tree
+          trace = tree[:root]
+          TRACE_TREES.delete(trace_id)
+        end
+
+        ActionCable.server.broadcast(
+          "rails_trace_viewer",
+          event: "trace_completed",
+          trace: trace
+        )
+      end
+    end
+  end
+end

data/lib/rails_trace_viewer/engine.rb

@@ -0,0 +1,32 @@
+module RailsTraceViewer
+  class Engine < ::Rails::Engine
+    isolate_namespace RailsTraceViewer
+
+    initializer "rails_trace_viewer.subscribe", after: :load_config_initializers do
+      ActiveSupport.on_load(:action_controller) do
+        RailsTraceViewer::Subscribers::ControllerSubscriber.attach
+      end
+
+      ActiveSupport.on_load(:active_record) do
+        RailsTraceViewer::Subscribers::SqlSubscriber.attach
+      end
+
+      ActiveSupport.on_load(:action_view) do
+        RailsTraceViewer::Subscribers::ViewSubscriber.attach
+      end
+
+      ActiveSupport.on_load(:active_job) do
+        RailsTraceViewer::Subscribers::ActiveJobSubscriber.attach
+      end
+
+      Rails.application.reloader.to_prepare do
+        RailsTraceViewer::Subscribers::SidekiqSubscriber.attach
+        RailsTraceViewer::Subscribers::MethodSubscriber.attach
+      end
+    end
+    
+    initializer "rails_trace_viewer.start_sweeper" do
+      RailsTraceViewer::Collector.start_sweeper!
+    end
+  end
+end

data/lib/rails_trace_viewer/job_link_registry.rb

@@ -0,0 +1,42 @@
+require "concurrent/map"
+
+module RailsTraceViewer
+  class JobLinkRegistry
+    @store   = Concurrent::Map.new
+    @pending = Concurrent::Map.new
+
+    class << self
+      def register(job, trace_id:, enqueue_node_id:)
+        job_id = job.job_id
+
+        @store[job_id] = {
+          trace_id: trace_id,
+          enqueue_node_id: enqueue_node_id
+        }
+
+        if @pending[job_id]
+          @pending[job_id].each do |callback|
+            callback.call(trace_id, enqueue_node_id)
+          end
+          @pending.delete(job_id)
+        end
+      end
+
+      def on_perform(job, &block)
+        job_id = job.job_id
+
+        if @store[job_id]
+          data = @store[job_id]
+          block.call(data[:trace_id], data[:enqueue_node_id])
+        else
+          @pending[job_id] ||= []
+          @pending[job_id] << block
+        end
+      end
+
+      def delete(job)
+        @store.delete(job.job_id)
+      end
+    end
+  end
+end

data/lib/rails_trace_viewer/route_resolver.rb

@@ -0,0 +1,20 @@
+module RailsTraceViewer
+  class RouteResolver
+    def self.resolve(path, method)
+      recognized = Rails.application.routes.recognize_path(path, method: method.downcase.to_sym)
+
+      route = Rails.application.routes.routes.find do |r|
+        r.defaults[:controller] == recognized[:controller] &&
+        r.defaults[:action] == recognized[:action]
+      end
+
+      {
+        name: route&.name || "(unnamed)",
+        verb: method,
+        path: path
+      }
+    rescue
+      nil
+    end
+  end
+end

data/lib/rails_trace_viewer/subscribers/active_job_subscriber.rb

@@ -0,0 +1,72 @@
+module RailsTraceViewer
+  module Subscribers
+    class ActiveJobSubscriber
+      def self.attach
+        return if @attached
+        @attached = true
+
+        safe_serialize = ->(data) { JSON.parse(data.to_json) rescue data.to_s }
+
+        ActiveSupport::Notifications.subscribe("enqueue.active_job") do |*_args, payload|
+          job = payload[:job]
+          next if job.class.include?(Sidekiq::Worker)
+
+          if RailsTraceViewer::TraceContext.active?
+            trace_id  = RailsTraceViewer::TraceContext.trace_id
+            parent_id = RailsTraceViewer::TraceContext.parent_id
+          else
+            trace_id  = SecureRandom.uuid
+            parent_id = nil
+            RailsTraceViewer::TraceContext.start!(trace_id)
+            RailsTraceViewer::Collector.start_trace(trace_id)
+          end
+
+          enqueue_node_id = SecureRandom.uuid
+
+          node = {
+            id: enqueue_node_id,
+            parent_id: parent_id,
+            type: "job_enqueue",
+            name: job.class.name,
+            source: "ActiveJob Enqueue",
+            queue: job.queue_name,
+            job_id: job.job_id,
+            arguments: safe_serialize.call(job.arguments),
+            scheduled_at: job.scheduled_at,
+            priority: job.priority,
+            children: []
+          }
+
+          RailsTraceViewer::JobLinkRegistry.register(job, trace_id: trace_id, enqueue_node_id: enqueue_node_id)
+          RailsTraceViewer::Collector.add_node(trace_id, node)
+        end
+
+        ActiveSupport::Notifications.subscribe("perform_start.active_job") do |*_args, payload|
+          job = payload[:job]
+
+          RailsTraceViewer::JobLinkRegistry.on_perform(job) do |trace_id, enqueue_node_id|
+            RailsTraceViewer::TraceContext.start_job!(trace_id)
+            RailsTraceViewer::Collector.start_trace(trace_id)
+
+            node_id = SecureRandom.uuid
+            node = {
+              id: node_id,
+              parent_id: enqueue_node_id,
+              type: "job_perform",
+              name: job.class.name,
+              source: "ActiveJob Perform",
+              job_id: job.job_id,
+              arguments: safe_serialize.call(job.arguments),
+              executions: job.executions,
+              children: []
+            }
+
+            RailsTraceViewer::Collector.add_node(trace_id, node)
+            RailsTraceViewer::TraceContext.push(node_id)
+            RailsTraceViewer::JobLinkRegistry.delete(job)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rails_trace_viewer/subscribers/controller_subscriber.rb

@@ -0,0 +1,51 @@
+module RailsTraceViewer
+  module Subscribers
+    class ControllerSubscriber
+      def self.attach
+        return if @attached
+        @attached = true
+        
+        ActiveSupport::Notifications.subscribe("start_processing.action_controller") do |*, payload|
+          controller_name = payload[:controller].to_s
+          next if controller_name.start_with?("RailsTraceViewer::")
+
+          trace_id = SecureRandom.uuid
+          RailsTraceViewer::TraceContext.start!(trace_id)
+          RailsTraceViewer::Collector.start_trace(trace_id)
+
+          safe_params = payload[:params].to_unsafe_h rescue payload[:params]
+          safe_params = safe_params.except("controller", "action") 
+
+          route = RailsTraceViewer::RouteResolver.resolve(payload[:path], payload[:method])
+          route_node_id = SecureRandom.uuid
+
+          route_node = {
+            id: route_node_id,
+            parent_id: nil,
+            type: "route",
+            name: "#{route[:verb]} #{route[:path]}",
+            source: payload[:path],
+            verb: route[:verb],
+            url_pattern: route[:path],
+            route_name: route[:name],
+            children: []
+          }
+          RailsTraceViewer::Collector.add_node(trace_id, route_node)
+
+          request_node = {
+            id: trace_id,
+            parent_id: route_node_id,
+            type: "request",
+            name: "#{payload[:controller]}##{payload[:action]}",
+            source: "#{payload[:controller]}.rb",
+            format: payload[:format],
+            params: safe_params,
+            children: []
+          }
+          RailsTraceViewer::Collector.add_node(trace_id, request_node)
+          RailsTraceViewer::TraceContext.push(trace_id)
+        end
+      end
+    end
+  end
+end

data/lib/rails_trace_viewer/subscribers/method_subscriber.rb

@@ -0,0 +1,59 @@
+module RailsTraceViewer
+  module Subscribers
+    class MethodSubscriber
+      @trace_point = nil
+
+      def self.attach
+        return if @trace_point
+
+        @trace_point = TracePoint.trace(:call, :return) do |tp|
+          path = tp.path.to_s
+          
+          is_app_path = path.start_with?(Rails.root.to_s) && 
+                        !path.include?("/vendor/") && 
+                        !path.include?("rails_trace_viewer")
+          is_active_job = (tp.defined_class < ApplicationJob) rescue false
+          is_sidekiq_worker = (tp.defined_class.include?(Sidekiq::Worker)) rescue false
+
+          next unless is_app_path || is_active_job || is_sidekiq_worker
+
+          unless RailsTraceViewer::TraceContext.active?
+             trace_id = SecureRandom.uuid
+             RailsTraceViewer::TraceContext.start!(trace_id)
+             RailsTraceViewer::Collector.start_trace(trace_id)
+             parent_id = nil
+          else
+             trace_id = RailsTraceViewer::TraceContext.trace_id
+             parent_id = RailsTraceViewer::TraceContext.parent_id
+          end
+
+          if tp.event == :call
+            node_id = SecureRandom.uuid
+            
+            class_name = tp.defined_class.name rescue "Anonymous"
+            is_singleton = tp.defined_class.singleton_class? rescue false
+            method_sig = "#{class_name}#{is_singleton ? '.' : '#'}#{tp.method_id}"
+
+            node = {
+              id: node_id,
+              parent_id: parent_id,
+              type: "method",
+              name: method_sig,
+              source: "#{path.sub(Rails.root.to_s, '')}:#{tp.lineno}",
+              file_path: path,
+              line_number: tp.lineno,
+              class_name: class_name,
+              method_name: tp.method_id,
+              children: []
+            }
+            RailsTraceViewer::Collector.add_node(trace_id, node)
+            RailsTraceViewer::TraceContext.push(node_id)
+
+          elsif tp.event == :return
+            RailsTraceViewer::TraceContext.pop
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rails_trace_viewer/subscribers/sidekiq_subscriber.rb

@@ -0,0 +1,111 @@
+module RailsTraceViewer
+  module Subscribers
+    class SidekiqSubscriber
+      def self.attach
+        return unless defined?(Sidekiq)
+        return if @attached
+        @attached = true
+
+        Sidekiq.configure_server do |config|
+          config.server_middleware { |chain| chain.add ServerMiddleware }
+          config.client_middleware { |chain| chain.add ClientMiddleware }
+        end
+        Sidekiq.configure_client do |config|
+          config.client_middleware { |chain| chain.add ClientMiddleware }
+        end
+      end
+
+      module SafeSerializer
+        def self.serialize(args)
+          return args unless args.is_a?(Array) || args.is_a?(Hash)
+          JSON.parse(args.to_json) rescue ["(Unserializable)"]
+        end
+      end
+
+      class ClientMiddleware
+        def call(worker_class, job, queue, redis_pool)
+          trace_id = nil
+          parent_id = nil
+          started_new_trace = false
+
+          if RailsTraceViewer::TraceContext.active?
+            trace_id  = RailsTraceViewer::TraceContext.trace_id
+            parent_id = RailsTraceViewer::TraceContext.parent_id
+          else
+            trace_id = SecureRandom.uuid
+            RailsTraceViewer::TraceContext.start!(trace_id)
+            RailsTraceViewer::Collector.start_trace(trace_id)
+            started_new_trace = true
+          end
+
+          enqueue_node_id = SecureRandom.uuid
+          job["_trace_id"] = trace_id
+          job["_enqueue_node_id"] = enqueue_node_id
+
+          safe_args = SafeSerializer.serialize(job["args"])
+
+          node = {
+            id: enqueue_node_id,
+            parent_id: parent_id,
+            type: "job_enqueue",
+            name: "#{worker_class} (Enqueue)",
+            source: "Sidekiq Client",
+            worker_class: worker_class.to_s,
+            queue: queue,
+            job_arguments: safe_args,
+            jid: job["jid"],
+            children: []
+          }
+
+          RailsTraceViewer::Collector.add_node(trace_id, node)
+          
+          yield
+        ensure
+          if started_new_trace
+            RailsTraceViewer::TraceContext.stop!
+          end
+        end
+      end
+
+      class ServerMiddleware
+        def call(worker, job, queue)
+          trace_id = job["_trace_id"]
+          parent_id = job["_enqueue_node_id"] || job["_parent_id"]
+          perform_node_id = nil
+
+          if trace_id
+            RailsTraceViewer::TraceContext.start_job!(trace_id)
+            RailsTraceViewer::Collector.start_trace(trace_id)
+
+            perform_node_id = SecureRandom.uuid
+            safe_args = SafeSerializer.serialize(job["args"])
+
+            node = {
+              id: perform_node_id,
+              parent_id: parent_id, 
+              type: "job_perform",
+              name: "#{worker.class} (Perform)",
+              source: "Sidekiq Server",
+              worker_class: worker.class.to_s,
+              queue: queue,
+              job_arguments: safe_args,
+              jid: job["jid"],
+              retry_count: job["retry_count"],
+              children: []
+            }
+
+            RailsTraceViewer::Collector.add_node(trace_id, node)
+            RailsTraceViewer::TraceContext.push(perform_node_id)
+          end
+
+          yield
+        ensure
+          if trace_id
+            RailsTraceViewer::TraceContext.pop if perform_node_id
+            RailsTraceViewer::TraceContext.stop!
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rails_trace_viewer/subscribers/sql_subscriber.rb

@@ -0,0 +1,43 @@
+module RailsTraceViewer
+  module Subscribers
+    class SqlSubscriber
+      def self.attach
+        return if @attached
+        @attached = true
+
+        ActiveSupport::Notifications.subscribe("sql.active_record") do |*_args, payload|
+          next unless RailsTraceViewer::TraceContext.active?
+          next if payload[:name] == "SCHEMA"
+
+          source = extract_app_source(caller)
+          next unless source
+
+          raw_binds = payload[:type_casted_binds]
+          if raw_binds.respond_to?(:call)
+            raw_binds = raw_binds.call
+          end
+
+          binds = (raw_binds || []).map { |val| val.to_s }
+
+          node = {
+            id: SecureRandom.uuid,
+            parent_id: RailsTraceViewer::TraceContext.parent_id,
+            type: "sql",
+            name: payload[:sql],
+            source: source,
+            full_sql: payload[:sql],
+            bind_values: binds,
+            connection_id: payload[:connection_id],
+            children: []
+          }
+
+          RailsTraceViewer::Collector.add_node(RailsTraceViewer::TraceContext.trace_id, node)
+        end
+      end
+
+      def self.extract_app_source(bt)
+        bt.find { |line| line.include?("/app/") }
+      end
+    end
+  end
+end

data/lib/rails_trace_viewer/subscribers/view_subscriber.rb

@@ -0,0 +1,34 @@
+module RailsTraceViewer
+  module Subscribers
+    class ViewSubscriber
+      EVENTS = ["render_template.action_view", "render_partial.action_view"]
+
+      def self.attach
+        return if @attached
+        @attached = true
+
+        EVENTS.each do |event|
+          ActiveSupport::Notifications.subscribe(event) do |*_args, payload|
+            next unless RailsTraceViewer::TraceContext.active?
+
+            file = payload[:identifier]
+            next unless file && file.include?(Rails.root.join("app").to_s)
+
+            node = {
+              id: SecureRandom.uuid,
+              parent_id: RailsTraceViewer::TraceContext.parent_id,
+              type: "view",
+              name: file.split("/app/").last,
+              source: file.sub(Rails.root.to_s, ''),
+              layout: payload[:layout],
+              full_path: file,
+              children: []
+            }
+
+            RailsTraceViewer::Collector.add_node(RailsTraceViewer::TraceContext.trace_id, node)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rails_trace_viewer/trace_context.rb

@@ -0,0 +1,56 @@
+module RailsTraceViewer
+  class TraceContext
+    THREAD_KEY = :rails_trace_viewer_context
+
+    class << self
+      def context
+        Thread.current[THREAD_KEY] ||= {
+          active: false,
+          trace_id: nil,
+          stack: []
+        }
+      end
+
+      def start!(trace_id)
+        ctx = context
+        ctx[:active] = true
+        ctx[:trace_id] = trace_id
+        ctx[:stack] = []
+      end
+
+      def start_job!(trace_id)
+        ctx = context
+        ctx[:active] = true
+        ctx[:trace_id] = trace_id
+        ctx[:stack] ||= []
+      end
+      
+      def stop!
+        ctx = context
+        ctx[:active] = false
+        ctx[:trace_id] = nil
+        ctx[:stack] = []
+      end
+
+      def active?
+        context[:active]
+      end
+
+      def trace_id
+        context[:trace_id]
+      end
+
+      def parent_id
+        context[:stack].last
+      end
+
+      def push(node_id)
+        context[:stack] << node_id
+      end
+
+      def pop
+        context[:stack].pop
+      end
+    end
+  end
+end

data/lib/rails_trace_viewer/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module RailsTraceViewer
+  VERSION = "0.1.0"
+end