Merge pull request #814 from karafka/feature/events-poll-nb-each

Add events_poll_nb_each and poll_nb_each for fiber scheduler integration

Author: mensfeld

CHANGELOG.md

@@ -3,8 +3,10 @@
 ## 0.25.1 (Unreleased)
 - [Enhancement] Use native ARM64 runners instead of QEMU emulation for Alpine musl aarch64 builds, improving build performance and reliability.
 - [Enhancement] Enable parallel compilation (`make -j$(nproc)`) for ARM64 Alpine musl builds.
-- [Enhancement] Add file descriptor API for fiber scheduler integration. Expose `queue_fd` and `background_queue_fd` on `Consumer`, `Producer`, and `Admin` to enable non-blocking monitoring with select/poll/epoll for integration with Ruby fiber schedulers (Falcon, Async) and custom event loops.
-- [Enhancement] Add non-blocking poll methods (`poll_nb`, `events_poll_nb`) that skip GVL release for efficient fiber scheduler integration when using `poll(0)`.
+- [Enhancement] Add file descriptor API for fiber scheduler integration. Expose `enable_queue_io_events` and `enable_background_queue_io_events` on `Consumer`, `Producer`, and `Admin` to enable non-blocking monitoring with select/poll/epoll for integration with Ruby fiber schedulers (Falcon, Async) and custom event loops.
+- [Enhancement] Add non-blocking poll methods (`poll_nb`, `events_poll_nb`) on `Consumer` that skip GVL release for efficient fiber scheduler integration when using `poll(0)`.
+- [Enhancement] Add `events_poll_nb_each` method on `Producer`, `Consumer`, and `Admin` for polling events in a single GVL/mutex session. Yields count after each iteration, caller returns `:stop` to break.
+- [Enhancement] Add `poll_nb_each` method on `Consumer` for non-blocking message polling with proper resource cleanup, yielding each message and supporting early termination via `:stop` return value.
 
 ## 0.25.0 (2026-01-20)
 - **[Deprecation]** `AbstractHandle#wait` parameter `max_wait_timeout:` (seconds) is deprecated in favor of `max_wait_timeout_ms:` (milliseconds). The old parameter still works but will be removed in v1.0.0.

lib/rdkafka/admin.rb

@@ -91,6 +91,45 @@ def enable_background_queue_io_events(fd, payload = "\x01")
       @native_kafka.enable_background_queue_io_events(fd, payload)
     end
 
+    # Polls for events in a non-blocking loop, yielding the count after each iteration.
+    #
+    # This method processes events (stats, errors, etc.) in a single GVL/mutex session,
+    # which is more efficient than repeated individual polls. It uses non-blocking polls
+    # internally (no GVL release between polls).
+    #
+    # Yields the count of events processed after each poll iteration, allowing the caller
+    # to implement timeout or other termination logic by returning `:stop`.
+    #
+    # @yield [count] Called after each poll iteration
+    # @yieldparam count [Integer] Number of events processed in this iteration
+    # @yieldreturn [Symbol, Object] Return `:stop` to break the loop, any other value continues
+    # @return [nil]
+    # @raise [Rdkafka::ClosedAdminError] if called on a closed admin client
+    #
+    # @note This method holds the inner lock until the queue is empty or `:stop` is returned.
+    #   Other admin operations will wait until this method returns.
+    # @note This method is thread-safe as it uses @native_kafka.with_inner synchronization
+    #
+    # @example Drain all pending events
+    #   admin.events_poll_nb_each { |_count| }
+    #
+    # @example With timeout control
+    #   deadline = monotonic_now + timeout_ms
+    #   admin.events_poll_nb_each do |_count|
+    #     :stop if monotonic_now >= deadline
+    #   end
+    def events_poll_nb_each
+      closed_admin_check(__method__)
+
+      @native_kafka.with_inner do |inner|
+        loop do
+          count = Rdkafka::Bindings.rd_kafka_poll_nb(inner, 0)
+          break if count.zero?
+          break if yield(count) == :stop
+        end
+      end
+    end
+
     # @return [Proc] finalizer proc for closing the admin
     # @private
     def finalizer

lib/rdkafka/consumer.rb

