dogstatsd-ruby 4.7.0 → 5.0.0

data/lib/datadog/statsd/connection.rb

@@ -3,30 +3,29 @@
 module Datadog
   class Statsd
     class Connection
-      def initialize(telemetry)
+      def initialize(telemetry: nil, logger: nil)
         @telemetry = telemetry
+        @logger = logger
       end
 
       # Close the underlying socket
       def close
-        @socket && @socket.close
+        begin
+          @socket && @socket.close if instance_variable_defined?(:@socket)
+        rescue StandardError => boom
+          logger.error { "Statsd: #{boom.class} #{boom}" } if logger
+        end
+        @socket = nil
       end
 
       def write(payload)
         logger.debug { "Statsd: #{payload}" } if logger
-        flush_telemetry = @telemetry.flush?
-        if flush_telemetry
-          payload += @telemetry.flush()
-        end
 
         send_message(payload)
 
-        if flush_telemetry
-          @telemetry.reset
-        end
+        telemetry.sent(packets: 1, bytes: payload.length) if telemetry
 
-        telemetry.bytes_sent += payload.length
-        telemetry.packets_sent += 1
+        true
       rescue StandardError => boom
         # Try once to reconnect if the socket has been closed
         retries ||= 1
@@ -36,21 +35,19 @@ module Datadog
            boom.is_a?(IOError) && boom.message =~ /closed stream/i)
           retries += 1
           begin
-            @socket = connect
+            close
             retry
           rescue StandardError => e
             boom = e
           end
         end
 
-        telemetry.bytes_dropped += payload.length
-        telemetry.packets_dropped += 1
+        telemetry.dropped(packets: 1, bytes: payload.length) if telemetry
         logger.error { "Statsd: #{boom.class} #{boom}" } if logger
         nil
       end
 
       private
-
       attr_reader :telemetry
       attr_reader :logger
 

data/lib/datadog/statsd/forwarder.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+module Datadog
+  class Statsd
+    class Forwarder
+      attr_reader :telemetry
+      attr_reader :transport_type
+
+      def initialize(
+        host: nil,
+        port: nil,
+        socket_path: nil,
+
+        buffer_max_payload_size: nil,
+        buffer_max_pool_size: nil,
+        buffer_overflowing_stategy: :drop,
+
+        telemetry_flush_interval: nil,
+        global_tags: [],
+
+        logger: nil
+      )
+        @transport_type = socket_path.nil? ? :udp : :uds
+
+        if telemetry_flush_interval
+          @telemetry = Telemetry.new(telemetry_flush_interval,
+            global_tags: global_tags,
+            transport_type: transport_type
+          )
+        end
+
+        @connection = case transport_type
+                      when :udp
+                        UDPConnection.new(host, port, logger: logger, telemetry: telemetry)
+                      when :uds
+                        UDSConnection.new(socket_path, logger: logger, telemetry: telemetry)
+                      end
+
+        # Initialize buffer
+        buffer_max_payload_size ||= (transport_type == :udp ? UDP_DEFAULT_BUFFER_SIZE : UDS_DEFAULT_BUFFER_SIZE)
+
+        if buffer_max_payload_size <= 0
+          raise ArgumentError, 'buffer_max_payload_size cannot be <= 0'
+        end
+
+        unless telemetry.nil? || telemetry.would_fit_in?(buffer_max_payload_size)
+          raise ArgumentError, "buffer_max_payload_size is not high enough to use telemetry (tags=(#{global_tags.inspect}))"
+        end
+
+        @buffer = MessageBuffer.new(@connection,
+          max_payload_size: buffer_max_payload_size,
+          max_pool_size: buffer_max_pool_size || DEFAULT_BUFFER_POOL_SIZE,
+          overflowing_stategy: buffer_overflowing_stategy,
+        )
+
+        @sender = Sender.new(buffer)
+        @sender.start
+      end
+
+      def send_message(message)
+        sender.add(message)
+
+        tick_telemetry
+      end
+
+      def sync_with_outbound_io
+        sender.rendez_vous
+      end
+
+      def flush(flush_telemetry: false, sync: false)
+        do_flush_telemetry if telemetry && flush_telemetry
+
+        sender.flush(sync: sync)
+      end
+
+      def host
+        return nil unless transport_type == :udp
+
+        connection.host
+      end
+
+      def port
+        return nil unless transport_type == :udp
+
+        connection.port
+      end
+
+      def socket_path
+        return nil unless transport_type == :uds
+
+        connection.socket_path
+      end
+
+      def close
+        sender.stop
+        connection.close
+      end
+
+      private
+      attr_reader :buffer
+      attr_reader :sender
+      attr_reader :connection
+
+      def do_flush_telemetry
+        telemetry_snapshot = telemetry.flush
+        telemetry.reset
+
+        telemetry_snapshot.each do |message|
+          sender.add(message)
+        end
+      end
+
+      def tick_telemetry
+        return nil unless telemetry
+
+        do_flush_telemetry if telemetry.should_flush?
+      end
+    end
+  end
+end

