Support only block call fasion strategy

cf. zendesk#846 (comment)

Author: abicky

README.md

@@ -747,47 +747,46 @@ 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 an object that implements `#user_data` and `#assign` as the argument `assignment_strategy` like below:
+You can use a custom assignment strategy like below:
 
 ```ruby
-class CustomAssignmentStrategy
-  def initialize(user_data)
-    @user_data = user_data
-  end
+Kafka::ConsumerGroup::Assignor.register_strategy(:custom) do |cluster:, members:, partitions:|
+  # strategy goes here
+end
+consumer = kafka.consumer(group_id: "some-group", assignment_strategy: :custom)
+```
 
-  # @return [String, nil]
-  def user_data
-    @user_data
-  end
+`members` is a hash mapping member ids to metadata, and `partitions` is a list of partitions the consumer group processes. The block should return a hash mapping member ids to partitions.
+For example, the following strategy assigns partitions randomly:
 
-  # 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 assign(cluster:, members:, partitions:)
-    ...
+```ruby
+Kafka::ConsumerGroup::Assignor.register_strategy(:custom) do |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
-Kafka::ConsumerGroup::Assignor.register_strategy(:custom, CustomAssignmentStrategy)
-
-strategy = CustomAssignmentStrategy.new("some-host-information")
-consumer = kafka.consumer(group_id: "some-group", assignment_strategy: strategy)
 ```
 
-If the strategy doesn't need user data and constructor arguments, you don't need to define a class:
+If the strategy needs user data, you can pass the option `user_data`, a proc that returns user data on each consumer:
 
-```
-Kafka::ConsumerGroup::Assignor.register_strategy(:another_custom) do |cluster:, members:, partitions:|
+```ruby
+Kafka::ConsumerGroup::Assignor.register_strategy(:custom, user_data: ->{ ... }) do |cluster:, members:, partitions:|
   # strategy goes here
 end
