ruby-kafka 0.1.3 → 0.1.4

data/lib/kafka/cluster.rb

@@ -0,0 +1,130 @@
+require "set"
+require "kafka/broker_pool"
+
+module Kafka
+
+  # A cluster represents the state of a Kafka cluster. It needs to be initialized
+  # with a non-empty list of seed brokers. The first seed broker that the cluster can connect
+  # to will be asked for the cluster metadata, which allows the cluster to map topic
+  # partitions to the current leader for those partitions.
+  class Cluster
+
+    # Initializes a Cluster with a set of seed brokers.
+    #
+    # The cluster will try to fetch cluster metadata from one of the brokers.
+    #
+    # @param seed_brokers [Array<String>]
+    # @param broker_pool [Kafka::BrokerPool]
+    # @param logger [Logger]
+    def initialize(seed_brokers:, broker_pool:, logger:)
+      if seed_brokers.empty?
+        raise ArgumentError, "At least one seed broker must be configured"
+      end
+
+      @logger = logger
+      @seed_brokers = seed_brokers
+      @broker_pool = broker_pool
+      @cluster_info = nil
+      @stale = true
+
+      # This is the set of topics we need metadata for. If empty, metadata for
+      # all topics will be fetched.
+      @target_topics = Set.new
+    end
+
+    def add_target_topics(topics)
+      new_topics = Set.new(topics) - @target_topics
+
+      unless new_topics.empty?
+        @logger.info "New topics added to target list: #{new_topics.to_a.join(', ')}"
+
+        @target_topics.merge(new_topics)
+
+        refresh_metadata!
+      end
+    end
+
+    def mark_as_stale!
+      @stale = true
+    end
+
+    def refresh_metadata!
+      @cluster_info = nil
+      cluster_info
+    end
+
+    def refresh_metadata_if_necessary!
+      refresh_metadata! if @stale
+    end
+
+    # Finds the broker acting as the leader of the given topic and partition.
+    #
+    # @param topic [String]
+    # @param partition [Integer]
+    # @return [Broker] the broker that's currently leader.
+    def get_leader(topic, partition)
+      connect_to_broker(get_leader_id(topic, partition))
+    end
+
+    def partitions_for(topic)
+      cluster_info.partitions_for(topic)
+    end
+
+    def topics
+      cluster_info.topics.map(&:topic_name)
+    end
+
+    def disconnect
+      @broker_pool.close
+    end
+
+    private
+
+    def get_leader_id(topic, partition)
+      cluster_info.find_leader_id(topic, partition)
+    end
+
+    def cluster_info
+      @cluster_info ||= fetch_cluster_info
+    end
+
+    # Fetches the cluster metadata.
+    #
+    # This is used to update the partition leadership information, among other things.
+    # The methods will go through each node listed in `seed_brokers`, connecting to the
+    # first one that is available. This node will be queried for the cluster metadata.
+    #
+    # @raise [ConnectionError] if none of the nodes in `seed_brokers` are available.
+    # @return [Protocol::MetadataResponse] the cluster metadata.
+    def fetch_cluster_info
+      @seed_brokers.each do |node|
+        @logger.info "Fetching cluster metadata from #{node}"
+
+        begin
+          host, port = node.split(":", 2)
+
+          broker = @broker_pool.connect(host, port.to_i)
+          cluster_info = broker.fetch_metadata(topics: @target_topics)
+
+          @stale = false
+
+          @logger.info "Discovered cluster metadata; nodes: #{cluster_info.brokers.join(', ')}"
+
+          return cluster_info
+        rescue Error => e
+          @logger.error "Failed to fetch metadata from #{node}: #{e}"
+        ensure
+          broker.disconnect unless broker.nil?
+        end
+      end
+
+      raise ConnectionError, "Could not connect to any of the seed brokers: #{@seed_brokers.join(', ')}"
+    end
+
+    def connect_to_broker(broker_id)
+      info = cluster_info.find_broker(broker_id)
+
+      @broker_pool.connect(info.host, info.port, node_id: info.node_id)
+    end
+  end
+end

data/lib/kafka/connection.rb

@@ -88,8 +88,6 @@ module Kafka
         wait_for_response(response_class, notification) unless response_class.nil?
       end
     rescue Errno::EPIPE, Errno::ECONNRESET, Errno::ETIMEDOUT, EOFError => e
-      @logger.error "Connection error: #{e}"
-
       close
 
       raise ConnectionError, "Connection error: #{e}"
@@ -98,7 +96,7 @@ module Kafka
     private
 
     def open
-      @logger.info "Opening connection to #{@host}:#{@port} with client id #{@client_id}..."
+      @logger.debug "Opening connection to #{@host}:#{@port} with client id #{@client_id}..."
 
       @socket = SocketWithTimeout.new(@host, @port, connect_timeout: @connect_timeout, timeout: @socket_timeout)
 

data/lib/kafka/fetch_operation.rb

