Rails Runner instrumentation

Signed-off-by: Marco Costa <marco.costa@datadoghq.com>

docs/GettingStarted.md

@@ -1369,7 +1369,9 @@ end
 
 The Rails integration will trace requests, database calls, templates rendering, and cache read/write/delete operations. The integration makes use of the Active Support Instrumentation, listening to the Notification API so that any operation instrumented by the API is traced.
 
-To enable the Rails instrumentation, create an initializer file in your `config/initializers` folder:
+To enable the Rails instrumentation, use the [Rails auto instrumentation instructions](#rails-or-hanami-applications).
+
+Alternatively, you can also create an initializer file in your `config/initializers` folder:
 
 ```ruby
 # config/initializers/datadog.rb

integration/apps/rails-seven/spec/integration/basic_spec.rb

@@ -44,4 +44,9 @@
       puts "      Webserver: #{json_result.fetch(:webserver_process)}"
     end
   end
+
+  context 'for Rails runner' do
+    subject { `bin/rails runner 'print "OK"'` }
+    it { expect { subject }.to output('OK').to_stdout }
+  end
 end

lib/datadog/tracing/contrib/analytics.rb

@@ -9,6 +9,11 @@ module Contrib
       module Analytics
         module_function
 
+        # Applies Analytics sampling rate, if applicable for this Contrib::Configuration.
+        def set_rate!(span, configuration)
+          set_sample_rate(span, configuration[:analytics_sample_rate]) if enabled?(configuration[:analytics_enabled])
+        end
+
         # Checks whether analytics should be enabled.
         # `flag` is a truthy/falsey value that represents a setting on the integration.
         def enabled?(flag = nil)

lib/datadog/tracing/contrib/rails/ext.rb

@@ -14,6 +14,15 @@ module Ext
           ENV_ANALYTICS_SAMPLE_RATE = 'DD_TRACE_RAILS_ANALYTICS_SAMPLE_RATE'
           ENV_DISABLE = 'DISABLE_DATADOG_RAILS'
 
+          SPAN_RUNNER_FILE = 'rails.runner.file'
+          SPAN_RUNNER_INLINE = 'rails.runner.inline'
+          SPAN_RUNNER_STDIN = 'rails.runner.stdin'
+          TAG_COMPONENT = 'rails'
+          TAG_OPERATION_FILE = 'runner.file'
+          TAG_OPERATION_INLINE = 'runner.inline'
+          TAG_OPERATION_STDIN = 'runner.stdin'
+          TAG_RUNNER_SOURCE = 'source'
+
           # @!visibility private
           MINIMUM_VERSION = Gem::Version.new('4')
         end

lib/datadog/tracing/contrib/rails/patcher.rb

@@ -5,6 +5,7 @@
 require_relative 'framework'
 require_relative 'log_injection'
 require_relative 'middlewares'
+require_relative 'runner'
 require_relative 'utils'
 require_relative '../semantic_logger/patcher'
 
@@ -28,6 +29,7 @@ def target_version
           def patch
             patch_before_initialize
             patch_after_initialize
+            patch_rails_runner
           end
 
           def patch_before_initialize
@@ -81,6 +83,11 @@ def after_initialize(app)
           def setup_tracer
             Contrib::Rails::Framework.setup
           end
+
+          # Instruments the `bin/rails runner` command.
+          def patch_rails_runner
+            ::Rails::Command.singleton_class.prepend(Command) if defined?(::Rails::Command)
+          end
         end
       end
     end

lib/datadog/tracing/contrib/rails/runner.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module Datadog
+  module Tracing
+    module Contrib
+      module Rails
+        # Limit the maximum size of the source code captured in the source tag.
+        MAX_TAG_VALUE_SIZE = 4096
+        private_constant :MAX_TAG_VALUE_SIZE
+
+        # Instruments the `bin/rails runner` command.
+        # This command executes the provided code with the host Rails application loaded.
+        # The command can be either:
+        # * `-`: for code provided through the STDIN.
+        # * File path: for code provided through a local file.
+        # * `inline code`: for code provided directly as a command line argument.
+        # @see https://guides.rubyonrails.org/v6.1/command_line.html#bin-rails-runner
+        module Runner
+          def runner(code_or_file = nil, *_command_argv)
+            if code_or_file == '-'
+              name = Ext::SPAN_RUNNER_STDIN
+              resource = nil
+              operation = Ext::TAG_OPERATION_STDIN
+              # The source is not yet available for STDIN, but it will be captured in `eval`.
+            elsif File.exist?(code_or_file)
+              name = Ext::SPAN_RUNNER_FILE
+              resource = code_or_file
+              operation = Ext::TAG_OPERATION_FILE
+              source = File.read(code_or_file)
+            else
+              name = Ext::SPAN_RUNNER_INLINE
+              resource = nil
+              operation = Ext::TAG_OPERATION_INLINE
+              source = code_or_file
+            end
+
+            Tracing.trace(
+              name,
+              service: Datadog.configuration.tracing[:rails][:service_name],
+              resource: resource,
+              tags: {
+                Tracing::Metadata::Ext::TAG_COMPONENT => Ext::TAG_COMPONENT,
+                Tracing::Metadata::Ext::TAG_OPERATION => operation,
+              }
+            ) do |span|
+              if source
+                span.set_tag(
+                  Ext::TAG_RUNNER_SOURCE,
+                  Core::Utils.truncate(source, MAX_TAG_VALUE_SIZE)
+                )
+              end
+              Contrib::Analytics.set_rate!(span, Datadog.configuration.tracing[:rails])
+
+              super
+            end
+          end
+
+          # Capture the executed source code when provided from STDIN.
+          def eval(*args)
+            span = Datadog::Tracing.active_span
+            if span.name == Ext::SPAN_RUNNER_STDIN
+              source = args[0]
+              span.set_tag(
+                Ext::TAG_RUNNER_SOURCE,
+                Core::Utils.truncate(source, MAX_TAG_VALUE_SIZE)
+              )
+            end
+
+            super
+          end
+
+          ruby2_keywords :eval if respond_to?(:ruby2_keywords, true)
+        end
+
+        # The instrumentation target, {Rails::Command::RunnerCommand} is only loaded
+        # right before `bin/rails runner` is executed. This means there's not much
+        # opportunity to patch it ahead of time.
+        # To ensure we can patch it successfully, we patch it's caller, {Rails::Command}
+        # and promptly patch {Rails::Command::RunnerCommand} when it is loaded.
+        module Command
+          def find_by_namespace(*args)
+            ret = super
+            # Patch RunnerCommand if it is loaded and not already patched.
+            if defined?(::Rails::Command::RunnerCommand) && !(::Rails::Command::RunnerCommand < Runner)
+              ::Rails::Command::RunnerCommand.prepend(Runner)
+            end
+            ret
+          end
+
+          ruby2_keywords :find_by_namespace if respond_to?(:ruby2_keywords, true)
+        end
+      end
+    end
+  end
+end

sig/datadog/tracing/contrib/rails/runner.rbs

@@ -0,0 +1,12 @@
+module Datadog
+  module Tracing
+    module Contrib
+      module Rails
+        module Runner
+        end
+        module Command
+        end
+      end
+    end
+  end
+end

spec/datadog/tracing/contrib/rails/runner_spec.rb

@@ -0,0 +1,147 @@
+# typed: false
+
+require_relative 'rails_helper'
+require_relative '../analytics_examples'
+
+RSpec.describe Datadog::Tracing::Contrib::Rails::Runner do
+  include_context 'Rails test application'
+
+  subject(:run) { ::Rails::Command.invoke 'runner', argv }
+  let(:argv) { [input] }
+  let(:input) {}
+  let(:source) { 'print "OK"' }
+
+  let(:configuration_options) { {} }
+  let(:span) do
+    expect(spans).to have(1).item
+    spans.first
+  end
+
+  before do
+    Datadog.configure do |c|
+      c.tracing.instrument :rails, **configuration_options
+    end
+
+    app
+  end
+
+  shared_context 'with a custom service name' do
+    context 'with a custom service name' do
+      let(:configuration_options) { { service_name: 'runner-name' } }
+
+      it 'sets the span service name' do
+        run
+        expect(span.service).to eq('runner-name')
+      end
+    end
+  end
+
+  shared_context 'with source code too long' do
+    context 'with source code too long' do
+      let(:source) { '123.to_i;' * 512  } # 4096-long string: 8 characters * 512
+
+      it 'truncates source tag to 4096 characters, with "..." at the end' do
+        run
+        expect(span.get_tag('source').size).to eq(4096)
+        expect(span.get_tag('source')).to start_with(source[0..(4096 - 3 - 1)]) # 3 fewer chars due to the appended '...'
+        expect(span.get_tag('source')).to end_with('...') # The appended '...'
+      end
+    end
+  end
+
+  context 'with a file path' do
+    around do |example|
+      Tempfile.open('empty-file') do |file|
+        @file_path = file.path
+        file.write(source)
+        file.flush
+
+        example.run
+      end
+    end
+
+    let(:file_path) { @file_path }
+    let(:input) { file_path }
+
+    it 'creates span for a file runner' do
+      expect { run }.to output('OK').to_stdout
+
+      expect(span.name).to eq('rails.runner.file')
+      expect(span.resource).to eq(file_path)
+      expect(span.service).to eq(tracer.default_service)
+      expect(span.get_tag('source')).to eq('print "OK"')
+      expect(span.get_tag('component')).to eq('rails')
+      expect(span.get_tag('operation')).to eq('runner.file')
+    end
+
+    include_context 'with a custom service name'
+    include_context 'with source code too long'
+
+    it_behaves_like 'analytics for integration', ignore_global_flag: false do
+      let(:source) { '' }
+      before { run }
+      let(:analytics_enabled_var) { Datadog::Tracing::Contrib::Rails::Ext::ENV_ANALYTICS_ENABLED }
+      let(:analytics_sample_rate_var) { Datadog::Tracing::Contrib::Rails::Ext::ENV_ANALYTICS_SAMPLE_RATE }
+    end
+  end
+
+  context 'from STDIN' do
+    around do |example|
+      begin
+        stdin = $stdin
+        $stdin = StringIO.new(source)
+        example.run
+      ensure
+        $stdin = stdin
+      end
+    end
+
+    let(:input) { '-' }
+
+    it 'creates span for an STDIN runner' do
+      expect { run }.to output('OK').to_stdout
+
+      expect(span.name).to eq('rails.runner.stdin')
+      expect(span.resource).to eq('rails.runner.stdin') # Fallback to span#name
+      expect(span.service).to eq(tracer.default_service)
+      expect(span.get_tag('source')).to eq('print "OK"')
+      expect(span.get_tag('component')).to eq('rails')
+      expect(span.get_tag('operation')).to eq('runner.stdin')
+    end
+
+    include_context 'with a custom service name'
+    include_context 'with source code too long'
+
+    it_behaves_like 'analytics for integration', ignore_global_flag: false do
+      let(:source) { '' }
+      before { run }
+      let(:analytics_enabled_var) { Datadog::Tracing::Contrib::Rails::Ext::ENV_ANALYTICS_ENABLED }
+      let(:analytics_sample_rate_var) { Datadog::Tracing::Contrib::Rails::Ext::ENV_ANALYTICS_SAMPLE_RATE }
+    end
+  end
+
+  context 'from an inline code snippet' do
+    let(:input) { source }
+
+    it 'creates span for an inline code snippet' do
+      expect { run }.to output('OK').to_stdout
+
+      expect(span.name).to eq('rails.runner.inline')
+      expect(span.resource).to eq('rails.runner.inline') # Fallback to span#name
+      expect(span.service).to eq(tracer.default_service)
+      expect(span.get_tag('source')).to eq('print "OK"')
+      expect(span.get_tag('component')).to eq('rails')
+      expect(span.get_tag('operation')).to eq('runner.inline')
+    end
+
+    include_context 'with a custom service name'
+    include_context 'with source code too long'
+
+    it_behaves_like 'analytics for integration', ignore_global_flag: false do
+      let(:source) { '' }
+      before { run }
+      let(:analytics_enabled_var) { Datadog::Tracing::Contrib::Rails::Ext::ENV_ANALYTICS_ENABLED }
+      let(:analytics_sample_rate_var) { Datadog::Tracing::Contrib::Rails::Ext::ENV_ANALYTICS_SAMPLE_RATE }
+    end
+  end
+end

spec/datadog/tracing/contrib/rails/support/base.rb

@@ -1,4 +1,5 @@
 require 'rails/all'
+require 'rails/command' # Loaded by the `bin/rails` script in a real Rails application
 
 require_relative 'controllers'
 require_relative 'models'