ruby-kafka 0.1.3 → 0.1.4

data/lib/kafka/producer.rb

@@ -1,7 +1,7 @@
 require "kafka/partitioner"
 require "kafka/message_buffer"
-require "kafka/protocol/message"
-require "kafka/transmission"
+require "kafka/produce_operation"
+require "kafka/pending_message"
 
 module Kafka
 
@@ -17,7 +17,7 @@ module Kafka
   #     producer = kafka.get_producer
   #
   # This is done in order to share a logger as well as a pool of broker connections across
-  # different producers. This also means that you don't need to pass the `broker_pool` and
+  # different producers. This also means that you don't need to pass the `cluster` and
   # `logger` options to `#get_producer`. See {#initialize} for the list of other options
   # you can pass in.
   #
@@ -86,8 +86,7 @@ module Kafka
 
     # Initializes a new Producer.
     #
-    # @param broker_pool [BrokerPool] the broker pool representing the cluster.
-    #   Typically passed in for you.
+    # @param cluster [Cluster] the cluster client. Typically passed in for you.
     #
     # @param logger [Logger] the logger that should be used. Typically passed
     #   in for you.
@@ -107,15 +106,20 @@ module Kafka
     # @param max_buffer_size [Integer] the number of messages allowed in the buffer
     #   before new writes will raise BufferOverflow exceptions.
     #
-    def initialize(broker_pool:, logger:, ack_timeout: 5, required_acks: 1, max_retries: 2, retry_backoff: 1, max_buffer_size: 1000)
-      @broker_pool = broker_pool
+    def initialize(cluster:, logger:, ack_timeout: 5, required_acks: 1, max_retries: 2, retry_backoff: 1, max_buffer_size: 1000)
+      @cluster = cluster
       @logger = logger
       @required_acks = required_acks
       @ack_timeout = ack_timeout
       @max_retries = max_retries
       @retry_backoff = retry_backoff
       @max_buffer_size = max_buffer_size
+
+      # A buffer organized by topic/partition.
       @buffer = MessageBuffer.new
+
+      # Messages added by `#produce` but not yet assigned a partition.
+      @pending_messages = []
     end
 
     # Produces a message to the specified topic. Note that messages are buffered in
@@ -152,18 +156,15 @@ module Kafka
         raise BufferOverflow, "Max buffer size #{@max_buffer_size} exceeded"
       end
 
-      if partition.nil?
-        # If no explicit partition key is specified we use the message key instead.
-        partition_key ||= key
-        partitioner = Partitioner.new(@broker_pool.partitions_for(topic))
-        partition = partitioner.partition_for_key(partition_key)
-      end
-
-      message = Protocol::Message.new(key: key, value: value)
-
-      @buffer.write(message, topic: topic, partition: partition)
+      @pending_messages << PendingMessage.new(
+        value: value,
+        key: key,
+        topic: topic,
+        partition: partition,
+        partition_key: partition_key,
+      )
 
-      partition
+      nil
     end
 
     # Sends all buffered messages to the Kafka brokers.
@@ -178,8 +179,12 @@ module Kafka
     def send_messages
       attempt = 0
 
