ruby-kafka 1.2.0 → 1.5.0

data/lib/kafka/cluster.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require "kafka/broker_pool"
+require "resolv"
 require "set"
 
 module Kafka
@@ -18,7 +19,8 @@ module Kafka
     # @param seed_brokers [Array<URI>]
     # @param broker_pool [Kafka::BrokerPool]
     # @param logger [Logger]
-    def initialize(seed_brokers:, broker_pool:, logger:)
+    # @param resolve_seed_brokers [Boolean] See {Kafka::Client#initialize}
+    def initialize(seed_brokers:, broker_pool:, logger:, resolve_seed_brokers: false)
       if seed_brokers.empty?
         raise ArgumentError, "At least one seed broker must be configured"
       end
@@ -26,6 +28,7 @@ module Kafka
       @logger = TaggedLogger.new(logger)
       @seed_brokers = seed_brokers
       @broker_pool = broker_pool
+      @resolve_seed_brokers = resolve_seed_brokers
       @cluster_info = nil
       @stale = true
 
@@ -117,7 +120,7 @@ module Kafka
 
     # Finds the broker acting as the coordinator of the given group.
     #
-    # @param group_id: [String]
+    # @param group_id [String]
     # @return [Broker] the broker that's currently coordinator.
     def get_group_coordinator(group_id:)
       @logger.debug "Getting group coordinator for `#{group_id}`"
@@ -127,7 +130,7 @@ module Kafka
 
     # Finds the broker acting as the coordinator of the given transaction.
     #
-    # @param transactional_id: [String]
+    # @param transactional_id [String]
     # @return [Broker] the broker that's currently coordinator.
     def get_transaction_coordinator(transactional_id:)
       @logger.debug "Getting transaction coordinator for `#{transactional_id}`"
@@ -418,32 +421,35 @@ module Kafka
     # @return [Protocol::MetadataResponse] the cluster metadata.
     def fetch_cluster_info
       errors = []
-
       @seed_brokers.shuffle.each do |node|
-        @logger.info "Fetching cluster metadata from #{node}"
-
-        begin
-          broker = @broker_pool.connect(node.hostname, node.port)
-          cluster_info = broker.fetch_metadata(topics: @target_topics)
-
-          if cluster_info.brokers.empty?
-            @logger.error "No brokers in cluster"
-          else
-            @logger.info "Discovered cluster metadata; nodes: #{cluster_info.brokers.join(', ')}"
-
-            @stale = false
-
-            return cluster_info
+        (@resolve_seed_brokers ? Resolv.getaddresses(node.hostname).shuffle : [node.hostname]).each do |hostname_or_ip|
+          node_info = node.to_s
+          node_info << " (#{hostname_or_ip})" if node.hostname != hostname_or_ip
+          @logger.info "Fetching cluster metadata from #{node_info}"
+
+          begin
+            broker = @broker_pool.connect(hostname_or_ip, node.port)
+            cluster_info = broker.fetch_metadata(topics: @target_topics)
+
+            if cluster_info.brokers.empty?
+              @logger.error "No brokers in cluster"
+            else
+              @logger.info "Discovered cluster metadata; nodes: #{cluster_info.brokers.join(', ')}"
+
+              @stale = false
+
+              return cluster_info
+            end
+          rescue Error => e
+            @logger.error "Failed to fetch metadata from #{node_info}: #{e}"
+            errors << [node_info, e]
+          ensure
+            broker.disconnect unless broker.nil?
           end
-        rescue Error => e
-          @logger.error "Failed to fetch metadata from #{node}: #{e}"
-          errors << [node, e]
-        ensure
-          broker.disconnect unless broker.nil?
         end
       end
 
-      error_description = errors.map {|node, exception| "- #{node}: #{exception}" }.join("\n")
+      error_description = errors.map {|node_info, exception| "- #{node_info}: #{exception}" }.join("\n")
 
       raise ConnectionError, "Could not connect to any of the seed brokers:\n#{error_description}"
     end

