ruby-kafka-aws-iam 1.4.1

data/lib/kafka/fetch_operation.rb

@@ -0,0 +1,115 @@
+# frozen_string_literal: true
+
+require "kafka/fetched_offset_resolver"
+require "kafka/fetched_batch_generator"
+
+module Kafka
+
+  # Fetches messages from one or more partitions.
+  #
+  #     operation = Kafka::FetchOperation.new(
+  #       cluster: cluster,
+  #       logger: logger,
+  #       min_bytes: 1,
+  #       max_wait_time: 10,
+  #     )
+  #
+  #     # These calls will schedule fetches from the specified topics/partitions.
+  #     operation.fetch_from_partition("greetings", 42, offset: :latest, max_bytes: 100000)
+  #     operation.fetch_from_partition("goodbyes", 13, offset: :latest, max_bytes: 100000)
+  #
+  #     operation.execute
+  #
+  class FetchOperation
+    def initialize(cluster:, logger:, min_bytes: 1, max_bytes: 10485760, max_wait_time: 5)
+      @cluster = cluster
+      @logger = TaggedLogger.new(logger)
+      @min_bytes = min_bytes
+      @max_bytes = max_bytes
+      @max_wait_time = max_wait_time
+      @topics = {}
+
+      @offset_resolver = Kafka::FetchedOffsetResolver.new(
+        logger: logger
+      )
+    end
+
+    def fetch_from_partition(topic, partition, offset: :latest, max_bytes: 1048576)
+      if offset == :earliest
+        offset = -2
+      elsif offset == :latest
+        offset = -1
+      end
+
+      @topics[topic] ||= {}
+      @topics[topic][partition] = {
+        fetch_offset: offset,
+        max_bytes: max_bytes,
+      }
+    end
+
+    def execute
+      @cluster.add_target_topics(@topics.keys)
+      @cluster.refresh_metadata_if_necessary!
+
+      topics_by_broker = {}
+
+      if @topics.none? {|topic, partitions| partitions.any? }
+        raise NoPartitionsToFetchFrom
+      end
+
+      @topics.each do |topic, partitions|
+        partitions.each do |partition, options|
+          broker = @cluster.get_leader(topic, partition)
+
+          topics_by_broker[broker] ||= {}
+          topics_by_broker[broker][topic] ||= {}
+          topics_by_broker[broker][topic][partition] = options
+        end
+      end
+
+      topics_by_broker.flat_map do |broker, topics|
+        @offset_resolver.resolve!(broker, topics)
+
+        options = {
+          max_wait_time: @max_wait_time * 1000, # Kafka expects ms, not secs
+          min_bytes: @min_bytes,
+          max_bytes: @max_bytes,
+          topics: topics,
+        }
+
+        response = broker.fetch_messages(**options)
+
+        response.topics.flat_map do |fetched_topic|
+          fetched_topic.partitions.map do |fetched_partition|
+            begin
+              Protocol.handle_error(fetched_partition.error_code)
+            rescue Kafka::OffsetOutOfRange => e
+              e.topic = fetched_topic.name
+              e.partition = fetched_partition.partition
+              e.offset = topics.fetch(e.topic).fetch(e.partition).fetch(:fetch_offset)
+
+              raise e
+            rescue Kafka::Error => e
+              topic = fetched_topic.name
+              partition = fetched_partition.partition
+              @logger.error "Failed to fetch from #{topic}/#{partition}: #{e.message}"
+              raise e
+            end
+
+            Kafka::FetchedBatchGenerator.new(
+              fetched_topic.name,
+              fetched_partition,
+              topics.fetch(fetched_topic.name).fetch(fetched_partition.partition).fetch(:fetch_offset),
+              logger: @logger
+            ).generate
+          end
+        end
+      end
+    rescue Kafka::ConnectionError, Kafka::LeaderNotAvailable, Kafka::NotLeaderForPartition
+      @cluster.mark_as_stale!
+
+      raise
+    end
+  end
+end

data/lib/kafka/fetched_batch.rb

