kafkat-onfocusio 0.3.1

data/lib/kafkat/command/set-replication-factor.rb

@@ -0,0 +1,173 @@
+module Kafkat
+  module Command
+
+    #
+    # Command to set the replication factor (RF) of a topic.
+    # The command accepts the topic name, the new desired replication factor,
+    # and, in case of an increase of the replication factor, a list of broker ids.
+    #
+    # When reducing the RF, a new partition assigment will be generated by
+    # removing the last partition replica of every replica set. The leader partition
+    # will not be removed from the replica set to prevent a leader election.
+    #
+    # When increasing the RF, a new partition assignment will be generated by allocating
+    # a new replica to every replica set. The new replica will be assigned to the provided
+    # broker list in a round robin fashion.
+    # If no broker id is specified on the command line, all brokers will be used.
+    #
+    #
+    class SetReplicationFactor < Base
+      register_as 'set-replication-factor'
+
+      usage 'set-replication-factor [topic] [--newrf <n>] [--brokers id[,id]]',
+            'Set the replication factor of'
+
+      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, "the comma-separated list of broker the new partitions must be assigned to", type: :string
+          opt :newrf, "the new replication factor", type: :integer, required: true
+        end
+
+        broker_ids = opts[:brokers] && opts[:brokers].split(',').map(&:to_i)
+        new_rf = opts[:newrf]
+
+        if new_rf < 1
+          puts "ERROR: replication factor is smaller than 1"
+          exit 1
+        end
+
+        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)
+            puts "ERROR: Broker #{id} is not currently active.\n"
+            exit 1
+          end
+        end
+
+        broker_count = broker_ids.size
+        if new_rf > broker_count
+          puts "ERROR: Replication factor is larger than number of brokers.\n"
+          exit 1
+        end
+
+        assignments = []
+        topics.each do |_, t|
+          current_rf = t.partitions[0].replicas.size
+          if new_rf < current_rf
+            warn_reduce_brokers if opts[:brokers]
+            assignments += reduce_rf(t, current_rf, new_rf)
+          elsif new_rf > current_rf
+            assignments += increase_rf(t, current_rf, new_rf, broker_ids)
+          end
+        end
+
+        # ****************
+        if assignments.empty?
+          puts "No partition reassignment required"
+        else
+          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
+
+
+      #
+      # For every partition, remove the last replica from the replica list.
+      # If the last replica is the leader, then the previous replica is removed instead.
+      #
+      def reduce_rf(topic, current_rf, new_rf)
+        delta_rf = current_rf - new_rf
+        if current_rf == 1
+          raise 'Current replication factor if 1. Cannot reduce further.'
+        end
+        unless delta_rf > 0
+          raise "New replication factor (#{new_rf}) must be smaller than current replication factor (#{current_rf})"
+        end
+        assignments = []
+        topic.partitions.map do |p|
+          new_replicas = p.replicas
+
+          (0...delta_rf).each do |_|
+            (0...new_replicas.size).each do |i|
+              if new_replicas[new_replicas.size-1-i] != p.leader
+                new_replicas.delete_at(new_replicas.size-1-i)
+                break
+              end
+            end
+          end
+
+          if new_replicas.size != new_rf
+            raise 'Unexpected state'
+          end
+          assignments << Assignment.new(topic.name, p.id, new_replicas)
+        end
+        assignments
+      end
+
+
+      #
+      # For every partition, filter out the brokers that already have a replica for this partition,
+      # then pick (new_rf - current_rf) brokers and assign them new replicas.
+      #
+      # The count of new replicas assigned to the brokers is maintained in order to uniformly assign new replicas.
+      #
+      def increase_rf(topic, current_rf, new_rf, brokers)
+        unless new_rf > current_rf
+          raise 'New replication factor must be greater than the current replication factor'
+        end
+
+        delta_rf = new_rf - current_rf
+        if delta_rf > brokers.size
+          raise "#{delta_rf} new replicas requested for topic #{p.topic_name} but only #{brokers.size} brokers available"
+        end
+
+        broker_counts = brokers.map { |b| {:id => b, :count => 0} }
+
+        assignments = []
+        topic.partitions.map do |p|
+          existing_replicas = p.replicas
+          pick_from = broker_counts.reject { |b| existing_replicas.include?(b[:id]) }
+          if delta_rf > pick_from.size
+            raise "Cannot create #{delta_rf} new replicas for partition #{p.topic_name}.#{p.id}, not enough brokers"
+          end
+          new_replicas = pick_from.sort { |a, b| a[:count] <=> b[:count] }[0...delta_rf]
+          new_replicas.each { |b| b[:count] += 1 }
+
+          final_replicas = existing_replicas + new_replicas.map { |b| b[:id] }
+
+          assignments << Assignment.new(topic.name, p.id, final_replicas)
+        end
+        assignments
+      end
+
+      def warn_reduce_brokers
+        return if @did_warn_reduce_brokers
+        puts "When reducing the replication factor the list of specified brokers is ignored."
+        puts "Once the replication factor is set, you can use the reassign command."
+        @did_warn_reduce_brokers = true
+      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 Interface::Admin::ExecutionFailedError
+          print result
+        end
+      end
+    end
+  end
+end

