ruby-kafka 0.7.0 → 0.7.1.beta1

data/lib/kafka/fetched_batch.rb

@@ -10,23 +10,31 @@ module Kafka
     # @return [Integer]
     attr_reader :partition
 
+    # @return [Integer]
+    attr_reader :last_offset
+
     # @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:)
+    def initialize(topic:, partition:, highwater_mark_offset:, messages:, last_offset: nil)
       @topic = topic
       @partition = partition
       @highwater_mark_offset = highwater_mark_offset
       @messages = messages
+      @last_offset = last_offset
     end
 
     def empty?
       @messages.empty?
     end
 
+    def unknown_last_offset?
+      @last_offset.nil?
+    end
+
     def first_offset
       if empty?
         nil
@@ -35,14 +43,6 @@ module Kafka
       end
     end
 
-    def last_offset
-      if empty?
-        highwater_mark_offset - 1
-      else
-        messages.last.offset
-      end
-    end
-
     def offset_lag
       if empty?
         0

data/lib/kafka/fetched_batch_generator.rb

@@ -0,0 +1,114 @@
+# 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, logger:)
+      @topic = topic
+      @fetched_partition = fetched_partition
+      @logger = logger
+    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
+          FetchedMessage.new(
+            message: message,
+            topic: @topic,
+            partition: @fetched_partition.partition
+          )
+        end
+      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
+      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
+        # 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|
+          unless record.is_control_record
+            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,
+        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_offset_resolver.rb

@@ -0,0 +1,48 @@
+# frozen_string_literal: true
+
+module Kafka
+  class FetchedOffsetResolver
+    def initialize(logger:)
+      @logger = 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

@@ -144,9 +144,9 @@ module Kafka
             highwater_mark_offset: batch.highwater_mark_offset,
             message_count: batch.messages.count,
           })
-
-          @next_offsets[batch.topic][batch.partition] = batch.last_offset + 1
         end
+
+        @next_offsets[batch.topic][batch.partition] = batch.last_offset + 1 unless batch.unknown_last_offset?
       end
 
       @queue << [:batches, batches]

data/lib/kafka/produce_operation.rb

@@ -30,8 +30,9 @@ module Kafka
   # * `:sent_message_count` – the number of messages that were successfully sent.
   #
   class ProduceOperation
-    def initialize(cluster:, buffer:, compressor:, required_acks:, ack_timeout:, logger:, instrumenter:)
+    def initialize(cluster:, transaction_manager:, buffer:, compressor:, required_acks:, ack_timeout:, logger:, instrumenter:)
       @cluster = cluster
+      @transaction_manager = transaction_manager
       @buffer = buffer
       @required_acks = required_acks
       @ack_timeout = ack_timeout
@@ -41,12 +42,23 @@ module Kafka
     end
 
     def execute
+      if (@transaction_manager.idempotent? || @transaction_manager.transactional?) && @required_acks != -1
+        raise 'You must set required_acks option to :all to use idempotent / transactional production'
+      end
+
+      if @transaction_manager.transactional? && !@transaction_manager.in_transaction?
+        raise "Produce operation can only be executed in a pending transaction"
+      end
+
       @instrumenter.instrument("send_messages.producer") do |notification|
         message_count = @buffer.size
 
         notification[:message_count] = message_count
 
         begin
+          if @transaction_manager.idempotent? || @transaction_manager.transactional?
+            @transaction_manager.init_producer_id
+          end
           send_buffered_messages
         ensure
           notification[:sent_message_count] = message_count - @buffer.size
@@ -58,6 +70,7 @@ module Kafka
 
     def send_buffered_messages
       messages_for_broker = {}
+      topic_partitions = {}
 
       @buffer.each do |topic, partition, messages|
         begin
@@ -65,6 +78,9 @@ module Kafka
 
           @logger.debug "Current leader for #{topic}/#{partition} is node #{broker}"
 
