|
| 1 | +require 'socket' |
| 2 | +require 'zlib' |
| 3 | + |
| 4 | +module GitHub |
| 5 | + # = Statsd: A Statsd client (https://github.com/etsy/statsd) |
| 6 | + # |
| 7 | + # @example Set up a global Statsd client for a server on localhost:8125 |
| 8 | + # $statsd = Statsd.new 'localhost', 8125 |
| 9 | + # @example Send some stats |
| 10 | + # $statsd.increment 'garets' |
| 11 | + # $statsd.timing 'glork', 320 |
| 12 | + # @example Use {#time} to time the execution of a block |
| 13 | + # $statsd.time('account.activate') { @account.activate! } |
| 14 | + # @example Create a namespaced statsd client and increment 'account.activate' |
| 15 | + # statsd = Statsd.new('localhost').tap{|sd| sd.namespace = 'account'} |
| 16 | + # statsd.increment 'activate' |
| 17 | + class Statsd |
| 18 | + class UDPClient |
| 19 | + attr_reader :sock |
| 20 | + |
| 21 | + def initialize(address, port = nil) |
| 22 | + address, port = address.split(':') if address.include?(':') |
| 23 | + addrinfo = Addrinfo.ip(address) |
| 24 | + |
| 25 | + @sock = UDPSocket.new(addrinfo.pfamily) |
| 26 | + @sock.connect(addrinfo.ip_address, port) |
| 27 | + end |
| 28 | + |
| 29 | + def send(msg) |
| 30 | + sock.write(msg) |
| 31 | + rescue SystemCallError |
| 32 | + nil |
| 33 | + end |
| 34 | + end |
| 35 | + |
| 36 | + class SecureUDPClient < UDPClient |
| 37 | + def initialize(address, port, key) |
| 38 | + super(address, port) |
| 39 | + @key = key |
| 40 | + end |
| 41 | + |
| 42 | + def send(msg) |
| 43 | + super(signed_payload(msg)) |
| 44 | + end |
| 45 | + |
| 46 | + private |
| 47 | + # defer loading openssl and securerandom unless needed. this shaves ~10ms off |
| 48 | + # of baseline require load time for environments that don't require message signing. |
| 49 | + def self.setup_openssl |
| 50 | + @sha256 ||= begin |
| 51 | + require 'securerandom' |
| 52 | + require 'openssl' |
| 53 | + OpenSSL::Digest::SHA256.new |
| 54 | + end |
| 55 | + end |
| 56 | + |
| 57 | + def signed_payload(message) |
| 58 | + sha256 = SecureUDPClient.setup_openssl |
| 59 | + payload = timestamp + nonce + message |
| 60 | + signature = OpenSSL::HMAC.digest(sha256, @key, payload) |
| 61 | + signature + payload |
| 62 | + end |
| 63 | + |
| 64 | + def timestamp |
| 65 | + [Time.now.to_i].pack("Q<") |
| 66 | + end |
| 67 | + |
| 68 | + def nonce |
| 69 | + SecureRandom.random_bytes(4) |
| 70 | + end |
| 71 | + end |
| 72 | + |
| 73 | + # A namespace to prepend to all statsd calls. |
| 74 | + attr_reader :namespace |
| 75 | + |
| 76 | + def namespace=(namespace) |
| 77 | + @namespace = namespace |
| 78 | + @prefix = namespace ? "#{@namespace}." : "".freeze |
| 79 | + end |
| 80 | + |
| 81 | + # All the endpoints where StatsD will report metrics |
| 82 | + attr_reader :shards |
| 83 | + |
| 84 | + # The client class used to initialize shard instances and send metrics. |
| 85 | + attr_reader :client_class |
| 86 | + |
| 87 | + #characters that will be replaced with _ in stat names |
| 88 | + RESERVED_CHARS_REGEX = /[\:\|\@]/ |
| 89 | + |
| 90 | + COUNTER_TYPE = "c".freeze |
| 91 | + TIMING_TYPE = "ms".freeze |
| 92 | + GAUGE_TYPE = "g".freeze |
| 93 | + HISTOGRAM_TYPE = "h".freeze |
| 94 | + |
| 95 | + def initialize(client_class = nil) |
| 96 | + @shards = [] |
| 97 | + @client_class = client_class || UDPClient |
| 98 | + self.namespace = nil |
| 99 | + end |
| 100 | + |
| 101 | + def self.simple(addr, port = nil) |
| 102 | + self.new.add_shard(addr, port) |
| 103 | + end |
| 104 | + |
| 105 | + def add_shard(*args) |
| 106 | + @shards << @client_class.new(*args) |
| 107 | + self |
| 108 | + end |
| 109 | + |
| 110 | + def enable_buffering(buffer_size = nil) |
| 111 | + return if @buffering |
| 112 | + @shards.map! { |client| Buffer.new(client, buffer_size) } |
| 113 | + @buffering = true |
| 114 | + end |
| 115 | + |
| 116 | + def disable_buffering |
| 117 | + return unless @buffering |
| 118 | + flush_all |
| 119 | + @shards.map! { |client| client.base_client } |
| 120 | + @buffering = false |
| 121 | + end |
| 122 | + |
| 123 | + def flush_all |
| 124 | + return unless @buffering |
| 125 | + @shards.each { |client| client.flush } |
| 126 | + end |
| 127 | + |
| 128 | + |
| 129 | + # Sends an increment (count = 1) for the given stat to the statsd server. |
| 130 | + # |
| 131 | + # @param stat (see #count) |
| 132 | + # @param sample_rate (see #count) |
| 133 | + # @see #count |
| 134 | + def increment(stat, sample_rate=1); count stat, 1, sample_rate end |
| 135 | + |
| 136 | + # Sends a decrement (count = -1) for the given stat to the statsd server. |
| 137 | + # |
| 138 | + # @param stat (see #count) |
| 139 | + # @param sample_rate (see #count) |
| 140 | + # @see #count |
| 141 | + def decrement(stat, sample_rate=1); count stat, -1, sample_rate end |
| 142 | + |
| 143 | + # Sends an arbitrary count for the given stat to the statsd server. |
| 144 | + # |
| 145 | + # @param [String] stat stat name |
| 146 | + # @param [Integer] count count |
| 147 | + # @param [Integer] sample_rate sample rate, 1 for always |
| 148 | + def count(stat, count, sample_rate=1); send stat, count, COUNTER_TYPE, sample_rate end |
| 149 | + |
| 150 | + # Sends an arbitary gauge value for the given stat to the statsd server. |
| 151 | + # |
| 152 | + # @param [String] stat stat name. |
| 153 | + # @param [Numeric] gauge value. |
| 154 | + # @example Report the current user count: |
| 155 | + # $statsd.gauge('user.count', User.count) |
| 156 | + def gauge(stat, value) |
| 157 | + send stat, value, GAUGE_TYPE |
| 158 | + end |
| 159 | + |
| 160 | + # Sends a timing (in ms) for the given stat to the statsd server. The |
| 161 | + # sample_rate determines what percentage of the time this report is sent. The |
| 162 | + # statsd server then uses the sample_rate to correctly track the average |
| 163 | + # timing for the stat. |
| 164 | + # |
| 165 | + # @param stat stat name |
| 166 | + # @param [Integer] ms timing in milliseconds |
| 167 | + # @param [Integer] sample_rate sample rate, 1 for always |
| 168 | + def timing(stat, ms, sample_rate=1); send stat, ms, TIMING_TYPE, sample_rate end |
| 169 | + |
| 170 | + # Reports execution time of the provided block using {#timing}. |
| 171 | + # |
| 172 | + # @param stat (see #timing) |
| 173 | + # @param sample_rate (see #timing) |
| 174 | + # @yield The operation to be timed |
| 175 | + # @see #timing |
| 176 | + # @example Report the time (in ms) taken to activate an account |
| 177 | + # $statsd.time('account.activate') { @account.activate! } |
| 178 | + def time(stat, sample_rate=1) |
| 179 | + start = Time.now |
| 180 | + result = yield |
| 181 | + timing(stat, ((Time.now - start) * 1000).round(5), sample_rate) |
| 182 | + result |
| 183 | + end |
| 184 | + |
| 185 | + # Sends a histogram measurement for the given stat to the statsd server. The |
| 186 | + # sample_rate determines what percentage of the time this report is sent. The |
| 187 | + # statsd server then uses the sample_rate to correctly track the average |
| 188 | + # for the stat. |
| 189 | + def histogram(stat, value, sample_rate=1); send stat, value, HISTOGRAM_TYPE, sample_rate end |
| 190 | + |
| 191 | + private |
| 192 | + def sampled(sample_rate) |
| 193 | + yield unless sample_rate < 1 and rand > sample_rate |
| 194 | + end |
| 195 | + |
| 196 | + def send(stat, delta, type, sample_rate=1) |
| 197 | + sampled(sample_rate) do |
| 198 | + stat = stat.to_s.dup |
| 199 | + stat.gsub!(/::/, ".".freeze) |
| 200 | + stat.gsub!(RESERVED_CHARS_REGEX, "_".freeze) |
| 201 | + |
| 202 | + msg = String.new |
| 203 | + msg << @prefix |
| 204 | + msg << stat |
| 205 | + msg << ":".freeze |
| 206 | + msg << delta.to_s |
| 207 | + msg << "|".freeze |
| 208 | + msg << type |
| 209 | + if sample_rate < 1 |
| 210 | + msg << "|@".freeze |
| 211 | + msg << sample_rate.to_s |
| 212 | + end |
| 213 | + |
| 214 | + shard = select_shard(stat) |
| 215 | + shard.send(msg) |
| 216 | + end |
| 217 | + end |
| 218 | + |
| 219 | + def select_shard(stat) |
| 220 | + if @shards.size == 1 |
| 221 | + @shards.first |
| 222 | + else |
| 223 | + @shards[Zlib.crc32(stat) % @shards.size] |
| 224 | + end |
| 225 | + end |
| 226 | + |
| 227 | + class Buffer |
| 228 | + DEFAULT_BUFFER_CAP = 512 |
| 229 | + |
| 230 | + attr_reader :base_client |
| 231 | + attr_accessor :flush_count |
| 232 | + |
| 233 | + def initialize(client, buffer_cap = nil) |
| 234 | + @base_client = client |
| 235 | + @buffer = String.new |
| 236 | + @buffer_cap = buffer_cap || DEFAULT_BUFFER_CAP |
| 237 | + @flush_count = 0 |
| 238 | + end |
| 239 | + |
| 240 | + def flush |
| 241 | + return unless @buffer.bytesize > 0 |
| 242 | + @base_client.send(@buffer) |
| 243 | + @buffer.clear |
| 244 | + @flush_count += 1 |
| 245 | + end |
| 246 | + |
| 247 | + def send(msg) |
| 248 | + flush if @buffer.bytesize + msg.bytesize >= @buffer_cap |
| 249 | + @buffer << msg |
| 250 | + @buffer << "\n".freeze |
| 251 | + nil |
| 252 | + end |
| 253 | + end |
| 254 | + end |
| 255 | +end |
0 commit comments