qubole-statsd-instrument 2.1.4

data/lib/statsd/instrument/backends/udp_backend.rb

@@ -0,0 +1,106 @@
+require 'monitor'
+
+module StatsD::Instrument::Backends
+  class UDPBackend < StatsD::Instrument::Backend
+
+    DEFAULT_IMPLEMENTATION = :statsd
+
+    include MonitorMixin
+
+    attr_reader :host, :port
+    attr_accessor :implementation
+
+    def initialize(server = nil, implementation = nil)
+      super()
+      self.server = server || "localhost:8125"
+      self.implementation = (implementation || DEFAULT_IMPLEMENTATION).to_sym
+    end
+
+    def collect_metric(metric)
+      unless implementation_supports_metric_type?(metric.type)
+        StatsD.logger.warn("[StatsD] Metric type #{metric.type.inspect} not supported on #{implementation} implementation.")
+        return false
+      end
+
+      if metric.sample_rate < 1.0 && rand > metric.sample_rate
+        return false
+      end
+
+      write_packet(generate_packet(metric))
+    end
+
+    def implementation_supports_metric_type?(type)
+      case type
+        when :h;  implementation == :datadog
+        when :kv; implementation == :statsite
+        else true
+      end
+    end
+
+    def server=(connection_string)
+      self.host, port = connection_string.split(':', 2)
+      self.port = port.to_i
+      invalidate_socket
+    end
+
+    def host=(host)
+      @host = host
+      invalidate_socket
+    end
+
+    def port=(port)
+      @port = port
+      invalidate_socket
+    end
+
+    def socket
+      if @socket.nil?
+        @socket = UDPSocket.new
+        @socket.connect(host, port)
+      end
+      @socket
+    end
+
+    def generate_packet(metric)
+      command = "#{metric.name}:#{metric.value}|#{metric.type}"
+      command << "|@#{metric.sample_rate}" if metric.sample_rate < 1 || (implementation == :statsite && metric.sample_rate > 1)
+      if metric.tags
+        if tags_supported?
+          if implementation == :datadog
+            command << "|##{metric.tags.join(',')}"
+          elsif implementation == :collectd
+            metric_tags = "#{metric.tags.join(',')}"
+            metric_tags = metric_tags.prepend("[") << "]"
+            command.prepend(metric_tags.gsub(":", "="))
+          end
+        else
+          StatsD.logger.warn("[StatsD] Tags are only supported on Datadog and CollectD implementation.")
+        end
+      end
+
+      command << "\n" if implementation == :statsite
+      command
+    end
+
+    def tags_supported?
+      implementation == :datadog || implementation == :collectd
+    end
+
+    def write_packet(command)
+      synchronize do
+        socket.send(command, 0) > 0
+      end
+    rescue ThreadError => e
+      # In cases where a TERM or KILL signal has been sent, and we send stats as
+      # part of a signal handler, locks cannot be acquired, so we do our best
+      # to try and send the command without a lock.
+      socket.send(command, 0) > 0
+    rescue SocketError, IOError, SystemCallError, Errno::ECONNREFUSED => e
+      StatsD.logger.error "[StatsD] #{e.class.name}: #{e.message}"
+    end
+
+    def invalidate_socket
+      @socket = nil
+    end
+  end
+end

data/lib/statsd/instrument/environment.rb

@@ -0,0 +1,54 @@
+require 'logger'
+
+# The environment module is used to detect, and initialize the environment in
+# which this library is active. It will use different default values based on the environment.
+module StatsD::Instrument::Environment
+  extend self
+
+  # Instantiates a default backend for the current environment.
+  #
+  # @return [StatsD::Instrument::Backend]
+  # @see #environment
+  def default_backend
+    case environment
+    when 'production', 'staging'
+      StatsD::Instrument::Backends::UDPBackend.new(ENV['STATSD_ADDR'], ENV['STATSD_IMPLEMENTATION'])
+    when 'test'
+      StatsD::Instrument::Backends::NullBackend.new
+    else
+      StatsD::Instrument::Backends::LoggerBackend.new(StatsD.logger)
+    end
+  end
+
+  # Detects the current environment, either by asking Rails, or by inspecting environment variables.
+  #
+  # - Within a Rails application, <tt>Rails.env</tt> is used.
+  # - It will check the following environment variables in order: <tt>RAILS_ENV</tt>, <tt>RACK_ENV</tt>, <tt>ENV</tt>.
+  # - If none of these are set, it will return <tt>development</tt>
+  #
+  # @return [String] The detected environment.
+  def environment
+    if defined?(Rails) && Rails.respond_to?(:env)
+      Rails.env.to_s
+    else
+      ENV['RAILS_ENV'] || ENV['RACK_ENV'] || ENV['ENV'] || 'development'
+    end
+  end
+
+  # Sets default values for sample rate and logger.
+  #
+  # - Default sample rate is set to the value in the STATSD_SAMPLE_RATE environment variable,
+  #   or 1.0 otherwise. See {StatsD#default_sample_rate}
+  # - {StatsD#logger} is set to a logger that send output to stderr.
+  #
+  # If you are including this library inside a Rails environment, additional initialization will
+  # be done as part of the {StatsD::Instrument::Railtie}.
+  #
+  # @return [void]
+  def setup
+    StatsD.default_sample_rate = ENV.fetch('STATSD_SAMPLE_RATE', 1.0).to_f
+    StatsD.logger = Logger.new($stderr)
+  end
+end
+
+StatsD::Instrument::Environment.setup

