waterdrop 2.8.15 → 2.9.0

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

@@ -20,7 +20,7 @@ module WaterDrop
           RdKafkaMetric = Struct.new(:type, :scope, :name, :key_location)
 
           # Namespace under which the DD metrics should be published
-          setting :namespace, default: 'waterdrop'
+          setting :namespace, default: "waterdrop"
 
           # Datadog client that we should use to publish the metrics
           setting :client
@@ -35,19 +35,19 @@ module WaterDrop
           # 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'),
+            RdKafkaMetric.new(:count, :root, "calls", "tx_d"),
+            RdKafkaMetric.new(:histogram, :root, "queue.size", "msg_cnt"),
 
             # 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])
+            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
@@ -78,13 +78,13 @@ module WaterDrop
           #
           # @param _event [Karafka::Core::Monitoring::Event]
           def on_error_occurred(_event)
-            count('error_occurred', 1, tags: default_tags)
+            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)
+            increment("acknowledged", tags: default_tags)
           end
 
           %i[
@@ -216,26 +216,26 @@ module WaterDrop
                 tags: default_tags
               )
             when :brokers
-              statistics.fetch('brokers').each_value do |broker_statistics|
+              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
+                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']}"]
+                  tags: default_tags + ["broker:#{broker_statistics["nodename"]}"]
                 )
               end
             when :topics
-              statistics.fetch('topics').each_value do |topic_statistics|
+              statistics.fetch("topics").each_value do |topic_statistics|
                 public_send(
                   metric.type,
                   metric.name,
                   topic_statistics.dig(*metric.key_location),
-                  tags: default_tags + ["topic:#{topic_statistics['topic']}"]
+                  tags: default_tags + ["topic:#{topic_statistics["topic"]}"]
                 )
               end
             else

data/lib/waterdrop/polling/config.rb

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+# WaterDrop main module
+module WaterDrop
+  # Namespace for FD-based polling components
+  module Polling
+    # Configuration for the global FD poller singleton
+    # These settings apply to all producers using FD polling mode
+    #
+    # @example Configure before creating any producers
+    #   WaterDrop::Polling::Config.setup do |config|
+    #     config.thread_priority = -1
+    #     config.poll_timeout = 500
+    #   end
+    class Config
+      extend ::Karafka::Core::Configurable
+
+      # Ruby thread priority for the poller thread
+      # Valid range: -3 to 3 (Ruby's thread priority range)
+      # Higher values = higher priority
+      setting :thread_priority, default: 0
+
+      # IO.select timeout in milliseconds
+      # Controls how often periodic polling happens when no FD events occur
+      # Lower values = more responsive OAuth/stats callbacks but higher CPU
+      setting :poll_timeout, default: 1_000
+
+      # Initial backoff delay in milliseconds after a polling error
+      setting :backoff_min, default: 100
+
+      # Maximum backoff delay in milliseconds after repeated errors
+      # Backoff doubles on each consecutive error up to this limit
+      setting :backoff_max, default: 30_000
+
+      class << self
+        # Configures the poller settings
+        # @yield [config] Configuration block
+        # @yieldparam config [Karafka::Core::Configurable::Node] config node
+        def setup
+          configure do |config|
+            yield(config)
+          end
+
+          Contracts::PollerConfig.new.validate!(
+            self.config.to_h,
+            Errors::ConfigurationInvalidError
+          )
+        end
+      end
+    end
+  end
+end

