waterdrop 2.0.7 → 2.6.11

data/lib/waterdrop/instrumentation/vendors/datadog/metrics_listener.rb

@@ -0,0 +1,210 @@
+# frozen_string_literal: true
+
+module WaterDrop
+  module Instrumentation
+    # Namespace for vendor specific instrumentation
+    module Vendors
+      # Datadog specific instrumentation
+      module Datadog
+        # Listener that can be used to subscribe to WaterDrop producer to receive stats via StatsD
+        # and/or Datadog
+        #
+        # @note You need to setup the `dogstatsd-ruby` client and assign it
+        class MetricsListener
+          include ::Karafka::Core::Configurable
+          extend Forwardable
+
+          def_delegators :config, :client, :rd_kafka_metrics, :namespace, :default_tags
+
+          # Value object for storing a single rdkafka metric publishing details
+          RdKafkaMetric = Struct.new(:type, :scope, :name, :key_location)
+
+          # Namespace under which the DD metrics should be published
+          setting :namespace, default: 'waterdrop'
+
+          # Datadog client that we should use to publish the metrics
+          setting :client
+
+          # Default tags we want to publish (for example hostname)
+          # Format as followed (example for hostname): `["host:#{Socket.gethostname}"]`
+          setting :default_tags, default: []
+
+          # All the rdkafka metrics we want to publish
+          #
+          # By default we publish quite a lot so this can be tuned
+          # Note, that the once with `_d` come from WaterDrop, not rdkafka or Kafka
+          setting :rd_kafka_metrics, default: [
+            # Client metrics
+            RdKafkaMetric.new(:count, :root, 'calls', 'tx_d'),
+            RdKafkaMetric.new(:histogram, :root, 'queue.size', 'msg_cnt_d'),
+
+            # Broker metrics
+            RdKafkaMetric.new(:count, :brokers, 'deliver.attempts', 'txretries_d'),
+            RdKafkaMetric.new(:count, :brokers, 'deliver.errors', 'txerrs_d'),
+            RdKafkaMetric.new(:count, :brokers, 'receive.errors', 'rxerrs_d'),
+            RdKafkaMetric.new(:gauge, :brokers, 'queue.latency.avg', %w[outbuf_latency avg]),
+            RdKafkaMetric.new(:gauge, :brokers, 'queue.latency.p95', %w[outbuf_latency p95]),
+            RdKafkaMetric.new(:gauge, :brokers, 'queue.latency.p99', %w[outbuf_latency p99]),
+            RdKafkaMetric.new(:gauge, :brokers, 'network.latency.avg', %w[rtt avg]),
+            RdKafkaMetric.new(:gauge, :brokers, 'network.latency.p95', %w[rtt p95]),
+            RdKafkaMetric.new(:gauge, :brokers, 'network.latency.p99', %w[rtt p99])
+          ].freeze
+
+          configure
+
+          # @param block [Proc] configuration block
+          def initialize(&block)
+            configure
+            setup(&block) if block
+          end
+
+          # @param block [Proc] configuration block
+          # @note We define this alias to be consistent with `WaterDrop#setup`
+          def setup(&block)
+            configure(&block)
+          end
+
+          # Hooks up to WaterDrop instrumentation for emitted statistics
+          #
+          # @param event [Karafka::Core::Monitoring::Event]
+          def on_statistics_emitted(event)
+            statistics = event[:statistics]
+
+            rd_kafka_metrics.each do |metric|
+              report_metric(metric, statistics)
+            end
+          end
+
+          # Increases the errors count by 1
+          #
+          # @param _event [Karafka::Core::Monitoring::Event]
+          def on_error_occurred(_event)
+            count('error_occurred', 1, tags: default_tags)
+          end
+
+          # Increases acknowledged messages counter
+          # @param _event [Karafka::Core::Monitoring::Event]
+          def on_message_acknowledged(_event)
+            increment('acknowledged', tags: default_tags)
+          end
+
+          %i[
+            produced_sync
+            produced_async
+          ].each do |event_scope|
+            class_eval <<~METHODS, __FILE__, __LINE__ + 1
+              # @param event [Karafka::Core::Monitoring::Event]
+              def on_message_#{event_scope}(event)
+                report_message(event[:message][:topic], :#{event_scope})
+              end
+
+              # @param event [Karafka::Core::Monitoring::Event]
+              def on_messages_#{event_scope}(event)
+                event[:messages].each do |message|
+                  report_message(message[:topic], :#{event_scope})
+                end
+              end
+            METHODS
+          end
+
+          # Reports the buffer usage when anything is added to the buffer
+          %i[
+            message_buffered
+            messages_buffered
+          ].each do |event_scope|
+            class_eval <<~METHODS, __FILE__, __LINE__ + 1
+              # @param event [Karafka::Core::Monitoring::Event]
+              def on_#{event_scope}(event)
+                histogram(
+                  'buffer.size',
+                  event[:buffer].size,
+                  tags: default_tags
+                )
+              end
+            METHODS
+          end
+
+          # Events that support many messages only
+          # Reports data flushing operation (production from the buffer)
+          %i[
+            flushed_sync
+            flushed_async
+          ].each do |event_scope|
+            class_eval <<~METHODS, __FILE__, __LINE__ + 1
+              # @param event [Karafka::Core::Monitoring::Event]
+              def on_buffer_#{event_scope}(event)
+                event[:messages].each do |message|
+                  report_message(message[:topic], :#{event_scope})
+                end
+              end
+            METHODS
+          end
+
+          private
+
+          %i[
+            count
+            gauge
+            histogram
+            increment
+            decrement
+          ].each do |metric_type|
+            class_eval <<~METHODS, __FILE__, __LINE__ + 1
+              def #{metric_type}(key, *args)
+                client.#{metric_type}(
+                  namespaced_metric(key),
+                  *args
+                )
+              end
+            METHODS
+          end
+
+          # Report that a message has been produced to a topic.
+          # @param topic [String] Kafka topic
+          # @param method_name [Symbol] method from which this message operation comes
+          def report_message(topic, method_name)
+            increment(method_name, tags: default_tags + ["topic:#{topic}"])
+          end
+
+          # Wraps metric name in listener's namespace
+          # @param metric_name [String] RdKafkaMetric name
+          # @return [String]
+          def namespaced_metric(metric_name)
+            "#{namespace}.#{metric_name}"
+          end
+
+          # Reports a given metric statistics to Datadog
+          # @param metric [RdKafkaMetric] metric value object
+          # @param statistics [Hash] hash with all the statistics emitted
+          def report_metric(metric, statistics)
+            case metric.scope
+            when :root
+              public_send(
+                metric.type,
+                metric.name,
+                statistics.fetch(*metric.key_location),
+                tags: default_tags
+              )
+            when :brokers
+              statistics.fetch('brokers').each_value do |broker_statistics|
+                # Skip bootstrap nodes
+                # Bootstrap nodes have nodeid -1, other nodes have positive
+                # node ids
+                next if broker_statistics['nodeid'] == -1
+
+                public_send(
+                  metric.type,
+                  metric.name,
+                  broker_statistics.dig(*metric.key_location),
+                  tags: default_tags + ["broker:#{broker_statistics['nodename']}"]
+                )
+              end
+            else
+              raise ArgumentError, metric.scope
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/waterdrop/middleware.rb

@@ -0,0 +1,50 @@
+# frozen_string_literal: true
+
+module WaterDrop
+  # Simple middleware layer for manipulating messages prior to their validation
+  class Middleware
+    def initialize
+      @mutex = Mutex.new
+      @steps = []
+    end
+
+    # Runs middleware on a single message prior to validation
+    #
+    # @param message [Hash] message hash
+    # @return [Hash] message hash. Either the same if transformed in place, or a copy if modified
+    #   into a new object.
+    # @note You need to decide yourself whether you don't use the message hash data anywhere else
+    #   and you want to save on memory by modifying it in place or do you want to do a deep copy
+    def run(message)
+      @steps.each do |step|
+        message = step.call(message)
+      end
+
+      message
+    end
+
+    # @param messages [Array<Hash>] messages on which we want to run middlewares
+    # @return [Array<Hash>] transformed messages
+    def run_many(messages)
+      messages.map do |message|
+        run(message)
+      end
+    end
+
+    # Register given middleware as the first one in the chain
+    # @param step [#call] step that needs to return the message
+    def prepend(step)
+      @mutex.synchronize do
+        @steps.prepend step
+      end
+    end
+
+    # Register given middleware as the last one in the chain
+    # @param step [#call] step that needs to return the message
+    def append(step)
+      @mutex.synchronize do
+        @steps.append step
+      end
+    end
+  end
+end

data/lib/{water_drop → waterdrop}/producer/async.rb

@@ -14,14 +14,30 @@ module WaterDrop
       # @raise [Errors::MessageInvalidError] When provided message details are invalid and the
       #   message could not be sent to Kafka
       def produce_async(message)
-        ensure_active!
+        message = middleware.run(message)
         validate_message!(message)
 
         @monitor.instrument(
           'message.produced_async',
           producer_id: id,
           message: message
-        ) { client.produce(**message) }
+        ) { produce(message) }
+      rescue *SUPPORTED_FLOW_ERRORS => e
+        # We use this syntax here because we want to preserve the original `#cause` when we
+        # instrument the error and there is no way to manually assign `#cause` value
+        begin
+          raise Errors::ProduceError, e.inspect
+        rescue Errors::ProduceError => ex
+          @monitor.instrument(
+            'error.occurred',
+            producer_id: id,
+            message: message,
+            error: ex,
+            type: 'message.produce_async'
+          )
+
+          raise ex
+        end
       end
 
       # Produces many messages to Kafka and does not wait for them to be delivered
@@ -35,7 +51,8 @@ module WaterDrop
       # @raise [Errors::MessageInvalidError] When any of the provided messages details are invalid
       #   and the message could not be sent to Kafka
       def produce_many_async(messages)
-        ensure_active!
+        dispatched = []
+        messages = middleware.run_many(messages)
         messages.each { |message| validate_message!(message) }
 
         @monitor.instrument(
@@ -43,8 +60,27 @@ module WaterDrop
           producer_id: id,
           messages: messages
         ) do
-          messages.map { |message| client.produce(**message) }
+          with_transaction_if_transactional do
+            messages.each do |message|
+              dispatched << produce(message)
+            end
+          end
+
+          dispatched
         end
+      rescue *SUPPORTED_FLOW_ERRORS => e
+        re_raised = Errors::ProduceManyError.new(dispatched, e.inspect)
+
+        @monitor.instrument(
+          'error.occurred',
+          producer_id: id,
+          messages: messages,
+          dispatched: dispatched,
+          error: re_raised,
+          type: 'messages.produce_many_async'
+        )
+
+        raise re_raised
       end
     end
   end

data/lib/{water_drop → waterdrop}/producer/buffer.rb

@@ -4,14 +4,6 @@ module WaterDrop
   class Producer
     # Component for buffered operations
     module Buffer
-      # Exceptions we catch when dispatching messages from a buffer
-      RESCUED_ERRORS = [
-        Rdkafka::RdkafkaError,
-        Rdkafka::Producer::DeliveryHandle::WaitTimeoutError
-      ].freeze
-
-      private_constant :RESCUED_ERRORS
-
       # Adds given message into the internal producer buffer without flushing it to Kafka
       #
       # @param message [Hash] hash that complies with the {Contracts::Message} contract
@@ -19,12 +11,15 @@ module WaterDrop
       #   message could not be sent to Kafka
       def buffer(message)
         ensure_active!
+
+        message = middleware.run(message)
         validate_message!(message)
 
         @monitor.instrument(
           'message.buffered',
           producer_id: id,
-          message: message
+          message: message,
+          buffer: @messages
         ) { @messages << message }
       end
 
@@ -36,12 +31,15 @@ module WaterDrop
       #   and the message could not be sent to Kafka
       def buffer_many(messages)
         ensure_active!
+
+        messages = middleware.run_many(messages)
         messages.each { |message| validate_message!(message) }
 
         @monitor.instrument(
           'messages.buffered',
           producer_id: id,
-          messages: messages
+          messages: messages,
+          buffer: @messages
         ) do
           messages.each { |message| @messages << message }
           messages
@@ -52,8 +50,6 @@ module WaterDrop
       # @return [Array<Rdkafka::Producer::DeliveryHandle>] delivery handles for messages that were
       #   flushed
       def flush_async
-        ensure_active!
-
         @monitor.instrument(
           'buffer.flushed_async',
           producer_id: id,
@@ -65,8 +61,6 @@ module WaterDrop
       # @return [Array<Rdkafka::Producer::DeliveryReport>] delivery reports for messages that were
       #   flushed
       def flush_sync
-        ensure_active!
-
         @monitor.instrument(
           'buffer.flushed_sync',
           producer_id: id,
@@ -80,33 +74,21 @@ module WaterDrop
       # @param sync [Boolean] should it flush in a sync way
       # @return [Array<Rdkafka::Producer::DeliveryHandle, Rdkafka::Producer::DeliveryReport>]
       #   delivery handles for async or delivery reports for sync
-      # @raise [Errors::FlushFailureError] when there was a failure in flushing
+      # @raise [Errors::ProduceManyError] when there was a failure in flushing
       # @note We use this method underneath to provide a different instrumentation for sync and
       #   async flushing within the public API
       def flush(sync)
         data_for_dispatch = nil
-        dispatched = []
 
         @buffer_mutex.synchronize do
           data_for_dispatch = @messages
           @messages = Concurrent::Array.new
         end
 
-        dispatched = data_for_dispatch.map { |message| client.produce(**message) }
-
-        return dispatched unless sync
-
-        dispatched.map do |handler|
-          handler.wait(
-            max_wait_timeout: @config.max_wait_timeout,
-            wait_timeout: @config.wait_timeout
-          )
-        end
-      rescue *RESCUED_ERRORS => e
-        key = sync ? 'buffer.flushed_sync.error' : 'buffer.flush_async.error'
-        @monitor.instrument(key, producer_id: id, error: e, dispatched: dispatched)
+        # Do nothing if nothing to flush
+        return data_for_dispatch if data_for_dispatch.empty?
 
-        raise Errors::FlushFailureError.new(dispatched)
+        sync ? produce_many_sync(data_for_dispatch) : produce_many_async(data_for_dispatch)
       end
     end
   end

data/lib/{water_drop → waterdrop}/producer/builder.rb

@@ -10,18 +10,13 @@ module WaterDrop
       # @return [Rdkafka::Producer, Producer::DummyClient] raw rdkafka producer or a dummy producer
       #   when we don't want to dispatch any messages
       def call(producer, config)
-        return DummyClient.new unless config.deliver
+        klass = config.client_class
+        # This allows us to have backwards compatibility.
+        # If it is the default client and delivery is set to false, we use dummy as we used to
+        # before `client_class` was introduced
+        klass = Clients::Dummy if klass == Clients::Rdkafka && !config.deliver
 
-        client = Rdkafka::Config.new(config.kafka.to_h).producer
-
-        # This callback is not global and is per client, thus we do not have to wrap it with a
-        # callbacks manager to make it work
-        client.delivery_callback = Instrumentation::Callbacks::Delivery.new(
-          producer.id,
-          config.monitor
-        )
-
-        client
+        klass.new(producer)
       end
     end
   end

data/lib/{water_drop → waterdrop}/producer/sync.rb

@@ -16,7 +16,7 @@ module WaterDrop
       # @raise [Errors::MessageInvalidError] When provided message details are invalid and the
       #   message could not be sent to Kafka
       def produce_sync(message)
-        ensure_active!
+        message = middleware.run(message)
         validate_message!(message)
 
         @monitor.instrument(
@@ -24,12 +24,23 @@ module WaterDrop
           producer_id: id,
           message: message
         ) do
-          client
-            .produce(**message)
-            .wait(
-              max_wait_timeout: @config.max_wait_timeout,
-              wait_timeout: @config.wait_timeout
-            )
+          wait(produce(message))
+        end
+      rescue *SUPPORTED_FLOW_ERRORS => e
+        # We use this syntax here because we want to preserve the original `#cause` when we
+        # instrument the error and there is no way to manually assign `#cause` value
+        begin
+          raise Errors::ProduceError, e.inspect
+        rescue Errors::ProduceError => ex
+          @monitor.instrument(
+            'error.occurred',
+            producer_id: id,
+            message: message,
+            error: ex,
+            type: 'message.produce_sync'
+          )
+
+          raise ex
         end
       end
 
@@ -46,19 +57,37 @@ module WaterDrop
       # @raise [Errors::MessageInvalidError] When any of the provided messages details are invalid
       #   and the message could not be sent to Kafka
       def produce_many_sync(messages)
-        ensure_active!
+        messages = middleware.run_many(messages)
         messages.each { |message| validate_message!(message) }
 
+        dispatched = []
+
         @monitor.instrument('messages.produced_sync', producer_id: id, messages: messages) do
-          messages
-            .map { |message| client.produce(**message) }
-            .map! do |handler|
-              handler.wait(
-                max_wait_timeout: @config.max_wait_timeout,
-                wait_timeout: @config.wait_timeout
-              )
+          with_transaction_if_transactional do
+            messages.each do |message|
+              dispatched << produce(message)
             end
+          end
+
+          dispatched.map! do |handler|
+            wait(handler)
+          end
+
+          dispatched
         end
+      rescue *SUPPORTED_FLOW_ERRORS => e
+        re_raised = Errors::ProduceManyError.new(dispatched, e.inspect)
+
+        @monitor.instrument(
+          'error.occurred',
+          producer_id: id,
+          messages: messages,
+          dispatched: dispatched,
+          error: re_raised,
+          type: 'messages.produce_many_sync'
+        )
+
+        raise re_raised
       end
     end
   end