codeclimate-poseidon 0.0.8

data/lib/poseidon/cluster_metadata.rb

@@ -0,0 +1,94 @@
+module Poseidon
+  # Encapsulates what we known about brokers, topics and partitions
+  # from Metadata API calls.
+  #
+  # @api private
+  class ClusterMetadata
+    attr_reader :brokers, :last_refreshed_at, :topic_metadata
+    def initialize
+      @brokers        = {}
+      @topic_metadata = {}
+      @last_refreshed_at = nil
+    end
+
+    # Update what we know about the cluter based on MetadataResponse
+    # 
+    # @param [MetadataResponse] topic_metadata_response
+    # @return nil
+    def update(topic_metadata_response)
+      update_brokers(topic_metadata_response.brokers)
+      update_topics(topic_metadata_response.topics)
+
+      @last_refreshed_at = Time.now
+      nil
+    end
+
+    # Do we have metadata for these topics already?
+    #
+    # @param [Enumberable<String>] topic_names A set of topics.
+    # @return [Boolean] true if we have metadata for all +topic_names+, otherwise false.
+    def have_metadata_for_topics?(topic_names)
+      topic_names.all? { |topic| @topic_metadata[topic] }
+    end
+
+    # Provides metadata for each topic
+    #
+    # @param [Enumerable<String>] topic_names Topics we should return metadata for
+    # @return [Hash<String,TopicMetadata>]
+    def metadata_for_topics(topic_names)
+      Hash[topic_names.map { |name| [name, @topic_metadata[name]] }]
+    end
+
+    # Provides a Broker object for +broker_id+. This corresponds to the
+    # broker ids in the TopicMetadata objects.
+    #
+    # @param [Integer] broker_id Broker id 
+    def broker(broker_id)
+      @brokers[broker_id]
+    end
+
+    # Return lead broker for topic and partition
+    def lead_broker_for_partition(topic_name, partition)
+      broker_id = @topic_metadata[topic_name].partition_leader(partition)
+      if broker_id
+        @brokers[broker_id]
+      else
+        nil
+      end
+    end
+
+    def topics
+      @topic_metadata.keys
+    end
+
+    def to_s
+      out = ""
+      @topic_metadata.each do |topic, metadata|
+        out << "Topic: #{topic}"
+        out << "-------------------------"
+        out << metadata.to_s
+      end
+      out
+    end
+
+    def reset
+      @brokers        = {}
+      @topic_metadata = {}
+    end
+
+    private
+    def update_topics(topics)
+      topics.each do |topic|
+        if topic.exists?
+          @topic_metadata[topic.name] = topic
+        end
+      end
+    end
+
+    def update_brokers(brokers)
+      brokers.each do |broker|
+        @brokers[broker.id] = broker
+      end
+    end
+  end
+end

data/lib/poseidon/compressed_value.rb

@@ -0,0 +1,23 @@
+module Poseidon
+  # @api private
+  class CompressedValue
+    def initialize(value, codec_id)
+      @value = value
+      @codec_id = codec_id
+    end
+
+    # Decompressed value
+    #
+    # Raises ??? if the compression codec is uknown
+    #
+    # @return [String] decompressed value
+    def decompressed
+      @decompressed ||= decompress
+    end
+
+    def compression_codec
+      Compression.find_codec(codec_id)
+    end
+    private
+  end
+end

data/lib/poseidon/compression.rb

@@ -0,0 +1,30 @@
+module Poseidon
+  # @api private
+  module Compression
+    class UnrecognizedCompressionCodec < StandardError; end
+
+    require "poseidon/compression/gzip_codec"
+    require "poseidon/compression/snappy_codec"
+
+    CODECS = {
+      #0 => no codec
+      1 => GzipCodec,
+      2 => SnappyCodec
+    }
+
+    # Fetches codec module for +codec_id+
+    # https://cwiki.apache.org/confluence/display/KAFKA/A+Guide+To+The+Kafka+Protocol#AGuideToTheKafkaProtocol-Compression
+    #
+    # @param [Integer] codec_id codec's as defined by the Kafka Protocol
+    # @return [Module] codec module for codec_id
+    #
+    # @private
+    def self.find_codec(codec_id)
+      codec = CODECS[codec_id]
+      if codec.nil?
+        raise UnrecognizedCompressionCodec, codec_id
+      end
+      codec
+    end
+  end
+end

data/lib/poseidon/compression/gzip_codec.rb

@@ -0,0 +1,23 @@
+module Poseidon
+  module Compression
+    module GzipCodec
+      def self.codec_id
+        1
+      end
+
+      def self.compress(s)
+        io = StringIO.new
+        io.set_encoding(Encoding::BINARY)
+        gz = Zlib::GzipWriter.new io, Zlib::DEFAULT_COMPRESSION, Zlib::DEFAULT_STRATEGY
+        gz.write s
+        gz.close
+        io.string
+      end
+
+      def self.decompress(s)
+        io = StringIO.new(s)
+        Zlib::GzipReader.new(io).read
+      end
+    end
+  end
+end

data/lib/poseidon/compression/snappy_codec.rb

@@ -0,0 +1,29 @@
+module Poseidon
+  module Compression
+    module SnappyCodec
+      def self.codec_id
+        2
+      end
+
+      def self.compress(s)
+        check!
+        Snappy.deflate(s)
+      end
+
+      def self.decompress(s)
+        check!
+        Snappy::Reader.new(StringIO.new(s)).read
+      end
+
+      def self.check!
+        @checked ||= begin
+          require 'snappy'
+          true
+        rescue LoadError
+          raise "Snappy compression is not available, please install the 'snappy' gem"
+        end
+      end
+
+    end
+  end
+end

data/lib/poseidon/connection.rb