data/lib/waterdrop/polling/latch.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+module WaterDrop
+  module Polling
+    # A thread-safe latch for synchronizing producer close operations
+    #
+    # When a producer is closed, two threads are involved:
+    # 1. The caller thread (user code calling producer.close)
+    # 2. The poller thread (background thread running IO.select)
+    #
+    # The close sequence:
+    # 1. Caller calls producer.close -> unregister_from_poller -> Poller#unregister
+    # 2. Poller#unregister signals via control pipe and calls state.wait_for_close (blocks on latch)
+    # 3. Poller thread receives control signal, drains queue, calls state.close
+    # 4. state.close releases the latch via release!
+    # 5. Caller's wait_for_close returns, unregister completes
+    #
+    # This ensures the producer is fully drained and removed from the poller
+    # before returning control to the caller, preventing race conditions.
+    class Latch
+      def initialize
+        @mutex = Mutex.new
+        @cv = ConditionVariable.new
+        @released = false
+      end
+
+      # Releases the latch and wakes any waiting threads
+      def release!
+        @mutex.synchronize do
+          @released = true
+          @cv.broadcast
+        end
+      end
+
+      # Waits until the latch is released
+      # Returns immediately if already released
+      def wait
+        @mutex.synchronize do
+          @cv.wait(@mutex) until @released
+        end
+      end
+
+      # @return [Boolean] whether the latch has been released
+      def released?
+        @released
+      end
+    end
+  end
+end

data/lib/waterdrop/polling/poller.rb

