Skip to content

Commit

Permalink
Add custom collectors to Prometheus registry (#16)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
dmexe authored May 18, 2021
1 parent 58df65f commit 50a0787
Show file tree
Hide file tree
Showing 5 changed files with 62 additions and 19 deletions.
4 changes: 4 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
Original file line number Diff line number Diff line change
@@ -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)
23 changes: 7 additions & 16 deletions lib/bm/instrumentations/management/server.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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>]
Expand All @@ -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]]
Expand Down
15 changes: 14 additions & 1 deletion lib/bm/instrumentations/puma/collector.rb
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand Down
5 changes: 3 additions & 2 deletions lib/puma/plugin/management_server.rb
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
require 'puma/dsl'
require 'puma/plugin'
require 'bm/instrumentations'
require 'bm/instrumentations/puma/collector'

module BM
module Instrumentations
Expand Down Expand Up @@ -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
Expand Down

0 comments on commit 50a0787

Please sign in to comment.