leffen-kafka-rb 0.0.15

data/lib/kafka/cli.rb

@@ -0,0 +1,170 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+require 'optparse'
+
+module Kafka
+  module CLI #:nodoc: all
+    extend self
+
+    def publish!
+      read_env
+      parse_args
+      validate_config
+      if config[:message]
+        push(config, config.delete(:message))
+      else
+        publish(config)
+      end
+    end
+
+
+    def subscribe!
+      read_env
+      parse_args
+      validate_config
+      subscribe(config)
+    end
+
+    def validate_config
+      if config[:help]
+        puts help
+        exit
+      end
+      config[:host] ||= IO::HOST
+      config[:port] ||= IO::PORT
+      config[:topic].is_a?(String) or raise "Missing topic"
+
+    rescue RuntimeError => e
+      puts e.message
+      puts help
+      exit
+    end
+
+    def parse_args(args = ARGV)
+      option_parser.parse(args)
+    end
+
+    def read_env(env = ENV)
+      config[:host] = env["KAFKA_HOST"] if env["KAFKA_HOST"]
+      config[:port] = env["KAFKA_PORT"].to_i if env["KAFKA_PORT"]
+      config[:topic] = env["KAFKA_TOPIC"] if env["KAFKA_TOPIC"]
+      config[:compression] = string_to_compression(env["KAFKA_COMPRESSION"]) if env["KAFKA_COMPRESSION"]
+    end
+
+    def config
+      @config ||= {:compression => string_to_compression("no")}
+    end
+
+    def help
+      option_parser.to_s
+    end
+
+    def option_parser
+      OptionParser.new do |opts|
+        opts.banner = "Usage: #{program_name} [options]"
+        opts.separator ""
+
+        opts.on("-h","--host HOST", "Set the kafka hostname") do |h|
+          config[:host] = h
+        end
+
+        opts.on("-p", "--port PORT", "Set the kafka port") do |port|
+          config[:port] = port.to_i
+        end
+
+        opts.on("-t", "--topic TOPIC", "Set the kafka topic") do |topic|
+          config[:topic] = topic
+        end
+
+        opts.on("-c", "--compression no|gzip|snappy", "Set the compression method") do |meth|
+          config[:compression] = string_to_compression(meth)
+        end if publish?
+
+        opts.on("-m","--message MESSAGE", "Message to send") do |msg|
+          config[:message] = msg
+        end if publish?
+
+        opts.separator ""
+
+        opts.on("--help", "show the help") do
+          config[:help] = true
+        end
+
+        opts.separator ""
+        opts.separator "You can set the host, port, topic and compression from the environment variables: KAFKA_HOST, KAFKA_PORT, KAFKA_TOPIC AND KAFKA_COMPRESSION"
+      end
+    end
+
+    def publish?
+      program_name == "leffen-kafka-publish"
+    end
+
+    def subscribe?
+      program_name == "kafka-subscribe"
+    end
+
+    def program_name(pn = $0)
+      File.basename(pn)
+    end
+
+    def string_to_compression(meth)
+      case meth
+      when "no" then Message::NO_COMPRESSION
+      when "gzip" then Message::GZIP_COMPRESSION
+      when "snappy" then Message::SNAPPY_COMPRESSION
+      else raise "No supported compression"
+      end
+    end
+
+    def push(options, message)
+      Producer.new(options).push(Message.new(message))
+    end
+
+    def publish(options)
+      trap(:INT){ exit }
+      producer = Producer.new(options)
+      loop do
+        publish_loop(producer)
+      end
+    end
+
+    def publish_loop(producer)
+      message = read_input
+      producer.push(Message.new(message))
+    end
+
+    def read_input
+      input = $stdin.gets
+      if input
+        input.strip
+      else
+        exit # gets return nil when eof
+      end
+
+    end
+
+    def subscribe(options)
+      trap(:INT){ exit }
+      consumer = Consumer.new(options)
+      consumer.loop do |messages|
+        messages.each do |message|
+          puts message.payload
+        end
+      end
+    end
+
+  end
+end

data/lib/kafka/consumer.rb

