Merge pull request #10 from lessonnine/TNT-2624/add-queue-time-measur…

…ement

Add request queue time middleware

.rubocop.yml

@@ -1,6 +1,5 @@
 require:
   - rubocop-performance
-  - rubocop-rspec
 
 AllCops:
   TargetRubyVersion: 2.6

CHANGELOG.md

@@ -6,9 +6,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased]
+
+## [1.0.0] - 2021-09-08
 ### Added
 - Release to Github Packages
 - Explicitly flush datadog metrics after publishing them
+- Middleware for measuring and tracking request queue time
 
 ### Changed
 - Replace `statsd.batch` with direct calls, as it aggregates metrics interally by default now.

Gemfile

@@ -7,8 +7,8 @@ gemspec
 
 gem "dogstatsd-ruby"
 
+gem "rack"
 gem "rake", "~> 12.0"
 gem "rspec", "~> 3.0"
 gem "rubocop", "~> 1.5"
 gem "rubocop-performance", "~> 1.9"
-gem "rubocop-rspec", "~> 2.0"

Gemfile.lock

@@ -16,6 +16,7 @@ GEM
       ast (~> 2.4.1)
     puma (5.4.0)
       nio4r (~> 2.0)
+    rack (2.2.3)
     rainbow (3.0.0)
     rake (12.3.3)
     regexp_parser (2.1.1)
@@ -47,9 +48,6 @@ GEM
     rubocop-performance (1.11.5)
       rubocop (>= 1.7.0, < 2.0)
       rubocop-ast (>= 0.4.0)
-    rubocop-rspec (2.4.0)
-      rubocop (~> 1.0)
-      rubocop-ast (>= 1.1.0)
     ruby-progressbar (1.11.0)
     unicode-display_width (2.0.0)
 
@@ -59,11 +57,11 @@ PLATFORMS
 DEPENDENCIES
   dogstatsd-ruby
   puma-plugin-telemetry!
+  rack
   rake (~> 12.0)
   rspec (~> 3.0)
   rubocop (~> 1.5)
   rubocop-performance (~> 1.9)
-  rubocop-rspec (~> 2.0)
 
 BUNDLED WITH
    2.1.4

README.md

@@ -78,6 +78,29 @@ Just be mindful that if the API takes long to call, it will slow down frequency
   config.add_target proc { |telemetry| puts telemetry.map { |k, v| "#{k}=#{v.inspect}" }.join(" ") }
 ```
 
+## Extra middleware
+
+This gems comes together with middleware for measuring request queue time, which will be reported in `request.env` and published to given statsd client.
+
+Example configuration:
+
+```ruby
+# in Gemfile add `require` part
+gem "puma-plugin-telemetry", require: ["rack/request_queue_time_middleware"]
+
+# in initializer, i.e. `request_queue_time.rb`
+Rails.application.config.middleware.insert_after(
+  0,
+  RequestQueueTimeMiddleware,
+  statsd: Datadog::Statsd.new(namespace: "ruby.puma", tags: { "app" => "accounts" })
+)
+
+Rails.application.config.log_tags ||= {}
+Rails.application.config.log_tags[:queue_time] = ->(req) { req.env[::RequestQueueTimeMiddleware::ENV_KEY] }
+```
+
+This will provide proper metric in Datadog and in logs as well. Logs can be transformed into log metrics and used for auto scaling purposes.
+
 ## Development
 
 After checking out the repo, run `bin/setup` to install dependencies. Then, run `rake spec` to run the tests. You can also run `bin/console` for an interactive prompt that will allow you to experiment.

lib/puma/plugin/telemetry/version.rb

@@ -3,7 +3,7 @@
 module Puma
   class Plugin
     module Telemetry
-      VERSION = "0.3.1"
+      VERSION = "1.0.0"
     end
   end
 end

lib/rack/request_queue_time_middleware.rb

@@ -0,0 +1,57 @@
+# frozen_string_literal: true
+
+# Measures the queue time (= time between receiving the request in downstream
+# load balancer and starting request in ruby process)
+class RequestQueueTimeMiddleware
+  ENV_KEY = "rack.request_queue_time"
+
+  def initialize(app, statsd:, process: Process)
+    @app = app
+    @statsd = statsd
+    @process = process
+  end
+
+  def call(env)
+    queue_time = measure_queue_time(env)
+
+    report_queue_time(queue_time)
+
+    env[ENV_KEY] = queue_time
+
+    @app.call(env)
+  end
+
+  private
+
+  def measure_queue_time(env)
+    start_time = queue_start(env)
+
+    return unless start_time
+
+    queue_time = request_start.to_f - start_time.to_f
+
+    queue_time unless queue_time.negative?
+  end
+
+  # Get the content of the x-amzn-trace-id header, the epoch time in seconds.
+  # see also: https://docs.aws.amazon.com/elasticloadbalancing/latest/application/load-balancer-request-tracing.html
+  def queue_start(env)
+    value = env["HTTP_X_AMZN_TRACE_ID"]
+    value&.split("Root=")&.last&.split("-")&.fetch(1)&.to_i(16)
+  end
+
+  def request_start
+    @process.clock_gettime(Process::CLOCK_REALTIME)
+  end
+
+  def report_queue_time(queue_time)
+    return if queue_time.nil?
+
+    @statsd.timing("queue.time", queue_time)
+
+    return unless defined?(Datadog) && Datadog.respond_to?(:tracer)
+
+    span = Datadog.tracer.active_root_span
+    span&.set_tag("request.queue_time", queue_time)
+  end
+end

spec/rack/request_queue_time_middleware_spec.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+require_relative "../../lib/rack/request_queue_time_middleware"
+require "rack"
+require "datadog/statsd"
+
+# Provide mock as Timecop doesn't support such case
+class MockedProcess
+  def initialize
+    @clock = Process.clock_gettime(Process::CLOCK_REALTIME)
+  end
+
+  def microseconds
+    @clock - @clock.floor
+  end
+
+  def amz_ago(ago)
+    (@clock - ago).to_i.to_s(16)
+  end
+
+  def clock_gettime(_arg)
+    @clock
+  end
+end
+
+RSpec.describe RequestQueueTimeMiddleware do
+  subject(:make_request) { request.get("/some/path", headers) }
+
+  let(:request) do
+    Rack::MockRequest.new(described_class.new(->(env) { [200, env, "Bazinga!"] },
+                                              statsd: statsd,
+                                              process: process))
+  end
+
+  let(:process) { MockedProcess.new }
+  let(:statsd) { instance_double(Datadog::Statsd) }
+  let(:headers) { { "HTTP_X_AMZN_TRACE_ID" => header } }
+  let(:header) do
+    %W[
+      Self=1-#{process.amz_ago(expected_duration - 6)}-12456789abcdef012345678
+      Root=1-#{process.amz_ago(expected_duration)}-abcdef012345678912345678
+    ].join(";")
+  end
+
+  let(:expected_duration) { 12 + process.microseconds }
+
+  context "when correct header" do
+    it "reports queue time" do
+      expect(statsd).to receive(:timing).with("queue.time", expected_duration)
+      expect(make_request.status).to eq(200)
+    end
+  end
+
+  context "when header missing" do
+    let(:headers) { {} }
+
+    it "doesn't report anything" do
+      expect(statsd).not_to receive(:timing)
+      expect(make_request.status).to eq(200)
+    end
+  end
+
+  context "when header in the future" do
+    let(:expected_duration) { -(12 + process.microseconds) }
+
+    it "doesn't report anything" do
+      expect(statsd).not_to receive(:timing)
+      expect(make_request.status).to eq(200)
+    end
+  end
+end