@@ -0,0 +1,58 @@
+# frozen_string_literal: true
+
+module Kafka
+
+  # An ordered sequence of messages fetched from a Kafka partition.
+  class FetchedBatch
+    # @return [String]
+    attr_reader :topic
+
+    # @return [Integer]
+    attr_reader :partition
+
+    # @return [Integer]
+    attr_reader :last_offset
+
+    # @return [Integer]
+    attr_reader :leader_epoch
+
+    # @return [Integer] the offset of the most recent message in the partition.
+    attr_reader :highwater_mark_offset
+
+    # @return [Array<Kafka::FetchedMessage>]
+    attr_accessor :messages
+
+    def initialize(topic:, partition:, highwater_mark_offset:, messages:, last_offset: nil, leader_epoch: nil)
+      @topic = topic
+      @partition = partition
+      @highwater_mark_offset = highwater_mark_offset
+      @messages = messages
+      @last_offset = last_offset
+      @leader_epoch = leader_epoch
+    end
+
+    def empty?
+      @messages.empty?
+    end
+
+    def unknown_last_offset?
+      @last_offset.nil?
+    end
+
+    def first_offset
+      if empty?
+        nil
+      else
+        messages.first.offset
+      end
+    end
+
+    def offset_lag
+      if empty?
+        0
+      else
+        (highwater_mark_offset - 1) - last_offset
+      end
+    end
+  end
+end

data/lib/kafka/fetched_batch_generator.rb

@@ -0,0 +1,120 @@
+# frozen_string_literal: true
+
+require "kafka/fetched_batch"
+
+module Kafka
+  class FetchedBatchGenerator
+    COMMITTED_TRANSACTION_SIGNAL = "\x00\x00\x00\x01".freeze
+    ABORTED_TRANSACTION_SIGNAL = "\x00\x00\x00\x00".freeze
+
+    def initialize(topic, fetched_partition, offset, logger:)
+      @topic = topic
+      @fetched_partition = fetched_partition
+      @logger = TaggedLogger.new(logger)
+      @offset = offset
+    end
+
+    def generate
+      if @fetched_partition.messages.empty?
+        empty_fetched_batch
+      elsif @fetched_partition.messages.first.is_a?(Kafka::Protocol::MessageSet)
+        extract_messages
+      else
+        extract_records
+      end
+    end
+
+    private
+
+    def empty_fetched_batch
+      FetchedBatch.new(
+        topic: @topic,
+        partition: @fetched_partition.partition,
+        last_offset: nil,
+        highwater_mark_offset: @fetched_partition.highwater_mark_offset,
+        messages: []
+      )
+    end
+
+    def extract_messages
+      last_offset = nil
+      messages = @fetched_partition.messages.flat_map do |message_set|
+        message_set.messages.map do |message|
+          last_offset = message.offset if last_offset.nil? || last_offset < message.offset
+          if message.offset >= @offset
+            FetchedMessage.new(
+              message: message,
+              topic: @topic,
+              partition: @fetched_partition.partition
+            )
+          end
+        end.compact
+      end
+      FetchedBatch.new(
+        topic: @topic,
+        partition: @fetched_partition.partition,
+        last_offset: last_offset,
+        highwater_mark_offset: @fetched_partition.highwater_mark_offset,
+        messages: messages
+      )
+    end
+
+    def extract_records
+      records = []
+      last_offset = nil
+      leader_epoch = nil
+      aborted_transactions = @fetched_partition.aborted_transactions.sort_by(&:first_offset)
+      aborted_producer_ids = {}
+
+      @fetched_partition.messages.each do |record_batch|
+        last_offset = record_batch.last_offset if last_offset.nil? || last_offset < record_batch.last_offset
+        leader_epoch = record_batch.partition_leader_epoch if leader_epoch.nil? || leader_epoch < record_batch.partition_leader_epoch
+        # Find the list of aborted producer IDs less than current offset
+        unless aborted_transactions.empty?
+          if aborted_transactions.first.first_offset <= record_batch.last_offset
+            aborted_transaction = aborted_transactions.shift
+            aborted_producer_ids[aborted_transaction.producer_id] = aborted_transaction.first_offset
+          end
+        end
+
+        if abort_marker?(record_batch)
+          # Abort marker, remove the producer from the aborted list
+          aborted_producer_ids.delete(record_batch.producer_id)
+        elsif aborted_producer_ids.key?(record_batch.producer_id) && record_batch.in_transaction
+          # Reject aborted record batch
+          @logger.info("Reject #{record_batch.records.size} aborted records of topic '#{@topic}', partition #{@fetched_partition.partition}, from offset #{record_batch.first_offset}")
+          next
+        end
+
+        record_batch.records.each do |record|
+          if !record.is_control_record && record.offset >= @offset
+            records << FetchedMessage.new(
+              message: record,
+              topic: @topic,
+              partition: @fetched_partition.partition
+            )
+          end
+        end
+      end
+
+      FetchedBatch.new(
+        topic: @topic,
+        partition: @fetched_partition.partition,
+        last_offset: last_offset,
+        leader_epoch: leader_epoch,
+        highwater_mark_offset: @fetched_partition.highwater_mark_offset,
+        messages: records
+      )
+    end
+
+    def abort_marker?(record_batch)
+      return false unless record_batch.is_control_batch
+
+      if record_batch.records.empty?
+        raise "Invalid control record batch at topic '#{@topic}', partition #{@fetched_partition}"
+      end
+
+      record_batch.records.first.key == ABORTED_TRANSACTION_SIGNAL
+    end
+  end
+end

