ruby-kafka 0.7.0 → 0.7.1.beta1

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/list_offset_request.rb

@@ -8,12 +8,14 @@ module Kafka
     #
     #     OffsetRequest => ReplicaId [TopicName [Partition Time MaxNumberOfOffsets]]
     #       ReplicaId => int32
+    #       IsolationLevel => int8
     #       TopicName => string
     #       Partition => int32
     #       Time => int64
-    #       MaxNumberOfOffsets => int32
     #
     class ListOffsetRequest
+      ISOLATION_READ_UNCOMMITTED = 0
+      ISOLATION_READ_COMMITTED = 1
 
       # @param topics [Hash]
       def initialize(topics:)
@@ -21,6 +23,10 @@ module Kafka
         @topics = topics
       end
 
+      def api_version
+        2
+      end
+
       def api_key
         LIST_OFFSET_API
       end
@@ -31,6 +37,7 @@ module Kafka
 
       def encode(encoder)
         encoder.write_int32(@replica_id)
+        encoder.write_int8(ISOLATION_READ_COMMITTED)
 
         encoder.write_array(@topics) do |topic, partitions|
           encoder.write_string(topic)
@@ -38,7 +45,6 @@ module Kafka
           encoder.write_array(partitions) do |partition|
             encoder.write_int32(partition.fetch(:partition))
             encoder.write_int64(partition.fetch(:time))
-            encoder.write_int32(partition.fetch(:max_offsets))
           end
         end
       end

data/lib/kafka/protocol/list_offset_response.rb

@@ -8,9 +8,11 @@ module Kafka
     # ## API Specification
     #
     #     OffsetResponse => [TopicName [PartitionOffsets]]
-    #       PartitionOffsets => Partition ErrorCode [Offset]
+    #       ThrottleTimeMS => int32
+    #       PartitionOffsets => Partition ErrorCode Timestamp Offset
     #       Partition => int32
     #       ErrorCode => int16
+    #       Timestamp => int64
     #       Offset => int64
     #
     class ListOffsetResponse
@@ -24,12 +26,13 @@ module Kafka
       end
 
       class PartitionOffsetInfo
-        attr_reader :partition, :error_code, :offsets
+        attr_reader :partition, :error_code, :timestamp, :offset
 
-        def initialize(partition:, error_code:, offsets:)
+        def initialize(partition:, error_code:, timestamp:, offset:)
           @partition = partition
           @error_code = error_code
-          @offsets = offsets
+          @timestamp = timestamp
+          @offset = offset
         end
       end
 
@@ -56,10 +59,11 @@ module Kafka
 
         Protocol.handle_error(partition_info.error_code)
 
-        partition_info.offsets.first
+        partition_info.offset
       end
 
       def self.decode(decoder)
+        _throttle_time_ms = decoder.int32
         topics = decoder.array do
           name = decoder.string
 
@@ -67,7 +71,8 @@ module Kafka
             PartitionOffsetInfo.new(
               partition: decoder.int32,
               error_code: decoder.int16,
-              offsets: decoder.array { decoder.int64 },
+              timestamp: decoder.int64,
+              offset: decoder.int64
             )
           end
 

data/lib/kafka/protocol/record.rb

@@ -10,6 +10,7 @@ module Kafka
         headers: {},
         attributes: 0,
         offset_delta: 0,
+        offset: 0,
         timestamp_delta: 0,
         create_time: Time.now,
         is_control_record: false
@@ -20,6 +21,7 @@ module Kafka
         @attributes = attributes
 
         @offset_delta = offset_delta
+        @offset = offset
         @timestamp_delta = timestamp_delta
         @create_time = create_time
         @is_control_record = is_control_record
@@ -47,6 +49,13 @@ module Kafka
         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)
 

data/lib/kafka/protocol/record_batch.rb

@@ -30,7 +30,7 @@ module Kafka
           first_sequence: 0,
           max_timestamp: Time.now
       )