data/lib/datadog/statsd/message_buffer.rb

@@ -0,0 +1,88 @@
+# frozen_string_literal: true
+
+module Datadog
+  class Statsd
+    class MessageBuffer
+      PAYLOAD_SIZE_TOLERANCE = 0.05
+
+      def initialize(connection,
+        max_payload_size: nil,
+        max_pool_size: DEFAULT_BUFFER_POOL_SIZE,
+        overflowing_stategy: :drop
+      )
+        raise ArgumentError, 'max_payload_size keyword argument must be provided' unless max_payload_size
+        raise ArgumentError, 'max_pool_size keyword argument must be provided' unless max_pool_size
+
+        @connection = connection
+        @max_payload_size = max_payload_size
+        @max_pool_size = max_pool_size
+        @overflowing_stategy = overflowing_stategy
+
+        @buffer = String.new
+        @message_count = 0
+      end
+
+      def add(message)
+        message_size = message.bytesize
+
+        return nil unless message_size > 0 # to avoid adding empty messages to the buffer
+        return nil unless ensure_sendable!(message_size)
+
+        flush if should_flush?(message_size)
+
+        buffer << "\n" unless buffer.empty?
+        buffer << message
+
+        @message_count += 1
+
+        # flush when we're pretty sure that we won't be able
+        # to add another message to the buffer
+        flush if preemptive_flush?
+
+        true
+      end
+
+      def flush
+        return if buffer.empty?
+
+        connection.write(buffer)
+
+        buffer.clear
+        @message_count = 0
+      end
+
+      private
+      attr :max_payload_size
+      attr :max_pool_size
+
+      attr :overflowing_stategy
+
+      attr :connection
+      attr :buffer
+
+      def should_flush?(message_size)
+        return true if buffer.bytesize + 1 + message_size >= max_payload_size
+
+        false
+      end
+
+      def preemptive_flush?
+        @message_count == max_pool_size || buffer.bytesize > bytesize_threshold
+      end
+
+      def ensure_sendable!(message_size)
+        return true if message_size <= max_payload_size
+
+        if overflowing_stategy == :raise
+          raise Error, 'Message too big for payload limit'
+        end
+
+        false
+      end
+
+      def bytesize_threshold
+        @bytesize_threshold ||= (max_payload_size - PAYLOAD_SIZE_TOLERANCE * max_payload_size).to_i
+      end
+    end
+  end
+end

data/lib/datadog/statsd/sender.rb

