kafkat 0.1.0 → 0.2.0

checksums.yaml

@@ -1,15 +1,15 @@
 ---
 !binary "U0hBMQ==":
   metadata.gz: !binary |-
-    NTY1OWNmMzNlMzg5MGUyYzAzYzIxOWNkODI4ZTQ1NTUzODc2NjQwYg==
+    NzIyMjdlYzlmM2JhZDI1YzAzZWM1ZTU2ZmU4NWVhNjI3NjY1NTkwMQ==
   data.tar.gz: !binary |-
-    OTkwMjE2ZDlhNDZmMmUzNGE2MGU5ZTFhMDliZjBlN2UxNmJhZTQ4Nw==
+    YjNkZDUzZGZlYmIzNGRiYTNkNTBhNjAwYWE2NmI5NzEwMWRmOTVkMw==
 SHA512:
   metadata.gz: !binary |-
-    YmU4NzZkZTJlNTk3NjkwNTcyZTc2YmZhZDg5ZmIwNjBjZGIzMDE1YWI5ZWUx
-    OTRkMDlhMDZiZjFhYTA3NjBkZjZjZWY1M2QzODlmNWQyZjM1YTdjZWIxMDFm
-    ZDQ2M2Y3NmE3YTBjZWYwZDJkMGM2OTdlY2FmYzM4YzIwYmRiZWQ=
+    NGQzMjRmOThmYmE3MDQ5ZDUyYTg4Zjg0ZTI3Zjc0MDRlOGY5NmYwMGNlZmI0
+    ZDlmZWE4OWM3MTI5MGE5NmU0MjU2NGQyYjk5NjNkZDFlZmRlMTllOWRkMDYx
+    NzhhNDcyYjIyZDU2M2E3NzZhYjE1ZGMyODI3ZjM4ZjJhODM5Mjk=
   data.tar.gz: !binary |-
-    YzgwZjZkMjQ2ZTBhNjM4YjFkZGIzMTE2ODlmMDlmODMyZjNlM2ZiMjE0MTgz
-    MWMzOGRhNzJmNDgwNTdhOTRkODBmMzdiODYyZGZlZjgxMTc5YzRmZjMwZTM4
-    ODljNmNmNDZiMGZlN2MyYzMyZmNlZGRmNzdmMzM5ZjcyODRkZDg=
+    MGIwYzM3NWMzMTZhNWJlMzdjY2JiNTRkYzNlMzBiZGQ4YmM2MWQ1MjI4ZWNm
+    Y2ExNGMxNTVjMTQ2Y2QzMWQzODlhZWE2YzljNTkyYTkyMWRiNzBiODkzZDdi
+    NjQ0OGE1MDlmNDYwY2M5NDZjNzJiYWUxOGE0NDdmMjY0ZjhhMjE=

data/CHANGELOG.md

@@ -0,0 +1,6 @@
+# CHANGE LOG
+
+
+
+## 0.2.0
+- add `cluster-restart` command

data/README.md

@@ -11,6 +11,11 @@ Simplified command-line administration for Kafka brokers.
 ## License & Attributions
 This project is released under the Apache License Version 2.0 (APLv2).
 
+## How to release
+
+- update the version number in `lib/kafkat/version.rb`
+- execute `bundle exec rake release`
+
 
 ## Usage
 
@@ -61,3 +66,5 @@ Here's a list of supported commands:
 ## Important Note
 
 The gem needs read/write access to the Kafka log directory for some operations (clean indexes).
+
+

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/drain.rb

@@ -1,6 +1,7 @@
 module Kafkat
   module Command
     class Drain < Base
+
       register_as 'drain'
 
       usage 'drain <broker id> [--topic <t>] [--brokers <ids>]',

data/lib/kafkat/interface/zookeeper.rb

@@ -12,6 +12,10 @@ module Kafkat
         @zk_path = config.zk_path
       end
 
+      def get_broker_ids
+        zk.children(brokers_path)
+      end
+
       def get_brokers(ids=nil)
         brokers = {}
         ids ||= zk.children(brokers_path)

data/lib/kafkat/version.rb

@@ -1,3 +1,3 @@
 module Kafkat
-  VERSION = '0.1.0'
+  VERSION = '0.2.0'
 end

data/spec/lib/kafkat/command/cluster_restart_spec.rb