data/lib/statsd/instrument/helpers.rb

@@ -0,0 +1,14 @@
+module StatsD::Instrument::Helpers
+  def capture_statsd_calls(&block)
+    mock_backend = StatsD::Instrument::Backends::CaptureBackend.new
+    old_backend, StatsD.backend = StatsD.backend, mock_backend
+    block.call
+    mock_backend.collected_metrics
+  ensure
+    if old_backend.kind_of?(StatsD::Instrument::Backends::CaptureBackend)
+      old_backend.collected_metrics.concat(mock_backend.collected_metrics)
+    end
+
+    StatsD.backend = old_backend
+  end
+end

data/lib/statsd/instrument/matchers.rb

@@ -0,0 +1,96 @@
+require 'rspec/expectations'
+require 'rspec/core/version'
+
+module StatsD::Instrument::Matchers
+  CUSTOM_MATCHERS = {
+    increment: :c,
+    measure: :ms,
+    gauge: :g,
+    histogram: :h,
+    set: :s,
+    key_value: :kv
+  }
+
+  class Matcher
+    include RSpec::Matchers::Composable if RSpec::Core::Version::STRING.start_with?('3')
+    include StatsD::Instrument::Helpers
+
+    def initialize(metric_type, metric_name, options = {})
+      @metric_type = metric_type
+      @metric_name = metric_name
+      @options = options
+    end
+
+    def matches?(block)
+      begin
+        expect_statsd_call(@metric_type, @metric_name, @options, &block)
+      rescue RSpec::Expectations::ExpectationNotMetError => e
+        @message = e.message
+
+        false
+      end
+    end
+
+    def failure_message
+      @message
+    end
+
+    def failure_message_when_negated
+      "No StatsD calls for metric #{@metric_name} expected."
+    end
+
+    def supports_block_expectations?
+      true
+    end
+
+    private
+
+    def expect_statsd_call(metric_type, metric_name, options, &block)
+      metrics = capture_statsd_calls(&block)
+      metrics = metrics.select { |m| m.type == metric_type && m.name == metric_name }
+
+      raise RSpec::Expectations::ExpectationNotMetError, "No StatsD calls for metric #{metric_name} were made." if metrics.empty?
+      raise RSpec::Expectations::ExpectationNotMetError, "The numbers of StatsD calls for metric #{metric_name} was unexpected. Expected #{options[:times].inspect}, got #{metrics.length}" if options[:times] && options[:times] != metrics.length
+
+      [:sample_rate, :value, :tags].each do |expectation|
+        next unless options[expectation]
+
+        num_matches = metrics.count do |m|
+          matcher = RSpec::Matchers::BuiltIn::Match.new(options[expectation])
+          matcher.matches?(m.public_send(expectation))
+        end
+
+        found = options[:times] ? num_matches == options[:times] : num_matches > 0
+
+        if !found
+          message = metric_information(metric_name, options, metrics, expectation)
+          raise RSpec::Expectations::ExpectationNotMetError, message
+        end
+      end
+
+      true
+    end
+
+    def metric_information(metric_name, options, metrics, expectation)
+      message = "expected StatsD #{expectation.inspect} for metric '#{metric_name}' to be called"
+
+      message += "\n  "
+      message += options[:times] ? "exactly #{options[:times]} times" : "at least once"
+      message += " with: #{options[expectation]}"
+
+      message += "\n  captured metric values: #{metrics.map(&expectation).join(', ')}"
+
+      message
+    end
+  end
+
+  CUSTOM_MATCHERS.each do |method_name, metric_type|
+    klass = Class.new(Matcher)
+
+    define_method "trigger_statsd_#{method_name}" do |metric_name, options = {}|
+      klass.new(metric_type, metric_name, options)
+    end
+
+    StatsD::Instrument::Matchers.const_set(method_name.capitalize, klass)
+  end
+end

