ruby-kafka 0.1.7 → 0.2.0

data/lib/kafka/protocol/join_group_response.rb

@@ -0,0 +1,31 @@
+module Kafka
+  module Protocol
+    class JoinGroupResponse
+      attr_reader :error_code
+
+      attr_reader :generation_id, :group_protocol
+
+      attr_reader :leader_id, :member_id, :members
+
+      def initialize(error_code:, generation_id:, group_protocol:, leader_id:, member_id:, members:)
+        @error_code = error_code
+        @generation_id = generation_id
+        @group_protocol = group_protocol
+        @leader_id = leader_id
+        @member_id = member_id
+        @members = members
+      end
+
+      def self.decode(decoder)
+        new(
+          error_code: decoder.int16,
+          generation_id: decoder.int32,
+          group_protocol: decoder.string,
+          leader_id: decoder.string,
+          member_id: decoder.string,
+          members: Hash[decoder.array { [decoder.string, decoder.bytes] }],
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/leave_group_request.rb

@@ -0,0 +1,23 @@
+module Kafka
+  module Protocol
+    class LeaveGroupRequest
+      def initialize(group_id:, member_id:)
+        @group_id = group_id
+        @member_id = member_id
+      end
+
+      def api_key
+        13
+      end
+
+      def response_class
+        LeaveGroupResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_string(@member_id)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/leave_group_response.rb

@@ -0,0 +1,15 @@
+module Kafka
+  module Protocol
+    class LeaveGroupResponse
+      attr_reader :error_code
+
+      def initialize(error_code:)
+        @error_code = error_code
+      end
+
+      def self.decode(decoder)
+        new(error_code: decoder.int16)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/member_assignment.rb

@@ -0,0 +1,40 @@
+module Kafka
+  module Protocol
+    class MemberAssignment
+      attr_reader :topics
+
+      def initialize(version: 0, topics: {}, user_data: nil)
+        @version = version
+        @topics = topics
+        @user_data = user_data
+      end
+
+      def assign(topic, partitions)
+        @topics[topic] ||= []
+        @topics[topic].concat(partitions)
+      end
+
+      def encode(encoder)
+        encoder.write_int16(@version)
+
+        encoder.write_array(@topics) do |topic, partitions|
+          encoder.write_string(topic)
+
+          encoder.write_array(partitions) do |partition|
+            encoder.write_int32(partition)
+          end
+        end
+
+        encoder.write_bytes(@user_data)
+      end
+
+      def self.decode(decoder)
+        new(
+          version: decoder.int16,
+          topics: Hash[decoder.array { [decoder.string, decoder.array { decoder.int32 }] }],
+          user_data: decoder.bytes,
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/message_set.rb

@@ -3,10 +3,8 @@ module Kafka
     class MessageSet
       attr_reader :messages
 
-      def initialize(messages: [], compression_codec: nil, compression_threshold: 1)
+      def initialize(messages: [])
         @messages = messages
-        @compression_codec = compression_codec
-        @compression_threshold = compression_threshold
       end
 
       def size
@@ -18,10 +16,10 @@ module Kafka
       end
 
       def encode(encoder)
-        if compress?
-          encode_with_compression(encoder)
-        else
-          encode_without_compression(encoder)
+        # Messages in a message set are *not* encoded as an array. Rather,
+        # they are written in sequence.
+        @messages.each do |message|
+          message.encode(encoder)
         end
       end
 
@@ -41,36 +39,6 @@ module Kafka
 
         new(messages: fetched_messages)
       end
-
-      private
-
-      def compress?
-        !@compression_codec.nil? && size >= @compression_threshold
-      end
-
-      def encode_with_compression(encoder)
-        codec = @compression_codec
-
-        buffer = StringIO.new
-        encode_without_compression(Encoder.new(buffer))
-        data = codec.compress(buffer.string)
-
-        wrapper_message = Protocol::Message.new(
-          value: data,
-          attributes: codec.codec_id,
-        )
-
-        message_set = MessageSet.new(messages: [wrapper_message])
-        message_set.encode(encoder)
-      end
-
-      def encode_without_compression(encoder)
-        # Messages in a message set are *not* encoded as an array. Rather,
-        # they are written in sequence.
-        @messages.each do |message|
-          message.encode(encoder)
-        end
-      end
     end
   end
 end

data/lib/kafka/protocol/metadata_response.rb

@@ -121,7 +121,11 @@ module Kafka
       # @param node_id [Integer] the node id of the broker.
       # @return [BrokerInfo] information about the broker.
       def find_broker(node_id)
-        @brokers.find {|broker| broker.node_id == node_id }
+        broker = @brokers.find {|broker| broker.node_id == node_id }
+
+        raise Kafka::Error, "No broker with id #{node_id}" if broker.nil?
+
+        broker
       end
 
       def partitions_for(topic_name)

data/lib/kafka/protocol/offset_commit_request.rb

@@ -0,0 +1,42 @@
+module Kafka
+  module Protocol
+    class OffsetCommitRequest
+      def api_key
+        8
+      end
+
+      def api_version
+        2
+      end
+
+      def response_class
+        OffsetCommitResponse
+      end
+
+      def initialize(group_id:, generation_id:, member_id:, retention_time: 0, offsets:)
+        @group_id = group_id
+        @generation_id = generation_id
+        @member_id = member_id
+        @retention_time = retention_time
+        @offsets = offsets
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_int32(@generation_id)
+        encoder.write_string(@member_id)
+        encoder.write_int64(@retention_time)
+
+        encoder.write_array(@offsets) do |topic, partitions|
+          encoder.write_string(topic)
+
+          encoder.write_array(partitions) do |partition, offset|
+            encoder.write_int32(partition)
+            encoder.write_int64(offset)
+            encoder.write_string(nil) # metadata
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/offset_commit_response.rb

@@ -0,0 +1,27 @@
+module Kafka
+  module Protocol
+    class OffsetCommitResponse
+      attr_reader :topics
+
+      def initialize(topics:)
+        @topics = topics
+      end
+
+      def self.decode(decoder)
+        topics = decoder.array {
+          topic = decoder.string
+          partitions = decoder.array {
+            partition = decoder.int32
+            error_code = decoder.int16
+
+            [partition, error_code]
+          }
+
+          [topic, Hash[partitions]]
+        }
+
+        new(topics: Hash[topics])
+      end
+    end
+  end
+end

data/lib/kafka/protocol/offset_fetch_request.rb

@@ -0,0 +1,34 @@
+module Kafka
+  module Protocol
+    class OffsetFetchRequest
+      def initialize(group_id:, topics:)
+        @group_id = group_id
+        @topics = topics
+      end
+
+      def api_key
+        9
+      end
+
+      def api_version
+        1
+      end
+
+      def response_class
+        OffsetFetchResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+
+        encoder.write_array(@topics) do |topic, partitions|
+          encoder.write_string(topic)
+
+          encoder.write_array(partitions) do |partition|
+            encoder.write_int32(partition)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/offset_fetch_response.rb

@@ -0,0 +1,51 @@
+module Kafka
+  module Protocol
+    class OffsetFetchResponse
+      class PartitionOffsetInfo
+        attr_reader :offset, :metadata, :error_code
+
+        def initialize(offset:, metadata:, error_code:)
+          @offset = offset
+          @metadata = metadata
+          @error_code = error_code
+        end
+      end
+
+      attr_reader :topics
+
+      def initialize(topics:)
+        @topics = topics
+      end
+
+      def offset_for(topic, partition)
+        offset_info = topics.fetch(topic).fetch(partition)
+
+        Protocol.handle_error(offset_info.error_code)
+
+        offset_info.offset
+      end
+
+      def self.decode(decoder)
+        topics = decoder.array {
+          topic = decoder.string
+
+          partitions = decoder.array {
+            partition = decoder.int32
+
+            info = PartitionOffsetInfo.new(
+              offset: decoder.int64,
+              metadata: decoder.string,
+              error_code: decoder.int16,
+            )
+
+            [partition, info]
+          }
+
+          [topic, Hash[partitions]]
+        }
+
+        new(topics: Hash[topics])
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sync_group_request.rb

@@ -0,0 +1,31 @@
+module Kafka
+  module Protocol
+    class SyncGroupRequest
+      def initialize(group_id:, generation_id:, member_id:, group_assignment: {})
+        @group_id = group_id
+        @generation_id = generation_id
+        @member_id = member_id
+        @group_assignment = group_assignment
+      end
+
+      def api_key
+        14
+      end
+
+      def response_class
+        SyncGroupResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_int32(@generation_id)
+        encoder.write_string(@member_id)
+
+        encoder.write_array(@group_assignment) do |member_id, member_assignment|
+          encoder.write_string(member_id)
+          encoder.write_bytes(Encoder.encode_with(member_assignment))
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/sync_group_response.rb

@@ -0,0 +1,21 @@
+require "kafka/protocol/member_assignment"
+
+module Kafka
+  module Protocol
+    class SyncGroupResponse
+      attr_reader :error_code, :member_assignment
+
+      def initialize(error_code:, member_assignment:)
+        @error_code = error_code
+        @member_assignment = member_assignment
+      end
+
+      def self.decode(decoder)
+        new(
+          error_code: decoder.int16,
+          member_assignment: MemberAssignment.decode(Decoder.from_string(decoder.bytes)),
+        )
+      end
+    end
+  end
+end

data/lib/kafka/round_robin_assignment_strategy.rb

@@ -0,0 +1,40 @@
+require "kafka/protocol/member_assignment"
+
+module Kafka
+
+  # A consumer group partition assignment strategy that assigns partitions to
+  # consumers in a round-robin fashion.
+  class RoundRobinAssignmentStrategy
+    def initialize(cluster:)
+      @cluster = cluster
+    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 = {}
+
+      members.each do |member_id|
+        group_assignment[member_id] = Protocol::MemberAssignment.new
+      end
+
+      topics.each do |topic|
+        partitions = @cluster.partitions_for(topic).map(&:partition_id)
+
+        partitions_per_member = partitions.group_by {|partition_id|
+          partition_id % members.count
+        }.values
+
+        members.zip(partitions_per_member).each do |member_id, member_partitions|
+          group_assignment[member_id].assign(topic, member_partitions)
+        end
+      end
+
+      group_assignment
+    end
+  end
+end