Add middleware to record queue time (#34172)

config/initializers/prometheus_exporter.rb

@@ -1,8 +1,10 @@
 # frozen_string_literal: true
 
 if ENV['MASTODON_PROMETHEUS_EXPORTER_ENABLED'] == 'true'
+  require 'prometheus_exporter'
+  require 'prometheus_exporter/middleware'
+
   if ENV['MASTODON_PROMETHEUS_EXPORTER_LOCAL'] == 'true'
-    require 'prometheus_exporter'
     require 'prometheus_exporter/server'
     require 'prometheus_exporter/client'
 
@@ -17,9 +19,11 @@ if ENV['MASTODON_PROMETHEUS_EXPORTER_ENABLED'] == 'true'
 
   if ENV['MASTODON_PROMETHEUS_EXPORTER_WEB_DETAILED_METRICS'] == 'true'
     # Optional, as those metrics might generate extra overhead and be redundant with what OTEL provides
-    require 'prometheus_exporter/middleware'
-
     # Per-action/controller request stats like HTTP status and timings
     Rails.application.middleware.unshift PrometheusExporter::Middleware
+  else
+    # Include stripped down version of PrometheusExporter::Middleware that only collects queue time
+    require 'mastodon/middleware/prometheus_queue_time'
+    Rails.application.middleware.unshift Mastodon::Middleware::PrometheusQueueTime
   end
 end

lib/mastodon/middleware/prometheus_queue_time.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Mastodon
+  module Middleware
+    class PrometheusQueueTime < ::PrometheusExporter::Middleware
+      # Overwrite to only collect the queue time metric
+      def call(env)
+        queue_time = measure_queue_time(env)
+
+        result = @app.call(env)
+
+        result
+      ensure
+        obj = {
+          type: 'web',
+          queue_time: queue_time,
+          default_labels: default_labels(env, result),
+        }
+        labels = custom_labels(env)
+        obj = obj.merge(custom_labels: labels) if labels
+
+        @client.send_json(obj)
+      end
+    end
+  end
+end

spec/lib/mastodon/middleware/prometheus_queue_time_spec.rb

@@ -0,0 +1,32 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+require 'prometheus_exporter'
+require 'prometheus_exporter/middleware'
+require 'mastodon/middleware/prometheus_queue_time'
+
+RSpec.describe Mastodon::Middleware::PrometheusQueueTime do
+  subject { described_class.new(app, client:) }
+
+  let(:app) do
+    proc { |_env| [200, {}, 'OK'] }
+  end
+  let(:client) do
+    instance_double(PrometheusExporter::Client, send_json: true)
+  end
+
+  describe '#call' do
+    let(:env) do
+      {
+        'HTTP_X_REQUEST_START' => "t=#{(Time.now.to_f * 1000).to_i}",
+      }
+    end
+
+    it 'reports a queue time to the client' do
+      subject.call(env)
+
+      expect(client).to have_received(:send_json)
+        .with(hash_including(queue_time: instance_of(Float)))
+    end
+  end
+end