-      transmission = Transmission.new(
-        broker_pool: @broker_pool,
+      # Make sure we get metadata for this topic.
+      target_topics = @pending_messages.map(&:topic).uniq
+      @cluster.add_target_topics(target_topics)
+
+      operation = ProduceOperation.new(
+        cluster: @cluster,
         buffer: @buffer,
         required_acks: @required_acks,
         ack_timeout: @ack_timeout,
@@ -187,21 +192,21 @@ module Kafka
       )
 
       loop do
-        @logger.info "Sending #{@buffer.size} messages"
-
         attempt += 1
-        transmission.send_messages
 
-        if @buffer.empty?
-          @logger.info "Successfully transmitted all messages"
+        @cluster.refresh_metadata_if_necessary!
+
+        assign_partitions!
+        operation.execute
+
+        if @pending_messages.empty? && @buffer.empty?
           break
         elsif attempt <= @max_retries
-          @logger.warn "Failed to transmit all messages, retry #{attempt} of #{@max_retries}"
-          @logger.info "Waiting #{@retry_backoff}s before retrying"
+          @logger.warn "Failed to send all messages; attempting retry #{attempt} of #{@max_retries} after #{@retry_backoff}s"
 
           sleep @retry_backoff
         else
-          @logger.error "Failed to transmit all messages; keeping remaining messages in buffer"
+          @logger.error "Failed to send all messages; keeping remaining messages in buffer"
           break
         end
       end
@@ -224,14 +229,43 @@ module Kafka
     #
     # @return [Integer] buffer size.
     def buffer_size
-      @buffer.size
+      @pending_messages.size + @buffer.size
     end
 
     # Closes all connections to the brokers.
     #
     # @return [nil]
     def shutdown
-      @broker_pool.shutdown
+      @cluster.disconnect
+    end
+
+    private
+
+    def assign_partitions!
+      until @pending_messages.empty?
+        # We want to keep the message in the first-stage buffer in case there's an error.
+        message = @pending_messages.first
+
+        partition = message.partition
+
+        if partition.nil?
+          partition_count = @cluster.partitions_for(message.topic).count
+          partition = Partitioner.partition_for_key(partition_count, message)
+        end
+
+        @buffer.write(
+          value: message.value,
+          key: message.key,
+          topic: message.topic,
+          partition: partition,
+        )
+
+        # Now it's safe to remove the message from the first-stage buffer.
+        @pending_messages.shift
+      end
+    rescue Kafka::Error => e
+      @logger.error "Failed to assign pending message to a partition: #{e}"
+      @cluster.mark_as_stale!
     end
   end
 end

data/lib/kafka/protocol.rb

@@ -1,7 +1,12 @@
 module Kafka
   module Protocol
+    # The replica id of non-brokers is always -1.
+    REPLICA_ID = -1
+
     APIS = {
       0 => :produce,
+      1 => :fetch,
+      2 => :list_offset,
       3 => :topic_metadata,
     }
 
@@ -45,3 +50,7 @@ require "kafka/protocol/topic_metadata_request"
 require "kafka/protocol/metadata_response"
 require "kafka/protocol/produce_request"
 require "kafka/protocol/produce_response"
+require "kafka/protocol/fetch_request"
+require "kafka/protocol/fetch_response"
+require "kafka/protocol/list_offset_request"
+require "kafka/protocol/list_offset_response"

data/lib/kafka/protocol/decoder.rb

@@ -5,6 +5,9 @@ module Kafka
     # from it. The Kafka protocol is not self-describing, so a client must call
     # these methods in just the right order for things to work.
     class Decoder
+      def self.from_string(str)
+        new(StringIO.new(str))
+      end
 
       # Initializes a new decoder.
       #
@@ -13,6 +16,10 @@ module Kafka
         @io = io
       end
 
+      def eof?
+        @io.eof?
+      end
+
       # Decodes an 8-bit integer from the IO object.
       #
       # @return [Integer]

data/lib/kafka/protocol/fetch_request.rb

@@ -0,0 +1,53 @@
+module Kafka
+  module Protocol
+
+    # A request to fetch messages from a given partition.
+    #
+    # ## API Specification
+    #
+    #     FetchRequest => ReplicaId MaxWaitTime MinBytes [TopicName [Partition FetchOffset MaxBytes]]
+    #       ReplicaId => int32
+    #       MaxWaitTime => int32
+    #       MinBytes => int32
+    #       TopicName => string
+    #       Partition => int32
+    #       FetchOffset => int64
+    #       MaxBytes => int32
+    #
+    class FetchRequest
+
+      # @param max_wait_time [Integer]
+      # @param min_bytes [Integer]
+      # @param topics [Hash]
+      def initialize(max_wait_time:, min_bytes:, topics:)
+        @replica_id = REPLICA_ID
+        @max_wait_time = max_wait_time
+        @min_bytes = min_bytes
+        @topics = topics
+      end
+
+      def api_key
+        1
+      end
+
+      def encode(encoder)
+        encoder.write_int32(@replica_id)
+        encoder.write_int32(@max_wait_time)
+        encoder.write_int32(@min_bytes)
+
+        encoder.write_array(@topics) do |topic, partitions|
+          encoder.write_string(topic)
+
+          encoder.write_array(partitions) do |partition, config|
+            fetch_offset = config.fetch(:fetch_offset)
+            max_bytes = config.fetch(:max_bytes)
+
+            encoder.write_int32(partition)
+            encoder.write_int64(fetch_offset)
+            encoder.write_int32(max_bytes)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kafka/protocol/fetch_response.rb

@@ -0,0 +1,75 @@
+require "kafka/protocol/message_set"
+
+module Kafka
+  module Protocol
+
+    # A response to a fetch request.
+    #
+    # ## API Specification
+    #
+    #     FetchResponse => [TopicName [Partition ErrorCode HighwaterMarkOffset MessageSetSize MessageSet]]
+    #       TopicName => string
+    #       Partition => int32
+    #       ErrorCode => int16
+    #       HighwaterMarkOffset => int64
+    #       MessageSetSize => int32
+    #
+    class FetchResponse
+      class FetchedPartition
+        attr_reader :partition, :error_code
+        attr_reader :highwater_mark_offset, :messages
+
+        def initialize(partition:, error_code:, highwater_mark_offset:, messages:)
+          @partition = partition
+          @error_code = error_code
+          @highwater_mark_offset = highwater_mark_offset
+          @messages = messages
+        end
+      end
+
+      class FetchedTopic
+        attr_reader :name, :partitions
+
+        def initialize(name:, partitions:)
+          @name = name
+          @partitions = partitions
+        end
+      end
+
+      attr_reader :topics
+
+      def initialize(topics: [])
+        @topics = topics
+      end
+
+      def self.decode(decoder)
+        topics = decoder.array do
+          topic_name = decoder.string
+
+          partitions = decoder.array do
+            partition = decoder.int32
+            error_code = decoder.int16
+            highwater_mark_offset = decoder.int64
+
+            message_set_decoder = Decoder.from_string(decoder.bytes)
+            message_set = MessageSet.decode(message_set_decoder)
+
+            FetchedPartition.new(
+              partition: partition,
+              error_code: error_code,
+              highwater_mark_offset: highwater_mark_offset,
+              messages: message_set.messages,
+            )
+          end
+
+          FetchedTopic.new(
+            name: topic_name,
+            partitions: partitions,
+          )
+        end
+
+        new(topics: topics)
+      end
+    end
+  end
+end

data/lib/kafka/protocol/list_offset_request.rb

@@ -0,0 +1,41 @@
+module Kafka
+  module Protocol
+    # A request to list the available offsets for a set of topics/partitions.
+    #
+    # ## API Specification
+    #
+    #     OffsetRequest => ReplicaId [TopicName [Partition Time MaxNumberOfOffsets]]
+    #       ReplicaId => int32
+    #       TopicName => string
+    #       Partition => int32
+    #       Time => int64
+    #       MaxNumberOfOffsets => int32
+    #
+    class ListOffsetRequest
+
+      # @param topics [Hash]
+      def initialize(topics:)
+        @replica_id = REPLICA_ID
+        @topics = topics
+      end
+
+      def api_key
+        2
+      end
+
+      def encode(encoder)
+        encoder.write_int32(@replica_id)
+
+        encoder.write_array(@topics) do |topic, partitions|
+          encoder.write_string(topic)
+
+          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
+    end
+  end
+end

data/lib/kafka/protocol/list_offset_response.rb

@@ -0,0 +1,82 @@
+module Kafka
+  module Protocol
+
+    # A response to a list offset request.
+    #
+    # ## API Specification
+    #
+    #     OffsetResponse => [TopicName [PartitionOffsets]]
+    #       PartitionOffsets => Partition ErrorCode [Offset]
+    #       Partition => int32
+    #       ErrorCode => int16
+    #       Offset => int64
+    #
+    class ListOffsetResponse
+      class TopicOffsetInfo
+        attr_reader :name, :partition_offsets
+
+        def initialize(name:, partition_offsets:)
+          @name = name
+          @partition_offsets = partition_offsets
+        end
+      end
+
+      class PartitionOffsetInfo
+        attr_reader :partition, :error_code, :offsets
+
+        def initialize(partition:, error_code:, offsets:)
+          @partition = partition
+          @error_code = error_code
+          @offsets = offsets
+        end
+      end
+
+      attr_reader :topics
+
+      def initialize(topics:)
+        @topics = topics
+      end
+
+      def offset_for(topic, partition)
+        topic_info = @topics.find {|t| t.name == topic }
+
+        if topic_info.nil?
+          raise UnknownTopicOrPartition, "Unknown topic #{topic}"
+        end
+
+        partition_info = topic_info
+          .partition_offsets
+          .find {|p| p.partition == partition }
+
+        if partition_info.nil?
+          raise UnknownTopicOrPartition, "Unknown partition #{topic}/#{partition}"
+        end
+
+        Protocol.handle_error(partition_info.error_code)
+
+        partition_info.offsets.first
+      end
+
+      def self.decode(decoder)
+        topics = decoder.array do
+          name = decoder.string
+
+          partition_offsets = decoder.array do
+            PartitionOffsetInfo.new(
+              partition: decoder.int32,
+              error_code: decoder.int16,
+              offsets: decoder.array { decoder.int64 },
+            )
+          end
+
+          TopicOffsetInfo.new(
+            name: name,
+            partition_offsets: partition_offsets
+          )
+        end
+
+        new(topics: topics)
+      end
+    end
+  end
+end