@@ -0,0 +1,415 @@
+# frozen_string_literal: true
+
+module WaterDrop
+  # Namespace for FD-based polling components
+  # Contains the global Poller singleton and State class for managing producer polling
+  module Polling
+    # Global poller singleton that manages a single polling thread for all FD-mode producers
+    # This replaces librdkafka's native background polling threads with a single Ruby thread
+    # that uses IO.select for efficient multiplexing
+    #
+    # Spawning one thread per producer is acceptable for 1-2 producers but in case of a system
+    # with several (transactional for example) the cost becomes bigger and bigger.
+    #
+    # This implementation handles things by being event-driven instead of GVL releasing blocking.
+    #
+    # @note Newly registered producers may experience up to 1 second delay before their first
+    #   poll cycle, as the poller thread only rebuilds its IO list when IO.select times out.
+    #   This is acceptable because producers are expected to be long-lived and the initial
+    #   connection overhead to Kafka typically exceeds this delay anyway.
+    class Poller
+      include Singleton
+      include ::Karafka::Core::Helpers::Time
+
+      # Make new public so users can create dedicated poller instances for isolation
+      # The singleton instance remains available via Poller.instance for the default behavior
+      public_class_method :new
+
+      # Mutex for thread-safe ID generation - initialized at class load time
+      # to avoid race conditions with lazy initialization
+      ID_MUTEX = Mutex.new
+
+      # Counter for generating unique poller IDs
+      @id_counter = 0
+
+      class << self
+        # Generates incremental IDs for poller instances (starting from 0)
+        # @return [Integer] next poller ID
+        def next_id
+          ID_MUTEX.synchronize do
+            id = @id_counter
+            @id_counter += 1
+            id
+          end
+        end
+      end
+
+      # @return [Integer] unique identifier for this poller instance
+      attr_reader :id
+
+      def initialize
+        @id = self.class.next_id
+        @mutex = Mutex.new
+        @producers = {}
+        @thread = nil
+        @shutdown = false
+        @pid = Process.pid
+
+        # Cached collections - rebuilt only when producers change
+        @cached_ios = []
+        @cached_io_to_state = {}
+        @cached_states = []
+        @cached_result = nil
+        @ios_dirty = true
+      end
+
+      # Checks if the current thread is the poller thread
+      # Used to detect when close is called from within a callback to avoid deadlock
+      # @return [Boolean] true if current thread is the poller thread
+      def in_poller_thread?
+        Thread.current == @thread
+      end
+
+      # Checks if the poller thread is alive
+      # @return [Boolean] true if the poller thread is running
+      def alive?
+        @thread&.alive? || false
+      end
+
+      # Returns the number of registered producers
+      # @return [Integer] number of producers
+      def count
+        @mutex.synchronize { @producers.size }
+      end
+
+      # Shuts down the poller and resets state
+      # @note This is primarily for testing to reset singleton state between tests
+      def shutdown!
+        @mutex.synchronize { @shutdown = true }
+
+        thread = @thread
+        if thread&.alive?
+          thread.join(5)
+          thread.kill if thread.alive?
+        end
+
+        @mutex.synchronize do
+          @producers.each_value { |state| state.close unless state.closed? }
+          @producers.clear
+          @thread = nil
+          @shutdown = false
+          @ios_dirty = true
+          @cached_ios = []
+          @cached_io_to_state = {}
+          @cached_states = []
+          @cached_result = nil
+          @poll_timeout_s = nil
+        end
+      end
+
+      # Registers a producer with the poller
+      # @param producer [WaterDrop::Producer] the producer instance
+      # @param client [Rdkafka::Producer] the rdkafka client
+      def register(producer, client)
+        ensure_same_process!
+
+        state = State.new(
+          producer.id,
+          client,
+          producer.monitor,
+          producer.config.polling.fd.max_time,
+          producer.config.polling.fd.periodic_poll_interval
+        )
+
+        @mutex.synchronize do
+          @producers[producer.id] = state
+          @ios_dirty = true
+          # Reset shutdown flag in case thread is exiting but hasn't yet
+          # This prevents race where new producer is closed by exiting thread
+          @shutdown = false
+          ensure_thread_running!
+        end
+
+        producer.monitor.instrument(
+          "poller.producer_registered",
+          producer_id: producer.id
+        )
+      end
+
+      # Unregisters a producer from the poller
+      # This method blocks until the producer is fully removed from the poller
+      # to prevent race conditions when disconnect/reconnect happens in quick succession
+      # This matches the threaded polling behavior which drains without timeout
+      # @param producer [WaterDrop::Producer] the producer instance
+      def unregister(producer)
+        state, thread = @mutex.synchronize { [@producers[producer.id], @thread] }
+
+        return unless state
+
+        # Signal the poller thread to handle removal
+        state.signal_close
+
+        # Wait for the state to be fully closed by the poller thread
+        # This prevents race conditions where a new registration with the same
+        # producer_id could be deleted by a pending close signal
+        # Skip waiting if called from within the poller thread itself (e.g., from a callback)
+        # to avoid deadlock - the poller thread can't wait for itself
+        # The cleanup will happen after the callback returns
+        state.wait_for_close unless Thread.current == thread
+
+        producer.monitor.instrument(
+          "poller.producer_unregistered",
+          producer_id: producer.id
+        )
+      end
+
+      private
+
+      # Ensures we're in the same process (for fork safety)
+      def ensure_same_process!
+        return if @pid == Process.pid
+
+        # Reset state after fork - parent's thread and producers are not valid in child
+        @mutex = Mutex.new
+        @producers = {}
+        @thread = nil
+        @shutdown = false
+        @pid = Process.pid
+        @cached_ios = []
+        @cached_io_to_state = {}
+        @cached_states = []
+        @cached_result = nil
+        @ios_dirty = true
+      end
+
+      # Ensures the polling thread is running
+      # Must be called within @mutex.synchronize
+      def ensure_thread_running!
+        return if @thread&.alive?
+
+        @shutdown = false
+        @thread = Thread.new { polling_loop }
+        @thread.name = "waterdrop.poller##{@id}"
+        @thread.priority = Config.config.thread_priority
+      end
+
+      # Main polling loop that runs in a dedicated thread
+      def polling_loop
+        backoff_ms = 0
+
+        loop do
+          break if @shutdown
+
+          # Apply backoff from previous error
+          if backoff_ms > 0
+            sleep(backoff_ms / 1_000.0)
+            backoff_ms = 0
+          end
+
+          # Collect readable IOs (queue FDs)
+          readable_ios, io_to_state = collect_readable_ios
+
+          # Exit when no producers registered
+          # New registrations will start a fresh thread via ensure_thread_running!
+          break if readable_ios.empty?
+
+          poll_with_select(readable_ios, io_to_state)
+        rescue => e
+          # Report error and apply exponential backoff to prevent spam
+          broadcast_error("poller.polling_loop", e)
+          backoff_ms =
+            if backoff_ms.zero?
+              Config.config.backoff_min
+            else
+              [backoff_ms * 2, Config.config.backoff_max].min
+            end
+        end
+      ensure
+        # Clear thread reference first so new registrations will start a fresh thread
+        # This prevents race where register sees old thread as alive during cleanup
+        @mutex.synchronize { @thread = nil }
+
+        # When the poller thread exits (error or clean shutdown), close all remaining states
+        # This releases any latches that might be waiting in unregister calls
+        close_all_states
+      end
+
+      # Broadcasts an error to all registered producers' monitors
+      # @param type [String] error type identifier
+      # @param error [Exception] the error to report
+      def broadcast_error(type, error)
+        @cached_states.each do |state|
+          state.monitor.instrument(
+            "error.occurred",
+            type: type,
+            error: error,
+            producer_id: state.producer_id
+          )
+        end
+      end
+
+      # Collects all IOs to monitor and builds a mapping from IO to State
+      # Uses cached arrays when possible to avoid allocations in the hot path
+      # @return [Array<Array<IO>, Hash{IO => State}, Array<State>>] tuple of ios, io-to-state map, states
+      def collect_readable_ios
+        # Fast path: return cached result if not dirty (no mutex needed)
+        # Safe because @cached_result is frozen and assigned atomically
+        return @cached_result unless @ios_dirty
+
+        @mutex.synchronize do
+          @cached_ios = []
+          @cached_io_to_state = {}
+          @cached_states = []
+
+          @producers.each_value do |state|
+            io = state.io
+            @cached_ios << io
+            @cached_io_to_state[io] = state
+            @cached_states << state
+          end
+
+          @cached_result = [@cached_ios, @cached_io_to_state, @cached_states].freeze
+          @ios_dirty = false
+        end
+
+        @cached_result
+      end
+
+      # Poll producers using IO.select for efficient multiplexing
+      # @param readable_ios [Array<IO>] IOs to monitor
+      # @param io_to_state [Hash{IO => State}] mapping from IO to state
+      def poll_with_select(readable_ios, io_to_state)
+        begin
+          ready = IO.select(readable_ios, nil, nil, poll_timeout_s)
+        rescue IOError, Errno::EBADF
+          # An IO was closed - mark dirty to rebuild on next iteration
+          @ios_dirty = true
+          return
+        end
+
+        if ready.nil?
+          # Timeout: poll ALL producers to ensure OAuth/stats fire
+          poll_all_producers
+        else
+          # FDs ready: handle close signals and poll active producers
+          any_polled = false
+
+          ready[0].each do |io|
+            state = io_to_state[io]
+            next unless state
+
+            # Drain the pipe first (clears librdkafka signals + our signals)
+            state.drain
+
+            # Check if this producer is closing (flag set before signal)
+            if state.closing?
+              handle_close_signal(state)
+            else
+              poll_producer(state)
+              # Check if callback signaled close while we were polling
+              # (e.g., user code closed producer from within delivery callback)
+              if state.closing?
+                handle_close_signal(state)
+              else
+                any_polled = true
+              end
+            end
+          end
+
+          # Check for stale producers when actively polling
+          # Skip when single producer (most common case) - no other producers to become stale
+          # (ensures OAuth/stats fire for idle producers when others are busy)
+          poll_stale_producers if any_polled && @cached_states.size > 1
+        end
+      end
+
+      # Polls all registered producers
+      # Called when IO.select times out to ensure periodic polling happens
+      # This ensures OAuth token refresh and statistics callbacks fire for all producers
+      def poll_all_producers
+        @cached_states.each { |state| poll_producer(state) }
+      end
+
+      # Polls producers that haven't been polled recently
+      # Called when processing continue signals to prevent starvation of idle producers
+      # when one producer is very busy
+      # Each State internally throttles the check to avoid excessive overhead
+      def poll_stale_producers
+        @cached_states.each do |state|
+          poll_producer(state) if state.needs_periodic_poll?
+        end
+      end
+
+      # Drains the producer's event queue by polling until empty or time quanta exceeded
+      # @param state [State] the producer state
+      def poll_producer(state)
+        # state.poll returns:
+        # - true when queue is empty (fully drained)
+        # - false when timeout hit (more events may remain)
+        drained = state.poll
+        state.mark_polled!
+
+        # Hit time limit but still have events - signal to continue polling
+        state.signal_continue unless drained
+      rescue Rdkafka::ClosedProducerError
+        # Producer was closed, will be cleaned up
+      end
+
+      # Handles a close signal from a producer
+      # @param state [State] the producer state
+      def handle_close_signal(state)
+        # Drain remaining events before closing
+        # This matches rdkafka's native polling thread behavior: keep polling until outq_len is zero
+        drain_producer_queue(state)
+
+        # Remove producer from registry and clean up
+        # If this was the last producer, signal shutdown to stop the thread immediately
+        @mutex.synchronize do
+          @producers.delete(state.producer_id)
+          @ios_dirty = true
+
+          # Stop thread immediately when last producer unregisters to prevent resource leakage
+          @shutdown = true if @producers.empty?
+        end
+
+        state.close
+      end
+
+      # Closes all remaining producer states
+      # Called when the poller thread exits to release any pending latches
+      # This prevents deadlocks if producers are waiting in unregister
+      def close_all_states
+        states = @mutex.synchronize do
+          to_close = @producers.values.dup
+          @producers.clear
+          @ios_dirty = true
+          to_close
+        end
+
+        states.each do |state|
+          state.close unless state.closed?
+        rescue
+          # Ignore errors during cleanup
+        end
+      end
+
+      # Drains the producer's event queue completely before closing
+      # Matches rdkafka's native polling thread behavior: keep polling until queue is empty
+      # @param state [State] the producer state
+      def drain_producer_queue(state)
+        loop do
+          break if state.queue_empty?
+
+          state.poll
+        end
+      rescue Rdkafka::ClosedProducerError
+        # Producer was already closed, nothing more to drain
+      end
+
+      # @return [Float] poll_timeout converted to seconds (cached)
+      def poll_timeout_s
+        @poll_timeout_s ||= Config.config.poll_timeout / 1_000.0
+      end
+    end
+  end
+end

