Use ruby-kafka 1.3.0 or later

ruby-kafka 1.3.0 includes the changes on
zendesk/ruby-kafka#846, and they make
ruby-kafka-ec2 much simpler.

Author: abicky

README.md

@@ -24,17 +24,17 @@ Or install it yourself as:
 
 ### Kafka::EC2::MixedInstanceAssignmentStrategy
 
-`Kafka::EC2::MixedInstanceAssignmentStrategy` is an assignor for auto-scaling groups with mixed instance policies. The throughputs of consumers usually depend on instance families and availability zones. For example, if your application writes data to a database, the throughputs of consumers running on the same availability zone as the writer DB instance is higher.
+`Kafka::EC2::MixedInstanceAssignmentStrategy` is an assignor for auto-scaling groups with mixed instance policies. The throughputs of consumers usually depend on instance families and availability zones. For example, if your application writes data to a database, the throughputs of consumers running on the same availability zone as that of the writer DB instance is higher.
 
-To assign more partitions to consumers with high throughputs, you have to define `Kafka::EC2::MixedInstanceAssignmentStrategyFactory` first like below:
+To assign more partitions to consumers with high throughputs, you have to initialize `Kafka::EC2::MixedInstanceAssignmentStrategy` first like below:
 
 ```ruby
 require "aws-sdk-rds"
 require "kafka"
 require "kafka/ec2"
 
 rds = Aws::RDS::Client.new(region: "ap-northeast-1")
-assignment_strategy_factory = Kafka::EC2::MixedInstanceAssignmentStrategyFactory.new(
+assignment_strategy = Kafka::EC2::MixedInstanceAssignmentStrategy.new(
   instance_family_weights: {
     "r4" => 1.00,
     "r5" => 1.20,
@@ -68,19 +68,17 @@ assignment_strategy_factory = Kafka::EC2::MixedInstanceAssignmentStrategyFactory
 
 In the preceding example, consumers running on c5 instances will have 1.5x as many partitions compared to consumers running on r4 instances. In a similar way, if the writer DB instance is in ap-northeast-1a, consumers in ap-northeast-1a will have 4x as many partitions compared to consumers in ap-northeast-1c.
 
-You can use `Kafka::EC2::MixedInstanceAssignmentStrategy` by specifying the factory to `Kafka::EC2.with_assignment_strategy_factory` and creating a consumer in the block:
+You can use `Kafka::EC2::MixedInstanceAssignmentStrategy` by specifying it to `Kafka#consumer`:
 
 
 ```ruby
-consumer = Kafka::EC2.with_assignment_strategy_factory(assignment_strategy_factory) do
-  kafka.consumer(group_id: ENV["KAFKA_CONSUMER_GROUP_ID"])
-end
+consumer = kafka.consumer(group_id: ENV["KAFKA_CONSUMER_GROUP_ID"], assignment_strategy: assignment_strategy)
 ```
 
 You can also specify weights for each combination of availability zones and instance families:
 
 ```ruby
-assignment_strategy_factory = Kafka::EC2::MixedInstanceAssignmentStrategyFactory.new(
+assignment_strategy = Kafka::EC2::MixedInstanceAssignmentStrategy.new(
   weights: ->() {
     db_cluster = rds.describe_db_clusters(filters: [
       { name: "db-cluster-id", values: [ENV["RDS_CLUSTER"]] },
@@ -121,7 +119,7 @@ assignment_strategy_factory = Kafka::EC2::MixedInstanceAssignmentStrategyFactory
 The strategy also has the option `partition_weights`. This is useful when the topic has some skewed partitions. Suppose the partition with ID 0 of the topic "foo" receives twice as many records as other partitions. To reduce the number of partitions assigned to the consumer that consumes the partition with ID 0, specify `partition_weights` like below:
 
 ```ruby
