ruby-kafka-aws-iam 1.4.1

data/lib/kafka/tagged_logger.rb

@@ -0,0 +1,77 @@
+# Basic implementation of a tagged logger that matches the API of
+# ActiveSupport::TaggedLogging.
+
+require 'delegate'
+require 'logger'
+
+module Kafka
+  class TaggedLogger < SimpleDelegator
+
+    %i(debug info warn error).each do |method|
+      define_method method do |msg_or_progname, &block|
+        if block_given?
+          super(msg_or_progname, &block)
+        else
+          super("#{tags_text}#{msg_or_progname}")
+        end
+      end
+    end
+
+    def tagged(*tags)
+      new_tags = push_tags(*tags)
+      yield self
+    ensure
+      pop_tags(new_tags.size)
+    end
+
+    def push_tags(*tags)
+      tags.flatten.reject { |t| t.nil? || t.empty? }.tap do |new_tags|
+        current_tags.concat new_tags
+      end
+    end
+
+    def pop_tags(size = 1)
+      current_tags.pop size
+    end
+
+    def clear_tags!
+      current_tags.clear
+    end
+
+    def current_tags
+      # We use our object ID here to avoid conflicting with other instances
+      thread_key = @thread_key ||= "kafka_tagged_logging_tags:#{object_id}".freeze
+      Thread.current[thread_key] ||= []
+    end
+
+    def tags_text
+      tags = current_tags
+      if tags.any?
+        tags.collect { |tag| "[#{tag}] " }.join
+      end
+    end
+
+    def self.new(logger_or_stream = nil)
+      # don't keep wrapping the same logger over and over again
+      return logger_or_stream if logger_or_stream.is_a?(TaggedLogger)
+      super
+    end
+
+    def initialize(logger_or_stream = nil)
+      logger = if %w(info debug warn error).all? { |s| logger_or_stream.respond_to?(s) }
+        logger_or_stream
+      elsif logger_or_stream
+        ::Logger.new(logger_or_stream)
+      else
+        ::Logger.new(nil)
+      end
+      super(logger)
+    end
+
+    def flush
+      clear_tags!
+      super if defined?(super)
+    end
+  end
+
+end

data/lib/kafka/transaction_manager.rb

