Merge pull request #271 from palazzem/rack-context

[rack] middleware refactoring; Context cleaning happens at the end of the request

docs/GettingStarted.md

@@ -182,20 +182,22 @@ To start using the middleware in your generic Rack application, add it to your `
 
     run app
 
-Experimental distributed tracing support is available for this library.
-You need to set the ``:distributed_tracing_enabled`` option to true, for example:
+The Rack middleware can be configured using the global configuration object:
 
     # config.ru example
+    require 'ddtrace'
     require 'ddtrace/contrib/rack/middlewares'
 
-    use Datadog::Contrib::Rack::TraceMiddleware, distributed_tracing_enabled: true
+    Datadog.configure do |c|
+      c.use :rack, service_name: 'api-intake', distributed_tracing: true
+    end
 
     app = proc do |env|
-      # trace and read 'x-datadog-trace-id' and 'x-datadog-parent-id'
       [ 200, {'Content-Type' => 'text/plain'}, "OK" ]
     end
 
-See [distributed tracing](#Distributed_Tracing) for details.
+In the example above, we've activated the Distributed Tracing flag, please
+see [distributed tracing](#Distributed_Tracing) for more details.
 
 #### Configure the tracer
 

lib/ddtrace/contrib/rack/middlewares.rb

@@ -23,52 +23,35 @@ class TraceMiddleware
         end
         option :distributed_tracing, default: false
 
-        def initialize(app, options = {})
-          # update options with our configuration, unless it's already available
-          [:tracer, :service_name, :distributed_tracing].each do |k|
-            Datadog.configuration[:rack][k] = options[k] unless options[k].nil?
-          end
-
+        def initialize(app)
           @app = app
         end
 
-        def configure
-          # ensure that the configuration is executed only once
-          return clean_context if @tracer && @service
-
-          # retrieve the current tracer and service
-          @tracer = Datadog.configuration[:rack][:tracer]
-          @service = Datadog.configuration[:rack][:service_name]
-          @distributed_tracing = Datadog.configuration[:rack][:distributed_tracing]
-        end
-
-        # rubocop:disable Metrics/MethodLength
         def call(env)
-          # configure the Rack middleware once
-          configure()
+          # retrieve integration settings
+          tracer = Datadog.configuration[:rack][:tracer]
+          service = Datadog.configuration[:rack][:service_name]
+          distributed_tracing = Datadog.configuration[:rack][:distributed_tracing]
 
           trace_options = {
-            service: @service,
+            service: service,
             resource: nil,
             span_type: Datadog::Ext::HTTP::TYPE
           }
 
-          request_span = nil
-          begin
-            if @distributed_tracing
-              context = HTTPPropagator.extract(env)
-              @tracer.provider.context = context if context.trace_id
-            end
-          ensure
-            # start a new request span and attach it to the current Rack environment;
-            # we must ensure that the span `resource` is set later
-            request_span = @tracer.trace('rack.request', trace_options)
+          if distributed_tracing
+            context = HTTPPropagator.extract(env)
+            tracer.provider.context = context if context.trace_id
           end
 
+          # start a new request span and attach it to the current Rack environment;
+          # we must ensure that the span `resource` is set later
+          request_span = tracer.trace('rack.request', trace_options)
           env[:datadog_rack_request_span] = request_span
 
           # call the rest of the stack
           status, headers, response = @app.call(env)
+          [status, headers, response]
 
         # rubocop:disable Lint/RescueException
         # Here we really want to catch *any* exception, not only StandardError,
@@ -80,7 +63,7 @@ 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)
+          request_span.set_error(e) unless request_span.nil?
           raise e
         ensure
           # the source of truth in Rack is the PATH_INFO key that holds the
@@ -113,18 +96,14 @@ def call(env)
             request_span.status = 1
           end
 
-          request_span.finish()
-
-          [status, headers, response]
-        end
-
-        private
+          # ensure the request_span is finished and the context reset;
+          # this assumes that the Rack middleware creates a root span
+          request_span.finish
 
-        # TODO: Remove this once we change how context propagation works. This
-        # ensures we clean thread-local variables on each HTTP request avoiding
-        # memory leaks.
-        def clean_context
-          @tracer.provider.context = Datadog::Context.new
+          # TODO: Remove this once we change how context propagation works. This
+          # ensures we clean thread-local variables on each HTTP request avoiding
+          # memory leaks.
+          tracer.provider.context = Datadog::Context.new
         end
       end
     end

test/contrib/grape/rack_app.rb

@@ -27,22 +27,30 @@ class BaseRackAPITest < MiniTest::Test
   include Rack::Test::Methods
 
   def app
-    tracer = @tracer
-
     # create a custom Rack application with the Rack middleware and a Grape API
     Rack::Builder.new do
-      use Datadog::Contrib::Rack::TraceMiddleware, tracer: tracer
+      use Datadog::Contrib::Rack::TraceMiddleware
       map '/api/' do
         run RackTestingAPI
       end
     end.to_app
   end
 
   def setup
-    # use a dummy tracer
-    @tracer = get_test_tracer()
-    pin = Datadog::Pin.get_from(::Grape)
-    pin.tracer = @tracer
     super
+    # store the configuration and use a DummyTracer
+    @tracer = get_test_tracer
+
+    Datadog.configure do |c|
+      c.use :grape
+      c.use :rack, tracer: @tracer
+    end
+  end
+
+  def teardown
+    super
+    # reset the configuration
+    Datadog.registry[:rack].reset_options!
+    Datadog.registry[:grape].reset_options!
   end
 end

test/contrib/rack/distributed_test.rb

@@ -14,13 +14,17 @@ def call(env)
   end
 end
 
-class DistributedTest < Minitest::Test
+class DistributedTest < RackBaseTest
   RACK_HOST = 'localhost'.freeze
   RACK_PORT = 9292
 
   def setup
-    @tracer = get_test_tracer
+    super
     @rack_port = RACK_PORT
+
+    Datadog.configure do |c|
+      c.use :rack, tracer: @tracer, distributed_tracing: true
+    end
   end
 
   def check_distributed(tracer, client, distributed, message)
@@ -54,7 +58,7 @@ def test_net_http_get
     tracer = @tracer
 
     app = Rack::Builder.new do
-      use Datadog::Contrib::Rack::TraceMiddleware, tracer: tracer, distributed_tracing: true
+      use Datadog::Contrib::Rack::TraceMiddleware
       run EchoApp.new
     end
 

test/contrib/rack/helpers.rb

@@ -16,7 +16,7 @@ def app
 
     # rubocop:disable Metrics/BlockLength
     Rack::Builder.new do
-      use Datadog::Contrib::Rack::TraceMiddleware, tracer: tracer
+      use Datadog::Contrib::Rack::TraceMiddleware
 
       map '/success/' do
         run(proc { |_env| [200, { 'Content-Type' => 'text/html' }, 'OK'] })
@@ -90,9 +90,18 @@ def app
   end
 
   def setup
-    # configure our Middleware with a DummyTracer
-    @tracer = get_test_tracer()
-    Datadog.configuration[:rack][:service_name] = 'rack'
     super
+    # store the configuration and use a DummyTracer
+    @tracer = get_test_tracer
+
+    Datadog.configure do |c|
+      c.use :rack, tracer: @tracer
+    end
+  end
+
+  def teardown
+    super
+    # reset the configuration
+    Datadog.registry[:rack].reset_options!
   end
 end

test/contrib/rack/middleware_test.rb

@@ -247,25 +247,14 @@ def test_middleware_context_cleaning
     assert_equal(0, @tracer.provider.context.trace.length)
     assert_equal(1, @tracer.writer.spans.length)
   end
-end
-
-class CustomTracerTest < RackBaseTest
-  def app
-    tracer = @tracer
-    service = 'custom-rack'
-
-    Rack::Builder.new do
-      use Datadog::Contrib::Rack::TraceMiddleware, tracer: tracer, service_name: service
 
-      map '/' do
-        run(proc { |_env| [200, { 'Content-Type' => 'text/html' }, 'OK'] })
-      end
+  def test_request_middleware_custom_service
+    # ensure the Rack request is properly traced with a custom service name
+    Datadog.configure do |c|
+      c.use :rack, service_name: 'custom-rack'
     end
-  end
 
-  def test_request_middleware_custom_service
-    # ensure the Rack request is properly traced
-    get '/'
+    get '/success/'
     assert last_response.ok?
 
     spans = @tracer.writer.spans()
@@ -278,35 +267,8 @@ def test_request_middleware_custom_service
     assert_equal('GET 200', span.resource)
     assert_equal('GET', span.get_tag('http.method'))
     assert_equal('200', span.get_tag('http.status_code'))
-    assert_equal('/', span.get_tag('http.url'))
+    assert_equal('/success/', span.get_tag('http.url'))
     assert_equal(0, span.status)
     assert_nil(span.parent)
   end
 end
-
-class RackBaseTest < Minitest::Test
-  def test_middleware_builder_defaults
-    previous_configuration = Datadog.registry[:rack].to_h
-    Datadog.registry[:rack].reset_options!
-
-    middleware = Datadog::Contrib::Rack::TraceMiddleware.new(proc {})
-    refute_nil(middleware)
-    assert_equal(Datadog.tracer, Datadog.configuration[:rack][:tracer])
-    assert_equal('rack', Datadog.configuration[:rack][:service_name])
-
-    Datadog.configuration.use(:rack, previous_configuration)
-  end
-
-  def test_middleware_builder
-    # it should set the tracer and the service
-    previous_configuration = Datadog.registry[:rack].to_h
-
-    tracer = get_test_tracer()
-    middleware = Datadog::Contrib::Rack::TraceMiddleware.new(proc {}, tracer: tracer, service_name: 'custom-rack')
-    refute_nil(middleware)
-    assert_equal(tracer, Datadog.configuration[:rack][:tracer])
-    assert_equal('custom-rack', Datadog.configuration[:rack][:service_name])
-
-    Datadog.configuration.use(:rack, previous_configuration)
-  end
-end