data/lib/kafka/consumer_group/assignor.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+require "kafka/protocol/member_assignment"
+
+module Kafka
+  class ConsumerGroup
+
+    # A consumer group partition assignor
+    class Assignor
+      Partition = Struct.new(:topic, :partition_id)
+
+      # @param cluster [Kafka::Cluster]
+      # @param strategy [Object] an object that implements #protocol_type,
+      #   #user_data, and #assign.
+      def initialize(cluster:, strategy:)
+        @cluster = cluster
+        @strategy = strategy
+      end
+
+      def protocol_name
+        @strategy.respond_to?(:protocol_name) ? @strategy.protocol_name : @strategy.class.to_s
+      end
+
+      def user_data
+        @strategy.user_data if @strategy.respond_to?(:user_data)
+      end
+
+      # Assign the topic partitions to the group members.
+      #
+      # @param members [Hash<String, Kafka::Protocol::JoinGroupResponse::Metadata>] a hash
+      #   mapping member ids to metadata.
+      # @param topics [Array<String>] topics
+      # @return [Hash<String, Kafka::Protocol::MemberAssignment>] a hash mapping member
+      #   ids to assignments.
+      def assign(members:, topics:)
+        topic_partitions = topics.flat_map do |topic|
+          begin
+            partition_ids = @cluster.partitions_for(topic).map(&:partition_id)
+          rescue UnknownTopicOrPartition
+            raise UnknownTopicOrPartition, "unknown topic #{topic}"
+          end
+          partition_ids.map {|partition_id| Partition.new(topic, partition_id) }
+        end
+
+        group_assignment = {}
+
+        members.each_key do |member_id|
+          group_assignment[member_id] = Protocol::MemberAssignment.new
+        end
+        @strategy.call(cluster: @cluster, members: members, partitions: topic_partitions).each do |member_id, partitions|
+          Array(partitions).each do |partition|
+            group_assignment[member_id].assign(partition.topic, [partition.partition_id])
+          end
+        end
+
+        group_assignment
+      rescue Kafka::LeaderNotAvailable
+        sleep 1
+        retry
+      end
+    end
+  end
+end

data/lib/kafka/consumer_group.rb

@@ -1,13 +1,14 @@
 # frozen_string_literal: true
 
 require "set"
+require "kafka/consumer_group/assignor"
 require "kafka/round_robin_assignment_strategy"
 
 module Kafka
   class ConsumerGroup
     attr_reader :assigned_partitions, :generation_id, :group_id
 
-    def initialize(cluster:, logger:, group_id:, session_timeout:, rebalance_timeout:, retention_time:, instrumenter:)
+    def initialize(cluster:, logger:, group_id:, session_timeout:, rebalance_timeout:, retention_time:, instrumenter:, assignment_strategy:)
       @cluster = cluster
       @logger = TaggedLogger.new(logger)
       @group_id = group_id
@@ -19,7 +20,10 @@ module Kafka
       @members = {}
       @topics = Set.new
       @assigned_partitions = {}
-      @assignment_strategy = RoundRobinAssignmentStrategy.new(cluster: @cluster)
+      @assignor = Assignor.new(
+        cluster: cluster,
+        strategy: assignment_strategy || RoundRobinAssignmentStrategy.new
+      )
       @retention_time = retention_time
     end
 
@@ -113,9 +117,12 @@ module Kafka
 
         Protocol.handle_error(response.error_code)
       end
-    rescue ConnectionError, UnknownMemberId, RebalanceInProgress, IllegalGeneration => e
+    rescue ConnectionError, UnknownMemberId, IllegalGeneration => e
       @logger.error "Error sending heartbeat: #{e}"
       raise HeartbeatError, e
+    rescue RebalanceInProgress => e
+      @logger.warn "Error sending heartbeat: #{e}"
+      raise HeartbeatError, e
     rescue NotCoordinatorForGroup
       @logger.error "Failed to find coordinator for group `#{@group_id}`; retrying..."
       sleep 1
