FEATURE: add customizable labels option to puma collector (#172) (#173)

FEATURE: add customizable labels option to puma collector (#172) (#173)

diff --git a/CHANGELOG b/CHANGELOG
index 2545ca4..9cdcfda 100644
--- a/CHANGELOG
+++ b/CHANGELOG
@@ -2,6 +2,7 @@ Unreleased
 
 - FEATURE: add job_name and queue_name labels to delayed job metrics
 - FEATURE: always scope puma metrics on hostname in collector
+- FEATURE: add customizable labels option to puma collector
 
 0.7.0 - 29-12-2020
 
diff --git a/README.md b/README.md
index 50cd6aa..cf3d1ea 100644
--- a/README.md
+++ b/README.md
@@ -538,7 +538,7 @@ end
 | Gauge | `puma_thread_pool_capacity_total` | Number of puma threads available at current scale           |
 | Gauge | `puma_max_threads_total`          | Number of puma threads at available at max scale            |
 
-All metrics may have a `phase` label.
+All metrics may have a `phase` label and all custom labels provided with the `labels` option.
 
 ### Resque metrics
 
diff --git a/lib/prometheus_exporter/instrumentation/puma.rb b/lib/prometheus_exporter/instrumentation/puma.rb
index b913a57..6d7cd48 100644
--- a/lib/prometheus_exporter/instrumentation/puma.rb
+++ b/lib/prometheus_exporter/instrumentation/puma.rb
@@ -5,8 +5,8 @@ require "json"
 # collects stats from puma
 module PrometheusExporter::Instrumentation
   class Puma
-    def self.start(client: nil, frequency: 30)
-      puma_collector = new
+    def self.start(client: nil, frequency: 30, labels: {})
+      puma_collector = new(labels)
       client ||= PrometheusExporter::Client.default
       Thread.new do
         while true
@@ -22,11 +22,16 @@ module PrometheusExporter::Instrumentation
       end
     end
 
+    def initialize(metric_labels = {})
+      @metric_labels = metric_labels
+    end
+
     def collect
       metric = {
         pid: pid,
         type: "puma",
-        hostname: ::PrometheusExporter.hostname
+        hostname: ::PrometheusExporter.hostname,
+        metric_labels: @metric_labels
       }
       collect_puma_stats(metric)
       metric
diff --git a/lib/prometheus_exporter/server/puma_collector.rb b/lib/prometheus_exporter/server/puma_collector.rb
index 238f422..f9af4cd 100644
--- a/lib/prometheus_exporter/server/puma_collector.rb
+++ b/lib/prometheus_exporter/server/puma_collector.rb
@@ -34,6 +34,9 @@ module PrometheusExporter::Server
         if m["custom_labels"]
           labels.merge!(m["custom_labels"])
         end
+        if m["metric_labels"]
+          labels.merge!(m["metric_labels"])
+        end
 
         PUMA_GAUGES.map do |k, help|
           k = k.to_s
diff --git a/test/server/collector_test.rb b/test/server/collector_test.rb
index 167aede..1454de0 100644
--- a/test/server/collector_test.rb
+++ b/test/server/collector_test.rb
@@ -490,6 +490,30 @@ class PrometheusCollectorTest < Minitest::Test
     mock_puma.verify
   end
 
+  def test_it_can_collect_puma_metrics_with_metric_labels
+    collector = PrometheusExporter::Server::Collector.new
+    client = PipedClient.new(collector, custom_labels: { service: 'service1' })
+
+    mock_puma = Minitest::Mock.new
+    mock_puma.expect(
+      :stats,
+      '{ "workers": 1, "phase": 0, "booted_workers": 1, "old_workers": 0, "worker_status": [{ "pid": 87819, "index": 0, "phase": 0, "booted": true, "last_checkin": "2018-10-16T11:50:31Z", "last_status": { "backlog":0, "running":8, "pool_capacity":32, "max_threads": 32 } }] }'
+    )
+
+    instrument = PrometheusExporter::Instrumentation::Puma.new({ foo: 'bar' })
+
+    Object.stub_const(:Puma, mock_puma) do
+      metric = instrument.collect
+      client.send_json metric
+    end
+
+    result = collector.prometheus_metrics_text
+    assert(result.include?('puma_booted_workers_total{phase="0",service="service1",foo="bar"} 1'), "has booted workers")
+    assert(result.include?('puma_request_backlog_total{phase="0",service="service1",foo="bar"} 0'), "has total backlog")
+    assert(result.include?('puma_thread_pool_capacity_total{phase="0",service="service1",foo="bar"} 32'), "has pool capacity")
+    mock_puma.verify
+  end
+
   def test_it_can_collect_resque_metrics
     collector = PrometheusExporter::Server::Collector.new
     client = PipedClient.new(collector, custom_labels: { service: 'service1' })

GitHub sha: 12f945d309722aedcc836ecd38eacf240a75213d

This commit appears in #173 which was merged by SamSaffron.