ruby-kafka-ec2 0.1.5 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 8dde731c3652090bf18202d68b916cbdcff9ed09673bd84d5f20470a37c63373
-  data.tar.gz: d1a95de4724b3b5f85230c55a70469cc5e6c1e6008423b83f74c415bf2c9d289
+  metadata.gz: e7732059807b7aad8dfe8df2758fa0dc9ad8f8063adaf9cf71b615c8384e74aa
+  data.tar.gz: feb725eb274ff28e6b3e5827f02d9c1b406c127f7109c6ec6a4001054665b125
 SHA512:
-  metadata.gz: f37b8fa41b773933aac85f170884adb75fc0e446faf9fb92c109aa039f5a869874194dbdf3a9099899e273ef8543f75c5f7aca0fd99cff1845bc43ac081bde50
-  data.tar.gz: 78bc5df7157441563d73e19f35804069ddbb2c1863bccfe2711594c27a6caed78a5209fdd2ca7f55a3cf2302bed326ec47302720bf4253a9f8b4df950e5a0d0f
+  metadata.gz: eade4b284de35a438d52b4f18928cb9e287fe65064c1a7baec50545627bd7996835150326e88499b2ee13d4499dd1aa1093471fdd829b8b2a545eecee1799a5b
+  data.tar.gz: fed5703514978a1986a720678d3590c29fcc9359802fe9b98a3a6632f7dcd28cf0f43151f19a8c5a1e1d9a609667a13598b3780589ff57480d928955f6219f8f

data/README.md

@@ -24,9 +24,9 @@ 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"
@@ -34,7 +34,7 @@ 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,

data/lib/kafka/ec2/mixed_instance_assignment_strategy.rb

@@ -6,10 +6,8 @@ require "kafka/protocol/member_assignment"
 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,82 +21,97 @@ module Kafka
       #   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
+        member_id_to_instance_id = {}
 
         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 = {}
+        instance_id_to_total_acceptable_partition_weight = Hash.new(0)
         instance_id_to_capacity.each do |instance_id, capacity|
           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
-            loop do
-              topic, partition = topic_partitions[last_index]
-              partition_weight = partition_weights.dig(topic, partition)
-              if last_index == topic_partitions.size || acceptable_partition_weight - partition_weight < 0
-                member_id_to_acceptable_partition_weight[member_id] = acceptable_partition_weight
-                break
-              end
-
-              group_assignment[member_id].assign(topic, [partition])
-              last_index += 1
+            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
+
+              member_id_to_partitions[member_id] << partition
               acceptable_partition_weight -= partition_weight
+
+              last_index += 1
             end
+
+            member_id_to_acceptable_partition_weight[member_id] = acceptable_partition_weight
+            instance_id_to_total_acceptable_partition_weight[instance_id] += acceptable_partition_weight
           end
         end
 