-consumer = kafka.consumer(group_id: "some-group", assignment_strategy: :another_custom)
 ```
 
+For example, the following strategy displays the IP address of each consumer:
+
+```ruby
+user_data_proc = ->{ Socket.ip_address_list.find(&:ipv4_private?).ip_address }
+Kafka::ConsumerGroup::Assignor.register_strategy(:custom, user_data: user_data_proc) do |cluster:, members:, partitions:|
+  members.each do |id, metadata|
+    puts "#{id}: #{metadata.user_data}"
+  end
+  ...
+end
+```
 
 ### Thread Safety
 

lib/kafka/consumer_group.rb

@@ -22,7 +22,7 @@ def initialize(cluster:, logger:, group_id:, session_timeout:, rebalance_timeout
       @assigned_partitions = {}
       @assignor = Assignor.new(
         cluster: cluster,
-        strategy: assignment_strategy || RoundRobinAssignmentStrategy.new
+        strategy: assignment_strategy || :roundrobin
       )
       @retention_time = retention_time
     end

lib/kafka/consumer_group/assignor.rb

@@ -7,30 +7,33 @@ class ConsumerGroup
 
     # A consumer group partition assignor
     class Assignor
+      Strategy = Struct.new(:assign, :user_data)
       Partition = Struct.new(:topic, :partition_id)
 
-      def self.register_strategy(name, strategy_class = nil, &block)
-        if strategy_classes[name.to_s]
+      # @param name [String]
+      # @param user_data [Proc, nil] a proc that returns user data as a string
+      # @yield [cluster:, members:, partitions:]
+      # @yieldparam cluster [Kafka::Cluster]
+      # @yieldparam members [Hash<String, Kafka::Protocol::JoinGroupResponse::Metadata>]
+      #   a hash mapping member ids to metadata
+      # @yieldparam partitions [Array<Kafka::ConsumerGroup::Assignor::Partition>]
+      #   a list of partitions the consumer group processes
+      # @yieldreturn [Hash<String, Array<Kafka::ConsumerGroup::Assignor::Partition>]
+      #   a hash mapping member ids to partitions.
+      def self.register_strategy(name, user_data: nil, &block)
+        if strategies[name.to_s]
           raise ArgumentError, "The strategy '#{name}' is already registered."
         end
 
-        unless strategy_class.nil? ^ !block_given?
-          raise "Either strategy class or block but not both must be specified."
+        unless block_given?
+          raise ArgumentError, "The block must be given."
         end
 
-        if block_given?
-          strategy_class = Class.new do
-            define_method(:assign) do |cluster:, members:, partitions:|
-              block.call(cluster: cluster, members: members, partitions: partitions)
-            end
-          end
-        end
-
-        strategy_classes[name.to_s] = strategy_class
+        strategies[name.to_s] = Strategy.new(block, user_data)
       end
 
-      def self.strategy_classes
-        @strategy_classes ||= {}
+      def self.strategies
+        @strategies ||= {}
       end
 
       attr_reader :protocol_name
@@ -40,25 +43,14 @@ def self.strategy_classes
       #   #user_data and #assign.
       def initialize(cluster:, strategy:)
         @cluster = cluster
-
-        case strategy
-        when String, Symbol
-          @protocol_name = strategy.to_s
-          klass = self.class.strategy_classes.fetch(@protocol_name) do
-            raise ArgumentError, "The strategy '#{@protocol_name}' is not registered."
-          end
-          @strategy = klass.new
-        else
-          @protocol_name, _ = self.class.strategy_classes.find {|name, klass| strategy.is_a?(klass) }
-          if @protocol_name.nil?
-            raise ArgumentError, "The strategy class '#{strategy.class}' is not registered."
-          end
-          @strategy = strategy
+        @protocol_name = strategy.to_s
+        @strategy = self.class.strategies.fetch(@protocol_name) do
+          raise ArgumentError, "The strategy '#{@protocol_name}' is not registered."
         end
       end
 
       def user_data
-        @strategy.user_data if @strategy.respond_to?(:user_data)
+        @strategy.user_data.call if @strategy.user_data
       end
 
       # Assign the topic partitions to the group members.
@@ -83,7 +75,8 @@ def assign(members:, topics:)
         members.each_key do |member_id|
           group_assignment[member_id] = Protocol::MemberAssignment.new
         end
-        @strategy.assign(cluster: @cluster, members: members, partitions: topic_partitions).each do |member_id, partitions|
+
+        @strategy.assign.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

lib/kafka/round_robin_assignment_strategy.rb

@@ -5,10 +5,6 @@ module Kafka
   # A consumer group partition assignment strategy that assigns partitions to
   # consumers in a round-robin fashion.
   class RoundRobinAssignmentStrategy
-    def user_data
-      nil
-    end
-
     # Assign the topic partitions to the group members.
     #
     # @param cluster [Kafka::Cluster]
@@ -31,4 +27,7 @@ def assign(cluster:, members:, partitions:)
 end
 
 require "kafka/consumer_group/assignor"
-Kafka::ConsumerGroup::Assignor.register_strategy(:roundrobin, Kafka::RoundRobinAssignmentStrategy)
+strategy = Kafka::RoundRobinAssignmentStrategy.new
+Kafka::ConsumerGroup::Assignor.register_strategy(:roundrobin) do |cluster:, members:, partitions:|
+  strategy.assign(cluster: cluster, members: members, partitions: partitions)
+end

spec/consumer_group/assignor_spec.rb

@@ -6,7 +6,7 @@
       around do |example|
         described_class.register_strategy(:custom) { |**kwargs| [] }
         example.run
-        described_class.strategy_classes.delete("custom")
+        described_class.strategies.delete("custom")
       end
 
       it do
@@ -17,79 +17,40 @@
 
   describe "#assign" do
     let(:cluster) { double(:cluster) }
-    let(:assignor) { described_class.new(cluster: cluster, strategy: strategy) }
+    let(:assignor) { described_class.new(cluster: cluster, strategy: :custom) }
 
     let(:members) { Hash[(0...10).map {|i| ["member#{i}", nil] }] }
     let(:topics) { ["greetings"] }
     let(:partitions) { (0...30).map {|i| double(:"partition#{i}", partition_id: i) } }
 
     before do
       allow(cluster).to receive(:partitions_for) { partitions }
-    end
-    after do
-      described_class.strategy_classes.delete("custom")
-    end
 
-    context "when the strategy is an object" do
-      let(:strategy) do
-        klass = Class.new do
-          def assign(cluster:, members:, partitions:)
-            assignment = {}
-            partition_count_per_member = (partitions.count.to_f / members.count).ceil
-            partitions.each_slice(partition_count_per_member).with_index do |chunk, index|
-              assignment[members.keys[index]] = chunk
-            end
-
-            assignment
-          end
+      described_class.register_strategy(:custom) do |cluster:, members:, partitions:|
+        assignment = {}
+        partition_count_per_member = (partitions.count.to_f / members.count).ceil
+        partitions.each_slice(partition_count_per_member).with_index do |chunk, index|
+          assignment[members.keys[index]] = chunk
         end
-        described_class.register_strategy(:custom, klass)
-
-        klass.new
-      end
-
-      it "assigns all partitions" do
-        assignments = assignor.assign(members: members, topics: topics)
-
-        partitions.each do |partition|
-          member = assignments.values.find {|assignment|
-            assignment.topics.find {|topic, partitions|
-              partitions.include?(partition.partition_id)
-            }
-          }
 
-          expect(member).to_not be_nil
-        end
+        assignment
       end
     end
+    after do
+      described_class.strategies.delete("custom")
+    end
 
-    context "when the strategy is a string" do
-      let(:strategy) { :custom }
-
-      before do
-        described_class.register_strategy(:custom) do |cluster:, members:, partitions:|
-          assignment = {}
-          partition_count_per_member = (partitions.count.to_f / members.count).ceil
-          partitions.each_slice(partition_count_per_member).with_index do |chunk, index|
-            assignment[members.keys[index]] = chunk
-          end
-
-          assignment
-        end
-      end
-
-      it "assigns all partitions" do
-        assignments = assignor.assign(members: members, topics: topics)
+    it "assigns all partitions" do
+      assignments = assignor.assign(members: members, topics: topics)
 
-        partitions.each do |partition|
-          member = assignments.values.find {|assignment|
-            assignment.topics.find {|topic, partitions|
-              partitions.include?(partition.partition_id)
-            }
+      partitions.each do |partition|
+        member = assignments.values.find {|assignment|
+          assignment.topics.find {|topic, partitions|
+            partitions.include?(partition.partition_id)
           }
+        }
 
-          expect(member).to_not be_nil
-        end
+        expect(member).to_not be_nil
       end
     end
   end

spec/functional/consumer_group_spec.rb

@@ -436,15 +436,16 @@
   end
 
   example "consuming messages with a custom assignment strategy" do
-    topic = create_random_topic(num_partitions: 1)
-    messages = (1..1000).to_a
+    num_partitions = 3
+    topic = create_random_topic(num_partitions: num_partitions)
+    messages = (1..100 * num_partitions).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)
+        producer.produce(i.to_s, topic: topic, partition: i % num_partitions)
       end
 
       producer.deliver_messages
@@ -453,35 +454,25 @@
     group_id = "test#{rand(1000)}"
 
     mutex = Mutex.new
-    received_messages = []
-
-    assignment_strategy_class = Class.new do
-      def initialize(weight)
-        @weight = weight
-      end
+    received_messages = Hash.new {|h, k| h[k] = [] }
 
-      def user_data
-        @weight.to_s
+    consumer_weights = {}
+    user_data_proc = -> do
+      consumer_weights[Thread.current.object_id] ||= (consumer_weights.size + 1).to_s
+    end
+    Kafka::ConsumerGroup::Assignor.register_strategy(:custom, user_data: user_data_proc) do |cluster:, members:, partitions:|
+      member_ids = members.flat_map {|id, metadata| [id] * metadata.user_data.to_i }
+      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
 
-      def assign(cluster:, members:, partitions:)
-        member_ids = members.flat_map {|id, metadata| [id] * metadata.user_data.to_i }
-        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
-
-        partitions_per_member
-      end
+      partitions_per_member
     end
-
-    Kafka::ConsumerGroup::Assignor.register_strategy(:custom, assignment_strategy_class)
     begin
-      consumers = 2.times.map do |i|
-        assignment_strategy = assignment_strategy_class.new(i + 1)
-
+      consumers = 2.times.map do
         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: assignment_strategy)
+        consumer = kafka.consumer(group_id: group_id, offset_retention_time: offset_retention_time, assignment_strategy: :custom)
         consumer.subscribe(topic)
         consumer
       end
@@ -490,9 +481,9 @@ def assign(cluster:, members:, partitions:)
         t = Thread.new do
           consumer.each_message do |message|
             mutex.synchronize do
-              received_messages << message
+              received_messages[consumer] << message
 
-              if received_messages.count == messages.count
+              if received_messages.values.flatten.count == messages.count
                 consumers.each(&:stop)
               end
             end
@@ -506,9 +497,10 @@ def assign(cluster:, members:, partitions:)
 
       threads.each(&:join)
 
-      expect(received_messages.map(&:value).map(&:to_i)).to match_array messages
+      expect(received_messages.values.flatten.map {|v| v.value.to_i }).to match_array messages
+      expect(received_messages.values.map(&:count)).to match_array [messages.count / 3, messages.count / 3 * 2]
     ensure
-      Kafka::ConsumerGroup::Assignor.strategy_classes.delete("custom")
+      Kafka::ConsumerGroup::Assignor.strategies.delete("custom")
     end
   end