diff --git a/lib/ddtrace/contrib/delayed_job/plugin.rb b/lib/ddtrace/contrib/delayed_job/plugin.rb index 21c112e79d5..73b4e09015b 100644 --- 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 index d2016f28309..51bdaf33dea 100644 --- 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 index b6583dae9fd..61bf0087330 100644 --- 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 diff --git a/sorbet/rbi/todo.rbi b/sorbet/rbi/todo.rbi index 86faebea00f..08e8c80eb83 100644 --- a/sorbet/rbi/todo.rbi +++ b/sorbet/rbi/todo.rbi @@ -101,6 +101,7 @@ module Datadog::Vendor::ActiveRecord::ConnectionHandling::DEFAULT_ENV; end module Datadog::Vendor::ActiveRecord::ConnectionHandling::Rails; end module EthonSupport::Ethon::Easy; end module Rack::Request; end +module Rack::BodyProxy; end module T::CompatibilityPatches::RSpecCompatibility::MethodDoubleExtensions; end module T::CompatibilityPatches::RSpecCompatibility::RecorderExtensions; end module T::InterfaceWrapper::Helpers; end