rdkafka 0.4.2 → 0.8.0

data/lib/rdkafka/consumer/headers.rb

@@ -0,0 +1,63 @@
+module Rdkafka
+  class Consumer
+    # A message headers
+    class Headers
+      # Reads a native kafka's message header into ruby's hash
+      #
+      # @return [Hash<String, String>] a message headers
+      #
+      # @raise [Rdkafka::RdkafkaError] when fail to read headers
+      #
+      # @private
+      def self.from_native(native_message)
+        headers_ptrptr = FFI::MemoryPointer.new(:pointer)
+        err = Rdkafka::Bindings.rd_kafka_message_headers(native_message, headers_ptrptr)
+
+        if err == Rdkafka::Bindings::RD_KAFKA_RESP_ERR__NOENT
+          return {}
+        elsif err != Rdkafka::Bindings::RD_KAFKA_RESP_ERR_NO_ERROR
+          raise Rdkafka::RdkafkaError.new(err, "Error reading message headers")
+        end
+
+        headers_ptr = headers_ptrptr.read_pointer
+
+        name_ptrptr = FFI::MemoryPointer.new(:pointer)
+        value_ptrptr = FFI::MemoryPointer.new(:pointer)
+        size_ptr = Rdkafka::Bindings::SizePtr.new
+        headers = {}
+
+        idx = 0
+        loop do
+          err = Rdkafka::Bindings.rd_kafka_header_get_all(
+            headers_ptr,
+            idx,
+            name_ptrptr,
+            value_ptrptr,
+            size_ptr
+          )
+
+          if err == Rdkafka::Bindings::RD_KAFKA_RESP_ERR__NOENT
+            break
+          elsif err != Rdkafka::Bindings::RD_KAFKA_RESP_ERR_NO_ERROR
+            raise Rdkafka::RdkafkaError.new(err, "Error reading a message header at index #{idx}")
+          end
+
+          name_ptr = name_ptrptr.read_pointer
+          name = name_ptr.respond_to?(:read_string_to_null) ? name_ptr.read_string_to_null : name_ptr.read_string
+
+          size = size_ptr[:value]
+
+          value_ptr = value_ptrptr.read_pointer
+
+          value = value_ptr.read_string(size)
+
+          headers[name.to_sym] = value
+
+          idx += 1
+        end
+
+        headers
+      end
+    end
+  end
+end

data/lib/rdkafka/consumer/message.rb

@@ -26,6 +26,9 @@ module Rdkafka
       # @return [Time, nil]
       attr_reader :timestamp
 
+      # @return [Hash<String, String>] a message headers
+      attr_reader :headers
+
       # @private
       def initialize(native_message)
         # Set topic
@@ -54,12 +57,16 @@ module Rdkafka
                      else
                        nil
                      end
+
+        @headers = Headers.from_native(native_message)
       end
 
       # Human readable representation of this message.
       # @return [String]
       def to_s
-        "<Message in '#{topic}' with key '#{truncate(key)}', payload '#{truncate(payload)}', partition #{partition}, offset #{offset}, timestamp #{timestamp}>"
+        is_headers = @headers.empty? ? "" : ", headers #{headers.size}"
+
+        "<Message in '#{topic}' with key '#{truncate(key)}', payload '#{truncate(payload)}', partition #{partition}, offset #{offset}, timestamp #{timestamp}#{is_headers}>"
       end
 
       def truncate(string)
@@ -69,6 +76,9 @@ module Rdkafka
           string
         end
       end
+
+      private
+
     end
   end
 end

data/lib/rdkafka/consumer/partition.rb

@@ -10,20 +10,25 @@ module Rdkafka
       # @return [Integer, nil]
       attr_reader :offset
 
+      # Partition's error code
+      # @return [Integer]
+      attr_reader :err
+
       # @private
-      def initialize(partition, offset)
+      def initialize(partition, offset, err = 0)
         @partition = partition
         @offset = offset
+        @err = err
       end
 
       # Human readable representation of this partition.
       # @return [String]
       def to_s
-        if offset.nil?
-          "<Partition #{partition} without offset>"
-        else
-          "<Partition #{partition} with offset #{offset}>"
-        end
+        message = "<Partition #{partition}"
+        message += " offset=#{offset}" if offset
+        message += " err=#{err}" if err != 0
+        message += ">"
+        message
       end
 
       # Human readable representation of this partition.

data/lib/rdkafka/consumer/topic_partition_list.rb

