kafkat 0.0.9

data/lib/kafkat/command/reassign.rb

@@ -0,0 +1,89 @@
+module Kafkat
+  module Command
+    class Reassign < Base
+      register_as 'reassign'
+
+      usage 'reassign [topic] [--brokers <ids>] [--replicas <n>]',
+            'Begin reassignment of partitions.'
+
+      def run
+        topic_name = ARGV.shift unless ARGV[0] && ARGV[0].start_with?('--')
+
+        all_brokers = zookeeper.get_brokers
+        topics = topic_name && zookeeper.get_topics([topic_name])
+        topics ||= zookeeper.get_topics
+
+        opts = Trollop.options do
+          opt :brokers, "replica set (broker IDs)", type: :string
+          opt :replicas, "number of replicas (count)", type: :integer
+        end
+
+        broker_ids = opts[:brokers] && opts[:brokers].split(',').map(&:to_i)
+        replica_count = opts[:replicas]
+
+        broker_ids ||= zookeeper.get_brokers.values.map(&:id)
+
+        all_brokers_id = all_brokers.values.map(&:id)
+        broker_ids.each do |id|
+          if !all_brokers_id.include?(id)
+            print "ERROR: Broker #{id} is not currently active.\n"
+            exit 1
+          end
+        end
+
+        # *** This logic is duplicated from Kakfa 0.8.1.1 ***
+
+        assignments = []
+        broker_count = broker_ids.size
+
+        topics.each do |_, t|
+          # This is how Kafka's AdminUtils determines these values.
+          partition_count = t.partitions.size
+          topic_replica_count = replica_count || t.partitions[0].replicas.size
+
+          if topic_replica_count > broker_count
+            print "ERROR: Replication factor is larger than brokers.\n"
+            exit 1
+          end
+
+          start_index = Random.rand(broker_count)
+          replica_shift = Random.rand(broker_count)
+
+          t.partitions.each do |p|
+            replica_shift += 1 if p.id > 0 && p.id % broker_count == 0
+            first_replica_index = (p.id + start_index) % broker_count
+
+            replicas = [broker_ids[first_replica_index]]
+
+            (0...topic_replica_count-1).each do |i|
+              shift = 1 + (replica_shift + i) % (broker_count - 1)
+              index = (first_replica_index + shift) % broker_count
+              replicas << broker_ids[index]
+            end
+
+            replicas.reverse!
+            assignments << Assignment.new(t.name, p.id, replicas)
+          end
+        end
+
+        # ****************
+
+        print "This operation executes the following assignments:\n\n"
+        print_assignment_header
+        assignments.each { |a| print_assignment(a) }
+        print "\n"
+
+        return unless agree("Proceed (y/n)?")
+
+        result = nil
+        begin
+          print "\nBeginning.\n"
+          result = admin.reassign!(assignments)
+          print "Started.\n"
+        rescue Admin::ExecutionFailedError
+          print result
+        end
+      end
+    end
+  end
+end

data/lib/kafkat/command/resign-rewrite.rb

@@ -0,0 +1,76 @@
+module Kafkat
+  module Command
+    class ResignForce < Base
+      register_as 'resign-rewrite'
+
+      usage 'resign-rewrite <broker id>',
+            'Forcibly rewrite leaderships to exclude a broker.'
+
+      usage 'resign-rewrite <broker id> --force',
+            'Same as above but proceed if there are no available ISRs.'
+
+      def run
+        broker_id = ARGV[0] && ARGV.shift.to_i
+        if broker_id.nil?
+          puts "You must specify a broker ID."
+          exit 1
+        end
+
+        opts = Trollop.options do
+          opt :force, "force"
+        end
+
+        print "This operation rewrites leaderships in ZK to exclude broker '#{broker_id}'.\n"
+        print "WARNING: This is a last resort. Try the 'shutdown' command first!\n\n".red
+
+        return unless agree("Proceed (y/n)?")
+
+        brokers = zookeeper.get_brokers
+        topics = zookeeper.get_topics
+        force = opts[:force]
+
+        ops = {}
+        topics.each do |_, t|
+          t.partitions.each do |p|
+            next if p.leader != broker_id
+
+            alternates = p.isr.reject { |i| i == broker_id }
+            new_leader_id = alternates.sample
+
+            if !new_leader_id && !force
+              print "Partition #{t.name}-#{p.id} has no other ISRs!\n"
+              exit 1
+            end
+
+            new_leader_id ||= -1
+            ops[p] = new_leader_id
+          end
+        end
+
+        print "\n"
+        print "Summary of the new assignments:\n\n"
+
+        print "Partition\tLeader\n"
+        ops.each do |p, lid|
+          print justify("#{p.topic_name}-#{p.id}")
+          print justify(lid.to_s)
+          print "\n"
+        end
+
+        begin
+          print "\nStarting.\n"
+          ops.each do |p, lid|
+            retryable(tries: 3, on: Interface::Zookeeper::WriteConflictError) do
+              zookeeper.write_leader(p, lid)
+            end
+          end
+        rescue Interface::Zookeeper::WriteConflictError => e
+          print "Failed to update leaderships in ZK. Try re-running.\n\n"
+          exit 1
+        end
+
+        print "Done.\n"
+      end
+    end
+  end
+end

