Implement Kafka::ConsumerGroup::Assignor.register_strategy

This method prevents us from using the same protocol name
and makes us register custom strategies easily.

Author: abicky

README.md

@@ -727,41 +727,48 @@ 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 `#protocol_name`, `#user_data`, and `#assign` as the argument `assignment_strategy` like below:
+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:
 
 ```ruby
 class CustomAssignmentStrategy
   def initialize(user_data)
     @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 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(members:, partitions:)
+  def assign(cluster:, members:, partitions:)
     ...
   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:
+
+```
+Kafka::ConsumerGroup::Assignor.register_strategy(:another_custom) do |cluster:, members:, partitions:|
+  # strategy goes here
+end
+consumer = kafka.consumer(group_id: "some-group", assignment_strategy: :another_custom)
+```
+
+
 ### Thread Safety
 
 You typically don't want to share a Kafka client object between threads, since the network communication is not synchronized. Furthermore, you should avoid using threads in a consumer unless you're very careful about waiting for all work to complete before returning from the `#each_message` or `#each_batch` block. This is because _checkpointing_ assumes that returning from the block means that the messages that have been yielded have been successfully processed.

lib/kafka/consumer_group/assignor.rb

@@ -9,20 +9,56 @@ class ConsumerGroup
     class Assignor
       Partition = Struct.new(:topic, :partition_id)
 
+      def self.register_strategy(name, strategy_class = nil, &block)
+        if strategy_classes[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."
+        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
+      end
+
+      def self.strategy_classes
+        @strategy_classes ||= {}
+      end
+
+      attr_reader :protocol_name
+
       # @param cluster [Kafka::Cluster]
-      # @param strategy [Object] an object that implements #protocol_type,
-      #   #user_data, and #assign.
+      # @param strategy [String, Symbol, Object]  a string, a symbol, or an object that implements
+      #   #user_data and #assign.
       def initialize(cluster:, strategy:)
         @cluster = cluster
-        @strategy = strategy
-      end
 
-      def protocol_name
-        @strategy.protocol_name
+        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
+        end
       end
 
       def user_data
-        @strategy.user_data
+        @strategy.user_data if @strategy.respond_to?(:user_data)
       end
 
       # Assign the topic partitions to the group members.
@@ -47,7 +83,7 @@ def assign(members:, topics:)
         members.each_key do |member_id|
           group_assignment[member_id] = Protocol::MemberAssignment.new
         end
-        @strategy.assign(members: members, partitions: topic_partitions).each do |member_id, partitions|
+        @strategy.assign(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,23 +5,20 @@ module Kafka
   # A consumer group partition assignment strategy that assigns partitions to
   # consumers in a round-robin fashion.
   class RoundRobinAssignmentStrategy
-    def protocol_name
-      "roundrobin"
-    end
-
     def user_data
       nil
     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 assign(members:, partitions:)
+    def assign(cluster:, members:, partitions:)
       member_ids = members.keys
       partitions_per_member = Hash.new {|h, k| h[k] = [] }
       partitions.each_with_index do |partition, index|
@@ -32,3 +29,6 @@ def assign(members:, partitions:)
     end
   end
 end
+
+require "kafka/consumer_group/assignor"
+Kafka::ConsumerGroup::Assignor.register_strategy(:roundrobin, Kafka::RoundRobinAssignmentStrategy)

spec/consumer_group/assignor_spec.rb

@@ -1,48 +1,96 @@
 # frozen_string_literal: true
 
 describe Kafka::ConsumerGroup::Assignor do
-  let(:cluster) { double(:cluster) }
-  let(:assignor) { described_class.new(cluster: cluster, strategy: strategy) }
-  let(:strategy) do
-    klass = Class.new do
-      def protocol_name
-        "test"
+  describe ".register_strategy" do
+    context "when another strategy is already registered with the same name" do
+      around do |example|
+        described_class.register_strategy(:custom) { |**kwargs| [] }
+        example.run
+        described_class.strategy_classes.delete("custom")
       end
 
-      def user_data
-        nil
+      it do
+        expect { described_class.register_strategy(:custom) { |**kwargs| [] } }.to raise_error(ArgumentError)
       end
+    end
+  end
+
+  describe "#assign" do
+    let(:cluster) { double(:cluster) }
+    let(:assignor) { described_class.new(cluster: cluster, strategy: strategy) }
+
+    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
 
-      def assign(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
+            assignment
+          end
         end
+        described_class.register_strategy(:custom, klass)
 
-        assignment
+        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
       end
     end
-    klass.new
-  end
 
-  it "assigns all partitions" do
-    members = Hash[(0...10).map {|i| ["member#{i}", nil] }]
-    topics = ["greetings"]
-    partitions = (0...30).map {|i| double(:"partition#{i}", partition_id: i) }
+    context "when the strategy is a string" do
+      let(:strategy) { :custom }
 
-    allow(cluster).to receive(:partitions_for) { partitions }
+      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
 
-    assignments = assignor.assign(members: members, topics: topics)
+          assignment
+        end
+      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)
-        }
-      }
+        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
+          expect(member).to_not be_nil
+        end
+      end
     end
   end
 end

