poseidon 0.0.1

data/lib/poseidon/producer.rb

@@ -0,0 +1,193 @@
+module Poseidon
+  # Provides a high level interface for sending messages to a cluster
+  # of Kafka brokers.
+  #
+  # ## Producer Creation
+  #
+  # Producer requires a broker list and a client_id:
+  #
+  #     producer = Producer.new(["broker1:port1", "broker2:port1"], "my_client_id",
+  #                             :type => :sync)
+  #
+  # The broker list is only used to bootstrap our knowledge of the cluster --
+  # it does not need to contain every broker. The client id should be unique
+  # across all clients in the cluster.
+  #
+  # ## Sending Messages
+  #
+  # Messages must have a topic before being sent:
+  #
+  #     messages = []
+  #     messages << Poseidon::MessageToSend.new("topic1", "Hello Word")
+  #     messages << Poseidon::MessageToSend.new("user_updates_topic", user.update, user.id)
+  #     producer.send_messages(messages)
+  #
+  # ## Producer Types
+  #
+  # There are two types of producers: sync and async. They can be specified
+  # via the :type option when creating a producer.
+  #
+  # ## Sync Producer
+  #
+  # The :sync producer blocks while sends messages to the cluster. The more
+  # messages you can send per #send_messages call the more efficient it will
+  # be.
+  #
+  # ## Compression
+  #
+  # When creating the producer you can specify a compression method:
+  #
+  #     producer = Producer.new(["broker1:port1"], "my_client_id",
+  #                             :type => :sync, :compression_codec => :gzip)
+  #
+  # If you don't specify which topics to compress it will compress all topics.
+  # You can specify a set of topics to compress when creating the producer:
+  #
+  #     producer = Producer.new(["broker1:port1"], "my_client_id",
+  #                             :type => :sync, :compression_codec => :gzip,
+  #                             :compressed_topics => ["compressed_topic_1"])
+  #
+  # ## Partitioning
+  #
+  # For keyless messages the producer will round-robin messages to all 
+  # _available_ partitions for at topic. This means that if we are unable to
+  # send messages to a specific broker we'll retry sending those to a different
+  # broker.
+  #
+  # However, if you specify a key when creating the message, the producer
+  # will choose a partition based on the key and only send to that partition.
+  # 
+  # ## Custom Partitioning
+  #
+  # You may also specify a custom partitioning scheme for messages by passing
+  # a Proc (or any object that responds to #call) to the Producer.  The proc
+  # must return a Fixnum >= 0 and less-than partition_count.
+  #
+  #     my_partitioner = Proc.new { |key, partition_count|  Zlib::crc32(key) % partition_count }
+  #
+  #     producer = Producer.new(["broker1:port1", "broker2:port1"], "my_client_id",
+  #                             :type => :sync, :partitioner => my_partitioner)
+  #
+  # @api public
+  class Producer
+    # @api private
+    VALID_OPTIONS = [
+      :type,
+      :compression_codec,
+      :compressed_topics,
+      :metadata_refresh_interval_ms,
+      :partitioner,
+      :max_send_retries,
+      :retry_backoff_ms,
+      :required_acks,
+      :ack_timeout_ms
+    ]
+
+    # @api private
+    OPTION_DEFAULTS = {
+      :type => :sync
+    }
+
+    # Returns a new Producer.
+    #
+    # @param [Array<String>] brokers An array of brokers in the form "host1:port1"
+    #
+    # @param [String] client_id A client_id used to indentify the producer.
+    #
+    # @param [Hash] options
+    #
+    # @option options [:sync / :async] :type (:sync)
+    #   Whether we should send messages right away or queue them and send
+    #   them in the background.
+    #
+    # @option options [:gzip / :snappy / :none] :compression_codec (:none)
+    #   Type of compression to use.
+    # 
+    # @option options [Enumberable<String>] :compressed_topics (nil)
+    #   Topics to compress.  If this is not specified we will compress all
+    #   topics provided that +:compression_codec+ is set.
+    #
+    # @option options [Integer: Seconds] :metadata_refresh_interval_ms (600_000)
+    #   How frequently we should update the topic metadata in milliseconds.
+    #
+    # @option options [#call, nil] :partitioner
+    #   Object which partitions messages based on key. 
+    #   Responds to #call(key, partition_count).
+    #
+    # @option options [Integer] :max_send_retries (3)
+    #   Number of times to retry sending of messages to a leader.
+    #
+    # @option options [Integer] :retry_backoff_ms (100)
+    #   The amount of time (in milliseconds) to wait before refreshing the metadata
+    #   after we are unable to send messages. 
+    #   Number of times to retry sending of messages to a leader.
+    #
+    # @option options [Integer] :required_acks (0)
+    #   The number of acks required per request.
+    #
+    # @option options [Integer] :request_timeout_ms (1500)
+    #   How long the producer waits for acks.
+    #
+    # @api public
+    def initialize(brokers, client_id, options = {})
+      options = options.dup
+      validate_options(options)
+
+      if !brokers.respond_to?(:each)
+        raise ArgumentError, "brokers must respond to #each"
+      end
+      @brokers    = brokers
+      @client_id  = client_id
+      @producer   = build_producer(options)
+      @shutdown   = false
+    end
+
+    # Send messages to the cluster.
+    #
+    # @param [Enumerable<MessageToSend>] messages
+    #   Messages must have a +topic+ set and may have a +key+ set.
+    #
+    # @return [Boolean]
+    #
+    # @api public
+    def send_messages(messages)
+      raise Errors::ProducerShutdownError if @shutdown
+      if !messages.respond_to?(:each)
+        raise ArgumentError, "messages must respond to #each"
+      end
+
+      @producer.send_messages(convert_to_messages_objects(messages))
+    end
+
+    # Closes all open connections to brokers
+    def shutdown
+      @shutdown = true
+      @producer.shutdown
+    end
+
+    private
+    def validate_options(options)
+      unknown_keys = options.keys - VALID_OPTIONS
+      if unknown_keys.any?
+        raise ArgumentError, "Unknown options: #{unknown_keys.inspect}"
+      end
+
+      @type = options.delete(:type) || :sync
+    end
+
+    def convert_to_messages_objects(messages)
+      messages.map do |m|
+        Message.new(:value => m.value, :topic => m.topic, :key => m.key)
+      end
+    end
+
+    def build_producer(options)
+      case @type
+      when :sync
+        SyncProducer.new(@client_id, @brokers, options)
+      when :async
+        raise "Not implemented yet"
+      end
+    end
+  end
+end

