kafka-consumer 0.0.1

data/lib/kazoo/cluster.rb

@@ -0,0 +1,82 @@
+module Kazoo
+  class Cluster
+
+    attr_reader :zookeeper, :chroot
+
+    def initialize(zookeeper, chroot: "")
+      @zookeeper, @chroot = zookeeper, chroot
+      @zk_mutex, @brokers_mutex, @topics_mutex, @consumergroups_mutex = Mutex.new, Mutex.new, Mutex.new, Mutex.new
+    end
+
+    def zk
+      @zk_mutex.synchronize do
+        @zk ||= Zookeeper.new(zookeeper)
+      end
+    end
+
+    def brokers
+      @brokers_mutex.synchronize do
+        @brokers ||= begin
+          brokers = zk.get_children(path: node_with_chroot("/brokers/ids"))
+          result, threads, mutex = {}, ThreadGroup.new, Mutex.new
+          brokers.fetch(:children).map do |id|
+            t = Thread.new do
+              broker_info = zk.get(path: node_with_chroot("/brokers/ids/#{id}"))
+              broker = Kazoo::Broker.from_json(self, id, JSON.parse(broker_info.fetch(:data)))
+              mutex.synchronize { result[id.to_i] = broker }
+            end
+            threads.add(t)
+          end
+          threads.list.each(&:join)
+          result
+        end
+      end
+    end
+
+    def consumergroups
+      @consumergroups ||= begin
+        consumers = zk.get_children(path: node_with_chroot("/consumers"))
+        consumers.fetch(:children).map { |name| Kazoo::Consumergroup.new(self, name) }
+      end
+    end
+
+    def topics
+      @topics_mutex.synchronize do
+        @topics ||= begin
+          topics = zk.get_children(path: node_with_chroot("/brokers/topics"))
+          result, threads, mutex = {}, ThreadGroup.new, Mutex.new
+          topics.fetch(:children).each do |name|
+            t = Thread.new do
+              topic_info = zk.get(path: node_with_chroot("/brokers/topics/#{name}"))
+              topic = Kazoo::Topic.from_json(self, name, JSON.parse(topic_info.fetch(:data)))
+              mutex.synchronize { result[name] = topic }
+            end
+            threads.add(t)
+          end
+          threads.list.each(&:join)
+          result
+        end
+      end
+    end
+
+    def partitions
+      topics.values.flat_map(&:partitions)
+    end
+
+    def reset_metadata
+      @topics, @brokers = nil, nil
+    end
+
+    def under_replicated?
+      partitions.any?(&:under_replicated?)
+    end
+
+    def node_with_chroot(path)
+      "#{@chroot}#{path}"
+    end
+
+    def close
+      zk.close
+    end
+  end
+end

data/lib/kazoo/consumergroup.rb

