[fix] rack middleware close trace via body proxy

DataDog · Nov 1, 2021 · 14c7616 · 14c7616
1 parent 8b97105
commit 14c7616
Show file tree

Hide file tree

Showing 5 changed files with 963 additions and 5 deletions.
diff --git a/Gemfile b/Gemfile
@@ -31,6 +31,7 @@ gem 'rspec-collection_matchers', '~> 1.1'
 gem 'rspec_junit_formatter', '>= 0.4.1'
 gem 'rspec_n', '~> 1.3' if RUBY_VERSION >= '2.4.0'
 gem 'ruby-prof', '~> 1.4' if RUBY_PLATFORM != 'java' && RUBY_VERSION >= '2.4.0'
+gem 'rack'
 if RUBY_VERSION >= '2.5.0'
   # Merging branch coverage results does not work for old, unsupported rubies.
   # We have a fix up for review, https://github.com/simplecov-ruby/simplecov/pull/972,

diff --git a/lib/ddtrace/contrib/delayed_job/plugin.rb b/lib/ddtrace/contrib/delayed_job/plugin.rb
@@ -9,7 +9,7 @@ module DelayedJob
       # DelayedJob plugin that instruments invoke_job hook
       class Plugin < Delayed::Plugin
         def self.instrument_invoke(job, &block)
-          return block.call(job) unless tracer && tracer.enabled
+          return yield(job) unless tracer && tracer.enabled
 
           tracer.trace(Ext::SPAN_JOB, service: configuration[:service_name], resource: job_name(job),
                                       on_error: configuration[:error_handler]) do |span|
@@ -29,7 +29,7 @@ def self.instrument_invoke(job, &block)
         end
 
         def self.instrument_enqueue(job, &block)
-          return block.call(job) unless tracer && tracer.enabled
+          return yield(job) unless tracer && tracer.enabled
 
           tracer.trace(Ext::SPAN_ENQUEUE, service: configuration[:client_service_name], resource: job_name(job)) do |span|
             set_sample_rate(span)

diff --git a/lib/ddtrace/contrib/rack/middlewares.rb b/lib/ddtrace/contrib/rack/middlewares.rb
@@ -86,7 +86,21 @@ def call(env)
 
           # call the rest of the stack
           status, headers, response = @app.call(env)
-          [status, headers, response]
+
+          # some middleware like Rack::CommonLogger wraps body via Rack::BodyProxy
+          # Rack::BodyProxy is a correct mechanism to finish span after all middleware call
+          # so Rack::CommonLogger will have trace_id & span_id tags
+
+          body = if defined?(::Rack::BodyProxy)
+                   ::Rack::BodyProxy.new(response) do
+                     close_trace(request_span, frontend_span, env, status, headers, response, original_env, tracer)
+                   end
+                 else
+                   close_trace(request_span, frontend_span, env, status, headers, response, original_env, tracer)
+                   response
+                 end
+
+          [status, headers, body]
 
         # rubocop:disable Lint/RescueException
         # Here we really want to catch *any* exception, not only StandardError,
@@ -98,9 +112,13 @@ def call(env)
           # catch exceptions that may be raised in the middleware chain
           # Note: if a middleware catches an Exception without re raising,
           # the Exception cannot be recorded here.
-          request_span.set_error(e) unless request_span.nil?
+          request_span.set_error(e) if request_span
+
+          close_trace(request_span, frontend_span, env, status, headers, response, original_env, tracer)
           raise e
-        ensure
+        end
+
+        def close_trace(request_span, frontend_span, env, status, headers, response, original_env, tracer)
           if request_span
             # Rack is a really low level interface and it doesn't provide any
             # advanced functionality like routers. Because of that, we assume that

diff --git a/lib/ddtrace/contrib/rack/patcher.rb b/lib/ddtrace/contrib/rack/patcher.rb
@@ -14,6 +14,7 @@ def target_version
 
         def patch
           # Patch middleware
+          require 'rack'
           require_relative 'middlewares'
         end
       end