@@ -75,6 +75,103 @@ def enable_background_queue_io_events(fd, payload = "\x01")
       @native_kafka.enable_background_queue_io_events(fd, payload)
     end
 
+    # Polls for events in a non-blocking loop, yielding the count after each iteration.
+    #
+    # This method processes events (stats, errors, etc.) in a single GVL/mutex session,
+    # which is more efficient than repeated individual polls. It uses non-blocking polls
+    # internally (no GVL release between polls).
+    #
+    # Yields the count of events processed after each poll iteration, allowing the caller
+    # to implement timeout or other termination logic by returning `:stop`.
+    #
+    # @yield [count] Called after each poll iteration
+    # @yieldparam count [Integer] Number of events processed in this iteration
+    # @yieldreturn [Symbol, Object] Return `:stop` to break the loop, any other value continues
+    # @return [nil]
+    # @raise [Rdkafka::ClosedConsumerError] if called on a closed consumer
+    #
+    # @note This method holds the inner lock until the queue is empty or `:stop` is returned.
+    #   Other consumer operations will wait until this method returns.
+    # @note This method is thread-safe as it uses @native_kafka.with_inner synchronization
+    # @note Do NOT use this if `consumer_poll_set` was set to `true`
+    #
+    # @example Drain all pending events
+    #   consumer.events_poll_nb_each { |_count| }
+    #
+    # @example With timeout control
+    #   deadline = monotonic_now + timeout_ms
+    #   consumer.events_poll_nb_each do |_count|
+    #     :stop if monotonic_now >= deadline
+    #   end
+    def events_poll_nb_each
+      closed_consumer_check(__method__)
+
+      @native_kafka.with_inner do |inner|
+        loop do
+          count = Rdkafka::Bindings.rd_kafka_poll_nb(inner, 0)
+          break if count.zero?
+          break if yield(count) == :stop
+        end
+      end
+    end
+
+    # Polls for messages in a non-blocking loop, yielding each message to the caller.
+    #
+    # This method processes messages in a single GVL/mutex session until the queue is empty
+    # or the caller returns `:stop`. It handles the message pointer lifecycle internally,
+    # ensuring proper cleanup via `rd_kafka_message_destroy`.
+    #
+    # @yield [message] Called for each message received
+    # @yieldparam message [Consumer::Message] The received message
+    # @yieldreturn [Symbol, Object] Return `:stop` to break the loop, any other value continues
+    # @return [nil]
+    # @raise [Rdkafka::ClosedConsumerError] if called on a closed consumer
+    # @raise [Rdkafka::RdkafkaError] if a Kafka error occurs while polling
+    #
+    # @note This method uses `rd_kafka_consumer_poll` to fetch messages, unlike
+    #   `events_poll_nb_each` which uses `rd_kafka_poll` for event callbacks (delivery reports,
+    #   statistics, etc.). For consumers, use this method to receive messages and
+    #   `events_poll_nb_each` for processing background events.
+    # @note This method holds the inner lock for the duration. Other consumer operations
+    #   will wait until this method returns.
+    # @note Timeout/max_messages logic should be implemented by the caller
+    #
+    # @example Process messages until queue is empty
+    #   consumer.poll_nb_each do |message|
+    #     process(message)
+    #   end
+    #
+    # @example Process with early termination
+    #   count = 0
+    #   consumer.poll_nb_each do |message|
+    #     process(message)
+    #     count += 1
+    #     :stop if count >= 10
+    #   end
+    def poll_nb_each
+      closed_consumer_check(__method__)
+
+      @native_kafka.with_inner do |inner|
+        loop do
+          message_ptr = Rdkafka::Bindings.rd_kafka_consumer_poll_nb(inner, 0)
+          break if message_ptr.null?
+
+          begin
+            native_message = Rdkafka::Bindings::Message.new(message_ptr)
+
+            if native_message[:err] != Rdkafka::Bindings::RD_KAFKA_RESP_ERR_NO_ERROR
+              raise Rdkafka::RdkafkaError.new(native_message[:err])
+            end
+
+            result = yield Consumer::Message.new(native_message)
+            break if result == :stop
+          ensure
+            Rdkafka::Bindings.rd_kafka_message_destroy(message_ptr)
+          end
+        end
+      end
+    end
+
     # @return [Proc] finalizer proc for closing the consumer
     # @private
     def finalizer

