ruby-kafka-aws-iam 1.4.1

data/lib/kafka/instrumenter.rb

@@ -0,0 +1,38 @@
+# frozen_string_literal: true
+
+module Kafka
+  class Instrumenter
+    NAMESPACE = "kafka"
+
+    def initialize(default_payload = {})
+      @default_payload = default_payload
+
+      if defined?(ActiveSupport::Notifications)
+        @backend = ActiveSupport::Notifications
+      else
+        @backend = nil
+      end
+    end
+
+    def instrument(event_name, payload = {}, &block)
+      if @backend
+        payload.update(@default_payload)
+
+        @backend.instrument("#{event_name}.#{NAMESPACE}", payload, &block)
+      else
+        block.call(payload) if block
+      end
+    end
+  end
+
+  class DecoratingInstrumenter
+    def initialize(backend, extra_payload = {})
+      @backend = backend
+      @extra_payload = extra_payload
+    end
+
+    def instrument(event_name, payload = {}, &block)
+      @backend.instrument(event_name, @extra_payload.merge(payload), &block)
+    end
+  end
+end

data/lib/kafka/interceptors.rb

@@ -0,0 +1,33 @@
+# frozen_string_literal: true
+
+module Kafka
+  # Holds a list of interceptors that implement `call`
+  # and wraps calls to a chain of custom interceptors.
+  class Interceptors
+    def initialize(interceptors:, logger:)
+      @interceptors = interceptors || []
+      @logger = TaggedLogger.new(logger)
+    end
+
+    # This method is called when the client produces a message or once the batches are fetched.
+    # The message returned from the first call is passed to the second interceptor call, and so on in an
+    # interceptor chain. This method does not throw exceptions.
+    #
+    # @param intercepted [Kafka::PendingMessage || Kafka::FetchedBatch] the produced message or
+    #   fetched batch.
+    #
+    # @return [Kafka::PendingMessage || Kafka::FetchedBatch] the intercepted message or batch
+    #   returned by the last interceptor.
+    def call(intercepted)
+      @interceptors.each do |interceptor|
+        begin
+          intercepted = interceptor.call(intercepted)
+        rescue Exception => e
+          @logger.warn "Error executing interceptor for topic: #{intercepted.topic} partition: #{intercepted.partition}: #{e.message}\n#{e.backtrace.join("\n")}"
+        end
+      end
+
+      intercepted
+    end
+  end
+end

data/lib/kafka/lz4_codec.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Kafka
+  class LZ4Codec
+    def codec_id
+      3
+    end
+
+    def produce_api_min_version
+      0
+    end
+
+    def load
+      require "extlz4"
+    rescue LoadError
+      raise LoadError, "using lz4 compression requires adding a dependency on the `extlz4` gem to your Gemfile."
+    end
+
+    def compress(data)
+      LZ4.encode(data)
+    end
+
+    def decompress(data)
+      LZ4.decode(data)
+    end
+  end
+end

data/lib/kafka/message_buffer.rb