@@ -0,0 +1,197 @@
+require 'spec_helper'
+
+module Kafkat
+  module ClusterRestart
+
+    describe Kafkat::ClusterRestart do
+      around(:all) do |example|
+        prev_home = ENV['HOME']
+        tmp = Dir.mktmpdir
+        ENV['HOME'] = tmp
+        begin
+          example.run
+        ensure
+          FileUtils.rm_rf tmp
+          ENV['HOME'] = prev_home
+        end
+      end
+
+      describe Session do
+
+        describe '#allBrokersRestarted?' do
+          context 'when some brokers have not been restarted' do
+            let (:session) {
+              Session.new('broker_states' => {'1' => Session::STATE_NOT_RESTARTED, '2' => Session::STATE_RESTARTED})
+            }
+
+            it do
+              expect(session.all_restarted?).to be false
+            end
+          end
+
+          context 'when all brokers have been restarted' do
+            let (:session) {
+              Session.new('broker_states' => {'1' => Session::STATE_RESTARTED, '2' => Session::STATE_RESTARTED})
+            }
+
+            it do
+              expect(session.all_restarted?).to be true
+            end
+          end
+        end
+
+        describe '#update_states!' do
+          let (:session) {
+            Session.new('broker_states' => {'1' => Session::STATE_NOT_RESTARTED, '2' => Session::STATE_RESTARTED})
+          }
+
+          it 'validates state' do
+            expect { session.update_states!('my_state', []) }.to raise_error UnknownStateError
+          end
+
+          it 'validates broker ids' do
+            expect { session.update_states!(Session::STATE_RESTARTED, ['1', '4']) }.to raise_error UnknownBrokerError
+          end
+
+          it 'changes the states' do
+            session.update_states!(Session::STATE_PENDING, ['1'])
+            expect(session.broker_states['1']).to eql(Session::STATE_PENDING)
+          end
+        end
+
+      end
+
+
+      describe Subcommands do
+        let(:p1) { Partition.new('topic1', 'p1', ['1', '2', '3'], '1', 1) }
+        let(:p2) { Partition.new('topic1', 'p2', ['1', '2', '3'], '2', 1) }
+        let(:p3) { Partition.new('topic1', 'p3', ['2', '3', '4'], '3', 1) }
+        let (:topics) {
+          {
+              'topic1' => Topic.new('topic1', [p1, p2, p3])
+          }
+        }
+        let(:zookeeper) { double('zookeeper') }
+        let(:broker_ids) { ['1', '2', '3', '4'] }
+        let(:broker_4) { Broker.new('4', 'i-xxxxxx.inst.aws.airbnb.com', 9092) }
+        let(:session) { Session.from_brokers(broker_ids) }
+
+        describe Subcommands::Next do
+
+          let(:next_command) { Subcommands::Next.new({}) }
+
+          it 'execute next with 4 brokers and 3 partitions' do
+            allow(zookeeper).to receive(:get_broker_ids).and_return(broker_ids)
+            allow(zookeeper).to receive(:get_broker).and_return(broker_4)
+            allow(zookeeper).to receive(:get_topics).and_return(topics)
+            allow(Session).to receive(:exists?).and_return(true)
+            allow(Session).to receive(:load!).and_return(session)
+            allow(session).to receive(:save!)
+            allow(next_command).to receive(:zookeeper).and_return(zookeeper)
+
+            expect(Session).to receive(:load!)
+            expect(session).to receive(:save!)
+
+            next_command.run
+            expect(next_command.session.broker_states['4']).to eq(Session::STATE_PENDING)
+          end
+        end
+
+        describe Subcommands::Good do
+          let(:good_command) { Subcommands::Good.new({}) }
+
+          let(:session){
+            Session.new('broker_states' => {'1' => Session::STATE_PENDING})
+          }
+
+          it 'set one broker to be restarted' do
+            allow(Session).to receive(:exists?).and_return(true)
+            allow(Session).to receive(:load!).and_return(session)
+            allow(session).to receive(:save!)
+
+            # expect(Session).to receive(:load!)
+            expect(session).to receive(:save!)
+            good_command.restart('1')
+            expect(good_command.session.broker_states['1']).to eq(Session::STATE_RESTARTED)
+          end
+        end
+      end
+
+      describe ClusterRestartHelper do
+        let(:p1) { Partition.new('topic1', 'p1', ['1', '2', '3'], '1', 1) }
+        let(:p2) { Partition.new('topic1', 'p2', ['1', '2', '3'], '2', 1) }
+        let(:p3) { Partition.new('topic1', 'p3', ['2', '3', '4'], '3', 1) }
+        let (:topics) {
+          {
+              'topic1' => Topic.new('topic1', [p1, p2, p3])
+          }
+        }
+        let(:zookeeper) { double('zookeeper') }
+        let(:broker_ids) { ['1', '2', '3', '4'] }
+
+        describe '#get_broker_to_leader_partition_mapping' do
+          it 'initialize a new mapping with 4 nodes' do
+            broker_to_partition = ClusterRestartHelper.get_broker_to_leader_partition_mapping(topics)
+
+            expect(broker_to_partition['1']).to eq([p1])
+            expect(broker_to_partition['2']).to eq([p2])
+            expect(broker_to_partition['3']).to eq([p3])
+            expect(broker_to_partition['4']).to eq([])
+          end
+        end
+
+
+        describe '#calculate_costs' do
+          context 'when no restarted brokers' do
+            it do
+              broker_to_partition = ClusterRestartHelper.get_broker_to_leader_partition_mapping(topics)
+              session = Session.from_brokers(broker_ids)
+
+              expect(ClusterRestartHelper.calculate_cost('1', broker_to_partition['1'], session)).to eq(3)
+              expect(ClusterRestartHelper.calculate_cost('2', broker_to_partition['2'], session)).to eq(3)
+              expect(ClusterRestartHelper.calculate_cost('3', broker_to_partition['3'], session)).to eq(3)
+              expect(ClusterRestartHelper.calculate_cost('4', broker_to_partition['4'], session)).to eq(0)
+            end
+          end
+
+          context 'when one broker has been restarted' do
+            it do
+              broker_to_partition = ClusterRestartHelper.get_broker_to_leader_partition_mapping(topics)
+              session = Session.from_brokers(broker_ids)
+              session.update_states!(Session::STATE_RESTARTED, ['4'])
+
+              expect(ClusterRestartHelper.calculate_cost('1', broker_to_partition['1'], session)).to eq(3)
+              expect(ClusterRestartHelper.calculate_cost('2', broker_to_partition['2'], session)).to eq(3)
+              expect(ClusterRestartHelper.calculate_cost('3', broker_to_partition['3'], session)).to eq(2)
+              expect(ClusterRestartHelper.calculate_cost('4', broker_to_partition['4'], session)).to eq(0)
+            end
+          end
+        end
+
+        describe '#select_broker_with_min_cost' do
+          context 'when no restarted brokers' do
+            it do
+              session = Session.from_brokers(broker_ids)
+
+              broker_id, cost = ClusterRestartHelper.select_broker_with_min_cost(session, topics)
+              expect(broker_id).to eq('4')
+              expect(cost).to eq(0)
+            end
+          end
+
+          context 'when next selection after one broker is restarted' do
+            it do
+              session = Session.from_brokers(broker_ids)
+              session.update_states!(Session::STATE_RESTARTED, ['4'])
+
+              broker_id, cost = ClusterRestartHelper.select_broker_with_min_cost(session, topics)
+              expect(broker_id).to eq('3')
+              expect(cost).to eq(2)
+            end
+          end
+        end
+      end
+    end
+  end
+end
+

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: kafkat
 version: !ruby/object:Gem::Version
