Support custom assignment strategy

This commit adds the feature that allows users to specify custom
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.

Author: abicky

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)
@@ -723,6 +724,45 @@ 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 implement a custom assignment strategy and use it by passing a proc object that create the strategy as the argument `assignment_strategy_builder` like below:
+
+```ruby
+class CustomAssignmentStrategy
+  def initialize(cluster, user_data)
+    @cluster = cluster
+    @user_data = user_data
+  end
+
+  # @return [String]
+  def protocol_name
+    ...
+  end
+
+  # @return [String, nil]
+  def user_data
+    @user_data
+  end
+
+  # Assign the topic partitions to the group members.
+  #
+  # @param members [Hash<String, Protocol::JoinGroupResponse::Metadata>] a hash
+  #   mapping member ids to metadata
+  # @param topics [Array<String>] topics
+  # @return [Hash<String, Protocol::MemberAssignment>] a hash mapping member
+  #   ids to assignments.
+  def assign(members:, topics:)
+    ...
+  end
+end
+
+assignment_strategy_builder = ->(cluster) do
+  CustomAssignmentStrategy.new(cluster, "some-host-information")
+end
+consumer = kafka.consumer(group_id: "some-group", assignment_strategy_builder: assignment_strategy_builder)
+```
 
 ### Thread Safety
 

lib/kafka/client.rb