lib/rdkafka/producer.rb

@@ -258,41 +258,41 @@ def queue_size
 
     alias_method :queue_length, :queue_size
 
-    # Drains the producer's event queue by continuously polling until empty or time limit reached.
+    # Polls for events in a non-blocking loop, yielding the count after each iteration.
     #
-    # This method is useful when you need to ensure delivery callbacks are processed within a
-    # bounded time, particularly when polling multiple producers from a single thread where
-    # fair scheduling is required to prevent starvation.
+    # This method processes delivery callbacks in a single GVL/mutex session, which is more
+    # efficient than repeated individual polls. It uses non-blocking polls internally
+    # (no GVL release between polls).
     #
-    # Uses non-blocking polls internally (no GVL release) for efficiency. The method holds
-    # a single `with_inner` lock for the duration, minimizing per-poll overhead when processing
-    # many events.
+    # Yields the count of events processed after each poll iteration, allowing the caller
+    # to implement timeout or other termination logic by returning `:stop`.
     #
-    # @param timeout_ms [Integer] maximum time to spend draining in milliseconds (default: 100)
-    # @return [Boolean] true if no more events to process, false if stopped due to time limit
+    # @yield [count] Called after each poll iteration
+    # @yieldparam count [Integer] Number of events processed in this iteration
+    # @yieldreturn [Symbol, Object] Return `:stop` to break the loop, any other value continues
+    # @return [nil]
     # @raise [Rdkafka::ClosedProducerError] if called on a closed producer
     #
-    # @note This method holds the inner lock for up to `timeout_ms`. Other producer operations
-    #   (produce, close, etc.) will wait until this method returns.
+    # @note This method holds the inner lock until the queue is empty or `:stop` is returned.
+    #   Other producer operations (produce, close, etc.) will wait until this method returns.
     # @note This method is thread-safe as it uses @native_kafka.with_inner synchronization
     #
-    # @example Basic usage - drain for up to 100ms
-    #   fully_drained = producer.poll_drain_nb
+    # @example Drain all pending callbacks
+    #   producer.events_poll_nb_each { |_count| }
     #
-    # @example Round-robin polling multiple producers fairly
-    #   producers.each do |producer|
-    #     fully_drained = producer.poll_drain_nb(10)
-    #     # If false, this producer has more pending events
+    # @example With timeout control
+    #   deadline = monotonic_now + timeout_ms
+    #   producer.events_poll_nb_each do |_count|
+    #     :stop if monotonic_now >= deadline
     #   end
-    def poll_drain_nb(timeout_ms = 100)
+    def events_poll_nb_each
       closed_producer_check(__method__)
 
       @native_kafka.with_inner do |inner|
-        deadline = monotonic_now_ms + timeout_ms
-
         loop do
-          break true if Rdkafka::Bindings.rd_kafka_poll_nb(inner, 0).zero?
-          break false if monotonic_now_ms >= deadline
+          count = Rdkafka::Bindings.rd_kafka_poll_nb(inner, 0)
+          break if count.zero?
+          break if yield(count) == :stop
         end
       end
     end

spec/lib/rdkafka/admin_spec.rb

@@ -972,6 +972,47 @@
     end
   end
 
+  describe "#events_poll_nb_each" do
+    it "does not raise when queue is empty" do
+      expect { admin.events_poll_nb_each { |_| } }.not_to raise_error
+    end
+
+    it "yields the count after each poll" do
+      counts = []
+      # Stub to return events, then zero
+      call_count = 0
+      allow(Rdkafka::Bindings).to receive(:rd_kafka_poll_nb) do
+        call_count += 1
+        (call_count <= 2) ? 1 : 0
+      end
+
+      admin.events_poll_nb_each { |count| counts << count }
+
+      expect(counts).to eq([1, 1])
+    end
+
+    it "stops when block returns :stop" do
+      iterations = 0
+      # Stub to always return events
+      allow(Rdkafka::Bindings).to receive(:rd_kafka_poll_nb).and_return(1)
+
+      admin.events_poll_nb_each do |_count|
+        iterations += 1
+        :stop if iterations >= 3
+      end
+
+      expect(iterations).to eq(3)
+    end
+
+    context "when admin is closed" do
+      before { admin.close }
+
+      it "raises ClosedAdminError" do
+        expect { admin.events_poll_nb_each { |_| } }.to raise_error(Rdkafka::ClosedAdminError, /events_poll_nb_each/)
+      end
+    end
+  end
+
   describe "file descriptor access for fiber scheduler integration" do
     let(:admin) { config.admin(run_polling_thread: false) }
 