-        if last_index < topic_partitions.size
-          member_id_to_acceptable_partition_weight.sort_by { |_, remaining| -remaining }.each do |member_id, _|
-            topic, partition = topic_partitions[last_index]
-            group_assignment[member_id].assign(topic, [partition])
-            last_index += 1
-            break if 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
+          partition = partitions[last_index]
+          member_id_to_partitions[member_id] << 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
@@ -113,12 +126,12 @@ module Kafka
         (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

data/lib/kafka/ec2/version.rb

@@ -1,5 +1,5 @@
 module Kafka
   class EC2
-    VERSION = "0.1.5"
+    VERSION = "0.2.0"
   end
 end

data/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

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: ruby-kafka-ec2
 version: !ruby/object:Gem::Version
-  version: 0.1.5
+  version: 0.2.0
 platform: ruby
 authors:
 - abicky
-autorequire: 
+autorequire:
 bindir: exe
 cert_chain: []
-date: 2020-10-08 00:00:00.000000000 Z
+date: 2022-03-29 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: ruby-kafka
@@ -89,10 +89,7 @@ files:
 - bin/console
 - bin/setup
 - lib/kafka/ec2.rb
-- lib/kafka/ec2/ext/consumer_group.rb
-- lib/kafka/ec2/ext/protocol/join_group_request.rb
 - lib/kafka/ec2/mixed_instance_assignment_strategy.rb
-- lib/kafka/ec2/mixed_instance_assignment_strategy_factory.rb
 - lib/kafka/ec2/version.rb
 - ruby-kafka-ec2.gemspec
 homepage: https://github.com/abicky/ruby-kafka-ec2
@@ -101,7 +98,7 @@ licenses:
 metadata:
   homepage_uri: https://github.com/abicky/ruby-kafka-ec2
   source_code_uri: https://github.com/abicky/ruby-kafka-ec2
-post_install_message: 
+post_install_message:
 rdoc_options: []
 require_paths:
 - lib
@@ -116,8 +113,8 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.0.3
-signing_key: 
+rubygems_version: 3.2.22
+signing_key:
 specification_version: 4
 summary: An extension of ruby-kafka for EC2
 test_files: []

data/lib/kafka/ec2/ext/consumer_group.rb

@@ -1,33 +0,0 @@
-# frozen_string_literal: true
-
-require "kafka/consumer_group"
-require "kafka/ec2/mixed_instance_assignment_strategy"
-
-module Kafka
-  class EC2
-    module Ext
-      module ConsumerGroup
-        def initialize(*args, **kwargs)
-          super
-          if Kafka::EC2.assignment_strategy_factory
-            @assignment_strategy = Kafka::EC2.assignment_strategy_factory.create(cluster: @cluster)
-          end
-          Kafka::EC2.assignment_strategy_classes[@group_id] = @assignment_strategy.class
-        end
-
-        def join_group
-          super
-          if Kafka::EC2.assignment_strategy_classes[@group_id] == Kafka::EC2::MixedInstanceAssignmentStrategy
-            @assignment_strategy.member_id_to_metadata = @members
-          end
-        end
-      end
-    end
-  end
-end
-
-module Kafka
-  class ConsumerGroup
-    prepend Kafka::EC2::Ext::ConsumerGroup
-  end
-end

data/lib/kafka/ec2/ext/protocol/join_group_request.rb

@@ -1,39 +0,0 @@
-# frozen_string_literal: true
-
-require "net/http"
-
-require "kafka/protocol/consumer_group_protocol"
-require "kafka/protocol/join_group_request"
-
-module Kafka
-  class EC2
-    module Ext
-      module Protocol
-        module JoinGroupRequest
-          def initialize(*args, topics: [], **kwargs)
-            super
-            if Kafka::EC2.assignment_strategy_classes[@group_id] == Kafka::EC2::MixedInstanceAssignmentStrategy
-              user_data = Net::HTTP.start("169.254.169.254", 80) do |http|
-                instance_id = http.get("/latest/meta-data/instance-id").body
-                instance_type = http.get("/latest/meta-data/instance-type").body
-                az = http.get("/latest/meta-data/placement/availability-zone").body
-                "|#{instance_id},#{instance_type},#{az}"
-              end
-              @group_protocols = {
-                "mixedinstance" => Kafka::Protocol::ConsumerGroupProtocol.new(topics: topics, user_data: user_data),
-              }
-            end
-          end
-        end
-      end
-    end
-  end
-end
-
-module Kafka
-  module Protocol
-    class JoinGroupRequest
-      prepend Kafka::EC2::Ext::Protocol::JoinGroupRequest
-    end
-  end
-end

data/lib/kafka/ec2/mixed_instance_assignment_strategy_factory.rb

@@ -1,30 +0,0 @@
-# frozen_string_literal: true
-
-require "kafka/ec2/mixed_instance_assignment_strategy"
-
-module Kafka
-  class EC2
-    class MixedInstanceAssignmentStrategyFactory
-      # @param instance_family_weights [Hash, Proc]
-      # @param availability_zone_weights [Hash, Proc]
-      # @param weights [Hash, Proc]
-      # @see Kafka::EC2::MixedInstanceAssignmentStrategy#initialize
-      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 create(cluster:)
-        Kafka::EC2::MixedInstanceAssignmentStrategy.new(
-          cluster: cluster,
-          instance_family_weights: @instance_family_weights,
-          availability_zone_weights: @availability_zone_weights,
-          weights: @weights,
-          partition_weights: @partition_weights,
-        )
-      end
-    end
-  end
-end