kafka-consumer 0.0.1

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA1:
+  metadata.gz: ad31d35f34d283a73bc8865fc90e1748ecc48ffd
+  data.tar.gz: 58fe93f719be96c3276604973b12887f42eb00f5
+SHA512:
+  metadata.gz: 1ac246b560fe2f8ed24f71d9f32a0323f7d79dbbd1643a52e17fb2e524c027166f87dc12ad136ddcee0343c012fe27932a82ed26513a7184bb030a8fc5d4a64b
+  data.tar.gz: 415dc58012fc48dba46a9d5ba34f6ea955fb78a194e3666c9d78f6e2767fdac76fd5decff989d7b4b52b1e31e054b3cca5b49abe337138ed9230ce3f82a92bdc

data/.gitignore

@@ -0,0 +1,14 @@
+/.bundle/
+/.yardoc
+/Gemfile.lock
+/_yardoc/
+/coverage/
+/doc/
+/pkg/
+/spec/reports/
+/tmp/
+*.bundle
+*.so
+*.o
+*.a
+mkmf.log

data/Gemfile

@@ -0,0 +1,4 @@
+source 'https://rubygems.org'
+gemspec
+
+gem 'snappy'

data/LICENSE.txt

@@ -0,0 +1,22 @@
+Copyright (c) 2015 Willem van Bergen
+
+MIT License
+
+Permission is hereby granted, free of charge, to any person obtaining
+a copy of this software and associated documentation files (the
+"Software"), to deal in the Software without restriction, including
+without limitation the rights to use, copy, modify, merge, publish,
+distribute, sublicense, and/or sell copies of the Software, and to
+permit persons to whom the Software is furnished to do so, subject to
+the following conditions:
+
+The above copyright notice and this permission notice shall be
+included in all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
+EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
+NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE
+LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION
+OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION
+WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.

data/README.md