data/lib/kafkat/command/shutdown.rb

@@ -0,0 +1,30 @@
+module Kafkat
+  module Command
+    class Resign < Base
+      register_as 'shutdown'
+
+      usage 'shutdown <broker id>',
+            'Gracefully remove leaderships from a broker (requires JMX).'
+
+      def run
+        broker_id = ARGV[0] && ARGV.shift.to_i
+        if broker_id.nil?
+          puts "You must specify a broker ID."
+          exit 1
+        end
+
+        print "This operation gracefully removes leaderships from broker '#{broker_id}'.\n"
+        return unless agree("Proceed (y/n)?")
+
+        result = nil
+        begin
+          print "\nBeginning shutdown.\n"
+          result = admin.shutdown!(broker_id)
+          print "Started.\n"
+        rescue Admin::ExecutionFailedError
+          print result
+        end
+      end
+    end
+  end
+end

data/lib/kafkat/command/topics.rb

@@ -0,0 +1,15 @@
+module Kafkat
+  module Command
+    class Topics < Base
+      register_as 'topics'
+
+      usage 'topics',
+            'Print all topics.'
+
+      def run
+        ts = zookeeper.get_topics
+        ts.each { |name, t| print_topic(t) }
+      end
+    end
+  end
+end

data/lib/kafkat/command.rb

@@ -0,0 +1,68 @@
+module Kafkat
+  module Command
+    class NotFoundError < StandardError; end
+
+    def self.all
+      @all ||= {}
+    end
+
+    def self.get(name)
+      klass = all[name.downcase]
+      raise NotFoundError if !klass
+      klass
+    end
+
+    class Base
+      include Formatting
+
+      attr_reader :config
+
+      class << self
+        attr_reader :command_name
+      end
+
+      def self.register_as(name)
+        @command_name = name
+        Command.all[name] = self
+      end
+
+      def self.usages
+        @usages ||= []
+      end
+
+      def self.usage(format, description)
+        usages << [format, description]
+      end
+
+      def initialize(config)
+        @config = config
+      end
+
+      def run
+        raise NotImplementedError
+      end
+
+      def admin
+        @admin ||= begin
+          Interface::Admin.new(config)
+        end
+      end
+
+      def zookeeper
+        @zookeeper ||= begin
+          Interface::Zookeeper.new(config)
+        end
+      end
+
+      def kafka_logs
+        @kafka_logs ||= begin
+          Interface::KafkaLogs.new(config)
+        end
+      end
+    end
+  end
+end
+
+# Require all of the commands.
+command_glob = File.expand_path("../command/*.rb", __FILE__)
+Dir[command_glob].each { |f| require f }

data/lib/kafkat/config.rb

@@ -0,0 +1,45 @@
+module Kafkat
+  class Config
+    CONFIG_PATHS = [
+      '~/.kafkatcfg',
+      '/etc/kafkatcfg'
+    ]
+
+    class NotFoundError < StandardError; end
+    class ParseError < StandardError; end
+
+    attr_reader :kafka_path
+    attr_reader :log_path
+    attr_reader :zk_path
+
+    def self.load!
+      string = nil
+      e = nil
+
+      CONFIG_PATHS.each do |rel_path|
+        begin
+          path = File.expand_path(rel_path)
+          string = File.read(path)
+          break
+        rescue => e
+        end
+      end
+
+      raise e if e && string.nil?
+
+      json = JSON.parse(string)
+      self.new(json)
+
+    rescue Errno::ENOENT
+      raise NotFoundError
+    rescue JSON::JSONError
+      raise ParseError
+    end
+
+    def initialize(json)
+      @kafka_path = json['kafka_path']
+      @log_path = json['log_path']
+      @zk_path = json['zk_path']
+    end
+  end
+end

data/lib/kafkat/interface/admin.rb

