Remove Kafka::ConsumerGroup::Assignor.register_strategy

See zendesk#846 (comment)
for the reason.

Author: abicky

README.md

@@ -747,47 +747,86 @@ 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 like below:
+You can use a custom assignment strategy by passing an object that implements `#call` as the argument `assignment_strategy` like below:
 
 ```ruby
-Kafka::ConsumerGroup::Assignor.register_strategy(:custom) do |cluster:, members:, partitions:|
-  # strategy goes here
+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
-consumer = kafka.consumer(group_id: "some-group", assignment_strategy: :custom)
+
+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 block should return a hash mapping member ids to partitions.
-For example, the following strategy assigns partitions randomly:
+`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
-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
+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 can pass the option `user_data`, a proc that returns user data on each consumer:
+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
-Kafka::ConsumerGroup::Assignor.register_strategy(:custom, user_data: ->{ ... }) do |cluster:, members:, partitions:|
-  # strategy goes here
+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
 ```
 
-For example, the following strategy displays the IP address of each consumer:
+Note that the strategy uses the class name as the default protocol name. You can change it by defining the method `protocol_name`:
 
 ```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}"
+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
 
 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.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 || :roundrobin
+        strategy: assignment_strategy || RoundRobinAssignmentStrategy.new
       )
       @retention_time = retention_time
     end

lib/kafka/consumer_group/assignor.rb

@@ -7,50 +7,22 @@ class ConsumerGroup
 
     # A consumer group partition assignor
     class Assignor
-      Strategy = Struct.new(:assign, :user_data)
       Partition = Struct.new(:topic, :partition_id)
 
-      # @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 block_given?
-          raise ArgumentError, "The block must be given."
-        end
-
-        strategies[name.to_s] = Strategy.new(block, user_data)
-      end
-
-      def self.strategies
-        @strategies ||= {}
-      end
-
-      attr_reader :protocol_name
-
       # @param cluster [Kafka::Cluster]
-      # @param strategy [String, Symbol, Object]  a string, a symbol, or an object that implements
-      #   #user_data and #assign.
+      # @param strategy [Object] an object that implements #protocol_type,
+      #   #user_data, and #assign.
       def initialize(cluster:, strategy:)
         @cluster = cluster
-        @protocol_name = strategy.to_s
-        @strategy = self.class.strategies.fetch(@protocol_name) do
-          raise ArgumentError, "The strategy '#{@protocol_name}' is not registered."
-        end
+        @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.call if @strategy.user_data
+        @strategy.user_data if @strategy.respond_to?(:user_data)
       end
 
       # Assign the topic partitions to the group members.
@@ -75,8 +47,7 @@ def assign(members:, topics:)
         members.each_key do |member_id|
           group_assignment[member_id] = Protocol::MemberAssignment.new
         end
-
-        @strategy.assign.call(cluster: @cluster, members: members, partitions: topic_partitions).each do |member_id, partitions|
+        @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

lib/kafka/round_robin_assignment_strategy.rb

@@ -5,6 +5,10 @@ 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
+
     # Assign the topic partitions to the group members.
     #
     # @param cluster [Kafka::Cluster]
@@ -14,7 +18,7 @@ class RoundRobinAssignmentStrategy
     #   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:)
+    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|
@@ -25,9 +29,3 @@ def assign(cluster:, members:, partitions:)
     end
   end
 end
-
-require "kafka/consumer_group/assignor"
-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

@@ -1,46 +1,72 @@
 # frozen_string_literal: true
 
 describe Kafka::ConsumerGroup::Assignor do
-  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.strategies.delete("custom")
+  let(:cluster) { double(:cluster) }
+  let(:assignor) { described_class.new(cluster: cluster, strategy: strategy) }
+
+  describe "#protocol_name" do
+    subject(:protocol_name) { assignor.protocol_name }
+
+    context "when the strategy has the method" do
+      let(:strategy) { double(protocol_name: "custom") }
+
+      it "returns the return value" do
+        expect(protocol_name).to eq "custom"
       end
+    end
+
+    context "when the strategy has the method" do
+      let(:strategy) { Class.new.new }
 
-      it do
-        expect { described_class.register_strategy(:custom) { |**kwargs| [] } }.to raise_error(ArgumentError)
+      it "returns the return value" do
+        expect(protocol_name).to match /\A#<Class:0x[0-9a-f]+>\z/
       end
     end
   end
 
-  describe "#assign" do
-    let(:cluster) { double(:cluster) }
-    let(:assignor) { described_class.new(cluster: cluster, strategy: :custom) }
+  describe "#user_data" do
+    subject(:user_data) { assignor.user_data }
 
-    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) } }
+    context "when the strategy has the method" do
+      let(:strategy) { double(user_data: "user_data") }
 
-    before do
-      allow(cluster).to receive(:partitions_for) { partitions }
+      it "returns the return value" do
+        expect(user_data).to eq "user_data"
+      end
+    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
+    context "when the strategy has the method" do
+      let(:strategy) { Class.new.new }
 
-        assignment
+      it "returns the return value" do
+        expect(user_data).to be_nil
       end
     end
-    after do
-      described_class.strategies.delete("custom")
+  end
+
+  describe "#assign" do
+    let(:strategy) do
+      klass = Class.new do
+        def call(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
+      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) }
+
+      allow(cluster).to receive(:partitions_for) { partitions }
+
       assignments = assignor.assign(members: members, topics: topics)
 
       partitions.each do |partition|