codeclimate-poseidon 0.0.8

data/lib/poseidon/producer.rb

@@ -0,0 +1,199 @@
+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 = [
+      :ack_timeout_ms,
+      :compressed_topics,
+      :compression_codec,
+      :max_send_retries,
+      :metadata_refresh_interval_ms,
+      :partitioner,
+      :retry_backoff_ms,
+      :required_acks,
+      :socket_timeout_ms,
+      :type,
+    ]
+
+    # @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: Milliseconds] :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] :ack_timeout_ms (1500)
+    #   How long the producer waits for acks.
+    #
+    # @option options [Integer] :socket_timeout_ms] (10000)
+    #   How long the producer socket waits for any reply from server.
+    #
+    # @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. Raises an exception if the producer fails to send the messages.
+    #
+    # @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 close
+      @shutdown = true
+      @producer.close
+    end
+
+    alias_method :shutdown, :close
+
+    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,37 @@
+module Poseidon
+  # @api private
+  class ProducerCompressionConfig
+    COMPRESSION_CODEC_MAP = {
+      :gzip   => Compression::GzipCodec,
+      :snappy => Compression::SnappyCodec,
+      :none   => nil
+    }
+
+    def initialize(compression_codec, compressed_topics)
+      if compression_codec
+        unless COMPRESSION_CODEC_MAP.has_key?(compression_codec)
+          raise ArgumentError, "Unknown compression codec: '#{compression_codec}' (accepted: #{COMPRESSION_CODEC_MAP.keys.inspect})"
+        end
+        @compression_codec = COMPRESSION_CODEC_MAP[compression_codec]
+      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.rb

@@ -0,0 +1,122 @@
+module Poseidon
+  # @api private
+  module Protocol
+    require "poseidon/protocol/protocol_struct"
+    require "poseidon/protocol/request_buffer"
+    require "poseidon/protocol/response_buffer"
+
+    API_KEYS = {
+      :produce => 0,
+      :fetch => 1,
+      :offset => 2,
+      :metadata => 3
+    }
+
+    # Request/Response Common Structures
+    RequestCommon = ProtocolStruct.new(:api_key => :int16,
+                                       :api_version => :int16,
+                                       :correlation_id => :int32,
+                                       :client_id => :string)
+    ResponseCommon = ProtocolStruct.new(:correlation_id => :int32)
+
+    # MessageSet Common Structure
+    MessageStruct = ProtocolStruct.new(:magic_type => :int8,
+                                       :attributes => :int8,
+                                       :key => :bytes,
+                                       :value => :bytes).prepend_size.prepend_crc32.truncatable
+    MessageWithOffsetStruct = ProtocolStruct.new(:offset => :int64,
+                                                 :message => MessageStruct)
+     
+    # When part of produce requests of fetch responses a MessageSet
+    # has a prepended size.  When a MessageSet is compressed and
+    # nested in a Message size is not prepended.
+    MessageSetStruct = ProtocolStruct.new(:messages => [Message]).
+                                          size_bound_array(:messages)
+    MessageSetStructWithSize = MessageSetStruct.dup.prepend_size
+
+    # Produce Request
+    MessagesForPartition = ProtocolStruct.new(:partition    => :int32,
+                                              :message_set  => MessageSet)
+    MessagesForTopic = ProtocolStruct.new(:topic => :string,
+                                          :messages_for_partitions =>
+                                            [MessagesForPartition])
+    ProduceRequest = ProtocolStruct.new(:common => RequestCommon,
+                                        :required_acks => :int16,
+                                        :timeout => :int32,
+                                        :messages_for_topics => [MessagesForTopic])
+
+    # Produce Response
+    ProducePartitionResponse = ProtocolStruct.new(:partition => :int32,
+                                                  :error => :int16,
+                                                  :offset => :int64)
+    ProduceTopicResponse = ProtocolStruct.new(:topic => :string,
+                                              :partitions => [ProducePartitionResponse])
+    ProduceResponse = ProtocolStruct.new(:common => ResponseCommon,
+                                         :topic_response => [ProduceTopicResponse])
+
+    # Fetch Request
+    PartitionFetch = ProtocolStruct.new(:partition => :int32,
+                                        :fetch_offset => :int64,
+                                        :max_bytes => :int32)
+    TopicFetch = ProtocolStruct.new(:topic => :string,
+                                    :partition_fetches => [PartitionFetch])
+    FetchRequest = ProtocolStruct.new(:common => RequestCommon,
+                                      :replica_id => :int32,
+                                      :max_wait_time => :int32,
+                                      :min_bytes => :int32,
+                                      :topic_fetches => [TopicFetch])
+
+    # Fetch Response
+    PartitionFetchResponse = ProtocolStruct.new(:partition => :int32,
+                                                :error => :int16,
+                                                :highwater_mark_offset => :int64,
+                                                :message_set => MessageSet)
+    TopicFetchResponse = ProtocolStruct.new(:topic => :string,
+                                            :partition_fetch_responses => [PartitionFetchResponse])
+    FetchResponse = ProtocolStruct.new(
+      :common => ResponseCommon,
+      :topic_fetch_responses => [TopicFetchResponse])
+
+    # Offset Request
+    PartitionOffsetRequest = ProtocolStruct.new(:partition => :int32,
+                                                :time => :int64,
+                                                :max_number_of_offsets => :int32)
+    TopicOffsetRequest = ProtocolStruct.new(
+      :topic => :string,
+      :partition_offset_requests => [PartitionOffsetRequest])
+    OffsetRequest = ProtocolStruct.new(:common => RequestCommon,
+                                       :replica_id => :int32,
+                                       :topic_offset_requests => [TopicOffsetRequest])
+
+    # Offset Response
+    Offset = ProtocolStruct.new(:offset => :int64)
+    PartitionOffset = ProtocolStruct.new(:partition => :int32,
+                                         :error => :int16,
+                                         :offsets => [Offset])
+    TopicOffsetResponse = ProtocolStruct.new(:topic => :string,
+                                             :partition_offsets => [PartitionOffset])
+    OffsetResponse = ProtocolStruct.new(
+      :common => ResponseCommon,
+      :topic_offset_responses => [TopicOffsetResponse])
+
+    # Metadata Request
+    MetadataRequest = ProtocolStruct.new( :common => RequestCommon,
+                                          :topic_names => [:string])
+
+    # Metadata Response
+    Broker = ProtocolStruct.new(:id => :int32,
+                                :host => :string,
+                                :port => :int32)
+    PartitionMetadata = ProtocolStruct.new(:error     => :int16,
+                                           :id        => :int32,
+                                           :leader    => :int32,
+                                           :replicas  => [:int32],
+                                           :isr       => [:int32])
+    TopicMetadataStruct = ProtocolStruct.new(:error => :int16,
+                                       :name => :string,
+                                       :partitions => [PartitionMetadata])
+    MetadataResponse = ProtocolStruct.new(:common => ResponseCommon,
+                                          :brokers => [Broker],
+                                          :topics => [TopicMetadata])
+  end
+end

data/lib/poseidon/protocol/protocol_struct.rb

@@ -0,0 +1,256 @@
+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?
+          if !have_header?(buffer)
+            @truncated = true
+            return
+          end
+
+          @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 have_header?(buffer)
+        if self.class.truncatable?
+          if self.class.prepend_crc32?
+            header_bytes = 8
+          else
+            header_bytes = 4
+          end
+
+          return buffer.bytes_remaining >= header_bytes
+        else
+          return true
+        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