Add custom collectors to Prometheus registry

The gem `prometheus-client` doesn't support a custom collectors [issue
90](prometheus/client_ruby#90). This PL add
missing functionality, but only if the management server is used.

Author: dmexe

CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+
+- Ability to use custom metric collectors
+
 ### Changed
 
 - __Breaking__ The class `BM::Instrumentations::Aws::Collector` turned into gem private,

lib/bm/instrumentations/internal/prometheus_registry_custom_collectors.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module BM
+  module Instrumentations
+    # Extends the Prometheus registry to support custom metric collectors
+    #
+    # @see https://github.com/prometheus/client_ruby/issues/90
+    module PrometheusRegistryCustomCollectors
+      # Initialize custom_collectors array
+      def initialize(*args, **kwargs)
+        super(*args, **kwargs)
+        @custom_collectors = []
+      end
+
+      # Registers an updater that poll and update metrics periodically
+      #
+      # @param collector [Proc]
+      def add_custom_collector(&collector)
+        @mutex.synchronize do
+          @custom_collectors << collector
+        end
+      end
+
+      # Invokes all registered custom collectors to update metrics
+      #
+      # @return [void]
+      def custom_collectors!
+        @custom_collectors.each(&:call)
+      end
+    end
+  end
+end
+
+Prometheus::Client::Registry.prepend(BM::Instrumentations::PrometheusRegistryCustomCollectors)

lib/bm/instrumentations/management/server.rb

@@ -16,18 +16,14 @@ module Management
       # @param host [String] is a bind address that a server uses for listening (default: `0.0.0.0`)
       # @param logger [Logger, nil] is a logger instance for notifications (default: `Logger.new($stdout)`)
       # @param registry [Prometheus::Client::Registry, nil] override a default Prometheus registry
-      # @param puma_launcher [Puma::Launcher, nil]
       #
       # @return [Server]
-      def self.server(port: nil, host: nil, logger: nil, registry: nil, puma_launcher: nil)
-        registry ||= ::Prometheus::Client.registry
-        puma_metrics = puma_launcher ? Puma::Collector.new(registry: registry, launcher: puma_launcher) : nil
+      def self.server(port: nil, host: nil, logger: nil, registry: nil)
         Server.new(
           port: port || 9990,
           host: host || '0.0.0.0',
           logger: logger || ::Logger.new($stdout, progname: Server.name),
-          registry: registry,
-          puma_metrics: puma_metrics
+          registry: registry || ::Prometheus::Client.registry
         )
       end
 
@@ -45,9 +41,8 @@ def self.server(port: nil, host: nil, logger: nil, registry: nil, puma_launcher:
       # @attr [Logger] logger
       # @attr [Prometheus::Client::Registry] registry
       # @attr [Puma::Events] events
-      # @attr [Puma::Collector, nil] puma_metrics
       class Server
-        attr_reader :host, :port, :logger, :registry, :events, :puma_metrics
+        attr_reader :host, :port, :logger, :registry, :events
 
         # The socket backlog value
         BACKLOG = 3
@@ -65,15 +60,13 @@ class Server
         # @param port [Integer]
         # @param logger [Logger]
         # @param registry [Prometheus::Client::Registry]
-        # @param puma_metrics [Puma::Collector, nil]
         #
         # @api private
-        def initialize(port:, host:, logger:, registry:, puma_metrics: nil)
+        def initialize(port:, host:, logger:, registry:)
           @host = host
           @port = port
           @logger = logger
           @registry = registry
-          @puma_metrics = puma_metrics
         end
 
         # Creates a management server backed by {Puma::Server} then bind and
@@ -119,7 +112,7 @@ def puma_options
         #
         # @return [#call] a frozen rack application
         def rack_app
-          ::Rack::ShowExceptions.new(ServeEndpoints.new(registry, puma_metrics)).freeze
+          ::Rack::ShowExceptions.new(ServeEndpoints.new(registry)).freeze
         end
 
         # Handles a running server instance
@@ -161,10 +154,8 @@ class ServeEndpoints
           NOT_FOUND = [404, PLAIN_TEXT, ['not found']].freeze
 
           # @param registry [Prometheus::Client::Registry]
-          # @param puma_metrics [Puma::Metrics, nil]
-          def initialize(registry, puma_metrics)
+          def initialize(registry)
             @registry = registry
-            @puma_metrics = puma_metrics
           end
 
           # @param env [Hash<String, Any>]
@@ -190,7 +181,7 @@ def call(env) # rubocop:disable Metrics/MethodLength
 
           # @return [(Integer, Hash<String, String>, Array<String>)]
           def metrics
-            @puma_metrics&.update
+            @registry.custom_collectors! if @registry.respond_to?(:custom_collectors!)
 
             text = Prometheus::Client::Formats::Text.marshal(@registry)
             [200, METRICS_TEXT, [text]]

lib/bm/instrumentations/puma/collector.rb

@@ -2,8 +2,9 @@
 
 require 'socket'
 require 'puma'
-require 'tcp_server_socket_backlog/tcp_server_socket_backlog'
 
+require 'tcp_server_socket_backlog/tcp_server_socket_backlog'
+require 'bm/instrumentations/internal/prometheus_registry_custom_collectors'
 require_relative 'metrics_collection'
 
 module BM
@@ -29,6 +30,11 @@ def initialize(registry:, launcher:)
           metrics_collection.server_version(::Puma::Server::VERSION)
         end
 
+        # @return [Proc]
+        def to_proc
+          -> { update }
+        end
+
         # Updates Puma metrics in the registry
         def update
           metrics_collection.update_stats(launcher.stats)
@@ -39,6 +45,13 @@ def update
             metrics_collection.update_backlog(listener: index, backlog: backlog) if backlog
           end
         end
+
+        # @param launcher [Puma::Launcher]
+        # @param registry [Prometheus::Client::Registry, nil]
+        def self.install(launcher, registry: nil)
+          registry ||= Prometheus::Client.registry
+          registry.add_custom_collector(&Collector.new(registry: registry, launcher: launcher))
+        end
       end
     end
   end

lib/puma/plugin/management_server.rb

@@ -3,6 +3,7 @@
 require 'puma/dsl'
 require 'puma/plugin'
 require 'bm/instrumentations'
+require 'bm/instrumentations/puma/collector'
 
 module BM
   module Instrumentations
@@ -40,13 +41,13 @@ def start(launcher) # rubocop:disable Metrics/MethodLength
       host: launcher.options[:management_server_host],
       port: launcher.options[:management_server_port],
       logger: launcher.options[:management_server_logger],
-      registry: launcher.options[:management_server_registry],
-      puma_launcher: launcher
+      registry: launcher.options[:management_server_registry]
     }
 
     # @type [Puma::Events]
     events = launcher.events
     server = BM::Instrumentations::Management.server(**args)
+    BM::Instrumentations::Puma::Collector.install(launcher, registry: args[:registry])
 
     events.on_booted do
       running = server.run