@@ -0,0 +1,104 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Kafka
+  class Consumer
+
+    include Kafka::IO
+
+    MAX_SIZE = 1024 * 1024        # 1 megabyte
+    DEFAULT_POLLING_INTERVAL = 2  # 2 seconds
+    MAX_OFFSETS = 1
+    LATEST_OFFSET = -1
+    EARLIEST_OFFSET = -2
+
+    attr_accessor :topic, :partition, :offset, :max_size, :request_type, :polling
+
+    def initialize(options = {})
+      self.topic        = options[:topic]     || "test"
+      self.partition    = options[:partition] || 0
+      self.host         = options[:host]      || HOST
+      self.port         = options[:port]      || PORT
+      self.offset       = options[:offset]
+      self.max_size     = options[:max_size]  || MAX_SIZE
+      self.polling      = options[:polling]   || DEFAULT_POLLING_INTERVAL
+      connect(host, port)
+    end
+
+    def loop(&block)
+      messages = []
+      while (true) do
+        messages = consume
+        block.call(messages) if messages && !messages.empty?
+        sleep(polling)
+      end
+    end
+
+    def consume
+      self.offset ||= fetch_latest_offset
+      send_consume_request
+      message_set = Kafka::Message.parse_from(read_data_response)
+      self.offset += message_set.size
+      message_set.messages
+    rescue SocketError
+      nil
+    end
+
+    def fetch_latest_offset
+      send_offsets_request(LATEST_OFFSET)
+      read_offsets_response
+    end
+
+    def fetch_earliest_offset
+      send_offsets_request(EARLIEST_OFFSET)
+      read_offsets_response
+    end
+
+    def send_offsets_request(offset)
+      write(encoded_request_size)
+      write(encode_request(Kafka::RequestType::OFFSETS, topic, partition, offset, MAX_OFFSETS))
+    end
+
+    def read_offsets_response
+      read_data_response[4,8].reverse.unpack('q')[0]
+    end
+
+    def send_consume_request
+      write(encoded_request_size)
+      write(encode_request(Kafka::RequestType::FETCH, topic, partition, offset, max_size))
+    end
+
+    def read_data_response
+      data_length = read(4).unpack("N").shift
+      data = read(data_length)
+      # TODO: inspect error code instead of skipping it
+      data[2, data.length]
+    end
+
+    def encoded_request_size
+      size = 2 + 2 + topic.length + 4 + 8 + 4
+      [size].pack("N")
+    end
+
+    def encode_request(request_type, topic, partition, offset, max_size)
+      request_type = [request_type].pack("n")
+      topic        = [topic.length].pack('n') + topic
+      partition    = [partition].pack("N")
+      offset       = [offset].pack("q").reverse # DIY 64bit big endian integer
+      max_size     = [max_size].pack("N")
+      request_type + topic + partition + offset + max_size
+    end
+  end
+end

data/lib/kafka/encoder.rb

@@ -0,0 +1,59 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+module Kafka
+  module Encoder
+    def self.message(message, compression = Message::NO_COMPRESSION)
+      message.encode(compression)
+    end
+
+    def self.message_block(topic, partition, messages, compression)
+      message_set = message_set(messages, compression)
+
+      topic     = [topic.length].pack("n") + topic
+      partition = [partition].pack("N")
+      messages  = [message_set.length].pack("N") + message_set
+
+      return topic + partition + messages
+    end
+
+    def self.message_set(messages, compression)
+      message_set = Array(messages).collect { |message|
+        self.message(message)
+      }.join("")
+      message_set = self.message(Message.new(message_set), compression) unless compression == Message::NO_COMPRESSION
+      message_set
+    end
+
+    def self.produce(topic, partition, messages, compression = Message::NO_COMPRESSION)
+      request = [RequestType::PRODUCE].pack("n")
+      data = request + self.message_block(topic, partition, messages, compression)
+
+      return [data.length].pack("N") + data
+    end
+
+    def self.multiproduce(producer_requests, compression = Message::NO_COMPRESSION)
+      part_set = Array(producer_requests).map { |req|
+        self.message_block(req.topic, req.partition, req.messages, compression)
+      }
+
+      request = [RequestType::MULTIPRODUCE].pack("n")
+      parts = [part_set.length].pack("n") + part_set.join("")
+      data = request + parts
+
+      return [data.length].pack("N") + data
+    end
+  end
+end

data/lib/kafka/error_codes.rb

