rdkafka 0.11.1 → 0.13.1

data/lib/rdkafka/producer.rb

@@ -1,4 +1,6 @@
-require "securerandom"
+# frozen_string_literal: true
+
+require "objspace"
 
 module Rdkafka
   # A producer for Kafka messages. To create a producer set up a {Config} and call {Config#producer producer} on that.
@@ -10,29 +12,22 @@ module Rdkafka
     attr_reader :delivery_callback
 
     # @private
-    def initialize(native_kafka)
-      @id = SecureRandom.uuid
-      @closing = false
+    # Returns the number of arguments accepted by the callback, by default this is nil.
+    #
+    # @return [Integer, nil]
+    attr_reader :delivery_callback_arity
+
+    # @private
+    def initialize(native_kafka, partitioner_name)
       @native_kafka = native_kafka
+      @partitioner_name = partitioner_name || "consistent_random"
 
-      # Makes sure, that the producer gets closed before it gets GCed by Ruby
-      ObjectSpace.define_finalizer(@id, proc { close })
-
-      # Start thread to poll client for delivery callbacks
-      @polling_thread = Thread.new do
-        loop do
-          Rdkafka::Bindings.rd_kafka_poll(@native_kafka, 250)
-          # Exit thread if closing and the poll queue is empty
-          if @closing && Rdkafka::Bindings.rd_kafka_outq_len(@native_kafka) == 0
-            break
-          end
-        end
-      end
-      @polling_thread.abort_on_exception = true
+      # Makes sure, that native kafka gets closed before it gets GCed by Ruby
+      ObjectSpace.define_finalizer(self, native_kafka.finalizer)
     end
 
     # Set a callback that will be called every time a message is successfully produced.
-    # The callback is called with a {DeliveryReport}
+    # The callback is called with a {DeliveryReport} and {DeliveryHandle}
     #
     # @param callback [Proc, #call] The callback
     #
@@ -40,20 +35,31 @@ module Rdkafka
     def delivery_callback=(callback)
       raise TypeError.new("Callback has to be callable") unless callback.respond_to?(:call)
       @delivery_callback = callback
+      @delivery_callback_arity = arity(callback)
     end
 
     # Close this producer and wait for the internal poll queue to empty.
     def close
-      ObjectSpace.undefine_finalizer(@id)
+      return if closed?
+      ObjectSpace.undefine_finalizer(self)
+      @native_kafka.close
+    end
 
-      return unless @native_kafka
+    # Whether this producer has closed
+    def closed?
+      @native_kafka.closed?
+    end
+
+    # Wait until all outstanding producer requests are completed, with the given timeout
+    # in seconds. Call this before closing a producer to ensure delivery of all messages.
+    #
+    # @param timeout_ms [Integer] how long should we wait for flush of all messages
+    def flush(timeout_ms=5_000)
+      closed_producer_check(__method__)
 
-      # Indicate to polling thread that we're closing
-      @closing = true
-      # Wait for the polling thread to finish up
-      @polling_thread.join
-      Rdkafka::Bindings.rd_kafka_destroy(@native_kafka)
-      @native_kafka = nil
+      @native_kafka.with_inner do |inner|
+        Rdkafka::Bindings.rd_kafka_flush(inner, timeout_ms)
+      end
     end
 
     # Partition count for a given topic.
@@ -62,10 +68,11 @@ module Rdkafka
     # @param topic [String] The topic name.
     #
     # @return partition count [Integer,nil]
-    #
     def partition_count(topic)
       closed_producer_check(__method__)
-      Rdkafka::Metadata.new(@native_kafka, topic).topics&.first[:partition_count]
+      @native_kafka.with_inner do |inner|
+        Rdkafka::Metadata.new(inner, topic).topics&.first[:partition_count]
+      end
     end
 
     # Produces a message to a Kafka topic. The message is added to rdkafka's queue, call {DeliveryHandle#wait wait} on the returned delivery handle to make sure it is delivered.
