rails-instrumentation 0.1.0

data/bin/console

@@ -0,0 +1,14 @@
+#!/usr/bin/env ruby
+
+require "bundler/setup"
+require "rails/instrumentation"
+
+# You can add fixtures and/or initialization code here to make experimenting
+# with your gem easier. You can also use a different console, if you like.
+
+# (If you use this, don't forget to add pry to your Gemfile!)
+# require "pry"
+# Pry.start
+
+require "irb"
+IRB.start(__FILE__)

data/bin/setup

@@ -0,0 +1,8 @@
+#!/usr/bin/env bash
+set -euo pipefail
+IFS=$'\n\t'
+set -vx
+
+bundle install
+
+# Do any other automated setup that you need to do here

data/lib/rails/instrumentation.rb

@@ -0,0 +1,60 @@
+require 'rails/instrumentation/version'
+require 'rails/instrumentation/patch'
+require 'rails/instrumentation/subscriber'
+require 'rails/instrumentation/subscribers/action_controller_subscriber'
+require 'rails/instrumentation/subscribers/action_view_subscriber'
+require 'rails/instrumentation/subscribers/active_record_subscriber'
+require 'rails/instrumentation/subscribers/active_support_subscriber'
+require 'rails/instrumentation/subscribers/action_mailer_subscriber'
+require 'rails/instrumentation/subscribers/active_job_subscriber'
+require 'rails/instrumentation/subscribers/action_cable_subscriber'
+require 'rails/instrumentation/subscribers/active_storage_subscriber'
+require 'rails/instrumentation/utils'
+
+require 'opentracing'
+
+module Rails
+  module Instrumentation
+    class Error < StandardError; end
+
+    TAGS = {
+      'component' => 'ruby-rails',
+      'instrumentation.version' => Rails::Instrumentation::VERSION
+    }.freeze
+
+    def self.instrument(tracer: OpenTracing.global_tracer,
+                        exclude_events: [])
+      @tracer = tracer
+
+      add_subscribers(exclude_events: exclude_events)
+      Patch.patch_process_action
+    end
+
+    def self.tracer
+      @tracer
+    end
+
+    def self.add_subscribers(exclude_events: [])
+      ActiveRecordSubscriber.subscribe(exclude_events: exclude_events)
+      ActionControllerSubscriber.subscribe(exclude_events: exclude_events)
+      ActionViewSubscriber.subscribe(exclude_events: exclude_events)
+      ActiveSupportSubscriber.subscribe(exclude_events: exclude_events)
+      ActionMailerSubscriber.subscribe(exclude_events: exclude_events)
+      ActiveJobSubscriber.subscribe(exclude_events: exclude_events)
+      ActionCableSubscriber.subscribe(exclude_events: exclude_events)
+      ActiveStorageSubscriber.subscribe(exclude_events: exclude_events)
+    end
+    private_class_method :add_subscribers
+
+    def self.uninstrument
+      ActiveRecordSubscriber.unsubscribe
+      ActionControllerSubscriber.unsubscribe
+      ActionViewSubscriber.unsubscribe
+      ActiveSupportSubscriber.unsubscribe
+      ActionMailerSubscriber.unsubscribe
+      ActiveJobSubscriber.unsubscribe
+      ActionCableSubscriber.unsubscribe
+      ActiveStorageSubscriber.unsubscribe
+    end
+  end
+end

data/lib/rails/instrumentation/patch.rb

