From d567f21d4f4c3feb09a08443cf75eb583904093b Mon Sep 17 00:00:00 2001
From: Eugen Rochko <eugen@zeonfederated.com>
Date: Thu, 26 Jan 2017 19:08:05 +0100
Subject: [PATCH] Improve StatsD instrumentation

---
 config/application.rb                  |  3 ---
 config/initializers/instrumentation.rb | 18 ++++++++++++++++++
 config/initializers/statsd.rb          | 19 ++++++-------------
 lib/statsd_monitor.rb                  | 11 -----------
 4 files changed, 24 insertions(+), 27 deletions(-)
 create mode 100644 config/initializers/instrumentation.rb
 delete mode 100644 lib/statsd_monitor.rb

diff --git a/config/application.rb b/config/application.rb
index d0b06bf95..e561d0473 100644
--- a/config/application.rb
+++ b/config/application.rb
@@ -3,7 +3,6 @@ require_relative 'boot'
 require 'rails/all'
 
 require_relative '../app/lib/exceptions'
-require_relative '../lib/statsd_monitor'
 
 # Require the gems listed in Gemfile, including any gems
 # you've limited to :test, :development, or :production.
@@ -31,8 +30,6 @@ module Mastodon
 
     config.active_job.queue_adapter = :sidekiq
 
-    config.middleware.insert(0, ::StatsDMonitor)
-
     config.middleware.insert_before 0, Rack::Cors do
       allow do
         origins  '*'
diff --git a/config/initializers/instrumentation.rb b/config/initializers/instrumentation.rb
new file mode 100644
index 000000000..8483f2be2
--- /dev/null
+++ b/config/initializers/instrumentation.rb
@@ -0,0 +1,18 @@
+# frozen_string_literal: true
+
+instrumentation_hostname = ENV.fetch('INSTRUMENTATION_HOSTNAME') { 'localhost' }
+
+ActiveSupport::Notifications.subscribe(/process_action.action_controller/) do |*args|
+  event      = ActiveSupport::Notifications::Event.new(*args)
+  controller = event.payload[:controller]
+  action     = event.payload[:action]
+  format     = event.payload[:format] || 'all'
+  format     = 'all' if format == '*/*'
+  status     = event.payload[:status]
+  key        = "#{controller}.#{action}.#{format}.#{instrumentation_hostname}"
+
+  ActiveSupport::Notifications.instrument :performance, action: :measure, measurement: "#{key}.total_duration", value: event.duration
+  ActiveSupport::Notifications.instrument :performance, action: :measure, measurement: "#{key}.db_time", value: event.payload[:db_runtime]
+  ActiveSupport::Notifications.instrument :performance, action: :measure, measurement: "#{key}.view_time", value: event.payload[:view_runtime]
+  ActiveSupport::Notifications.instrument :performance, measurement: "#{key}.status.#{status}"
+end
diff --git a/config/initializers/statsd.rb b/config/initializers/statsd.rb
index c9c754e7f..4e772a5ed 100644
--- a/config/initializers/statsd.rb
+++ b/config/initializers/statsd.rb
@@ -3,18 +3,11 @@
 StatsD.prefix              = 'mastodon'
 StatsD.default_sample_rate = 1
 
-StatsDMonitor.extend(StatsD::Instrument)
-StatsDMonitor.statsd_measure(:call, 'request.duration')
+ActiveSupport::Notifications.subscribe(/performance/) do |name, _start, _finish, _id, payload|
+  action      = payload[:action] || :increment
+  measurement = payload[:measurement]
+  value       = payload[:value]
+  key_name    = "#{name}.#{measurement}"
 
-STATSD_REQUEST_METRICS = {
-  'request.status.success'               => 200,
-  'request.status.not_found'             => 404,
-  'request.status.too_many_requests'     => 429,
-  'request.status.internal_server_error' => 500,
-}.freeze
-
-STATSD_REQUEST_METRICS.each do |name, code|
-  StatsDMonitor.statsd_count_if(:call, name) do |status, _env, _body|
-    status.to_i == code
-  end
+  StatsD.send(action.to_s, key_name, (value || 1))
 end
diff --git a/lib/statsd_monitor.rb b/lib/statsd_monitor.rb
deleted file mode 100644
index e48ce6541..000000000
--- a/lib/statsd_monitor.rb
+++ /dev/null
@@ -1,11 +0,0 @@
-# frozen_string_literal: true
-
-class StatsDMonitor
-  def initialize(app)
-    @app = app
-  end
-
-  def call(env)
-    @app.call(env)
-  end
-end