@@ -0,0 +1,127 @@
+module Kafka
+
+  # Fetches messages from one or more partitions.
+  #
+  #     operation = Kafka::FetchOperation.new(
+  #       cluster: cluster,
+  #       logger: logger,
+  #       min_bytes: 1,
+  #       max_wait_time: 10,
+  #     )
+  #
+  #     # These calls will schedule fetches from the specified topics/partitions.
+  #     operation.fetch_from_partition("greetings", 42, offset: :latest, max_bytes: 100000)
+  #     operation.fetch_from_partition("goodbyes", 13, offset: :latest, max_bytes: 100000)
+  #
+  #     operation.execute
+  #
+  class FetchOperation
+    def initialize(cluster:, logger:, min_bytes:, max_wait_time:)
+      @cluster = cluster
+      @logger = logger
+      @min_bytes = min_bytes
+      @max_wait_time = max_wait_time
+      @topics = {}
+    end
+
+    def fetch_from_partition(topic, partition, offset:, max_bytes:)
+      if offset == :earliest
+        offset = -2
+      elsif offset == :latest
+        offset = -1
+      end
+
+      @topics[topic] ||= {}
+      @topics[topic][partition] = {
+        fetch_offset: offset,
+        max_bytes: max_bytes,
+      }
+    end
+
+    def execute
+      @cluster.add_target_topics(@topics.keys)
+      @cluster.refresh_metadata_if_necessary!
+
+      topics_by_broker = {}
+
+      @topics.each do |topic, partitions|
+        partitions.each do |partition, options|
+          broker = @cluster.get_leader(topic, partition)
+
+          topics_by_broker[broker] ||= {}
+          topics_by_broker[broker][topic] ||= {}
+          topics_by_broker[broker][topic][partition] = options
+        end
+      end
+
+      topics_by_broker.flat_map {|broker, topics|
+        resolve_offsets(broker, topics)
+
+        options = {
+          max_wait_time: @max_wait_time * 1000, # Kafka expects ms, not secs
+          min_bytes: @min_bytes,
+          topics: topics,
+        }
+
+        response = broker.fetch_messages(**options)
+
+        response.topics.flat_map {|fetched_topic|
+          fetched_topic.partitions.flat_map {|fetched_partition|
+            Protocol.handle_error(fetched_partition.error_code)
+
+            fetched_partition.messages.map {|offset, message|
+              FetchedMessage.new(
+                value: message.value,
+                key: message.key,
+                topic: fetched_topic.name,
+                partition: fetched_partition.partition,
+                offset: offset,
+              )
+            }
+          }
+        }
+      }
+    rescue Kafka::LeaderNotAvailable, Kafka::NotLeaderForPartition
+      @cluster.mark_as_stale!
+
+      raise
+    end
+
+    private
+
+    def resolve_offsets(broker, 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,
+            max_offsets: 1,
+          }
+        end
+      end
+
+      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
+  end
+end

data/lib/kafka/fetched_message.rb

@@ -0,0 +1,27 @@
+module Kafka
+  class FetchedMessage
+
+    # @return [String] the value of the message.
+    attr_reader :value
+
+    # @return [String] the key of the message.
+    attr_reader :key
+
+    # @return [String] the name of the topic that the message was written to.
+    attr_reader :topic
+
+    # @return [Integer] the partition number that the message was written to.
+    attr_reader :partition
+
+    # @return [Integer] the offset of the message in the partition.
+    attr_reader :offset
+
+    def initialize(value:, key:, topic:, partition:, offset:)
+      @value = value
+      @key = key
+      @topic = topic
+      @partition = partition
+      @offset = offset
+    end
+  end
+end

data/lib/kafka/instrumentation.rb

@@ -6,7 +6,7 @@ module Kafka
   end
 
   if defined?(ActiveSupport::Notifications)
-    Instrumentation = ActiveSupport::Notifications 
+    Instrumentation = ActiveSupport::Notifications
   else
     Instrumentation = NullInstrumentation
   end

data/lib/kafka/message_buffer.rb

@@ -1,3 +1,5 @@
+require "kafka/protocol/message"
+
 module Kafka
 
   # Buffers messages for specific topics/partitions.
@@ -11,8 +13,9 @@ module Kafka
       @size = 0
     end
 
-    def write(message, topic:, partition:)
+    def write(value:, key:, topic:, partition:)
       @size += 1
+      message = Protocol::Message.new(key: key, value: value)
       buffer_for(topic, partition) << message
     end
 
@@ -50,6 +53,10 @@ module Kafka
       @buffer.delete(topic) if @buffer[topic].empty?
     end
 
+    def message_count_for_partition(topic:, partition:)
+      buffer_for(topic, partition).count
+    end
+
     # Clears messages across all topics and partitions.
     #
     # @return [nil]

data/lib/kafka/partitioner.rb

@@ -4,24 +4,29 @@ module Kafka
 
   # Assigns partitions to messages.
   class Partitioner
-    def initialize(partitions)
-      @partitions = partitions
-    end
 