spec/functional/consumer_group_spec.rb

@@ -460,15 +460,11 @@ def initialize(weight)
         @weight = weight
       end
 
-      def protocol_name
-        "custom"
-      end
-
       def user_data
         @weight.to_s
       end
 
-      def assign(members:, partitions:)
+      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|
@@ -479,36 +475,41 @@ def assign(members:, partitions:)
       end
     end
 
-    consumers = 2.times.map do |i|
-      assignment_strategy = assignment_strategy_class.new(i + 1)
+    Kafka::ConsumerGroup::Assignor.register_strategy(:custom, assignment_strategy_class)
+    begin
+      consumers = 2.times.map do |i|
+        assignment_strategy = assignment_strategy_class.new(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: assignment_strategy)
-      consumer.subscribe(topic)
-      consumer
-    end
+        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.subscribe(topic)
+        consumer
+      end
 
-    threads = consumers.map do |consumer|
-      t = Thread.new do
-        consumer.each_message do |message|
-          mutex.synchronize do
-            received_messages << message
+      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)
+              if received_messages.count == messages.count
+                consumers.each(&:stop)
+              end
             end
           end
         end
-      end
 
-      t.abort_on_exception = true
+        t.abort_on_exception = true
 
-      t
-    end
+        t
+      end
 
-    threads.each(&:join)
+      threads.each(&:join)
 
-    expect(received_messages.map(&:value).map(&:to_i)).to match_array messages
+      expect(received_messages.map(&:value).map(&:to_i)).to match_array messages
+    ensure
+      Kafka::ConsumerGroup::Assignor.strategy_classes.delete("custom")
+    end
   end
 
   def wait_until(timeout:)

spec/round_robin_assignment_strategy_spec.rb

@@ -7,7 +7,7 @@
     members = Hash[(0...10).map {|i| ["member#{i}", nil] }]
     partitions = (0...30).map {|i| double(:"partition#{i}", topic: "greetings", partition_id: i) }
 
-    assignments = strategy.assign(members: members, partitions: partitions)
+    assignments = strategy.assign(cluster: nil, members: members, partitions: partitions)
 
     partitions.each do |partition|
       member = assignments.values.find {|assigned_partitions|
@@ -27,7 +27,7 @@
       double(:"partition#{i}", topic: topic, partition_id: i)
     }
 
-    assignments = strategy.assign(members: members, partitions: partitions)
+    assignments = strategy.assign(cluster: nil, members: members, partitions: partitions)
 
     partitions.each do |partition|
       member = assignments.values.find {|assigned_partitions|
@@ -85,7 +85,7 @@
           }
         }
 
-        assignments = strategy.assign(members: members, partitions: partitions)
+        assignments = strategy.assign(cluster: nil, members: members, partitions: partitions)
 
         expect_all_partitions_assigned(topics, assignments)
         expect_even_assignments(topics, assignments)