data/lib/statsd/instrument/metric.rb

@@ -0,0 +1,117 @@
+# The Metric class represents a metric sample to be send by a backend.
+#
+# @!attribute type
+#   @return [Symbol] The metric type. Must be one of {StatsD::Instrument::Metric::TYPES}
+# @!attribute name
+#   @return [String] The name of the metric. {StatsD#prefix} will automatically be applied
+#     to the metric in the constructor, unless the <tt>:no_prefix</tt> option is set.
+# @!attribute value
+#   @see #default_value
+#   @return [Numeric, String] The value to collect for the metric. Depending on the metric
+#     type, <tt>value</tt> can be a string, integer, or float.
+# @!attribute sample_rate
+#   The sample rate to use for the metric. How the sample rate is handled differs per backend.
+#   The UDP backend will actually sample metric submissions based on the sample rate, while
+#   the logger backend will just include the sample rate in its output for debugging purposes.
+#   @see StatsD#default_sample_rate
+#   @return [Float] The sample rate to use for this metric. This should be a value between
+#     0 and 1. If not set, it will use the default sample rate set to {StatsD#default_sample_rate}.
+# @!attribute tags
+#   The tags to associate with the metric.
+#   @note Only the Datadog implementation supports tags.
+#   @see .normalize_tags
+#   @return [Array<String>, Hash<String, String>, nil] the tags to associate with the metric.
+#     You can either specify the tags as an array of strings, or a Hash of key/value pairs.
+#
+# @see StatsD The StatsD module contains methods that generate metric instances.
+# @see StatsD::Instrument::Backend A StatsD::Instrument::Backend is used to collect metrics.
+#
+class StatsD::Instrument::Metric
+
+  attr_accessor :type, :name, :value, :sample_rate, :tags
+
+  # Initializes a new metric instance.
+  # Normally, you don't want to call this method directly, but use one of the metric collection
+  # methods on the {StatsD} module.
+  #
+  # @option options [Symbol] :type The type of the metric.
+  # @option options [String] :name The name of the metric without prefix.
+  # @option options [Boolean] :no_prefix Set to <tt>true</tt> if you don't want to apply {StatsD#prefix}
+  # @option options [Numeric, String, nil] :value The value to collect for the metric. If set to
+  #   <tt>nil>/tt>, {#default_value} will be used.
+  # @option options [Numeric, nil] :sample_rate The sample rate to use. If not set, it will use
+  #   {StatsD#default_sample_rate}.
+  # @option options [Array<String>, Hash<String, String>, nil] :tags The tags to apply to this metric.
+  #   See {.normalize_tags} for more information.
+  def initialize(options = {})
+    @type = options[:type] or raise ArgumentError, "Metric :type is required."
+    @name = options[:name] or raise ArgumentError, "Metric :name is required."
+    @name = normalize_name(@name)
+    @name = StatsD.prefix ? "#{StatsD.prefix}.#{@name}" : @name unless options[:no_prefix]
+
+    @value       = options[:value] || default_value
+    @sample_rate = options[:sample_rate] || StatsD.default_sample_rate
+    @tags        = StatsD::Instrument::Metric.normalize_tags(options[:tags])
+  end
+
+  # The default value for this metric, which will be used if it is not set.
+  #
+  # A default value is only defined for counter metrics (<tt>1</tt>). For all other
+  # metric types, this emthod will raise an <tt>ArgumentError</tt>.
+  #
+  # @return [Numeric, String] The default value for this metric.
+  # @raise ArgumentError if the metric type doesn't have a default value
+  def default_value
+    case type
+      when :c; 1
+      else raise ArgumentError, "A value is required for metric type #{type.inspect}."
+    end
+  end
+
+  # @private
+  # @return [String]
+  def to_s
+    str = "#{TYPES[type]} #{name}:#{value}"
+    str << " @#{sample_rate}" if sample_rate != 1.0
+    str << " " << tags.map { |t| "##{t}"}.join(' ') if tags
+    str
+  end
+
+  # @private
+  # @return [String]
+  def inspect
+    "#<StatsD::Instrument::Metric #{self.to_s}>"
+  end
+
+  # The metric types that are supported by this library. Note that every StatsD server
+  # implementation only supports a subset of them.
+  TYPES = {
+    c:  'increment',
+    ms: 'measure',
+    g:  'gauge',
+    h:  'histogram',
+    kv: 'key/value',
+    s:  'set',
+  }
+
+  # Strip metric names of special characters used by StatsD line protocol, replace with underscore
+  #
+  # @param name [String]
+  # @return [String]
+  def normalize_name(name)
+    name.tr(':|@'.freeze, '_')
+  end
+
+  # Utility function to convert tags to the canonical form.
+  #
+  # - Tags specified as key value pairs will be converted into an array
+  # - Tags are normalized to only use word characters and underscores.
+  #
+  # @param tags [Array<String>, Hash<String, String>, nil] Tags specified in any form.
+  # @return [Array<String>, nil] the list of tags in canonical form.
+  def self.normalize_tags(tags)
+    return unless tags
+    tags = tags.map { |k, v| k.to_s + ":".freeze + v.to_s } if tags.is_a?(Hash)
+    tags.map { |tag| tag.tr('|,'.freeze, ''.freeze) }
+  end
+end