@@ -0,0 +1,306 @@
+# frozen_string_literal: true
+
+require 'kafka/transaction_state_machine'
+
+module Kafka
+  class TransactionManager
+    DEFAULT_TRANSACTION_TIMEOUT = 60 # 60 seconds
+    TRANSACTION_RESULT_COMMIT = true
+    TRANSACTION_RESULT_ABORT = false
+
+    attr_reader :producer_id, :producer_epoch, :transactional_id
+
+    def initialize(
+      cluster:,
+      logger:,
+      idempotent: false,
+      transactional: false,
+      transactional_id: nil,
+      transactional_timeout: DEFAULT_TRANSACTION_TIMEOUT
+    )
+      @cluster = cluster
+      @logger = TaggedLogger.new(logger)
+
+      @transactional = transactional
+      @transactional_id = transactional_id
+      @transactional_timeout = transactional_timeout
+      @transaction_state = Kafka::TransactionStateMachine.new(logger: logger)
+      @transaction_partitions = {}
+
+      # If transactional mode is enabled, idempotent must be enabled
+      @idempotent = transactional || idempotent
+
+      @producer_id = -1
+      @producer_epoch = 0
+
+      @sequences = {}
+    end
+
+    def idempotent?
+      @idempotent == true
+    end
+
+    def transactional?
+      @transactional == true && !@transactional_id.nil?
+    end
+
+    def init_producer_id(force = false)
+      return if @producer_id >= 0 && !force
+
+      response = transaction_coordinator.init_producer_id(
+        transactional_id: @transactional_id,
+        transactional_timeout: @transactional_timeout
+      )
+      Protocol.handle_error(response.error_code)
+
+      # Reset producer id
+      @producer_id = response.producer_id
+      @producer_epoch = response.producer_epoch
+
+      # Reset sequence
+      @sequences = {}
+
+      @logger.debug "Current Producer ID is #{@producer_id} and Producer Epoch is #{@producer_epoch}"
+    end
+
+    def next_sequence_for(topic, partition)
+      @sequences[topic] ||= {}
+      @sequences[topic][partition] ||= 0
+    end
+
+    def update_sequence_for(topic, partition, sequence)
+      @sequences[topic] ||= {}
+      @sequences[topic][partition] = sequence
+    end
+
+    def init_transactions
+      force_transactional!
+      unless @transaction_state.uninitialized?
+        @logger.warn("Transaction already initialized!")
+        return
+      end
+      init_producer_id(true)
+      @transaction_partitions = {}
+      @transaction_state.transition_to!(TransactionStateMachine::READY)
+
+      @logger.info "Transaction #{@transactional_id} is initialized, Producer ID: #{@producer_id} (Epoch #{@producer_epoch})"
+
+      nil
+    rescue
+      @transaction_state.transition_to!(TransactionStateMachine::ERROR)
+      raise
+    end
+
+    def add_partitions_to_transaction(topic_partitions)
+      force_transactional!
+
+      if @transaction_state.uninitialized?
+        raise Kafka::InvalidTxnStateError, 'Transaction is uninitialized'
+      end
+
+      # Extract newly created partitions
+      new_topic_partitions = {}
+      topic_partitions.each do |topic, partitions|
+        partitions.each do |partition|
+          @transaction_partitions[topic] ||= {}
+          if !@transaction_partitions[topic][partition]
+            new_topic_partitions[topic] ||= []
+            new_topic_partitions[topic] << partition
+
+            @logger.info "Adding parition #{topic}/#{partition}  to transaction #{@transactional_id}, Producer ID: #{@producer_id} (Epoch #{@producer_epoch})"
+          end
+        end
+      end
+
+      unless new_topic_partitions.empty?
+        response = transaction_coordinator.add_partitions_to_txn(
+          transactional_id: @transactional_id,
+          producer_id: @producer_id,
+          producer_epoch: @producer_epoch,
+          topics: new_topic_partitions
+        )
+
+        # Update added topic partitions
+        response.errors.each do |tp|
+          tp.partitions.each do |p|
+            Protocol.handle_error(p.error_code)
+            @transaction_partitions[tp.topic] ||= {}
+            @transaction_partitions[tp.topic][p.partition] = true
+          end
+        end
+      end
+
+      nil
+    rescue
+      @transaction_state.transition_to!(TransactionStateMachine::ERROR)
+      raise
+    end
+
+    def begin_transaction
+      force_transactional!
+      raise Kafka::InvalidTxnStateError, 'Transaction has already started' if @transaction_state.in_transaction?
+      raise Kafka::InvalidTxnStateError, 'Transaction is not ready' unless @transaction_state.ready?
+      @transaction_state.transition_to!(TransactionStateMachine::IN_TRANSACTION)
+
+      @logger.info "Begin transaction #{@transactional_id}, Producer ID: #{@producer_id} (Epoch #{@producer_epoch})"
+
+      nil
+    rescue
+      @transaction_state.transition_to!(TransactionStateMachine::ERROR)
+      raise
+    end
+
+    def commit_transaction
+      force_transactional!
+
+      if @transaction_state.committing_transaction?
+        @logger.warn("Transaction is being committed")
+        return
+      end
+
+      unless @transaction_state.in_transaction?
+        raise Kafka::InvalidTxnStateError, 'Transaction is not valid to commit'
+      end
+
+      @transaction_state.transition_to!(TransactionStateMachine::COMMITTING_TRANSACTION)
+
+      @logger.info "Commiting transaction #{@transactional_id}, Producer ID: #{@producer_id} (Epoch #{@producer_epoch})"
+
+      response = transaction_coordinator.end_txn(
+        transactional_id: @transactional_id,
+        producer_id: @producer_id,
+        producer_epoch: @producer_epoch,
+        transaction_result: TRANSACTION_RESULT_COMMIT
+      )
+      Protocol.handle_error(response.error_code)
+
+      @logger.info "Transaction #{@transactional_id} is committed, Producer ID: #{@producer_id} (Epoch #{@producer_epoch})"
+      complete_transaction
+
+      nil
+    rescue
+      @transaction_state.transition_to!(TransactionStateMachine::ERROR)
+      raise
+    end
+
+    def abort_transaction
+      force_transactional!
+
+      if @transaction_state.aborting_transaction?
+        @logger.warn("Transaction is being aborted")
+        return
+      end
+
+      unless @transaction_state.in_transaction?
+        @logger.warn('Aborting transaction that was never opened on brokers')
+        return
+      end
+
+      @transaction_state.transition_to!(TransactionStateMachine::ABORTING_TRANSACTION)
+
+      @logger.info "Aborting transaction #{@transactional_id}, Producer ID: #{@producer_id} (Epoch #{@producer_epoch})"
+
+      response = transaction_coordinator.end_txn(
+        transactional_id: @transactional_id,
+        producer_id: @producer_id,
+        producer_epoch: @producer_epoch,
+        transaction_result: TRANSACTION_RESULT_ABORT
+      )
+      Protocol.handle_error(response.error_code)
+
+      @logger.info "Transaction #{@transactional_id} is aborted, Producer ID: #{@producer_id} (Epoch #{@producer_epoch})"
+
+      complete_transaction
+
+      nil
+    rescue
+      @transaction_state.transition_to!(TransactionStateMachine::ERROR)
+      raise
+    end
+
+    def send_offsets_to_txn(offsets:, group_id:)
+      force_transactional!
+
+      unless @transaction_state.in_transaction?
+        raise Kafka::InvalidTxnStateError, 'Transaction is not valid to send offsets'
+      end
+
+      add_response = transaction_coordinator.add_offsets_to_txn(
+        transactional_id: @transactional_id,
+        producer_id: @producer_id,
+        producer_epoch: @producer_epoch,
+        group_id: group_id
+      )
+      Protocol.handle_error(add_response.error_code)
+
+      send_response = group_coordinator(group_id: group_id).txn_offset_commit(
+        transactional_id: @transactional_id,
+        group_id: group_id,
+        producer_id: @producer_id,
+        producer_epoch: @producer_epoch,
+        offsets: offsets
+      )
+      send_response.errors.each do |tp|
+        tp.partitions.each do |partition|
+          Protocol.handle_error(partition.error_code)
+        end
+      end
+
+      nil
+    rescue
+      @transaction_state.transition_to!(TransactionStateMachine::ERROR)
+      raise
+    end
+
+    def in_transaction?
+      @transaction_state.in_transaction?
+    end
+
+    def error?
+      @transaction_state.error?
+    end
+
+    def ready?
+      @transaction_state.ready?
+    end
+
+    def close
+      if in_transaction?
+        @logger.warn("Aborting pending transaction ...")
+        abort_transaction
+      elsif @transaction_state.aborting_transaction? || @transaction_state.committing_transaction?
+        @logger.warn("Transaction is finishing. Sleeping until finish!")
+        sleep 5
+      end
+    end
+
+    private
+
+    def force_transactional!
+      unless transactional?
+        raise Kafka::InvalidTxnStateError, 'Please turn on transactional mode to use transaction'
+      end
+
+      if @transactional_id.nil? || @transactional_id.empty?
+        raise Kafka::InvalidTxnStateError, 'Please provide a transaction_id to use transactional mode'
+      end
+    end
+
+    def transaction_coordinator
+      @cluster.get_transaction_coordinator(
+        transactional_id: @transactional_id
+      )
+    end
+
+    def group_coordinator(group_id:)
+      @cluster.get_group_coordinator(
+        group_id: group_id
+      )
+    end
+
+    def complete_transaction
+      @transaction_state.transition_to!(TransactionStateMachine::READY)
+      @transaction_partitions = {}
+    end
+  end
+end