@@ -106,7 +113,7 @@ module Rdkafka
       if partition_key
         partition_count = partition_count(topic)
         # If the topic is not present, set to -1
-        partition = Rdkafka::Bindings.partitioner(partition_key, partition_count) if partition_count
+        partition = Rdkafka::Bindings.partitioner(partition_key, partition_count, @partitioner_name) if partition_count
       end
 
       # If partition is nil, use -1 to let librdafka set the partition randomly or
@@ -156,10 +163,12 @@ module Rdkafka
       args << :int << Rdkafka::Bindings::RD_KAFKA_VTYPE_END
 
       # Produce the message
-      response = Rdkafka::Bindings.rd_kafka_producev(
-        @native_kafka,
-        *args
-      )
+      response = @native_kafka.with_inner do |inner|
+        Rdkafka::Bindings.rd_kafka_producev(
+          inner,
+          *args
+        )
+      end
 
       # Raise error if the produce call was not successful
       if response != 0
@@ -170,13 +179,22 @@ module Rdkafka
       delivery_handle
     end
 
-    # @private
-    def call_delivery_callback(delivery_handle)
-      @delivery_callback.call(delivery_handle) if @delivery_callback
+    def call_delivery_callback(delivery_report, delivery_handle)
+      return unless @delivery_callback
+
+      args = [delivery_report, delivery_handle].take(@delivery_callback_arity)
+      @delivery_callback.call(*args)
+    end
+
+    def arity(callback)
+      return callback.arity if callback.respond_to?(:arity)
+
+      callback.method(:call).arity
     end
 
+    private
     def closed_producer_check(method)
-      raise Rdkafka::ClosedProducerError.new(method) if @native_kafka.nil?
+      raise Rdkafka::ClosedProducerError.new(method) if closed?
     end
   end
 end

data/lib/rdkafka/version.rb

@@ -1,5 +1,7 @@
+# frozen_string_literal: true
+
 module Rdkafka
-  VERSION = "0.11.1"
-  LIBRDKAFKA_VERSION = "1.8.2"
-  LIBRDKAFKA_SOURCE_SHA256 = "6a747d293a7a4613bd2897e28e8791476fbe1ae7361f2530a876e0fd483482a6"
+  VERSION = "0.13.1"
+  LIBRDKAFKA_VERSION = "2.0.2"
+  LIBRDKAFKA_SOURCE_SHA256 = "f321bcb1e015a34114c83cf1aa7b99ee260236aab096b85c003170c90a47ca9d"
 end

data/lib/rdkafka.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "rdkafka/version"
 
 require "rdkafka/abstract_handle"
@@ -16,6 +18,7 @@ require "rdkafka/consumer/partition"
 require "rdkafka/consumer/topic_partition_list"
 require "rdkafka/error"
 require "rdkafka/metadata"
+require "rdkafka/native_kafka"
 require "rdkafka/producer"
 require "rdkafka/producer/delivery_handle"
 require "rdkafka/producer/delivery_report"

data/rdkafka.gemspec

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require File.expand_path('lib/rdkafka/version', __dir__)
 
 Gem::Specification.new do |gem|

data/spec/rdkafka/abstract_handle_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::AbstractHandle do

data/spec/rdkafka/admin/create_topic_handle_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::Admin::CreateTopicHandle do

data/spec/rdkafka/admin/create_topic_report_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::Admin::CreateTopicReport do

data/spec/rdkafka/admin/delete_topic_handle_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::Admin::DeleteTopicHandle do

data/spec/rdkafka/admin/delete_topic_report_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::Admin::DeleteTopicReport do

data/spec/rdkafka/admin_spec.rb

@@ -1,9 +1,11 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 require "ostruct"
 
 describe Rdkafka::Admin do
-  let(:config)   { rdkafka_config }
-  let(:admin)    { config.admin }
+  let(:config) { rdkafka_config }
+  let(:admin)  { config.admin }
 
   after do
     # Registry should always end up being empty
