ruby-kafka 1.2.0 → 1.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0f8391cc7b1989cb5f669796bc4ad647b77d882e6506fae42bab18acb8a6bcc6
-  data.tar.gz: 012baaff5d2cc9eb17e3a7b7342f49f7c905a5f91d26078fa0ecf2f0fa81a2ad
+  metadata.gz: b16e9e52014e784610725bb2ba5c5a431694ce6da7878b6c1ff024504d6ef6c4
+  data.tar.gz: '039026011e9cd5e5dce59aed4f53b2a8970aa8025fe3b008620bdef8a802bbf5'
 SHA512:
-  metadata.gz: 7f4e9302ca0ab41a6fded75f95ed866d959a2a027b90051ed7f3d7fba573aa63e57be7692d004dbeb3bacd99fe44b24188f81b0f3bed50f68e1da5189262271f
-  data.tar.gz: 7266bdd50e66a7ab9b3c71025468af0a4bea778fd312c41d4c699e2609420c11c28accb3246407ce98330bf521c109691bd1d4943dd532a44ee6732f1a410922
+  metadata.gz: 1d13c2032f4bd38e09a714fe40d59467ff61b2c51e0a604daf7f9e1a3431af32f77e62433b863862a68e6bc74f53ff4eceed7e88145acc102613ac6fa2600b9f
+  data.tar.gz: f424e6e2bee5f318766880f3b7eb338ce706783ea647908902a8a6d409fcf8f9ac3f0afe7ca8a932710a9d8f192e44a25a346da6639861fb7d377f3dd3257309

data/CHANGELOG.md

@@ -4,6 +4,11 @@ Changes and additions to the library will be listed here.
 
 ## Unreleased
 