data/lib/kafka/transaction_state_machine.rb

@@ -0,0 +1,72 @@
+# frozen_string_literal: true
+
+module Kafka
+  class TransactionStateMachine
+    class InvalidTransitionError < StandardError; end
+    class InvalidStateError < StandardError; end
+
+    STATES = [
+      UNINITIALIZED          = :uninitialized,
+      READY                  = :ready,
+      IN_TRANSACTION         = :in_trasaction,
+      COMMITTING_TRANSACTION = :committing_transaction,
+      ABORTING_TRANSACTION   = :aborting_transaction,
+      ERROR                  = :error
+    ]
+
+    TRANSITIONS = {
+      UNINITIALIZED          => [READY, ERROR],
+      READY                  => [UNINITIALIZED, COMMITTING_TRANSACTION, ABORTING_TRANSACTION],
+      IN_TRANSACTION         => [READY],
+      COMMITTING_TRANSACTION => [IN_TRANSACTION],
+      ABORTING_TRANSACTION   => [IN_TRANSACTION],
+      # Any states can transition to error state
+      ERROR                  => STATES
+    }
+
+    def initialize(logger:)
+      @state = UNINITIALIZED
+      @mutex = Mutex.new
+      @logger = TaggedLogger.new(logger)
+    end
+
+    def transition_to!(next_state)
+      raise InvalidStateError unless STATES.include?(next_state)
+      unless TRANSITIONS[next_state].include?(@state)
+        raise InvalidTransitionError, "Could not transition from state '#{@state}' to state '#{next_state}'"
+      end
+      @logger.debug("Transaction state changed to '#{next_state}'!")
+      @mutex.synchronize { @state = next_state }
+    end
+
+    def uninitialized?
+      in_state?(UNINITIALIZED)
+    end
+
+    def ready?
+      in_state?(READY)
+    end
+
+    def in_transaction?
+      in_state?(IN_TRANSACTION)
+    end
+
+    def committing_transaction?
+      in_state?(COMMITTING_TRANSACTION)
+    end
+
+    def aborting_transaction?
+      in_state?(ABORTING_TRANSACTION)
+    end
+
+    def error?
+      in_state?(ERROR)
+    end
+
+    private
+
+    def in_state?(state)
+      @mutex.synchronize { @state == state }
+    end
+  end
+end

data/lib/kafka/version.rb

@@ -0,0 +1,5 @@
+# frozen_string_literal: true
+
+module Kafka
+  VERSION = "1.4.1"
+end

data/lib/kafka/zstd_codec.rb

@@ -0,0 +1,27 @@
+# frozen_string_literal: true
+
+module Kafka
+  class ZstdCodec
+    def codec_id
+      4
+    end
+
+    def produce_api_min_version
+      7
+    end
+
+    def load
+      require "zstd-ruby"
+    rescue LoadError
+      raise LoadError, "using zstd compression requires adding a dependency on the `zstd-ruby` gem to your Gemfile."
+    end
+
+    def compress(data)
+      Zstd.compress(data)
+    end
+
+    def decompress(data)
+      Zstd.decompress(data)
+    end
+  end
+end