ruby-kafka-aws-iam 1.4.1

data/lib/kafka/protocol/encoder.rb

@@ -0,0 +1,184 @@
+# frozen_string_literal: true
+
+require "stringio"
+
+module Kafka
+  module Protocol
+    # An encoder wraps an IO object, making it easy to write specific data types
+    # to it.
+    class Encoder
+      # Initializes a new encoder.
+      #
+      # @param io [IO] an object that acts as an IO.
+      def initialize(io)
+        @io = io
+        @io.set_encoding(Encoding::BINARY)
+      end
+
+      # Writes bytes directly to the IO object.
+      #
+      # @param bytes [String]
+      # @return [nil]
+      def write(bytes)
+        @io.write(bytes)
+
+        nil
+      end
+
+      # Writes an 8-bit boolean to the IO object.
+      #
+      # @param boolean [Boolean]
+      # @return [nil]
+      def write_boolean(boolean)
+        boolean ? write_int8(1) : write_int8(0)
+      end
+
+      # Writes an 8-bit integer to the IO object.
+      #
+      # @param int [Integer]
+      # @return [nil]
+      def write_int8(int)
+        write([int].pack("C"))
+      end
+
+      # Writes a 16-bit integer to the IO object.
+      #
+      # @param int [Integer]
+      # @return [nil]
+      def write_int16(int)
+        write([int].pack("s>"))
+      end
+
+      # Writes a 32-bit integer to the IO object.
+      #
+      # @param int [Integer]
+      # @return [nil]
+      def write_int32(int)
+        write([int].pack("l>"))
+      end
+
+      # Writes a 64-bit integer to the IO object.
+      #
+      # @param int [Integer]
+      # @return [nil]
+      def write_int64(int)
+        write([int].pack("q>"))
+      end
+
+      # Writes an array to the IO object.
+      #
+      # Each item in the specified array will be yielded to the provided block;
+      # it's the responsibility of the block to write those items using the
+      # encoder.
+      #
+      # @param array [Array]
+      # @return [nil]
+      def write_array(array, &block)
+        if array.nil?
+          # An array can be null, which is different from it being empty.
+          write_int32(-1)
+        else
+          write_int32(array.size)
+          array.each(&block)
+        end
+      end
+
+      # Writes an array to the IO object.
+      # Just like #write_array, unless the size is under varint format
+      #
+      # @param array [Array]
+      # @return [nil]
+      def write_varint_array(array, &block)
+        if array.nil?
+          write_varint(-1)
+        else
+          write_varint(array.size)
+          array.each(&block)
+        end
+      end
+
+      # Writes a string to the IO object.
+      #
+      # @param string [String]
+      # @return [nil]
+      def write_string(string)
+        if string.nil?
+          write_int16(-1)
+        else
+          write_int16(string.bytesize)
+          write(string)
+        end
+      end
+
+      # Writes a string to the IO object, the size is under varint format
+      #
+      # @param string [String]
+      # @return [nil]
+      def write_varint_string(string)
+        if string.nil?
+          write_varint(-1)
+        else
+          write_varint(string.bytesize)
+          write(string)
+        end
+      end
+
+      # Writes an integer under varints serializing to the IO object.
+      # https://developers.google.com/protocol-buffers/docs/encoding#varints
+      #
+      # @param int [Integer]
+      # @return [nil]
+      def write_varint(int)
+        int = int << 1
+        int = ~int | 1 if int < 0
+
+        chunks = []
+        while int >> 7 != 0
+          chunks << (int & 0x7f | 0x80)
+          int >>= 7
+        end
+        chunks << int
+        write(chunks.pack("C*"))
+      end
+
+      # Writes a byte string to the IO object.
+      #
+      # @param bytes [String]
+      # @return [nil]
+      def write_bytes(bytes)
+        if bytes.nil?
+          write_int32(-1)
+        else
+          write_int32(bytes.bytesize)
+          write(bytes)
+        end
+      end
+
+      # Writes a byte string to the IO object, the size is under varint format
+      #
+      # @param bytes [String]
+      # @return [nil]
+      def write_varint_bytes(bytes)
+        if bytes.nil?
+          write_varint(-1)
+        else
+          write_varint(bytes.bytesize)
+          write(bytes)
+        end
+      end
+
+      # Encodes an object into a new buffer.
+      #
+      # @param object [#encode] the object that will encode itself.
+      # @return [String] the encoded data.
+      def self.encode_with(object)
+        buffer = StringIO.new
+        encoder = new(buffer)
+
+        object.encode(encoder)
+
+        buffer.string
+      end
+    end
+  end
+end

