From 28e4971fd96e7ad12272439f0d833eb780248fa2 Mon Sep 17 00:00:00 2001
From: Ivan Giuliani <ivan@gocardless.com>
Date: Mon, 19 Jul 2021 13:23:06 +0100
Subject: [PATCH] Add timeout on HTTP requests

By default, neither webrick or rack will timeout http requests. This
means that if a request to `/metrics` or `/queue/metrics` takes longer
than the `scrape_timeout`, we risk backing up requests for metrics and
eventually overloading the worker.

This adds rack-timeout to force the request to exit after the default
timeout of 30 seconds (but it can be changed either by setting
`RACK_TIMEOUT_SERVICE_TIMEOUT` to a different value or by the
`--metrics-request-timeout` options flag when starting the worker).
---
 bin/que     | 20 ++++++++++++++------
 que.gemspec |  1 +
 2 files changed, 15 insertions(+), 6 deletions(-)

diff --git a/bin/que b/bin/que
index 1e54e9a4..65c5d25e 100755
--- a/bin/que
+++ b/bin/que
@@ -9,6 +9,7 @@ require "ostruct"
 require "que"
 require "rack"
 require "prometheus/middleware/exporter"
+require "rack-timeout"
 
 $stdout.sync = true
 
@@ -54,6 +55,10 @@ OptionParser.new do |opts|
     options.timeout = timeout
   end
 
+  opts.on("--metrics-request-timeout [TIMEOUT]", Integer, "Max duration (in seconds) for metrics HTTP requests to be served") do |timeout|
+    options.metrics_request_timeout = timeout
+  end
+
   opts.on("-v", "--version", "Show Que version") do
     require "que"
     $stdout.puts "Que version #{Que::Version}"
@@ -83,12 +88,13 @@ rescue LoadError
   $stdout.puts "Could not load file '#{file}'"
 end
 
-log_level      = options.log_level     || ENV["QUE_LOG_LEVEL"]           || "info"
-queue_name     = options.queue_name    || ENV["QUE_QUEUE"]               || Que::Worker::DEFAULT_QUEUE
-wake_interval  = options.wake_interval || ENV["QUE_WAKE_INTERVAL"]&.to_f || Que::Worker::DEFAULT_WAKE_INTERVAL
-cursor_expiry  = options.cursor_expiry || wake_interval
-worker_count   = options.worker_count  || 1
-timeout        = options.timeout
+log_level       = options.log_level     || ENV["QUE_LOG_LEVEL"]           || "info"
+queue_name      = options.queue_name    || ENV["QUE_QUEUE"]               || Que::Worker::DEFAULT_QUEUE
+wake_interval   = options.wake_interval || ENV["QUE_WAKE_INTERVAL"]&.to_f || Que::Worker::DEFAULT_WAKE_INTERVAL
+cursor_expiry   = options.cursor_expiry || wake_interval
+worker_count    = options.worker_count  || 1
+timeout         = options.timeout
+metrics_timeout = options.metrics_request_timeout || ENV["RACK_TIMEOUT_SERVICE_TIMEOUT"] || 30
 
 Que.logger ||= Logger.new(STDOUT)
 
@@ -130,6 +136,7 @@ if options.metrics_port
 
     app = Rack::URLMap.new(
       "/" => Rack::Builder.new do
+        use Rack::Timeout, service_timeout: metrics_timeout
         use Que::Middleware::WorkerCollector, worker_group: worker_group
         use Prometheus::Middleware::Exporter
 
@@ -138,6 +145,7 @@ if options.metrics_port
       "/queue" => Rack::Builder.new do
         registry = Prometheus::Client::Registry.new
 
+        use Rack::Timeout, service_timeout: metrics_timeout
         use Que::Middleware::QueueCollector, registry: registry
         use Prometheus::Middleware::Exporter, registry: registry
 
diff --git a/que.gemspec b/que.gemspec
index 420da6d4..0ef3da15 100644
--- a/que.gemspec
+++ b/que.gemspec
@@ -28,6 +28,7 @@ Gem::Specification.new do |spec|
   spec.add_dependency "prometheus-client"
 
   spec.add_dependency "rack", "~> 2.0"
+  spec.add_dependency "rack-timeout", "~> 0.6"
 
   spec.add_runtime_dependency "activesupport"
 end