data/lib/waterdrop/polling/queue_pipe.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module WaterDrop
+  module Polling
+    # A pipe connected to librdkafka's queue event notification system
+    # When events (delivery reports, statistics) arrive, librdkafka writes to the pipe
+    # allowing IO.select to wake up immediately
+    #
+    # This pipe is also used by WaterDrop to signal:
+    # - Continue: when poll hits time limit but more events remain
+    # - Close: when producer is being closed (combined with @closing flag in State)
+    #
+    # Reusing the same pipe reduces file descriptors and IO.select monitoring overhead
+    class QueuePipe
+      # @return [IO] the readable end of the pipe for use with IO.select
+      attr_reader :reader
+
+      # Creates a new queue pipe and connects it to the client's event queue
+      # @param client [Rdkafka::Producer] the rdkafka client
+      # @raise [StandardError] if enable_queue_io_events fails
+      def initialize(client)
+        @reader, @writer = IO.pipe
+
+        # Tell librdkafka to write to our pipe when events arrive on the main queue
+        client.enable_queue_io_events(@writer.fileno)
+      end
+
+      # Signals by writing a byte to the pipe
+      # Used to wake IO.select for continue/close signals
+      # Thread-safe and non-blocking; silently ignores errors
+      def signal
+        @writer.write_nonblock("W", exception: false)
+      rescue IOError, Errno::EBADF
+        # Pipe closed
+      end
+
+      # Drains all pending bytes from the pipe
+      # Called after IO.select returns to clear the notification
+      # Uses a single large read to drain in one syscall (pipe buffers are typically 64KB)
+      def drain
+        @reader.read_nonblock(1_048_576, exception: false)
+      rescue IOError, Errno::EBADF
+        # Pipe closed during drain
+      end
+
+      # Closes both ends of the pipe
+      def close
+        close_io(@reader)
+        close_io(@writer)
+      end
+
+      private
+
+      # Safely closes an IO object
+      # @param io [IO] the IO to close
+      def close_io(io)
+        io.close
+      rescue IOError, Errno::EBADF
+        # Already closed, ignore
+      end
+    end
+  end
+end