data/lib/kafka/protocol/end_txn_request.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class EndTxnRequest
+      def initialize(transactional_id:, producer_id:, producer_epoch:, transaction_result:)
+        @transactional_id = transactional_id
+        @producer_id = producer_id
+        @producer_epoch = producer_epoch
+        @transaction_result = transaction_result
+      end
+
+      def api_key
+        END_TXN_API
+      end
+
+      def response_class
+        EndTxnResposne
+      end
+
+      def encode(encoder)
+        encoder.write_string(@transactional_id)
+        encoder.write_int64(@producer_id)
+        encoder.write_int16(@producer_epoch)
+        encoder.write_boolean(@transaction_result)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/end_txn_response.rb

@@ -0,0 +1,19 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class EndTxnResposne
+      attr_reader :error_code
+
+      def initialize(error_code:)
+        @error_code = error_code
+      end
+
+      def self.decode(decoder)
+        _throttle_time_ms = decoder.int32
+        error_code = decoder.int16
+        new(error_code: error_code)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/fetch_request.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+
+    # A request to fetch messages from a given partition.
+    #
+    # ## API Specification
+    #
+    #     FetchRequest => ReplicaId MaxWaitTime MinBytes MaxBytes IsolationLevel [TopicName [Partition FetchOffset MaxBytes]]
+    #       ReplicaId => int32
+    #       MaxWaitTime => int32
+    #       MinBytes => int32
+    #       MaxBytes => int32
+    #       IsolationLevel => int8
+    #       TopicName => string
+    #       Partition => int32
+    #       FetchOffset => int64
+    #       MaxBytes => int32
+    #
+    class FetchRequest
+      ISOLATION_READ_UNCOMMITTED = 0
+      ISOLATION_READ_COMMITTED = 1
+
+      # @param max_wait_time [Integer]
+      # @param min_bytes [Integer]
+      # @param topics [Hash]
+      def initialize(max_wait_time:, min_bytes:, max_bytes:, topics:)
+        @replica_id = REPLICA_ID
+        @max_wait_time = max_wait_time
+        @min_bytes = min_bytes
+        @max_bytes = max_bytes
+        @topics = topics
+      end
+
+      def api_key
+        FETCH_API
+      end
+
+      def api_version
+        4
+      end
+
+      def response_class
+        Protocol::FetchResponse
+      end
+
+      def encode(encoder)
+        encoder.write_int32(@replica_id)
+        encoder.write_int32(@max_wait_time)
+        encoder.write_int32(@min_bytes)
+        encoder.write_int32(@max_bytes)
+        encoder.write_int8(ISOLATION_READ_COMMITTED)
+
+        encoder.write_array(@topics) do |topic, partitions|
+          encoder.write_string(topic)
+
+          encoder.write_array(partitions) do |partition, config|
+            fetch_offset = config.fetch(:fetch_offset)
+            max_bytes = config.fetch(:max_bytes)
+
+            encoder.write_int32(partition)
+            encoder.write_int64(fetch_offset)
+            encoder.write_int32(max_bytes)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/fetch_response.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+require "kafka/protocol/message_set"
+require "kafka/protocol/record_batch"
+
+module Kafka
+  module Protocol
+
+    # A response to a fetch request.
+    #
+    # ## API Specification
+    #
+    #     FetchResponse => ThrottleTimeMS [TopicName [Partition ErrorCode HighwaterMarkOffset LastStableOffset [AbortedTransaction] Records]]
+    #       ThrottleTimeMS => int32
+    #       TopicName => string
+    #       Partition => int32
+    #       ErrorCode => int16
+    #       HighwaterMarkOffset => int64
+    #       LastStableOffset => int64
+    #       MessageSetSize => int32
+    #       AbortedTransaction => [
+    #             ProducerId => int64
+    #             FirstOffset => int64
+    #       ]
+    #
+    class FetchResponse
+      MAGIC_BYTE_OFFSET = 16
+      MAGIC_BYTE_LENGTH = 1
+
+      class FetchedPartition
+        attr_reader :partition, :error_code
+        attr_reader :highwater_mark_offset, :last_stable_offset, :aborted_transactions, :messages
+
+        def initialize(partition:, error_code:, highwater_mark_offset:, last_stable_offset:, aborted_transactions:, messages:)
+          @partition = partition
+          @error_code = error_code
+          @highwater_mark_offset = highwater_mark_offset
+          @messages = messages
+          @last_stable_offset = last_stable_offset
+          @aborted_transactions = aborted_transactions
+        end
+      end
+
+      class FetchedTopic
+        attr_reader :name, :partitions
+
+        def initialize(name:, partitions:)
+          @name = name
+          @partitions = partitions
+        end
+      end
+
+      class AbortedTransaction
+        attr_reader :producer_id, :first_offset
+
+        def initialize(producer_id:, first_offset:)
+          @producer_id = producer_id
+          @first_offset = first_offset
+        end
+      end
+
+      attr_reader :topics
+
+      def initialize(topics: [], throttle_time_ms: 0)
+        @topics = topics
+        @throttle_time_ms = throttle_time_ms
+      end
+
+      def self.decode(decoder)
+        throttle_time_ms = decoder.int32
+
+        topics = decoder.array do
+          topic_name = decoder.string
+
+          partitions = decoder.array do
+            partition = decoder.int32
+            error_code = decoder.int16
+            highwater_mark_offset = decoder.int64
+            last_stable_offset = decoder.int64
+
+            aborted_transactions = decoder.array do
+              producer_id = decoder.int64
+              first_offset = decoder.int64
+              AbortedTransaction.new(
+                producer_id: producer_id,
+                first_offset: first_offset
+              )
+            end
+
+            messages_raw = decoder.bytes
+            messages = []
+
+            if !messages_raw.nil? && !messages_raw.empty?
+              messages_decoder = Decoder.from_string(messages_raw)
+
+              magic_byte = messages_decoder.peek(MAGIC_BYTE_OFFSET, MAGIC_BYTE_LENGTH)[0].to_i
+              if magic_byte == RecordBatch::MAGIC_BYTE
+                until messages_decoder.eof?
+                  begin
+                    record_batch = RecordBatch.decode(messages_decoder)
+                    messages << record_batch
+                  rescue InsufficientDataMessage
+                    if messages.length > 0
+                      break
+                    else
+                      raise
+                    end
+                  end
+                end
+              else
+                message_set = MessageSet.decode(messages_decoder)
+                messages << message_set
+              end
+            end
+
+            FetchedPartition.new(
+              partition: partition,
+              error_code: error_code,
+              highwater_mark_offset: highwater_mark_offset,
+              last_stable_offset: last_stable_offset,
+              aborted_transactions: aborted_transactions,
+              messages: messages
+            )
+          end
+
+          FetchedTopic.new(
+            name: topic_name,
+            partitions: partitions,
+          )
+        end
+
+        new(topics: topics, throttle_time_ms: throttle_time_ms)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/find_coordinator_request.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class FindCoordinatorRequest
+      def initialize(coordinator_key:, coordinator_type:)
+        @coordinator_key  = coordinator_key
+        @coordinator_type = coordinator_type
+      end
+
+      def api_key
+        FIND_COORDINATOR_API
+      end
+
+      def api_version
+        1
+      end
+
+      def encode(encoder)
+        encoder.write_string(@coordinator_key)
+        encoder.write_int8(@coordinator_type)
+      end
+
+      def response_class
+        FindCoordinatorResponse
+      end
+    end
+  end
+end