spec/lib/rdkafka/consumer_spec.rb

@@ -1439,6 +1439,109 @@ def collect(name, list)
     end
   end
 
+  describe "#events_poll_nb_each" do
+    it "does not raise when queue is empty" do
+      expect { consumer.events_poll_nb_each { |_| } }.not_to raise_error
+    end
+
+    it "yields the count after each poll" do
+      counts = []
+      # Stub to return events, then zero
+      call_count = 0
+      allow(Rdkafka::Bindings).to receive(:rd_kafka_poll_nb) do
+        call_count += 1
+        (call_count <= 2) ? 1 : 0
+      end
+
+      consumer.events_poll_nb_each { |count| counts << count }
+
+      expect(counts).to eq([1, 1])
+    end
+
+    it "stops when block returns :stop" do
+      iterations = 0
+      # Stub to always return events
+      allow(Rdkafka::Bindings).to receive(:rd_kafka_poll_nb).and_return(1)
+
+      consumer.events_poll_nb_each do |_count|
+        iterations += 1
+        :stop if iterations >= 3
+      end
+
+      expect(iterations).to eq(3)
+    end
+
+    context "when consumer is closed" do
+      before { consumer.close }
+
+      it "raises ClosedConsumerError" do
+        expect { consumer.events_poll_nb_each { |_| } }.to raise_error(Rdkafka::ClosedConsumerError, /events_poll_nb_each/)
+      end
+    end
+  end
+
+  describe "#poll_nb_each" do
+    it "does not raise when queue is empty" do
+      consumer.subscribe(TestTopics.consume_test_topic)
+      # Give it a moment to subscribe
+      sleep 0.5
+
+      messages = []
+      consumer.poll_nb_each { |msg| messages << msg }
+      expect(messages).to be_a(Array)
+    end
+
+    it "yields messages and respects :stop" do
+      topic = TestTopics.consume_test_topic
+      consumer.subscribe(topic)
+
+      # Produce some messages
+      5.times { |i| producer.produce(topic: topic, payload: "poll_nb_each test #{i}") }
+      producer.flush
+
+      # Use blocking poll first to ensure consumer is ready and messages are fetched
+      # poll_nb_each is non-blocking so we need to ensure messages are available first
+      deadline = Time.now + 30
+      first_message = nil
+      while Time.now < deadline && first_message.nil?
+        first_message = consumer.poll(100)
+      end
+
+      # Now test that :stop works - we should get exactly one more message then stop
+      # (we already consumed one with blocking poll above)
+      messages = []
+      consumer.poll_nb_each do |message|
+        messages << message
+        :stop if messages.size >= 1
+      end
+
+      # Should have stopped after exactly 1 message (we got 1 via blocking poll earlier)
+      expect(messages.size).to eq(1)
+    end
+
+    it "properly cleans up message pointers" do
+      topic = TestTopics.consume_test_topic
+      consumer.subscribe(topic)
+
+      producer.produce(topic: topic, payload: "cleanup test")
+      producer.flush
+      sleep 2
+
+      # This should not leak memory - message_destroy is called in ensure
+      expect {
+        consumer.poll_nb_each { |_| }
+      }.not_to raise_error
+    end
+
+    context "when consumer is closed" do
+      before { consumer.close }
+
+      it "raises ClosedConsumerError" do
+        expect { consumer.poll_nb_each { |_| } }.to raise_error(Rdkafka::ClosedConsumerError, /poll_nb_each/)
+      end
+    end
+  end
+
   describe "file descriptor access for fiber scheduler integration" do
     it "enables IO events on consumer queue" do
       consumer.subscribe(TestTopics.consume_test_topic)