ktl 1.0.0-java

data/lib/ktl/cluster_stats_task.rb

@@ -0,0 +1,41 @@
+# encoding: utf-8
+
+module Ktl
+  class ClusterStatsTask
+    def initialize(zk_client, shell)
+      @zk_client = zk_client
+      @shell = shell
+    end
+
+    def execute
+      brokers = @zk_client.brokers
+      partitions = @zk_client.all_partitions
+      topics = extract_topics(partitions)
+      leaders = @zk_client.leader_and_isr_for(partitions)
+      ownership = broker_ownership(leaders)
+      @shell.say 'Cluster status:'
+      @shell.say '  topics: %d (%d partitions)' % [topics.size, partitions.size]
+      @shell.say '  brokers: %d' % [brokers.size]
+      brokers.foreach do |broker|
+        leader_for = ownership[broker.id]
+        share = leader_for.fdiv(partitions.size.to_f) * 100
+        @shell.say '    - %d leader for %d partitions (%.2f %%)' % [broker.id, leader_for, share]
+      end
+    end
+
+    private
+
+    def extract_topics(partitions)
+      partitions.map(proc { |tp| tp.topic }, CanBuildFrom).to_seq
+    end
+
+    def broker_ownership(leaders)
+      result = Hash.new(0)
+      leaders.foreach do |item|
+        leader = item.last.leader_and_isr.leader
+        result[leader] += 1
+      end
+      result
+    end
+  end
+end

data/lib/ktl/command.rb

@@ -0,0 +1,33 @@
+# encoding: utf-8
+
+module Ktl
+  class Command < Thor
+
+    java_import 'java.io.ByteArrayOutputStream'
+
+    private
+
+    def with_zk_client
+      zk_client = ZookeeperClient.new(options.zookeeper).setup
+      yield zk_client
+    rescue => e
+      logger.error '%s (%s)' % [e.message, e.class.name]
+      logger.debug e.backtrace.join($/)
+      $stderr.puts '%s (%s)' % [e.message, e.class.name]
+      $stderr.puts e.backtrace.join($/)
+    ensure
+      zk_client.close if zk_client
+    end
+
+    def logger
+      @logger ||= Logger.new($stdout).tap do |log|
+        log.formatter = ShellFormater.new(shell)
+      end
+    end
+
+    def silence_scala(&block)
+      baos = ByteArrayOutputStream.new
+      Scala::Console.with_out(baos) { block.call }
+    end
+  end
+end

data/lib/ktl/decommission_plan.rb