data/lib/kafka/fetched_message.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module Kafka
+  class FetchedMessage
+    # @return [String] the name of the topic that the message was written to.
+    attr_reader :topic
+
+    # @return [Integer] the partition number that the message was written to.
+    attr_reader :partition
+
+    def initialize(message:, topic:, partition:)
+      @message = message
+      @topic = topic
+      @partition = partition
+    end
+
+    # @return [String] the value of the message.
+    def value
+      @message.value
+    end
+
+    # @return [String] the key of the message.
+    def key
+      @message.key
+    end
+
+    # @return [Integer] the offset of the message in the partition.
+    def offset
+      @message.offset
+    end
+
+    # @return [Time] the timestamp of the message.
+    def create_time
+      @message.create_time
+    end
+
+    # @return [Hash<String, String>] the headers of the message.
+    def headers
+      @message.headers
+    end
+
+    # @return [Boolean] whether this record is a control record
+    def is_control_record
+      @message.is_control_record
+    end
+
+  end
+end

data/lib/kafka/fetched_offset_resolver.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module Kafka
+  class FetchedOffsetResolver
+    def initialize(logger:)
+      @logger = TaggedLogger.new(logger)
+    end
+
+    def resolve!(broker, topics)
+      pending_topics = filter_pending_topics(topics)
+      return topics if pending_topics.empty?
+
+      response = broker.list_offsets(topics: pending_topics)
+
+      pending_topics.each do |topic, partitions|
+        partitions.each do |options|
+          partition = options.fetch(:partition)
+          resolved_offset = response.offset_for(topic, partition)
+
+          @logger.debug "Offset for #{topic}/#{partition} is #{resolved_offset.inspect}"
+
+          topics[topic][partition][:fetch_offset] = resolved_offset || 0
+        end
+      end
+    end
+
+    private
+
+    def filter_pending_topics(topics)
+      pending_topics = {}
+      topics.each do |topic, partitions|
+        partitions.each do |partition, options|
+          offset = options.fetch(:fetch_offset)
+          next if offset >= 0
+
+          @logger.debug "Resolving offset `#{offset}` for #{topic}/#{partition}..."
+
+          pending_topics[topic] ||= []
+          pending_topics[topic] << {
+            partition: partition,
+            time: offset
+          }
+        end
+      end
+      pending_topics
+    end
+  end
+end

data/lib/kafka/fetcher.rb