data/lib/kafkat/command/topics.rb

@@ -0,0 +1,16 @@
+module Kafkat
+  module Command
+    class Topics < Base
+      register_as 'topics'
+
+      usage 'topics',
+            'Print all topics.'
+
+      def run
+        topic_names = zookeeper.get_topic_names
+
+        topic_names.each { |name| print_topic_name(name) }
+      end
+    end
+  end
+end

data/lib/kafkat/command/verify-reassign.rb

@@ -0,0 +1,18 @@
+module Kafkat
+  module Command
+    class VerifyReassign < Base
+      register_as 'verify-reassign'
+
+      usage 'verify-reassign reassign_YYYY-MM-DDThh:mm:ssZ.json',
+            'Verify reassignment of partitions.'
+
+      def run
+        file_name = ARGV.shift
+
+        all_brokers = zookeeper.get_brokers
+
+        puts admin.verify_reassign(file_name)
+      end
+    end
+  end
+end

data/lib/kafkat/command/verify-replicas.rb

@@ -0,0 +1,92 @@
+module Kafkat
+  module Command
+    class VerifyReplicas < Base
+      register_as 'verify-replicas'
+
+      usage 'verify-replicas  [--topics] [--broker <id>] [--print-details] [--print-summary]',
+            'Check if all partitions in a topic have same number of replicas.'
+
+      def run
+        opts = Trollop.options do
+          opt :topics, "topic names", type: :string
+          opt :broker, "broker ID", type: :string
+          opt :print_details, "show replica size of mismatched partitions", :default => false
+          opt :print_summary, "show summary of mismatched partitions", :default => false
+        end
+
+        topic_names = opts[:topics]
+        print_details = opts[:print_details]
+        print_summary = opts[:print_summary]
+
+        if topic_names
+          topics_list = topic_names.split(',')
+          topics = zookeeper.get_topics(topics_list)
+        end
+        topics ||= zookeeper.get_topics
+        broker = opts[:broker] && opts[:broker].to_i
+
+        partition_replica_size, partition_replica_size_stat = verify_replicas(broker, topics)
+
+        print_summary = !print_details || print_summary
+        print_mismatched_partitions(partition_replica_size, partition_replica_size_stat, print_details, print_summary)
+      end
+
+      def verify_replicas(broker, topics)
+        partition_replica_size = {}
+        partition_replica_size_stat = {}
+
+        topics.each do |_, t|
+          partition_replica_size[t.name] = {}
+          partition_replica_size_stat[t.name] = {}
+
+          t.partitions.each do |p|
+            replica_size = p.replicas.length
+
+            next if broker && !p.replicas.include?(broker)
+
+            partition_replica_size_stat[t.name][replica_size] ||= 0
+            partition_replica_size_stat[t.name][replica_size] += 1
+
+            partition_replica_size[t.name][p.id] = replica_size
+          end
+
+        end
+
+        return partition_replica_size, partition_replica_size_stat
+      end
+
+      def print_mismatched_partitions(partition_replica_size, partition_replica_size_stat, print_details, print_summary)
+        topic_column_width = partition_replica_size.keys.max_by(&:length).length
+        if print_details
+          printf "%-#{topic_column_width}s %-10s %-15s %-20s\n", "topic", "partition", "replica_size", "replication_factor"
+
+          partition_replica_size.each do |topic_name, partition|
+            replication_factor = partition_replica_size_stat[topic_name].key(partition_replica_size_stat[topic_name].values.max)
+
+            partition.each do |id, replica_size|
+              if replica_size != replication_factor
+                printf "%-#{topic_column_width}s %-10d %-15d %-20d\n", topic_name, id, replica_size, replication_factor
+              end
+            end
+          end
+        end
+
+        if print_summary
+          printf "%-#{topic_column_width}s %-15s %-10s %-15s %-20s\n", "topic", "replica_size", "count", "percentage", "replication_factor"
+          partition_replica_size_stat.each do |topic_name, partition|
+            if partition.values.size > 1
+              replication_factor = partition_replica_size_stat[topic_name].key(partition_replica_size_stat[topic_name].values.max)
+              num_partitions = 0.0
+              partition.each { |key, value| num_partitions += value }
+
+              partition.each do |replica_size, cnt|
+                printf "%-#{topic_column_width}s %-15d %-10d %-15d %-20d\n", topic_name, replica_size, cnt, (cnt * 100 /num_partitions)
+                                                                                           .to_i, replication_factor
+              end
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/kafkat/command.rb