+          topic_partitions[topic] ||= Set.new
+          topic_partitions[topic].add(partition)
+
           messages_for_broker[broker] ||= MessageBuffer.new
           messages_for_broker[broker].concat(messages, topic: topic, partition: partition)
         rescue Kafka::Error => e
@@ -81,25 +97,40 @@ module Kafka
         end
       end
 
+      # Add topic and partition to transaction
+      if @transaction_manager.transactional?
+        @transaction_manager.add_partitions_to_transaction(topic_partitions)
+      end
+
       messages_for_broker.each do |broker, message_buffer|
         begin
           @logger.info "Sending #{message_buffer.size} messages to #{broker}"
 
-          messages_for_topics = {}
+          records_for_topics = {}
 
           message_buffer.each do |topic, partition, records|
-            record_batch = Protocol::RecordBatch.new(records: records)
-            messages_for_topics[topic] ||= {}
-            messages_for_topics[topic][partition] = record_batch
+            record_batch = Protocol::RecordBatch.new(
+              records: records,
+              first_sequence: @transaction_manager.next_sequence_for(
+                topic, partition
+              ),
+              in_transaction: @transaction_manager.transactional?,
+              producer_id: @transaction_manager.producer_id,
+              producer_epoch: @transaction_manager.producer_epoch
+            )
+            records_for_topics[topic] ||= {}
+            records_for_topics[topic][partition] = record_batch
           end
 
           response = broker.produce(
-            messages_for_topics: messages_for_topics,
+            messages_for_topics: records_for_topics,
+            compressor: @compressor,
             required_acks: @required_acks,
             timeout: @ack_timeout * 1000, # Kafka expects the timeout in milliseconds.
+            transactional_id: @transaction_manager.transactional_id
           )
 
-          handle_response(broker, response) if response
+          handle_response(broker, response, records_for_topics) if response
         rescue ConnectionError => e
           @logger.error "Could not connect to broker #{broker}: #{e}"
 
@@ -109,11 +140,12 @@ module Kafka
       end
     end
 
-    def handle_response(broker, response)
+    def handle_response(broker, response, records_for_topics)
       response.each_partition do |topic_info, partition_info|
         topic = topic_info.topic
         partition = partition_info.partition
-        messages = @buffer.messages_for(topic: topic, partition: partition)
+        record_batch = records_for_topics[topic][partition]
+        records = record_batch.records
         ack_time = Time.now
 
         begin
@@ -128,14 +160,20 @@ module Kafka
             raise e
           end
 
-          messages.each_with_index do |message, index|
+          if @transaction_manager.idempotent? || @transaction_manager.transactional?
+            @transaction_manager.update_sequence_for(
+              topic, partition, record_batch.first_sequence + record_batch.size
+            )
+          end
+
+          records.each_with_index do |record, index|
             @instrumenter.instrument("ack_message.producer", {
-              key: message.key,
-              value: message.value,
+              key: record.key,
+              value: record.value,
               topic: topic,
               partition: partition,
               offset: partition_info.offset + index,
-              delay: ack_time - message.create_time,
+              delay: ack_time - record.create_time,
             })
           end
         rescue Kafka::CorruptMessage
@@ -156,7 +194,7 @@ module Kafka
         rescue Kafka::NotEnoughReplicasAfterAppend
           @logger.error "Messages written, but to fewer in-sync replicas than required for #{topic}/#{partition}"
         else
-          @logger.debug "Successfully appended #{messages.count} messages to #{topic}/#{partition} on #{broker}"
+          @logger.debug "Successfully appended #{records.count} messages to #{topic}/#{partition} on #{broker}"
 
           # The messages were successfully written; clear them from the buffer.
           @buffer.clear_messages(topic: topic, partition: partition)

data/lib/kafka/producer.rb

@@ -9,7 +9,6 @@ require "kafka/pending_message"
 require "kafka/compressor"
 
 module Kafka
-
   # Allows sending messages to a Kafka cluster.
   #
   # Typically you won't instantiate this class yourself, but rather have {Kafka::Client}
@@ -126,9 +125,11 @@ module Kafka
   #     end
   #
   class Producer