@@ -144,6 +151,8 @@ module Kafka
           rebalance_timeout: @rebalance_timeout,
           member_id: @member_id,
           topics: @topics,
+          protocol_name: @assignor.protocol_name,
+          user_data: @assignor.user_data,
         )
 
         Protocol.handle_error(response.error_code)
@@ -180,9 +189,14 @@ module Kafka
       if group_leader?
         @logger.info "Chosen as leader of group `#{@group_id}`"
 
-        group_assignment = @assignment_strategy.assign(
-          members: @members.keys,
-          topics: @topics,
+        topics = Set.new
+        @members.each do |_member, metadata|
+          metadata.topics.each { |t| topics.add(t) }
+        end
+
+        group_assignment = @assignor.assign(
+          members: @members,
+          topics: topics,
         )
       end
 

data/lib/kafka/crc32_hash.rb

@@ -0,0 +1,15 @@
+# frozen_string_literal: true
+
+require "zlib"
+
+module Kafka
+  class Crc32Hash
+
+    # crc32 is supported natively
+    def load; end
+
+    def hash(value)
+      Zlib.crc32(value)
+    end
+  end
+end

data/lib/kafka/digest.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require "kafka/crc32_hash"
+require "kafka/murmur2_hash"
+
+module Kafka
+  module Digest
+    FUNCTIONS_BY_NAME = {
+      :crc32 => Crc32Hash.new,
+      :murmur2 => Murmur2Hash.new
+    }.freeze
+
+    def self.find_digest(name)
+      digest = FUNCTIONS_BY_NAME.fetch(name) do
+        raise LoadError, "Unknown hash function #{name}"
+      end
+
+      digest.load
+      digest
+    end
+  end
+end

data/lib/kafka/murmur2_hash.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Kafka
+  class Murmur2Hash
+    SEED = [0x9747b28c].pack('L')
+
+    def load
+      require 'digest/murmurhash'
+    rescue LoadError
+      raise LoadError, "using murmur2 hashing requires adding a dependency on the `digest-murmurhash` gem to your Gemfile."
+    end
+
+    def hash(value)
+      ::Digest::MurmurHash2.rawdigest(value, SEED) & 0x7fffffff
+    end
+  end
+end

data/lib/kafka/partitioner.rb

@@ -1,11 +1,16 @@
 # frozen_string_literal: true
 
-require "zlib"
+require "kafka/digest"
 
 module Kafka
 
   # Assigns partitions to messages.
   class Partitioner
+    # @param hash_function [Symbol, nil] the algorithm used to compute a messages
+    #   destination partition. Default is :crc32
+    def initialize(hash_function: nil)
+      @digest = Digest.find_digest(hash_function || :crc32)
+    end
 
     # Assigns a partition number based on a partition key. If no explicit
     # partition key is provided, the message key will be used instead.
@@ -28,7 +33,7 @@ module Kafka
       if key.nil?
         rand(partition_count)
       else
-        Zlib.crc32(key) % partition_count
+        @digest.hash(key) % partition_count
       end
     end
   end

data/lib/kafka/protocol/add_offsets_to_txn_response.rb

@@ -1,3 +1,5 @@
+# frozen_string_literal: true
+
 module Kafka
   module Protocol
     class AddOffsetsToTxnResponse

data/lib/kafka/protocol/encoder.rb

@@ -126,7 +126,7 @@ module Kafka
       # Writes an integer under varints serializing to the IO object.
       # https://developers.google.com/protocol-buffers/docs/encoding#varints
       #
-      # @param string [Integer]
+      # @param int [Integer]
       # @return [nil]
       def write_varint(int)
         int = int << 1

data/lib/kafka/protocol/join_group_request.rb

@@ -7,14 +7,14 @@ module Kafka
     class JoinGroupRequest
       PROTOCOL_TYPE = "consumer"
 
-      def initialize(group_id:, session_timeout:, rebalance_timeout:, member_id:, topics: [])
+      def initialize(group_id:, session_timeout:, rebalance_timeout:, member_id:, topics: [], protocol_name:, user_data: nil)
         @group_id = group_id
         @session_timeout = session_timeout * 1000 # Kafka wants ms.
         @rebalance_timeout = rebalance_timeout * 1000 # Kafka wants ms.
         @member_id = member_id || ""
         @protocol_type = PROTOCOL_TYPE
         @group_protocols = {
-          "roundrobin" => ConsumerGroupProtocol.new(topics: topics),
+          protocol_name => ConsumerGroupProtocol.new(topics: topics, user_data: user_data),
         }
       end
 

data/lib/kafka/protocol/join_group_response.rb

@@ -3,6 +3,8 @@
 module Kafka
   module Protocol
     class JoinGroupResponse
+      Metadata = Struct.new(:version, :topics, :user_data)
+
       attr_reader :error_code
 
       attr_reader :generation_id, :group_protocol
@@ -25,7 +27,13 @@ module Kafka
           group_protocol: decoder.string,
           leader_id: decoder.string,
           member_id: decoder.string,
-          members: Hash[decoder.array { [decoder.string, decoder.bytes] }],
+          members: Hash[
+            decoder.array do
+              member_id = decoder.string
+              d = Decoder.from_string(decoder.bytes)
+              [member_id, Metadata.new(d.int16, d.array { d.string }, d.bytes)]
+            end
+          ],
         )
       end
     end