@@ -0,0 +1,229 @@
+module Kazoo
+  class Consumergroup
+    attr_reader :cluster, :name
+
+    def initialize(cluster, name)
+      @cluster, @name = cluster, name
+    end
+
+    def create
+      cluster.zk.create(path: cluster.node_with_chroot("/consumers/#{name}"))
+      cluster.zk.create(path: cluster.node_with_chroot("/consumers/#{name}/ids"))
+      cluster.zk.create(path: cluster.node_with_chroot("/consumers/#{name}/owners"))
+      cluster.zk.create(path: cluster.node_with_chroot("/consumers/#{name}/offsets"))
+    end
+
+    def exists?
+      stat = cluster.zk.stat(path: cluster.node_with_chroot("/consumers/#{name}"))
+      stat.fetch(:stat).exists?
+    end
+
+
+    def instantiate(id: nil)
+      Instance.new(self, id: id)
+    end
+
+    def instances
+      instances = cluster.zk.get_children(path: cluster.node_with_chroot("/consumers/#{name}/ids"))
+      instances.fetch(:children).map { |id| Instance.new(self, id: id) }
+    end
+
+    def watch_instances(&block)
+      cb = Zookeeper::Callbacks::WatcherCallback.create(&block)
+      result = cluster.zk.get_children(
+        path: cluster.node_with_chroot("/consumers/#{name}/ids"),
+        watcher: cb,
+      )
+
+      instances = result.fetch(:children).map { |id| Instance.new(self, id: id) }
+
+      if result.fetch(:rc) != Zookeeper::Constants::ZOK
+        raise Kazoo::Error, "Failed to watch instances. Error code result[:rc]"
+      end
+
+      [instances, cb]
+    end
+
+
+    def watch_partition_claim(partition, &block)
+      cb = Zookeeper::Callbacks::WatcherCallback.create(&block)
+
+      result = cluster.zk.get(
+        path:    cluster.node_with_chroot("/consumers/#{name}/owners/#{partition.topic.name}/#{partition.id}"),
+        watcher: cb
+      )
+
+      case result.fetch(:rc)
+      when Zookeeper::Constants::ZNONODE # Nobody is claiming this partition yet
+        [nil, nil]
+      when Zookeeper::Constants::ZOK
+        [Kazoo::Consumergroup::Instance.new(self, id: result.fetch(:data)), cb]
+      else
+        raise Kazoo::Error, "Failed set watch for partition claim of #{partition.topic.name}/#{partition.id}. Error code: #{result.fetch(:rc)}"
+      end
+    end
+
+    def retrieve_offset(partition)
+      result = cluster.zk.get(path: cluster.node_with_chroot("/consumers/#{name}/offsets/#{partition.topic.name}/#{partition.id}"))
+      case result.fetch(:rc)
+        when Zookeeper::Constants::ZOK;
+          result.fetch(:data).to_i
+        when Zookeeper::Constants::ZNONODE;
+          nil
+        else
+          raise Kazoo::Error, "Failed to retrieve offset for partition #{partition.topic.name}/#{partition.id}. Error code: #{result.fetch(:rc)}"
+      end
+    end
+
+    def commit_offset(partition, offset)
+      result = cluster.zk.set(
+        path: cluster.node_with_chroot("/consumers/#{name}/offsets/#{partition.topic.name}/#{partition.id}"),
+        data: (offset + 1).to_s
+      )
+
+      if result.fetch(:rc) == Zookeeper::Constants::ZNONODE
+        result = cluster.zk.create(path: cluster.node_with_chroot("/consumers/#{name}/offsets/#{partition.topic.name}"))
+        case result.fetch(:rc)
+          when Zookeeper::Constants::ZOK, Zookeeper::Constants::ZNODEEXISTS
+        else
+          raise Kazoo::Error, "Failed to commit offset #{offset} for partition #{partition.topic.name}/#{partition.id}. Error code: #{result.fetch(:rc)}"
+        end
+
+        result = cluster.zk.create(
+          path: cluster.node_with_chroot("/consumers/#{name}/offsets/#{partition.topic.name}/#{partition.id}"),
+          data: (offset + 1).to_s
+        )
+      end
+
+      if result.fetch(:rc) != Zookeeper::Constants::ZOK
+        raise Kazoo::Error, "Failed to commit offset #{offset} for partition #{partition.topic.name}/#{partition.id}. Error code: #{result.fetch(:rc)}"
+      end
+    end
+
+    def reset_offsets
+      result = cluster.zk.get_children(path: cluster.node_with_chroot("/consumers/#{name}/offsets"))
+      raise Kazoo::Error unless result.fetch(:rc) == Zookeeper::Constants::ZOK
+
+      result.fetch(:children).each do |topic|
+        result = cluster.zk.get_children(path: cluster.node_with_chroot("/consumers/#{name}/offsets/#{topic}"))
+        raise Kazoo::Error unless result.fetch(:rc) == Zookeeper::Constants::ZOK
+
+        result.fetch(:children).each do |partition|
+          cluster.zk.delete(path: cluster.node_with_chroot("/consumers/#{name}/offsets/#{topic}/#{partition}"))
+          raise Kazoo::Error unless result.fetch(:rc) == Zookeeper::Constants::ZOK
+        end
+
+        cluster.zk.delete(path: cluster.node_with_chroot("/consumers/#{name}/offsets/#{topic}"))
+        raise Kazoo::Error unless result.fetch(:rc) == Zookeeper::Constants::ZOK
+      end
+    end
+
+    def inspect
+      "#<Kazoo::Consumergroup name=#{name}>"
+    end
+
+    def eql?(other)
+      other.kind_of?(Kazoo::Consumergroup) && cluster == other.cluster && name == other.name
+    end
+
+    alias_method :==, :eql?
+
+    def hash
+      [cluster, name].hash
+    end
+
+    class Instance
+
+      def self.generate_id
+        "#{Socket.gethostname}:#{SecureRandom.uuid}"
+      end
+
+      attr_reader :group, :id
+
+      def initialize(group, id: nil)
+        @group = group
+        @id = id || self.class.generate_id
+      end
+
+      def registered?
+        stat = cluster.zk.stat(path: cluster.node_with_chroot("/consumers/#{group.name}/ids/#{id}"))
+        stat.fetch(:stat).exists?
+      end
+
+      def register(subscription)
+        result = cluster.zk.create(
+          path: cluster.node_with_chroot("/consumers/#{group.name}/ids/#{id}"),
+          ephemeral: true,
+          data: JSON.generate({
+            version: 1,
+            timestamp: Time.now.to_i,
+            pattern: "static",
+            subscription: Hash[*subscription.flat_map { |topic| [topic.name, 1] } ]
+          })
+        )
+
+        if result.fetch(:rc) != Zookeeper::Constants::ZOK
+          raise Kazoo::ConsumerInstanceRegistrationFailed, "Failed to register instance #{id} for consumer group #{group.name}! Error code: #{result.fetch(:rc)}"
+        end
+
+        subscription.each do |topic|
+          stat = cluster.zk.stat(path: cluster.node_with_chroot("/consumers/#{group.name}/owners/#{topic.name}"))
+          unless stat.fetch(:stat).exists?
+            result = cluster.zk.create(path: cluster.node_with_chroot("/consumers/#{group.name}/owners/#{topic.name}"))
+            if result.fetch(:rc) != Zookeeper::Constants::ZOK
+              raise Kazoo::ConsumerInstanceRegistrationFailed, "Failed to register subscription of #{topic.name} for consumer group #{group.name}! Error code: #{result.fetch(:rc)}"
+            end
+          end
+        end
+      end
+
+      def deregister
+        cluster.zk.delete(path: cluster.node_with_chroot("/consumers/#{group.name}/ids/#{id}"))
+      end
+
+      def claim_partition(partition)
+        result = cluster.zk.create(
+          path: cluster.node_with_chroot("/consumers/#{group.name}/owners/#{partition.topic.name}/#{partition.id}"),
+          ephemeral: true,
+          data: id,
+        )
+
+        case result.fetch(:rc)
+        when Zookeeper::Constants::ZOK
+          return true
+        when Zookeeper::Constants::ZNODEEXISTS
+          raise Kazoo::PartitionAlreadyClaimed, "Partition #{partition.topic.name}/#{partition.id} is already claimed!"
+        else
+          raise Kazoo::Error, "Failed to claim partition #{partition.topic.name}/#{partition.id}. Error code: #{result.fetch(:rc)}"
+        end
+      end
+
+      def release_partition(partition)
+        result = cluster.zk.delete(path: cluster.node_with_chroot("/consumers/#{group.name}/owners/#{partition.topic.name}/#{partition.id}"))
+        if result.fetch(:rc) != Zookeeper::Constants::ZOK
+          raise Kazoo::Error, "Failed to release partition #{partition.topic.name}/#{partition.id}. Error code: #{result.fetch(:rc)}"
+        end
+      end
+
+      def inspect
+        "#<Kazoo::Consumergroup::Instance group=#{group.name} id=#{id}>"
+      end
+
+      def hash
+        [group, id].hash
+      end
+
+      def eql?(other)
+        other.kind_of?(Kazoo::Consumergroup::Instance) && group == other.group && id == other.id
+      end
+
+      alias_method :==, :eql?
+
+      private
+
+      def cluster
+        group.cluster
+      end
+    end
+  end
+end

