Rack instrumentation: support and test with v3+

lib/new_relic/agent/instrumentation/rack/chain.rb

@@ -21,8 +21,16 @@ def to_app_with_newrelic_deferred_dependency_detection
           if ::NewRelic::Agent::Instrumentation::RackHelpers.middleware_instrumentation_enabled?
             ::NewRelic::Agent.logger.info("Installing #{builder_class} middleware instrumentation")
 
-            def run_with_newrelic(app, *args)
-              run_with_tracing(app) { |wrapped_app| run_without_newrelic(wrapped_app, *args) }
+            def run_with_newrelic(app = nil, *args, &block)
+              app_or_block = app || block
+              run_with_tracing(app_or_block) do |wrapped|
+                # Rack::Builder#run for Rack v3+ supports a block, and does not
+                # support args. Whether a block or an app is provided, that
+                # callable object will be wrapped into a MiddlewareProxy
+                # instance. That proxy instance must then be passed to
+                # run_without_newrelic as the app argument.
+                block ? run_without_newrelic(wrapped, &nil) : run_without_newrelic(wrapped, *args)
+              end
             end
 
             alias_method(:run_without_newrelic, :run)

lib/new_relic/agent/instrumentation/rack/prepend.rb

@@ -23,8 +23,15 @@ def to_app
         with_deferred_dependency_detection { super }
       end
 
-      def run(app, *args)
-        run_with_tracing(app) { |wrapped_app| super(wrapped_app, *args) }
+      def run(app = nil, *args, &block)
+        app_or_block = app || block
+        run_with_tracing(app_or_block) do |wrapped|
+          # Rack::Builder#run for Rack v3+ supports a block, and does not
+          # support args. Whether a block or an app is provided, that callable
+          # object will be wrapped into a MiddlewareProxy instance. That
+          # proxy instance must then be passed to super as the app argument.
+          block ? super(wrapped, &nil) : super(wrapped, *args)
+        end
       end
 
       def use(middleware_class, *args, &blk)

test/multiverse/suites/rack/Envfile

@@ -14,10 +14,12 @@ PUMA_VERSIONS = [
 ]
 
 def gem_list(puma_version = nil)
-  rack_version = RUBY_VERSION >= '2.3.0' ? '~> 2.2.3' : '~> 2.1.4'
+  # try to test with the latest Rack version by not specifying a version
+  # constraint at all, but provide one for older Rubies not supporting Rack v3+
+  rack_version = ", '~> 2.1.4'" if RUBY_VERSION < '2.3.0'
   <<-RB
     gem 'puma'#{puma_version}
-    gem 'rack', '#{rack_version}'
+    gem 'rack'#{rack_version}
     gem 'rack-test'
     #{ruby3_gem_webrick}
   RB

test/multiverse/suites/rack/rack_builder_test.rb

@@ -112,6 +112,21 @@ def instance.middleware_instrumentation_enabled?; false; end
     end
   end
 
+  # Rack::Builder#run supports blocks as of v3.0.0
+  def test_run_with_tracing_with_the_original_run_method_is_given_a_block
+    skip 'Requires Rack v3+' unless defined?(Rack::RELEASE) && Rack::RELEASE >= '3.0.0'
+
+    builder = Rack::Builder.new
+    payload = [200, {}, ['Hello, Mr. Toad!']]
+    ::NewRelic::Agent::Instrumentation::RackHelpers.stub :middleware_instrumentation_enabled?, true do
+      builder.run { |_env| payload }
+    end
+    run = builder.instance_variable_get(:@run)
+
+    assert_kind_of NewRelic::Agent::Instrumentation::MiddlewareProxy, run
+    assert_equal payload, run.call({})
+  end
+
   def test_use_with_tracing
     skip 'Requires MiniTest v5+' unless MiniTest::Unit::VERSION > '5.0'