kafkat-onfocusio 0.3.1

data/lib/kafkat/interface/zookeeper.rb

@@ -0,0 +1,178 @@
+require 'thread'
+require 'zk'
+
+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_broker_ids
+        zk.children(brokers_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_topic_names()
+        return zk.children(topics_path)
+      end
+
+      def get_topics(names=nil)
+        error_msgs = {}
+        topics = {}
+
+        if names == nil
+          pool.with_connection do |cnx|
+            names = cnx.children(topics_path)
+          end
+        end
+
+        threads = names.map do |name|
+          Thread.new do
+            begin
+              topics[name] = get_topic(name)
+            rescue => e
+              error_msgs[name] = e
+            end
+          end
+        end
+        threads.map(&:join)
+
+        unless error_msgs.empty?
+          STDERR.print "ERROR: zk cmds failed on get_topics: \n#{error_msgs.values.join("\n")}\n"
+          exit 1
+        end
+        topics
+      end
+
+      def get_topic(name)
+        partition_queue = Queue.new
+        path1 = topic_path(name)
+        path2 = topic_partitions_path(name)
+
+        partitions = []
+        topic_string = pool.with_connection { |cnx| cnx.get(path1).first }
+        partition_ids = pool.with_connection { |cnx| cnx.children(path2) }
+
+        topic_json = JSON.parse(topic_string)
+
+        threads = partition_ids.map do |id|
+          id = id.to_i
+
+          Thread.new do
+            path3 = topic_partition_state_path(name, id)
+            partition_string = pool.with_connection { |cnx| cnx.get(path3).first }
+            partition_json = JSON.parse(partition_string)
+            replicas = topic_json['partitions'][id.to_s]
+            leader = partition_json['leader']
+            isr = partition_json['isr']
+
+            partition_queue << Partition.new(name, id, replicas, leader, isr)
+          end
+        end
+        threads.map(&:join)
+
+        until partition_queue.empty? do
+          partitions << partition_queue.pop
+        end
+
+        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 pool
+        @pool ||= ZK.new_pool(zk_path, :min_clients => 10, :max_clients => 300, :timeout => 1)
+      end
+
+      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/command_io.rb

@@ -0,0 +1,21 @@
+module Kafkat
+  module CommandIO
+    def prompt_and_execute_assignments(assignments)
+      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 Interface::Admin::ExecutionFailedError
+        print result
+      end
+    end
+  end
+end

data/lib/kafkat/utility/formatting.rb

@@ -0,0 +1,68 @@
+module Kafkat
+  module Formatting
+    def justify(field, width=2)
+      field = field.to_s
+      count = [width - (field.length / 8), 1].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_name(topic_name)
+      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/logging.rb

@@ -0,0 +1,7 @@
+module Kafkat
+  module Logging
+   def print_err(message)
+    STDERR.print message
+   end
+  end
+end

data/lib/kafkat/utility.rb

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

data/lib/kafkat/version.rb

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

data/lib/kafkat.rb

@@ -0,0 +1,14 @@
+require 'zk'
+require 'json'
+require 'trollop'
+require 'retryable'
+require 'highline/import'
+require 'colored'
+
+require 'kafkat/version'
+require 'kafkat/utility'
+require 'kafkat/config'
+require 'kafkat/cluster'
+require 'kafkat/interface'
+require 'kafkat/command'
+require 'kafkat/cli'

data/spec/factories/topic.rb

@@ -0,0 +1,53 @@
+module Kafkat
+  FactoryGirl.define do
+    factory :topic, class:Topic do
+      name "topic_name"
+
+      factory :topic_with_one_empty_broker do
+        partitions {[Partition.new(name, 0, [0], 0, [0]),
+                     Partition.new(name, 1, [1], 1, [1]),
+                     Partition.new(name, 2, [0], 0, [0]),
+                     Partition.new(name, 3, [0], 0, [0]),
+                     Partition.new(name, 4, [1], 1, [1])]}
+      end
+
+      factory :topic_rep_factor_one do
+        partitions {[Partition.new(name, 0, [0], 0, [0]),
+                     Partition.new(name, 1, [1], 1, [1]),
+                     Partition.new(name, 2, [2], 2, [2]),
+                     Partition.new(name, 3, [0], 0, [0]),
+                     Partition.new(name, 4, [1], 1, [1])]}
+      end
+
+      factory :topic_rep_factor_two do
+        partitions {[Partition.new(name, 0, [0, 1], 0, [0]),
+                     Partition.new(name, 1, [0, 2], 2, [2]),
+                     Partition.new(name, 2, [1, 2], 1, [1]),
+                     Partition.new(name, 3, [0, 1], 0, [0]),
+                     Partition.new(name, 4, [0, 2], 2, [2])]}
+      end
+
+      factory :topic_rep_factor_three do
+        partitions {[Partition.new(name, 0, [0, 1, 2], 0, [0]),
+                     Partition.new(name, 1, [0, 1, 2], 1, [1]),
+                     Partition.new(name, 2, [0, 1, 2], 2, [2]),
+                     Partition.new(name, 3, [0, 1, 2], 0, [0]),
+                     Partition.new(name, 4, [0, 1, 2], 1, [1])]}
+      end
+
+      factory :topic_rep_factor_three_with_four_replicas_in_partition1 do
+        name "topic_name1"
+        partitions  [Partition.new("topic_name1", 0, [0, 1, 2], 0, [0]),
+                     Partition.new("topic_name1", 1, [0, 1, 2, 6], 1, [1]),
+                     Partition.new("topic_name1", 2, [0, 1, 2], 2, [2])]
+      end
+
+      factory :topic2_rep_factor_three do
+        name "topic_name2"
+        partitions [Partition.new("topic_name2", 0, [3, 4, 5], 0, [0]),
+                     Partition.new("topic_name2", 1, [3, 4, 5], 0, [0]),
+                     Partition.new("topic_name2", 2, [3, 4, 5], 1, [1])]
+      end
+    end
+  end
+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
+

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

@@ -0,0 +1,59 @@
+require 'spec_helper'
+
+module Kafkat
+  RSpec.describe Command::Drain do
+    let(:drain) { Command::Drain.new({}) }
+    let(:broker_id) { 0 }
+    let(:destination_broker_ids) { [1, 2] }
+
+    context 'three nodes with replication factor 1' do
+      let(:topic_rep_factor_one) { FactoryGirl.build(:topic_rep_factor_one) }
+
+      it 'should put replicas to broker with lowest number of replicas' do
+        assignments = drain.generate_assignments(broker_id,
+                                                 {"topic_name" => topic_rep_factor_one},
+                                                 destination_broker_ids)
+        expect(assignments).to have_exactly(2).items
+        expect(assignments[0].replicas).to eq([2])
+        expect(assignments[1].replicas).to eq([1])
+      end
+    end
+
+    context 'three nodes with replication factor 2' do
+      let(:topic_rep_factor_two) { FactoryGirl.build(:topic_rep_factor_two) }
+      it 'should put replicas to broker with lowest number of replicas' do
+        assignments = drain.generate_assignments(broker_id,
+                                                 {"topic_name" => topic_rep_factor_two},
+                                                 destination_broker_ids)
+        expect(assignments).to have_exactly(4).items
+        expect(assignments[0].replicas).to eq([2, 1])
+        expect(assignments[1].replicas).to eq([1, 2])
+        expect(assignments[2].replicas).to eq([2, 1])
+        expect(assignments[3].replicas).to eq([1, 2])
+      end
+    end
+
+    context 'not enough brokers to keep all replicas' do
+      let(:topic_rep_factor_three) { FactoryGirl.build(:topic_rep_factor_three) }
+
+      it 'should raise SystemExit' do
+        expect do
+          drain.generate_assignments(broker_id,
+                                     {"topic_name" => topic_rep_factor_three},
+                                     destination_broker_ids)
+        end.to raise_error(SystemExit)
+      end
+    end
+
+    context 'one destination broker is empty' do
+      let(:topic_with_one_empty_broker) { FactoryGirl.build(:topic_with_one_empty_broker) }
+      it 'should not raise exception' do
+        expect do
+          drain.generate_assignments(broker_id,
+                                     {"topic_name" => topic_with_one_empty_broker},
+                                     destination_broker_ids)
+        end.not_to raise_error
+      end
+    end
+  end
+end

data/spec/lib/kafkat/command/verify-replicas_spec.rb

@@ -0,0 +1,50 @@
+require 'spec_helper'
+
+module Kafkat
+  RSpec.describe Command::VerifyReplicas do
+    let(:verify_replicas) { Command::VerifyReplicas.new({}) }
+
+    context 'two topics with replication factor 3' do
+      let(:topic_rep_factor_three_with_four_replicas_in_partition1) {
+        FactoryGirl.build(:topic_rep_factor_three_with_four_replicas_in_partition1) }
+      let(:topic2_rep_factor_three) { FactoryGirl.build(:topic2_rep_factor_three) }
+
+      it 'should return empty mismatched partitions for all brokers' do
+        partition_replica_size, partition_replica_size_stat = verify_replicas.verify_replicas(nil,
+                                                                                              {"topic_name2" => topic2_rep_factor_three})
+
+        expect(partition_replica_size).to eq({"topic_name2" => {0 => 3, 1 => 3, 2 => 3}})
+        expect(partition_replica_size_stat).to eq({"topic_name2" => {3 => 3}})
+      end
+
+      it 'should return topic 1 partition 1 for all brokers' do
+        partition_replica_size, partition_replica_size_stat = verify_replicas.verify_replicas(nil,
+                                                                                              {"topic_name1" => topic_rep_factor_three_with_four_replicas_in_partition1,
+                                                                                               "topic_name2" => topic2_rep_factor_three})
+
+        expect(partition_replica_size).to eq({"topic_name1" => {0 => 3, 1 => 4, 2 => 3}, "topic_name2" => {0 => 3, 1 => 3, 2 => 3}})
+        expect(partition_replica_size_stat).to eq({"topic_name1" => {3 => 2, 4 => 1}, "topic_name2" => {3 => 3}})
+      end
+
+      it 'should return topic 1 partition 1 for broker 6' do
+        partition_replica_size, partition_replica_size_stat = verify_replicas.verify_replicas(6,
+                                                                                              {"topic_name1" => topic_rep_factor_three_with_four_replicas_in_partition1,
+                                                                                               "topic_name2" => topic2_rep_factor_three})
+
+        expect(partition_replica_size).to eq({"topic_name1" => {1 => 4}, "topic_name2" => {}})
+        expect(partition_replica_size_stat).to eq({"topic_name1" => {4 => 1}, "topic_name2" => {}})
+      end
+
+      it 'should return empty mismatched partition for broker 3' do
+        partition_replica_size, partition_replica_size_stat = verify_replicas.verify_replicas(3,
+                                                                                              {"topic_name1" => topic_rep_factor_three_with_four_replicas_in_partition1,
+                                                                                               "topic_name2" => topic2_rep_factor_three})
+
+        expect(partition_replica_size).to eq({"topic_name1" => {}, "topic_name2" => {0 => 3, 1 => 3, 2 => 3}})
+        expect(partition_replica_size_stat).to eq({"topic_name1" => {}, "topic_name2" => {3 => 3}})
+      end
+    end
+
+
+  end
+end