@@ -0,0 +1,90 @@
+require 'tempfile'
+
+module Kafkat
+  module Interface
+    class Admin
+      class ExecutionFailedError < StandardError; end
+
+      attr_reader :kafka_path
+      attr_reader :zk_path
+
+      def initialize(config)
+        @kafka_path = config.kafka_path
+        @zk_path = config.zk_path
+      end
+
+      def elect_leaders!(partitions)
+        file = Tempfile.new('kafkat-partitions.json')
+
+        json_partitions = []
+        partitions.each do |p|
+          json_partitions << {
+            'topic' => p.topic_name,
+            'partition' => p.id
+          }
+        end
+
+        json = {'partitions' => json_partitions}
+        file.write(JSON.dump(json))
+        file.close
+
+        run_tool(
+          'kafka-preferred-replica-election',
+            '--path-to-json-file', file.path
+        )
+      ensure
+        file.unlink
+      end
+
+      def reassign!(assignments)
+        file = Tempfile.new('kafkat-partitions.json')
+
+        json_partitions = []
+        assignments.each do |a|
+          json_partitions << {
+            'topic' => a.topic_name,
+            'partition' => a.partition_id,
+            'replicas' => a.replicas
+          }
+        end
+
+        json = {
+          'partitions' => json_partitions,
+          'version' => 1
+        }
+
+        file.write(JSON.dump(json))
+        file.close
+
+        run_tool(
+          'kafka-reassign-partitions',
+            '--execute',
+            '--reassignment-json-file', file.path
+        )
+      ensure
+        file.unlink
+      end
+
+      def shutdown!(broker_id, options={})
+        args = ['--broker', broker_id]
+        args += ['--num.retries', options[:retries]] if options[:retries]
+        args += ['--retry.interval.ms', option[:interval]] if options[:interval]
+
+        run_tool(
+          'kafka-run-class',
+            'kafka.admin.ShutdownBroker',
+            *args
+        )
+      end
+
+      def run_tool(name, *args)
+        path = File.join(kafka_path, "bin/#{name}.sh")
+        args += ['--zookeeper', "\"#{zk_path}\""]
+        args_string = args.join(' ')
+        result = `#{path} #{args_string}`
+        raise ExecutionFailedError if $?.to_i > 0
+        result
+      end
+    end
+  end
+end

data/lib/kafkat/interface/kafka_logs.rb

@@ -0,0 +1,54 @@
+module Kafkat
+  module Interface
+    class KafkaLogs
+      UNTRUNCATED_SIZE = 10 * 1024 * 1024 # 1MB
+
+      class NoLogsError < StandardError; end
+      class KafkaRunningError < StandardError; end
+
+      attr_reader :log_path
+
+      def initialize(config)
+        @log_path = config.log_path
+      end
+
+      def clean_indexes!
+        check_exists
+
+        to_remove = []
+        lock_for_write do
+          index_glob = File.join(log_path, '**/*.index')
+          Dir[index_glob].each do |index_path|
+            size = File.size(index_path)
+            to_remove << index_path if size == UNTRUNCATED_SIZE
+          end
+        end
+
+        to_remove.each do |path|
+          print "Removing #{path}.\n"
+          File.unlink(path)
+        end
+
+        to_remove.size
+      end
+
+      private
+
+      def check_exists
+        raise NoLogsError unless File.exists?(log_path)
+      end
+
+      def lock_for_write
+        File.open(lockfile_path, File::CREAT) do |lockfile|
+          locked = lockfile.flock(File::LOCK_EX | File::LOCK_NB)
+          raise KafkaRunningError unless locked
+          yield
+        end
+      end
+
+      def lockfile_path
+        File.join(log_path, '.lock')
+      end
+    end
+  end
+end

data/lib/kafkat/interface/zookeeper.rb