-    # Assigns a partition number based on a key.
+    # 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 key [String, nil] the key to base the partition assignment on, or nil.
+    # @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 partition_for_key(key)
+    def self.partition_for_key(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(@partitions.count)
+        rand(partition_count)
       else
-        Zlib.crc32(key) % @partitions.count
+        Zlib.crc32(key) % partition_count
       end
     end
   end

data/lib/kafka/pending_message.rb

@@ -0,0 +1,13 @@
+module Kafka
+  class PendingMessage
+    attr_reader :value, :key, :topic, :partition, :partition_key
+
+    def initialize(value:, key:, topic:, partition:, partition_key:)
+      @key = key
+      @value = value
+      @topic = topic
+      @partition = partition
+      @partition_key = partition_key
+    end
+  end
+end

data/lib/kafka/produce_operation.rb

@@ -0,0 +1,116 @@
+module Kafka
+  # A produce operation attempts to send all messages in a buffer to the Kafka cluster.
+  # Since topics and partitions are spread among all brokers in a cluster, this usually
+  # involves sending requests to several or all of the brokers.
+  #
+  # ## Instrumentation
+  #
+  # When executing the operation, an `append_message_set.kafka` notification will be
+  # emitted for each message set that was successfully appended to a topic partition.
+  # The following keys will be found in the payload:
+  #
+  # * `:topic` — the topic that was written to.
+  # * `:partition` — the partition that the message set was appended to.
+  # * `:offset` — the offset of the first message in the message set.
+  # * `:message_count` — the number of messages that were appended.
+  #
+  # If there was an error appending the message set, the key `:exception` will be set
+  # in the payload. In that case, the message set will most likely not have been
+  # appended and will possibly be retried later. Check this key before reporting the
+  # operation as successful.
+  #
+  class ProduceOperation
+    def initialize(cluster:, buffer:, required_acks:, ack_timeout:, logger:)
+      @cluster = cluster
+      @buffer = buffer
+      @required_acks = required_acks
+      @ack_timeout = ack_timeout
+      @logger = logger
+    end
+
+    def execute
+      messages_for_broker = {}
+
+      @buffer.each do |topic, partition, messages|
+        begin
+          broker = @cluster.get_leader(topic, partition)
+
+          @logger.debug "Current leader for #{topic}/#{partition} is node #{broker}"
+
+          messages_for_broker[broker] ||= MessageBuffer.new
+          messages_for_broker[broker].concat(messages, topic: topic, partition: partition)
+        rescue Kafka::Error => e
+          @logger.error "Could not connect to leader for partition #{topic}/#{partition}: #{e}"
+
+          # We can't send the messages right now, so we'll just keep them in the buffer.
+          # We'll mark the cluster as stale in order to force a metadata refresh.
+          @cluster.mark_as_stale!
+        end
+      end
+
+      messages_for_broker.each do |broker, message_set|
+        begin
+          @logger.info "Sending #{message_set.size} messages to #{broker}"
+
+          response = broker.produce(
+            messages_for_topics: message_set.to_h,
+            required_acks: @required_acks,
+            timeout: @ack_timeout * 1000, # Kafka expects the timeout in milliseconds.
+          )
+
+          handle_response(response) if response
+        rescue ConnectionError => e
+          @logger.error "Could not connect to broker #{broker}: #{e}"
+
+          # Mark the cluster as stale in order to force a cluster metadata refresh.
+          @cluster.mark_as_stale!
+        end
+      end
+    end
+
+    private
+
+    def handle_response(response)
+      response.each_partition do |topic_info, partition_info|
+        topic = topic_info.topic
+        partition = partition_info.partition
+        offset = partition_info.offset
+        message_count = @buffer.message_count_for_partition(topic: topic, partition: partition)
+
+        begin
+          payload = {
+            topic: topic,
+            partition: partition,
+            offset: offset,
+            message_count: message_count,
+          }
+
+          Instrumentation.instrument("append_message_set.kafka", payload) do
+            Protocol.handle_error(partition_info.error_code)
+          end
+        rescue Kafka::CorruptMessage
+          @logger.error "Corrupt message when writing to #{topic}/#{partition}"
+        rescue Kafka::UnknownTopicOrPartition
+          @logger.error "Unknown topic or partition #{topic}/#{partition}"
+        rescue Kafka::LeaderNotAvailable
+          @logger.error "Leader currently not available for #{topic}/#{partition}"
+          @cluster.mark_as_stale!
+        rescue Kafka::NotLeaderForPartition
+          @logger.error "Broker not currently leader for #{topic}/#{partition}"
+          @cluster.mark_as_stale!
+        rescue Kafka::RequestTimedOut
+          @logger.error "Timed out while writing to #{topic}/#{partition}"
+        rescue Kafka::NotEnoughReplicas
+          @logger.error "Not enough in-sync replicas for #{topic}/#{partition}"
+        rescue Kafka::NotEnoughReplicasAfterAppend
+          @logger.error "Messages written, but to fewer in-sync replicas than required for #{topic}/#{partition}"
+        else
+          @logger.debug "Successfully appended #{message_count} messages to #{topic}/#{partition} at offset #{offset}"
+
+          # The messages were successfully written; clear them from the buffer.
+          @buffer.clear_messages(topic: topic, partition: partition)
+        end
+      end
+    end
+  end
+end