data/lib/kafka/protocol/record_batch.rb

@@ -77,7 +77,7 @@ module Kafka
         record_batch_encoder.write_int8(MAGIC_BYTE)
 
         body = encode_record_batch_body
-        crc = Digest::CRC32c.checksum(body)
+        crc = ::Digest::CRC32c.checksum(body)
 
         record_batch_encoder.write_int32(crc)
         record_batch_encoder.write(body)
@@ -213,7 +213,7 @@ module Kafka
       end
 
       def mark_control_record
-        if in_transaction && is_control_batch
+        if is_control_batch
           record = @records.first
           record.is_control_record = true unless record.nil?
         end

data/lib/kafka/protocol/sasl_handshake_request.rb

@@ -8,7 +8,7 @@ module Kafka
 
     class SaslHandshakeRequest
 
-      SUPPORTED_MECHANISMS = %w(GSSAPI PLAIN SCRAM-SHA-256 SCRAM-SHA-512 OAUTHBEARER)
+      SUPPORTED_MECHANISMS = %w(AWS_MSK_IAM GSSAPI PLAIN SCRAM-SHA-256 SCRAM-SHA-512 OAUTHBEARER)
 
       def initialize(mechanism)
         unless SUPPORTED_MECHANISMS.include?(mechanism)

data/lib/kafka/protocol/sync_group_response.rb

@@ -13,9 +13,12 @@ module Kafka
       end
 
       def self.decode(decoder)
+        error_code = decoder.int16
+        member_assignment_bytes = decoder.bytes
+
         new(
-          error_code: decoder.int16,
-          member_assignment: MemberAssignment.decode(Decoder.from_string(decoder.bytes)),
+          error_code: error_code,
+          member_assignment: member_assignment_bytes ? MemberAssignment.decode(Decoder.from_string(member_assignment_bytes)) : nil
         )
       end
     end

data/lib/kafka/protocol/txn_offset_commit_response.rb

@@ -1,17 +1,46 @@
+# frozen_string_literal: true
+
 module Kafka
   module Protocol
     class TxnOffsetCommitResponse
+      class PartitionError
+        attr_reader :partition, :error_code
+
+        def initialize(partition:, error_code:)
+          @partition = partition
+          @error_code = error_code
+        end
+      end
+
+      class TopicPartitionsError
+        attr_reader :topic, :partitions
+
+        def initialize(topic:, partitions:)
+          @topic = topic
+          @partitions = partitions
+        end
+      end
 