data/lib/kafka/protocol/find_coordinator_response.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class FindCoordinatorResponse
+      attr_reader :error_code, :error_message
+
+      attr_reader :coordinator_id, :coordinator_host, :coordinator_port
+
+      def initialize(error_code:, error_message:, coordinator_id:, coordinator_host:, coordinator_port:)
+        @error_code = error_code
+        @coordinator_id = coordinator_id
+        @coordinator_host = coordinator_host
+        @coordinator_port = coordinator_port
+      end
+
+      def self.decode(decoder)
+        _throttle_time_ms = decoder.int32
+        new(
+          error_code: decoder.int16,
+          error_message: decoder.string,
+          coordinator_id: decoder.int32,
+          coordinator_host: decoder.string,
+          coordinator_port: decoder.int32,
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/heartbeat_request.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class HeartbeatRequest
+      def initialize(group_id:, generation_id:, member_id:)
+        @group_id = group_id
+        @generation_id = generation_id
+        @member_id = member_id
+      end
+
+      def api_key
+        HEARTBEAT_API
+      end
+
+      def response_class
+        HeartbeatResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_int32(@generation_id)
+        encoder.write_string(@member_id)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/heartbeat_response.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class HeartbeatResponse
+      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/init_producer_id_request.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class InitProducerIDRequest
+      def initialize(transactional_id: nil, transactional_timeout:)
+        @transactional_id = transactional_id
+        @transactional_timeout = transactional_timeout
+      end
+
+      def api_key
+        INIT_PRODUCER_ID_API
+      end
+
+      def response_class
+        InitProducerIDResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@transactional_id)
+        # Timeout is in ms unit
+        encoder.write_int32(@transactional_timeout * 1000)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/init_producer_id_response.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class InitProducerIDResponse
+      attr_reader :error_code, :producer_id, :producer_epoch
+
+      def initialize(error_code:, producer_id:, producer_epoch:)
+        @error_code = error_code
+        @producer_id = producer_id
+        @producer_epoch = producer_epoch
+      end
+
+      def self.decode(decoder)
+        _throttle_time_ms = decoder.int32
+        error_code = decoder.int16
+        producer_id = decoder.int64
+        producer_epoch = decoder.int16
+        new(
+          error_code: error_code,
+          producer_id: producer_id,
+          producer_epoch: producer_epoch
+        )
+      end
+    end
+  end
+end

