class ActionDispatch::ServerTiming::Subscriber

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 initialize
  @mutex = Mutex.new
end

def unsubscribe
  @mutex.synchronize do
    ActiveSupport::Notifications.unsubscribe @subscriber
    @subscriber = nil
  end
end