data/lib/statsd/instrument/metric_expectation.rb

@@ -0,0 +1,67 @@
+# @private
+class StatsD::Instrument::MetricExpectation
+
+  attr_accessor :times, :type, :name, :value, :sample_rate, :tags
+  attr_reader :ignore_tags
+
+  def initialize(options = {})
+    @type = options[:type] or raise ArgumentError, "Metric :type is required."
+    @name = options[:name] or raise ArgumentError, "Metric :name is required."
+    @name = StatsD.prefix ? "#{StatsD.prefix}.#{@name}" : @name unless options[:no_prefix]
+    @tags = StatsD::Instrument::Metric.normalize_tags(options[:tags])
+    @times = options[:times] or raise ArgumentError, "Metric :times is required."
+    @sample_rate = options[:sample_rate]
+    @value = options[:value]
+    @ignore_tags = StatsD::Instrument::Metric.normalize_tags(options[:ignore_tags])
+  end
+
+  def matches(actual_metric)
+    return false if sample_rate && sample_rate != actual_metric.sample_rate
+    return false if value && value != actual_metric.value
+
+    if tags
+
+      expected_tags = Set.new(tags)
+      actual_tags = Set.new(actual_metric.tags)
+
+      if ignore_tags
+        ignored_tags = Set.new(ignore_tags) - expected_tags
+        actual_tags -= ignored_tags
+
+        if ignore_tags.is_a?(Array)
+          actual_tags.delete_if{ |key| ignore_tags.include?(key.split(":").first) }
+        end
+      end
+
+      return expected_tags.subset?(actual_tags)
+    end
+    true
+  end
+
+  def default_value
+    case type
+      when :c; 1
+    end
+  end
+
+  TYPES = {
+      c:  'increment',
+      ms: 'measure',
+      g:  'gauge',
+      h:  'histogram',
+      kv: 'key/value',
+      s:  'set',
+  }
+
+  def to_s
+    str = "#{TYPES[type]} #{name}:#{value}"
+    str << " @#{sample_rate}" if sample_rate != 1.0
+    str << " " << tags.map { |t| "##{t}"}.join(' ') if tags
+    str << " times:#{times}" if times > 1
+    str
+  end
+
+  def inspect
+    "#<StatsD::Instrument::MetricExpectation #{self.to_s}>"
+  end
+end

data/lib/statsd/instrument/railtie.rb

@@ -0,0 +1,14 @@
+# This Railtie runs some initializers that will set the logger to <tt>Rails#logger</tt>,
+# and will initialize the {StatsD#backend} based on the Rails environment.
+#
+# @see StatsD::Instrument::Environment
+class StatsD::Instrument::Railtie < Rails::Railtie
+
+  initializer 'statsd-instrument.use_rails_logger' do
+    ::StatsD.logger = Rails.logger
+  end
+
+  initializer 'statsd-instrument.setup_backend', after: 'statsd-instrument.use_rails_logger' do
+    ::StatsD.backend = ::StatsD::Instrument::Environment.default_backend
+  end
+end

data/lib/statsd/instrument/version.rb

@@ -0,0 +1,5 @@
+module StatsD
+  module Instrument
+    VERSION = "2.1.4"
+  end
+end