@@ -0,0 +1,87 @@
+# frozen_string_literal: true
+
+require "kafka/protocol/message"
+
+module Kafka
+
+  # Buffers messages for specific topics/partitions.
+  class MessageBuffer
+    include Enumerable
+
+    attr_reader :size, :bytesize
+
+    def initialize
+      @buffer = {}
+      @size = 0
+      @bytesize = 0
+    end
+
+    def write(value:, key:, topic:, partition:, create_time: Time.now, headers: {})
+      message = Protocol::Record.new(key: key, value: value, create_time: create_time, headers: headers)
+
+      buffer_for(topic, partition) << message
+
+      @size += 1
+      @bytesize += message.bytesize
+    end
+
+    def concat(messages, topic:, partition:)
+      buffer_for(topic, partition).concat(messages)
+
+      @size += messages.count
+      @bytesize += messages.map(&:bytesize).reduce(0, :+)
+    end
+
+    def to_h
+      @buffer
+    end
+
+    def empty?
+      @buffer.empty?
+    end
+
+    def each
+      @buffer.each do |topic, messages_for_topic|
+        messages_for_topic.each do |partition, messages_for_partition|
+          yield topic, partition, messages_for_partition
+        end
+      end
+    end
+
+    # Clears buffered messages for the given topic and partition.
+    #
+    # @param topic [String] the name of the topic.
+    # @param partition [Integer] the partition id.
+    #
+    # @return [nil]
+    def clear_messages(topic:, partition:)
+      return unless @buffer.key?(topic) && @buffer[topic].key?(partition)
+
+      @size -= @buffer[topic][partition].count
+      @bytesize -= @buffer[topic][partition].map(&:bytesize).reduce(0, :+)
+
+      @buffer[topic].delete(partition)
+      @buffer.delete(topic) if @buffer[topic].empty?
+    end
+
+    def messages_for(topic:, partition:)
+      buffer_for(topic, partition)
+    end
+
+    # Clears messages across all topics and partitions.
+    #
+    # @return [nil]
+    def clear
+      @buffer = {}
+      @size = 0
+      @bytesize = 0
+    end
+
+    private
+
+    def buffer_for(topic, partition)
+      @buffer[topic] ||= {}
+      @buffer[topic][partition] ||= []
+    end
+  end
+end

data/lib/kafka/murmur2_hash.rb

@@ -0,0 +1,17 @@
+# frozen_string_literal: true
+
+module Kafka
+  class Murmur2Hash
+    SEED = [0x9747b28c].pack('L')
+
+    def load
+      require 'digest/murmurhash'
+    rescue LoadError
+      raise LoadError, "using murmur2 hashing requires adding a dependency on the `digest-murmurhash` gem to your Gemfile."
+    end
+
+    def hash(value)
+      ::Digest::MurmurHash2.rawdigest(value, SEED) & 0x7fffffff
+    end
+  end
+end

data/lib/kafka/offset_manager.rb