@@ -0,0 +1,35 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+# 
+#    http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+module Kafka
+  module ErrorCodes
+    NO_ERROR                = 0
+    OFFSET_OUT_OF_RANGE     = 1
+    INVALID_MESSAGE_CODE    = 2
+    WRONG_PARTITION_CODE    = 3
+    INVALID_RETCH_SIZE_CODE = 4
+
+    STRINGS = {
+      0 => 'No error',
+      1 => 'Offset out of range',
+      2 => 'Invalid message code',
+      3 => 'Wrong partition code',
+      4 => 'Invalid retch size code',
+    }
+
+    def self.to_s(code)
+      STRINGS[code] || 'Unknown error'
+    end
+  end
+end

data/lib/kafka/io.rb

@@ -0,0 +1,57 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+# 
+#    http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+module Kafka
+  module IO
+    attr_accessor :socket, :host, :port, :compression
+
+    HOST = "localhost"
+    PORT = 9092
+
+    def connect(host, port)
+      raise ArgumentError, "No host or port specified" unless host && port
+      self.host = host
+      self.port = port
+      self.socket = TCPSocket.new(host, port)
+    end
+
+    def reconnect
+      self.socket = TCPSocket.new(self.host, self.port)
+    rescue
+      self.disconnect
+      raise
+    end
+
+    def disconnect
+      self.socket.close rescue nil
+      self.socket = nil
+    end
+
+    def read(length)
+      self.socket.read(length) || raise(SocketError, "no data")
+    rescue
+      self.disconnect
+      raise SocketError, "cannot read: #{$!.message}"
+    end
+
+    def write(data)
+      self.reconnect unless self.socket
+      self.socket.write(data)
+    rescue
+      self.disconnect
+      raise SocketError, "cannot write: #{$!.message}"
+    end
+
+  end
+end

data/lib/kafka/message.rb