@@ -174,7 +176,6 @@ describe Rdkafka::Admin do
       end
     end
 
-
     it "deletes a topic that was newly created" do
       create_topic_handle = admin.create_topic(topic_name, topic_partition_count, topic_replication_factor)
       create_topic_report = create_topic_handle.wait(max_wait_timeout: 15.0)

data/spec/rdkafka/bindings_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 require 'zlib'
 
@@ -76,6 +78,13 @@ describe Rdkafka::Bindings do
       result_2 = (Zlib.crc32(partition_key) % partition_count)
       expect(result_1).to eq(result_2)
     end
+
+    it "should return the partition calculated by the specified partitioner" do
+      result_1 = Rdkafka::Bindings.partitioner(partition_key, partition_count, "murmur2")
+      ptr = FFI::MemoryPointer.from_string(partition_key)
+      result_2 = Rdkafka::Bindings.rd_kafka_msg_partitioner_murmur2(nil, ptr, partition_key.size, partition_count, nil, nil)
+      expect(result_1).to eq(result_2)
+    end
   end
 
   describe "stats callback" do

data/spec/rdkafka/callbacks_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::Callbacks do

data/spec/rdkafka/config_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::Config do
@@ -148,11 +150,24 @@ describe Rdkafka::Config do
       }.to raise_error(Rdkafka::Config::ConfigError, "No such configuration property: \"invalid.key\"")
     end
 
+    it "allows string partitioner key" do
+      expect(Rdkafka::Producer).to receive(:new).with(kind_of(Rdkafka::NativeKafka), "murmur2").and_call_original
+      config = Rdkafka::Config.new("partitioner" => "murmur2")
+      config.producer.close
+    end
+
+    it "allows symbol partitioner key" do
+      expect(Rdkafka::Producer).to receive(:new).with(kind_of(Rdkafka::NativeKafka), "murmur2").and_call_original
+      config = Rdkafka::Config.new(:partitioner => "murmur2")
+      config.producer.close
+    end
+
     it "should allow configuring zstd compression" do
       config = Rdkafka::Config.new('compression.codec' => 'zstd')
       begin
-        expect(config.producer).to be_a Rdkafka::Producer
-        config.producer.close
+        producer = config.producer
+        expect(producer).to be_a Rdkafka::Producer
+        producer.close
       rescue Rdkafka::Config::ConfigError => ex
         pending "Zstd compression not supported on this machine"
         raise ex

data/spec/rdkafka/consumer/headers_spec.rb

