kafkat-onfocusio 0.3.1

data/lib/kafkat/command/cluster_restart.rb

@@ -0,0 +1,336 @@
+module Kafkat
+  module Command
+    class ClusterRestart < Base
+
+      register_as 'cluster-restart'
+
+      usage 'cluster-restart help', 'Determine the server restart sequence for kafka'
+
+      def run
+        subcommand_name = ARGV.shift || 'help'
+        begin
+          subcommand_class = ['Kafkat', 'ClusterRestart', 'Subcommands', subcommand_name.capitalize].inject(Object) do |mod, class_name|
+            mod.const_get(class_name)
+          end
+          subcommand_class.new(config).run
+        rescue NameError
+          print "ERROR: Unknown command #{subcommand_name}"
+        end
+      end
+    end
+  end
+end
+
+module Kafkat
+  module ClusterRestart
+    module Subcommands
+
+      class Help < ::Kafkat::Command::Base
+        def run
+          puts 'cluster-restart help                Print Help and exit'
+          puts 'cluster-restart reset               Clean up the restart state'
+          puts 'cluster-restart start               Initialize the cluster-restart session for the cluster'
+          puts 'cluster-restart next                Calculate the next broker to restart based on the current state'
+          puts 'cluster-restart good <broker>       Mark this broker as successfully restarted'
+          puts 'cluster-restart log                 Print the state of the brokers'
+          puts 'cluster-restart restore <file>      Start a new session and restore the state defined in that file'
+        end
+      end
+
+      class Start < ::Kafkat::Command::Base
+
+        attr_reader :session
+
+        def run
+          if Session.exists?
+            puts "ERROR: A session is already started"
+            puts "\n[Action] Please run 'next' or 'reset' commands"
+            exit 1
+          end
+
+          print "Starting a new Cluster-Restart session.\n"
+
+          @session = Session.from_zookeepers(zookeeper)
+          @session.save!
+
+          puts "\n[Action] Please run 'next' to select the broker with lowest restarting cost"
+        end
+      end
+
+      class Reset < ::Kafkat::Command::Base
+
+        def run
+          if Session.exists?
+            Session.reset!
+          end
+          puts "Session reset"
+          puts "\n[Action] Please run 'start' to start the session"
+        end
+      end
+
+      class Restore < ::Kafkat::Command::Base
+
+        attr_reader :session
+
+        def run
+          if Session.exists?
+            puts "ERROR: A session is already started"
+            puts "\n[Action] Please run 'next' or 'reset' commands"
+            exit 1
+          end
+
+          file_name = ARGV[0]
+          @session = Session.load!(file_name)
+          @session.save!
+          puts "Session restored"
+          puts "\m[Action] Please run 'next' to select the broker with lowest restarting cost"
+        end
+      end
+
+      class Next < ::Kafkat::Command::Base
+
+        attr_reader :session, :topics
+
+        def run
+          unless Session.exists?
+            puts "ERROR: no session in progress"
+            puts "\n[Action] Please run 'start' command"
+            exit 1
+          end
+
+          @session = Session.load!
+          if @session.all_restarted?
+            puts "All the brokers have been restarted"
+          else
+            pendings = @session.pending_brokers
+            if pendings.size > 1
+              puts "ERROR Illegal state: multiple brokers are in Pending state"
+              exit 1
+            elsif pendings.size == 1
+              next_broker = pendings[0]
+              puts "Broker #{next_broker} is in Pending state"
+            else
+              @topics = zookeeper.get_topics
+              next_broker, cost = ClusterRestartHelper.select_broker_with_min_cost(session, topics)
+              @session.update_states!(Session::STATE_PENDING, [next_broker])
+              @session.save!
+              puts "The next broker is: #{next_broker}"
+            end
+            puts "\n[Action-1] Restart broker #{next_broker} aka #{zookeeper.get_broker(next_broker).host}"
+            puts "\n[Action-2] Run 'good #{next_broker}' to mark it as restarted."
+          end
+        end
+      end
+
+      class Log < ::Kafkat::Command::Base
+
+        attr_reader :session
+
+        def run
+          unless Session.exists?
+            puts "ERROR: no session in progress"
+            puts "\n[Action] Please run 'start' command"
+            exit 1
+          end
+
+          @session = Session.load!
+          puts JSON.pretty_generate(@session.to_h)
+        end
+      end
+
+      class Good < ::Kafkat::Command::Base
+
+        attr_reader :session
+
+        def run
+          unless Session.exists?
+            puts "ERROR: no session in progress"
+            puts "\n[Action] Please run 'start' command"
+            exit 1
+          end
+
+          broker_id = ARGV[0]
+          if broker_id.nil?
+            puts "ERROR You must specify a broker id"
+            exit 1
+          end
+          restart(broker_id)
+          puts "Broker #{broker_id} has been marked as restarted"
+          puts "\n[Action] Please run 'next' to select the broker with lowest restarting cost"
+        end
+
+        def restart(broker_id)
+          @session = Session.load!
+          begin
+            if session.pending?(broker_id)
+              session.update_states!(Session::STATE_RESTARTED, [broker_id])
+              session.save!
+            else
+              puts "ERROR Broker state is #{session.state(broker_id)}"
+              exit 1
+            end
+          rescue UnknownBrokerError => e
+            puts "ERROR #{e.to_s}"
+            exit 1
+          end
+        end
+      end
+    end
+
+    class UnknownBrokerError < StandardError;
+    end
+    class UnknownStateError < StandardError;
+    end
+
+    class ClusterRestartHelper
+
+      def self.select_broker_with_min_cost(session, topics)
+        broker_to_partition = get_broker_to_leader_partition_mapping(topics)
+        broker_restart_cost = Hash.new(0)
+        session.broker_states.each do |broker_id, state|
+          if state == Session::STATE_NOT_RESTARTED
+            current_cost = calculate_cost(broker_id, broker_to_partition[broker_id], session)
+            broker_restart_cost[broker_id] = current_cost if current_cost != nil
+          end
+        end
+
+        # Sort by cost first, and then broker_id
+        broker_restart_cost.min_by { |broker_id, cost| [cost, broker_id] }
+      end
+
+      def self.get_broker_to_leader_partition_mapping(topics)
+        broker_to_partitions = Hash.new { |h, key| h[key] = [] }
+
+        topics.values.flat_map { |topic| topic.partitions }
+            .each do |partition|
+          broker_to_partitions[partition.leader] << partition
+        end
+        broker_to_partitions
+      end
+
+      def self.calculate_cost(broker_id, partitions, session)
+        raise UnknownBrokerError, "Unknown broker #{broker_id}" unless session.broker_states.key?(broker_id)
+        partitions.find_all { |partition| partition.leader == broker_id }
+            .reduce(0) do |cost, partition|
+          cost += partition.replicas.length
+          cost -= partition.replicas.find_all { |replica| session.restarted?(replica) }.size
+          cost
+        end
+      end
+    end
+
+
+    class Session
+
+      SESSION_PATH = '~/kafkat_cluster_restart_session.json'
+      STATE_RESTARTED = 'restarted' # use String instead of Symbol to facilitate JSON ser/deser
+      STATE_NOT_RESTARTED = 'not_restarted'
+      STATE_PENDING = 'pending'
+      STATES= [STATE_NOT_RESTARTED, STATE_RESTARTED, STATE_PENDING]
+
+      class NotFoundError < StandardError;
+      end
+      class ParseError < StandardError;
+      end
+
+      attr_reader :broker_states
+
+      def self.exists?
+        File.file?(File.expand_path(SESSION_PATH))
+      end
+
+      def self.load!(session_file = SESSION_PATH)
+        path = File.expand_path(session_file)
+        string = File.read(path)
+
+        json = JSON.parse(string)
+        self.new(json)
+
+      rescue Errno::ENOENT
+        raise NotFoundError
+      rescue JSON::JSONError
+        raise ParseError
+      end
+
+      def self.reset!(session_file = SESSION_PATH)
+        path = File.expand_path(session_file)
+        File.delete(path)
+      end
+
+      def self.from_zookeepers(zookeeper)
+        broker_ids = zookeeper.get_broker_ids
+        Session.from_brokers(broker_ids)
+      end
+
+      def self.from_brokers(brokers)
+        states = brokers.each_with_object({}) { |id, h| h[id] = STATE_NOT_RESTARTED }
+        Session.new('broker_states' => states)
+      end
+
+      def initialize(data = {})
+        @broker_states = data['broker_states'] || {}
+      end
+
+      def save!(session_file = SESSION_PATH)
+        File.open(File.expand_path(session_file), 'w') do |f|
+          f.puts JSON.pretty_generate(self.to_h)
+        end
+      end
+
+      def update_states!(state, ids)
+        state = state.to_s if state.is_a?(Symbol)
+        unless STATES.include?(state)
+          raise UnknownStateError, "Unknown State #{state}"
+        end
+
+        intersection = ids & broker_states.keys
+        unless intersection == ids
+          raise UnknownBrokerError, "Unknown brokers: #{(ids - intersection).join(', ')}"
+        end
+
+        ids.each { |id| broker_states[id] = state }
+        self
+      end
+
+
+      def state(broker_id)
+        raise UnknownBrokerError, "Unknown broker: #{broker_id}" unless @broker_states.key?(broker_id)
+        broker_states[broker_id]
+      end
+
+      def state?(broker_id, state)
+        raise UnknownBrokerError, "Unknown broker: #{broker_id}" unless @broker_states.key?(broker_id)
+        raise UnknownStateError, "Unknown state: #{state}" unless STATES.include?(state)
+        @broker_states[broker_id] == state
+      end
+
+      def pending?(broker_id)
+        state?(broker_id, STATE_PENDING)
+      end
+
+      def not_restarted?(broker_id)
+        state?(broker_id, STATE_NOT_RESTARTED)
+      end
+
+      def restarted?(broker_id)
+        state?(broker_id, STATE_RESTARTED)
+      end
+
+      def all_restarted?
+        @broker_states.values.all? { |state| state == STATE_RESTARTED }
+      end
+
+      def pending_brokers
+        broker_states.keys.find_all do |broker_id|
+          broker_states[broker_id] == STATE_PENDING
+        end
+      end
+
+      def to_h
+        {
+            :broker_states => broker_states,
+        }
+      end
+    end
+  end
+end

