From 5361e8e9a7092ca081ff0d5126daf8950bb2c792 Mon Sep 17 00:00:00 2001
From: Oleg Pudeyev <156273877+p-datadog@users.noreply.github.com>
Date: Tue, 14 Jan 2025 15:07:33 -0500
Subject: [PATCH] DEBUG-3182 DI railtie (#4272)

---
 Steepfile                                     |  1 +
 integration/apps/hanami/Gemfile               | 10 +-----
 integration/apps/rails-five/Gemfile           |  2 +-
 .../app/controllers/di_controller.rb          |  6 ++++
 integration/apps/rails-five/bin/run           |  8 ++---
 .../rails-five/config/initializers/datadog.rb |  6 ++++
 integration/apps/rails-five/config/routes.rb  |  2 ++
 .../rails-five/spec/integration/di_spec.rb    | 36 +++++++++++++++++++
 .../spec/support/integration_helper.rb        | 19 ++++++++++
 integration/apps/rails-seven/Gemfile          |  2 +-
 .../app/controllers/di_controller.rb          |  6 ++++
 integration/apps/rails-seven/bin/run          |  8 ++---
 .../config/initializers/datadog.rb            |  7 +++-
 integration/apps/rails-seven/config/routes.rb |  2 ++
 .../rails-seven/spec/integration/di_spec.rb   | 36 +++++++++++++++++++
 .../spec/support/integration_helper.rb        | 19 ++++++++++
 integration/apps/rails-six/Gemfile            | 11 ++++--
 .../app/controllers/di_controller.rb          |  6 ++++
 integration/apps/rails-six/bin/run            |  8 ++---
 .../rails-six/config/initializers/datadog.rb  |  6 ++++
 integration/apps/rails-six/config/routes.rb   |  2 ++
 .../rails-six/spec/integration/di_spec.rb     | 36 +++++++++++++++++++
 .../spec/support/integration_helper.rb        | 19 ++++++++++
 .../images/include/datadog/demo_env.rb        | 14 +++++++-
 lib/datadog/auto_instrument.rb                |  3 ++
 lib/datadog/di.rb                             | 19 +++-------
 lib/datadog/di/contrib.rb                     | 26 ++++++++++++++
 lib/datadog/di/contrib/railtie.rb             | 15 ++++++++
 sig/datadog/di/contrib.rbs                    |  9 +++++
 sig/datadog/di/contrib/railtie.rbs            |  8 +++++
 30 files changed, 310 insertions(+), 42 deletions(-)
 create mode 100644 integration/apps/rails-five/app/controllers/di_controller.rb
 create mode 100644 integration/apps/rails-five/spec/integration/di_spec.rb
 create mode 100644 integration/apps/rails-seven/app/controllers/di_controller.rb
 create mode 100644 integration/apps/rails-seven/spec/integration/di_spec.rb
 create mode 100644 integration/apps/rails-six/app/controllers/di_controller.rb
 create mode 100644 integration/apps/rails-six/spec/integration/di_spec.rb
 create mode 100644 lib/datadog/di/contrib.rb
 create mode 100644 lib/datadog/di/contrib/railtie.rb
 create mode 100644 sig/datadog/di/contrib.rbs
 create mode 100644 sig/datadog/di/contrib/railtie.rbs

diff --git a/Steepfile b/Steepfile
index b0151312b7d..d0cb048720a 100644
--- a/Steepfile
+++ b/Steepfile
@@ -91,6 +91,7 @@ target :datadog do
   ignore 'lib/datadog/core/workers/queue.rb'
   ignore 'lib/datadog/core/workers/runtime_metrics.rb'
   ignore 'lib/datadog/di/configuration/settings.rb'
+  ignore 'lib/datadog/di/contrib/railtie.rb'
   ignore 'lib/datadog/kit/appsec/events.rb' # disabled because of https://github.com/soutaro/steep/issues/701
   ignore 'lib/datadog/kit/identity.rb'      # disabled because of https://github.com/soutaro/steep/issues/701
   ignore 'lib/datadog/opentelemetry.rb'
diff --git a/integration/apps/hanami/Gemfile b/integration/apps/hanami/Gemfile
index 5305085d12a..d5f0279b658 100644
--- a/integration/apps/hanami/Gemfile
+++ b/integration/apps/hanami/Gemfile
@@ -15,15 +15,7 @@ gem 'unicorn'
 gem 'webrick'
 gem 'pry-byebug'
 
-gem_spec = Datadog::DemoEnv.gem_spec('datadog')
-req = {require: 'datadog/auto_instrument'}
-opts = if gem_spec.last.is_a?(Hash)
-  gem_spec.pop.merge(req)
-else
-  req
-end
-gem_spec << opts
-gem 'datadog', *gem_spec
+gem *Datadog::DemoEnv.gem_datadog_auto_instrument
 gem 'google-protobuf', '~> 3.0'
 
 group :development do
diff --git a/integration/apps/rails-five/Gemfile b/integration/apps/rails-five/Gemfile
index d50445bff50..246fbd9480f 100644
--- a/integration/apps/rails-five/Gemfile
+++ b/integration/apps/rails-five/Gemfile
@@ -25,7 +25,7 @@ gem 'unicorn'
 gem 'loofah', '2.19.1'
 
 # Choose correct specs for 'datadog' demo environment
-gem 'datadog', *Datadog::DemoEnv.gem_spec('datadog')
+gem *Datadog::DemoEnv.gem_datadog_auto_instrument
 
 gem 'dogstatsd-ruby'
 gem 'ffi'
diff --git a/integration/apps/rails-five/app/controllers/di_controller.rb b/integration/apps/rails-five/app/controllers/di_controller.rb
new file mode 100644
index 00000000000..beabc41dd06
--- /dev/null
+++ b/integration/apps/rails-five/app/controllers/di_controller.rb
@@ -0,0 +1,6 @@
+class DiController < ApplicationController
+  def ar_serializer
+    test = Test.create!
+    render json: Datadog::DI.component.serializer.serialize_value(test)
+  end
+end
diff --git a/integration/apps/rails-five/bin/run b/integration/apps/rails-five/bin/run
index c0c594994c7..70f814253c7 100755
--- a/integration/apps/rails-five/bin/run
+++ b/integration/apps/rails-five/bin/run
@@ -8,13 +8,13 @@ puts "\n== Starting application process =="
 process = (ARGV[0] || Datadog::DemoEnv.process)
 command = case process
           when 'puma'
-            "bundle exec ddprofrb exec puma -C /app/config/puma.rb"
+            "bundle exec puma -C /app/config/puma.rb"
           when 'unicorn'
-            "bundle exec ddprofrb exec unicorn -c /app/config/unicorn.rb"
+            "bundle exec unicorn -c /app/config/unicorn.rb"
           when 'console'
-            "bundle exec ddprofrb exec rails c"
+            "bundle exec rails c"
           when 'irb'
-            "bundle exec ddprofrb exec irb"
+            "bundle exec irb"
           when nil, ''
             abort("\n== ERROR: Must specify a application process! ==")
           else
diff --git a/integration/apps/rails-five/config/initializers/datadog.rb b/integration/apps/rails-five/config/initializers/datadog.rb
index f3060959c06..8b8f91c1a67 100644
--- a/integration/apps/rails-five/config/initializers/datadog.rb
+++ b/integration/apps/rails-five/config/initializers/datadog.rb
@@ -28,4 +28,10 @@
       c.profiling.exporter.transport = Datadog::DemoEnv.profiler_file_transport
     end
   end
+
+  if c.respond_to?(:dynamic_instrumentation)
+    c.remote.enabled = true
+    c.dynamic_instrumentation.enabled = true
+    c.dynamic_instrumentation.internal.development = true
+  end
 end
diff --git a/integration/apps/rails-five/config/routes.rb b/integration/apps/rails-five/config/routes.rb
index b9731693eb6..8e4d1893ede 100644
--- a/integration/apps/rails-five/config/routes.rb
+++ b/integration/apps/rails-five/config/routes.rb
@@ -7,6 +7,8 @@
   get 'basic/default', to: 'basic#default'
   get 'basic/fibonacci', to: 'basic#fibonacci'
 
+  get 'di/ar_serializer', to: 'di#ar_serializer'
+
   # Job test scenarios
   post 'jobs', to: 'jobs#create'
 end
diff --git a/integration/apps/rails-five/spec/integration/di_spec.rb b/integration/apps/rails-five/spec/integration/di_spec.rb
new file mode 100644
index 00000000000..a2d2db72499
--- /dev/null
+++ b/integration/apps/rails-five/spec/integration/di_spec.rb
@@ -0,0 +1,36 @@
+require 'spec_helper'
+require 'json'
+
+RSpec.describe 'Dynamic Instrumentation' do
+  include_context 'integration test'
+  di_test
+
+  describe 'ActiveRecord integration' do
+    let(:response) { get('di/ar_serializer') }
+    subject { JSON.parse(response.body) }
+
+    it 'is loaded' do
+      expect(response.code).to eq '200'
+
+      # If AR integration is loaded, this output will be the result of
+      # the custom serializer.
+      # If AR integration is not loaded, the output here will have a bunch of
+      # internal AR fields but not the attributes themselves.
+      expect(subject).to match(
+        {"type"=>"Test",
+         "entries"=>
+          [[{"type"=>"Symbol", "value"=>"attributes"},
+            {"type"=>"Hash",
+             "entries"=>
+              [[{"type"=>"String", "value"=>"id"}, {"type"=>"Integer", "value"=>String}],
+               [{"type"=>"String", "value"=>"version"}, {"type"=>"NilClass", "isNull"=>true}],
+               [{"type"=>"String", "value"=>"data"}, {"type"=>"NilClass", "isNull"=>true}],
+               [{"type"=>"String", "value"=>"created_at"},
+                {"type"=>"Time", "value"=>String}],
+               [{"type"=>"String", "value"=>"updated_at"},
+                {"type"=>"Time", "value"=>String}]]}],
+           [{"type"=>"Symbol", "value"=>"new_record"}, {"type"=>"FalseClass", "value"=>"false"}]]}
+      )
+    end
+  end
+end
diff --git a/integration/apps/rails-five/spec/support/integration_helper.rb b/integration/apps/rails-five/spec/support/integration_helper.rb
index b0260327e6d..37871c5b065 100644
--- a/integration/apps/rails-five/spec/support/integration_helper.rb
+++ b/integration/apps/rails-five/spec/support/integration_helper.rb
@@ -19,4 +19,23 @@ def get(path)
       Net::HTTP.get_response(uri)
     end
   end
+
+  module ClassMethods
+    def di_test
+      if RUBY_ENGINE == 'jruby'
+        before(:all) do
+          skip "Dynamic instrumentation is not supported on JRuby"
+        end
+      end
+      if RUBY_VERSION < "2.6"
+        before(:all) do
+          skip "Dynamic instrumentation requires Ruby 2.6 or higher"
+        end
+      end
+    end
+  end
+
+  def self.included(base)
+    base.extend(ClassMethods)
+  end
 end
diff --git a/integration/apps/rails-seven/Gemfile b/integration/apps/rails-seven/Gemfile
index 85b7d128063..d28e118b82c 100644
--- a/integration/apps/rails-seven/Gemfile
+++ b/integration/apps/rails-seven/Gemfile
@@ -7,7 +7,7 @@ git_source(:github) { |repo| "https://github.com/#{repo}.git" }
 gem "rails", "~> 7.0.5"
 
 # Choose correct specs for 'datadog' demo environment
-gem 'datadog', *Datadog::DemoEnv.gem_spec('datadog')
+gem *Datadog::DemoEnv.gem_datadog_auto_instrument
 
 gem 'dogstatsd-ruby'
 
diff --git a/integration/apps/rails-seven/app/controllers/di_controller.rb b/integration/apps/rails-seven/app/controllers/di_controller.rb
new file mode 100644
index 00000000000..beabc41dd06
--- /dev/null
+++ b/integration/apps/rails-seven/app/controllers/di_controller.rb
@@ -0,0 +1,6 @@
+class DiController < ApplicationController
+  def ar_serializer
+    test = Test.create!
+    render json: Datadog::DI.component.serializer.serialize_value(test)
+  end
+end
diff --git a/integration/apps/rails-seven/bin/run b/integration/apps/rails-seven/bin/run
index c0c594994c7..70f814253c7 100755
--- a/integration/apps/rails-seven/bin/run
+++ b/integration/apps/rails-seven/bin/run
@@ -8,13 +8,13 @@ puts "\n== Starting application process =="
 process = (ARGV[0] || Datadog::DemoEnv.process)
 command = case process
           when 'puma'
-            "bundle exec ddprofrb exec puma -C /app/config/puma.rb"
+            "bundle exec puma -C /app/config/puma.rb"
           when 'unicorn'
-            "bundle exec ddprofrb exec unicorn -c /app/config/unicorn.rb"
+            "bundle exec unicorn -c /app/config/unicorn.rb"
           when 'console'
-            "bundle exec ddprofrb exec rails c"
+            "bundle exec rails c"
           when 'irb'
-            "bundle exec ddprofrb exec irb"
+            "bundle exec irb"
           when nil, ''
             abort("\n== ERROR: Must specify a application process! ==")
           else
diff --git a/integration/apps/rails-seven/config/initializers/datadog.rb b/integration/apps/rails-seven/config/initializers/datadog.rb
index 89e276870ea..d2749049b11 100644
--- a/integration/apps/rails-seven/config/initializers/datadog.rb
+++ b/integration/apps/rails-seven/config/initializers/datadog.rb
@@ -1,6 +1,5 @@
 require 'datadog/statsd'
 require 'datadog'
-require 'datadog/appsec'
 
 Datadog.configure do |c|
   c.env = 'integration'
@@ -26,4 +25,10 @@
     # Reconfigure transport to write pprof to file
     c.profiling.exporter.transport = Datadog::DemoEnv.profiler_file_transport
   end
+
+  if c.respond_to?(:dynamic_instrumentation)
+    c.remote.enabled = true
+    c.dynamic_instrumentation.enabled = true
+    c.dynamic_instrumentation.internal.development = true
+  end
 end
diff --git a/integration/apps/rails-seven/config/routes.rb b/integration/apps/rails-seven/config/routes.rb
index 4f012b42a9d..06a88068d6d 100644
--- a/integration/apps/rails-seven/config/routes.rb
+++ b/integration/apps/rails-seven/config/routes.rb
@@ -10,6 +10,8 @@
   get 'basic/fibonacci', to: 'basic#fibonacci'
   get 'basic/boom', to: 'basic#boom'
 
+  get 'di/ar_serializer', to: 'di#ar_serializer'
+
   # Job test scenarios
   post 'jobs', to: 'jobs#create'
 end
diff --git a/integration/apps/rails-seven/spec/integration/di_spec.rb b/integration/apps/rails-seven/spec/integration/di_spec.rb
new file mode 100644
index 00000000000..9f37c14ffbe
--- /dev/null
+++ b/integration/apps/rails-seven/spec/integration/di_spec.rb
@@ -0,0 +1,36 @@
+require 'spec_helper'
+require 'json'
+
+RSpec.describe 'Dynamic Instrumentation' do
+  include_context 'integration test'
+  di_test
+
+  describe 'ActiveRecord integration' do
+    let(:response) { get('di/ar_serializer') }
+    subject { JSON.parse(response.body) }
+
+    it 'is loaded' do
+      expect(response.code).to eq '200'
+
+      # If AR integration is loaded, this output will be the result of
+      # the custom serializer.
+      # If AR integration is not loaded, the output here will have a bunch of
+      # internal AR fields but not the attributes themselves.
+      expect(subject).to match(
+        {"type"=>"Test",
+         "entries"=>
+          [[{"type"=>"Symbol", "value"=>"attributes"},
+            {"type"=>"Hash",
+             "entries"=>
+              [[{"type"=>"String", "value"=>"id"}, {"type"=>"Integer", "value"=>String}],
+               [{"type"=>"String", "value"=>"version"}, {"type"=>"NilClass", "isNull"=>true}],
+               [{"type"=>"String", "value"=>"data"}, {"type"=>"NilClass", "isNull"=>true}],
+               [{"type"=>"String", "value"=>"created_at"},
+                {"type"=>"ActiveSupport::TimeWithZone", "value"=>String}],
+               [{"type"=>"String", "value"=>"updated_at"},
+                {"type"=>"ActiveSupport::TimeWithZone", "value"=>String}]]}],
+           [{"type"=>"Symbol", "value"=>"new_record"}, {"type"=>"FalseClass", "value"=>"false"}]]}
+      )
+    end
+  end
+end
diff --git a/integration/apps/rails-seven/spec/support/integration_helper.rb b/integration/apps/rails-seven/spec/support/integration_helper.rb
index b0260327e6d..37871c5b065 100644
--- a/integration/apps/rails-seven/spec/support/integration_helper.rb
+++ b/integration/apps/rails-seven/spec/support/integration_helper.rb
@@ -19,4 +19,23 @@ def get(path)
       Net::HTTP.get_response(uri)
     end
   end
+
+  module ClassMethods
+    def di_test
+      if RUBY_ENGINE == 'jruby'
+        before(:all) do
+          skip "Dynamic instrumentation is not supported on JRuby"
+        end
+      end
+      if RUBY_VERSION < "2.6"
+        before(:all) do
+          skip "Dynamic instrumentation requires Ruby 2.6 or higher"
+        end
+      end
+    end
+  end
+
+  def self.included(base)
+    base.extend(ClassMethods)
+  end
 end
diff --git a/integration/apps/rails-six/Gemfile b/integration/apps/rails-six/Gemfile
index 9dcce283d06..7056563881a 100644
--- a/integration/apps/rails-six/Gemfile
+++ b/integration/apps/rails-six/Gemfile
@@ -18,7 +18,7 @@ gem 'puma'
 gem 'unicorn'
 
 # Choose correct specs for 'datadog' demo environment
-gem 'datadog', *Datadog::DemoEnv.gem_spec('datadog')
+gem *Datadog::DemoEnv.gem_datadog_auto_instrument
 
 gem 'dogstatsd-ruby'
 gem 'ffi'
@@ -90,7 +90,14 @@ end
 
 group :test, :development do
   gem 'byebug', platform: :ruby
-  gem 'mock_redis'
+  # mock_redis 0.47.0+ requires redis 5 or higher
+  # mock_redis 0.42.0+ requires Ruby 3.0 or higher
+  # mock_redis 0.37.0+ requires Ruby 2.7 or higher
+  if RUBY_VERSION >= '2.7'
+    gem 'mock_redis', '~> 0.37.0'
+  else
+    gem 'mock_redis', '< 0.37.0'
+  end
   gem 'parallel_tests'
 
   gem 'listen'
diff --git a/integration/apps/rails-six/app/controllers/di_controller.rb b/integration/apps/rails-six/app/controllers/di_controller.rb
new file mode 100644
index 00000000000..beabc41dd06
--- /dev/null
+++ b/integration/apps/rails-six/app/controllers/di_controller.rb
@@ -0,0 +1,6 @@
+class DiController < ApplicationController
+  def ar_serializer
+    test = Test.create!
+    render json: Datadog::DI.component.serializer.serialize_value(test)
+  end
+end
diff --git a/integration/apps/rails-six/bin/run b/integration/apps/rails-six/bin/run
index c0c594994c7..70f814253c7 100755
--- a/integration/apps/rails-six/bin/run
+++ b/integration/apps/rails-six/bin/run
@@ -8,13 +8,13 @@ puts "\n== Starting application process =="
 process = (ARGV[0] || Datadog::DemoEnv.process)
 command = case process
           when 'puma'
-            "bundle exec ddprofrb exec puma -C /app/config/puma.rb"
+            "bundle exec puma -C /app/config/puma.rb"
           when 'unicorn'
-            "bundle exec ddprofrb exec unicorn -c /app/config/unicorn.rb"
+            "bundle exec unicorn -c /app/config/unicorn.rb"
           when 'console'
-            "bundle exec ddprofrb exec rails c"
+            "bundle exec rails c"
           when 'irb'
-            "bundle exec ddprofrb exec irb"
+            "bundle exec irb"
           when nil, ''
             abort("\n== ERROR: Must specify a application process! ==")
           else
diff --git a/integration/apps/rails-six/config/initializers/datadog.rb b/integration/apps/rails-six/config/initializers/datadog.rb
index 42f9582a4e7..4a323bb6be1 100644
--- a/integration/apps/rails-six/config/initializers/datadog.rb
+++ b/integration/apps/rails-six/config/initializers/datadog.rb
@@ -28,4 +28,10 @@
       c.profiling.exporter.transport = Datadog::DemoEnv.profiler_file_transport
     end
   end
+
+  if c.respond_to?(:dynamic_instrumentation)
+    c.remote.enabled = true
+    c.dynamic_instrumentation.enabled = true
+    c.dynamic_instrumentation.internal.development = true
+  end
 end
diff --git a/integration/apps/rails-six/config/routes.rb b/integration/apps/rails-six/config/routes.rb
index 4f012b42a9d..06a88068d6d 100644
--- a/integration/apps/rails-six/config/routes.rb
+++ b/integration/apps/rails-six/config/routes.rb
@@ -10,6 +10,8 @@
   get 'basic/fibonacci', to: 'basic#fibonacci'
   get 'basic/boom', to: 'basic#boom'
 
+  get 'di/ar_serializer', to: 'di#ar_serializer'
+
   # Job test scenarios
   post 'jobs', to: 'jobs#create'
 end
diff --git a/integration/apps/rails-six/spec/integration/di_spec.rb b/integration/apps/rails-six/spec/integration/di_spec.rb
new file mode 100644
index 00000000000..a2d2db72499
--- /dev/null
+++ b/integration/apps/rails-six/spec/integration/di_spec.rb
@@ -0,0 +1,36 @@
+require 'spec_helper'
+require 'json'
+
+RSpec.describe 'Dynamic Instrumentation' do
+  include_context 'integration test'
+  di_test
+
+  describe 'ActiveRecord integration' do
+    let(:response) { get('di/ar_serializer') }
+    subject { JSON.parse(response.body) }
+
+    it 'is loaded' do
+      expect(response.code).to eq '200'
+
+      # If AR integration is loaded, this output will be the result of
+      # the custom serializer.
+      # If AR integration is not loaded, the output here will have a bunch of
+      # internal AR fields but not the attributes themselves.
+      expect(subject).to match(
+        {"type"=>"Test",
+         "entries"=>
+          [[{"type"=>"Symbol", "value"=>"attributes"},
+            {"type"=>"Hash",
+             "entries"=>
+              [[{"type"=>"String", "value"=>"id"}, {"type"=>"Integer", "value"=>String}],
+               [{"type"=>"String", "value"=>"version"}, {"type"=>"NilClass", "isNull"=>true}],
+               [{"type"=>"String", "value"=>"data"}, {"type"=>"NilClass", "isNull"=>true}],
+               [{"type"=>"String", "value"=>"created_at"},
+                {"type"=>"Time", "value"=>String}],
+               [{"type"=>"String", "value"=>"updated_at"},
+                {"type"=>"Time", "value"=>String}]]}],
+           [{"type"=>"Symbol", "value"=>"new_record"}, {"type"=>"FalseClass", "value"=>"false"}]]}
+      )
+    end
+  end
+end
diff --git a/integration/apps/rails-six/spec/support/integration_helper.rb b/integration/apps/rails-six/spec/support/integration_helper.rb
index b0260327e6d..37871c5b065 100644
--- a/integration/apps/rails-six/spec/support/integration_helper.rb
+++ b/integration/apps/rails-six/spec/support/integration_helper.rb
@@ -19,4 +19,23 @@ def get(path)
       Net::HTTP.get_response(uri)
     end
   end
+
+  module ClassMethods
+    def di_test
+      if RUBY_ENGINE == 'jruby'
+        before(:all) do
+          skip "Dynamic instrumentation is not supported on JRuby"
+        end
+      end
+      if RUBY_VERSION < "2.6"
+        before(:all) do
+          skip "Dynamic instrumentation requires Ruby 2.6 or higher"
+        end
+      end
+    end
+  end
+
+  def self.included(base)
+    base.extend(ClassMethods)
+  end
 end
diff --git a/integration/images/include/datadog/demo_env.rb b/integration/images/include/datadog/demo_env.rb
index 516f2a0d479..46190a078cd 100644
--- a/integration/images/include/datadog/demo_env.rb
+++ b/integration/images/include/datadog/demo_env.rb
@@ -3,7 +3,7 @@ module Datadog
   module DemoEnv
     module_function
 
-    def gem_spec(gem_name, defaults = {})
+    def gem_spec(gem_name)
       args = if local_gem(gem_name)
                [local_gem(gem_name)]
              elsif git_gem(gem_name)
@@ -17,6 +17,18 @@ def gem_spec(gem_name, defaults = {})
       args
     end
 
+    def gem_datadog_auto_instrument
+      gem_spec = gem_spec('datadog')
+      req = {require: 'datadog/auto_instrument'}
+      opts = if gem_spec.last.is_a?(Hash)
+        gem_spec.pop.merge(req)
+      else
+        req
+      end
+      gem_spec << opts
+      ['datadog', *gem_spec]
+    end
+
     def gem_env_name(gem_name)
       gem_name.upcase.tr('-', '_')
     end
diff --git a/lib/datadog/auto_instrument.rb b/lib/datadog/auto_instrument.rb
index e02305a0009..264d53611cd 100644
--- a/lib/datadog/auto_instrument.rb
+++ b/lib/datadog/auto_instrument.rb
@@ -6,6 +6,9 @@
 require_relative '../datadog'
 require_relative 'tracing/contrib/auto_instrument'
 
+# DI is not loaded on Ruby 2.5 and JRuby
+Datadog::DI::Contrib.load_now_or_later if defined?(Datadog::DI::Contrib)
+
 Datadog::Profiling.start_if_enabled
 
 module Datadog
diff --git a/lib/datadog/di.rb b/lib/datadog/di.rb
index 221c32864f6..ff185b09318 100644
--- a/lib/datadog/di.rb
+++ b/lib/datadog/di.rb
@@ -18,21 +18,6 @@
 require_relative 'di/transport'
 require_relative 'di/utils'
 
-if defined?(ActiveRecord::Base)
-  # The third-party library integrations need to be loaded after the
-  # third-party libraries are loaded. Tracing and appsec use Railtie
-  # to delay integrations until all of the application's dependencies
-  # are loaded, when running under Rails. We should do the same here in
-  # principle, however DI currently only has an ActiveRecord integration
-  # and AR should be loaded before any application code is loaded, being
-  # part of Rails, therefore for now we should be OK to just require the
-  # AR integration from here.
-  #
-  # TODO this require might need to be delayed via Rails post-initialization
-  # logic?
-  require_relative 'di/contrib/active_record'
-end
-
 module Datadog
   # Namespace for Datadog dynamic instrumentation.
   #
@@ -75,3 +60,7 @@ def component
   # for line probes to work) activate tracking in an initializer.
   Datadog::DI.activate_tracking
 end
+
+require_relative 'di/contrib'
+
+Datadog::DI::Contrib.load_now_or_later
diff --git a/lib/datadog/di/contrib.rb b/lib/datadog/di/contrib.rb
new file mode 100644
index 00000000000..8098b5befb0
--- /dev/null
+++ b/lib/datadog/di/contrib.rb
@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+require_relative '../core/contrib/rails/utils'
+
+module Datadog
+  module DI
+    module Contrib
+      module_function def load_now_or_later
+        if Datadog::Core::Contrib::Rails::Utils.railtie_supported?
+          require_relative 'contrib/railtie'
+        else
+          load_now
+        end
+      end
+
+      # This method can be called more than once, to attempt to load
+      # DI components that depend on third-party libraries after additional
+      # dependencies are loaded (or potentially loaded).
+      module_function def load_now
+        if defined?(ActiveRecord::Base)
+          require_relative 'contrib/active_record'
+        end
+      end
+    end
+  end
+end
diff --git a/lib/datadog/di/contrib/railtie.rb b/lib/datadog/di/contrib/railtie.rb
new file mode 100644
index 00000000000..b0a3f57d8d2
--- /dev/null
+++ b/lib/datadog/di/contrib/railtie.rb
@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Datadog
+  module DI
+    module Contrib
+      # Railtie class initializes dynamic instrumentation contrib code
+      # in Rails environments.
+      class Railtie < Rails::Railtie
+        initializer 'datadog.dynamic_instrumentation.initialize' do |app|
+          Contrib.load_now
+        end
+      end
+    end
+  end
+end
diff --git a/sig/datadog/di/contrib.rbs b/sig/datadog/di/contrib.rbs
new file mode 100644
index 00000000000..80ac578ffac
--- /dev/null
+++ b/sig/datadog/di/contrib.rbs
@@ -0,0 +1,9 @@
+module Datadog
+  module DI
+    module Contrib
+      def self?.load_now_or_later: () -> void
+
+      def self?.load_now: () -> void
+    end
+  end
+end
diff --git a/sig/datadog/di/contrib/railtie.rbs b/sig/datadog/di/contrib/railtie.rbs
new file mode 100644
index 00000000000..473a68df3ca
--- /dev/null
+++ b/sig/datadog/di/contrib/railtie.rbs
@@ -0,0 +1,8 @@
+module Datadog
+  module DI
+    module Contrib
+      class Railtie < Rails::Railtie
+      end
+    end
+  end
+end