@@ -0,0 +1,209 @@
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License.  You may obtain a copy of the License at
+# 
+#    http://www.apache.org/licenses/LICENSE-2.0
+# 
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+module Kafka
+
+  # A message. The format of a message is as follows:
+  #
+  # 4 byte big-endian int: length of message in bytes (including the rest of
+  #                        the header, but excluding the length field itself)
+  # 1 byte: "magic" identifier (format version number)
+  #
+  # If the magic byte == 0, there is one more header field:
+  #
+  # 4 byte big-endian int: CRC32 checksum of the payload
+  #
+  # If the magic byte == 1, there are two more header fields:
+  #
+  # 1 byte: "attributes" (flags for compression, codec etc)
+  # 4 byte big-endian int: CRC32 checksum of the payload
+  #
+  # All following bytes are the payload.
+  class Message
+
+    MAGIC_IDENTIFIER_DEFAULT = 0
+    MAGIC_IDENTIFIER_COMPRESSION = 1
+    NO_COMPRESSION = 0
+    GZIP_COMPRESSION = 1
+    SNAPPY_COMPRESSION = 2
+    BASIC_MESSAGE_HEADER = 'NC'.freeze
+    VERSION_0_HEADER = 'N'.freeze
+    VERSION_1_HEADER = 'CN'.freeze
+    COMPRESSION_CODEC_MASK = 0x03
+
+    attr_accessor :magic, :checksum, :payload
+
+    def initialize(payload = nil, magic = MAGIC_IDENTIFIER_DEFAULT, checksum = nil)
+      self.magic       = magic
+      self.payload     = payload || ""
+      self.checksum    = checksum || self.calculate_checksum
+      @compression = NO_COMPRESSION
+    end
+
+    def calculate_checksum
+      Zlib.crc32(self.payload)
+    end
+
+    def valid?
+      self.checksum == calculate_checksum
+    end
+
+    # Takes a byte string containing one or more messages; returns a MessageSet
+    # with the messages parsed from the string, and the number of bytes
+    # consumed from the string.
+    def self.parse_from(data)
+      messages = []
+      bytes_processed = 0
+
+      while bytes_processed <= data.length - 5 # 5 = size of BASIC_MESSAGE_HEADER
+        message_size, magic = data[bytes_processed, 5].unpack(BASIC_MESSAGE_HEADER)
+        break if bytes_processed + message_size + 4 > data.length # message is truncated
+
+        case magic
+        when MAGIC_IDENTIFIER_DEFAULT
+          # |  0  |  1  |  2  |  3  |  4  |  5  |  6  |  7  |  8  |  9      ...
+          # |                       |     |                       |
+          # |      message_size     |magic|        checksum       | payload ...
+          payload_size = message_size - 5 # 5 = sizeof(magic) + sizeof(checksum)
+          checksum = data[bytes_processed + 5, 4].unpack(VERSION_0_HEADER).shift
+          payload  = data[bytes_processed + 9, payload_size]
+          messages << Kafka::Message.new(payload, magic, checksum)
+
+        when MAGIC_IDENTIFIER_COMPRESSION
+          # |  0  |  1  |  2  |  3  |  4  |  5  |  6  |  7  |  8  |  9  | 10      ...
+          # |                       |     |     |                       |
+          # |         size          |magic|attrs|        checksum       | payload ...
+          payload_size = message_size - 6 # 6 = sizeof(magic) + sizeof(attrs) + sizeof(checksum)
+          attributes, checksum = data[bytes_processed + 5, 5].unpack(VERSION_1_HEADER)
+          payload = data[bytes_processed + 10, payload_size]
+
+          case attributes & COMPRESSION_CODEC_MASK
+          when NO_COMPRESSION # a single uncompressed message
+            messages << Kafka::Message.new(payload, magic, checksum)
+          when GZIP_COMPRESSION # a gzip-compressed message set -- parse recursively
+            uncompressed = Zlib::GzipReader.new(StringIO.new(payload)).read
+            message_set = parse_from(uncompressed)
+            raise 'malformed compressed message' if message_set.size != uncompressed.size
+            messages.concat(message_set.messages)
+          when SNAPPY_COMPRESSION # a snappy-compresses message set -- parse recursively
+            ensure_snappy! do
+              uncompressed = Snappy::Reader.new(StringIO.new(payload)).read
+              message_set = parse_from(uncompressed)
+              raise 'malformed compressed message' if message_set.size != uncompressed.size
+              messages.concat(message_set.messages)
+            end
+          else
+            # https://cwiki.apache.org/confluence/display/KAFKA/Compression
+            raise "Unsupported Kafka compression codec: #{attributes & COMPRESSION_CODEC_MASK}"
+          end
+
+        else
+          raise "Unsupported Kafka message version: magic number #{magic}"
+        end
+
+        bytes_processed += message_size + 4 # 4 = sizeof(message_size)
+      end
+
+      MessageSet.new(bytes_processed, messages)
+    end
+
+    def encode(compression = NO_COMPRESSION)
+      @compression = compression
+
+      self.payload = asciify_payload
+      self.payload = compress_payload if compression?
+
+      data = magic_and_compression + [calculate_checksum].pack("N") + payload
+      [data.length].pack("N") + data
+    end
+
+
+    # Encapsulates a list of Kafka messages (as Kafka::Message objects in the
+    # +messages+ attribute) and their total serialized size in bytes (the +size+
+    # attribute).
+    class MessageSet < Struct.new(:size, :messages); end
+
+    def self.ensure_snappy!
+      if Object.const_defined? "Snappy"
+        yield
+      else
+        fail "Snappy not available!"
+      end
+    end
+
+    def ensure_snappy! &block
+      self.class.ensure_snappy! &block
+    end
+
+    private
+
+    attr_reader :compression
+
+    def compression?
+      compression != NO_COMPRESSION
+    end
+
+    def magic_and_compression
+      if compression?
+        [MAGIC_IDENTIFIER_COMPRESSION, compression].pack("CC")
+      else
+        [MAGIC_IDENTIFIER_DEFAULT].pack("C")
+      end
+    end
+
+    def asciify_payload
+      if RUBY_VERSION[0, 3] == "1.8"
+        payload
+      else
+        payload.to_s.force_encoding(Encoding::ASCII_8BIT)
+      end
+    end
+
+    def compress_payload
+      case compression
+        when GZIP_COMPRESSION
+          gzip
+        when SNAPPY_COMPRESSION
+          snappy
+      end
+    end
+
+    def gzip
+      with_buffer do |buffer|
+        gz = Zlib::GzipWriter.new buffer, nil, nil
+        gz.write payload
+        gz.close
+      end
+    end
+
+    def snappy
+      ensure_snappy! do
+        with_buffer do |buffer|
+          Snappy::Writer.new buffer do |w|
+            w << payload
+          end
+        end
+      end
+    end
+
+    def with_buffer
+      buffer = StringIO.new
+      buffer.set_encoding Encoding::ASCII_8BIT unless RUBY_VERSION =~ /^1\.8/
+      yield buffer if block_given?
+      buffer.rewind
+      buffer.string
+    end
+  end
+end
+