1 files changed, 93 insertions, 0 deletions
diff --git a/lib/gitlab/metrics/samplers/puma_sampler.rb b/lib/gitlab/metrics/samplers/puma_sampler.rb
new file mode 100644
index 00000000000..25e40c70230
--- /dev/null
+++ b/lib/gitlab/metrics/samplers/puma_sampler.rb
@@ -0,0 +1,93 @@
+# frozen_string_literal: true
+
+require 'puma/state_file'
+
+module Gitlab
+  module Metrics
+    module Samplers
+      class PumaSampler < BaseSampler
+        def metrics
+          @metrics ||= init_metrics
+        end
+
+        def init_metrics
+          {
+            puma_workers:            ::Gitlab::Metrics.gauge(:puma_workers, 'Total number of workers'),
+            puma_running_workers:    ::Gitlab::Metrics.gauge(:puma_running_workers, 'Number of active workers'),
+            puma_stale_workers:      ::Gitlab::Metrics.gauge(:puma_stale_workers, 'Number of stale workers'),
+            puma_phase:              ::Gitlab::Metrics.gauge(:puma_phase, 'Phase number (increased during phased restarts)'),
+            puma_running:            ::Gitlab::Metrics.gauge(:puma_running, 'Number of running threads'),
+            puma_queued_connections: ::Gitlab::Metrics.gauge(:puma_queued_connections, 'Number of connections in that worker\'s "todo" set waiting for a worker thread'),
+            puma_active_connections: ::Gitlab::Metrics.gauge(:puma_active_connections, 'Number of threads processing a request'),
+            puma_pool_capacity:      ::Gitlab::Metrics.gauge(:puma_pool_capacity, 'Number of requests the worker is capable of taking right now'),
+            puma_max_threads:        ::Gitlab::Metrics.gauge(:puma_max_threads, 'Maximum number of worker threads'),
+            puma_idle_threads:       ::Gitlab::Metrics.gauge(:puma_idle_threads, 'Number of spawned threads which are not processing a request')
+          }
+        end
+
+        def sample
+          json_stats = puma_stats
+          return unless json_stats
+
+          stats = JSON.parse(json_stats)
+
+          if cluster?(stats)
+            sample_cluster(stats)
+          else
+            sample_single_worker(stats)
+          end
+        end
+
+        private
+
+        def puma_stats
+          Puma.stats
+        rescue NoMethodError
+          Rails.logger.info "PumaSampler: stats are not available yet, waiting for Puma to boot"
+          nil
+        end
+
+        def sample_cluster(stats)
+          set_master_metrics(stats)
+
+          stats['worker_status'].each do |worker|
+            last_status = worker['last_status']
+            labels = { worker: "worker_#{worker['index']}" }
+
+            metrics[:puma_phase].set(labels, worker['phase'])
+            set_worker_metrics(last_status, labels) if last_status.present?
+          end
+        end
+
+        def sample_single_worker(stats)
+          metrics[:puma_workers].set({}, 1)
+          metrics[:puma_running_workers].set({}, 1)
+
+          set_worker_metrics(stats)
+        end
+
+        def cluster?(stats)
+          stats['worker_status'].present?
+        end
+
+        def set_master_metrics(stats)
+          labels = { worker: "master" }
+
+          metrics[:puma_workers].set(labels, stats['workers'])
+          metrics[:puma_running_workers].set(labels, stats['booted_workers'])
+          metrics[:puma_stale_workers].set(labels, stats['old_workers'])
+          metrics[:puma_phase].set(labels, stats['phase'])
+        end
+
+        def set_worker_metrics(stats, labels = {})
+          metrics[:puma_running].set(labels, stats['running'])
+          metrics[:puma_queued_connections].set(labels, stats['backlog'])
+          metrics[:puma_active_connections].set(labels, stats['max_threads'] - stats['pool_capacity'])
+          metrics[:puma_pool_capacity].set(labels, stats['pool_capacity'])
+          metrics[:puma_max_threads].set(labels, stats['max_threads'])
+          metrics[:puma_idle_threads].set(labels, stats['running'] + stats['pool_capacity'] - stats['max_threads'])
+        end
+      end
+    end
+  end
+end