ruby-kafka-ec2 0.1.2 → 0.1.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d731fe0da3282726a8bb898f98f0ab3fe0133463662d9e7bc3c536c965bfc278
-  data.tar.gz: e260eb39399056cc48cdf8b7b985b037cdb60eafc8f495771d7949e7eb581a52
+  metadata.gz: 66c71213189c16f43593597889adfb2ef3d0f4757cbf6ae7bb600310a2f88855
+  data.tar.gz: 71a0256485b92b88ed891e76bd3c3cbd85420b70332835cd33d4aeaf748e207d
 SHA512:
-  metadata.gz: 3f31a1d280d2f49b864fbc4068a4f198952180101973d2a93c708ff621238f83476dac72e8c302427f3413a6a9761f81b5a703b8020aee3edba5bf2a42e0a967
-  data.tar.gz: 9d4cda171d25754e42bda7a90c4d7b9c8a77fa8fc984a5546ccb1cbc10b4ccf10999f75c6dbaa7ecd4fdecdde285a3ccec0ead4071326564f6527382853edbe0
+  metadata.gz: c88ff1e2fe4ebd92fe6b9a13a87fd5e9582c09228ae17b29cdbd4c0186f83e22d6c96bb6f193fbd5ae2ed624979fe3c33a80c6db88e72a0800b2f878ffbfc7b1
+  data.tar.gz: 78cd8b945be174b64cdd261686683521160ed7c2c7334c6586726d1803abea4e0aad96d8470304ef324fe7b04281b7048d1591794771178c2bf71d0597f07ca7

data/README.md

@@ -118,6 +118,17 @@ 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(
+  partition_weights: {
+    "foo" => {
+      0 => 2,
+    },
+  }
+)
+```
 
 ## Development
 

data/lib/kafka/ec2/mixed_instance_assignment_strategy.rb

@@ -23,11 +23,12 @@ 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: {})
+      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.
@@ -41,6 +42,7 @@ module Kafka
         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
@@ -50,6 +52,7 @@ module Kafka
 
           instance_id, instance_type, az = member_id_to_metadata[member_id].split(",")
           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
@@ -64,25 +67,48 @@ 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_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]
-          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
+
+              group_assignment[member_id].assign(topic, [partition])
+              acceptable_partition_weight -= partition_weight
 
-          last_index += partition_count
+              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_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])
+        while last_index < topic_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_weight = partition_weights.dig(topic, partition)
+          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
@@ -97,7 +123,18 @@ module Kafka
         instance_family, _ = instance_type.split(".")
 
         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)
+        (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

@@ -9,9 +9,11 @@ module Kafka
       # @param availability_zone_weights [Hash, Proc]
       # @param weights [Hash, Proc]
       # @see Kafka::EC2::MixedInstanceAssignmentStrategy#initialize
-      def initialize(instance_family_weights: {}, availability_zone_weights: {}, weights: {})
+      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:)
@@ -19,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.2"
+    VERSION = "0.1.7"
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: ruby-kafka-ec2
 version: !ruby/object:Gem::Version
-  version: 0.1.2
+  version: 0.1.7
 platform: ruby
 authors:
 - abicky
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2020-07-05 00:00:00.000000000 Z
+date: 2021-03-16 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.1.4
 signing_key: 
 specification_version: 4
 summary: An extension of ruby-kafka for EC2