client.rb in dawanda-statsd-client-0.1.6

- old
+ new

@@ -3,29 +3,41 @@
 require 'yaml'
 
 class Statsd
   class << self 
     def increment(metric, options={})
-      value = options.is_a?(Fixnum) ? options : (options[:by] || 1)
-      client.update_stats(metric, value*factor, (options[:sample_rate] || 1))
+      if options.is_a?(Fixnum)
+        value = options
+        sample_rate = 1
+      else
+        value = (options[:by] || 1)
+        sample_rate = (options[:sample_rate] || 1)
+      end
+      client.update_stats(metric, value*factor, sample_rate)
     end
     alias_method :inc, :increment
   
     def decrement(metric, options={})
-      value = options.is_a?(Fixnum) ? options : (options[:by] || 1)
-      client.update_stats(metric, value*factor*(-1), (options[:sample_rate] || 1))
+      if options.is_a?(Fixnum)
+        value = options
+        sample_rate = 1
+      else
+        value = (options[:by] || 1)
+        sample_rate = (options[:sample_rate] || 1)
+      end
+      client.update_stats(metric, value*factor*(-1), sample_rate)
     end
     alias_method :dec, :decrement
 
     def timing(metric, value)
       client.timing(metric, value)
     end
     alias_method :time, :timing
     
     def client
       return Statsd::DummyClient if deactivated?
-      @client ||= Statsd::Client.new(host, port)
+      @client ||= Statsd::Client.new(host, port, tcp?)
     end
   
     def config
       home = File.expand_path('~')
       files = ["#{home}/.statsd-client.yml", '/etc/statsd-client.yml']
@@ -40,10 +52,14 @@
     end
   
     def port
       config['port'] || 8125
     end
+
+    def tcp?
+      config['tcp']
+    end
   
     # statds reports with default configs 1/10 of actual value
     def factor
       config['factor'] || 1
     end
@@ -62,18 +78,19 @@
   end
   
   
   class Client
     VERSION = File.read( File.join(File.dirname(__FILE__),'..', '..', 'VERSION') ).strip
-    attr_reader :host, :port
+    attr_reader :host, :port, :tcp
 
     # Initializes a Statsd client.
     #
     # @param [String] host
     # @param [Integer] port
-    def initialize(host = 'localhost', port = 8125)
-      @host, @port = host, port
+    # @param [Boolean] tcp
+    def initialize(host = 'localhost', port = 8125, tcp = false)
+      @host, @port, @tcp = host, port, tcp
     end
 
     # Sends timing statistics.
     #
     # @param [Array, String] stats name of statistic(s) being updated
@@ -121,10 +138,11 @@
     end
 
     private
 
     def send(data, sample_rate = 1)
+      #puts "sending #{data} with sample #{sample_rate}"
       sampled_data = {}
       
       if sample_rate < 1
         if Kernel.rand <= sample_rate
           data.each do |k,v|
@@ -133,15 +151,23 @@
         end
       else
         sampled_data = data
       end
 
-      socket = UDPSocket.new
-
+      if self.tcp
+        socket = TCPSocket.new( self.host, self.port)
+      else
+        socket = UDPSocket.new
+      end
+      
       begin
         sampled_data.each do |k,v|
           message = [k,v].join(':')
-          socket.send(message, 0, self.host, self.port)
+          if self.tcp
+            socket.send(message, 0)
+          else
+            socket.send(message, 0, self.host, self.port)            
+          end
         end
       rescue Exception => e
         puts "Unexpected error: #{e}"
       ensure
         socket.close