@@ -0,0 +1,147 @@
+require 'thread'
+
+module Kafkat
+  module Interface
+    class Zookeeper
+      class NotFoundError < StandardError; end
+      class WriteConflictError < StandardError; end
+
+      attr_reader :zk_path
+
+      def initialize(config)
+        @zk_path = config.zk_path
+      end
+
+      def get_brokers(ids=nil)
+        brokers = {}
+        ids ||= zk.children(brokers_path)
+
+        threads = ids.map do |id|
+          id = id.to_i
+          Thread.new do
+            begin
+              brokers[id] = get_broker(id)
+            rescue
+            end
+          end
+        end
+        threads.map(&:join)
+
+        brokers
+      end
+
+      def get_broker(id)
+        path = broker_path(id)
+        string = zk.get(path).first
+        json = JSON.parse(string)
+        host, port = json['host'], json['port']
+        Broker.new(id, host, port)
+      rescue ZK::Exceptions::NoNode
+        raise NotFoundError
+      end
+
+      def get_topics(names=nil)
+        topics = {}
+        names ||= zk.children(topics_path)
+
+        threads = names.map do |name|
+          Thread.new do
+            begin
+              topics[name] = get_topic(name)
+            rescue => e
+            end
+          end
+        end
+        threads.map(&:join)
+
+        topics
+      end
+
+      def get_topic(name)
+        path1 = topic_path(name)
+        topic_string = zk.get(path1).first
+        topic_json = JSON.parse(topic_string)
+
+        partitions = []
+        path2 = topic_partitions_path(name)
+
+        threads = zk.children(path2).map do |id|
+          id = id.to_i
+          Thread.new do
+            path3 = topic_partition_state_path(name, id)
+            partition_string = zk.get(path3).first
+            partition_json = JSON.parse(partition_string)
+
+            replicas = topic_json['partitions'][id.to_s]
+            leader = partition_json['leader']
+            isr = partition_json['isr']
+
+            partitions << Partition.new(name, id, replicas, leader, isr)
+          end
+        end
+        threads.map(&:join)
+
+        partitions.sort_by!(&:id)
+        Topic.new(name, partitions)
+      rescue ZK::Exceptions::NoNode
+        raise NotFoundError
+      end
+
+      def get_controller
+        string = zk.get(controller_path).first
+        controller_json = JSON.parse(string)
+        controller_id = controller_json['brokerid']
+        get_broker(controller_id)
+      rescue ZK::Exceptions::NoNode
+        raise NotFoundError
+      end
+
+      def write_leader(partition, broker_id)
+        path = topic_partition_state_path(partition.topic_name, partition.id)
+        string, stat = zk.get(path)
+
+        partition_json = JSON.parse(string)
+        partition_json['leader'] = broker_id
+        new_string = JSON.dump(partition_json)
+
+        unless zk.set(path, new_string, version: stat.version)
+          raise ChangedDuringUpdateError
+        end
+      end
+
+      private
+
+      def zk
+        @zk ||= ZK.new(zk_path)
+      end
+
+      def brokers_path
+        '/brokers/ids'
+      end
+
+      def broker_path(id)
+        "/brokers/ids/#{id}"
+      end
+
+      def topics_path
+        '/brokers/topics'
+      end
+
+      def topic_path(name)
+        "/brokers/topics/#{name}"
+      end
+
+      def topic_partitions_path(name)
+        "/brokers/topics/#{name}/partitions"
+      end
+
+      def topic_partition_state_path(name, id)
+        "/brokers/topics/#{name}/partitions/#{id}/state"
+      end
+
+      def controller_path
+        "/controller"
+      end
+    end
+  end
+end

data/lib/kafkat/interface.rb

@@ -0,0 +1,3 @@
+require 'kafkat/interface/admin'
+require 'kafkat/interface/kafka_logs'
+require 'kafkat/interface/zookeeper'

data/lib/kafkat/utility/formatting.rb

@@ -0,0 +1,63 @@
+module Kafkat
+  module Formatting
+    def justify(field, width=2)
+      field = field.to_s
+      count = [width - (field.length / 8), 0].max
+      field + "\t" * count
+    end
+
+    def print_broker(broker)
+      print justify(broker.id)
+      print justify("#{broker.host}:#{broker.port}")
+      print "\n"
+    end
+
+    def print_broker_header
+      print justify('Broker')
+      print justify('Socket')
+      print "\n"
+    end
+
+    def print_topic(topic)
+      print justify(topic.name)
+      print "\n"
+    end
+
+    def print_topic_header
+      print justify('Topic')
+      print "\n"
+    end
+
+    def print_partition(partition)
+      print justify(partition.topic_name)
+      print justify(partition.id)
+      print justify(partition.leader || 'none')
+      print justify(partition.replicas.inspect, 7)
+      print justify(partition.isr.inspect, 7)
+      print "\n"
+    end
+
+    def print_partition_header
+      print justify('Topic')
+      print justify('Partition')
+      print justify('Leader')
+      print justify('Replicas', 7)
+      print justify('ISRs', 7)
+      print "\n"
+    end
+
+    def print_assignment(assignment)
+      print justify(assignment.topic_name)
+      print justify(assignment.partition_id)
+      print justify(assignment.replicas.inspect)
+      print "\n"
+    end
+
+    def print_assignment_header
+      print justify('Topic')
+      print justify('Partition')
+      print justify('Replicas')
+      print "\n"
+    end
+  end
+end

data/lib/kafkat/utility.rb

@@ -0,0 +1 @@
+require 'kafkat/utility/formatting'

data/lib/kafkat/version.rb

@@ -0,0 +1,3 @@
+module Kafkat
+  VERSION = '0.0.9'
+end