data/lib/poseidon/producer_compression_config.rb

@@ -0,0 +1,36 @@
+module Poseidon
+  # @api private
+  class ProducerCompressionConfig
+    COMPRESSION_CODEC_MAP = {
+      :gzip   => Compression::GzipCodec,
+      :snappy => Compression::SnappyCodec
+    }
+
+    def initialize(compression_codec, compressed_topics)
+      if compression_codec
+        @compression_codec = COMPRESSION_CODEC_MAP[compression_codec]
+        if @compression_codec.nil?
+          raise ArgumentError, "Uknown compression codec: '#{compression_codec}' (accepted: #{COMPRESSION_CODEC_MAP.keys.inspect})"
+        end
+      else
+        @compression_codec = nil
+      end
+
+      if compressed_topics
+        @compressed_topics = Set.new(compressed_topics)
+      else
+        @compressed_topics = nil
+      end
+    end
+
+    def compression_codec_for_topic(topic)
+      return false if @compression_codec.nil?
+
+      if @compressed_topics.nil? || (@compressed_topics && @compressed_topics.include?(topic))
+        @compression_codec
+      else
+        false
+      end
+    end
+  end
+end

data/lib/poseidon/protocol/protocol_struct.rb

@@ -0,0 +1,238 @@
+module Poseidon
+  module Protocol
+    class ProtocolStruct < Struct
+      class EncodingError < StandardError;end
+      class DecodingError < StandardError;end
+
+      def self.new(hash)
+        klass = super(*hash.keys)
+        klass.type_map = hash
+        klass
+      end
+
+      def self.type_map=(type_map)
+        @type_map = type_map
+      end
+
+      def self.type_map
+        @type_map 
+      end
+
+      def self.prepend_size
+        @prepend_size = true
+        self
+      end
+
+      def self.prepend_crc32
+        @prepend_crc32 = true
+        self
+      end
+
+      def self.truncatable
+        @truncatable = true
+        self
+      end
+
+      def self.prepend_size?
+        @prepend_size
+      end
+
+      def self.prepend_crc32?
+        @prepend_crc32
+      end
+
+      def self.truncatable?
+        @truncatable
+      end
+
+      def self.size_bound_array(member)
+        @size_bound_members ||= []
+        @size_bound_members  << member
+        self
+      end
+
+      def self.size_bound_array?(member)
+        @size_bound_members ||= []
+        @size_bound_members.include?(member)
+      end
+
+      # Recursively find all objects with errors
+      def objects_with_errors
+        children = []
+        each_pair do |member, value|
+          case value
+          when Array
+            value.each do |v|
+              if v.respond_to?(:objects_with_errors)
+                children << v
+              end
+            end
+          else
+            if value.respond_to?(:objects_with_errors)
+              children << value
+            end
+          end
+        end
+
+        children_with_errors = children.map(&:objects_with_errors).flatten
+        if members.include?(:error) && self[:error] != Errors::NO_ERROR_CODE
+          children_with_errors + [self]
+        else
+          children_with_errors
+        end
+      end
+
+      def raise_error
+        raise error_class if error_class
+      end
+
+      def error_class
+        Errors::ERROR_CODES[self[:error]]
+      end
+
+      def raise_error_if_one_exists
+        objects_with_errors.each do |object|
+          object.raise_error
+        end
+      end
+
+      def write(buffer)
+        maybe_prepend_size(buffer) do
+          maybe_prepend_crc32(buffer) do
+            each_pair do |member, value|
+              begin
+                write_member(buffer, member, value)
+              rescue
+                raise EncodingError, "Error writting #{member} in #{self.class} (#{$!.class}: #{$!.message})"
+              end
+            end
+          end
+        end
+      end
+
+      def maybe_prepend_size(buffer)
+        if self.class.prepend_size?
+          buffer.prepend_size do
+            yield
+          end
+        else
+          yield
+        end
+      end
+
+      def maybe_prepend_crc32(buffer)
+        if self.class.prepend_crc32?
+          buffer.prepend_crc32 do
+            yield
+          end
+        else
+          yield
+        end
+      end
+
+      def write_member(buffer, member, value)
+        case type = type_map[member]
+        when Array
+          buffer.int32(value.size) unless self.class.size_bound_array?(member)
+          value.each { |v| write_type(buffer, type.first, v) }
+        else
+          write_type(buffer, type, value)
+        end
+      end
+
+      def write_type(buffer, type, value)
+        case type
+        when Symbol
+          buffer.send(type, value)
+        else
+          value.write(buffer)
+        end
+      end
+
+      # Populate struct from buffer based on members and their type definition.
+      def self.read(buffer)
+        s = new
+        s.read(buffer)
+        s
+      end
+
+      def read(buffer)
+        if self.class.prepend_size?
+          @size = buffer.int32
+
+          if self.class.prepend_crc32?
+            @crc32 = buffer.int32
+            @computed_crc32 = [Zlib::crc32(buffer.peek(@size-4))].pack("l>").unpack("l>").first
+            if @crc32 != @computed_crc32
+              @checksum_failed = true
+            end
+            expected_bytes_remaining = @size - 4
+          else
+            expected_bytes_remaining = @size
+          end
+
+
+          if self.class.truncatable? && expected_bytes_remaining > buffer.bytes_remaining
+            @truncated = true
+            return
+          end
+        end
+
+        members.each do |member|
+          begin
+            self[member] = read_member(buffer, member)
+          rescue DecodingError
+            # Just reraise instead of producing a crazy nested exception
+            raise
+          rescue
+            raise DecodingError, "Error while reading #{member} in #{self.class} (#{$!.class}: #{$!.message}))"
+          end
+        end
+      end
+
+      def read_member(buffer, member)
+        case type = type_map[member]
+        when Array
+          if self.class.size_bound_array?(member)
+            if @size
+              array_buffer = ResponseBuffer.new(buffer.read(@size))
+            else
+              array_buffer = buffer
+            end
+
+            array = []
+            while !array_buffer.eof? && (v = read_type(array_buffer, type.first))
+              array << v
+            end
+            array
+          else
+            buffer.int32.times.map { read_type(buffer, type.first) }
+          end
+        else
+          read_type(buffer, type)
+        end
+      end
+
+      def read_type(buffer, type)
+        case type
+        when Symbol
+          buffer.send(type)
+        else
+          type.read(buffer)
+        end
+      end
+
+      def type_map
+        self.class.type_map
+      end
+
+      def checksum_failed?
+        @checksum_failed
+      end
+
+      def truncated?
+        @truncated
+      end
+    end
+  end
+end

