ruby-kafka-custom 0.7.7.26

data/lib/kafka/protocol/member_assignment.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+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.rb

@@ -0,0 +1,172 @@
+# frozen_string_literal: true
+
+require "stringio"
+require "zlib"
+
+module Kafka
+  module Protocol
+
+    # ## API Specification
+    #
+    #     Message => Crc MagicByte Attributes Timestamp Key Value
+    #         Crc => int32
+    #         MagicByte => int8
+    #         Attributes => int8
+    #         Timestamp => int64, in ms
+    #         Key => bytes
+    #         Value => bytes
+    #
+    class Message
+      MAGIC_BYTE = 1
+
+      attr_reader :key, :value, :codec_id, :offset
+
+      attr_reader :bytesize, :create_time
+
+      def initialize(value:, key: nil, create_time: Time.now, codec_id: 0, offset: -1)
+        @key = key
+        @value = value
+        @codec_id = codec_id
+        @offset = offset
+        @create_time = create_time
+
+        @bytesize = @key.to_s.bytesize + @value.to_s.bytesize
+      end
+
+      def encode(encoder)
+        data = encode_with_crc
+
+        encoder.write_int64(offset)
+        encoder.write_bytes(data)
+      end
+
+      def ==(other)
+        @key == other.key &&
+          @value == other.value &&
+          @codec_id == other.codec_id &&
+          @offset == other.offset
+      end
+
+      def compressed?
+        @codec_id != 0
+      end
+
+      # @return [Array<Kafka::Protocol::Message>]
+      def decompress
+        codec = Compression.find_codec_by_id(@codec_id)
+
+        # For some weird reason we need to cut out the first 20 bytes.
+        data = codec.decompress(value)
+        message_set_decoder = Decoder.from_string(data)
+        message_set = MessageSet.decode(message_set_decoder)
+
+        correct_offsets(message_set.messages)
+      end
+
+      def self.decode(decoder)
+        offset = decoder.int64
+        message_decoder = Decoder.from_string(decoder.bytes)
+
+        _crc = message_decoder.int32
+        magic_byte = message_decoder.int8
+        attributes = message_decoder.int8
+
+        # The magic byte indicates the message format version. There are situations
+        # where an old message format can be returned from a newer version of Kafka,
+        # because old messages are not necessarily rewritten on upgrades.
+        case magic_byte
+        when 0
+          # No timestamp in the pre-0.10 message format.
+          timestamp = nil
+        when 1
+          timestamp = message_decoder.int64
+
+          # If the timestamp is set to zero, it's because the message has been upgraded
+          # from the Kafka 0.9 disk format to the Kafka 0.10 format. The former didn't
+          # have a timestamp attribute, so we'll just set the timestamp to nil.
+          timestamp = nil if timestamp.zero?
+        else
+          raise Kafka::Error, "Invalid magic byte: #{magic_byte}"
+        end
+
+        key = message_decoder.bytes
+        value = message_decoder.bytes
+
+        # The codec id is encoded in the three least significant bits of the
+        # attributes.
+        codec_id = attributes & 0b111
+
+        # The timestamp will be nil if the message was written in the Kafka 0.9 log format.
+        create_time = timestamp && Time.at(timestamp / 1000.0)
+
+        new(key: key, value: value, codec_id: codec_id, offset: offset, create_time: create_time)
+      end
+
+      # Ensure the backward compatibility of Message format from Kafka 0.11.x
+      def is_control_record
+        false
+      end
+
+      def headers
+        {}
+      end
+
+      private
+
+      # Offsets may be relative with regards to wrapped message offset, but there are special cases.
+      #
+      # Cases when client will receive corrected offsets:
+      #   - When fetch request is version 0, kafka will correct relative offset on broker side before replying fetch response
+      #   - When messages is stored in 0.9 format on disk (broker configured to do so).
+      #
+      # All other cases, compressed inner messages should have relative offset, with below attributes:
+      #   - The container message should have the 'real' offset
+      #   - The container message's offset should be the 'real' offset of the last message in the compressed batch
+      def correct_offsets(messages)
+        max_relative_offset = messages.last.offset
+
+        # The offsets are already correct, do nothing.
+        return messages if max_relative_offset == offset
+
+        # The contained messages have relative offsets, and needs to be corrected.
+        base_offset = offset - max_relative_offset
+
+        messages.map do |message|
+          Message.new(
+            offset: message.offset + base_offset,
+            value: message.value,
+            key: message.key,
+            create_time: message.create_time,
+            codec_id: message.codec_id
+          )
+        end
+      end
+
+      def encode_with_crc
+        buffer = StringIO.new
+        encoder = Encoder.new(buffer)
+
+        data = encode_without_crc
+        crc = Zlib.crc32(data)
+
+        encoder.write_int32(crc)
+        encoder.write(data)
+
+        buffer.string
+      end
+
+      def encode_without_crc
+        buffer = StringIO.new
+        encoder = Encoder.new(buffer)
+
+        encoder.write_int8(MAGIC_BYTE)
+        encoder.write_int8(@codec_id)
+        encoder.write_int64((@create_time.to_f * 1000).to_i)
+        encoder.write_bytes(@key)
+        encoder.write_bytes(@value)
+
+        buffer.string
+      end
+    end
+  end
+end

