ruby-kafka-aws-iam 1.4.1

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, :replicas
+
+        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

data/lib/kafka/protocol/offset_commit_request.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class OffsetCommitRequest
+      # This value signals to the broker that its default configuration should be used.
+      DEFAULT_RETENTION_TIME = -1
+
+      def api_key
+        OFFSET_COMMIT_API
+      end
+
+      def api_version
+        2
+      end
+
+      def response_class
+        OffsetCommitResponse
+      end
+
+      def initialize(group_id:, generation_id:, member_id:, retention_time: DEFAULT_RETENTION_TIME, 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,29 @@
+# frozen_string_literal: true
+
+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,38 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class OffsetFetchRequest
+      def initialize(group_id:, topics:)
+        @group_id = group_id
+        @topics = topics
+      end
+
+      def api_key
+        OFFSET_FETCH_API
+      end
+
+      # setting topics to nil fetches all offsets for a consumer group
+      # and that feature is only available in API version 2+
+      def api_version
+        @topics.nil? ? 2 : 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,56 @@
+# frozen_string_literal: true
+
+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, nil)
+
+        if offset_info
+          Protocol.handle_error(offset_info.error_code)
+          offset_info.offset
+        else
+          -1
+        end
+      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/produce_request.rb

@@ -0,0 +1,94 @@
+# frozen_string_literal: true
+
+require "stringio"
+
+module Kafka
+  module Protocol
+
+    # A produce request sends a message set to the server.
+    #
+    # ## API Specification
+    #
+    #     ProduceRequest => RequiredAcks Timeout [TopicName [Partition MessageSetSize MessageSet]]
+    #         RequiredAcks => int16
+    #         Timeout => int32
+    #         Partition => int32
+    #         MessageSetSize => int32
+    #
+    #     MessageSet => [Offset MessageSize Message]
+    #         Offset => int64
+    #         MessageSize => int32
+    #
+    #     Message => Crc MagicByte Attributes Key Value
+    #         Crc => int32
+    #         MagicByte => int8
+    #         Attributes => int8
+    #         Key => bytes
+    #         Value => bytes
+    #
+    class ProduceRequest
+      API_MIN_VERSION = 3
+
+      attr_reader :transactional_id, :required_acks, :timeout, :messages_for_topics, :compressor
+
+      # @param required_acks [Integer]
+      # @param timeout [Integer]
+      # @param messages_for_topics [Hash]
+      def initialize(transactional_id: nil, required_acks:, timeout:, messages_for_topics:, compressor: nil)
+        @transactional_id = transactional_id
+        @required_acks = required_acks
+        @timeout = timeout
+        @messages_for_topics = messages_for_topics
+        @compressor = compressor
+      end
+
+      def api_key
+        PRODUCE_API
+      end
+
+      def api_version
+        compressor.codec.nil? ? API_MIN_VERSION : [compressor.codec.produce_api_min_version, API_MIN_VERSION].max
+      end
+
+      def response_class
+        requires_acks? ? Protocol::ProduceResponse : nil
+      end
+
+      # Whether this request requires any acknowledgements at all. If no acknowledgements
+      # are required, the server will not send back a response at all.
+      #
+      # @return [Boolean] true if acknowledgements are required, false otherwise.
+      def requires_acks?
+        @required_acks != 0
+      end
+
+      def encode(encoder)
+        encoder.write_string(@transactional_id)
+        encoder.write_int16(@required_acks)
+        encoder.write_int32(@timeout)
+
+        encoder.write_array(@messages_for_topics) do |topic, messages_for_partition|
+          encoder.write_string(topic)
+
+          encoder.write_array(messages_for_partition) do |partition, record_batch|
+            encoder.write_int32(partition)
+
+            record_batch.fulfill_relative_data
+            encoded_record_batch = compress(record_batch)
+            encoder.write_bytes(encoded_record_batch)
+          end
+        end
+      end
+
+      private
+
+      def compress(record_batch)
+        if @compressor.nil?
+          Protocol::Encoder.encode_with(record_batch)
+        else
+          @compressor.compress(record_batch)
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/produce_response.rb

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class ProduceResponse
+      class TopicInfo
+        attr_reader :topic, :partitions
+
+        def initialize(topic:, partitions:)
+          @topic = topic
+          @partitions = partitions
+        end
+      end
+
+      class PartitionInfo
+        attr_reader :partition, :error_code, :offset, :timestamp
+
+        def initialize(partition:, error_code:, offset:, timestamp:)
+          @partition = partition
+          @error_code = error_code
+          @offset = offset
+          @timestamp = timestamp
+        end
+      end
+
+      attr_reader :topics, :throttle_time_ms
+
+      def initialize(topics: [], throttle_time_ms: 0)
+        @topics = topics
+        @throttle_time_ms = throttle_time_ms
+      end
+
+      def each_partition
+        @topics.each do |topic_info|
+          topic_info.partitions.each do |partition_info|
+            yield topic_info, partition_info
+          end
+        end
+      end
+
+      def self.decode(decoder)
+        topics = decoder.array do
+          topic = decoder.string
+
+          partitions = decoder.array do
+            PartitionInfo.new(
+              partition: decoder.int32,
+              error_code: decoder.int16,
+              offset: decoder.int64,
+              timestamp: Time.at(decoder.int64 / 1000.0),
+            )
+          end
+
+          TopicInfo.new(topic: topic, partitions: partitions)
+        end
+
+        throttle_time_ms = decoder.int32
+
+        new(topics: topics, throttle_time_ms: throttle_time_ms)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/record.rb