@@ -4,7 +4,7 @@ module Rdkafka
     class TopicPartitionList
       # Create a topic partition list.
       #
-      # @param data [Hash<String => [nil,Partition]>] The topic and partion data or nil to create an empty list
+      # @param data [Hash{String => nil,Partition}] The topic and partition data or nil to create an empty list
       #
       # @return [TopicPartitionList]
       def initialize(data=nil)
@@ -54,7 +54,7 @@ module Rdkafka
           if partitions.is_a? Integer
             partitions = (0..partitions - 1)
           end
-          @data[topic.to_s] = partitions.map { |p| Partition.new(p, nil) }
+          @data[topic.to_s] = partitions.map { |p| Partition.new(p, nil, 0) }
         end
       end
 
@@ -71,7 +71,7 @@ module Rdkafka
 
       # Return a `Hash` with the topics as keys and and an array of partition information as the value if present.
       #
-      # @return [Hash<String, [Array<Partition>, nil]>]
+      # @return [Hash{String => Array<Partition>,nil}]
       def to_h
         @data
       end
@@ -88,7 +88,7 @@ module Rdkafka
 
       # Create a new topic partition list based of a native one.
       #
-      # @param pointer [FFI::Pointer] Optional pointer to an existing native list. Its contents will be copied and afterwards it will be destroyed.
+      # @param pointer [FFI::Pointer] Optional pointer to an existing native list. Its contents will be copied.
       #
       # @return [TopicPartitionList]
       #
@@ -106,12 +106,12 @@ module Rdkafka
             data[elem[:topic]] = nil
           else
             partitions = data[elem[:topic]] || []
-            offset = if elem[:offset] == -1001
+            offset = if elem[:offset] == Rdkafka::Bindings::RD_KAFKA_OFFSET_INVALID
                        nil
                      else
                        elem[:offset]
                      end
-            partition = Partition.new(elem[:partition], offset)
+            partition = Partition.new(elem[:partition], offset, elem[:err])
             partitions.push(partition)
             data[elem[:topic]] = partitions
           end
@@ -119,42 +119,45 @@ module Rdkafka
 
         # Return the created object
         TopicPartitionList.new(data)
-      ensure
-        # Destroy the tpl
-        Rdkafka::Bindings.rd_kafka_topic_partition_list_destroy(pointer)
       end
 
-      # Create a native tpl with the contents of this object added
+      # Create a native tpl with the contents of this object added.
       #
+      # The pointer will be cleaned by `rd_kafka_topic_partition_list_destroy` when GC releases it.
+      #
+      # @return [FFI::Pointer]
       # @private
       def to_native_tpl
-        Rdkafka::Bindings.rd_kafka_topic_partition_list_new(count).tap do |tpl|
-          @data.each do |topic, partitions|
-            if partitions
-              partitions.each do |p|
-                Rdkafka::Bindings.rd_kafka_topic_partition_list_add(
-                  tpl,
-                  topic,
-                  p.partition
-                )
-                if p.offset
-                  Rdkafka::Bindings.rd_kafka_topic_partition_list_set_offset(
-                    tpl,
-                    topic,
-                    p.partition,
-                    p.offset
-                  )
-                end
-              end
-            else
+        tpl = Rdkafka::Bindings.rd_kafka_topic_partition_list_new(count)
+
+        @data.each do |topic, partitions|
+          if partitions
+            partitions.each do |p|
               Rdkafka::Bindings.rd_kafka_topic_partition_list_add(
                 tpl,
                 topic,
-                -1
+                p.partition
               )
+
+              if p.offset
+                Rdkafka::Bindings.rd_kafka_topic_partition_list_set_offset(
+                  tpl,
+                  topic,
+                  p.partition,
+                  p.offset
+                )
+              end
             end
+          else
+            Rdkafka::Bindings.rd_kafka_topic_partition_list_add(
+              tpl,
+              topic,
+              -1
+            )
           end
         end
+
+        tpl
       end
     end
   end

data/lib/rdkafka/error.rb

@@ -39,5 +39,22 @@ module Rdkafka
     def is_partition_eof?
       code == :partition_eof
     end