@@ -0,0 +1,259 @@
+# frozen_string_literal: true
+
+module Kafka
+
+  # Manages a consumer's position in partitions, figures out where to resume processing
+  # from, etc.
+  class OffsetManager
+
+    # The default broker setting for offsets.retention.minutes is 1440.
+    DEFAULT_RETENTION_TIME = 1440 * 60
+
+    def initialize(cluster:, group:, fetcher:, logger:, commit_interval:, commit_threshold:, offset_retention_time:)
+      @cluster = cluster
+      @group = group
+      @fetcher = fetcher
+      @logger = TaggedLogger.new(logger)
+      @commit_interval = commit_interval
+      @commit_threshold = commit_threshold
+
+      @uncommitted_offsets = 0
+      @processed_offsets = {}
+      @default_offsets = {}
+      @committed_offsets = nil
+      @resolved_offsets = {}
+      @last_commit = Time.now
+      @last_recommit = nil
+      @recommit_interval = (offset_retention_time || DEFAULT_RETENTION_TIME) / 2
+    end
+
+    # Set the default offset for a topic.
+    #
+    # When the consumer is started for the first time, or in cases where it gets stuck and
+    # has to reset its position, it must start either with the earliest messages or with
+    # the latest, skipping to the very end of each partition.
+    #
+    # @param topic [String] the name of the topic.
+    # @param default_offset [Symbol] either `:earliest` or `:latest`.
+    # @return [nil]
+    def set_default_offset(topic, default_offset)
+      @default_offsets[topic] = default_offset
+    end
+
+    # Mark a message as having been processed.
+    #
+    # When offsets are committed, the message's offset will be stored in Kafka so
+    # that we can resume from this point at a later time.
+    #
+    # @param topic [String] the name of the topic.
+    # @param partition [Integer] the partition number.
+    # @param offset [Integer] the offset of the message that should be marked as processed.
+    # @return [nil]
+    def mark_as_processed(topic, partition, offset)
+      unless @group.assigned_to?(topic, partition)
+        @logger.debug "Not marking #{topic}/#{partition}:#{offset} as processed for partition not assigned to this consumer."
+        return
+      end
+      @processed_offsets[topic] ||= {}
+
+      last_processed_offset = @processed_offsets[topic][partition] || -1
+      if last_processed_offset > offset + 1
+        @logger.debug "Not overwriting newer offset #{topic}/#{partition}:#{last_processed_offset - 1} with older #{offset}"
+        return
+      end
+
+      @uncommitted_offsets += 1
+
+      # The committed offset should always be the offset of the next message that the
+      # application will read, thus adding one to the last message processed.
+      @processed_offsets[topic][partition] = offset + 1
+      @logger.debug "Marking #{topic}/#{partition}:#{offset} as processed"
+    end
+
+    # Move the consumer's position in the partition back to the configured default
+    # offset, either the first or latest in the partition.
+    #
+    # @param topic [String] the name of the topic.
+    # @param partition [Integer] the partition number.
+    # @return [nil]
+    def seek_to_default(topic, partition)
+      # Remove any cached offset, in case things have changed broker-side.
+      clear_resolved_offset(topic)
+
+      offset = resolve_offset(topic, partition)
+
+      seek_to(topic, partition, offset)
+    end
+
+    # Move the consumer's position in the partition to the specified offset.
+    #
+    # @param topic [String] the name of the topic.
+    # @param partition [Integer] the partition number.
+    # @param offset [Integer] the offset that the consumer position should be moved to.
+    # @return [nil]
+    def seek_to(topic, partition, offset)
+      @processed_offsets[topic] ||= {}
+      @processed_offsets[topic][partition] = offset
+
+      @fetcher.seek(topic, partition, offset)
+    end
+
+    # Return the next offset that should be fetched for the specified partition.
+    #
+    # @param topic [String] the name of the topic.
+    # @param partition [Integer] the partition number.
+    # @return [Integer] the next offset that should be fetched.
+    def next_offset_for(topic, partition)
+      offset = @processed_offsets.fetch(topic, {}).fetch(partition) {
+        committed_offset_for(topic, partition)
+      }
+
+      # A negative offset means that no offset has been committed, so we need to
+      # resolve the default offset for the topic.
+      if offset < 0
+        resolve_offset(topic, partition)
+      else
+        # The next offset is the last offset.
+        offset
+      end
+    end
+
+    # Commit offsets of messages that have been marked as processed.
+    #
+    # If `recommit` is set to true, we will also commit the existing positions
+    # even if no messages have been processed on a partition. This is done
+    # in order to avoid the offset information expiring in cases where messages
+    # are very rare -- it's essentially a keep-alive.
+    #
+    # @param recommit [Boolean] whether to recommit offsets that have already been
+    #   committed.
+    # @return [nil]
+    def commit_offsets(recommit = false)
+      offsets = offsets_to_commit(recommit)
+      unless offsets.empty?
+        @logger.debug "Committing offsets#{recommit ? ' with recommit' : ''}: #{prettify_offsets(offsets)}"
+
+        @group.commit_offsets(offsets)
+
+        @last_commit = Time.now
+        @last_recommit = Time.now if recommit
+
+        @uncommitted_offsets = 0
+        @committed_offsets = nil
+      end
+    end
+
+    # Commit offsets if necessary, according to the offset commit policy specified
+    # when initializing the class.
+    #
+    # @return [nil]
+    def commit_offsets_if_necessary
+      recommit = recommit_timeout_reached?
+      if recommit || commit_timeout_reached? || commit_threshold_reached?
+        commit_offsets(recommit)
+      end
+    end
+
+    # Clear all stored offset information.
+    #
+    # @return [nil]
+    def clear_offsets
+      @processed_offsets.clear
+      @resolved_offsets.clear
+
+      # Clear the cached commits from the brokers.
+      @committed_offsets = nil
+    end
+
+    # Clear stored offset information for all partitions except those specified
+    # in `excluded`.
+    #
+    #     offset_manager.clear_offsets_excluding("my-topic" => [1, 2, 3])
+    #
+    # @return [nil]
+    def clear_offsets_excluding(excluded)
+      # Clear all offsets that aren't in `excluded`.
+      @processed_offsets.each do |topic, partitions|
+        partitions.keep_if do |partition, _|
+          excluded.fetch(topic, []).include?(partition)
+        end
+      end
+
+      # Clear the cached commits from the brokers.
+      @committed_offsets = nil
+      @resolved_offsets.clear
+    end
+
+    private
+
+    def clear_resolved_offset(topic)
+      @resolved_offsets.delete(topic)
+    end
+
+    def resolve_offset(topic, partition)
+      @resolved_offsets[topic] ||= fetch_resolved_offsets(topic)
+      @resolved_offsets[topic].fetch(partition)
+    end
+
+    def fetch_resolved_offsets(topic)
+      default_offset = @default_offsets.fetch(topic)
+      partitions = @group.assigned_partitions.fetch(topic)
+
+      @cluster.resolve_offsets(topic, partitions, default_offset)
+    end
+
+    def seconds_since(time)
+      Time.now - time
+    end
+
+    def seconds_since_last_commit
+      seconds_since(@last_commit)
+    end
+
+    def committed_offsets
+      @committed_offsets ||= @group.fetch_offsets
+    end
+
+    def committed_offset_for(topic, partition)
+      committed_offsets.offset_for(topic, partition)
+    end
+
+    def offsets_to_commit(recommit = false)
+      if recommit
+        offsets_to_recommit.merge!(@processed_offsets) do |_topic, committed, processed|
+          committed.merge!(processed)
+        end
+      else
+        @processed_offsets
+      end
+    end
+
+    def offsets_to_recommit
+      committed_offsets.topics.each_with_object({}) do |(topic, partition_info), offsets|
+        topic_offsets = partition_info.keys.each_with_object({}) do |partition, partition_map|
+          offset = committed_offsets.offset_for(topic, partition)
+          partition_map[partition] = offset unless offset == -1
+        end
+        offsets[topic] = topic_offsets unless topic_offsets.empty?
+      end
+    end
+
+    def recommit_timeout_reached?
+      @last_recommit.nil? || seconds_since(@last_recommit) >= @recommit_interval
+    end
+
+    def commit_timeout_reached?
+      @commit_interval != 0 && seconds_since_last_commit >= @commit_interval
+    end
+
+    def commit_threshold_reached?
+      @commit_threshold != 0 && @uncommitted_offsets >= @commit_threshold
+    end
+
+    def prettify_offsets(offsets)
+      offsets.flat_map do |topic, partitions|
+        partitions.map { |partition, offset| "#{topic}/#{partition}:#{offset}" }
+      end.join(', ')
+    end
+  end
+end

