Fix failing specs

.gitignore

@@ -0,0 +1,2 @@
+.rvmrc
+

Rakefile

@@ -21,7 +21,7 @@ require 'rspec/core/rake_task'
 
 GEM = 'kafka-rb'
 GEM_NAME = 'Kafka Client'
-GEM_VERSION = '0.0.5'
+GEM_VERSION = '0.0.7'
 AUTHORS = ['Alejandro Crosa']
 EMAIL = "alejandrocrosa@gmail.com"
 HOMEPAGE = "http://github.com/acrosa/kafka-rb"

kafka-rb.gemspec

@@ -16,7 +16,7 @@
 
 Gem::Specification.new do |s|
   s.name = %q{kafka-rb}
-  s.version = "0.0.6"
+  s.version = "0.0.7"
 
   s.required_rubygems_version = Gem::Requirement.new(">= 0") if s.respond_to? :required_rubygems_version=
   s.authors = ["Alejandro Crosa"]
@@ -25,7 +25,7 @@ Gem::Specification.new do |s|
   s.description = %q{kafka-rb allows you to produce and consume messages using the Kafka distributed publish/subscribe messaging service.}
   s.email = %q{alejandrocrosa@gmail.com}
   s.extra_rdoc_files = ["LICENSE"]
-  s.files = ["LICENSE", "README.md", "Rakefile", "lib/kafka", "lib/kafka/batch.rb", "lib/kafka/consumer.rb", "lib/kafka/io.rb", "lib/kafka/message.rb", "lib/kafka/producer.rb", "lib/kafka/request_type.rb", "lib/kafka/error_codes.rb", "lib/kafka.rb", "spec/batch_spec.rb", "spec/consumer_spec.rb", "spec/io_spec.rb", "spec/kafka_spec.rb", "spec/message_spec.rb", "spec/producer_spec.rb", "spec/spec_helper.rb"]
+  s.files = ["LICENSE", "README.md", "Rakefile", "lib/kafka", "lib/kafka/batch.rb", "lib/kafka/consumer.rb", "lib/kafka/consumer_registry.rb", "lib/kafka/io.rb", "lib/kafka/message.rb", "lib/kafka/producer.rb", "lib/kafka/request_type.rb", "lib/kafka/error_codes.rb", "lib/kafka.rb", "spec/batch_spec.rb", "spec/consumer_spec.rb", "spec/io_spec.rb", "spec/kafka_spec.rb", "spec/message_spec.rb", "spec/producer_spec.rb", "spec/spec_helper.rb"]
   s.homepage = %q{http://github.com/acrosa/kafka-rb}
   s.require_paths = ["lib"]
   s.rubygems_version = %q{1.3.7}
@@ -43,4 +43,6 @@ Gem::Specification.new do |s|
   else
     s.add_dependency(%q<rspec>, [">= 0"])
   end
+
+  s.add_dependency(%q<zk>, [">= 0.8.5"])
 end

lib/kafka.rb

@@ -25,6 +25,7 @@
 require File.join(File.dirname(__FILE__), "kafka", "message")
 require File.join(File.dirname(__FILE__), "kafka", "producer")
 require File.join(File.dirname(__FILE__), "kafka", "consumer")
+require File.join(File.dirname(__FILE__), "kafka", "consumer_registry")
 
 module Kafka
 end

lib/kafka/consumer.rb

@@ -20,8 +20,9 @@ class Consumer
     CONSUME_REQUEST_TYPE = Kafka::RequestType::FETCH
     MAX_SIZE = 1048576 # 1 MB
     DEFAULT_POLLING_INTERVAL = 2 # 2 seconds
+    DEFAULT_GROUP_ID = "test"
 
-    attr_accessor :topic, :partition, :offset, :max_size, :request_type, :polling
+    attr_accessor :topic, :partition, :offset, :max_size, :request_type, :polling, :group_id, :consumer_id, :zk_connect
 
     def initialize(options = {})
       self.topic        = options[:topic]        || "test"
@@ -32,6 +33,15 @@ def initialize(options = {})
       self.max_size     = options[:max_size]     || MAX_SIZE
       self.request_type = options[:request_type] || CONSUME_REQUEST_TYPE
       self.polling      = options[:polling]      || DEFAULT_POLLING_INTERVAL
+      self.group_id     = options[:group_id]     || DEFAULT_GROUP_ID
+      self.consumer_id  = options[:consumer_id]  || Consumer.generate_id(self.group_id)
+      self.zk_connect   = options[:zk_connect]
+
+      unless self.zk_connect.nil?
+        @consumer_registry = Kafka::ConsumerRegistry.new(self.zk_connect)
+        @consumer_registry.register(self.group_id, self.consumer_id, self.topic)
+      end
+
       self.connect(self.host, self.port)
     end
 
@@ -94,5 +104,11 @@ def parse_message_set_from(data)
       self.offset += processed
       messages
     end
+
+    private
+
+    def self.generate_id(group_id)
+      "#{group_id}_#{Socket.gethostname}-#{Time.now.usec}"
+    end
   end
 end

lib/kafka/consumer_registry.rb

@@ -0,0 +1,61 @@
+# This class handles the consumers interaction with zookeeper
+#
+# Directories:
+# 1. Consumer id registry:
+# /consumers/[group_id]/ids[consumer_id] -> topic1,...topicN
+# A consumer has a unique consumer id within a consumer group. A consumer registers its id as an ephemeral znode
+# and puts all topics that it subscribes to as the value of the znode. The znode is deleted when the client is gone.
+# A consumer subscribes to event changes of the consumer id registry within its group.
+#
+# The consumer id is picked up from configuration, instead of the sequential id assigned by ZK. Generated sequential
+# ids are hard to recover during temporary connection loss to ZK, since it's difficult for the client to figure out
+# whether the creation of a sequential znode has succeeded or not. More details can be found at
+# (http://wiki.apache.org/hadoop/ZooKeeper/ErrorHandling)
+#
+# 2. Broker node registry:
+# /brokers/[0...N] --> { "host" : "host:port",
+#                        "topics" : {"topic1": ["partition1" ... "partitionN"], ...,
+#                                    "topicN": ["partition1" ... "partitionN"] } }
+# This is a list of all present broker brokers. A unique logical node id is configured on each broker node. A broker
+# node registers itself on start-up and creates a znode with the logical node id under /brokers. The value of the znode
+# is a JSON String that contains (1) the host name and the port the broker is listening to, (2) a list of topics that
+# the broker serves, (3) a list of logical partitions assigned to each topic on the broker.
+# A consumer subscribes to event changes of the broker node registry.
+#
+# 3. Partition owner registry:
+# /consumers/[group_id]/owner/[topic]/[broker_id-partition_id] --> consumer_node_id
+# This stores the mapping before broker partitions and consumers. Each partition is owned by a unique consumer
+# within a consumer group. The mapping is reestablished after each rebalancing.
+#
+# 4. Consumer offset tracking:
+# /consumers/[group_id]/offsets/[topic]/[broker_id-partition_id] --> offset_counter_value
+# Each consumer tracks the offset of the latest message consumed for each partition.
+#
+#/
+
+require 'zk'
+require 'json'
+
+module Kafka
+  class ConsumerRegistry
+    def initialize(connect='localhost:9092')
+      @connect = connect
+    end
+
+    def register(group_id, consumer_id, topic)
+      dir = "/consumers/#{group_id}/ids"
+      keeper.mkdir_p(dir)
+
+      keeper.create(
+        "#{dir}/#{consumer_id}", 
+        {topic => 1}.to_json, 
+        {:ephemeral => true}
+      )
+    end
+
+    def keeper
+      @keeper = ZK.new(@connect) if @keeper.nil?
+      @keeper
+    end
+  end
+end

spec/consumer_spec.rb

@@ -64,6 +64,17 @@
       Consumer::MAX_SIZE.should eql(1048576)
       @consumer.max_size.should eql(1048576)
     end
+
+    it "should generate a consumer id on initialize" do
+      Time.stub!(:now).and_return(Time.new(0))
+      Socket.stub!(:gethostname).and_return("foohost")
+      consumer = Consumer.new({:group_id => "foogroup"})
+      consumer.consumer_id.should eql("foogroup_foohost-0")
+    end
+
+    it "should have a group id" do
+      @consumer.should respond_to(:group_id)
+    end
 
     it "should return the size of the request" do
       @consumer.request_size.should eql(24)

spec/producer_spec.rb

@@ -1,3 +1,5 @@
+# encoding: utf-8
+
 # 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.
@@ -12,7 +14,6 @@
 # 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.
-# encoding: utf-8
 require File.dirname(__FILE__) + '/spec_helper'
 
 describe Producer do
@@ -64,11 +65,13 @@
         message = Kafka::Message.new("ümlaut")
         encoded = @producer.encode(message)
         data = [encoded.size].pack("N") + encoded
+        payload = Kafka::Message.parse_from(data).payload
+
         if RUBY_VERSION[0,3] == "1.8" # Use old iconv on Ruby 1.8 for encoding
           ic = Iconv.new('UTF-8//IGNORE', 'UTF-8')
-          ic.iconv(Kafka::Message.parse_from(data).payload).should eql("ümlaut")
+          ic.iconv(payload).should eql("ümlaut")
         else
-          Kafka::Message.parse_from(data).payload.force_encoding(Encoding::ASCII_8BIT).should eql("ümlaut")
+          payload.force_encoding(Encoding::UTF_8).should eql("ümlaut")
         end
       end
     end