data/lib/kazoo/partition.rb

@@ -0,0 +1,62 @@
+module Kazoo
+  class Partition
+    attr_reader :topic, :id, :replicas
+
+    def initialize(topic, id, replicas: nil)
+      @topic, @id, @replicas = topic, id, replicas
+      @mutex = Mutex.new
+    end
+
+    def cluster
+      topic.cluster
+    end
+
+    def replication_factor
+      replicas.length
+    end
+
+    def leader
+      @mutex.synchronize do
+        refresh_state if @leader.nil?
+        @leader
+      end
+    end
+
+    def isr
+      @mutex.synchronize do
+        refresh_state if @isr.nil?
+        @isr
+      end
+    end
+
+    def under_replicated?
+      isr.length < replication_factor
+    end
+
+    def inspect
+      "#<Kazoo::Partition #{topic.name}/#{id}>"
+    end
+
+    def eql?(other)
+      other.kind_of?(Kazoo::Partition) && topic == other.topic && id == other.id
+    end
+
+    alias_method :==, :eql?
+
+    def hash
+      [topic, id].hash
+    end
+
+    protected
+
+    def refresh_state
+      state_json = cluster.zk.get(path: cluster.node_with_chroot("/brokers/topics/#{topic.name}/partitions/#{id}/state"))
+      set_state(JSON.parse(state_json.fetch(:data)))
+    end
+
+    def set_state(json)
+      @leader = cluster.brokers.fetch(json.fetch('leader'))
+      @isr = json.fetch('isr').map { |r| cluster.brokers.fetch(r) }
+    end
+  end
+end