@@ -0,0 +1,224 @@
+# frozen_string_literal: true
+
+require "kafka/fetch_operation"
+
+module Kafka
+  class Fetcher
+    attr_reader :queue, :max_wait_time
+
+    def initialize(cluster:, logger:, instrumenter:, max_queue_size:, group:)
+      @cluster = cluster
+      @logger = TaggedLogger.new(logger)
+      @instrumenter = instrumenter
+      @max_queue_size = max_queue_size
+      @group = group
+
+      @queue = Queue.new
+      @commands = Queue.new
+      @next_offsets = Hash.new { |h, k| h[k] = {} }
+
+      # We are only running when someone calls start.
+      @running = false
+
+      # Long poll until at least this many bytes can be fetched.
+      @min_bytes = 1
+
+      # Long poll at most this number of seconds.
+      @max_wait_time = 1
+
+      # The maximum number of bytes to fetch for any given fetch request.
+      @max_bytes = 10485760
+
+      # The maximum number of bytes to fetch per partition, by topic.
+      @max_bytes_per_partition = {}
+
+      # An incrementing counter used to synchronize resets between the
+      # foreground and background thread.
+      @current_reset_counter = 0
+    end
+
+    def subscribe(topic, max_bytes_per_partition:)
+      @commands << [:subscribe, [topic, max_bytes_per_partition]]
+    end
+
+    def seek(topic, partition, offset)
+      @commands << [:seek, [topic, partition, offset]]
+    end
+
+    def configure(min_bytes:, max_bytes:, max_wait_time:)
+      @commands << [:configure, [min_bytes, max_bytes, max_wait_time]]
+    end
+
+    def start
+      return if @running
+
+      @running = true
+
+      @thread = Thread.new do
+        while @running
+          loop
+        end
+        @logger.info "#{@group} Fetcher thread exited."
+      end
+      @thread.abort_on_exception = true
+    end
+
+    def stop
+      return unless @running
+      @commands << [:stop, []]
+      @thread.join
+    end
+
+    def reset
+      @current_reset_counter = current_reset_counter + 1
+      @commands << [:reset]
+    end
+
+    def data?
+      !@queue.empty?
+    end
+
+    def poll
+      tag, message, reset_counter = @queue.deq
+
+      # Batches are tagged with the current reset counter value. If the batch
+      # has a reset_counter < current_reset_counter, we know it was fetched
+      # prior to the most recent reset and should be discarded.
+      if tag == :batches && message.any? && current_reset_counter > reset_counter
+        @logger.warn "Skipping stale messages buffered prior to reset"
+        return tag, []
+      end
+
+      return [tag, message]
+    end
+
+    private
+
+    attr_reader :current_reset_counter
+
+    def loop
+      @logger.push_tags(@group.to_s)
+      @instrumenter.instrument("loop.fetcher", {
+        queue_size: @queue.size,
+      })
+
+      return unless @running
+
+      if !@commands.empty?
+        cmd, args = @commands.deq
+
+        @logger.debug "Handling fetcher command: #{cmd}"
+
+        send("handle_#{cmd}", *args)
+      elsif @queue.size < @max_queue_size
+        step
+      else
+        @logger.info "Reached max fetcher queue size (#{@max_queue_size}), sleeping 1s"
+        sleep 1
+      end
+    ensure
+      @logger.pop_tags
+    end
+
+    def handle_configure(min_bytes, max_bytes, max_wait_time)
+      @min_bytes = min_bytes
+      @max_bytes = max_bytes
+      @max_wait_time = max_wait_time
+    end
+
+    def handle_reset
+      @next_offsets.clear
+      @queue.clear
+    end
+
+    def handle_stop(*)
+      @running = false
+      @commands.clear
+
+      # After stopping, we need to reconfigure the topics and partitions to fetch
+      # from. Otherwise we'd keep fetching from a bunch of partitions we may no
+      # longer be assigned.
+      handle_reset
+    end
+
+    def handle_subscribe(topic, max_bytes_per_partition)
+      @logger.info "Will fetch at most #{max_bytes_per_partition} bytes at a time per partition from #{topic}"
+      @max_bytes_per_partition[topic] = max_bytes_per_partition
+    end
+
+    def handle_seek(topic, partition, offset)
+      @instrumenter.instrument('seek.consumer',
+                               group_id: @group.group_id,
+                               topic: topic,
+                               partition: partition,
+                               offset: offset)
+      @logger.info "Seeking #{topic}/#{partition} to offset #{offset}"
+      @next_offsets[topic][partition] = offset
+    end
+
+    def step
+      batches = fetch_batches
+
+      batches.each do |batch|
+        unless batch.empty?
+          @instrumenter.instrument("fetch_batch.consumer", {
+            topic: batch.topic,
+            partition: batch.partition,
+            offset_lag: batch.offset_lag,
+            highwater_mark_offset: batch.highwater_mark_offset,
+            message_count: batch.messages.count,
+          })
+        end
+
+        @next_offsets[batch.topic][batch.partition] = batch.last_offset + 1 unless batch.unknown_last_offset?
+      end
+
+      @queue << [:batches, batches, current_reset_counter]
+    rescue Kafka::NoPartitionsToFetchFrom
+      @logger.warn "No partitions to fetch from, sleeping for 1s"
+      sleep 1
+    rescue Kafka::Error => e
+      @queue << [:exception, e]
+    end
+
+    def fetch_batches
+      @logger.debug "Fetching batches"
+
+      operation = FetchOperation.new(
+        cluster: @cluster,
+        logger: @logger,
+        min_bytes: @min_bytes,
+        max_bytes: @max_bytes,
+        max_wait_time: @max_wait_time,
+      )
+
+      @next_offsets.each do |topic, partitions|
+        # Fetch at most this many bytes from any single partition.
+        max_bytes = @max_bytes_per_partition[topic]
+
+        partitions.each do |partition, offset|
+          operation.fetch_from_partition(topic, partition, offset: offset, max_bytes: max_bytes)
+        end
+      end
+
+      operation.execute
+    rescue UnknownTopicOrPartition
+      @logger.error "Failed to fetch from some partitions. Maybe a rebalance has happened? Refreshing cluster info."
+
+      # Our cluster information has become stale, we need to refresh it.
+      @cluster.refresh_metadata!
+
+      # Don't overwhelm the brokers in case this keeps happening.
+      sleep 10
+
+      retry
+    rescue NoPartitionsToFetchFrom
+      backoff = @max_wait_time > 0 ? @max_wait_time : 1
+
+      @logger.info "There are no partitions to fetch from, sleeping for #{backoff}s"
+      sleep backoff
+
+      []
+    end
+  end
+end