@@ -350,6 +350,8 @@ def async_producer(delivery_interval: 0, delivery_threshold: 0, max_queue_size:
     #   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_builder [Proc] a procedure that implements
+    #   `call(Kafka::Cluster)` and creates a assignment strategy.
     # @return [Consumer]
     def consumer(
         group_id:,
@@ -361,7 +363,8 @@ def consumer(
         offset_retention_time: nil,
         fetcher_max_queue_size: 100,
         refresh_topic_interval: 0,
-        interceptors: []
+        interceptors: [],
+        assignment_strategy_builder: nil
     )
       cluster = initialize_cluster
 
@@ -372,6 +375,12 @@ def consumer(
       # The Kafka protocol expects the retention time to be in ms.
       retention_time = (offset_retention_time && offset_retention_time * 1_000) || -1
 
+      if assignment_strategy_builder
+        assignment_strategy = assignment_strategy_builder.call(cluster)
+      else
+        assignment_strategy = RoundRobinAssignmentStrategy.new(cluster: cluster)
+      end
+
       group = ConsumerGroup.new(
         cluster: cluster,
         logger: @logger,
@@ -380,6 +389,7 @@ def consumer(
         rebalance_timeout: rebalance_timeout,
         retention_time: retention_time,
         instrumenter: instrumenter,
+        assignment_strategy: assignment_strategy,
       )
 
       fetcher = Fetcher.new(

lib/kafka/consumer_group.rb

@@ -7,7 +7,7 @@ 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 +19,7 @@ def initialize(cluster:, logger:, group_id:, session_timeout:, rebalance_timeout
       @members = {}
       @topics = Set.new
       @assigned_partitions = {}
-      @assignment_strategy = RoundRobinAssignmentStrategy.new(cluster: @cluster)
+      @assignment_strategy = assignment_strategy
       @retention_time = retention_time
     end
 
@@ -144,6 +144,8 @@ def join_group
           rebalance_timeout: @rebalance_timeout,
           member_id: @member_id,
           topics: @topics,
+          protocol_name: @assignment_strategy.protocol_name,
+          user_data: @assignment_strategy.user_data,
         )
 
         Protocol.handle_error(response.error_code)
@@ -181,7 +183,7 @@ def synchronize
         @logger.info "Chosen as leader of group `#{@group_id}`"
 
         group_assignment = @assignment_strategy.assign(
-          members: @members.keys,
+          members: @members,
           topics: @topics,
         )
       end

lib/kafka/protocol/join_group_request.rb

@@ -7,14 +7,14 @@ module Protocol
     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
 

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 @@ def self.decode(decoder)
           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

lib/kafka/round_robin_assignment_strategy.rb

@@ -11,16 +11,25 @@ def initialize(cluster:)
       @cluster = cluster
     end
 
+    def protocol_name
+      "roundrobin"
+    end
+
+    def user_data
+      nil
+    end
+
     # Assign the topic partitions to the group members.
     #
-    # @param members [Array<String>] member ids
+    # @param members [Hash<String, Protocol::JoinGroupResponse::Metadata>] a hash
+    #   mapping member ids to metadata
     # @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|
+      members.each_key do |member_id|
         group_assignment[member_id] = Protocol::MemberAssignment.new
       end
 
@@ -37,7 +46,7 @@ def assign(members:, topics:)
         index % members.count
       end.values
 
-      members.zip(partitions_per_member).each do |member_id, member_partitions|
+      members.keys.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])

spec/functional/consumer_group_spec.rb

@@ -435,6 +435,93 @@
     end
   end
 
+  example "consuming messages with a custom assignment strategy" do
+    topic = create_random_topic(num_partitions: 1)
+    messages = (1..1000).to_a
+
+    begin
+      kafka = Kafka.new(kafka_brokers, client_id: "test")
+      producer = kafka.producer
+
+      messages.each do |i|
+        producer.produce(i.to_s, topic: topic, partition: 0)
+      end
+
+      producer.deliver_messages
+    end
+
+    group_id = "test#{rand(1000)}"
+
+    mutex = Mutex.new
+    received_messages = []
+
+    assignment_strategy_class = Class.new do
+      def initialize(cluster, weight)
+        @cluster = cluster
+        @weight = weight
+      end
+
+      def protocol_name
+        "custom"
+      end
+
+      def user_data
+        @weight.to_s
+      end
+
+      def assign(members:, topics:)
+        group_assignment = members.each_key.with_object({}) do |member_id, assignment|
+          assignment[member_id] = Kafka::Protocol::MemberAssignment.new
+        end
+
+        member_ids = members.flat_map { |id, metadata| [id] * metadata.user_data.to_i }
+        partition_index = 0
+        topics.each_with_index do |topic, i|
+          @cluster.partitions_for(topic).each_with_index do |partition, j|
+            member_id = member_ids[partition_index % member_ids.size]
+            group_assignment[member_id].assign(topic, [partition.partition_id])
+            partition_index += 1
+          end
+        end
+
+        group_assignment
+      end
+    end
+
+    consumers = 2.times.map do |i|
+      assignment_strategy_builder = ->(cluster) {
+        assignment_strategy_class.new(cluster, i + 1)
+      }
+
+      kafka = Kafka.new(kafka_brokers, client_id: "test", logger: logger)
+      consumer = kafka.consumer(group_id: group_id, offset_retention_time: offset_retention_time, assignment_strategy_builder: assignment_strategy_builder)
+      consumer.subscribe(topic)
+      consumer
+    end
+
+    threads = consumers.map do |consumer|
+      t = Thread.new do
+        consumer.each_message do |message|
+          mutex.synchronize do
+            received_messages << message
+
+            if received_messages.count == messages.count
+              consumers.each(&:stop)
+            end
+          end
+        end
+      end
+
+      t.abort_on_exception = true
+
+      t
+    end
+
+    threads.each(&:join)
+
+    expect(received_messages.map(&:value).map(&:to_i)).to match_array messages
+  end
+
   def wait_until(timeout:)
     Timeout.timeout(timeout) do
       sleep 0.5 until yield

spec/round_robin_assignment_strategy_spec.rb

@@ -5,7 +5,7 @@
   let(:strategy) { described_class.new(cluster: cluster) }
 
   it "assigns all partitions" do
-    members = (0...10).map {|i| "member#{i}" }
+    members = Hash[(0...10).map {|i| ["member#{i}", nil] }]
     topics = ["greetings"]
     partitions = (0...30).map {|i| double(:"partition#{i}", partition_id: i) }
 
@@ -28,7 +28,7 @@
     cluster = double(:cluster)
     strategy = described_class.new(cluster: cluster)
 
-    members = (0...10).map {|i| "member#{i}" }
+    members = Hash[(0...10).map {|i| ["member#{i}", nil] }]
     topics = ["topic1", "topic2"]
     partitions = (0...5).map {|i| double(:"partition#{i}", partition_id: i) }
 
@@ -57,32 +57,32 @@
     {
       name: "uneven topics",
       topics: { "topic1" => [0], "topic2" => (0..50).to_a },
-      members: ["member1", "member2"],
+      members: { "member1" => nil, "member2" => nil },
     },
     {
       name: "only one partition",
       topics: { "topic1" => [0] },
-      members: ["member1", "member2"],
+      members: { "member1" => nil, "member2" => nil },
     },
     {
       name: "lots of partitions",
       topics: { "topic1" => (0..100).to_a },
-      members: ["member1"]
+      members: { "member1" => nil },
     },
     {
       name: "lots of members",
       topics: { "topic1" => (0..10).to_a, "topic2" => (0..10).to_a },
-      members: (0..50).map { |i| "member#{i}" }
+      members: Hash[(0..50).map { |i| ["member#{i}", nil] }]
     },
     {
       name: "odd number of partitions",
       topics: { "topic1" => (0..14).to_a },
-      members: ["member1", "member2"]
+      members: { "member1" => nil, "member2" => nil },
     },
     {
       name: "five topics, 10 partitions, 3 consumers",
       topics: { "topic1" => [0, 1], "topic2" => [0, 1], "topic3" => [0, 1], "topic4" => [0, 1], "topic5" => [0, 1] },
-      members: ["member1", "member2", "member3"]
+      members: { "member1" => nil, "member2" => nil, "member3" => nil },
     }
   ].each do |name:, topics:, members:|
       it name do