@@ -0,0 +1,62 @@
+# frozen_string_literal: true
+
+require "spec_helper"
+
+describe Rdkafka::Consumer::Headers do
+  let(:headers) do
+    { # Note String keys!
+      "version" => "2.1.3",
+      "type" => "String"
+    }
+  end
+  let(:native_message) { double('native message') }
+  let(:headers_ptr) { double('headers pointer') }
+
+  describe '.from_native' do
+    before do
+      expect(Rdkafka::Bindings).to receive(:rd_kafka_message_headers).with(native_message, anything) do |_, headers_ptrptr|
+        expect(headers_ptrptr).to receive(:read_pointer).and_return(headers_ptr)
+        Rdkafka::Bindings::RD_KAFKA_RESP_ERR_NO_ERROR
+      end
+
+      expect(Rdkafka::Bindings).to \
+        receive(:rd_kafka_header_get_all)
+          .with(headers_ptr, 0, anything, anything, anything) do |_, _, name_ptrptr, value_ptrptr, size_ptr|
+              expect(name_ptrptr).to receive(:read_pointer).and_return(double("pointer 0", read_string_to_null: headers.keys[0]))
+              expect(size_ptr).to receive(:[]).with(:value).and_return(headers.keys[0].size)
+              expect(value_ptrptr).to receive(:read_pointer).and_return(double("value pointer 0", read_string: headers.values[0]))
+              Rdkafka::Bindings::RD_KAFKA_RESP_ERR_NO_ERROR
+            end
+
+      expect(Rdkafka::Bindings).to \
+        receive(:rd_kafka_header_get_all)
+          .with(headers_ptr, 1, anything, anything, anything) do |_, _, name_ptrptr, value_ptrptr, size_ptr|
+              expect(name_ptrptr).to receive(:read_pointer).and_return(double("pointer 1", read_string_to_null: headers.keys[1]))
+              expect(size_ptr).to receive(:[]).with(:value).and_return(headers.keys[1].size)
+              expect(value_ptrptr).to receive(:read_pointer).and_return(double("value pointer 1", read_string: headers.values[1]))
+              Rdkafka::Bindings::RD_KAFKA_RESP_ERR_NO_ERROR
+            end
+
+      expect(Rdkafka::Bindings).to \
+        receive(:rd_kafka_header_get_all)
+          .with(headers_ptr, 2, anything, anything, anything)
+          .and_return(Rdkafka::Bindings::RD_KAFKA_RESP_ERR__NOENT)
+    end
+
+    subject { described_class.from_native(native_message) }
+
+    it { is_expected.to eq(headers) }
+    it { is_expected.to be_frozen }
+
+    it 'allows String key' do
+      expect(subject['version']).to eq("2.1.3")
+    end
+
+    it 'allows Symbol key, but warns' do
+      expect(Kernel).to \
+        receive(:warn).with("rdkafka deprecation warning: header access with Symbol key :version treated as a String. " \
+                            "Please change your code to use String keys to avoid this warning. Symbol keys will break in version 1.")
+      expect(subject[:version]).to eq("2.1.3")
+    end
+  end
+end

data/spec/rdkafka/consumer/message_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::Consumer::Message do

data/spec/rdkafka/consumer/partition_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::Consumer::Partition do

data/spec/rdkafka/consumer/topic_partition_list_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 
 describe Rdkafka::Consumer::TopicPartitionList do

data/spec/rdkafka/consumer_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 require "ostruct"
 require 'securerandom'
@@ -53,7 +55,7 @@ describe Rdkafka::Consumer do
 
   describe "#pause and #resume" do
     context "subscription" do
-      let(:timeout) { 1000 }
+      let(:timeout) { 2000 }
 
       before { consumer.subscribe("consume_test_topic") }
       after { consumer.unsubscribe }
@@ -241,7 +243,7 @@ describe Rdkafka::Consumer do
 
     it "should return the assignment when subscribed" do
       # Make sure there's a message