data/lib/kazoo/topic.rb

@@ -0,0 +1,46 @@
+module Kazoo
+  class Topic
+
+    attr_reader :cluster, :name
+    attr_accessor :partitions
+
+    def initialize(cluster, name, partitions: nil)
+      @cluster, @name, @partitions = cluster, name, partitions
+    end
+
+    def self.from_json(cluster, name, json)
+      topic = new(cluster, name)
+      topic.partitions = json.fetch('partitions').map do |(id, replicas)|
+        topic.partition(id.to_i, replicas: replicas.map { |b| cluster.brokers[b] })
+      end.sort_by(&:id)
+
+      return topic
+    end
+
+    def partition(*args)
+      Kazoo::Partition.new(self, *args)
+    end
+
+    def replication_factor
+      partitions.map(&:replication_factor).min
+    end
+
+    def under_replicated?
+      partitions.any?(:under_replicated?)
+    end
+
+    def inspect
+      "#<Kazoo::Topic #{name}>"
+    end
+
+    def eql?(other)
+      other.kind_of?(Kazoo::Topic) && cluster == other.cluster && name == other.name
+    end
+
+    alias_method :==, :eql?
+
+    def hash
+      [cluster, name].hash
+    end
+  end
+end

data/lib/kazoo/version.rb

@@ -0,0 +1,3 @@
+module Kazoo
+  VERSION = "0.1.0"
+end

data/lib/kazoo.rb

@@ -0,0 +1,19 @@
+require 'zookeeper'
+require 'json'
+require 'thread'
+require 'socket'
+require 'securerandom'
+
+module Kazoo
+  Error = Class.new(StandardError)
+  ConsumerInstanceRegistrationFailed = Class.new(Kazoo::Error)
+  PartitionAlreadyClaimed = Class.new(Kazoo::Error)
+  ReleasePartitionFailure = Class.new(Kazoo::Error)
+end
+
+require 'kazoo/cluster'
+require 'kazoo/broker'
+require 'kazoo/topic'
+require 'kazoo/partition'
+require 'kazoo/consumergroup'
+require 'kazoo/version'

data/test/broker_test.rb

