ruby-kafka-ec2 0.1.1 → 0.1.6

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d8620cd2b3622df81ce8105a45930522ff9c9b433efbcba2c356c52bd7bdc7df
-  data.tar.gz: 5ccbcce906e3f000d445fe3a2135f4da2d7f7dbbbe3c7267849fc59bd9aabc9d
+  metadata.gz: b87f937f29a107174e7a972b79e2a717780bdd28b28fa37e3601633ff1f075b0
+  data.tar.gz: 413c79ff75ccb9f9304dd5c57dfb2fb1f26b0b673eeae9cb7e7d86cb68d65e74
 SHA512:
-  metadata.gz: d085d3c8a0bcd5a3e0ab00f1b89d3cba53af8f5a33e8bf78bbfa5e03e96c027a3a747ef34a730400e4567e3b6533615fc6477c57eb4164da2a3d9f1b43ab71cb
-  data.tar.gz: a1343b55e46e61346f1b5b179b15cfcc838ecfb39650063369fa9dd92c0b786b3872bcf0330bc7039ca66a856e9f3657e3249a6dc927460bb2f158186fb4930c
+  metadata.gz: 22dfa6ddf4e0e0f8eb5db7bd6cb9c54b25edb175186bba989b32df698e23c705754c8dff057e7d82df480a3835bea110a363177b36739bca73432206227ac680
+  data.tar.gz: 1deda3ac8da64c41f9d598e26049b667aec2d81956a98dc7ded279c8a13e517d7dc4042ab34f9f4e8a05365be46c570e274207bd5bd430210b7de458972cace3

data/README.md

@@ -77,6 +77,59 @@ consumer = Kafka::EC2.with_assignment_strategy_factory(assignment_strategy_facto
 end
 ```
 
+You can also specify weights for each combination of availability zones and instance families:
+
+```ruby
+assignment_strategy_factory = Kafka::EC2::MixedInstanceAssignmentStrategyFactory.new(
+  weights: ->() {
+    db_cluster = rds.describe_db_clusters(filters: [
+      { name: "db-cluster-id", values: [ENV["RDS_CLUSTER"]] },
+    ]).db_clusters.first
+    db_instance_id = db_cluster.db_cluster_members.find { |m| m.is_cluster_writer }.db_instance_identifier
+    db_instance = rds.describe_db_instances(filters: [
+      { name: "db-cluster-id", values: [ENV["RDS_CLUSTER"]] },
+      { name: "db-instance-id", values: [db_instance_id] },
+    ]).db_instances.first
+
+    weights_for_writer_az = {
+      "r4" => 1.00,
+      "r5" => 1.20,
+      "m5" => 1.35,
+      "c5" => 1.50,
+    }
+    weights_for_other_az = {
+      "r4" => 0.40,
+      "r5" => 0.70,
+      "m5" => 0.80,
+      "c5" => 1.00,
+    }
+    if db_instance.availability_zone == "ap-northeast-1a"
+      {
+        "ap-northeast-1a" => weights_for_writer_az,
+        "ap-northeast-1c" => weights_for_other_az,
+      }
+    else
+      {
+        "ap-northeast-1a" => weights_for_other_az,
+        "ap-northeast-1c" => weights_for_writer_az,,
+      }
+    end
+  },
+)
+```
+
+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(
+  partition_weights: {
+    "foo" => {
+      0 => 2,
+    },
+  }
+)
+```
+
 ## Development
 
 After checking out the repo, run `bin/setup` to install dependencies. Then, run `rake spec` to run the tests. You can also run `bin/console` for an interactive prompt that will allow you to experiment.

data/lib/kafka/ec2/mixed_instance_assignment_strategy.rb

@@ -18,10 +18,17 @@ module Kafka
       #   is the availability zone 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"
       #   is called.
-      def initialize(cluster:, instance_family_weights:, availability_zone_weights:)
+      # @param weights [Hash{String => Hash{String => Numeric}}, Proc] a hash whose the key
+      #   is the availability zone or the instance family and whose value is the hash like
+      #   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
         @instance_family_weights = instance_family_weights
         @availability_zone_weights = availability_zone_weights
+        @weights = weights
+        @partition_weights = partition_weights
       end
 
       # Assign the topic partitions to the group members.
@@ -38,12 +45,13 @@ module Kafka
 
         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(",")
           instance_id_to_member_ids[instance_id] << member_id
-          capacity = calculate_capacity(instance_type, az, instance_family_to_capacity, az_to_capacity)
+          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
@@ -57,27 +65,40 @@ module Kafka
           Array.new(partitions.count) { topic }.zip(partitions)
         end
 
-        partition_count_per_capacity = topic_partitions.size / total_capacity
+        partition_weights = build_partition_weights(topics)
+        partition_weight_per_capacity = topic_partitions.sum { |topic, partition| partition_weights.dig(topic, partition) } / total_capacity
+
         last_index = 0
-        instance_id_to_capacity.sort_by { |_, capacity| -capacity }.each do |instance_id, capacity|
-          partition_count = (capacity * partition_count_per_capacity).round
+        member_id_to_acceptable_partition_weight = {}
+        instance_id_to_capacity.each do |instance_id, capacity|
           member_ids = instance_id_to_member_ids[instance_id]
-          topic_partitions[last_index, partition_count]&.each_with_index do |(topic, partition), index|
-            member_id = member_ids[index % member_ids.size]
-            group_assignment[member_id].assign(topic, [partition])
-          end
+          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)
+              break if acceptable_partition_weight - partition_weight < 0
 