data/lib/kafkat/command/controller.rb

@@ -0,0 +1,18 @@
+module Kafkat
+  module Command
+    class Controller < Base
+      register_as 'controller'
+
+      usage 'controller',
+            'Print the current controller.'
+
+      def run
+        c = zookeeper.get_controller
+        print "The current controller is '#{c.id}' (#{c.host}:#{c.port}).\n"
+      rescue Interface::Zookeeper::NotFoundError
+        print "ERROR: Couldn't determine the current controller.\n"
+        exit 1
+      end
+    end
+  end
+end

data/lib/kafkat/command/drain.rb

@@ -0,0 +1,109 @@
+module Kafkat
+  module Command
+    class Drain < Base
+
+      register_as 'drain'
+
+      usage 'drain <broker id> [--topic <t>] [--brokers <ids>]',
+            'Reassign partitions from a specific broker to destination brokers.'
+
+      # For each partition (of specified topic) on the source broker, the command is to
+      # assign the partition to one of the destination brokers that does not already have
+      # this partition, along with existing brokers to achieve minimal movement of data.
+      # To help distribute data evenly, if there are more than one destination brokers
+      # meet the requirement, the command will always choose the brokers with the lowest
+      # number of partitions of the involving topic.
+      #
+      # In order to find out the broker with lowest number of partitions, the command maintain
+      # a hash table with broker id as key and number of partitions as value. The hash table
+      # will be updated along with assignment.
+      def run
+        source_broker = ARGV[0] && ARGV.shift.to_i
+        if source_broker.nil?
+          puts "You must specify a broker ID."
+          exit 1
+        end
+
+        opts = Trollop.options do
+          opt :brokers, "destination broker IDs", type: :string
+          opt :topic,   "topic name to reassign", type: :string
+        end
+
+        topic_name = opts[:topic]
+        topics = topic_name && zookeeper.get_topics([topic_name])
+        topics ||= zookeeper.get_topics
+
+        destination_brokers = opts[:brokers] && opts[:brokers].split(',').map(&:to_i)
+        destination_brokers ||= zookeeper.get_brokers.values.map(&:id)
+        destination_brokers.delete(source_broker)
+
+        active_brokers = zookeeper.get_brokers.values.map(&:id)
+
+        unless (inactive_brokers = destination_brokers - active_brokers).empty?
+          print "ERROR: Broker #{inactive_brokers} are not currently active.\n"
+          exit 1
+        end
+
+        assignments =
+          generate_assignments(source_broker, topics, destination_brokers)
+
+        print "Num of topics got from zookeeper: #{topics.length}\n"
+        print "Num of partitions in the assignment: #{assignments.size}\n"
+        prompt_and_execute_assignments(assignments)
+      end
+
+      def generate_assignments(source_broker, topics, destination_brokers)
+
+        assignments = []
+        topics.each do |_, t|
+          partitions_by_broker = build_partitions_by_broker(t, destination_brokers)
+
+          t.partitions.each do |p|
+            if p.replicas.include? source_broker
+              replicas_size = p.replicas.length
+              replicas = p.replicas - [source_broker]
+              source_broker_is_leader = p.replicas.first == source_broker
+              potential_broker_ids = destination_brokers - replicas
+              if potential_broker_ids.empty?
+                print "ERROR: Not enough destination brokers to reassign topic \"#{t.name}\".\n"
+                exit 1
+              end
+
+              num_partitions_on_potential_broker =
+                partitions_by_broker.select { |id, _| potential_broker_ids.include? id }
+              assigned_broker_id = num_partitions_on_potential_broker.min_by{ |id, num| num }[0]
+              if source_broker_is_leader
+                replicas.unshift(assigned_broker_id)
+              else
+                replicas << assigned_broker_id
+              end
+              partitions_by_broker[assigned_broker_id] += 1
+
+              if replicas.length != replicas_size
+                STDERR.print "ERROR: Number of replicas changes after reassignment topic: #{t.name}, partition: #{p.id} \n"
+                exit 1
+              end
+
+              assignments << Assignment.new(t.name, p.id, replicas)
+            end
+          end
+        end
+
+        assignments
+      end
+
+      # Build a hash map from broker id to number of partitions on it to facilitate
+      # finding the broker with lowest number of partitions to help balance brokers.
+      def build_partitions_by_broker(topic, destination_brokers)
+        partitions_by_broker = Hash.new(0)
+        destination_brokers.each { |id| partitions_by_broker[id] = 0 }
+        topic.partitions.each do |p|
+          p.replicas.each do |r|
+            partitions_by_broker[r] += 1
+          end
+        end
+        partitions_by_broker
+      end
+    end
+  end
+end

