karafka-rdkafka 0.12.3 → 0.13.0

data/lib/rdkafka/producer.rb

@@ -1,8 +1,15 @@
+# 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.
   class Producer
+    # Cache partitions count for 30 seconds
+    PARTITIONS_COUNT_TTL = 30
+
+    private_constant :PARTITIONS_COUNT_TTL
+
     # @private
     # Returns the current delivery callback, by default this is nil.
     #
@@ -16,12 +23,32 @@ module Rdkafka
     attr_reader :delivery_callback_arity
 
     # @private
-    def initialize(client, partitioner_name)
-      @client = client
+    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(self, client.finalizer)
+      # Makes sure, that native kafka gets closed before it gets GCed by Ruby
+      ObjectSpace.define_finalizer(self, native_kafka.finalizer)
+
+      @_partitions_count_cache = Hash.new do |cache, topic|
+        topic_metadata = nil
+
+        @native_kafka.with_inner do |inner|
+          topic_metadata = ::Rdkafka::Metadata.new(inner, topic).topics&.first
+        end
+
+        cache[topic] = [
+          monotonic_now,
+          topic_metadata ? topic_metadata[:partition_count] : nil
+        ]
+      end
+    end
+
+    # @return [String] producer name
+    def name
+      @name ||= @native_kafka.with_inner do |inner|
+        ::Rdkafka::Bindings.rd_kafka_name(inner)
+      end
     end
 
     # Set a callback that will be called every time a message is successfully produced.
@@ -38,9 +65,26 @@ module Rdkafka
 
     # Close this producer and wait for the internal poll queue to empty.
     def close
+      return if closed?
       ObjectSpace.undefine_finalizer(self)
+      @native_kafka.close
+    end
 
-      @client.close
+    # 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__)
+
+      @native_kafka.with_inner do |inner|
+        Rdkafka::Bindings.rd_kafka_flush(inner, timeout_ms)
+      end
     end
 
     # Partition count for a given topic.
@@ -50,9 +94,20 @@ module Rdkafka
     #
     # @return partition count [Integer,nil]
     #
+    # We cache the partition count for a given topic for given time
+    # This prevents us in case someone uses `partition_key` from querying for the count with
+    # each message. Instead we query once every 30 seconds at most
+    #
+    # @param topic [String] topic name
+    # @return [Integer] partition count for a given topic
     def partition_count(topic)
       closed_producer_check(__method__)
-      Rdkafka::Metadata.new(@client.native, topic).topics&.first[:partition_count]
+
+      @_partitions_count_cache.delete_if do |_, cached|
+        monotonic_now - cached.first > PARTITIONS_COUNT_TTL
+      end
+
+      @_partitions_count_cache[topic].last
     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.
@@ -143,10 +198,12 @@ module Rdkafka
       args << :int << Rdkafka::Bindings::RD_KAFKA_VTYPE_END
 
       # Produce the message
-      response = Rdkafka::Bindings.rd_kafka_producev(
-        @client.native,
-        *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
@@ -157,7 +214,6 @@ module Rdkafka
       delivery_handle
     end
 
-    # @private
     def call_delivery_callback(delivery_report, delivery_handle)
       return unless @delivery_callback
 
@@ -171,8 +227,15 @@ module Rdkafka
       callback.method(:call).arity
     end
 
+    private
+
+    def monotonic_now
+      # needed because Time.now can go backwards
+      Process.clock_gettime(Process::CLOCK_MONOTONIC)
+    end
+
     def closed_producer_check(method)
-      raise Rdkafka::ClosedProducerError.new(method) if @client.closed?
+      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.12.3"
+  VERSION = "0.13.0"
   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"
@@ -18,7 +20,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/client"
 require "rdkafka/producer/delivery_handle"
 require "rdkafka/producer/delivery_report"

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'
 

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'
@@ -9,6 +11,10 @@ describe Rdkafka::Consumer do
   after { consumer.close }
   after { producer.close }
 
+  describe '#name' do
+    it { expect(consumer.name).to include('rdkafka#consumer-') }
+  end
+
   describe "#subscribe, #unsubscribe and #subscription" do
     it "should subscribe, unsubscribe and return the subscription" do
       expect(consumer.subscription).to be_empty
@@ -53,7 +59,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 }
@@ -268,6 +274,28 @@ describe Rdkafka::Consumer do
     end
   end
 
+  describe '#assignment_lost?' do
+    it "should not return true as we do have an assignment" do
+      consumer.subscribe("consume_test_topic")
+      expected_subscription = Rdkafka::Consumer::TopicPartitionList.new.tap do |list|
+        list.add_topic("consume_test_topic")
+      end
+
+      expect(consumer.assignment_lost?).to eq false
+      consumer.unsubscribe
+    end
+
+    it "should not return true after voluntary unsubscribing" do
+      consumer.subscribe("consume_test_topic")
+      expected_subscription = Rdkafka::Consumer::TopicPartitionList.new.tap do |list|
+        list.add_topic("consume_test_topic")
+      end
+
+      consumer.unsubscribe
+      expect(consumer.assignment_lost?).to eq false
+    end
+  end
+
   describe "#close" do
     it "should close a consumer" do
       consumer.subscribe("consume_test_topic")
@@ -593,7 +621,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 +631,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 +739,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 +764,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 +778,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 +821,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
@@ -861,6 +906,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
 
@@ -902,6 +948,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
@@ -916,11 +963,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
 
@@ -944,12 +991,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
@@ -962,18 +1009,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
@@ -993,6 +1028,7 @@ describe Rdkafka::Consumer do
         :assignment              => nil,
         :committed               => [],
         :query_watermark_offsets => [ nil, nil ],
+        :assignment_lost?        => []
     }.each do |method, args|
       it "raises an exception if #{method} is called" do
         expect {
@@ -1005,4 +1041,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

data/spec/rdkafka/error_spec.rb

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

data/spec/rdkafka/metadata_spec.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 require "spec_helper"
 require "securerandom"