codeclimate-poseidon 0.0.8

data/lib/poseidon/message_conductor.rb

@@ -0,0 +1,86 @@
+module Poseidon
+  # @api private
+  class MessageConductor
+    NO_PARTITION  = -1
+    NO_BROKER     = -1
+    # Create a new message conductor
+    #
+    # @param [Hash<String,TopicMetadata>] topics_metadata
+    #   Metadata for all topics this conductor may receive.
+    # @param [Object] partitioner
+    #   Custom partitioner
+    def initialize(cluster_metadata, partitioner)
+      @cluster_metadata   = cluster_metadata
+      @partitioner        = partitioner
+
+      # Don't always start from partition 0
+      @partition_counter  = rand(65536)
+    end
+
+    # Determines which partition a message should be sent to.
+    #
+    # @param [String] topic
+    #   Topic we are sending this message to
+    #
+    # @param [Object] key
+    #   Key for this message, may be nil
+    #
+    # @return [Integer,Integer] 
+    #   partition_id and broker_id to which this message should be sent
+    def destination(topic, key = nil)
+      topic_metadata = topic_metadatas[topic]
+      if topic_metadata && topic_metadata.leader_available?
+        partition_id = determine_partition(topic_metadata, key)
+        broker_id    = topic_metadata.partition_leader(partition_id) || NO_BROKER
+      else
+        partition_id  = NO_PARTITION
+        broker_id     = NO_BROKER
+      end
+
+      return partition_id, broker_id
+    end
+
+    private
+
+    def topic_metadatas
+      @cluster_metadata.topic_metadata
+    end
+
+    def determine_partition(topic_metadata, key)
+      if key
+        partition_for_keyed_message(topic_metadata, key)
+      else
+        partition_for_keyless_message(topic_metadata)
+      end
+    end
+
+    def partition_for_keyed_message(topic_metadata, key)
+      partition_count  = topic_metadata.partition_count
+      if @partitioner
+        partition_id     = @partitioner.call(key, partition_count)
+
+        if partition_id >= partition_count
+          raise Errors::InvalidPartitionError, "partitioner (#{@partitioner.inspect}) requested #{partition_id} while only #{partition_count} partitions exist"
+        end
+      else
+        partition_id = Zlib::crc32(key) % partition_count
+      end
+
+      partition_id
+    end
+
+    def partition_for_keyless_message(topic_metadata)
+      partition_count = topic_metadata.available_partition_count
+
+      if partition_count > 0
+        topic_metadata.available_partitions[next_partition_counter % partition_count].id
+      else
+        NO_PARTITION
+      end
+    end
+
+    def next_partition_counter
+      @partition_counter += 1
+    end
+  end
+end

data/lib/poseidon/message_set.rb

@@ -0,0 +1,80 @@
+module Poseidon
+  # @api private
+  class MessageSet
+    # Build a message set object from a binary encoded message set
+    #
+    # @param [String] string binary encoded message set
+    # @return [MessageSet]
+
+
+    def self.read(buffer)
+      ms = MessageSet.new
+      ms.struct = Protocol::MessageSetStructWithSize.read(buffer)
+      ms
+    end
+
+    def self.read_without_size(buffer)
+      ms = MessageSet.new
+      ms.struct = Protocol::MessageSetStruct.read(buffer)
+      ms
+    end
+
+    attr_accessor :struct
+    def initialize(messages = [])
+      self.struct = Protocol::MessageSetStructWithSize.new(messages)
+    end
+
+    def ==(other)
+      eql?(other)
+    end
+
+    def eql?(other)
+      struct.eql?(other.struct)
+    end
+
+    def objects_with_errors
+      struct.objects_with_errors
+    end
+
+    def write(buffer)
+      struct.write(buffer)
+    end
+
+    def <<(message)
+      struct.messages << message
+    end
+
+    def messages
+      struct.messages
+    end
+
+    def compress(codec)
+      MessageSet.new([to_compressed_message(codec)])
+    end
+
+    # Builds an array of Message objects from the MessageStruct objects. 
+    # Decompressing messages if necessary.
+    #
+    # @return [Array<Message>]
+    def flatten
+      messages = struct.messages.map do |message| 
+        if message.compressed?
+          s = message.decompressed_value
+          MessageSet.read_without_size(Protocol::ResponseBuffer.new(s)).flatten
+        else
+          message
+        end
+      end.flatten
+    end
+
+    private
+    def to_compressed_message(codec)
+      buffer = Protocol::RequestBuffer.new
+      struct.write(buffer)
+
+      value = codec.compress(buffer.to_s[4..-1])
+      Message.new(:value => value, :attributes => codec.codec_id)
+    end
+
+  end
+end