-          last_index += partition_count
-        end
+              group_assignment[member_id].assign(topic, [partition])
+              acceptable_partition_weight -= partition_weight
+
+              last_index += 1
+            end
 
-        if last_index < topic_partitions.size
-          member_ids = instance_id_to_member_ids.values.flatten
-          topic_partitions[last_index, topic_partitions.size].each_with_index do |(topic, partition), index|
-            member_id = member_ids[index % member_ids.size]
-            group_assignment[member_id].assign(topic, [partition])
+            member_id_to_acceptable_partition_weight[member_id] = acceptable_partition_weight
           end
         end
 
+        while last_index < topic_partitions.size
+          member_id, _ = member_id_to_acceptable_partition_weight.max_by { |_, remaining| remaining }
+          topic, partition = topic_partitions[last_index]
+          group_assignment[member_id].assign(topic, [partition])
+
+          member_id_to_acceptable_partition_weight[member_id] -= partition_weights.dig(topic, partition)
+
+          last_index += 1
+        end
+
         group_assignment
       rescue Kafka::LeaderNotAvailable
         sleep 1
@@ -86,9 +107,22 @@ module Kafka
 
       private
 
-      def calculate_capacity(instance_type, az, instance_family_to_capacity, az_to_capacity)
+      def calculate_capacity(instance_type, az, instance_family_to_capacity, az_to_capacity, weights)
         instance_family, _ = instance_type.split(".")
-        instance_family_to_capacity.fetch(instance_family, 1) * az_to_capacity.fetch(az, 1)
+
+        capacity = weights.dig(az, instance_family) || weights.dig(instance_family, az)
+        (capacity || instance_family_to_capacity.fetch(instance_family, 1) * az_to_capacity.fetch(az, 1)).to_f
+      end
+
+      def build_partition_weights(topics)
+        # 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
+        end
+
+        weights
       end
     end
   end

data/lib/kafka/ec2/mixed_instance_assignment_strategy_factory.rb

@@ -7,9 +7,13 @@ module Kafka
     class MixedInstanceAssignmentStrategyFactory
       # @param instance_family_weights [Hash, Proc]
       # @param availability_zone_weights [Hash, Proc]
-      def initialize(instance_family_weights: {}, availability_zone_weights: {})
+      # @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:)
@@ -17,6 +21,8 @@ module Kafka
           cluster: cluster,
           instance_family_weights: @instance_family_weights,
           availability_zone_weights: @availability_zone_weights,
+          weights: @weights,
+          partition_weights: @partition_weights,
         )
       end
     end

data/lib/kafka/ec2/version.rb

@@ -1,5 +1,5 @@
 module Kafka
   class EC2
-    VERSION = "0.1.1"
+    VERSION = "0.1.6"
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: ruby-kafka-ec2
 version: !ruby/object:Gem::Version
-  version: 0.1.1
+  version: 0.1.6
 platform: ruby
 authors:
 - abicky
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2020-06-23 00:00:00.000000000 Z
+date: 2020-10-08 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: ruby-kafka
@@ -116,7 +116,7 @@ required_rubygems_version: !ruby/object:Gem::Requirement
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.1.2
+rubygems_version: 3.0.3
 signing_key: 
 specification_version: 4
 summary: An extension of ruby-kafka for EC2