+    class AbortTransaction < StandardError; end
 
-    def initialize(cluster:, logger:, instrumenter:, compressor:, ack_timeout:, required_acks:, max_retries:, retry_backoff:, max_buffer_size:, max_buffer_bytesize:)
+    def initialize(cluster:, transaction_manager:, logger:, instrumenter:, compressor:, ack_timeout:, required_acks:, max_retries:, retry_backoff:, max_buffer_size:, max_buffer_bytesize:)
       @cluster = cluster
+      @transaction_manager = transaction_manager
       @logger = logger
       @instrumenter = instrumenter
       @required_acks = required_acks == :all ? -1 : required_acks
@@ -201,6 +202,12 @@ module Kafka
           "Cannot produce to #{topic}, max buffer bytesize (#{@max_buffer_bytesize} bytes) reached"
       end
 
+      # If the producer is in transactional mode, all the message production
+      # must be used when the producer is currently in transaction
+      if @transaction_manager.transactional? && !@transaction_manager.in_transaction?
+        raise 'You must trigger begin_transaction before producing messages'
+      end
+
       @target_topics.add(topic)
       @pending_message_queue.write(message)
 
@@ -267,9 +274,81 @@ module Kafka
     #
     # @return [nil]
     def shutdown
+      @transaction_manager.close
       @cluster.disconnect
     end
 
+    # Initializes the producer to ready for future transactions. This method
+    # should be triggered once, before any tranactions are created.
+    #
+    # @return [nil]
+    def init_transactions
+      @transaction_manager.init_transactions
+    end
+
+    # Mark the beginning of a transaction. This method transitions the state
+    # of the transaction trantiions to IN_TRANSACTION.
+    #
+    # All producing operations can only be executed while the transation is
+    # in this state. The records are persisted by Kafka brokers, but not visible
+    # the consumers until the #commit_transaction method is trigger. After a
+    # timeout period without committed, the transaction is timeout and
+    # considered as aborted.
+    #
+    # @return [nil]
+    def begin_transaction
+      @transaction_manager.begin_transaction
+    end
+
+    # This method commits the pending transaction, marks all the produced
+    # records committed. After that, they are visible to the consumers.
+    #
+    # This method can only be called if and only if the current transaction
+    # is at IN_TRANSACTION state.
+    #
+    # @return [nil]
+    def commit_transaction
+      @transaction_manager.commit_transaction
+    end
+
+    # This method abort the pending transaction, marks all the produced
+    # records aborted. All the records will be wiped out by the brokers and the
+    # cosumers don't have a chance to consume those messages, except they enable
+    # consuming uncommitted option.
+    #
+    # This method can only be called if and only if the current transaction
+    # is at IN_TRANSACTION state.
+    #
+    # @return [nil]
+    def abort_transaction
+      @transaction_manager.abort_transaction
+    end
+
+    # Syntactic sugar to enable easier transaction usage. Do the following steps
+    #
+    # - Start the transaction (with Producer#begin_transaction)
+    # - Yield the given block
+    # - Commit the transaction (with Producer#commit_transaction)
+    #
+    # If the block raises exception, the transaction is automatically aborted
+    # *before* bubble up the exception.
+    #
+    # If the block raises Kafka::Producer::AbortTransaction indicator exception,
+    # it aborts the transaction silently, without throwing up that exception.
+    #
+    # @return [nil]
+    def transaction
+      raise 'This method requires a block' unless block_given?
+      begin_transaction
+      yield
+      commit_transaction
+    rescue Kafka::Producer::AbortTransaction
+      abort_transaction
+    rescue
+      abort_transaction
+      raise
+    end
+
     private
 
     def deliver_messages_with_retries(notification)
@@ -279,6 +358,7 @@ module Kafka
 
       operation = ProduceOperation.new(
         cluster: @cluster,
+        transaction_manager: @transaction_manager,
         buffer: @buffer,
         required_acks: @required_acks,
         ack_timeout: @ack_timeout,