+
+    # Error comparison
+    def ==(another_error)
+       another_error.is_a?(self.class) && (self.to_s == another_error.to_s)
+    end
+  end
+
+  # Error with topic partition list returned by the underlying rdkafka library.
+  class RdkafkaTopicPartitionListError < RdkafkaError
+    # @return [TopicPartitionList]
+    attr_reader :topic_partition_list
+
+    # @private
+    def initialize(response, topic_partition_list, message_prefix=nil)
+      super(response, message_prefix)
+      @topic_partition_list = topic_partition_list
+    end
   end
 end

data/lib/rdkafka/metadata.rb

@@ -0,0 +1,91 @@
+module Rdkafka
+  class Metadata
+    attr_reader :brokers, :topics
+
+    def initialize(native_client, topic_name = nil)
+      native_topic = if topic_name
+        Rdkafka::Bindings.rd_kafka_topic_new(native_client, topic_name, nil)
+      end
+
+      ptr = FFI::MemoryPointer.new(:pointer)
+
+      # Retrieve metadata flag is 0/1 for single/multiple topics.
+      topic_flag = topic_name ? 1 : 0
+
+      # Retrieve the Metadata
+      result = Rdkafka::Bindings.rd_kafka_metadata(native_client, topic_flag, native_topic, ptr, 250)
+
+      # Error Handling
+      Rdkafka::Error.new(result) unless result.zero?
+
+      metadata_from_native(ptr.read_pointer)
+    ensure
+      Rdkafka::Bindings.rd_kafka_topic_destroy(native_topic) if topic_name
+      Rdkafka::Bindings.rd_kafka_metadata_destroy(ptr.read_pointer)
+    end
+
+    private
+
+    def metadata_from_native(ptr)
+      metadata = Metadata.new(ptr)
+      @brokers = Array.new(metadata[:brokers_count]) do |i|
+        BrokerMetadata.new(metadata[:brokers_metadata] + (i * BrokerMetadata.size)).to_h
+      end
+
+      @topics = Array.new(metadata[:topics_count]) do |i|
+        topic = TopicMetadata.new(metadata[:topics_metadata] + (i * TopicMetadata.size))
+        Rdkafka::Error.new(topic[:rd_kafka_resp_err]) unless topic[:rd_kafka_resp_err].zero?
+
+        partitions = Array.new(topic[:partition_count]) do |j|
+          partition = PartitionMetadata.new(topic[:partitions_metadata] + (j * PartitionMetadata.size))
+          Rdkafka::Error.new(partition[:rd_kafka_resp_err]) unless partition[:rd_kafka_resp_err].zero?
+          partition.to_h
+        end
+        topic.to_h.merge!(partitions: partitions)
+      end
+    end
+
+    class CustomFFIStruct < FFI::Struct
+      def to_h
+        members.each_with_object({}) do |mem, hsh|
+          val = self.[](mem)
+          next if val.is_a?(FFI::Pointer) || mem == :rd_kafka_resp_err
+
+          hsh[mem] = self.[](mem)
+        end
+      end
+    end
+
+    class Metadata < CustomFFIStruct
+      layout :brokers_count, :int,
+             :brokers_metadata, :pointer,
+             :topics_count, :int,
+             :topics_metadata, :pointer,
+             :broker_id, :int32,
+             :broker_name, :string
+    end
+
+    class BrokerMetadata < CustomFFIStruct
+      layout :broker_id, :int32,
+             :broker_name, :string,
+             :broker_port, :int
+    end
+
+    class TopicMetadata < CustomFFIStruct
+      layout :topic_name, :string,
+             :partition_count, :int,
+             :partitions_metadata, :pointer,
+             :rd_kafka_resp_err, :int
+    end
+
+    class PartitionMetadata < CustomFFIStruct
+      layout :partition_id, :int32,
+             :rd_kafka_resp_err, :int,
+             :leader, :int32,
+             :replica_count, :int,
+             :replicas, :pointer,
+             :in_sync_replica_brokers, :int,
+             :isrs, :pointer
+    end
+  end
+end

data/lib/rdkafka/producer.rb

@@ -2,7 +2,10 @@ module Rdkafka
   # A producer for Kafka messages. To create a producer set up a {Config} and call {Config#producer producer} on that.
   class Producer
     # @private
-    @delivery_callback = nil
+    # Returns the current delivery callback, by default this is nil.
+    #
+    # @return [Proc, nil]
+    attr_reader :delivery_callback
 
     # @private
     def initialize(native_kafka)
@@ -32,36 +35,45 @@ module Rdkafka
       @delivery_callback = callback
     end
 
