Merge pull request #45964 from jhawthorn/server_timing_safety

Fix ServerTiming in Threads, use single subscriber

actionpack/lib/action_dispatch/middleware/server_timing.rb

@@ -6,30 +6,71 @@ module ActionDispatch
   class ServerTiming
     SERVER_TIMING_HEADER = "Server-Timing"
 
+    class Subscriber # :nodoc:
+      include Singleton
+      KEY = :action_dispatch_server_timing_events
+
+      def initialize
+        @mutex = Mutex.new
+      end
+
+      def call(event)
+        if events = ActiveSupport::IsolatedExecutionState[KEY]
+          events << event
+        end
+      end
+
+      def collect_events
+        events = []
+        ActiveSupport::IsolatedExecutionState[KEY] = events
+        yield
+        events
+      ensure
+        ActiveSupport::IsolatedExecutionState.delete(KEY)
+      end
+
+      def ensure_subscribed
+        @mutex.synchronize do
+          # Subscribe to all events, except those beginning with "!"
+          # Ideally we would be more selective of what is being measured
+          @subscriber ||= ActiveSupport::Notifications.subscribe(/\A[^!]/, self)
+        end
+      end
+
+      def unsubscribe
+        @mutex.synchronize do
+          ActiveSupport::Notifications.unsubscribe @subscriber
+          @subscriber = nil
+        end
+      end
+    end
+
+    def self.unsubscribe # :nodoc:
+      Subscriber.instance.unsubscribe
+    end
+
     def initialize(app)
       @app = app
+      @subscriber = Subscriber.instance
+      @subscriber.ensure_subscribed
     end
 
     def call(env)
-      events = []
-      subscriber = ActiveSupport::Notifications.subscribe(/.*/) do |*args|
-        events << ActiveSupport::Notifications::Event.new(*args)
+      response = nil
+      events = @subscriber.collect_events do
+        response = @app.call(env)
       end
 
-      status, headers, body = begin
-        @app.call(env)
-      ensure
-        ActiveSupport::Notifications.unsubscribe(subscriber)
-      end
+      headers = response[1]
 
       header_info = events.group_by(&:name).map do |event_name, events_collection|
-        "#{event_name};dur=#{events_collection.sum(&:duration)}"
+        "%s;dur=%.2f" % [event_name, events_collection.sum(&:duration)]
       end
 
       header_info.prepend(headers[SERVER_TIMING_HEADER]) if headers[SERVER_TIMING_HEADER].present?
       headers[SERVER_TIMING_HEADER] = header_info.join(", ")
 
-      [ status, headers, body ]
+      response
     end
   end
 end

actionpack/test/dispatch/server_timing_test.rb

@@ -26,6 +26,13 @@ def create
     @middlewares = [ActionDispatch::ServerTiming]
   end
 
+  teardown do
+    # Avoid leaking subscription into other tests
+    # This will break any active instance of the middleware, but we don't
+    # expect there to be any outside of this file.
+    ActionDispatch::ServerTiming.unsubscribe
+  end
+
   test "server timing header is included in the response" do
     with_test_route_set do
       get "/"
@@ -48,6 +55,43 @@ def create
     end
   end
 
+  test "events are tracked by thread" do
+    barrier = Concurrent::CyclicBarrier.new(2)
+
+    stub_app = -> (env) {
+      env["proc"].call
+      [200, {}, "ok"]
+    }
+    app = ActionDispatch::ServerTiming.new(stub_app)
+
+    t1 = Thread.new do
+      app.call({ "proc" => -> {
+        barrier.wait
+        barrier.wait
+      } })
+    end
+
+    t2 = Thread.new do
+      barrier.wait
+
+      response = app.call({ "proc" => -> {
+        ActiveSupport::Notifications.instrument("custom.event") do
+          true
+        end
+      } })
+
+      barrier.wait
+
+      response
+    end
+
+    headers1 = t1.value[1]
+    headers2 = t2.value[1]
+
+    assert_match(/custom.event;dur=\w+/, headers2["Server-Timing"])
+    assert_no_match(/custom.event;dur=\w+/, headers1["Server-Timing"])
+  end
+
   test "does not overwrite existing header values" do
     @middlewares << Class.new do
       def initialize(app)
@@ -68,13 +112,6 @@ def call(env)
     end
   end
 
-  test "ensures it doesn't leak subscriptions when the app crashes" do
-    with_test_route_set do
-      post "/"
-      assert_not ActiveSupport::Notifications.notifier.listening?("custom.event")
-    end
-  end
-
   private
     def with_test_route_set
       with_routing do |set|