-      report = producer.produce(
+      producer.produce(
         topic:     "consume_test_topic",
         payload:   "payload 1",
         key:       "key 1",
@@ -272,7 +274,7 @@ describe Rdkafka::Consumer do
     it "should close a consumer" do
       consumer.subscribe("consume_test_topic")
       100.times do |i|
-        report = producer.produce(
+        producer.produce(
           topic:     "consume_test_topic",
           payload:   "payload #{i}",
           key:       "key #{i}",
@@ -284,12 +286,35 @@ describe Rdkafka::Consumer do
         consumer.poll(100)
       }.to raise_error(Rdkafka::ClosedConsumerError, /poll/)
     end
+
+    context 'when there are outgoing operations in other threads' do
+      it 'should wait and not crash' do
+        times = []
+
+        # Run a long running poll
+        thread = Thread.new do
+          times << Time.now
+          consumer.subscribe("empty_test_topic")
+          times << Time.now
+          consumer.poll(1_000)
+          times << Time.now
+        end
+
+        # Make sure it starts before we close
+        sleep(0.1)
+        consumer.close
+        close_time = Time.now
+        thread.join
+
+        times.each { |op_time| expect(op_time).to be < close_time }
+      end
+    end
   end
 
   describe "#commit, #committed and #store_offset" do
     # Make sure there's a stored offset
     let!(:report) do
-      report = producer.produce(
+      producer.produce(
         topic:     "consume_test_topic",
         payload:   "payload 1",
         key:       "key 1",
@@ -593,7 +618,7 @@ describe Rdkafka::Consumer do
   end
 
   describe "#poll with headers" do
-    it "should return message with headers" do
+    it "should return message with headers using string keys (when produced with symbol keys)" do
       report = producer.produce(
         topic:     "consume_test_topic",
         key:       "key headers",
@@ -603,7 +628,20 @@ describe Rdkafka::Consumer do
       message = wait_for_message(topic: "consume_test_topic", consumer: consumer, delivery_report: report)
       expect(message).to be
       expect(message.key).to eq('key headers')
-      expect(message.headers).to include(foo: 'bar')
+      expect(message.headers).to include('foo' => 'bar')
+    end
+
+    it "should return message with headers using string keys (when produced with string keys)" do
+      report = producer.produce(
+        topic:     "consume_test_topic",
+        key:       "key headers",
+        headers:   { 'foo' => 'bar' }
+      ).wait
+
+      message = wait_for_message(topic: "consume_test_topic", consumer: consumer, delivery_report: report)
+      expect(message).to be
+      expect(message.key).to eq('key headers')
+      expect(message.headers).to include('foo' => 'bar')
     end
 
     it "should return message with no headers" do
@@ -698,7 +736,7 @@ describe Rdkafka::Consumer do
       n.times do |i|
         handles << producer.produce(
           topic:     topic_name,
-          payload:   Time.new.to_f.to_s,
+          payload:   i % 10 == 0 ? nil : Time.new.to_f.to_s,
           key:       i.to_s,
           partition: 0
         )
@@ -723,7 +761,8 @@ describe Rdkafka::Consumer do
       #
       # This is, in effect, an integration test and the subsequent specs are
       # unit tests.
-      create_topic_handle = rdkafka_config.admin.create_topic(topic_name, 1, 1)
+      admin = rdkafka_config.admin
+      create_topic_handle = admin.create_topic(topic_name, 1, 1)
       create_topic_handle.wait(max_wait_timeout: 15.0)
       consumer.subscribe(topic_name)
       produce_n 42
@@ -736,6 +775,7 @@ describe Rdkafka::Consumer do
       expect(all_yields.flatten.size).to eq 42
       expect(all_yields.size).to be > 4
       expect(all_yields.flatten.map(&:key)).to eq (0..41).map { |x| x.to_s }
+      admin.close
     end
 
     it "should batch poll results and yield arrays of messages" do
@@ -778,13 +818,15 @@ describe Rdkafka::Consumer do
     end
 
     it "should yield [] if nothing is received before the timeout" do
-      create_topic_handle = rdkafka_config.admin.create_topic(topic_name, 1, 1)
+      admin = rdkafka_config.admin
+      create_topic_handle = admin.create_topic(topic_name, 1, 1)
       create_topic_handle.wait(max_wait_timeout: 15.0)
       consumer.subscribe(topic_name)
       consumer.each_batch do |batch|
         expect(batch).to eq([])
         break
       end
+      admin.close
     end
 
     it "should yield batchs of max_items in size if messages are already fetched" do
@@ -831,7 +873,6 @@ describe Rdkafka::Consumer do
         )
         consumer = config.consumer
         consumer.subscribe(topic_name)
-        loop_count = 0
         batches_yielded = []
         exceptions_yielded = []
         each_batch_iterations = 0
@@ -862,6 +903,7 @@ describe Rdkafka::Consumer do
         expect(batches_yielded.first.size).to eq 2
         expect(exceptions_yielded.flatten.size).to eq 1
         expect(exceptions_yielded.flatten.first).to be_instance_of(Rdkafka::RdkafkaError)
+        consumer.close
       end
     end
 
@@ -875,7 +917,6 @@ describe Rdkafka::Consumer do
         )
         consumer = config.consumer
         consumer.subscribe(topic_name)
-        loop_count = 0
         batches_yielded = []
         exceptions_yielded = []
         each_batch_iterations = 0
@@ -904,6 +945,7 @@ describe Rdkafka::Consumer do
         expect(each_batch_iterations).to eq 0
         expect(batches_yielded.size).to eq 0
         expect(exceptions_yielded.size).to eq 0
+        consumer.close
       end
     end
   end
@@ -918,11 +960,11 @@ describe Rdkafka::Consumer do
     context "with a working listener" do
       let(:listener) do
         Struct.new(:queue) do
-          def on_partitions_assigned(consumer, list)
+          def on_partitions_assigned(list)
             collect(:assign, list)
           end
 
-          def on_partitions_revoked(consumer, list)
+          def on_partitions_revoked(list)
             collect(:revoke, list)
           end
 
@@ -946,12 +988,12 @@ describe Rdkafka::Consumer do
     context "with a broken listener" do
       let(:listener) do
         Struct.new(:queue) do
-          def on_partitions_assigned(consumer, list)
+          def on_partitions_assigned(list)
             queue << :assigned
             raise 'boom'
           end
 
-          def on_partitions_revoked(consumer, list)
+          def on_partitions_revoked(list)
             queue << :revoked
             raise 'boom'
           end
@@ -964,18 +1006,6 @@ describe Rdkafka::Consumer do
         expect(listener.queue).to eq([:assigned, :revoked])
       end
     end
-
-    def notify_listener(listener)
-      # 1. subscribe and poll
-      consumer.subscribe("consume_test_topic")
-      wait_for_assignment(consumer)
-      consumer.poll(100)
-
-      # 2. unsubscribe
-      consumer.unsubscribe
-      wait_for_unassignment(consumer)
-      consumer.close
-    end
   end
 
   context "methods that should not be called after a consumer has been closed" do
@@ -1007,4 +1037,70 @@ describe Rdkafka::Consumer do
       end
     end
   end
+
+  it "provides a finalizer that closes the native kafka client" do
+    expect(consumer.closed?).to eq(false)
+
+    consumer.finalizer.call("some-ignored-object-id")
+
+    expect(consumer.closed?).to eq(true)
+  end
+
+  context "when the rebalance protocol is cooperative" do
+    let(:consumer) do
+      config = rdkafka_consumer_config(
+        {
+          :"partition.assignment.strategy" => "cooperative-sticky",
+          :"debug" => "consumer",
+        }
+      )
+      config.consumer_rebalance_listener = listener
+      config.consumer
+    end
+
+    let(:listener) do
+      Struct.new(:queue) do
+        def on_partitions_assigned(list)
+          collect(:assign, list)
+        end
+
+        def on_partitions_revoked(list)
+          collect(:revoke, list)
+        end
+
+        def collect(name, list)
+          partitions = list.to_h.map { |key, values| [key, values.map(&:partition)] }.flatten
+          queue << ([name] + partitions)
+        end
+      end.new([])
+    end
+
+    it "should be able to assign and unassign partitions using the cooperative partition assignment APIs" do
+      notify_listener(listener) do
+        handles = []
+        10.times do
+          handles << producer.produce(
+            topic:     "consume_test_topic",
+            payload:   "payload 1",
+            key:       "key 1",
+            partition: 0
+          )
+        end
+        handles.each(&:wait)
+
+        consumer.subscribe("consume_test_topic")
+        # Check the first 10 messages. Then close the consumer, which
+        # should break the each loop.
+        consumer.each_with_index do |message, i|
+          expect(message).to be_a Rdkafka::Consumer::Message
+          break if i == 10
+        end
+      end
+
+      expect(listener.queue).to eq([
+        [:assign, "consume_test_topic", 0, 1, 2],
+        [:revoke, "consume_test_topic", 0, 1, 2]
+      ])
+    end
+  end
 end