-    # Returns the current delivery callback, by default this is nil.
-    #
-    # @return [Proc, nil]
-    def delivery_callback
-      @delivery_callback
-    end
-
     # Close this producer and wait for the internal poll queue to empty.
     def close
+      return unless @native_kafka
+
       # 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
+    end
+
+    # Partition count for a given topic.
+    # NOTE: If 'allow.auto.create.topics' is set to true in the broker, the topic will be auto-created after returning nil.
+    #
+    # @param topic [String] The topic name.
+    #
+    # @return partition count [Integer,nil]
+    #
+    def partition_count(topic)
+      Rdkafka::Metadata.new(@native_kafka, topic).topics&.select { |x| x[:topic_name] == topic }&.dig(0, :partition_count)
     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.
     #
     # When no partition is specified the underlying Kafka library picks a partition based on the key. If no key is specified, a random partition will be used.
-    # When a timestamp is provided this is used instead of the autogenerated timestamp.
+    # When a timestamp is provided this is used instead of the auto-generated timestamp.
     #
     # @param topic [String] The topic to produce to
     # @param payload [String,nil] The message's payload
     # @param key [String] The message's key
     # @param partition [Integer,nil] Optional partition to produce to
     # @param timestamp [Time,Integer,nil] Optional timestamp of this message. Integer timestamp is in milliseconds since Jan 1 1970.
+    # @param headers [Hash<String,String>] Optional message headers
     #
     # @raise [RdkafkaError] When adding the message to rdkafka's queue failed
     #
     # @return [DeliveryHandle] Delivery handle that can be used to wait for the result of producing this message
-    def produce(topic:, payload: nil, key: nil, partition: nil, timestamp: nil)
+    def produce(topic:, payload: nil, key: nil, partition: nil, partition_key: nil, timestamp: nil, headers: nil)
       # Start by checking and converting the input
 
       # Get payload length
@@ -78,9 +90,15 @@ module Rdkafka
                    key.bytesize
                  end
 
-      # If partition is nil use -1 to let Kafka set the partition based
-      # on the key/randomly if there is no key
-      partition = -1 if partition.nil?
+      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
+      end
+
+      # If partition is nil, use -1 to let librdafka set the partition randomly or
+      # based on the key when present.
+      partition ||= -1
 
       # If timestamp is nil use 0 and let Kafka set one. If an integer or time
       # use it.
@@ -101,9 +119,7 @@ module Rdkafka
       delivery_handle[:offset] = -1
       DeliveryHandle.register(delivery_handle.to_ptr.address, delivery_handle)
 
-      # Produce the message
-      response = Rdkafka::Bindings.rd_kafka_producev(
-        @native_kafka,
+      args = [
         :int, Rdkafka::Bindings::RD_KAFKA_VTYPE_TOPIC, :string, topic,
         :int, Rdkafka::Bindings::RD_KAFKA_VTYPE_MSGFLAGS, :int, Rdkafka::Bindings::RD_KAFKA_MSG_F_COPY,
         :int, Rdkafka::Bindings::RD_KAFKA_VTYPE_VALUE, :buffer_in, payload, :size_t, payload_size,
@@ -111,10 +127,30 @@ module Rdkafka
         :int, Rdkafka::Bindings::RD_KAFKA_VTYPE_PARTITION, :int32, partition,
         :int, Rdkafka::Bindings::RD_KAFKA_VTYPE_TIMESTAMP, :int64, raw_timestamp,
         :int, Rdkafka::Bindings::RD_KAFKA_VTYPE_OPAQUE, :pointer, delivery_handle,
-        :int, Rdkafka::Bindings::RD_KAFKA_VTYPE_END
+      ]
+
+      if headers
+        headers.each do |key0, value0|
+          key = key0.to_s
+          value = value0.to_s
+          args += [
+            :int, Rdkafka::Bindings::RD_KAFKA_VTYPE_HEADER,
+            :string, key,
+            :pointer, value,
+            :size_t, value.bytes.size
+          ]
+        end
+      end
+
+      args += [:int, Rdkafka::Bindings::RD_KAFKA_VTYPE_END]
+
+      # Produce the message
+      response = Rdkafka::Bindings.rd_kafka_producev(
+        @native_kafka,
+        *args
       )
 
-      # Raise error if the produce call was not successfull
+      # Raise error if the produce call was not successful
       if response != 0
         DeliveryHandle.remove(delivery_handle.to_ptr.address)
         raise RdkafkaError.new(response)