poseidon 0.0.1

data/lib/poseidon/message.rb

@@ -0,0 +1,151 @@
+module Poseidon
+  # The Message class is used by both Producer and Consumer classes.
+  #
+  # = Basic usage
+  #
+  #    message = Poseidon::Message.new(:value => "hello", 
+  #                                    :key => "user:123", 
+  #                                    :topic => "salutations")
+  #
+  # = Sending a message
+  #
+  # When sending a message you must set the topic for the message, this
+  # can be done during creation or afterwards.
+  #
+  # = Compression
+  #
+  # In normal usage you should never have to worry about compressed
+  # Message objects. When producing the producer takes care of 
+  # compressing the messages and when fetching the fetcher will 
+  # return them decompressed.
+  #
+  # @api private
+  class Message
+    # Last 3 bits are used to indicate compression
+    COMPRESSION_MASK = 0x7
+    MAGIC_TYPE = 0
+
+    # Build a new Message object from its binary representation
+    #
+    # @param [ResponseBuffer] buffer
+    #   a response buffer containing binary data representing a message.
+    #
+    # @return [Message]
+    def self.read(buffer)
+      m = Message.new
+      m.struct = Protocol::MessageWithOffsetStruct.read(buffer)
+
+      # Return nil if the message is truncated.
+      if m.struct.message.truncated?
+        return nil
+      end
+
+      if m.struct.message.checksum_failed?
+        raise Errors::ChecksumError
+      end
+      m
+    end
+
+    attr_accessor :struct, :topic
+
+    # Create a new message object
+    # 
+    # @param [Hash] options
+    #
+    # @option options [String] :value (nil)
+    #   The messages value. Optional.
+    #
+    # @option options [String] :key (nil)
+    #   The messages key. Optional.
+    #
+    # @option options [String] :topic (nil)
+    #   The topic we should send this message to. Optional.
+    #
+    # @option options [String] :attributes (nil)
+    #   Attributes field for the message currently only idicates
+    #   whether or not the message is compressed.
+    def initialize(options = {})
+      build_struct(options)
+
+      @topic = options.delete(:topic)
+
+      if options.any?
+        raise ArgumentError, "Unknown options: #{options.keys.inspect}"
+      end
+    end
+
+    def ==(other)
+      eql?(other)
+    end
+
+    def eql?(other)
+      struct.eql?(other.struct)
+    end
+
+    def objects_with_errors
+      struct.objects_with_errors
+    end
+
+    # Write a binary representation of the message to buffer
+    #
+    # @param [RequestBuffer] buffer
+    # @return [nil]
+    def write(buffer)
+      @struct.write(buffer)
+      nil
+    end
+
+    # @return [String] the Message's key
+    def key
+      @struct.message.key
+    end
+
+    # @return [String] the Message's value
+    def value
+      @struct.message.value
+    end
+
+    # @return [Integer] the Message's offset
+    def offset
+      @struct.offset
+    end
+
+    # Is the value compressed?
+    #
+    # @return [Boolean] 
+    def compressed?
+      compression_codec_id > 0
+    end
+
+    # Decompressed value
+    #
+    # @return [String] decompressed value
+    def decompressed_value
+      compression_codec.decompress(value)
+    end
+
+    private
+    def attributes
+      @struct.message.attributes
+    end
+
+    def compression_codec
+      Compression.find_codec(compression_codec_id)
+    end
+
+    def compression_codec_id
+      attributes & COMPRESSION_MASK
+    end
+
+    def build_struct(options)
+      message_struct = Protocol::MessageStruct.new(
+        MAGIC_TYPE,
+        options.delete(:attributes) || 0,
+        options.delete(:key),
+        options.delete(:value)
+      )
+      struct = Protocol::MessageWithOffsetStruct.new(options.delete(:offset) || 0, message_struct)
+      self.struct = struct
+    end
+  end
+end

data/lib/poseidon/message_conductor.rb

@@ -0,0 +1,84 @@
+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
+      @partition_counter  = -1 
+    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.partitions[partition_id].leader || 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
+        next_partition_counter % partition_count
+      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,39 @@
+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
+  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_for_broker)
+      @messages -= messages_for_broker.messages
+    end
+
+    def all_sent?
+      !@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,154 @@
+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
+
+    # 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.
+    #   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 [:max_bytes] Maximum number of bytes to fetch
+    #   Default: 1048576 (1MB)
+    # @option options [:max_wait_ms] 
+    #   How long to block until the server sends us data.
+    #   Default: 100 (100ms)
+    # @option options [:min_bytes] Smallest amount of data the server should send us.
+    #   Default: 0 (Send us data as soon as it is ready)
+    #
+    # @api public
+    def initialize(client_id, host, port, topic, partition, offset, options = {})
+      @connection = Connection.new(host, port, client_id)
+      @topic = topic
+      @partition = partition
+      if Symbol === offset
+        raise ArgumentError, "Unknown special offset type: #{offset}" unless [:earliest_offset, :latest_offset].include?(offset)
+      end
+      @offset = offset
+      handle_options(options)
+    end
+
+    # Fetch messages from the broker.
+    #
+    # @param [Hash] options
+    #
+    # @option options [:max_bytes]
+    #   Maximum number of bytes to fetch
+    # @option options [:max_wait_ms]
+    #   How long to block until the server sends us data.
+    # @option options [:min_bytes] 
+    #   Smallest amount of data the server should send us.
+    #
+    # @api public
+    def fetch(options = {})
+      fetch_max_wait = options[:max_wait] || max_wait_ms
+      fetch_max_bytes = options[:max_bytes] || max_bytes
+      fetch_min_bytes = options[:min_bytes] || min_bytes
+
+      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
+      if partition_response.error != Errors::NO_ERROR_CODE
+        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
+
+    private
+    def handle_options(options)
+      @max_bytes    = options.delete(:max_bytes) || 1024*1024
+      @min_bytes    = options.delete(:min_bytes) || 0
+      @max_wait_ms  = options.delete(:max_wait_ms) || 10_000
+      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
+
+      if @offset == :earliest_offset
+        @offset = -2
+      elsif @offset == :latest_offset
+        @offset = -1
+      end
+
+      topic_offset_responses = @connection.offset(build_topic_offset_request)
+      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
+      if offset_struct.nil?
+        @offset = 0
+      else
+        @offset = offset_struct.offset
+      end
+    end
+
+    def build_topic_offset_request
+      partition_offset_request = Protocol::PartitionOffsetRequest.new(
+        @partition,
+        @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