lib/statsd/instrument.rb in statsd-instrument-1.2.0 vs lib/statsd/instrument.rb in statsd-instrument-1.3.0

@@ -10,15 +10,16 @@
 
 
 module StatsD
   class << self
     attr_accessor :host, :port, :mode, :logger, :enabled, :default_sample_rate,
-                  :prefix
+                  :prefix, :implementation
   end
   self.enabled = true
   self.default_sample_rate = 1
-  
+  self.implementation = :statsd
+
   TimeoutClass = defined?(::SystemTimer) ? ::SystemTimer : ::Timeout
 
   # StatsD.server = 'localhost:1234'
   def self.server=(conn)
     self.host, port = conn.split(':')
@@ -26,21 +27,21 @@
   end
 
   module Instrument
     def statsd_measure(method, name)
       add_to_method(method, name, :measure) do |old_method, new_method, metric_name, *args|
-        define_method(new_method) do |*args|
-          StatsD.measure(metric_name) { send(old_method, *args) }
+        define_method(new_method) do |*args, &block|
+          StatsD.measure(metric_name) { send(old_method, *args, &block) }
         end
       end
     end
 
     def statsd_count_success(method, name)
       add_to_method(method, name, :count_success) do |old_method, new_method, metric_name|
-        define_method(new_method) do |*args|
+        define_method(new_method) do |*args, &block|
           begin
-            truthiness = result = send(old_method, *args)
+            truthiness = result = send(old_method, *args, &block)
           rescue
             truthiness = false
             raise
           else
             truthiness = (yield(result) rescue false) if block_given?
@@ -52,13 +53,13 @@
       end
     end
 
     def statsd_count_if(method, name)
       add_to_method(method, name, :count_if) do |old_method, new_method, metric_name|
-        define_method(new_method) do |*args|
+        define_method(new_method) do |*args, &block|
           begin
-            truthiness = result = send(old_method, *args)
+            truthiness = result = send(old_method, *args, &block)
           rescue
             truthiness = false
             raise
           else
             truthiness = (yield(result) rescue false) if block_given?
@@ -70,13 +71,13 @@
       end
     end
 
     def statsd_count(method, name)
       add_to_method(method, name, :count) do |old_method, new_method, metric_name|
-        define_method(new_method) do |*args|
+        define_method(new_method) do |*args, &block|
           StatsD.increment(metric_name)
-          send(old_method, *args)
+          send(old_method, *args, &block)
         end
       end
     end
 
     private
@@ -100,11 +101,11 @@
 
   # glork:320|ms
   def self.measure(key, milli = nil)
     result = nil
     ms = milli || Benchmark.ms do
-      result = yield 
+      result = yield
     end
 
     write(key, ms, :ms)
     result
   end
@@ -112,13 +113,14 @@
   # gorets:1|c
   def self.increment(key, delta = 1, sample_rate = default_sample_rate)
     write(key, delta, :incr, sample_rate)
   end
 
-  #gaugor:333|g
-  def self.gauge(key, value, sample_rate = default_sample_rate)
-    write(key, value, :g, sample_rate)
+  # gaugor:333|g
+  # guagor:1234|kv|@1339864935 (statsite)
+  def self.gauge(key, value, sample_rate_or_epoch = default_sample_rate)
+    write(key, value, :g, sample_rate_or_epoch)
   end
 
   private
 
   def self.socket
@@ -134,13 +136,14 @@
     when :incr
       command << '|c'
     when :ms
       command << '|ms'
     when :g
-      command << '|g'
+      command << (self.implementation == :statsite ? '|kv' : '|g')
     end
 
-    command << "|@#{sample_rate}" if sample_rate < 1
+    command << "|@#{sample_rate}" if sample_rate < 1 || (self.implementation == :statsite && sample_rate > 1)
+    command << "\n" if self.implementation == :statsite
 
     if mode.to_s == 'production'
       socket_wrapper { socket.send(command, 0, host, port) }
     else
       logger.info "[StatsD] #{command}"