Merge pull request #21 from github/namespace

Namespace under GitHub

.gitignore

@@ -40,3 +40,4 @@ pkg
 #
 # For vim:
 #*.swp
+

Gemfile

@@ -0,0 +1,2 @@
+source "https://rubygems.org"
+gemspec

Gemfile.lock

@@ -0,0 +1,21 @@
+PATH
+  remote: .
+  specs:
+    statsd-ruby (0.4.0.github)
+
+GEM
+  remote: https://rubygems.org/
+  specs:
+    minitest (5.9.0)
+    rake (11.2.2)
+
+PLATFORMS
+  ruby
+
+DEPENDENCIES
+  minitest (~> 5.9)
+  rake (~> 11.2)
+  statsd-ruby!
+
+BUNDLED WITH
+   1.11.2

README.rdoc

@@ -5,7 +5,7 @@ A Ruby statsd client (https://github.com/etsy/statsd)
 = Installing
 
 Bundler:
-  gem "statsd-ruby", :require => "statsd"
+  gem "statsd-ruby", :require => "github/statsd"
 
 = Testing
 
@@ -14,7 +14,7 @@ Run the specs with <tt>rake spec</tt>
 Run the specs and include live integration specs with <tt>LIVE=true rake spec</tt>. Note: This will test over a real UDP socket.
 
 == Contributing to statsd
- 
+
 * Check out the latest master to make sure the feature hasn't been implemented or the bug hasn't been fixed yet
 * Check out the issue tracker to make sure someone already hasn't requested it and/or contributed it
 * Fork the project

Rakefile

@@ -1,39 +1,11 @@
 require 'rubygems'
 require 'rake'
 
-require 'jeweler'
-Jeweler::Tasks.new do |gem|
-  # gem is a Gem::Specification... see http://docs.rubygems.org/read/chapter/20 for more options
-  gem.name = "statsd-ruby"
-  gem.homepage = "http://github.com/reinh/statsd"
-  gem.license = "MIT"
-  gem.summary = %Q{A Statsd client in Ruby}
-  gem.description = %Q{A Statsd client in Ruby}
-  gem.email = "[email protected]"
-  gem.authors = ["Rein Henrichs"]
-  gem.add_development_dependency "minitest", ">= 0"
-  gem.add_development_dependency "yard", "~> 0.6.0"
-  gem.add_development_dependency "jeweler", "~> 1.5.2"
-  gem.add_development_dependency "rcov", ">= 0"
-end
-Jeweler::RubygemsDotOrgTasks.new
-
 require 'rake/testtask'
 Rake::TestTask.new(:spec) do |spec|
   spec.libs << 'lib' << 'spec'
   spec.pattern = 'spec/**/*_spec.rb'
   spec.verbose = true
 end
 
-require 'rcov/rcovtask'
-Rcov::RcovTask.new do |spec|
-  spec.libs << 'lib' << 'spec'
-  spec.pattern = 'spec/**/*_spec.rb'
-  spec.verbose = true
-  spec.rcov_opts << "--exclude spec,gems"
-end
-
 task :default => :spec
-
-require 'yard'
-YARD::Rake::YardocTask.new

lib/github/statsd.rb

@@ -0,0 +1,255 @@
+require 'socket'
+require 'zlib'
+
+module GitHub
+  # = Statsd: A Statsd client (https://github.com/etsy/statsd)
+  #
+  # @example Set up a global Statsd client for a server on localhost:8125
+  #   $statsd = Statsd.new 'localhost', 8125
+  # @example Send some stats
+  #   $statsd.increment 'garets'
+  #   $statsd.timing 'glork', 320
+  # @example Use {#time} to time the execution of a block
+  #   $statsd.time('account.activate') { @account.activate! }
+  # @example Create a namespaced statsd client and increment 'account.activate'
+  #   statsd = Statsd.new('localhost').tap{|sd| sd.namespace = 'account'}
+  #   statsd.increment 'activate'
+  class Statsd
+    class UDPClient
+      attr_reader :sock
+
+      def initialize(address, port = nil)
+        address, port = address.split(':') if address.include?(':')
+        addrinfo = Addrinfo.ip(address)
+
+        @sock = UDPSocket.new(addrinfo.pfamily)
+        @sock.connect(addrinfo.ip_address, port)
+      end
+
+      def send(msg)
+        sock.write(msg)
+      rescue SystemCallError
+        nil
+      end
+    end
+
+    class SecureUDPClient < UDPClient
+      def initialize(address, port, key)
+        super(address, port)
+        @key = key
+      end
+
+      def send(msg)
+        super(signed_payload(msg))
+      end
+
+      private
+      # defer loading openssl and securerandom unless needed. this shaves ~10ms off
+      # of baseline require load time for environments that don't require message signing.
+      def self.setup_openssl
+        @sha256 ||= begin
+          require 'securerandom'
+          require 'openssl'
+          OpenSSL::Digest::SHA256.new
+        end
+      end
+
+      def signed_payload(message)
+        sha256 = SecureUDPClient.setup_openssl
+        payload = timestamp + nonce + message
+        signature = OpenSSL::HMAC.digest(sha256, @key, payload)
+        signature + payload
+      end
+
+      def timestamp
+        [Time.now.to_i].pack("Q<")
+      end
+
+      def nonce
+        SecureRandom.random_bytes(4)
+      end
+    end
+
+    # A namespace to prepend to all statsd calls.
+    attr_reader :namespace
+
+    def namespace=(namespace)
+      @namespace = namespace
+      @prefix = namespace ? "#{@namespace}." : "".freeze
+    end
+
+    # All the endpoints where StatsD will report metrics
+    attr_reader :shards
+
+    # The client class used to initialize shard instances and send metrics.
+    attr_reader :client_class
+
+    #characters that will be replaced with _ in stat names
+    RESERVED_CHARS_REGEX = /[\:\|\@]/
+
+    COUNTER_TYPE = "c".freeze
+    TIMING_TYPE = "ms".freeze
+    GAUGE_TYPE = "g".freeze
+    HISTOGRAM_TYPE = "h".freeze
+
+    def initialize(client_class = nil)
+      @shards = []
+      @client_class = client_class || UDPClient
+      self.namespace = nil
+    end
+
+    def self.simple(addr, port = nil)
+      self.new.add_shard(addr, port)
+    end
+
+    def add_shard(*args)
+      @shards << @client_class.new(*args)
+      self
+    end
+
+    def enable_buffering(buffer_size = nil)
+      return if @buffering
+      @shards.map! { |client| Buffer.new(client, buffer_size) }
+      @buffering = true
+    end
+
+    def disable_buffering
+      return unless @buffering
+      flush_all
+      @shards.map! { |client| client.base_client }
+      @buffering = false
+    end
+
+    def flush_all
+      return unless @buffering
+      @shards.each { |client| client.flush }
+    end
+
+
+    # Sends an increment (count = 1) for the given stat to the statsd server.
+    #
+    # @param stat (see #count)
+    # @param sample_rate (see #count)
+    # @see #count
+    def increment(stat, sample_rate=1); count stat, 1, sample_rate end
+
+    # Sends a decrement (count = -1) for the given stat to the statsd server.
+    #
+    # @param stat (see #count)
+    # @param sample_rate (see #count)
+    # @see #count
+    def decrement(stat, sample_rate=1); count stat, -1, sample_rate end
+
+    # Sends an arbitrary count for the given stat to the statsd server.
+    #
+    # @param [String] stat stat name
+    # @param [Integer] count count
+    # @param [Integer] sample_rate sample rate, 1 for always
+    def count(stat, count, sample_rate=1); send stat, count, COUNTER_TYPE, sample_rate end
+
+    # Sends an arbitary gauge value for the given stat to the statsd server.
+    #
+    # @param [String] stat stat name.
+    # @param [Numeric] gauge value.
+    # @example Report the current user count:
+    #   $statsd.gauge('user.count', User.count)
+    def gauge(stat, value)
+      send stat, value, GAUGE_TYPE
+    end
+
+    # Sends a timing (in ms) for the given stat to the statsd server. The
+    # sample_rate determines what percentage of the time this report is sent. The
+    # statsd server then uses the sample_rate to correctly track the average
+    # timing for the stat.
+    #
+    # @param stat stat name
+    # @param [Integer] ms timing in milliseconds
+    # @param [Integer] sample_rate sample rate, 1 for always
+    def timing(stat, ms, sample_rate=1); send stat, ms, TIMING_TYPE, sample_rate end
+
+    # Reports execution time of the provided block using {#timing}.
+    #
+    # @param stat (see #timing)
+    # @param sample_rate (see #timing)
+    # @yield The operation to be timed
+    # @see #timing
+    # @example Report the time (in ms) taken to activate an account
+    #   $statsd.time('account.activate') { @account.activate! }
+    def time(stat, sample_rate=1)
+      start = Time.now
+      result = yield
+      timing(stat, ((Time.now - start) * 1000).round(5), sample_rate)
+      result
+    end
+
+    # Sends a histogram measurement for the given stat to the statsd server. The
+    # sample_rate determines what percentage of the time this report is sent. The
+    # statsd server then uses the sample_rate to correctly track the average
+    # for the stat.
+    def histogram(stat, value, sample_rate=1); send stat, value, HISTOGRAM_TYPE, sample_rate end
+
+    private
+    def sampled(sample_rate)
+      yield unless sample_rate < 1 and rand > sample_rate
+    end
+
+    def send(stat, delta, type, sample_rate=1)
+      sampled(sample_rate) do
+        stat = stat.to_s.dup
+        stat.gsub!(/::/, ".".freeze)
+        stat.gsub!(RESERVED_CHARS_REGEX, "_".freeze)
+
+        msg = String.new
+        msg << @prefix
+        msg << stat
+        msg << ":".freeze
+        msg << delta.to_s
+        msg << "|".freeze
+        msg << type
+        if sample_rate < 1
+          msg << "|@".freeze
+          msg << sample_rate.to_s
+        end
+
+        shard = select_shard(stat)
+        shard.send(msg)
+      end
+    end
+
+    def select_shard(stat)
+      if @shards.size == 1
+        @shards.first
+      else
+        @shards[Zlib.crc32(stat) % @shards.size]
+      end
+    end
+
+    class Buffer
+      DEFAULT_BUFFER_CAP = 512
+
+      attr_reader :base_client
+      attr_accessor :flush_count
+
+      def initialize(client, buffer_cap = nil)
+        @base_client = client
+        @buffer = String.new
+        @buffer_cap = buffer_cap || DEFAULT_BUFFER_CAP
+        @flush_count = 0
+      end
+
+      def flush
+        return unless @buffer.bytesize > 0
+        @base_client.send(@buffer)
+        @buffer.clear
+        @flush_count += 1
+      end
+
+      def send(msg)
+        flush if @buffer.bytesize + msg.bytesize >= @buffer_cap
+        @buffer << msg
+        @buffer << "\n".freeze
+        nil
+      end
+    end
+  end
+end