data/lib/poseidon/protocol/request_buffer.rb

@@ -0,0 +1,78 @@
+module Poseidon
+  module Protocol
+    # RequestBuffer allows you to build a Binary string for API requests
+    #
+    # API parallels the primitive types described on the wiki, with some
+    # sugar for prepending message sizes and checksums.
+    # (https://cwiki.apache.org/confluence/display/KAFKA/A+Guide+To+The+Kafka+Protocol#AGuideToTheKafkaProtocol-ProtocolPrimitiveTypes) 
+    class RequestBuffer
+      def initialize
+        @s = ''.encode("ASCII-8BIT")
+      end
+
+      def append(string)
+        string = string.dup
+        string.force_encoding("ASCII-8BIT")
+        @s << string
+        nil
+      end
+
+      def int8(int8)
+        append([int8].pack("C"))
+      end
+
+      def int16(int16)
+        append([int16].pack("s>"))
+      end
+
+      def int32(int32)
+        append([int32].pack("l>"))
+      end
+
+      def int64(int64)
+        append([int64].pack("q>"))
+      end
+
+      # Add a string
+      # 
+      # @param [String] string
+      def string(string)
+        if string.nil?
+          int16(-1)
+        else
+          int16(string.size)
+          append(string)
+        end
+      end
+
+      def bytes(string)
+        if string.nil?
+          int32(-1)
+        else
+          int32(string.size)
+          append(string)
+        end
+      end
+
+      def prepend_crc32
+        checksum_pos = @s.size
+        @s += " "
+        yield
+        @s[checksum_pos] = [Zlib::crc32(@s[(checksum_pos+1)..-1])].pack("N")
+        nil
+      end
+
+      def prepend_size
+        size_pos = @s.size
+        @s += " "
+        yield
+        @s[size_pos] = [(@s.size-1) - size_pos].pack("N")
+        nil
+      end
+
+      def to_s
+        @s
+      end
+    end
+  end
+end