data/lib/kafka/protocol/join_group_request.rb

@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require "kafka/protocol/consumer_group_protocol"
+
+module Kafka
+  module Protocol
+    class JoinGroupRequest
+      PROTOCOL_TYPE = "consumer"
+
+      def initialize(group_id:, session_timeout:, rebalance_timeout:, member_id:, topics: [], protocol_name:, user_data: nil)
+        @group_id = group_id
+        @session_timeout = session_timeout * 1000 # Kafka wants ms.
+        @rebalance_timeout = rebalance_timeout * 1000 # Kafka wants ms.
+        @member_id = member_id || ""
+        @protocol_type = PROTOCOL_TYPE
+        @group_protocols = {
+          protocol_name => ConsumerGroupProtocol.new(topics: topics, user_data: user_data),
+        }
+      end
+
+      def api_key
+        JOIN_GROUP_API
+      end
+
+      def api_version
+        1
+      end
+
+      def response_class
+        JoinGroupResponse
+      end
+
+      def encode(encoder)
+        encoder.write_string(@group_id)
+        encoder.write_int32(@session_timeout)
+        encoder.write_int32(@rebalance_timeout)
+        encoder.write_string(@member_id)
+        encoder.write_string(@protocol_type)
+
+        encoder.write_array(@group_protocols) do |name, metadata|
+          encoder.write_string(name)
+          encoder.write_bytes(Encoder.encode_with(metadata))
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/join_group_response.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module Kafka
+  module Protocol
+    class JoinGroupResponse
+      Metadata = Struct.new(:version, :topics, :user_data)
+
+      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 do
+              member_id = decoder.string
+              d = Decoder.from_string(decoder.bytes)
+              [member_id, Metadata.new(d.int16, d.array { d.string }, d.bytes)]
+            end
+          ],
+        )
+      end
+    end
+  end
+end