Merge pull request #4779 from DataDog/fix-rack

Rack: Ensure proxy span always finishes

Co-authored-by: marcotc <[email protected]>

Author: dd-mergequeue[bot]

lib/datadog/tracing/contrib/rack/request_queue.rb

@@ -24,6 +24,7 @@ def get_request_start(env, now = Core::Utils::Time.now.utc)
             # nginx header is seconds in the format "t=1512379167.574"
             # apache header is microseconds in the format "t=1570633834463123"
             # heroku header is milliseconds in the format "1570634024294"
+            # @see https://github.com/heroku/vegur/blob/65d168f757e0ddb448f41cfb9e4b0281c747378d/README.md?plain=1#L383-L384
             time_string = header.to_s.delete('^0-9')
             return if time_string.nil?
 

lib/datadog/tracing/contrib/rack/trace_proxy_middleware.rb

@@ -43,7 +43,13 @@ def call(env, configuration)
             # excluding the time spent processing the request itself
             queue_span.finish
 
-            yield.tap { request_span.finish }
+            yield
+          ensure
+            # Ensure that the spans are finished even if an exception is raised.
+            # **This is very important** to prevent the trace from leaking between requests,
+            # especially because `queue_span` is normally a root span.
+            queue_span&.finish
+            request_span&.finish
           end
         end
       end

spec/datadog/tracing/contrib/rack/trace_proxy_middleware_spec.rb

@@ -5,15 +5,13 @@
 RSpec.describe Datadog::Tracing::Contrib::Rack::TraceProxyMiddleware do
   describe '#call' do
     let(:service) { 'nginx' }
+    let(:env) { { 'HTTP_X_REQUEST_START' => timestamp.to_i * 1000 } }
 
     context 'when given timestamp' do
-      let(:timestamp) { Time.now.utc }
+      let(:timestamp) { Time.at(1757000000) }
 
       context 'when request_queuing: true' do
         it 'behaves like request_queuing: :exclude_request' do
-          env = double
-          expect(Datadog::Tracing::Contrib::Rack::QueueTime).to receive(:get_request_start).with(env).and_return(timestamp)
-
           result = described_class.call(env, request_queuing: true, web_service_name: service) { :success }
 
           expect(result).to eq :success
@@ -41,13 +39,25 @@
           expect(queue_span.get_tag('span.kind')).to eq('proxy')
           expect(queue_span).to be_measured
         end
+
+        context 'when the request fails' do
+          it 'finishes the spans even if an exception is raised' do
+            expect do
+              described_class.call(env, request_queuing: true, web_service_name: service) { raise 'error' }
+            end.to raise_error('error')
+
+            expect(spans).to have(2).items
+            queue_span, request_span = spans
+            expect(request_span).to be_root_span
+            expect(queue_span.parent_id).to eq(request_span.id)
+            expect(request_span).to be_finished
+            expect(queue_span).to be_finished
+          end
+        end
       end
 
       context 'when request_queuing: false' do
         it 'does not create spans' do
-          env = double
-          allow(Datadog::Tracing::Contrib::Rack::QueueTime).to receive(:get_request_start).with(env).and_return(timestamp)
-
           result = described_class.call(env, request_queuing: false, web_service_name: service) { :success }
 
           expect(result).to eq :success
@@ -57,14 +67,11 @@
       end
     end
 
-    context 'when given withouht timestamp' do
+    context 'when given without timestamp' do
       let(:timestamp) { nil }
 
       context 'when request_queuing: true' do
         it 'does not create spans' do
-          env = double
-          expect(Datadog::Tracing::Contrib::Rack::QueueTime).to receive(:get_request_start).with(env).and_return(timestamp)
-
           result = described_class.call(env, request_queuing: true, web_service_name: service) { :success }
 
           expect(result).to eq :success
@@ -75,9 +82,6 @@
 
       context 'when request_queuing: false' do
         it 'does not create spans' do
-          env = double
-          allow(Datadog::Tracing::Contrib::Rack::QueueTime).to receive(:get_request_start).with(env).and_return(timestamp)
-
           result = described_class.call(env, request_queuing: false, web_service_name: service) { :success }
 
           expect(result).to eq :success