+## 1.3.0
+
+- Support custom assignment strategy (#846).
+- Improved Exceptions in TransactionManager (#862).
+
 ## 1.2.0
 
 - Add producer consumer interceptors (#837).

data/README.md

@@ -26,6 +26,7 @@ A Ruby client library for [Apache Kafka](http://kafka.apache.org/), a distribute
         4. [Shutting Down a Consumer](#shutting-down-a-consumer)
         5. [Consuming Messages in Batches](#consuming-messages-in-batches)
         6. [Balancing Throughput and Latency](#balancing-throughput-and-latency)
+        7. [Customizing Partition Assignment Strategy](#customizing-partition-assignment-strategy)
     4. [Thread Safety](#thread-safety)
     5. [Logging](#logging)
     6. [Instrumentation](#instrumentation)
@@ -743,6 +744,88 @@ consumer.each_message do |message|
 end
 ```
 
+#### Customizing Partition Assignment Strategy
+
+In some cases, you might want to assign more partitions to some consumers. For example, in applications inserting some records to a database, the consumers running on hosts nearby the database can process more messages than the consumers running on other hosts.
+You can use a custom assignment strategy by passing an object that implements `#call` as the argument `assignment_strategy` like below:
+
+```ruby
+class CustomAssignmentStrategy
+  def initialize(user_data)
+    @user_data = user_data
+  end
+
+  # Assign the topic partitions to the group members.
+  #
+  # @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:)
+    ...
+  end
+end
+
+strategy = CustomAssignmentStrategy.new("some-host-information")
+consumer = kafka.consumer(group_id: "some-group", assignment_strategy: strategy)
+```
+
+`members` is a hash mapping member IDs to metadata, and partitions is a list of partitions the consumer group processes. The method `call` must return a hash mapping member IDs to partitions. For example, the following strategy assigns partitions randomly:
+
+```ruby
+class RandomAssignmentStrategy
+  def call(cluster:, members:, partitions:)
+    member_ids = members.keys
+    partitions.each_with_object(Hash.new {|h, k| h[k] = [] }) do |partition, partitions_per_member|
+      partitions_per_member[member_ids[rand(member_ids.count)]] << partition
+    end
+  end
+end
+```
+
+If the strategy needs user data, you should define the method `user_data` that returns user data on each consumer. For example, the following strategy uses the consumers' IP addresses as user data:
+
+```ruby
+class NetworkTopologyAssignmentStrategy
+  def user_data
+    Socket.ip_address_list.find(&:ipv4_private?).ip_address
+  end
+
+  def call(cluster:, members:, partitions:)
+    # Display the pair of the member ID and IP address
+    members.each do |id, metadata|
+      puts "#{id}: #{metadata.user_data}"
+    end
+
+    # Assign partitions considering the network topology
+    ...
+  end
+end
+```
+
+Note that the strategy uses the class name as the default protocol name. You can change it by defining the method `protocol_name`:
+
+```ruby
+class NetworkTopologyAssignmentStrategy
+  def protocol_name
+    "networktopology"
+  end
+
+  def user_data
+    Socket.ip_address_list.find(&:ipv4_private?).ip_address
+  end
+
+  def call(cluster:, members:, partitions:)
+    ...
+  end
+end
+```
+
+As the method `call` might receive different user data from what it expects, you should avoid using the same protocol name as another strategy that uses different user data.
+
 
 ### Thread Safety
 
@@ -972,7 +1055,7 @@ This configures the store to look up CA certificates from the system default cer
 
 In order to authenticate the client to the cluster, you need to pass in a certificate and key created for the client and trusted by the brokers.
 
-**NOTE**: You can disable hostname validation by passing `verify_hostname: false`.
+**NOTE**: You can disable hostname validation by passing `ssl_verify_hostname: false`.
 
 ```ruby
 kafka = Kafka.new(
@@ -981,6 +1064,7 @@ kafka = Kafka.new(
   ssl_client_cert: File.read('my_client_cert.pem'),
   ssl_client_cert_key: File.read('my_client_cert_key.pem'),
   ssl_client_cert_key_password: 'my_client_cert_key_password',
+  ssl_verify_hostname: false,
   # ...
 )
 ```

data/lib/kafka/async_producer.rb

@@ -59,8 +59,6 @@ module Kafka
   #     producer.shutdown
   #
   class AsyncProducer
-    THREAD_MUTEX = Mutex.new
-
     # Initializes a new AsyncProducer.
     #
     # @param sync_producer [Kafka::Producer] the synchronous producer that should
@@ -94,6 +92,8 @@ module Kafka
 
       # The timer will no-op if the delivery interval is zero.
       @timer = Timer.new(queue: @queue, interval: delivery_interval)
+
+      @thread_mutex = Mutex.new
     end
 
     # Produces a message to the specified topic.
@@ -131,6 +131,8 @@ module Kafka
     # @see Kafka::Producer#deliver_messages
     # @return [nil]
     def deliver_messages
+      ensure_threads_running!
+
       @queue << [:deliver_messages, nil]
 
       nil
@@ -142,6 +144,8 @@ module Kafka
     # @see Kafka::Producer#shutdown
     # @return [nil]
     def shutdown
+      ensure_threads_running!
+
       @timer_thread && @timer_thread.exit
       @queue << [:shutdown, nil]
       @worker_thread && @worker_thread.join
@@ -152,17 +156,22 @@ module Kafka
     private
 
     def ensure_threads_running!
-      THREAD_MUTEX.synchronize do
-        @worker_thread = nil unless @worker_thread && @worker_thread.alive?
-        @worker_thread ||= Thread.new { @worker.run }
-      end
+      return if worker_thread_alive? && timer_thread_alive?
 
-      THREAD_MUTEX.synchronize do
-        @timer_thread = nil unless @timer_thread && @timer_thread.alive?
-        @timer_thread ||= Thread.new { @timer.run }
+      @thread_mutex.synchronize do
+        @worker_thread = Thread.new { @worker.run } unless worker_thread_alive?
+        @timer_thread = Thread.new { @timer.run } unless timer_thread_alive?
       end
     end
 
+    def worker_thread_alive?
+      !!@worker_thread && @worker_thread.alive?
+    end
+
+    def timer_thread_alive?
+      !!@timer_thread && @timer_thread.alive?
+    end
+
     def buffer_overflow(topic, message)
       @instrumenter.instrument("buffer_overflow.async_producer", {
         topic: topic,
@@ -208,7 +217,7 @@ module Kafka
 
           case operation
           when :produce
-            produce(*payload)
+            produce(payload[0], **payload[1])
             deliver_messages if threshold_reached?
           when :deliver_messages
             deliver_messages

data/lib/kafka/client.rb

@@ -357,6 +357,8 @@ module Kafka
     #   If it is n (n > 0), the topic list will be refreshed every n seconds
     # @param interceptors [Array<Object>] a list of consumer interceptors that implement
     #   `call(Kafka::FetchedBatch)`.
+    # @param assignment_strategy [Object] a partition assignment strategy that
+    #   implements `protocol_type()`, `user_data()`, and `assign(members:, partitions:)`
     # @return [Consumer]
     def consumer(
         group_id:,
@@ -368,7 +370,8 @@ module Kafka
         offset_retention_time: nil,
         fetcher_max_queue_size: 100,
         refresh_topic_interval: 0,
-        interceptors: []
+        interceptors: [],
+        assignment_strategy: nil
     )
       cluster = initialize_cluster
 
@@ -387,6 +390,7 @@ module Kafka
         rebalance_timeout: rebalance_timeout,
         retention_time: retention_time,
         instrumenter: instrumenter,
+        assignment_strategy: assignment_strategy
       )
 
       fetcher = Fetcher.new(

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,8 +189,8 @@ module Kafka
       if group_leader?
         @logger.info "Chosen as leader of group `#{@group_id}`"
 
-        group_assignment = @assignment_strategy.assign(
-          members: @members.keys,
+        group_assignment = @assignor.assign(
+          members: @members,
           topics: @topics,
         )
       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/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/round_robin_assignment_strategy.rb

@@ -1,54 +1,31 @@
 # 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.
   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}"
-        end
-        Array.new(partitions.count) { topic }.zip(partitions)
-      end
-
-      partitions_per_member = topic_partitions.group_by.with_index do |_, index|
-        index % members.count
-      end.values
-
-      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
+    # @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:)
+      member_ids = members.keys
+      partitions_per_member = Hash.new {|h, k| h[k] = [] }
+      partitions.each_with_index do |partition, index|
+        partitions_per_member[member_ids[index % member_ids.count]] << partition
       end
 
-      group_assignment
-    rescue Kafka::LeaderNotAvailable
-      sleep 1
-      retry
+      partitions_per_member
     end
   end
 end

data/lib/kafka/transaction_manager.rb

@@ -95,7 +95,7 @@ module Kafka
       force_transactional!
 
       if @transaction_state.uninitialized?
-        raise 'Transaction is uninitialized'
+        raise Kafka::InvalidTxnStateError, 'Transaction is uninitialized'
       end
 
       # Extract newly created partitions
@@ -138,8 +138,8 @@ module Kafka
 
     def begin_transaction
       force_transactional!
-      raise 'Transaction has already started' if @transaction_state.in_transaction?
-      raise 'Transaction is not ready' unless @transaction_state.ready?
+      raise Kafka::InvalidTxnStateError, 'Transaction has already started' if @transaction_state.in_transaction?
+      raise Kafka::InvalidTxnStateError, 'Transaction is not ready' unless @transaction_state.ready?
       @transaction_state.transition_to!(TransactionStateMachine::IN_TRANSACTION)
 
       @logger.info "Begin transaction #{@transactional_id}, Producer ID: #{@producer_id} (Epoch #{@producer_epoch})"
@@ -159,7 +159,7 @@ module Kafka
       end
 
       unless @transaction_state.in_transaction?
-        raise 'Transaction is not valid to commit'
+        raise Kafka::InvalidTxnStateError, 'Transaction is not valid to commit'
       end
 
       @transaction_state.transition_to!(TransactionStateMachine::COMMITTING_TRANSACTION)
@@ -192,7 +192,8 @@ module Kafka
       end
 
       unless @transaction_state.in_transaction?
-        raise 'Transaction is not valid to abort'
+        @logger.warn('Aborting transaction that was never opened on brokers')
+        return
       end
 
       @transaction_state.transition_to!(TransactionStateMachine::ABORTING_TRANSACTION)
@@ -221,7 +222,7 @@ module Kafka
       force_transactional!
 
       unless @transaction_state.in_transaction?
-        raise 'Transaction is not valid to send offsets'
+        raise Kafka::InvalidTxnStateError, 'Transaction is not valid to send offsets'
       end
 
       add_response = transaction_coordinator.add_offsets_to_txn(
@@ -250,6 +251,10 @@ module Kafka
       @transaction_state.error?
     end
 
+    def ready?
+      @transaction_state.ready?
+    end
+
     def close
       if in_transaction?
         @logger.warn("Aborting pending transaction ...")
@@ -264,11 +269,11 @@ module Kafka
 
     def force_transactional!
       unless transactional?
-        raise 'Please turn on transactional mode to use transaction'
+        raise Kafka::InvalidTxnStateError, 'Please turn on transactional mode to use transaction'
       end
 
       if @transactional_id.nil? || @transactional_id.empty?
-        raise 'Please provide a transaction_id to use transactional mode'
+        raise Kafka::InvalidTxnStateError, 'Please provide a transaction_id to use transactional mode'
       end
     end
 

data/lib/kafka/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Kafka
-  VERSION = "1.2.0"
+  VERSION = "1.3.0"
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: ruby-kafka
 version: !ruby/object:Gem::Version
-  version: 1.2.0
+  version: 1.3.0
 platform: ruby
 authors:
 - Daniel Schierbeck
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2020-08-03 00:00:00.000000000 Z
+date: 2020-10-14 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: digest-crc
@@ -376,6 +376,7 @@ files:
 - lib/kafka/connection_builder.rb
 - lib/kafka/consumer.rb
 - lib/kafka/consumer_group.rb
+- lib/kafka/consumer_group/assignor.rb
 - lib/kafka/datadog.rb
 - lib/kafka/fetch_operation.rb
 - lib/kafka/fetched_batch.rb
@@ -494,8 +495,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubyforge_project: 
-rubygems_version: 2.7.6
+rubygems_version: 3.1.2
 signing_key: 
 specification_version: 4
 summary: A client library for the Kafka distributed commit log.