ruby-kafka-ec2 0.1.3 → 0.1.4

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 0e6d39fb8a782c1075f0454a2bc0479fafb12d380ad0bd78ba938d218af57e23
-  data.tar.gz: 2f339e9dc2ea096376b53fb4f9413a46da5f3bdcd7ff86c6b4854f3d8928cbf7
+  metadata.gz: 41d57e6a8d322551d668003fbba48e43507a4d54e73f3dd2d4918fe55bb6aaad
+  data.tar.gz: 1440f7a70cdbc3f47a524fc1847bcb34cbc1d55dbef5c04660c384c073018434
 SHA512:
-  metadata.gz: f77547fc1fb5b92b12b12857114218b0d3d055863ffe8c3a972a3255cbc781d6e441ae11ce5e8772add6d5be26113e5f874a3a333da0c1a5d6483e0b187141cd
-  data.tar.gz: 1a87e14365f4a49b139787d55aae75ac8e4131a4ed6fd766a9a55a297bcf4f613a2c5ecf062f84fb5fd69cfc792d2d623f8aef2a1285f2126f768da1e75566af
+  metadata.gz: 9a731656c0e0862559ea88c5395d547886cc738bb39ba4ae779056ea2d7376c8654bc20698f470b13f03bb79cea5384e3d78047c4a6703dfba43de830f54dfee
+  data.tar.gz: e78667585584f906f81c33a7eb09a254fc59034680b7ec4f1bcb39a0cb81e1b865d5adf1f17c55ddae0e179f680b309d592f04525d386549e10ba842eef2677c

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.
@@ -64,24 +65,35 @@ 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
+            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
 
-          last_index += partition_count
+              group_assignment[member_id].assign(topic, [partition])
+              last_index += 1
+              acceptable_partition_weight -= partition_weight
+            end
+          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]
+          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
           end
         end
 
@@ -97,7 +109,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,10 +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:)
@@ -21,6 +22,7 @@ module Kafka
           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.3"
+    VERSION = "0.1.4"
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: ruby-kafka-ec2
 version: !ruby/object:Gem::Version
-  version: 0.1.3
+  version: 0.1.4
 platform: ruby
 authors:
 - abicky
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2020-07-07 00:00:00.000000000 Z
+date: 2020-10-07 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