@@ -0,0 +1,43 @@
+# Kafka::Consumer
+
+High-level Kafka consumer for Ruby. Uses Zookeeper to manage load balancing, failover, and offset management.
+
+A consumer group consists of multiple instances of the same consumer. Every instance registers itself in
+Zookeeper. Based on the number of instances that are registered, an instance will start consuming some or
+all of the partitions of the topics it wants to consume.
+
+The distribution algorithm will make sure that every partition is only consumed by one consumer
+instance at a time. It uses Zookeeper watches to be notified of new consumer instances coming
+online or going offline, which will trigger a redistribition of all the partitions that are consumed.
+
+Periodically, it will commit the last processed offset of every partition to Zookeeper. Whenever a
+new consumer starts, it will resume consumingevery partition at the last committed offset. This implements
+an **at least once guarantee**, so it is possible that you end up consuming the same message more than once.
+It's your responsibility to deal with this if that is a problem for you, e.g. by using idempotent operations.
+
+## Usage
+
+First, add `kafka-consumer` to your **Gemfile**, and run `bundle install.
+If your messages are snappy-compressed, add the `snappy` gem as well.
+
+``` ruby
+zookeeper = "zk1:2181,zk2:2181,zk3:2181"
+name      = "access-log-processor"
+topics    = ["access_log"]
+
+consumer = Kafka::Consumer.new(name, topics, zookeeper: zookeeper)
+
+Signal.trap("INT") { consumer.interrupt }
+
+consumer.each do |message|
+  # process message
+end
+```
+
+## Contributing
+
+1. Fork it ( https://github.com/wvanbergen/kafka-consumer/fork )
+2. Create your feature branch (`git checkout -b my-new-feature`)
+3. Commit your changes (`git commit -am 'Add some feature'`)
+4. Push to the branch (`git push origin my-new-feature`)
+5. Create a new Pull Request

data/Rakefile

@@ -0,0 +1,43 @@
+require "bundler/gem_tasks"
+require "kafka/consumer"
+require "rake/testtask"
+require "benchmark"
+
+Rake::TestTask.new do |t|
+  t.libs = ["lib", "test"]
+  t.test_files = FileList['test/*_test.rb']
+end
+
+namespace :kafka do
+  task :consumer do
+    zookeeper = ENV["ZOOKEEPER"] or raise "Specify the ZOOKEEPER connection string."
+    name      = ENV["NAME"]      or raise "Specify NAME to name the consumergroup."
+    topics    = ENV["TOPICS"]    or raise "Specify the TOPICS you want to consume. Use comma as separator."
+
+    consumer = Kafka::Consumer.new(name, topics.split(','), zookeeper: zookeeper, initial_offset: :earliest_offset)
+
+    Signal.trap("TERM") { puts "TERM received"; consumer.interrupt }
+    Signal.trap("INT")  { puts "INT received";  consumer.interrupt }
+
+    counter = 0
+    duration = Benchmark.realtime do
+      consumer.each do |event|
+        counter += 1
+        print "Consumed #{counter} messages.\n" if counter % 1000 == 0
+      end
+    end
+
+    puts
+    puts "%d messages consumed in %0.3fs (%0.3f msg/s)" % [counter, duration, counter.to_f / duration]
+  end
+
+  namespace :consumer do
+    task :reset do
+      zookeeper = ENV["ZOOKEEPER"] or raise "Specify the ZOOKEEPER connection string."
+      name      = ENV["NAME"]      or raise "Specify NAME to name the consumergroup."
+
+      consumer = Kafka::Consumer.new(name, [], zookeeper: zookeeper)
+      consumer.group.reset_offsets
+    end
+  end
+end

data/bin/kazoo

@@ -0,0 +1,14 @@
+#!/usr/bin/env ruby
+$LOAD_PATH.unshift(File.expand_path('../../lib', __FILE__))
+require 'kazoo/cli'
+
+begin
+  ENV["THOR_DEBUG"] = "1"
+  Kazoo::CLI.start(ARGV)
+rescue Thor::UndefinedCommandError, Thor::UnknownArgumentError, Thor::AmbiguousCommandError, Thor::InvocationError => e
+  $stderr.puts(e.message)
+  exit(64)
+rescue Thor::Error => e
+  $stderr.puts(e.message)
+  exit(1)
+end

data/kafka-consumer.gemspec

@@ -0,0 +1,28 @@
+# coding: utf-8
+lib = File.expand_path('../lib', __FILE__)
+$LOAD_PATH.unshift(lib) unless $LOAD_PATH.include?(lib)
+require 'kafka/consumer/version'
+
+Gem::Specification.new do |spec|
+  spec.name          = "kafka-consumer"
+  spec.version       = Kafka::Consumer::VERSION
+  spec.authors       = ["Willem van Bergen"]
+  spec.email         = ["willem@vanbergen.org"]
+  spec.summary       = %q{High-level consumer for Kafka}
+  spec.description   = %q{High-level consumer for Kafka. Implements the Zookeeper-backed consumer implementation that offers offset management, load balancing and automatic failovers.}
+  spec.homepage      = "https://github.com/wvanbergen/kafka-consumer"
+  spec.license       = "MIT"
+
+  spec.files         = `git ls-files -z`.split("\x0")
+  spec.executables   = spec.files.grep(%r{^bin/}) { |f| File.basename(f) }
+  spec.test_files    = spec.files.grep(%r{^(test|spec|features)/})
+  spec.require_paths = ["lib"]
+
+  spec.add_development_dependency "bundler", "~> 1.7"
+  spec.add_development_dependency "rake", "~> 10.0"
+  spec.add_development_dependency "minitest", "~> 5.0"
+  spec.add_development_dependency "mocha", "~> 1.0"
+
+  spec.add_runtime_dependency "poseidon", "~> 0.0.5"
+  spec.add_runtime_dependency "zookeeper", "~> 1.4"
+end

data/lib/kafka/consumer/message.rb

@@ -0,0 +1,12 @@
+module Kafka
+  class Consumer
+    class Message
+      attr_reader :topic, :partition, :offset, :key, :value
+
+      def initialize(topic, partition, fetched_message)
+        @topic, @partition = topic, partition
+        @key, @value, @offset = fetched_message.key, fetched_message.value, fetched_message.offset
+      end
+    end
+  end
+end

data/lib/kafka/consumer/partition_consumer.rb

@@ -0,0 +1,146 @@
+module Kafka
+  class Consumer
+    class PartitionConsumer
+
+      attr_reader :consumer, :partition, :handler, :max_wait_ms, :initial_offset,
+                  :commit_interval, :last_processed_offset, :last_committed_offset
+
+      def initialize(consumer, partition, handler: nil, max_wait_ms: 100, initial_offset: :latest_offset, commit_interval: 5.0)
+        @consumer, @partition, @handler = consumer, partition, handler
+        @initial_offset, @max_wait_ms, @commit_interval = initial_offset, max_wait_ms, commit_interval
+
+        @commit_mutex = Mutex.new
+
+        @consumer_thread = Thread.new do
+          Thread.current.abort_on_exception = true
+          manage_partition_consumer
+        end
+
+        Thread.new do
+          Thread.current.abort_on_exception = true
+          background_committer
+        end
+      end
+
+      def wait
+        @consumer_thread.join if @consumer_thread.alive?
+      end
+
+      def interrupt
+        @consumer_thread[:interrupted] = true
+        continue
+      end
+
+      def interrupted?
+        @consumer_thread[:interrupted]
+      end
+
+      def stop
+        interrupt
+        wait
+        consumer.logger.info "Consumer for #{partition.topic.name}/#{partition.id} stopped."
+      end
+
+      def continue
+        @consumer_thread.run if @consumer_thread.status == 'sleep'
+      end
+
+      def claim_partition
+        consumer.logger.info "Claiming partition #{partition.topic.name}/#{partition.id}..."
+        begin
+          other_instance, change = consumer.group.watch_partition_claim(partition) { continue }
+          if other_instance.nil?
+            consumer.instance.claim_partition(partition)
+          elsif other_instance == consumer.instance
+            raise Kazoo::Error, "Already claimed this partition myself. That should not happen"
+          else
+            consumer.logger.warn "Partition #{partition.topic.name}/#{partition.id} is still claimed by instance #{other_instance.id}. Waiting for the claim to be released..."
+            Thread.stop unless change.completed?
+
+            return false if interrupted?
+            raise Kazoo::PartitionAlreadyClaimed
+          end
+        rescue Kazoo::PartitionAlreadyClaimed
+          retry unless interrupted?
+        end
+
+        true
+      end
+
+      def commit_last_offset
+        @commit_mutex.synchronize do
+          if last_processed_offset && (last_committed_offset.nil? || last_committed_offset < last_processed_offset)
+            consumer.group.commit_offset(partition, last_processed_offset)
+            @last_committed_offset = last_processed_offset + 1
+          end
+        end
+      end
+
+      def background_committer
+        until interrupted?
+          commit_last_offset
+          sleep(commit_interval)
+        end
+      end
+
+      def manage_partition_consumer
+        # First, we will try to claim the partition in Zookeeper to ensure there's
+        # only one consumer for it simultaneously.
+        if claim_partition
+          @last_committed_offset = consumer.group.retrieve_offset(partition)
+          case start_offset = last_committed_offset || initial_offset
+          when :earliest_offset, -2
+            consumer.logger.info "Starting consumer for #{partition.topic.name}/#{partition.id} at the earliest available offset..."
+          when :latest_offset, -1
+            consumer.logger.info "Starting consumer for #{partition.topic.name}/#{partition.id} for new messages..."
+          else
+            consumer.logger.info "Starting consumer for #{partition.topic.name}/#{partition.id} at offset #{start_offset}..."
+          end
+
+          begin
+            pc = Poseidon::PartitionConsumer.consumer_for_partition(
+                    consumer.group.name,
+                    consumer.cluster.brokers.values.map(&:addr),
+                    partition.topic.name,
+                    partition.id,
+                    start_offset
+            )
+
+            until interrupted?
+              pc.fetch(max_wait_ms: max_wait_ms).each do |message|
+                message = Message.new(partition.topic.name, partition.id, message)
+                handler.call(message)
+                @last_processed_offset = message.offset
+              end
+            end
+
+          rescue Poseidon::Errors::OffsetOutOfRange
+            pc.close
+
+            consumer.logger.warn "Offset #{start_offset} is no longer available for #{partition.topic.name}/#{partition.id}!"
+            case initial_offset
+            when :earliest_offset, -2
+              consumer.logger.warn "Instead, start consuming #{partition.topic.name}/#{partition.id} at the earliest available offset."
+            when :latest_offset, -1
+              consumer.logger.warn "Instead, start onsuming #{partition.topic.name}/#{partition.id} for new messages only."
+            end
+
+            start_offset = initial_offset
+            retry
+
+          ensure
+            consumer.logger.debug "Stopping consumer for #{partition.topic.name}/#{partition.id}..."
+            pc.close
+          end
+
+
+          commit_last_offset
+          consumer.logger.info "Committed offset #{last_committed_offset - 1} for #{partition.topic.name}/#{partition.id}..." if last_committed_offset
+
+          consumer.instance.release_partition(partition)
+          consumer.logger.debug "Released claim for partition #{partition.topic.name}/#{partition.id}."
+        end
+      end
+    end
+  end
+end

data/lib/kafka/consumer/version.rb

@@ -0,0 +1,5 @@
+module Kafka
+  class Consumer
+    VERSION = "0.0.1"
+  end
+end

data/lib/kafka/consumer.rb

@@ -0,0 +1,184 @@
+require "kazoo"
+require "poseidon"
+require "thread"
+require "logger"
+
+require "kafka/consumer/message"
+require "kafka/consumer/partition_consumer"
+require "kafka/consumer/version"
+
+module Kafka
+  class Consumer
+    BACKPRESSURE_MESSAGE_LIMIT = 1000
+
+    include Enumerable
+
+    attr_reader :subscription,
+        :cluster, :group, :instance,
+        :max_wait_ms, :initial_offset,
+        :logger
+
+    def initialize(name, subscription, zookeeper: [], chroot: '', max_wait_ms: 200, initial_offset: :latest_offset, logger: nil)
+      @name, @subscription = name, subscription
+      @max_wait_ms, @initial_offset = max_wait_ms, initial_offset
+      @logger = logger || Logger.new($stdout)
+
+      @cluster = Kazoo::Cluster.new(zookeeper, chroot: chroot)
+      @group = Kazoo::Consumergroup.new(@cluster, name)
+      @group.create unless @group.exists?
+
+      @instance = @group.instantiate
+      @instance.register(topics)
+    end
+
+    def name
+      group.name
+    end
+
+    def id
+      instance.id
+    end
+
+    def topics
+      @topics ||= begin
+        topic_names = Array(subscription)
+        topic_names.map { |topic_name| cluster.topics.fetch(topic_name) }
+      end
+    end
+
+    def partitions
+      topics.flat_map(&:partitions).sort_by { |partition| [partition.leader.id, partition.topic.name, partition.id] }
+    end
+
+    def interrupt
+      Thread.new do
+        Thread.current.abort_on_exception = true
+
+        logger.info "Stopping partition consumers..."
+        @consumer_manager[:interrupted] = true
+
+        # Make sure to wake up the manager thread, so it can shut down
+        continue
+      end
+    end
+
+    def interrupted?
+      @consumer_manager[:interrupted]
+    end
+
+    def stop
+      interrupt
+      wait
+    end
+
+    def wait
+      @consumer_manager.join if @consumer_manager.alive?
+    end
+
+    def dead?
+      @consumer_manager.status == false
+    end
+
+    def each(&block)
+      mutex = Mutex.new
+
+      handler = lambda do |message|
+        mutex.synchronize do
+          block.call(message)
+        end
+      end
+
+      @consumer_manager = Thread.new do
+        Thread.current.abort_on_exception = true
+        manage_partition_consumers(handler)
+      end
+
+      wait
+    end
+
+    def self.distribute_partitions(instances, partitions)
+      return {} if instances.empty?
+      partitions_per_instance = partitions.length.to_f / instances.length.to_f
+
+      partitions.group_by.with_index do |partition, index|
+        instance_index = index.fdiv(partitions_per_instance).floor
+        instances[instance_index]
+      end
+    end
+
+    private
+
+    def continue
+      @consumer_manager.run if @consumer_manager.status == 'sleep'
+    end
+
+    def manage_partition_consumers(handler)
+      logger.info "Registered for #{group.name} as #{instance.id}"
+
+      @partition_consumers = {}
+
+      until interrupted?
+        running_instances, change = group.watch_instances { continue }
+        logger.info "#{running_instances.length} instances have been registered: #{running_instances.map(&:id).join(', ')}."
+
+        # Distribute the partitions over the running instances. Afterwards, we can see
+        # what partitions are assigned to this particular instance. Because all instances
+        # run the same algorithm on the same sorted lists of instances and partitions,
+        # all instances should be in agreement of the distribtion.
+        distributed_partitions = self.class.distribute_partitions(running_instances, partitions)
+        my_partitions = distributed_partitions[@instance]
+
+        logger.info "Claiming #{my_partitions.length} out of #{partitions.length} partitions."
+
+        # based onw hat partitions we should be consuming and the partitions
+        # that we already are consuming, figure out what partition consumers
+        # to stop and start
+        partitions_to_stop  = @partition_consumers.keys - my_partitions
+        partitions_to_start = my_partitions - @partition_consumers.keys
+
+        # Stop the partition consumers we should no longer be running in parallel
+        if partitions_to_stop.length > 0
+          logger.info "Stopping #{partitions_to_stop.length} out of #{@partition_consumers.length} partition consumers."
+
+          threads = []
+          partitions_to_stop.each do |partition|
+            partition_consumer = @partition_consumers.delete(partition)
+            threads << Thread.new { partition_consumer.stop }
+          end
+          threads.each(&:join)
+        end
+
+        # Start all the partition consumers we are missing.
+        if partitions_to_start.length > 0
+          logger.info "Starting #{partitions_to_start.length} new partition consumers."
+
+          partitions_to_start.each do |partition|
+            @partition_consumers[partition] = PartitionConsumer.new(self, partition,
+                max_wait_ms: max_wait_ms, initial_offset: initial_offset, handler: handler)
+          end
+        end
+
+        unless change.completed?
+          logger.debug "Suspended consumer manager thread."
+          Thread.stop
+          logger.debug "Consumer manager thread woke up..."
+        end
+      end
+
+      logger.debug "Consumer interrupted."
+
+      # Stop all running partition consumers
+      threads = []
+      @partition_consumers.each_value do |partition_consumer|
+        threads << Thread.new { partition_consumer.stop }
+      end
+      threads.each(&:join)
+
+      # Deregister the instance. This should trigger a rebalance in all the remaining instances.
+      @instance.deregister
+      logger.debug "Consumer group instance #{instance.id} was deregistered"
+
+      cluster.close
+    end
+  end
+end

data/lib/kazoo/broker.rb

@@ -0,0 +1,68 @@
+module Kazoo
+  class Broker
+    attr_reader :cluster, :id, :host, :port, :jmx_port
+
+    def initialize(cluster, id, host, port, jmx_port: nil)
+      @cluster = cluster
+      @id, @host, @port = id, host, port
+      @jmx_port = jmx_port
+    end
+
+    def led_partitions
+      result, threads, mutex = [], ThreadGroup.new, Mutex.new
+      cluster.partitions.each do |partition|
+        t = Thread.new do
+          select = partition.leader == self
+          mutex.synchronize { result << partition } if select
+        end
+        threads.add(t)
+      end
+      threads.list.each(&:join)
+      result
+    end
+
+    def replicated_partitions
+      result, threads, mutex = [], ThreadGroup.new, Mutex.new
+      cluster.partitions.each do |partition|
+        t = Thread.new do
+          select = partition.replicas.include?(self)
+          mutex.synchronize { result << partition } if select
+        end
+        threads.add(t)
+      end
+      threads.list.each(&:join)
+      result
+    end
+
+    def critical?(replicas: 1)
+      result, threads, mutex = false, ThreadGroup.new, Mutex.new
+      replicated_partitions.each do |partition|
+        t = Thread.new do
+          isr = partition.isr.reject { |r| r == self }
+          mutex.synchronize { result = true if isr.length < replicas }
+        end
+        threads.add(t)
+      end
+      threads.list.each(&:join)
+      result
+    end
+
+    def addr
+      "#{host}:#{port}"
+    end
+
+    def eql?(other)
+      other.is_a?(Kazoo::Broker) && other.cluster == self.cluster && other.id == self.id
+    end
+
+    alias_method :==, :eql?
+
+    def hash
+      [self.cluster, self.id].hash
+    end
+
+    def self.from_json(cluster, id, json)
+      new(cluster, id.to_i, json.fetch('host'), json.fetch('port'), jmx_port: json.fetch('jmx_port', nil))
+    end
+  end
+end

data/lib/kazoo/cli.rb

@@ -0,0 +1,79 @@
+require 'kazoo'
+require 'thor'
+
+module Kazoo
+  class CLI < Thor
+    class_option :zookeeper, :type => :string, :default => ENV['ZOOKEEPER_PEERS']
+    class_option :chroot, :type => :string, :default => ""
+
+    desc "cluster", "Describes the Kafka cluster as registered in Zookeeper"
+    def cluster
+      validate_class_options!
+
+      kafka_cluster.brokers.values.sort_by(&:id).each do |broker|
+        $stdout.puts "#{broker.id}:\t#{broker.addr}\t(hosts #{broker.replicated_partitions.length} partitions, leads #{broker.led_partitions.length})"
+      end
+    end
+
+    desc "topics", "Lists all topics in the cluster"
+    def topics
+      validate_class_options!
+
+      kafka_cluster.topics.values.sort_by(&:name).each do |topic|
+        $stdout.puts topic.name
+      end
+    end
+
+    option :topic, :type => :string
+    desc "partitions", "Lists partitions"
+    def partitions
+      validate_class_options!
+
+      topics = kafka_cluster.topics.values
+      topics.select! { |t| t.name == options[:topic] } if options[:topic]
+      topics.sort_by!(&:name)
+
+      topics.each do |topic|
+        topic.partitions.each do |partition|
+          $stdout.puts "#{partition.topic.name}/#{partition.id}\tReplicas: #{partition.replicas.map(&:id).join(",")}"
+        end
+      end
+    end
+
+    option :replicas, :type => :numeric, :default => 1
+    desc "critical <broker>", "Determine whether a broker is critical"
+    def critical(broker_name)
+      validate_class_options!
+
+      if broker(broker_name).critical?(replicas: options[:replicas])
+        raise Thor::Error, "WARNING: broker #{broker_name} is critical and cannot be stopped safely!"
+      else
+        $stdout.puts "Broker #{broker_name} is non-critical and can be stopped safely."
+      end
+    end
+
+
+    private
+
+    def validate_class_options!
+      if options[:zookeeper].nil? || options[:zookeeper] == ''
+        raise Thor::InvocationError, "Please supply --zookeeper argument, or set the ZOOKEEPER_PEERS environment variable"
+      end
+    end
+
+    def broker(name_or_id)
+      broker = if name_or_id =~ /\A\d+\z/
+        kafka_cluster.brokers[name_or_id.to_i]
+      else
+        kafka_cluster.brokers.values.detect { |b| b.addr == name_or_id } || cluster.brokers.values.detect { |b| b.host == name_or_id }
+      end
+
+      raise Thor::InvocationError, "Broker #{name_or_id.inspect} not found!" if broker.nil?
+      broker
+    end
+
+    def kafka_cluster
+      @kafka_cluster ||= Kazoo::Cluster.new(options[:zookeeper], chroot: options[:chroot])
+    end
+  end
+end