data/lib/kafka/partitioner.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+require "kafka/digest"
+
+module Kafka
+
+  # Assigns partitions to messages.
+  class Partitioner
+    # @param hash_function [Symbol, nil] the algorithm used to compute a messages
+    #   destination partition. Default is :crc32
+    def initialize(hash_function: nil)
+      @digest = Digest.find_digest(hash_function || :crc32)
+    end
+
+    # Assigns a partition number based on a partition key. If no explicit
+    # partition key is provided, the message key will be used instead.
+    #
+    # If the key is nil, then a random partition is selected. Otherwise, a digest
+    # of the key is used to deterministically find a partition. As long as the
+    # number of partitions doesn't change, the same key will always be assigned
+    # to the same partition.
+    #
+    # @param partition_count [Integer] the number of partitions in the topic.
+    # @param message [Kafka::PendingMessage] the message that should be assigned
+    #   a partition.
+    # @return [Integer] the partition number.
+    def call(partition_count, message)
+      raise ArgumentError if partition_count == 0
+
+      # If no explicit partition key is specified we use the message key instead.
+      key = message.partition_key || message.key
+
+      if key.nil?
+        rand(partition_count)
+      else
+        @digest.hash(key) % partition_count
+      end
+    end
+  end
+end

data/lib/kafka/pause.rb

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module Kafka
+  # Manages the pause state of a partition.
+  #
+  # The processing of messages in a partition can be paused, e.g. if there was
+  # an exception during processing. This could be caused by a downstream service
+  # not being available. A typical way of solving such an issue is to back off
+  # for a little while and then try again. In order to do that, _pause_ the
+  # partition.
+  class Pause
+    def initialize(clock: Time)
+      @clock = clock
+      @started_at = nil
+      @pauses = 0
+      @timeout = nil
+      @max_timeout = nil
+      @exponential_backoff = false
+    end
+
+    # Mark the partition as paused.
+    #
+    # If exponential backoff is enabled, each subsequent pause of a partition will
+    # cause a doubling of the actual timeout, i.e. for pause number _n_, the actual
+    # timeout will be _2^n * timeout_.
+    #
+    # Only when {#reset!} is called is this state cleared.
+    #
+    # @param timeout [nil, Integer] if specified, the partition will automatically
+    #   resume after this many seconds.
+    # @param exponential_backoff [Boolean] whether to enable exponential timeouts.
+    def pause!(timeout: nil, max_timeout: nil, exponential_backoff: false)
+      @started_at = @clock.now
+      @timeout = timeout
+      @max_timeout = max_timeout
+      @exponential_backoff = exponential_backoff
+      @pauses += 1
+    end
+
+    # Resumes the partition.
+    #
+    # The number of pauses is still retained, and if the partition is paused again
+    # it may be with an exponential backoff.
+    def resume!
+      @started_at = nil
+      @timeout = nil
+      @max_timeout = nil
+    end
+
+    # Whether the partition is currently paused. The pause may have expired, in which
+    # case {#expired?} should be checked as well.
+    def paused?
+      # This is nil if we're not currently paused.
+      !@started_at.nil?
+    end
+
+    def pause_duration
+      if paused?
+        Time.now - @started_at
+      else
+        0
+      end
+    end
+
+    # Whether the pause has expired.
+    def expired?
+      # We never expire the pause if timeout is nil.
+      return false if @timeout.nil?
+
+      # Have we passed the end of the pause duration?
+      @clock.now >= ends_at
+    end
+
+    # Resets the pause state, ensuring that the next pause is not exponential.
+    def reset!
+      @pauses = 0
+    end
+
+    private
+
+    def ends_at
+      # Apply an exponential backoff to the timeout.
+      backoff_factor = @exponential_backoff ? 2**(@pauses - 1) : 1
+      timeout = backoff_factor * @timeout
+
+      # If set, don't allow a timeout longer than max_timeout.
+      timeout = @max_timeout if @max_timeout && timeout > @max_timeout
+
+      @started_at + timeout
+    end
+  end
+end