data/lib/poseidon/protocol/response_buffer.rb

@@ -0,0 +1,72 @@
+module Poseidon
+  module Protocol
+    class ResponseBuffer
+      def initialize(response)
+        @s = response
+        @pos = 0
+      end
+
+      def int8
+        byte = @s.slice(@pos, 1).unpack("C").first
+        @pos += 1
+        byte
+      end
+
+      def int16
+        short = @s.slice(@pos, 2).unpack("s>").first
+        @pos += 2
+        short
+      end
+
+      def int32
+        int = @s.slice(@pos, 4).unpack("l>").first
+        @pos += 4
+        int
+      end
+
+      def int64
+        long = @s.slice(@pos, 8).unpack("q>").first
+        @pos += 8
+        long
+      end
+
+      def string
+        len = int16
+        string = @s.slice(@pos, len)
+        @pos += len
+        string
+      end
+
+      def read(bytes)
+        data = @s.slice(@pos, bytes)
+        @pos += bytes
+        data
+      end
+
+      def peek(bytes)
+        @s.slice(@pos, bytes)
+      end
+
+      def bytes
+        n = int32
+        if n == -1
+          return nil
+        else
+          read(n)
+        end
+      end
+
+      def bytes_remaining
+        @s.size - @pos
+      end
+
+      def eof?
+        @pos == @s.size
+      end
+
+      def to_s
+        @s
+      end
+    end
+  end
+end