@@ -0,0 +1,42 @@
+module Rails
+  module Instrumentation
+    module Patch
+      def self.patch_process_action(klass: ::ActionController::Instrumentation)
+        klass.class_eval do
+          alias_method :process_action_original, :process_action
+
+          def process_action(method_name, *args)
+            # this naming scheme 'method.class' is how we ensure that the notification in the
+            # subscriber is the same one
+            name = "#{method_name}.#{self.class.name}"
+            puts ::Rails::Instrumentation.tracer
+            scope = ::Rails::Instrumentation.tracer.start_active_span(name)
+
+            # skip adding tags here. Getting the complete set of information is
+            # easiest in the notification
+
+            process_action_original(method_name, *args)
+          rescue Error => error
+            if scope
+              scope.span.set_tag('error', true)
+              scope.span.log_kv(key: 'message', value: error.message)
+            end
+
+            raise
+          ensure
+            scope.close
+          end
+        end
+      end
+
+      def self.restore_process_action(klass: ::ActionController::Instrumentation)
+        puts klass.respond_to? :process_action_original, true
+        klass.class_eval do
+          remove_method :process_action
+          alias_method :process_action, :process_action_original
+          remove_method :process_action_original
+        end
+      end
+    end
+  end
+end

data/lib/rails/instrumentation/subscriber.rb

@@ -0,0 +1,41 @@
+module Rails
+  module Instrumentation
+    module Subscriber
+      def self.included(base)
+        base.extend ClassMethods
+      end
+
+      module ClassMethods
+        def subscribe(exclude_events: [])
+          @subscriber_mutex = Mutex.new if @subscriber_mutex.nil?
+
+          # clear
+          unsubscribe
+          @subscribers = []
+
+          @subscriber_mutex.synchronize do
+            self::EVENTS.each do |event_name|
+              full_name = "#{event_name}.#{self::EVENT_NAMESPACE}"
+
+              next if exclude_events.include? full_name
+
+              @subscribers << Utils.register_subscriber(full_name: full_name,
+                                                        event_name: event_name,
+                                                        handler_module: self)
+            end
+          end
+        end
+
+        def unsubscribe
+          return if @subscribers.nil? || @subscriber_mutex.nil?
+
+          @subscriber_mutex.synchronize do
+            @subscribers.each do |subscriber|
+              ::ActiveSupport::Notifications.unsubscribe(subscriber)
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/rails/instrumentation/subscribers/action_cable_subscriber.rb