@@ -0,0 +1,45 @@
+require 'test_helper'
+
+class BrokerTest < Minitest::Test
+  include MockCluster
+
+  def setup
+    @cluster = mock_cluster
+  end
+
+  def test_broker_critical?
+    refute @cluster.brokers[1].critical?(replicas: 1), "We have 2 in-sync replicas for everything so we can lose 1."
+    assert @cluster.brokers[2].critical?(replicas: 2), "We only have 2 replicas so we can never lose 2."
+
+    # Simulate losing a broker from the ISR for a partition.
+    # This partition lives on broker 1 and 3
+    @cluster.topics['test.4'].partitions[2].expects(:isr).returns([@cluster.brokers[1]])
+
+    assert @cluster.brokers[1].critical?(replicas: 1), "Final remaining broker for this partition's ISR set, cannot lose"
+    refute @cluster.brokers[2].critical?(replicas: 1), "Not related to the under-replicated partitions"
+    refute @cluster.brokers[3].critical?(replicas: 1), "Already down, so not critical"
+  end
+
+  def test_from_json
+    json_payload = '{"jmx_port":9999,"timestamp":"1431719964125","host":"kafka03.example.com","version":1,"port":9092}'
+    broker = Kazoo::Broker.from_json(mock('cluster'), 3, JSON.parse(json_payload))
+
+    assert_equal 3, broker.id
+    assert_equal 'kafka03.example.com', broker.host
+    assert_equal 9092, broker.port
+    assert_equal 9999, broker.jmx_port
+    assert_equal "kafka03.example.com:9092", broker.addr
+  end
+
+  def test_replicated_partitions
+    assert_equal 3, @cluster.brokers[1].replicated_partitions.length
+    assert_equal 4, @cluster.brokers[2].replicated_partitions.length
+    assert_equal 3, @cluster.brokers[3].replicated_partitions.length
+  end
+
+  def test_led_partitions
+    assert_equal 2, @cluster.brokers[1].led_partitions.length
+    assert_equal 2, @cluster.brokers[2].led_partitions.length
+    assert_equal 1, @cluster.brokers[3].led_partitions.length
+  end
+end

data/test/cluster_test.rb

@@ -0,0 +1,16 @@
+require 'test_helper'
+
+class ClusterTest < Minitest::Test
+  include MockCluster
+
+  def setup
+    @cluster = mock_cluster
+  end
+
+  def test_cluster_under_replicated?
+    refute @cluster.under_replicated?
+
+    @cluster.topics['test.4'].partitions[2].expects(:isr).returns([@cluster.brokers[1]])
+    assert @cluster.under_replicated?
+  end
+end

data/test/partition_distribution_test.rb