-assignment_strategy_factory = Kafka::EC2::MixedInstanceAssignmentStrategyFactory.new(
+assignment_strategy = Kafka::EC2::MixedInstanceAssignmentStrategy.new(
   partition_weights: {
     "foo" => {
       0 => 2,

lib/kafka/ec2.rb

@@ -1,23 +1,7 @@
-require "kafka/ec2/ext/consumer_group"
-require "kafka/ec2/ext/protocol/join_group_request"
-require "kafka/ec2/mixed_instance_assignment_strategy_factory"
+require "kafka/ec2/mixed_instance_assignment_strategy"
 require "kafka/ec2/version"
 
 module Kafka
   class EC2
-    class << self
-      attr_reader :assignment_strategy_factory
-
-      def with_assignment_strategy_factory(factory)
-        @assignment_strategy_factory = factory
-        yield
-      ensure
-        @assignment_strategy_factory = nil
-      end
-
-      def assignment_strategy_classes
-        @assignment_strategy_classes ||= {}
-      end
-    end
   end
 end

lib/kafka/ec2/ext/consumer_group.rb

@@ -6,10 +6,8 @@
 module Kafka
   class EC2
     class MixedInstanceAssignmentStrategy
-      # metadata is a byte sequence created by Kafka::Protocol::ConsumerGroupProtocol.encode
-      attr_accessor :member_id_to_metadata
+      DELIMITER = ","
 
-      # @param cluster [Kafka::Cluster]
       # @param instance_family_weights [Hash{String => Numeric}, Proc] a hash whose the key
       #   is the instance family and whose value is the weight. If the object is a proc,
       #   it must returns such a hash and the proc is called every time the method "assign"
@@ -23,22 +21,35 @@ class MixedInstanceAssignmentStrategy
       #   instance_family_weights or availability_zone_weights. If the object is a proc,
       #   it must returns such a hash and the proc is called every time the method "assign"
       #   is called.
-      def initialize(cluster:, instance_family_weights: {}, availability_zone_weights: {}, weights: {}, partition_weights: {})
-        @cluster = cluster
+      def initialize(instance_family_weights: {}, availability_zone_weights: {}, weights: {}, partition_weights: {})
         @instance_family_weights = instance_family_weights
         @availability_zone_weights = availability_zone_weights
         @weights = weights
         @partition_weights = partition_weights
       end
 
+      def protocol_name
+        "mixedinstance"
+      end
+
+      def user_data
+        Net::HTTP.start("169.254.169.254", 80) do |http|
+          [
+            http.get("/latest/meta-data/instance-id").body,
+            http.get("/latest/meta-data/instance-type").body,
+            http.get("/latest/meta-data/placement/availability-zone").body,
+          ].join(DELIMITER)
+        end
+      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 = {}
+      def call(cluster:, members:, partitions:)
+        member_id_to_partitions = Hash.new { |h, k| h[k] = [] }
         instance_id_to_capacity = Hash.new(0)
         instance_id_to_member_ids = Hash.new { |h, k| h[k] = [] }
         total_capacity = 0
@@ -47,28 +58,17 @@ def assign(members:, topics:)
         instance_family_to_capacity = @instance_family_weights.is_a?(Proc) ? @instance_family_weights.call() : @instance_family_weights
         az_to_capacity = @availability_zone_weights.is_a?(Proc) ? @availability_zone_weights.call() : @availability_zone_weights
         weights = @weights.is_a?(Proc) ? @weights.call() : @weights
-        members.each do |member_id|
-          group_assignment[member_id] = Protocol::MemberAssignment.new
-
-          instance_id, instance_type, az = member_id_to_metadata[member_id].split(",")
+        members.each do |member_id, metadata|
+          instance_id, instance_type, az = metadata.user_data.split(DELIMITER)
           instance_id_to_member_ids[instance_id] << member_id
           member_id_to_instance_id[member_id] = instance_id
           capacity = calculate_capacity(instance_type, az, instance_family_to_capacity, az_to_capacity, weights)
           instance_id_to_capacity[instance_id] += capacity
           total_capacity += capacity
         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
-
-        partition_weights = build_partition_weights(topics)
-        partition_weight_per_capacity = topic_partitions.sum { |topic, partition| partition_weights.dig(topic, partition) } / total_capacity
+        partition_weights = build_partition_weights(partitions)
+        partition_weight_per_capacity = partitions.sum { |partition| partition_weights.dig(partition.topic, partition.partition_id) } / total_capacity
 
         last_index = 0
         member_id_to_acceptable_partition_weight = {}
@@ -77,12 +77,12 @@ def assign(members:, topics:)
           member_ids = instance_id_to_member_ids[instance_id]
           member_ids.each do |member_id|
             acceptable_partition_weight = capacity * partition_weight_per_capacity / member_ids.size
-            while last_index < topic_partitions.size
-              topic, partition = topic_partitions[last_index]
-              partition_weight = partition_weights.dig(topic, partition)
+            while last_index < partitions.size
+              partition = partitions[last_index]
+              partition_weight = partition_weights.dig(partition.topic, partition.partition_id)
               break if acceptable_partition_weight - partition_weight < 0
 
-              group_assignment[member_id].assign(topic, [partition])
+              member_id_to_partitions[member_id] << partition
               acceptable_partition_weight -= partition_weight
 
               last_index += 1
@@ -93,25 +93,25 @@ def assign(members:, topics:)
           end
         end
 
-        while last_index < topic_partitions.size
+        while last_index < partitions.size
           max_acceptable_partition_weight = member_id_to_acceptable_partition_weight.values.max
           member_ids = member_id_to_acceptable_partition_weight.select { |_, w| w == max_acceptable_partition_weight }.keys
           if member_ids.size == 1
             member_id = member_ids.first
           else
             member_id =  member_ids.max_by { |id| instance_id_to_total_acceptable_partition_weight[member_id_to_instance_id[id]] }
           end
-          topic, partition = topic_partitions[last_index]
-          group_assignment[member_id].assign(topic, [partition])
+          partition = partitions[last_index]
+          member_id_to_partitions[member_id] << partition
 
-          partition_weight = partition_weights.dig(topic, partition)
+          partition_weight = partition_weights.dig(partition.topic, partition.partition_id)
           member_id_to_acceptable_partition_weight[member_id] -= partition_weight
           instance_id_to_total_acceptable_partition_weight[member_id_to_instance_id[member_id]] -= partition_weight
 
           last_index += 1
         end
 
-        group_assignment
+        member_id_to_partitions
       rescue Kafka::LeaderNotAvailable
         sleep 1
         retry
@@ -126,12 +126,12 @@ def calculate_capacity(instance_type, az, instance_family_to_capacity, az_to_cap
         (capacity || instance_family_to_capacity.fetch(instance_family, 1) * az_to_capacity.fetch(az, 1)).to_f
       end
 
-      def build_partition_weights(topics)
+      def build_partition_weights(partitions)
         # Duplicate the weights to not destruct @partition_weights or the return value of @partition_weights
-        weights = (@partition_weights.is_a?(Proc) ? @partition_weights.call() : @partition_weights).dup
-        topics.each do |t|
-          weights[t] = weights[t].dup || {}
-          weights[t].default = 1
+        weights = (@partition_weights.is_a?(Proc) ? @partition_weights.call : @partition_weights).dup
+        partitions.map(&:topic).uniq.each do |topic|
+          weights[topic] = weights[topic].dup || {}
+          weights[topic].default = 1
         end
 
         weights