ralphrodkey-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,126 @@
+# 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)
+
+      error_code = parse_error_code(data[0,2])
+      case error_code
+      when -1
+        raise Kafka::UnknownException.new
+      when 1
+        raise Kafka::OffsetOutOfRangeException.new
+      when 2
+        raise Kafka::InvalidMessageCodeException.new
+      when 3
+        raise Kafka::WrongPartitionException.new
+      when 4
+        raise Kafka::InvalidFetchSizeException.new
+      end
+
+      data[2, data.length]
+    end
+
+    def parse_error_code(bytes)
+      unsigned = bytes.unpack("n").shift
+      length = 16
+      max = 2**length-1
+      mid = 2**(length-1)
+      (unsigned >= mid) ? -((unsigned ^ max) + 1) : unsigned
+    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,53 @@
+# 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
+    UNKNOWN_ERROR           = -1
+    NO_ERROR                = 0
+    OFFSET_OUT_OF_RANGE     = 1
+    INVALID_MESSAGE_CODE    = 2
+    WRONG_PARTITION_CODE    = 3
+    INVALID_RETCH_SIZE_CODE = 4
+
+    STRINGS = {
+      -1 => 'Unknown error',
+      0 => 'No error',
+      1 => 'Offset out of range',
+      2 => 'Invalid message code',
+      3 => 'Wrong partition code',
+      4 => 'Invalid fetch size code',
+    }
+
+    def self.to_s(code)
+      STRINGS[code] || 'Unknown error'
+    end
+  end
+
+  class UnknownException < IOError
+  end
+
+  class OffsetOutOfRangeException < IOError
+  end
+
+  class InvalidMessageCodeException < IOError
+  end
+
+  class WrongPartitionException < IOError
+  end
+
+  class InvalidFetchSizeException < IOError
+  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