data/lib/kafka/gzip_codec.rb

@@ -0,0 +1,34 @@
+# frozen_string_literal: true
+
+module Kafka
+  class GzipCodec
+    def codec_id
+      1
+    end
+
+    def produce_api_min_version
+      0
+    end
+
+    def load
+      require "zlib"
+    end
+
+    def compress(data)
+      buffer = StringIO.new
+      buffer.set_encoding(Encoding::BINARY)
+
+      writer = Zlib::GzipWriter.new(buffer, Zlib::DEFAULT_COMPRESSION, Zlib::DEFAULT_STRATEGY)
+      writer.write(data)
+      writer.close
+
+      buffer.string
+    end
+
+    def decompress(data)
+      buffer = StringIO.new(data)
+      reader = Zlib::GzipReader.new(buffer)
+      reader.read
+    end
+  end
+end

data/lib/kafka/heartbeat.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Kafka
+  class Heartbeat
+    def initialize(group:, interval:, instrumenter:)
+      @group = group
+      @interval = interval
+      @last_heartbeat = Time.now
+      @instrumenter = instrumenter
+    end
+
+    def trigger!
+      @instrumenter.instrument('heartbeat.consumer',
+                               group_id: @group.group_id,
+                               topic_partitions: @group.assigned_partitions) do
+        @group.heartbeat
+        @last_heartbeat = Time.now
+      end
+    end
+
+    def trigger
+      trigger! if Time.now > @last_heartbeat + @interval
+    end
+  end
+end