@@ -0,0 +1,70 @@
+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
+      include CommandIO
+      include Kafkat::Logging
+
+      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,50 @@
+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
+    attr_reader :json_files_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']
+      @json_files_path = json['json_files_path']
+      if !@json_files_path || !File.exist?(@json_files_path)
+        raise ArgumentError, "The directory \"json_files_path\": \"#{@json_files_path}\" does not exit."
+      end
+    end
+  end
+end

data/lib/kafkat/interface/admin.rb

@@ -0,0 +1,115 @@
+require 'time'
+
+module Kafkat
+  module Interface
+    class Admin
+      class ExecutionFailedError < StandardError; end
+
+      attr_reader :kafka_path
+      attr_reader :zk_path
+      attr_reader :json_files_path
+
+      def initialize(config)
+        @kafka_path = config.kafka_path
+        @zk_path = config.zk_path
+        @json_files_path = config.json_files_path
+      end
+
+      def elect_leaders!(partitions)
+        file = File.new File.join(@json_files_path, "elect-leaders_#{Time.now.xmlschema}.json"), "w"
+
+        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
+
+        puts "Using JSON file: " + file.path
+
+        run_tool(
+          'kafka-preferred-replica-election',
+            '--path-to-json-file', file.path
+        )
+      end
+
+      def reassign!(assignments)
+        file_name = "reassign_#{Time.now.xmlschema}.json"
+        file = File.new File.join(@json_files_path, file_name), "w"
+
+        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
+
+        puts "Using JSON file: " + file.path
+        puts "Run this command to check the status: kafkat verify-reassign #{file_name}"
+
+        run_tool(
+          'kafka-reassign-partitions',
+            '--execute',
+            '--reassignment-json-file', file.path
+        )
+      end
+
+      def verify_reassign(file_name)
+        file =
+          if File.exist? file_name
+            File.new file_name
+          else
+            File.new File.join(@json_files_path, file_name)
+          end
+
+        puts "Using JSON file: " + file.path
+
+        run_tool(
+          'kafka-reassign-partitions',
+            '--verify',
+            '--reassignment-json-file', file.path
+        )
+      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")
+        # The scripts in the Confluent package does not have .sh extensions
+        if !File.exist? path
+          path = File.join(kafka_path, "bin/#{name}")
+        end
+        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