data/lib/poseidon/message_to_send.rb

@@ -0,0 +1,33 @@
+module Poseidon
+  # A message we want to send to Kafka. Comprised of the
+  # topic we want to send it to, the body of the message
+  # and an optional key.
+  #
+  #     mts = Poseidon::MessageToSend.new("topic", "value", "opt_key")
+  #
+  # @api public
+  class MessageToSend
+    attr_reader :value, :key, :topic
+
+    # Create a new message for sending to a Kafka broker.
+    #
+    # @param [String] topic
+    #   Topic this message should be sent to.
+    #
+    # @param [String] value
+    #   Value of the message we want to send. 
+    #
+    # @param [String] key
+    #   Optional. Message's key, used to route a message
+    #   to a specific broker.  Otherwise, messages will be
+    #   sent to brokers in a round-robin manner.
+    #
+    # @api public
+    def initialize(topic, value, key = nil)
+      raise ArgumentError, "Must provide a non-nil topic" if topic.nil?
+      @topic  = topic
+      @value  = value
+      @key    = key
+    end
+  end
+end

data/lib/poseidon/messages_for_broker.rb

@@ -0,0 +1,56 @@
+module Poseidon
+  # Messages that should be sent to a particular broker.
+  # @api private
+  class MessagesForBroker
+    attr_reader :broker_id, :messages
+
+    def initialize(broker_id)
+      @broker_id = broker_id
+      @topics = {}
+      @messages = []
+    end
+
+    # Add a messages for this broker
+    def add(message, partition_id)
+      @messages << message
+
+      @topics[message.topic] ||= {}
+      @topics[message.topic][partition_id] ||= []
+      @topics[message.topic][partition_id] << message
+    end
+
+    # Build protocol objects for this broker!
+    def build_protocol_objects(compression_config)
+      @topics.map do |topic, messages_by_partition|
+        codec = compression_config.compression_codec_for_topic(topic)
+
+        messages_for_partitions = messages_by_partition.map do |partition, messages|
+          message_set = MessageSet.new(messages)
+          if codec
+            Protocol::MessagesForPartition.new(partition, message_set.compress(codec))
+          else
+            Protocol::MessagesForPartition.new(partition, message_set)
+          end
+        end
+        Protocol::MessagesForTopic.new(topic, messages_for_partitions)
+      end
+    end
+
+    # We can always retry these errors because they mean none of the kafka brokers persisted the message
+    ALWAYS_RETRYABLE = [Poseidon::Errors::LeaderNotAvailable, Poseidon::Errors::NotLeaderForPartition]
+
+    def successfully_sent(producer_response)
+      failed = []
+      producer_response.topic_response.each do |topic_response|
+        topic_response.partitions.each do |partition|
+          if ALWAYS_RETRYABLE.include?(partition.error_class)
+            Poseidon.logger.debug { "Received #{partition.error_class} when attempting to send messages to #{topic_response.topic} on #{partition.partition}" }
+            failed.push(*@topics[topic_response.topic][partition.partition])
+          end
+        end
+      end
+
+      return @messages - failed
+    end
+  end
+end

data/lib/poseidon/messages_to_send.rb

@@ -0,0 +1,47 @@
+module Poseidon
+  # A set of messages that we need to send to the cluster. May be used
+  # across multiple send attempts.
+  #
+  # If a custom partitioner is not used than a messages are distributed
+  # in round robin fasion to each partition with an available leader.
+  #
+  # @api private
+  class MessagesToSend
+    class InvalidPartitionError < StandardError; end
+    attr_reader :topic_set, :messages
+
+    # Create a new messages to send object.
+    #
+    # @param [Array<Message>] messages List of messages we want to send.
+    # @param [ClusterMetadta] cluster_metadata
+    def initialize(messages, cluster_metadata)
+      @messages         = messages
+      @cluster_metadata = cluster_metadata
+
+      build_topic_set
+    end
+
+    def needs_metadata?
+      !@cluster_metadata.have_metadata_for_topics?(topic_set)
+    end
+
+    def messages_for_brokers(message_conductor)
+      topic_metadatas = @cluster_metadata.metadata_for_topics(topic_set)
+      MessagesToSendBatch.new(@messages, message_conductor).messages_for_brokers
+    end
+
+    def successfully_sent(messages_sent)
+      @messages -= messages_sent
+    end
+
+    def pending_messages?
+      @messages.any?
+    end
+
+    private
+    def build_topic_set
+      @topic_set = Set.new
+      @messages.each { |m| @topic_set.add(m.topic) }
+    end
+  end
+end

