dogstatsd-ruby 3.3.0 → 5.6.1

data/lib/datadog/statsd/message_buffer.rb

@@ -0,0 +1,105 @@
+# 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,
+        serializer:
+      )
+        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
+        @serializer = serializer
+
+        @buffer = String.new
+        clear_buffer
+      end
+
+      def add(message)
+        # Serializes the message if it hasn't been already. Part of the
+        # delay_serialization feature.
+        if message.is_a?(Array)
+          message = @serializer.to_stat(*message[0], **message[1])
+        end
+
+        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 reset
+        clear_buffer
+        connection.reset_telemetry
+      end
+
+      def flush
+        return if buffer.empty?
+
+        connection.write(buffer)
+        clear_buffer
+      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 clear_buffer
+        buffer.clear
+        @message_count = 0
+      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,184 @@
+# frozen_string_literal: true
+
+module Datadog
+  class Statsd
+    # Sender is using a companion thread to flush and pack messages
+    # in a `MessageBuffer`.
+    # The communication with this thread is done using a `Queue`.
+    # If the thread is dead, it is starting a new one to avoid having a blocked
+    # Sender with no companion thread to communicate with (most of the time, having
+    # a dead companion thread means that a fork just happened and that we are
+    # running in the child process).
+    class Sender
+      CLOSEABLE_QUEUES = Queue.instance_methods.include?(:close)
+
+      def initialize(message_buffer, telemetry: nil, queue_size: UDP_DEFAULT_BUFFER_SIZE, logger: nil, flush_interval: nil, queue_class: Queue, thread_class: Thread)
+        @message_buffer = message_buffer
+        @telemetry = telemetry
+        @queue_size = queue_size
+        @logger = logger
+        @mx = Mutex.new
+        @queue_class = queue_class
+        @thread_class = thread_class
+        @flush_timer = if flush_interval
+          Datadog::Statsd::Timer.new(flush_interval) { flush(sync: true) }
+        else
+          nil
+        end
+      end
+
+      def flush(sync: false)
+        # keep a copy around in case another thread is calling #stop while this method is running
+        current_message_queue = message_queue
+
+        # don't try to flush if there is no message_queue instantiated or
+        # no companion thread running
+        if !current_message_queue
+          @logger.debug { "Statsd: can't flush: no message queue ready" } if @logger
+          return
+        end
+        if !sender_thread.alive?
+          @logger.debug { "Statsd: can't flush: no sender_thread alive" } if @logger
+          return
+        end
+
+        current_message_queue.push(:flush)
+        rendez_vous if sync
+      end
+
+      def rendez_vous
+        # could happen if #start hasn't be called
+        return unless message_queue
+
+        # Initialize and get the thread's sync queue
+        queue = (@thread_class.current[:statsd_sync_queue] ||= @queue_class.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
+
+        # if the thread does not exist, we assume we are running in a forked process,
+        # empty the message queue and message buffers (these messages belong to
+        # the parent process) and spawn a new companion thread.
+        if !sender_thread.alive?
+          @mx.synchronize {
+            # a call from another thread has already re-created
+            # the companion thread before this one acquired the lock
+            break if sender_thread.alive?
+            @logger.debug { "Statsd: companion thread is dead, re-creating one" } if @logger
+
+            message_queue.close if CLOSEABLE_QUEUES
+            @message_queue = nil
+            message_buffer.reset
+            start
+            @flush_timer.start if @flush_timer && @flush_timer.stop?
+          }
+        end
+
+        if message_queue.length <= @queue_size
+          message_queue << message
+        else
+          if @telemetry
+            bytesize = message.respond_to?(:bytesize) ? message.bytesize : 0
+            @telemetry.dropped_queue(packets: 1, bytes: bytesize)
+          end
+        end
+      end
+
+      def start
+        raise ArgumentError, 'Sender already started' if message_queue
+
+        # initialize a new message queue for the background thread
+        @message_queue = @queue_class.new
+        # start background thread
+        @sender_thread = @thread_class.new(&method(:send_loop))
+        @sender_thread.name = "Statsd Sender" unless Gem::Version.new(RUBY_VERSION) < Gem::Version.new('2.3')
+        @flush_timer.start if @flush_timer
+      end
+
+      if CLOSEABLE_QUEUES
+        # when calling stop, make sure that no other threads is trying
+        # to close the sender nor trying to continue to `#add` more message
+        # into the sender.
+        def stop(join_worker: true)
+          @flush_timer.stop if @flush_timer
+
+          message_queue = @message_queue
+          message_queue.close if message_queue
+
+          sender_thread = @sender_thread
+          sender_thread.join if sender_thread && join_worker
+        end
+      else
+        # when calling stop, make sure that no other threads is trying
+        # to close the sender nor trying to continue to `#add` more message
+        # into the sender.
+        def stop(join_worker: true)
+          @flush_timer.stop if @flush_timer
+
+          message_queue = @message_queue
+          message_queue << :close if message_queue
+
+          sender_thread = @sender_thread
+          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_class
+              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_class
+              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/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

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

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+# require 'forwardable'
+
+module Datadog
+  class Statsd
+    module Serialization
+      class Serializer
+        def initialize(prefix: nil, global_tags: [])
+          @stat_serializer = StatSerializer.new(prefix, global_tags: global_tags)
+          @service_check_serializer = ServiceCheckSerializer.new(global_tags: global_tags)
+          @event_serializer = EventSerializer.new(global_tags: global_tags)
+        end
+
+        # using *args would make new allocations
+        def to_stat(name, delta, type, tags: [], sample_rate: 1)
+          stat_serializer.format(name, delta, type, tags: tags, sample_rate: sample_rate)
+        end
+
+        # using *args would make new allocations
+        def to_service_check(name, status, options = EMPTY_OPTIONS)
+          service_check_serializer.format(name, status, options)
+        end
+
+        # using *args would make new allocations
+        def to_event(title, text, options = EMPTY_OPTIONS)
+          event_serializer.format(title, text, options)
+        end
+
+        def global_tags
+          stat_serializer.global_tags
+        end
+
+        protected
+        attr_reader :stat_serializer
+        attr_reader :service_check_serializer
+        attr_reader :event_serializer
+      end
+    end
+  end
+end

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

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+module Datadog
+  class Statsd
+    module Serialization
+      class ServiceCheckSerializer
+        SERVICE_CHECK_BASIC_OPTIONS = {
+          timestamp: 'd:',
+          hostname:  'h:',
+        }.freeze
+
+        def initialize(global_tags: [])
+          @tag_serializer = TagSerializer.new(global_tags)
+        end
+
+        def format(name, status, options = EMPTY_OPTIONS)
+          String.new.tap do |service_check|
+            # line basics
+            service_check << "_sc"
+            service_check << "|"
+            service_check << name.to_s
+            service_check << "|"
+            service_check << status.to_s
+
+            # we are serializing the generic service check options
+            # before serializing specialized options that need edge-cases
+            SERVICE_CHECK_BASIC_OPTIONS.each do |option_key, shortcut|
+              if value = options[option_key]
+                service_check << '|'
+                service_check << shortcut
+                service_check << value.to_s.delete('|')
+              end
+            end
+
+            if message = options[:message]
+              service_check << '|m:'
+              service_check << escape_message(message)
+            end
+
+            # also returns the global tags from serializer
+            if tags = tag_serializer.format(options[:tags])
+              service_check << '|#'
+              service_check << tags
+            end
+          end
+        end
+
+        protected
+        attr_reader :tag_serializer
+
+        def escape_message(message)
+          message.delete('|').tap do |m|
+            m.gsub!("\n", '\n')
+            m.gsub!('m:', 'm\:')
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module Datadog
+  class Statsd
+    module Serialization
+      class StatSerializer
+        def initialize(prefix, global_tags: [])
+          @prefix = prefix
+          @prefix_str = prefix.to_s
+          @tag_serializer = TagSerializer.new(global_tags)
+        end
+
+        def format(name, delta, type, tags: [], sample_rate: 1)
+          name = formated_name(name)
+
+          if sample_rate != 1
+            if tags_list = tag_serializer.format(tags)
+              "#{@prefix_str}#{name}:#{delta}|#{type}|@#{sample_rate}|##{tags_list}"
+            else
+              "#{@prefix_str}#{name}:#{delta}|#{type}|@#{sample_rate}"
+            end
+          else
+            if tags_list = tag_serializer.format(tags)
+              "#{@prefix_str}#{name}:#{delta}|#{type}|##{tags_list}"
+            else
+              "#{@prefix_str}#{name}:#{delta}|#{type}"
+            end
+          end
+        end
+
+        def global_tags
+          tag_serializer.global_tags
+        end
+
+        private
+
+        attr_reader :prefix
+        attr_reader :tag_serializer
+
+        def formated_name(name)
+          if name.is_a?(String)
+            # DEV: gsub is faster than dup.gsub!
+            formated = name.gsub('::', '.')
+          else
+            formated = name.to_s
+            formated.gsub!('::', '.')
+          end
+
+          formated.tr!(':|@', '_')
+          formated
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,96 @@
+# frozen_string_literal: true
+
+module Datadog
+  class Statsd
+    module Serialization
+      class TagSerializer
+        def initialize(global_tags = [], env = ENV)
+          # Convert to hash
+          global_tags = to_tags_hash(global_tags)
+
+          # Merge with default tags
+          global_tags = default_tags(env).merge(global_tags)
+
+          # Convert to tag list and set
+          @global_tags = to_tags_list(global_tags)
+          if @global_tags.any?
+            @global_tags_formatted = @global_tags.join(',')
+          else
+            @global_tags_formatted = nil
+          end
+        end
+
+        def format(message_tags)
+          if !message_tags || message_tags.empty?
+            return @global_tags_formatted
+          end
+
+          tags = if @global_tags_formatted
+                   [@global_tags_formatted, to_tags_list(message_tags)]
+                 else
+                   to_tags_list(message_tags)
+                 end
+
+          tags.join(',')
+        end
+
+        attr_reader :global_tags
+
+        private
+
+        def to_tags_hash(tags)
+          case tags
+          when Hash
+            tags.dup
+          when Array
+            Hash[
+              tags.map do |string|
+                tokens = string.split(':')
+                tokens << nil if tokens.length == 1
+                tokens.length == 2 ? tokens : nil
+              end.compact
+            ]
+          else
+            {}
+          end
+        end
+
+        def to_tags_list(tags)
+          case tags
+          when Hash
+            tags.map do |name, value|
+              if value
+                escape_tag_content("#{name}:#{value}")
+              else
+                escape_tag_content(name)
+              end
+            end
+          when Array
+            tags.map { |tag| escape_tag_content(tag) }
+          else
+            []
+          end
+        end
+
+        def escape_tag_content(tag)
+          tag.to_s.delete('|,')
+        end
+
+        def dd_tags(env = ENV)
+          return {} unless dd_tags = env['DD_TAGS']
+
+          to_tags_hash(dd_tags.split(','))
+        end
+
+        def default_tags(env = ENV)
+          dd_tags(env).tap do |tags|
+            tags['dd.internal.entity_id'] = env['DD_ENTITY_ID'] if env.key?('DD_ENTITY_ID')
+            tags['env'] = env['DD_ENV'] if env.key?('DD_ENV')
+            tags['service'] = env['DD_SERVICE'] if env.key?('DD_SERVICE')
+            tags['version'] = env['DD_VERSION'] if env.key?('DD_VERSION')
+          end
+        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'