@@ -0,0 +1,110 @@
+# frozen_string_literal: true
+
+module Datadog
+  class Statsd
+    class Sender
+      CLOSEABLE_QUEUES = Queue.instance_methods.include?(:close)
+
+      def initialize(message_buffer)
+        @message_buffer = message_buffer
+      end
+
+      def flush(sync: false)
+        raise ArgumentError, 'Start sender first' unless message_queue
+
+        message_queue.push(:flush)
+
+        rendez_vous if sync
+      end
+
+      def rendez_vous
+        # Initialize and get the thread's sync queue
+        queue = (Thread.current[:statsd_sync_queue] ||= Queue.new)
+        # tell sender-thread to notify us in the current
+        # thread's queue
+        message_queue.push(queue)
+        # wait for the sender thread to send a message
+        # once the flush is done
+        queue.pop
+      end
+
+      def add(message)
+        raise ArgumentError, 'Start sender first' unless message_queue
+
+        message_queue << message
+      end
+
+      def start
+        raise ArgumentError, 'Sender already started' if message_queue
+
+        # initialize message queue for background thread
+        @message_queue = Queue.new
+        # start background thread
+        @sender_thread = Thread.new(&method(:send_loop))
+      end
+
+      if CLOSEABLE_QUEUES
+        def stop(join_worker: true)
+          message_queue.close if message_queue
+          sender_thread.join if sender_thread && join_worker
+        end
+      else
+        def stop(join_worker: true)
+          message_queue << :close if message_queue
+          sender_thread.join if sender_thread && join_worker
+        end
+      end
+
+      private
+
+      attr_reader :message_buffer
+
+      attr_reader :message_queue
+      attr_reader :sender_thread
+
+      if CLOSEABLE_QUEUES
+        def send_loop
+          until (message = message_queue.pop).nil? && message_queue.closed?
+            # skip if message is nil, e.g. when message_queue
+            # is empty and closed
+            next unless message
+
+            case message
+            when :flush
+              message_buffer.flush
+            when Queue
+              message.push(:go_on)
+            else
+              message_buffer.add(message)
+            end
+          end
+
+          @message_queue = nil
+          @sender_thread = nil
+        end
+      else
+        def send_loop
+          loop do
+            message = message_queue.pop
+
+            next unless message
+
+            case message
+            when :close
+              break
+            when :flush
+              message_buffer.flush
+            when Queue
+              message.push(:go_on)
+            else
+              message_buffer.add(message)
+            end
+          end
+
+          @message_queue = nil
+          @sender_thread = nil
+        end
+      end
+    end
+  end
+end

data/lib/datadog/statsd/serialization.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+module Datadog
+  class Statsd
+    module Serialization
+    end
+  end
+end
+
+require_relative 'serialization/tag_serializer'
+require_relative 'serialization/service_check_serializer'
+require_relative 'serialization/event_serializer'
+require_relative 'serialization/stat_serializer'
+
+require_relative 'serialization/serializer'

data/lib/datadog/statsd/serialization/event_serializer.rb

@@ -0,0 +1,71 @@
+# frozen_string_literal: true
+
+module Datadog
+  class Statsd
+    module Serialization
+      class EventSerializer
+        EVENT_BASIC_OPTIONS = {
+          date_happened:    'd:',
+          hostname:         'h:',
+          aggregation_key:  'k:',
+          priority:         'p:',
+          source_type_name: 's:',
+          alert_type:       't:',
+        }.freeze
+
+        def initialize(global_tags: [])
+          @tag_serializer = TagSerializer.new(global_tags)
+        end
+
+        def format(title, text, options = EMPTY_OPTIONS)
+          title = escape(title)
+          text = escape(text)
+
+          String.new.tap do |event|
+            event << '_e{'
+            event << title.bytesize.to_s
+            event << ','
+            event << text.bytesize.to_s
+            event << '}:'
+            event << title
+            event << '|'
+            event << text
+
+            # we are serializing the generic service check options
+            # before serializing specialized options that need edge-cases
+            EVENT_BASIC_OPTIONS.each do |option_key, shortcut|
+              if value = options[option_key]
+                event << '|'
+                event << shortcut
+                event << value.to_s.delete('|')
+              end
+            end
+
+            # also returns the global tags from serializer
+            if tags = tag_serializer.format(options[:tags])
+              event << '|#'
+              event << tags
+            end
+
+            if event.bytesize > MAX_EVENT_SIZE
+              if options[:truncate_if_too_long]
+                event.slice!(MAX_EVENT_SIZE..event.length)
+              else
+                raise "Event #{title} payload is too big (more that 8KB), event discarded"
+              end
+            end
+          end
+        end
+
+        protected
+        attr_reader :tag_serializer
+
+        def escape(text)
+          text.delete('|').tap do |t|
+            t.gsub!("\n", '\n')
+          end
+        end
+      end
+    end
+  end
+end