data/lib/kafkat/command/elect-leaders.rb

@@ -0,0 +1,31 @@
+module Kafkat
+  module Command
+    class ElectLeaders < Base
+      register_as 'elect-leaders'
+
+      usage 'elect-leaders [topic]',
+            'Begin election of the preferred leaders.'
+
+      def run
+        topic_name = ARGV[0] && ARGV.shift
+        topic_names = topic_name && [topic_name]
+
+        topics = zookeeper.get_topics(topic_names)
+        partitions = topics.values.map(&:partitions).flatten
+
+        topics_s = topic_name ? "'#{topic_name}'" : "all topics"
+        print "This operation elects the preferred replicas for #{topics_s}.\n"
+        return unless agree("Proceed (y/n)?")
+
+        result = nil
+        begin
+          print "\nBeginning.\n"
+          result = admin.elect_leaders!(partitions)
+          print "Started.\n"
+        rescue Interface::Admin::ExecutionFailedError
+          print result
+        end
+      end
+    end
+  end
+end

data/lib/kafkat/command/partitions.rb

@@ -0,0 +1,50 @@
+module Kafkat
+  module Command
+    class Describe < Base
+      register_as 'partitions'
+
+      usage 'partitions [topic]',
+            'Print partitions by topic.'
+      usage 'partitions [topic] --under-replicated',
+            'Print partitions by topic (only under-replicated).'
+      usage 'partitions [topic] --unavailable',
+            'Print partitions by topic (only unavailable).'
+
+      def run
+        topic_name = ARGV.shift unless ARGV[0] && ARGV[0].start_with?('--')
+        topic_names = topic_name && [topic_name]
+
+        @options = Trollop.options do
+          opt :under_replicated, "only under-replicated"
+          opt :unavailable, "only unavailable"
+        end
+
+        brokers = zookeeper.get_brokers
+        topics = zookeeper.get_topics(topic_names)
+
+        print_partition_header
+        topics.each do |name, t|
+          t.partitions.each do |p|
+            print_partition(p) if selected?(p, brokers)
+          end
+        end
+      end
+
+      private
+
+      def selected?(partition, brokers)
+        return partition.under_replicated? if only_under_replicated?
+        return !partition.has_leader?(brokers) if only_unavailable?
+        true
+      end
+
+      def only_under_replicated?
+        !!@options[:under_replicated]
+      end
+
+      def only_unavailable?
+        !!@options[:unavailable]
+      end
+    end
+  end
+end

data/lib/kafkat/command/reassign.rb

@@ -0,0 +1,80 @@
+module Kafkat
+  module Command
+    class Reassign < Base
+      register_as 'reassign'
+
+      usage 'reassign [topics] [--brokers <ids>] [--replicas <n>]',
+            'Begin reassignment of partitions.'
+
+      def run
+        topic_names = ARGV.shift unless ARGV[0] && ARGV[0].start_with?('--')
+
+        all_brokers = zookeeper.get_brokers
+
+        topics = nil
+        if topic_names
+           topics_list = topic_names.split(',')
+           topics = zookeeper.get_topics(topics_list)
+        end
+        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 (#{topic_replica_count}) is larger than brokers (#{broker_count}).\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
+
+        # ****************
+
+        prompt_and_execute_assignments(assignments)
+      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