data/lib/kafka/protocol/message_set.rb

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class MessageSet
+      attr_reader :messages
+
+      def initialize(messages: [])
+        @messages = messages
+      end
+
+      def size
+        @messages.size
+      end
+
+      def ==(other)
+        messages == other.messages
+      end
+
+      def encode(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
+
+      def self.decode(decoder)
+        fetched_messages = []
+
+        until decoder.eof?
+          begin
+            message = Message.decode(decoder)
+
+            if message.compressed?
+              fetched_messages.concat(message.decompress)
+            else
+              fetched_messages << message
+            end
+          rescue EOFError
+            if fetched_messages.empty?
+              # If the first message in the set is truncated, it's likely because the
+              # message is larger than the maximum size that we have asked for.
+              raise MessageTooLargeToRead
+            else
+              # We tried to decode a partial message at the end of the set; just skip it.
+            end
+          end
+        end
+
+        new(messages: fetched_messages)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/metadata_request.rb

@@ -0,0 +1,31 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class MetadataRequest
+
+      # A request for cluster metadata.
+      #
+      # @param topics [Array<String>]
+      def initialize(topics: [])
+        @topics = topics
+      end
+
+      def api_key
+        TOPIC_METADATA_API
+      end
+
+      def api_version
+        1
+      end
+
+      def response_class
+        Protocol::MetadataResponse
+      end
+
+      def encode(encoder)
+        encoder.write_array(@topics) {|topic| encoder.write_string(topic) }
+      end
+    end
+  end
+end

data/lib/kafka/protocol/metadata_response.rb

@@ -0,0 +1,185 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+
+    # A response to a {MetadataRequest}.
+    #
+    # The response contains information on the brokers, topics, and partitions in
+    # the cluster.
+    #
+    # * For each broker a node id, host, and port is provided.
+    # * For each topic partition the node id of the broker acting as partition leader,
+    #   as well as a list of node ids for the set of replicas, are given. The `isr` list is
+    #   the subset of replicas that are "in sync", i.e. have fully caught up with the
+    #   leader.
+    #
+    # ## API Specification
+    #
+    #     MetadataResponse => [Broker][TopicMetadata]
+    #       Broker => NodeId Host Port  (any number of brokers may be returned)
+    #         NodeId => int32
+    #         Host => string
+    #         Port => int32
+    #
+    #       TopicMetadata => TopicErrorCode TopicName [PartitionMetadata]
+    #         TopicErrorCode => int16
+    #
+    #       PartitionMetadata => PartitionErrorCode PartitionId Leader Replicas Isr
+    #         PartitionErrorCode => int16
+    #         PartitionId => int32
+    #         Leader => int32
+    #         Replicas => [int32]
+    #         Isr => [int32]
+    #
+    class MetadataResponse
+      class PartitionMetadata
+        attr_reader :partition_id, :leader
+
+        attr_reader :partition_error_code
+
+        def initialize(partition_error_code:, partition_id:, leader:, replicas: [], isr: [])
+          @partition_error_code = partition_error_code
+          @partition_id = partition_id
+          @leader = leader
+          @replicas = replicas
+          @isr = isr
+        end
+      end
+
+      class TopicMetadata
+        # @return [String] the name of the topic
+        attr_reader :topic_name
+
+        # @return [Array<PartitionMetadata>] the partitions in the topic.
+        attr_reader :partitions
+
+        attr_reader :topic_error_code
+
+        def initialize(topic_error_code: 0, topic_name:, partitions:)
+          @topic_error_code = topic_error_code
+          @topic_name = topic_name
+          @partitions = partitions
+        end
+      end
+
+      # @return [Array<Kafka::BrokerInfo>] the list of brokers in the cluster.
+      attr_reader :brokers
+
+      # @return [Array<TopicMetadata>] the list of topics in the cluster.
+      attr_reader :topics
+
+      # @return [Integer] The broker id of the controller broker.
+      attr_reader :controller_id
+
+      def initialize(brokers:, controller_id:, topics:)
+        @brokers = brokers
+        @controller_id = controller_id
+        @topics = topics
+      end
+
+      # Finds the node id of the broker that is acting as leader for the given topic
+      # and partition per this metadata.
+      #
+      # @param topic [String] the name of the topic.
+      # @param partition [Integer] the partition number.
+      # @return [Integer] the node id of the leader.
+      def find_leader_id(topic, partition)
+        topic_info = @topics.find {|t| t.topic_name == topic }
+
+        if topic_info.nil?
+          raise UnknownTopicOrPartition, "no topic #{topic}"
+        end
+
+        Protocol.handle_error(topic_info.topic_error_code)
+
+        partition_info = topic_info.partitions.find {|p| p.partition_id == partition }
+
+        if partition_info.nil?
+          raise UnknownTopicOrPartition, "no partition #{partition} in topic #{topic}"
+        end
+
+        begin
+          Protocol.handle_error(partition_info.partition_error_code)
+        rescue ReplicaNotAvailable
+          # This error can be safely ignored per the protocol specification.
+        end
+
+        partition_info.leader
+      end
+
+      # Finds the broker info for the given node id.
+      #
+      # @param node_id [Integer] the node id of the broker.
+      # @return [Kafka::BrokerInfo] information about the broker.
+      def find_broker(node_id)
+        broker = @brokers.find {|b| b.node_id == node_id }
+
+        raise Kafka::NoSuchBroker, "No broker with id #{node_id}" if broker.nil?
+
+        broker
+      end
+
+      def controller_broker
+        find_broker(controller_id)
+      end
+
+      def partitions_for(topic_name)
+        topic = @topics.find {|t| t.topic_name == topic_name }
+
+        if topic.nil?
+          raise UnknownTopicOrPartition, "unknown topic #{topic_name}"
+        end
+
+        Protocol.handle_error(topic.topic_error_code)
+
+        topic.partitions
+      end
+
+      # Decodes a MetadataResponse from a {Decoder} containing response data.
+      #
+      # @param decoder [Decoder]
+      # @return [MetadataResponse] the metadata response.
+      def self.decode(decoder)
+        brokers = decoder.array do
+          node_id = decoder.int32
+          host = decoder.string
+          port = decoder.int32
+          _rack = decoder.string
+
+          BrokerInfo.new(
+            node_id: node_id,
+            host: host,
+            port: port
+          )
+        end
+
+        controller_id = decoder.int32
+
+        topics = decoder.array do
+          topic_error_code = decoder.int16
+          topic_name = decoder.string
+          _is_internal = decoder.boolean
+
+          partitions = decoder.array do
+            PartitionMetadata.new(
+              partition_error_code: decoder.int16,
+              partition_id: decoder.int32,
+              leader: decoder.int32,
+              replicas: decoder.array { decoder.int32 },
+              isr: decoder.array { decoder.int32 },
+            )
+          end
+
+          TopicMetadata.new(
+            topic_error_code: topic_error_code,
+            topic_name: topic_name,
+            partitions: partitions,
+          )
+        end
+
+        new(brokers: brokers, controller_id: controller_id, topics: topics)
+      end
+    end
+  end
+end