@@ -0,0 +1,169 @@
+module Poseidon
+  # High level internal interface to a remote broker. Provides access to
+  # the broker API.
+  # @api private
+  class Connection
+    include Protocol
+
+    class ConnectionFailedError < StandardError; end
+    class TimeoutException < Exception; end
+
+    API_VERSION = 0
+    REPLICA_ID = -1 # Replica id is always -1 for non-brokers
+
+    # @yieldparam [Connection]
+    def self.open(host, port, client_id, socket_timeout_ms, &block)
+      connection = new(host, port, client_id, socket_timeout_ms)
+
+      yield connection
+    ensure
+      connection.close
+    end
+
+    attr_reader :host, :port
+
+    # Create a new connection
+    #
+    # @param [String] host Host to connect to
+    # @param [Integer] port Port broker listens on
+    # @param [String] client_id Unique across processes?
+    def initialize(host, port, client_id, socket_timeout_ms)
+      @host = host
+      @port = port
+
+      @client_id = client_id
+      @socket_timeout_ms = socket_timeout_ms
+    end
+
+    # Close broker connection
+    def close
+      @socket && @socket.close
+    end
+
+    # Execute a produce call
+    #
+    # @param [Integer] required_acks
+    # @param [Integer] timeout
+    # @param [Array<Protocol::MessagesForTopics>] messages_for_topics Messages to send
+    # @return [ProduceResponse]
+    def produce(required_acks, timeout, messages_for_topics)
+      ensure_connected
+      req = ProduceRequest.new( request_common(:produce),
+                                required_acks,
+                                timeout,
+                                messages_for_topics) 
+      send_request(req)
+      if required_acks != 0
+        read_response(ProduceResponse)
+      else
+        true
+      end
+    end
+
+    # Execute a fetch call
+    #
+    # @param [Integer] max_wait_time
+    # @param [Integer] min_bytes
+    # @param [Integer] topic_fetches
+    def fetch(max_wait_time, min_bytes, topic_fetches)
+      ensure_connected
+      req = FetchRequest.new( request_common(:fetch),
+                                REPLICA_ID,
+                                max_wait_time,
+                                min_bytes,
+                                topic_fetches) 
+      send_request(req)
+      read_response(FetchResponse)
+    end
+
+    def offset(offset_topic_requests)
+      ensure_connected
+      req = OffsetRequest.new(request_common(:offset),
+                              REPLICA_ID,
+                              offset_topic_requests)
+      send_request(req)
+      read_response(OffsetResponse).topic_offset_responses
+    end
+
+    # Fetch metadata for +topic_names+
+    #
+    # @param [Enumberable<String>] topic_names
+    #   A list of topics to retrive metadata for
+    # @return [TopicMetadataResponse] metadata for the topics
+    def topic_metadata(topic_names)
+      ensure_connected
+      req = MetadataRequest.new( request_common(:metadata),
+                                 topic_names)
+      send_request(req)
+      read_response(MetadataResponse)
+    end
+
+    private
+    def ensure_connected
+      if @socket.nil? || @socket.closed?
+        begin
+          @socket = TCPSocket.new(@host, @port)
+        rescue SystemCallError
+          raise_connection_failed_error
+        end
+      end
+    end
+
+    def read_response(response_class)
+      r = ensure_read_or_timeout(4)
+      if r.nil?
+        raise_connection_failed_error
+      end
+      n = r.unpack("N").first
+      s = ensure_read_or_timeout(n)
+      buffer = Protocol::ResponseBuffer.new(s)
+      response_class.read(buffer)
+    rescue Errno::ECONNRESET, SocketError, TimeoutException
+      @socket = nil
+      raise_connection_failed_error
+    end
+
+    def ensure_read_or_timeout(maxlen)
+      if IO.select([@socket], nil, nil, @socket_timeout_ms / 1000.0)
+         @socket.read(maxlen)
+      else
+         raise TimeoutException.new
+      end
+    end
+
+    def send_request(request)
+      buffer = Protocol::RequestBuffer.new
+      request.write(buffer)
+      ensure_write_or_timeout([buffer.to_s.bytesize].pack("N") + buffer.to_s)
+    rescue Errno::EPIPE, Errno::ECONNRESET, TimeoutException
+      @socket = nil
+      raise_connection_failed_error
+    end
+
+    def ensure_write_or_timeout(data)
+      if IO.select(nil, [@socket], nil, @socket_timeout_ms / 1000.0)
+        @socket.write(data)
+      else
+        raise TimeoutException.new
+      end
+    end
+
+    def request_common(request_type)
+      RequestCommon.new(
+        API_KEYS[request_type],
+        API_VERSION,
+        next_correlation_id,
+        @client_id
+      )
+    end
+
+    def next_correlation_id
+      @correlation_id ||= 0
+      @correlation_id  += 1
+    end
+
+    def raise_connection_failed_error
+      raise ConnectionFailedError, "Failed to connect to #{@host}:#{@port}"
+    end
+  end
+end

data/lib/poseidon/fetched_message.rb

@@ -0,0 +1,37 @@
+module Poseidon
+
+  # A message fetched from a Kafka broker.
+  #
+  # ```
+  # fetched_messages = consumer.fetch
+  # fetched_messages.each do |fm|
+  #   puts "Topic: #{fm.topic}"
+  #   puts "Value #{fm.value}"
+  #   puts "Key: #{fm.key}"
+  #   puts "Offset: #{fm.offset}"
+  # end
+  # ```
+  #
+  # @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, keys will be
+  #   sent to brokers in a round-robin manner.
+  #
+  # @api public
+  class FetchedMessage
+    attr_reader :value, :key, :topic, :offset
+
+    def initialize(topic, value, key, offset)
+      @topic  = topic
+      @value  = value
+      @key    = key
+      @offset = offset
+    end
+  end
+end

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