ruby-kafka-temp-fork 0.0.1

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

data/lib/kafka/protocol/record_batch.rb

@@ -0,0 +1,223 @@
+require 'bigdecimal'
+require 'digest/crc32'
+require 'kafka/protocol/record'
+
+module Kafka
+  module Protocol
+    class RecordBatch
+      MAGIC_BYTE = 2
+      # The size of metadata before the real record data
+      RECORD_BATCH_OVERHEAD = 49
+      # Masks to extract information from attributes
+      CODEC_ID_MASK = 0b00000111
+      IN_TRANSACTION_MASK = 0b00010000
+      IS_CONTROL_BATCH_MASK = 0b00100000
+      TIMESTAMP_TYPE_MASK = 0b001000
+
+      attr_reader :records, :first_offset, :first_timestamp, :partition_leader_epoch, :in_transaction, :is_control_batch, :last_offset_delta, :max_timestamp, :producer_id, :producer_epoch, :first_sequence
+
+      attr_accessor :codec_id
+
+      def initialize(
+          records: [],
+          first_offset: 0,
+          first_timestamp: Time.now,
+          partition_leader_epoch: 0,
+          codec_id: 0,
+          in_transaction: false,
+          is_control_batch: false,
+          last_offset_delta: 0,
+          producer_id: -1,
+          producer_epoch: 0,
+          first_sequence: 0,
+          max_timestamp: Time.now
+      )
+        @records = Array(records)
+        @first_offset = first_offset
+        @first_timestamp = first_timestamp
+        @codec_id = codec_id
+
+        # Records verification
+        @last_offset_delta = last_offset_delta
+        @max_timestamp = max_timestamp
+
+        # Transaction information
+        @producer_id = producer_id
+        @producer_epoch = producer_epoch
+
+        @first_sequence = first_sequence
+        @partition_leader_epoch = partition_leader_epoch
+        @in_transaction = in_transaction
+        @is_control_batch = is_control_batch
+
+        mark_control_record
+      end
+
+      def size
+        @records.size
+      end
+
+      def last_offset
+        @first_offset + @last_offset_delta
+      end
+
+      def attributes
+        0x0000 | @codec_id |
+          (@in_transaction ? IN_TRANSACTION_MASK : 0x0) |
+          (@is_control_batch ? IS_CONTROL_BATCH_MASK : 0x0)
+      end
+
+      def encode(encoder)
+        encoder.write_int64(@first_offset)
+
+        record_batch_buffer = StringIO.new
+        record_batch_encoder = Encoder.new(record_batch_buffer)
+
+        record_batch_encoder.write_int32(@partition_leader_epoch)
+        record_batch_encoder.write_int8(MAGIC_BYTE)
+
+        body = encode_record_batch_body
+        crc = ::Digest::CRC32c.checksum(body)
+
+        record_batch_encoder.write_int32(crc)
+        record_batch_encoder.write(body)
+
+        encoder.write_bytes(record_batch_buffer.string)
+      end
+
+      def encode_record_batch_body
+        buffer = StringIO.new
+        encoder = Encoder.new(buffer)
+
+        encoder.write_int16(attributes)
+        encoder.write_int32(@last_offset_delta)
+        encoder.write_int64((@first_timestamp.to_f * 1000).to_i)
+        encoder.write_int64((@max_timestamp.to_f * 1000).to_i)
+
+        encoder.write_int64(@producer_id)
+        encoder.write_int16(@producer_epoch)
+        encoder.write_int32(@first_sequence)
+
+        encoder.write_int32(@records.length)
+
+        records_array = encode_record_array
+        if compressed?
+          codec = Compression.find_codec_by_id(@codec_id)
+          records_array = codec.compress(records_array)
+        end
+        encoder.write(records_array)
+
+        buffer.string
+      end
+
+      def encode_record_array
+        buffer = StringIO.new
+        encoder = Encoder.new(buffer)
+        @records.each do |record|
+          record.encode(encoder)
+        end
+        buffer.string
+      end
+
+      def compressed?
+        @codec_id != 0
+      end
+
+      def fulfill_relative_data
+        first_record = records.min_by { |record| record.create_time }
+        @first_timestamp = first_record.nil? ? Time.now : first_record.create_time
+
+        last_record = records.max_by { |record| record.create_time }
+        @max_timestamp = last_record.nil? ? Time.now : last_record.create_time
+
+        records.each_with_index do |record, index|
+          record.offset_delta = index
+          record.timestamp_delta = ((record.create_time - first_timestamp) * 1000).to_i
+        end
+        @last_offset_delta = records.length - 1
+      end
+
+      def ==(other)
+        records == other.records &&
+          first_offset == other.first_offset &&
+          partition_leader_epoch == other.partition_leader_epoch &&
+          in_transaction == other.in_transaction &&
+          is_control_batch == other.is_control_batch &&
+          last_offset_delta == other.last_offset_delta &&
+          producer_id == other.producer_id &&
+          producer_epoch == other.producer_epoch &&
+          first_sequence == other.first_sequence
+      end
+
+      def self.decode(decoder)
+        first_offset = decoder.int64
+
+        record_batch_raw = decoder.bytes
+        record_batch_decoder = Decoder.from_string(record_batch_raw)
+
+        partition_leader_epoch = record_batch_decoder.int32
+        # Currently, the magic byte is used to distingush legacy MessageSet and
+        # RecordBatch. Therefore, we don't care about magic byte here yet.
+        _magic_byte = record_batch_decoder.int8
+        _crc = record_batch_decoder.int32
+
+        attributes = record_batch_decoder.int16
+        codec_id = attributes & CODEC_ID_MASK
+        in_transaction = (attributes & IN_TRANSACTION_MASK) > 0
+        is_control_batch = (attributes & IS_CONTROL_BATCH_MASK) > 0
+        log_append_time = (attributes & TIMESTAMP_TYPE_MASK) != 0
+
+        last_offset_delta = record_batch_decoder.int32
+        first_timestamp = Time.at(record_batch_decoder.int64 / BigDecimal(1000))
+        max_timestamp = Time.at(record_batch_decoder.int64 / BigDecimal(1000))
+
+        producer_id = record_batch_decoder.int64
+        producer_epoch = record_batch_decoder.int16
+        first_sequence = record_batch_decoder.int32
+
+        records_array_length = record_batch_decoder.int32
+        records_array_raw = record_batch_decoder.read(
+          record_batch_raw.size - RECORD_BATCH_OVERHEAD
+        )
+        if codec_id != 0
+          codec = Compression.find_codec_by_id(codec_id)
+          records_array_raw = codec.decompress(records_array_raw)
+        end
+
+        records_array_decoder = Decoder.from_string(records_array_raw)
+        records_array = []
+        until records_array_decoder.eof?
+          record = Record.decode(records_array_decoder)
+          record.offset = first_offset + record.offset_delta
+          record.create_time = log_append_time && max_timestamp ? max_timestamp : first_timestamp + record.timestamp_delta / BigDecimal(1000)
+          records_array << record
+        end
+
+        raise InsufficientDataMessage if records_array.length != records_array_length
+
+        new(
+          records: records_array,
+          first_offset: first_offset,
+          first_timestamp: first_timestamp,
+          partition_leader_epoch: partition_leader_epoch,
+          in_transaction: in_transaction,
+          is_control_batch: is_control_batch,
+          last_offset_delta: last_offset_delta,
+          producer_id: producer_id,
+          producer_epoch: producer_epoch,
+          first_sequence: first_sequence,
+          max_timestamp: max_timestamp
+        )
+      rescue EOFError
+        raise InsufficientDataMessage, 'Partial trailing record detected!'
+      end
+
+      def mark_control_record
+        if is_control_batch
+          record = @records.first
+          record.is_control_record = true unless record.nil?
+        end
+      end
+    end
+  end
+end