-      attr_reader :error_code
+      attr_reader :errors
 
-      def initialize(error_code:)
-        @error_code = error_code
+      def initialize(errors:)
+        @errors = errors
       end
 
       def self.decode(decoder)
         _throttle_time_ms = decoder.int32
-        error_code = decoder.int16
-        new(error_code: error_code)
+        errors = decoder.array do
+          TopicPartitionsError.new(
+            topic: decoder.string,
+            partitions: decoder.array do
+              PartitionError.new(
+                partition: decoder.int32,
+                error_code: decoder.int16
+              )
+            end
+          )
+        end
+        new(errors: errors)
       end
     end
   end

data/lib/kafka/round_robin_assignment_strategy.rb

@@ -1,54 +1,52 @@
-# frozen_string_literal: true
-
-require "kafka/protocol/member_assignment"
-
 module Kafka
 
-  # A consumer group partition assignment strategy that assigns partitions to
-  # consumers in a round-robin fashion.
+  # A round robin assignment strategy inpired on the
+  # original java client round robin assignor. It's capable
+  # of handling identical as well as different topic subscriptions
+  # accross the same consumer group.
   class RoundRobinAssignmentStrategy
-    def initialize(cluster:)
-      @cluster = cluster
+    def protocol_name
+      "roundrobin"
     end
 
     # Assign the topic partitions to the group members.
     #
-    # @param members [Array<String>] member ids
-    # @param topics [Array<String>] topics
-    # @return [Hash<String, Protocol::MemberAssignment>] a hash mapping member
-    #   ids to assignments.
-    def assign(members:, topics:)
-      group_assignment = {}
-
-      members.each do |member_id|
-        group_assignment[member_id] = Protocol::MemberAssignment.new
-      end
-
-      topic_partitions = topics.flat_map do |topic|
-        begin
-          partitions = @cluster.partitions_for(topic).map(&:partition_id)
-        rescue UnknownTopicOrPartition
-          raise UnknownTopicOrPartition, "unknown topic #{topic}"
+    # @param cluster [Kafka::Cluster]
+    # @param members [Hash<String, Kafka::Protocol::JoinGroupResponse::Metadata>] a hash
+    #   mapping member ids to metadata
+    # @param partitions [Array<Kafka::ConsumerGroup::Assignor::Partition>] a list of
+    #   partitions the consumer group processes
+    # @return [Hash<String, Array<Kafka::ConsumerGroup::Assignor::Partition>] a hash
+    #   mapping member ids to partitions.
+    def call(cluster:, members:, partitions:)
+      partitions_per_member = Hash.new {|h, k| h[k] = [] }
+      relevant_partitions = valid_sorted_partitions(members, partitions)
+      members_ids = members.keys
+      iterator = (0...members.size).cycle
+      idx = iterator.next
+
+      relevant_partitions.each do |partition|
+        topic = partition.topic
+
+        while !members[members_ids[idx]].topics.include?(topic)
+          idx = iterator.next
         end
-        Array.new(partitions.count) { topic }.zip(partitions)
+
+        partitions_per_member[members_ids[idx]] << partition
+        idx = iterator.next
       end
 
-      partitions_per_member = topic_partitions.group_by.with_index do |_, index|
-        index % members.count
-      end.values
+      partitions_per_member
+    end
 
-      members.zip(partitions_per_member).each do |member_id, member_partitions|
-        unless member_partitions.nil?
-          member_partitions.each do |topic, partition|
-            group_assignment[member_id].assign(topic, [partition])
-          end
-        end
-      end
+    def valid_sorted_partitions(members, partitions)
+      subscribed_topics = members.map do |id, metadata|
+        metadata && metadata.topics
+      end.flatten.compact
 
-      group_assignment
-    rescue Kafka::LeaderNotAvailable
-      sleep 1
-      retry
+      partitions
+        .select { |partition| subscribed_topics.include?(partition.topic) }
+        .sort_by { |partition| partition.topic }
     end
   end
 end