data/lib/kafka/pending_message.rb

@@ -0,0 +1,29 @@
+# frozen_string_literal: true
+
+module Kafka
+  class PendingMessage
+    attr_reader :value, :key, :headers, :topic, :partition, :partition_key, :create_time, :bytesize
+
+    def initialize(value:, key:, headers: {}, topic:, partition:, partition_key:, create_time:)
+      @value = value
+      @key = key
+      @headers = headers
+      @topic = topic
+      @partition = partition
+      @partition_key = partition_key
+      @create_time = create_time
+      @bytesize = key.to_s.bytesize + value.to_s.bytesize
+    end
+
+    def ==(other)
+      @value == other.value &&
+        @key == other.key &&
+        @topic == other.topic &&
+        @headers == other.headers &&
+        @partition == other.partition &&
+        @partition_key == other.partition_key &&
+        @create_time == other.create_time &&
+        @bytesize == other.bytesize
+    end
+  end
+end

data/lib/kafka/pending_message_queue.rb

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module Kafka
+
+  class PendingMessageQueue
+    attr_reader :size, :bytesize
+
+    def initialize
+      clear
+    end
+
+    def write(message)
+      @messages << message
+      @size += 1
+      @bytesize += message.bytesize
+    end
+
+    def empty?
+      @messages.empty?
+    end
+
+    def clear
+      @messages = []
+      @size = 0
+      @bytesize = 0
+    end
+
+    def replace(messages)
+      clear
+      messages.each {|message| write(message) }
+    end
+
+    # Yields each message in the queue.
+    #
+    # @yieldparam [PendingMessage] message
+    # @return [nil]
+    def each(&block)
+      @messages.each(&block)
+    end
+  end
+end