-        @records = records
+        @records = Array(records)
         @first_offset = first_offset
         @first_timestamp = first_timestamp
         @codec_id = codec_id
@@ -55,6 +55,10 @@ module Kafka
         @records.size
       end
 
+      def last_offset
+        @first_offset + @last_offset_delta
+      end
+
       def attributes
         0x0000 | @codec_id |
           (@in_transaction ? IN_TRANSACTION_MASK : 0x0) |
@@ -131,6 +135,18 @@ module Kafka
         @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
 

data/lib/kafka/ssl_context.rb

@@ -4,21 +4,35 @@ require "openssl"
 
 module Kafka
   module SslContext
+    CLIENT_CERT_DELIMITER = "\n-----END CERTIFICATE-----\n"
 
-    def self.build(ca_cert_file_path: nil, ca_cert: nil, client_cert: nil, client_cert_key: nil, ca_certs_from_system: nil)
-      return nil unless ca_cert_file_path || ca_cert || client_cert || client_cert_key || ca_certs_from_system
+    def self.build(ca_cert_file_path: nil, ca_cert: nil, client_cert: nil, client_cert_key: nil, client_cert_chain: nil, ca_certs_from_system: nil)
+      return nil unless ca_cert_file_path || ca_cert || client_cert || client_cert_key || client_cert_chain || ca_certs_from_system
 
       ssl_context = OpenSSL::SSL::SSLContext.new
 
       if client_cert && client_cert_key
-        ssl_context.set_params(
+        context_params = {
           cert: OpenSSL::X509::Certificate.new(client_cert),
-          key: OpenSSL::PKey.read(client_cert_key)
-        )
+          key: OpenSSL::PKey.read(client_cert_key),
+        }
+        if client_cert_chain
+          certs = []
+          client_cert_chain.split(CLIENT_CERT_DELIMITER).each do |cert|
+            cert += CLIENT_CERT_DELIMITER
+            certs << OpenSSL::X509::Certificate.new(cert)
+          end
+          context_params[:extra_chain_cert] = certs
+        end
+        ssl_context.set_params(context_params)
       elsif client_cert && !client_cert_key
         raise ArgumentError, "Kafka client initialized with `ssl_client_cert` but no `ssl_client_cert_key`. Please provide both."
       elsif !client_cert && client_cert_key
         raise ArgumentError, "Kafka client initialized with `ssl_client_cert_key`, but no `ssl_client_cert`. Please provide both."
+      elsif client_cert_chain && !client_cert
+        raise ArgumentError, "Kafka client initialized with `ssl_client_cert_chain`, but no `ssl_client_cert`. Please provide cert, key and chain."
+      elsif client_cert_chain && !client_cert_key
+        raise ArgumentError, "Kafka client initialized with `ssl_client_cert_chain`, but no `ssl_client_cert_key`. Please provide cert, key and chain."
       end
 
       if ca_cert || ca_cert_file_path || ca_certs_from_system

data/lib/kafka/transaction_manager.rb

@@ -0,0 +1,261 @@
+# 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 = 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 '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 'Transaction has already started' if @transaction_state.in_transaction?
+      raise '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 '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?
+        raise 'Transaction is not valid to abort'
+      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 in_transaction?
+      @transaction_state.in_transaction?
+    end
+
+    def error?
+      @transaction_state.error?
+    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 'Please turn on transactional mode to use transaction'
+      end
+
+      if @transactional_id.nil? || @transactional_id.empty?
+        raise 'Please provide a transaction_id to use transactional mode'
+      end
+    end
+
+    def transaction_coordinator
+      @cluster.get_transaction_coordinator(
+        transactional_id: @transactional_id
+      )
+    end
+
+    def complete_transaction
+      @transaction_state.transition_to!(TransactionStateMachine::READY)
+      @transaction_partitions = {}
+    end
+  end
+end