Use Module builder pattern for method instrumentation -- part 2

Gemfile

@@ -1,2 +1,4 @@
 source "https://rubygems.org"
 gemspec
+
+gem "pry-byebug"

lib/statsd/instrument.rb

@@ -56,19 +56,21 @@ def self.generate_metric_name(metric_name, callee, *args)
       metric_name.respond_to?(:call) ? metric_name.call(callee, args).gsub('::', '.') : metric_name.gsub('::', '.')
     end
 
+    def self.duration
+      start = current_timestamp
+      yield
+      current_timestamp - start
+    end
+
     if Process.respond_to?(:clock_gettime)
       # @private
-      def self.duration
-        start = Process.clock_gettime(Process::CLOCK_MONOTONIC)
-        yield
-        Process.clock_gettime(Process::CLOCK_MONOTONIC) - start
+      def self.current_timestamp
+        Process.clock_gettime(Process::CLOCK_MONOTONIC)
       end
     else
       # @private
-      def self.duration
-        start = Time.now
-        yield
-        Time.now - start
+      def self.current_timestamp
+        Time.now
       end
     end
 
@@ -206,8 +208,43 @@ def statsd_remove_measure(method, name)
       remove_from_method(method, name, :measure)
     end
 
+    def self.count_method(
+      method_name,
+      on_success: DEFAULT_ON_SUCCESS,
+      on_exception: DEFAULT_ON_EXCEPTION,
+      **metric_options
+    )
+      metric = StatsD::Instrument::Metric.new(metric_options.merge(type: :c))
+
+      StatsD::Instrument::MethodCounter.new(
+        method_name,
+        metric,
+        on_success: on_success,
+        on_exception: on_exception,
+      )
+    end
+
+    def self.measure_method(
+      method_name,
+      on_success: DEFAULT_ON_SUCCESS,
+      on_exception: DEFAULT_ON_EXCEPTION,
+      **metric_options
+    )
+      metric = StatsD::Instrument::Metric.new(metric_options.merge(type: :ms, value: 0))
+
+      StatsD::Instrument::MethodMeasurer.new(
+        method_name,
+        metric,
+        on_success: on_success,
+        on_exception: on_exception,
+      )
+    end
+
     private
 
+    DEFAULT_ON_SUCCESS = -> (_metric, _result) {}
+    DEFAULT_ON_EXCEPTION = -> (metric, _ex) { metric.discard }
+
     def statsd_instrumentation_for(method, name, action)
       unless statsd_instrumentations.key?([method, name, action])
         mod = Module.new do
@@ -403,5 +440,7 @@ def collect_metric(options)
 require 'statsd/instrument/helpers'
 require 'statsd/instrument/assertions'
 require 'statsd/instrument/metric_expectation'
+require 'statsd/instrument/method_counter'
+require 'statsd/instrument/method_measurer'
 require 'statsd/instrument/matchers' if defined?(::RSpec)
 require 'statsd/instrument/railtie' if defined?(Rails)

lib/statsd/instrument/method_counter.rb

@@ -0,0 +1,67 @@
+module StatsD
+  module Instrument
+    class MethodCounter < Module
+      attr_reader(:method_name)
+
+      def initialize(
+        method_name,
+        metric,
+        on_success:,
+        on_exception:
+      )
+        @method_name = method_name
+        @metric = metric
+        @on_success = on_success
+        @on_exception = on_exception
+      end
+
+      def prepend_features(mod)
+        super(mod)
+        preserve_visibility(mod, method_name) do
+          define_method_lifecycle(@metric, @on_success, @on_exception)
+        end
+      end
+
+      def inspect
+        "#<#{self.class.name}[#{method_name.inspect}]>"
+      end
+
+      private
+
+      def define_method_lifecycle(metric, on_success, on_exception)
+        define_method(method_name) do |*args, &block|
+          begin
+            result = super(*args, &block)
+            on_success.call(metric, result)
+            result
+          rescue => ex
+            begin
+              on_exception.call(metric, ex)
+            ensure
+              raise(ex)
+            end
+          ensure
+            StatsD.backend.collect_metric(metric) unless metric.discarded?
+          end
+        end
+      end
+
+      def preserve_visibility(mod, method_name)
+        original_visibility = method_visibility(mod, method_name)
+        yield
+        __send__(original_visibility, method_name)
+      end
+
+      def method_visibility(mod, method)
+        case
+        when mod.private_method_defined?(method)
+          :private
+        when mod.protected_method_defined?(method)
+          :protected
+        else
+          :public
+        end
+      end
+    end
+  end
+end

lib/statsd/instrument/method_measurer.rb

@@ -0,0 +1,28 @@
+module StatsD
+  module Instrument
+    class MethodMeasurer < MethodCounter
+      private
+
+      def define_method_lifecycle(metric, on_success, on_exception)
+        define_method(method_name) do |*args, &block|
+          start = StatsD::Instrument.current_timestamp
+          begin
+            result = super(*args, &block)
+            metric.value = StatsD::Instrument.current_timestamp - start
+            on_success.call(metric, result)
+            result
+          rescue => ex
+            metric.value = StatsD::Instrument.current_timestamp - start
+            begin
+              on_exception.call(metric, ex)
+            ensure
+              raise(ex)
+            end
+          ensure
+            StatsD.backend.collect_metric(metric) unless metric.discarded?
+          end
+        end
+      end
+    end
+  end
+end

lib/statsd/instrument/metric.rb

@@ -47,7 +47,7 @@ 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]
-
+    @discarded   = false
     @value       = options[:value] || default_value
     @sample_rate = options[:sample_rate] || StatsD.default_sample_rate
     @tags        = StatsD::Instrument::Metric.normalize_tags(options[:tags])
@@ -82,6 +82,14 @@ def inspect
     "#<StatsD::Instrument::Metric #{self.to_s}>"
   end
 
+  def discard
+    @discarded = true
+  end
+
+  def discarded?
+    !!@discarded
+  end
+
   # The metric types that are supported by this library. Note that every StatsD server
   # implementation only supports a subset of them.
   TYPES = {