@@ -0,0 +1,88 @@
+module Kafka
+  module Protocol
+    class Record
+      attr_reader :key, :value, :headers, :attributes, :bytesize
+      attr_accessor :offset_delta, :timestamp_delta, :offset, :create_time, :is_control_record
+
+      def initialize(
+        key: nil,
+        value:,
+        headers: {},
+        attributes: 0,
+        offset_delta: 0,
+        offset: 0,
+        timestamp_delta: 0,
+        create_time: Time.now,
+        is_control_record: false
+      )
+        @key = key
+        @value = value
+        @headers = headers
+        @attributes = attributes
+
+        @offset_delta = offset_delta
+        @offset = offset
+        @timestamp_delta = timestamp_delta
+        @create_time = create_time
+        @is_control_record = is_control_record
+
+        @bytesize = @key.to_s.bytesize + @value.to_s.bytesize
+      end
+
+      def encode(encoder)
+        record_buffer = StringIO.new
+
+        record_encoder = Encoder.new(record_buffer)
+
+        record_encoder.write_int8(@attributes)
+        record_encoder.write_varint(@timestamp_delta)
+        record_encoder.write_varint(@offset_delta)
+
+        record_encoder.write_varint_string(@key)
+        record_encoder.write_varint_bytes(@value)
+
+        record_encoder.write_varint_array(@headers.to_a) do |header_key, header_value|
+          record_encoder.write_varint_string(header_key.to_s)
+          record_encoder.write_varint_bytes(header_value.to_s)
+        end
+
+        encoder.write_varint_bytes(record_buffer.string)
+      end
+
+      def ==(other)
+        offset_delta == other.offset_delta &&
+          timestamp_delta == other.timestamp_delta &&
+          offset == other.offset &&
+          is_control_record == other.is_control_record
+      end
+
+      def self.decode(decoder)
+        record_decoder = Decoder.from_string(decoder.varint_bytes)
+
+        attributes = record_decoder.int8
+        timestamp_delta = record_decoder.varint
+        offset_delta = record_decoder.varint
+
+        key = record_decoder.varint_string
+        value = record_decoder.varint_bytes
+
+        headers = {}
+        record_decoder.varint_array do
+          header_key = record_decoder.varint_string
+          header_value = record_decoder.varint_bytes
+
+          headers[header_key] = header_value
+        end
+
+        new(
+          key: key,
+          value: value,
+          headers: headers,
+          attributes: attributes,
+          offset_delta: offset_delta,
+          timestamp_delta: timestamp_delta
+        )
+      end
+    end
+  end
+end