data/lib/poseidon/messages_to_send_batch.rb

@@ -0,0 +1,27 @@
+module Poseidon
+  # A batch of messages for an individual send attempt to the cluster.
+  # @api private
+  class MessagesToSendBatch
+    def initialize(messages, message_conductor)
+      @messages           = messages
+      @message_conductor  = message_conductor
+    end
+
+    # Groups messages by broker and preps them for transmission.
+    #
+    # @return [Array<MessagesForBroker>]
+    def messages_for_brokers
+      messages_for_broker_ids = {}
+      @messages.each do |message|
+        partition_id, broker_id = @message_conductor.destination(message.topic,
+                                                                 message.key)
+
+        # Create a nested hash to group messages by broker_id, topic, partition.
+        messages_for_broker_ids[broker_id] ||= MessagesForBroker.new(broker_id)
+        messages_for_broker_ids[broker_id].add(message, partition_id)
+      end
+
+      messages_for_broker_ids.values
+    end
+  end
+end

data/lib/poseidon/partition_consumer.rb

@@ -0,0 +1,225 @@
+module Poseidon
+  # A primitive Kafka Consumer which operates on a specific broker, topic and partition.
+  #
+  # Example in the README.
+  #
+  # @api public
+  class PartitionConsumer
+    # The offset of the latest message the broker recieved for this partition.
+    # Useful for knowning how far behind the consumer is. This value is only
+    # as recent as the last fetch call.
+    attr_reader :highwater_mark
+
+    attr_reader :host, :port
+
+    attr_reader :offset
+
+    attr_reader :topic
+
+    # Returns a consumer pointing at the lead broker for the partition.
+    #
+    # Eventually this will be replaced by higher level consumer functionality,
+    # this is a stop-gap.
+    #
+    def self.consumer_for_partition(client_id, seed_brokers, topic, partition, offset, options = {})
+
+      broker = BrokerPool.open(client_id, seed_brokers, options[:socket_timeout_ms] || 10_000) do |broker_pool|
+        cluster_metadata = ClusterMetadata.new
+        cluster_metadata.update(broker_pool.fetch_metadata([topic]))
+
+        cluster_metadata.lead_broker_for_partition(topic, partition)
+      end
+
+      new(client_id, broker.host, broker.port, topic, partition, offset, options)
+    end
+
+    # Create a new consumer which reads the specified topic and partition from
+    # the host.
+    #
+    # @param [String] client_id  Used to identify this client should be unique.
+    # @param [String] host
+    # @param [Integer] port 
+    # @param [String] topic Topic to read from
+    # @param [Integer] partition Partitions are zero indexed.
+    # @param [Integer,Symbol] offset 
+    #   Offset to start reading from. A negative offset can also be passed.
+    #   There are a couple special offsets which can be passed as symbols:
+    #     :earliest_offset       Start reading from the first offset the server has.
+    #     :latest_offset         Start reading from the latest offset the server has.
+    #
+    # @param [Hash] options
+    #   Theses options can all be overridden in each individual fetch command.
+    #
+    # @option options [Integer] :max_bytes
+    #   Maximum number of bytes to fetch
+    #   Default: 1048576 (1MB)
+    #
+    # @option options [Integer] :max_wait_ms
+    #   How long to block until the server sends us data.
+    #   NOTE: This is only enforced if min_bytes is > 0.
+    #   Default: 100 (100ms)
+    #
+    # @option options [Integer] :min_bytes
+    #   Smallest amount of data the server should send us.
+    #   Default: 1 (Send us data as soon as it is ready)
+    #
+    # @option options [Integer] :socket_timeout_ms
+    #   How long to wait for reply from server. Should be higher than max_wait_ms.
+    #   Default: 10000 (10s)
+    #
+    # @api public
+    def initialize(client_id, host, port, topic, partition, offset, options = {})
+      @host = host
+      @port = port
+
+      handle_options(options)
+
+      @connection = Connection.new(host, port, client_id, @socket_timeout_ms)
+      @topic = topic
+      @partition = partition
+      if Symbol === offset
+        raise ArgumentError, "Unknown special offset type: #{offset}" unless [:earliest_offset, :latest_offset].include?(offset)
+      end
+      @offset = offset
+    end
+
+    # Fetch messages from the broker.
+    #
+    # @param [Hash] options
+    #
+    # @option options [Integer] :max_bytes
+    #   Maximum number of bytes to fetch
+    #
+    # @option options [Integer] :max_wait_ms
+    #   How long to block until the server sends us data.
+    #
+    # @option options [Integer] :min_bytes
+    #   Smallest amount of data the server should send us.
+    #
+    # @api public
+    def fetch(options = {})
+      fetch_max_wait = options.delete(:max_wait_ms) || max_wait_ms
+      fetch_max_bytes = options.delete(:max_bytes) || max_bytes
+      fetch_min_bytes = options.delete(:min_bytes) || min_bytes
+
+      if options.keys.any?
+        raise ArgumentError, "Unknown options: #{options.keys.inspect}"
+      end
+
+      topic_fetches = build_topic_fetch_request(fetch_max_bytes)
+      fetch_response = @connection.fetch(fetch_max_wait, fetch_min_bytes, topic_fetches)
+      topic_response = fetch_response.topic_fetch_responses.first 
+      partition_response = topic_response.partition_fetch_responses.first
+
+      unless partition_response.error == Errors::NO_ERROR_CODE
+        if @offset < 0 &&
+          Errors::ERROR_CODES[partition_response.error] == Errors::OffsetOutOfRange
+          @offset = :earliest_offset
+          return fetch(options)
+        end
+
+        raise Errors::ERROR_CODES[partition_response.error]
+      else
+        @highwater_mark = partition_response.highwater_mark_offset
+        messages = partition_response.message_set.flatten.map do |m|
+          FetchedMessage.new(topic_response.topic, m.value, m.key, m.offset)
+        end
+        if messages.any?
+          @offset = messages.last.offset + 1
+        end
+        messages
+      end
+    end
+
+    # @return [Integer] next offset we will fetch
+    #
+    # @api public
+    def next_offset
+      resolve_offset_if_necessary
+      @offset
+    end
+
+    # Close the connection to the kafka broker
+    #
+    # @return [Nil]
+    #
+    # @api public
+    def close
+      @connection.close
+      nil
+    end
+
+    private
+    def handle_options(options)
+      @max_bytes         = options.delete(:max_bytes) || 1024*1024
+      @min_bytes         = options.delete(:min_bytes) || 1
+      @max_wait_ms       = options.delete(:max_wait_ms) || 10_000
+      @socket_timeout_ms = options.delete(:socket_timeout_ms) || @max_wait_ms + 10_000
+
+      if @socket_timeout_ms < @max_wait_ms
+        raise ArgumentError, "Setting socket_timeout_ms should be higher than max_wait_ms"
+      end
+
+      if options.keys.any?
+        raise ArgumentError, "Unknown options: #{options.keys.inspect}"
+      end
+    end
+
+    def max_wait_ms
+      @max_wait_ms
+    end
+
+    def max_bytes
+      @max_bytes
+    end
+
+    def min_bytes
+      @min_bytes
+    end
+
+    def resolve_offset_if_necessary
+      return unless Symbol === @offset || @offset < 0
+
+      protocol_offset = case @offset
+      when :earliest_offset
+        -2
+      when :latest_offset
+        -1
+      else
+        -1
+      end
+
+      topic_offset_responses = @connection.offset(build_topic_offset_request(protocol_offset))
+      partition_offsets = topic_offset_responses.first.partition_offsets
+      if partition_offsets.first.error != Errors::NO_ERROR_CODE
+        raise Errors::ERROR_CODES[partition_offsets.first.error]
+      end
+
+      offset_struct = partition_offsets.first.offsets.first
+
+      @offset = if offset_struct.nil?
+        0
+      elsif @offset.kind_of?(Fixnum) && @offset < 0
+        offset_struct.offset + @offset
+      else
+        offset_struct.offset
+      end
+    end
+
+    def build_topic_offset_request(protocol_offset)
+      partition_offset_request = Protocol::PartitionOffsetRequest.new(
+        @partition,
+        protocol_offset,
+        max_number_of_offsets = 1)
+        
+      [Protocol::TopicOffsetRequest.new(topic, [partition_offset_request])]
+    end
+
+    def build_topic_fetch_request(max_bytes)
+      partition_fetches = [Protocol::PartitionFetch.new(@partition,
+                                                        next_offset,
+                                                        max_bytes)]
+      topic_fetches = [Protocol::TopicFetch.new(topic, partition_fetches)]
+    end
+  end
+end