-  version: 0.1.0
+  version: 0.2.0
 platform: ruby
 authors:
 - Nelson Gauthier
 autorequire: 
 bindir: bin
 cert_chain: []
-date: 2016-07-11 00:00:00.000000000 Z
+date: 2016-10-19 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: zk
@@ -214,6 +214,7 @@ files:
 - .rspec
 - .simplecov
 - .travis.yml
+- CHANGELOG.md
 - Gemfile
 - LICENSE.txt
 - README.md
@@ -231,6 +232,7 @@ files:
 - lib/kafkat/command.rb
 - lib/kafkat/command/brokers.rb
 - lib/kafkat/command/clean-indexes.rb
+- lib/kafkat/command/cluster_restart.rb
 - lib/kafkat/command/controller.rb
 - lib/kafkat/command/drain.rb
 - lib/kafkat/command/elect-leaders.rb
@@ -245,11 +247,13 @@ files:
 - lib/kafkat/interface/admin.rb
 - lib/kafkat/interface/kafka_logs.rb
 - lib/kafkat/interface/zookeeper.rb
+- lib/kafkat/reboot.rb
 - lib/kafkat/utility.rb
 - lib/kafkat/utility/command_io.rb
 - lib/kafkat/utility/formatting.rb
 - lib/kafkat/version.rb
 - spec/factories/topic.rb
+- spec/lib/kafkat/command/cluster_restart_spec.rb
 - spec/lib/kafkat/command/drain_spec.rb
 - spec/spec_helper.rb
 homepage: https://github.com/airbnb/kafkat
@@ -278,5 +282,6 @@ specification_version: 4
 summary: Simplified command-line administration for Kafka brokers
 test_files:
 - spec/factories/topic.rb
+- spec/lib/kafkat/command/cluster_restart_spec.rb
 - spec/lib/kafkat/command/drain_spec.rb
 - spec/spec_helper.rb