From 4a6cf67c4692a7be9d5fdc7083a9bda7956d3cad Mon Sep 17 00:00:00 2001
From: David Roetzel <david@roetzel.de>
Date: Fri, 14 Mar 2025 14:52:04 +0100
Subject: [PATCH] Add middleware to record queue time (#34172)

---
 config/initializers/prometheus_exporter.rb    | 10 ++++--
 .../middleware/prometheus_queue_time.rb       | 26 +++++++++++++++
 .../middleware/prometheus_queue_time_spec.rb  | 32 +++++++++++++++++++
 3 files changed, 65 insertions(+), 3 deletions(-)
 create mode 100644 lib/mastodon/middleware/prometheus_queue_time.rb
 create mode 100644 spec/lib/mastodon/middleware/prometheus_queue_time_spec.rb

diff --git a/config/initializers/prometheus_exporter.rb b/config/initializers/prometheus_exporter.rb
index fab095658f..fab08ceebb 100644
--- a/config/initializers/prometheus_exporter.rb
+++ b/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
diff --git a/lib/mastodon/middleware/prometheus_queue_time.rb b/lib/mastodon/middleware/prometheus_queue_time.rb
new file mode 100644
index 0000000000..fae171612d
--- /dev/null
+++ b/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
diff --git a/spec/lib/mastodon/middleware/prometheus_queue_time_spec.rb b/spec/lib/mastodon/middleware/prometheus_queue_time_spec.rb
new file mode 100644
index 0000000000..eaab93772d
--- /dev/null
+++ b/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