@@ -0,0 +1,68 @@
+# encoding: utf-8
+
+module Ktl
+  class DecommissionPlan
+    def initialize(zk_client, broker_id)
+      @zk_client = zk_client
+      @broker_id = broker_id
+      @replicas_count = Hash.new(0)
+      @leaders_count = Hash.new(0)
+    end
+
+    def generate
+      plan = Scala::Collection::Map.empty
+      brokers = @zk_client.broker_ids
+      brokers = brokers - @broker_id
+      partitions = @zk_client.all_partitions
+      topics = topics_from(partitions)
+      assignments = @zk_client.replica_assignment_for_topics(topics)
+      count_leaders_and_replicas(assignments)
+      partitions = ScalaEnumerable.new(partitions).sort_by { |tp| tp.topic + tp.partition.to_s }
+      partitions.each do |tp|
+        replicas = assignments[tp]
+        if replicas.contains?(@broker_id)
+          if brokers.size >= replicas.size
+            brokers_diff = ScalaEnumerable.new(brokers.diff(replicas)).sort
+            broker_index = replicas.index_of(@broker_id)
+            new_broker = elect_new_broker(broker_index, brokers_diff)
+            new_replicas = replicas.updated(broker_index, new_broker, CanBuildFrom)
+            plan += Scala::Tuple.new(tp, new_replicas)
+          else
+            raise InsufficientBrokersRemainingError, %(#{brokers.size} remaining brokers, #{replicas.size} replicas needed)
+          end
+        end
+      end
+      plan
+    end
+
+    private
+
+    def elect_new_broker(broker_index, diff)
+      if broker_index.zero?
+        new_broker = diff.min_by { |broker| @leaders_count[broker] }
+        @leaders_count[new_broker] += 1
+      else
+        new_broker = diff.min_by { |broker| @replicas_count[broker] }
+        @replicas_count[new_broker] += 1
+      end
+      new_broker
+    end
+
+    def count_leaders_and_replicas(assignments)
+      assignments.foreach do |assignment|
+        replicas = ScalaEnumerable.new(assignment.last)
+        replicas.each_with_index do |broker, index|
+          if index.zero?
+            @leaders_count[broker] += 1
+          else
+            @replicas_count[broker] += 1
+          end
+        end
+      end
+    end
+
+    def topics_from(partitions)
+      partitions.map(proc { |tp| tp.topic }, CanBuildFrom).to_seq
+    end
+  end
+end

data/lib/ktl/migration_plan.rb

@@ -0,0 +1,27 @@
+# encoding: utf-8
+
+module Ktl
+  class MigrationPlan
+    def initialize(zk_client, old_leader, new_leader)
+      @zk_client = zk_client
+      @old_leader = old_leader.to_java
+      @new_leader = new_leader.to_java
+    end
+
+    def generate
+      plan = Scala::Collection::Map.empty
+      topics = @zk_client.all_topics
+      assignments = ScalaEnumerable.new(@zk_client.replica_assignment_for_topics(topics))
+      assignments.each do |item|
+        topic_partition = item.first
+        replicas = item.last
+        if replicas.contains?(@old_leader)
+          index = replicas.index_of(@old_leader)
+          new_replicas = replicas.updated(index, @new_leader, CanBuildFrom)
+          plan += Scala::Tuple.new(topic_partition, new_replicas)
+        end
+      end
+      plan
+    end
+  end
+end

data/lib/ktl/reassigner.rb

@@ -0,0 +1,133 @@
+# encoding: utf-8
+
+module Ktl
+  class Reassigner
+    def initialize(zk_client, options={})
+      @zk_client = zk_client
+      @limit = options[:limit]
+      @overflow_path = '/ktl/overflow'
+      @state_path = '/ktl/reassign'
+      @logger = options[:logger] || NullLogger.new
+      @log_assignments = !!options[:log_assignments]
+    end
+
+    def reassignment_in_progress?
+      partitions = @zk_client.partitions_being_reassigned
+      partitions.size > 0
+    end
+
+    def overflow?
+      overflow_znodes = @zk_client.get_children(@overflow_path)
+      overflow_znodes.size > 0
+    rescue ZkClient::Exception::ZkNoNodeException
+      false
+    end
+
+    def load_overflow
+      overflow = Scala::Collection::Map.empty
+      overflow_nodes = @zk_client.get_children(@overflow_path)
+      overflow_nodes.foreach do |index|
+        overflow_json = @zk_client.read_data(overflow_path(index)).first
+        data = parse_reassignment_json(overflow_json)
+        overflow = overflow.send('++', data)
+      end
+      delete_previous_overflow
+      overflow
+    end
+
+    def execute(reassignment)
+      reassignments = split(reassignment, @limit)
+      actual_reassignment = reassignments.shift
+      if @log_assignments
+        Scala::Collection::JavaConversions.as_java_iterable(actual_reassignment).each do |pr|
+          topic_and_partition, replicas = pr.elements
+          brokers = Scala::Collection::JavaConversions.as_java_iterable(replicas).to_a
+          @logger.info "Assigning #{topic_and_partition.topic},#{topic_and_partition.partition} to #{brokers.join(',')}"
+        end
+      end
+      json = reassignment_json(actual_reassignment)
+      @zk_client.reassign_partitions(json)
+      manage_overflow(reassignments)
+      manage_progress_state(actual_reassignment)
+    end
+
+    private
+
+    JSON_MAX_SIZE = 1024**2
+
+    def manage_progress_state(reassignment)
+      delete_previous_state
+      json = reassignment_json(reassignment)
+      @zk_client.create_znode(@state_path, json)
+    end
+
+    def delete_previous_state
+      if @zk_client.exists?(@state_path)
+        @zk_client.delete_znode(@state_path)
+      end
+    end
+
+    def delete_previous_overflow
+      overflow = @zk_client.get_children(@overflow_path)
+      overflow.foreach do |index|
+        @zk_client.delete_znode(overflow_path(index))
+      end
+    rescue ZkClient::Exception::ZkNoNodeException
+    end
+
+    def manage_overflow(reassignments)
+      delete_previous_overflow
+      empty_map = Scala::Collection::Map.empty
+      overflow = reassignments.reduce(empty_map) do |acc, data|
+        acc.send('++', data)
+      end
+      if overflow.size > 0
+        write_overflow(split(overflow))
+      end
+    end
+
+    def overflow_path(index)
+      [@overflow_path, index].join('/')
+    end
+
+    def write_overflow(reassignments)
+      reassignments.each_with_index do |reassignment, index|
+        overflow_json = reassignment_json(reassignment)
+        @zk_client.create_znode(overflow_path(index), overflow_json)
+      end
+    end
+
+    def reassignment_json(reassignment)
+      zk_utils.format_as_reassignment_json(reassignment)
+    end
+
+    def parse_reassignment_json(json)
+      zk_utils.parse_partition_reassignment_data(json)
+    end
+
+    def zk_utils
+      @zk_utils ||= Kafka::Utils::ZkUtils.new(nil, nil, false)
+    end
+
+    def maybe_split_by_limit(reassignment, limit=nil)
+      if limit
+        splitted = ScalaEnumerable.new(reassignment.grouped(limit)).map(&:seq)
+      else
+        splitted = [reassignment]
+      end
+    end
+
+    def split(reassignment, limit=nil)
+      splitted = maybe_split_by_limit(reassignment, limit)
+      bytesize = reassignment_json(splitted.first).bytesize
+      while bytesize > JSON_MAX_SIZE do
+        splitted = splitted.flat_map do |s|
+          group_size = s.size.fdiv(2).round
+          ScalaEnumerable.new(s.grouped(group_size)).map(&:seq)
+        end
+        bytesize = reassignment_json(splitted.first).bytesize
+      end
+      splitted
+    end
+  end
+end

data/lib/ktl/reassignment_progress.rb

@@ -0,0 +1,57 @@
+# encoding: utf-8
+
+module Ktl
+  class ReassignmentProgress
+    def initialize(zk_client, options={})
+      @zk_client = zk_client
+      @utils = options[:utils] || Kafka::Utils::ZkUtils
+      @logger = options[:logger] || NullLogger.new
+      @state_path = '/ktl/reassign'
+      @options = options
+    end
+
+    def display(shell)
+      in_progress = reassignment_in_progress
+      original = original_reassignment
+      if in_progress && !in_progress.empty?
+        original_size, remaining_size = original.size, in_progress.size
+        done_percentage = (original_size - remaining_size).fdiv(original_size) * 100
+        @logger.info 'remaining partitions to reassign: %d (%.2f%% done)' % [remaining_size, done_percentage]
+        if @options[:verbose]
+          shell.print_table(table_data(in_progress), indent: 2 + 6)
+        end
+      else
+        @logger.info 'no partitions remaining to reassign'
+      end
+    end
+
+    private
+
+    def table_data(reassignments)
+      topics = reassignments.group_by { |r| r['topic'] }
+      table = topics.map do |t, r|
+        reassignments = r.sort_by { |r| r['partition'] }
+        reassignments = reassignments.map { |r| '%d => %s' % [r['partition'], r['replicas'].inspect] }.join(', ')
+        [t, reassignments]
+      end.sort_by(&:first)
+      table.unshift(%w[topic assignments])
+      table
+    end
+
+    def reassignment_in_progress
+      read_json(@utils.reassign_partitions_path).fetch('partitions')
+    rescue ZkClient::Exception::ZkNoNodeException
+      {}
+    end
+
+    def original_reassignment
+      read_json(@state_path).fetch('partitions')
+    rescue ZkClient::Exception::ZkNoNodeException
+      {}
+    end
+
+    def read_json(path)
+      JSON.parse(@zk_client.read_data(path).first)
+    end
+  end
+end

data/lib/ktl/reassignment_task.rb

@@ -0,0 +1,46 @@
+# encoding: utf-8
+
+module Ktl
+  class ReassignmentTask
+    def initialize(reassigner, plan, shell, options={})
+      @reassigner = reassigner
+      @plan = plan
+      @shell = shell
+      @logger = options[:logger] || NullLogger.new
+    end
+
+    def execute(dryrun = false)
+      if @reassigner.reassignment_in_progress?
+        @logger.warn 'reassignment already in progress, exiting'
+      else
+        if use_overflow?
+          @logger.info 'loading overflow data'
+          reassignment = @reassigner.load_overflow
+        else
+          @logger.info 'generating a new reassignment plan'
+          reassignment = @plan.generate
+        end
+        if reassignment.size > 0
+          @logger.info 'reassigning %d partitions' % reassignment.size
+          if dryrun
+            @logger.info 'dryrun detected, skipping reassignment'
+          else
+            @reassigner.execute(reassignment)
+          end
+        else
+          @logger.warn 'empty reassignment, ignoring'
+        end
+      end
+    end
+
+    private
+
+    def use_overflow?
+      if @reassigner.overflow?
+        @logger.info 'overflow from previous reassignment found, use? [y/n]'
+        @shell.yes? ' ' * 8 << '>'
+      end
+    end
+  end
+end
+

data/lib/ktl/shell_formatter.rb

@@ -0,0 +1,32 @@
+# encoding: utf-8
+
+module Ktl
+  class ShellFormater < Logger::Formatter
+    def initialize(shell)
+      @shell = shell
+      @padding = ' ' * 2
+    end
+
+    def call(severity, time, progname, msg)
+      severity_part = sprintf('%-5s', severity.downcase)
+      severity_part = @shell.set_color(severity_part, color_for[severity])
+      line = %(#{@padding}#{severity_part} #{msg2str(msg)})
+      line << NEWLINE unless line.end_with?(NEWLINE)
+      line
+    end
+
+    private
+
+    NEWLINE = "\n".freeze
+
+    def color_for
+      @color_for ||= {
+        'DEBUG' => :cyan,
+        'INFO' => :magenta,
+        'WARN' => :yellow,
+        'ERROR' => :red,
+        'FATAL' => :red,
+      }
+    end
+  end
+end