@@ -0,0 +1,65 @@
+module Rails
+  module Instrumentation
+    module ActionCableSubscriber
+      include Subscriber
+
+      EVENT_NAMESPACE = 'action_cable'.freeze
+
+      EVENTS = %w[
+        perform_action
+        transmit
+        transmit_subscription_confirmation
+        transmit_subscription_rejection
+        broadcast
+      ].freeze
+
+      class << self
+        def perform_action(event)
+          tags = {
+            'channel_class' => event.payload[:channel_class],
+            'action' => event.payload[:action],
+            'data' => event.payload[:data]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def transmit(event)
+          tags = {
+            'channel_class' => event.payload[:channel_class],
+            'data' => event.payload[:data],
+            'via' => event.payload[:via]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def transmit_subscription_confirmation(event)
+          tags = {
+            'channel_class' => event.payload[:channel_class]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def transmit_subscription_rejection(event)
+          tags = {
+            'channel_class' => event.payload[:channel_class]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def broadcast(event)
+          tags = {
+            'broadcasting' => event.payload[:broadcasting],
+            'message' => event.payload[:message],
+            'coder' => event.payload[:coder]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+      end
+    end
+  end
+end

data/lib/rails/instrumentation/subscribers/action_controller_subscriber.rb

@@ -0,0 +1,158 @@
+module Rails
+  module Instrumentation
+    module ActionControllerSubscriber
+      include Subscriber
+
+      EVENT_NAMESPACE = 'action_controller'.freeze
+
+      EVENTS = %w[
+        write_fragment
+        read_fragment
+        expire_fragment
+        exist_fragment?
+        write_page
+        expire_page
+        start_processing
+        process_action
+        send_file
+        send_data
+        redirect_to
+        halted_callback
+        unpermitted_parameters
+      ].freeze
+
+      class << self
+        def write_fragment(event)
+          tags = {
+            'key.write' => event.payload[:key]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def read_fragment(event)
+          tags = {
+            'key.read' => event.payload[:key]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def expire_fragment(event)
+          tags = {
+            'key.expire' => event.payload[:key]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def exist_fragment?(event)
+          tags = {
+            'key.exist' => event.payload[:key]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def write_page(event)
+          tags = {
+            'path.write' => event.payload[:path]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def expire_page(event)
+          tags = {
+            'path.expire' => event.payload[:path]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def start_processing(event)
+          tags = {
+            'controller' => event.payload[:controller],
+            'controller.action' => event.payload[:action],
+            'request.params' => event.payload[:params],
+            'request.format' => event.payload[:format],
+            'http.method' => event.payload[:method],
+            'http.url' => event.payload[:path]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def process_action(event)
+          span_name = "#{event.payload[:action]}.#{event.payload[:controller]}"
+
+          tags = {
+            'controller' => event.payload[:controller],
+            'controller.action' => event.payload[:action],
+            'request.params' => event.payload[:params],
+            'request.format' => event.payload[:format],
+            'http.method' => event.payload[:method],
+            'http.url' => event.payload[:path],
+            'http.status_code' => event.payload[:status],
+            'view.runtime.ms' => event.payload[:view_runtime],
+            'db.runtime.ms' => event.payload[:db_runtime]
+          }
+
+          # Only append these tags onto the active span created by the patched 'process_action'
+          # Otherwise, create a new span for this notification as usual
+          active_span = ::Rails::Instrumentation.tracer.active_span
+          if active_span && active_span.operation_name == span_name
+            tags.each do |key, value|
+              ::Rails::Instrumentation.tracer.active_span.set_tag(key, value)
+            end
+          else
+            Utils.trace_notification(event: event, tags: tags)
+          end
+        end
+
+        def send_file(event)
+          tags = {
+            'path.send' => event.payload[:path]
+          }
+
+          # there may be additional keys in the payload. It might be worth
+          # trying to tag them too
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def send_data(event)
+          # no defined keys, but user keys may be passed in. Might want to add
+          # them at some point
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def redirect_to(event)
+          tags = {
+            'http.status_code' => event.payload[:status],
+            'redirect.url' => event.payload[:location]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def halted_callback(event)
+          tags = {
+            'filter' => event.payload[:filter]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def unpermitted_parameters(event)
+          tags = {
+            'unpermitted_keys' => event.payload[:keys]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+      end
+    end
+  end
+end

data/lib/rails/instrumentation/subscribers/action_mailer_subscriber.rb

@@ -0,0 +1,59 @@
+module Rails
+  module Instrumentation
+    module ActionMailerSubscriber
+      include Subscriber
+
+      EVENT_NAMESPACE = 'action_mailer'.freeze
+
+      EVENTS = %w[
+        receive
+        deliver
+        process
+      ].freeze
+
+      class << self
+        def receive(event)
+          tags = {
+            'mailer' => event.payload[:mailer],
+            'message.id' => event.payload[:message_id],
+            'message.subject' => event.payload[:subject],
+            'message.to' => event.payload[:to],
+            'message.from' => event.payload[:from],
+            'message.bcc' => event.payload[:bcc],
+            'message.cc' => event.payload[:cc],
+            'message.date' => event.payload[:date],
+            'message.body' => event.payload[:mail]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def deliver(event)
+          tags = {
+            'mailer' => event.payload[:mailer],
+            'message.id' => event.payload[:message_id],
+            'message.subject' => event.payload[:subject],
+            'message.to' => event.payload[:to],
+            'message.from' => event.payload[:from],
+            'message.bcc' => event.payload[:bcc],
+            'message.cc' => event.payload[:cc],
+            'message.date' => event.payload[:date],
+            'message.body' => event.payload[:mail]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+
+        def process(event)
+          tags = {
+            'mailer' => event.payload[:mailer],
+            'action' => event.payload[:action],
+            'args' => event.payload[:args]
+          }
+
+          Utils.trace_notification(event: event, tags: tags)
+        end
+      end
+    end
+  end
+end