@@ -0,0 +1,117 @@
+require 'test_helper'
+
+class PartitionDistributionTest < Minitest::Test
+
+  def setup
+    @cluster = Object.new
+    @group = Kazoo::Consumergroup.new(@cluster, 'test.ruby')
+    @topic = Kazoo::Topic.new(@cluster, 'test.4')
+    @partitions = [
+      Kazoo::Partition.new(@topic, 0),
+      Kazoo::Partition.new(@topic, 1),
+      Kazoo::Partition.new(@topic, 2),
+      Kazoo::Partition.new(@topic, 3),
+    ]
+  end
+
+  def test_single_consumer_gets_everything
+    instances = [Kazoo::Consumergroup::Instance.new(@group, id: '1')]
+    instance_copy = Kazoo::Consumergroup::Instance.new(@group, id: '1')
+
+    distribution = Kafka::Consumer.distribute_partitions(instances, @partitions)
+    assert_equal Set.new(instances), Set.new(distribution.keys)
+    assert_equal Set.new(@partitions), Set.new(distribution.values.flatten)
+
+    assert_equal @partitions, distribution[instance_copy]
+  end
+
+  def test_two_consumers_split_the_load
+    instances = [
+      Kazoo::Consumergroup::Instance.new(@group, id: '1'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '2'),
+    ]
+
+    distribution = Kafka::Consumer.distribute_partitions(instances, @partitions)
+    assert_equal Set.new(instances), Set.new(distribution.keys)
+    assert_equal Set.new(@partitions), Set.new(distribution.values.flatten)
+
+    assert_equal 2, distribution[instances[0]].length
+    assert_equal 2, distribution[instances[1]].length
+  end
+
+
+  def test_three_consumers_split_the_load
+    instances = [
+      Kazoo::Consumergroup::Instance.new(@group, id: '1'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '2'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '3'),
+    ]
+
+    distribution = Kafka::Consumer.distribute_partitions(instances, @partitions)
+    assert_equal Set.new(instances), Set.new(distribution.keys)
+    assert_equal Set.new(@partitions), Set.new(distribution.values.flatten)
+
+    assert_equal 2, distribution[instances[0]].length
+    assert_equal 1, distribution[instances[1]].length
+    assert_equal 1, distribution[instances[2]].length
+  end
+
+  def test_four_consumers_split_the_load
+    instances = [
+      Kazoo::Consumergroup::Instance.new(@group, id: '1'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '2'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '3'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '4'),
+    ]
+
+    distribution = Kafka::Consumer.distribute_partitions(instances, @partitions)
+    assert_equal Set.new(instances), Set.new(distribution.keys)
+    assert_equal Set.new(@partitions), Set.new(distribution.values.flatten)
+
+    assert_equal 1, distribution[instances[0]].length
+    assert_equal 1, distribution[instances[1]].length
+    assert_equal 1, distribution[instances[2]].length
+    assert_equal 1, distribution[instances[3]].length
+  end
+
+  def test_four_consumers_split_the_load_and_one_is_standby
+    instances = [
+      Kazoo::Consumergroup::Instance.new(@group, id: '1'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '2'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '3'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '4'),
+      Kazoo::Consumergroup::Instance.new(@group, id: '5'),
+    ]
+
+    distribution = Kafka::Consumer.distribute_partitions(instances, @partitions)
+    assert_equal Set.new(instances[0..3]), Set.new(distribution.keys)
+    assert_equal Set.new(@partitions), Set.new(distribution.values.flatten)
+
+    assert_equal 1, distribution[instances[0]].length
+    assert_equal 1, distribution[instances[1]].length
+    assert_equal 1, distribution[instances[2]].length
+    assert_equal 1, distribution[instances[3]].length
+
+    assert_nil distribution[instances[4]]
+  end
+
+  def test_assign_everything_with_random_number_of_partitions_and_instances
+    partitions = (0 .. rand(500)).map { |i| Kazoo::Partition.new(@topic, i) }
+    instances = (0 .. rand(100)).map { |i| Kazoo::Consumergroup::Instance.new(@group, id: i.to_s) }
+
+    distribution = Kafka::Consumer.distribute_partitions(instances, partitions)
+    assert_equal [partitions.length, instances.length].min, distribution.keys.length
+    assert_equal Set.new(partitions), Set.new(distribution.values.flatten)
+  end
+
+  def test_assign_zero_partitions
+    instances = [Kazoo::Consumergroup::Instance.new(@group, id: '1')]
+    distribution = Kafka::Consumer.distribute_partitions(instances, [])
+    assert distribution.empty?
+  end
+
+  def test_assign_to_zero_instances
+    distribution = Kafka::Consumer.distribute_partitions([], @partitions)
+    assert distribution.empty?
+  end
+end

data/test/partition_test.rb

@@ -0,0 +1,25 @@
+require 'test_helper'
+
+class PartitionTest < Minitest::Test
+  include MockCluster
+
+  def setup
+    @cluster = mock_cluster
+  end
+
+  def test_replication_factor
+    assert_equal 2, @cluster.topics['test.1'].partitions[0].replication_factor
+  end
+
+  def test_state
+    partition = @cluster.topics['test.1'].partitions[0]
+    partition.unstub(:leader)
+    partition.unstub(:isr)
+
+    json_payload = '{"controller_epoch":157,"leader":1,"version":1,"leader_epoch":8,"isr":[3,2,1]}'
+    @cluster.zk.expects(:get).with(path: "/brokers/topics/test.1/partitions/0/state").returns(data: json_payload)
+
+    assert_equal 1, partition.leader.id
+    assert_equal [3,2,1], partition.isr.map(&:id)
+  end
+end