statsd: Allow buffering for the clients

vmg · vmg · commit dba7e7ac0261 · 2016-01-13T09:50:56.000+01:00
The new `Buffer` wrapper allow wrapping any UDP clients in a "buffering
mode" that batches incoming metrics into a buffer and eventually flushes
them to the statsd-compatible server.

The receiving server *must* support the multi-metric StatsD protocol
(namely, parsing several metrics in the same packet, separated by
newlines). Both the original StatsD and now Brubeck support this. By
sending several metrics in a single packet, we reduce the amount of
network operations (both in the client and the server) and the amount of
context switches to the kernel UDP stack, significantly increasing the
throughput of the server implementation.

The default maximum packet size has been set at 512 bytes, which is a
"safe size" to reliably transmit an atomic packet through the internet:
The MTU on the internet is 576, and the size of the IPv4 header is 20
bytes, and the UDP header 8 bytes.  This leaves 548 bytes available for
user data. We cap it at 512 to add some headspace (most other UDP-based
protocols, such as DNS, do the same thing).
diff --git a/lib/statsd.rb b/lib/statsd.rb
@@ -14,6 +14,24 @@
 #   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)
@@ -51,24 +69,6 @@ def nonce
     end
   end
 
-  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 => boom
-      nil
-    end
-  end
-
   # A namespace to prepend to all statsd calls.
   attr_reader :namespace
 
@@ -103,6 +103,25 @@ def add_shard(*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)
@@ -200,4 +219,29 @@ def select_shard(stat)
       @shards[Zlib.crc32(stat) % @shards.size]
     end
   end
+
+  class Buffer
+    DEFAULT_BUFFER_CAP = 512
+
+    attr_reader :base_client
+
+    def initialize(client, buffer_cap = nil)
+      @base_client = client
+      @buffer = String.new
+      @buffer_cap = buffer_cap || DEFAULT_BUFFER_CAP
+    end
+
+    def flush
+      return unless @buffer.bytesize > 0
+      @base_client.send(@buffer)
+      @buffer.clear
+    end
+
+    def send(msg)
+      flush if @buffer.bytesize + msg.bytesize >= @buffer_cap
+      @buffer << msg
+      @buffer << "\n".freeze
+      nil
+    end
+  end
 end
