ktl 1.0.0-java

data/lib/ktl/shuffle_plan.rb

@@ -0,0 +1,145 @@
+# encoding: utf-8
+
+module Ktl
+  class ShufflePlan
+    def initialize(zk_client, options = {})
+      @zk_client = zk_client
+      @options = options
+      @logger = options[:logger] || NullLogger.new
+      @log_plan = !!options[:log_plan]
+    end
+
+    def generate
+      topics = @zk_client.all_topics
+      if (filter = @options[:filter])
+        topics = topics.filter { |t| !!t.match(filter) }
+      end
+      topics_partitions = ScalaEnumerable.new(@zk_client.partitions_for_topics(topics))
+      topics_partitions = topics_partitions.sort_by(&:first)
+      replica_assignments = @zk_client.replica_assignment_for_topics(topics)
+      brokers = select_brokers
+      reassignment_plan = Scala::Collection::Map.empty
+      topics_partitions.each do |tp|
+        topic, partitions = tp.elements
+        nr_replicas = @options[:replication_factor] || replica_assignments.apply(Kafka::TopicAndPartition.new(topic, 0)).size
+        assignment = assign_replicas_to_brokers(topic, brokers, partitions.size, nr_replicas)
+        assignment.each do |pr|
+          partition, replicas = pr.elements
+          topic_partition = Kafka::TopicAndPartition.new(topic, partition)
+          current_assignment = replica_assignments.apply(topic_partition)
+          unless current_assignment == replicas
+            @logger.info "Moving #{topic_partition.topic},#{topic_partition.partition} from #{current_assignment} to #{replicas}" if @log_plan
+            reassignment_plan += Scala::Tuple.new(topic_partition, replicas)
+          end
+        end
+      end
+      reassignment_plan
+    end
+
+    def generate_for_new_topic(topic, partition_count)
+      brokers = select_brokers
+      nr_replicas = @options[:replication_factor] || 1
+      assignment = assign_replicas_to_brokers(topic, brokers, partition_count, nr_replicas)
+      assignment.map do |pr|
+        partition, replicas = pr.elements
+        Scala::Collection::JavaConversions.as_java_iterable(replicas).to_a
+      end
+    end
+
+    private
+
+    def select_brokers
+      brokers = @options[:brokers] ? Array(@options[:brokers]).map(&:to_i) : ScalaEnumerable.new(@zk_client.broker_ids).to_a
+      brokers -= Array(@options[:blacklist]).map(&:to_i) if @options[:blacklist]
+      brokers
+    end
+
+    def assign_replicas_to_brokers(topic, brokers, partition_count, replica_count)
+      broker_metadatas = brokers.map { |x| Kafka::Admin::BrokerMetadata.new(x.to_java(:int), Scala::Option[nil]) }
+      broker_metadatas = Scala::Collection::JavaConversions.as_scala_iterable(broker_metadatas).to_seq
+      Kafka::Admin.assign_replicas_to_brokers(broker_metadatas, partition_count, replica_count)
+    rescue Kafka::Admin::AdminOperationException => e
+      raise ArgumentError, sprintf('%s (%s)', e.message, e.class.name), e.backtrace
+    end
+  end
+
+  class RendezvousShufflePlan < ShufflePlan
+    def assign_replicas_to_brokers(topic, brokers, partition_count, replica_count)
+      if replica_count > brokers.size
+        raise ArgumentError, sprintf('replication factor: %i larger than available brokers: %i', replica_count, brokers.size)
+      end
+      result = []
+      partition_count.times do |partition|
+        sorted = brokers.sort_by do |broker|
+          key = [partition, topic, broker].pack('l<a*l<')
+          Java::OrgJrubyUtil::MurmurHash.hash32(key.to_java_bytes, 0, key.bytesize, SEED)
+        end
+        selected = sorted.take(replica_count)
+        result.push(Scala::Tuple.new(partition, Scala::Collection::JavaConversions.as_scala_iterable(selected).to_list))
+      end
+      result
+    end
+
+    private
+
+    SEED = 1683520333
+  end
+
+  class RackAwareShufflePlan < RendezvousShufflePlan
+    def initialize(*args)
+      super
+      @rack_mappings = {}
+    end
+
+    def assign_replicas_to_brokers(topic, brokers, partition_count, replica_count)
+      if replica_count > brokers.size
+        raise ArgumentError, sprintf('replication factor: %i larger than available brokers: %i', replica_count, brokers.size)
+      end
+      result = []
+      racks = brokers.each_with_object({}) do |broker, acc|
+        rack = rack_for(broker)
+        acc[rack] ||= []
+        acc[rack] << broker
+      end
+      partition_count.times do |partition|
+        first_sorted = racks.flat_map do |rack, rack_brokers|
+          hashed_brokers = rack_brokers.map do |broker|
+            key = [partition, topic, broker].pack('l<a*l<')
+            {id: broker, hash: Java::OrgJrubyUtil::MurmurHash.hash32(key.to_java_bytes, 0, key.bytesize, SEED)}
+          end.sort_by do |broker|
+            broker[:hash]
+          end
+          hashed_brokers.each_with_index do |broker, index|
+            broker[:index] = index
+          end
+        end
+        sorted = first_sorted.sort_by do |broker|
+          [broker[:index], broker[:hash], broker[:id]]
+        end
+        selected = sorted.take(replica_count).map {|broker| broker[:id]}
+        result.push(Scala::Tuple.new(partition, Scala::Collection::JavaConversions.as_scala_iterable(selected).to_list))
+      end
+      result
+    end
+
+    private
+
+    def rack_for(broker_id)
+      unless @rack_mappings[broker_id]
+        broker_metadata = Kafka::Admin.get_broker_metadatas(@zk_client, [broker_id]).first
+        rack = broker_metadata.rack
+        unless rack.isDefined
+          raise "Broker #{broker_metadata.id} is missing rack information, unable to create rack aware shuffle plan."
+        end
+        @rack_mappings[broker_id] = rack.get
+      end
+      @rack_mappings[broker_id]
+    rescue Java::KafkaAdmin::AdminOperationException => e
+      if e.message.match '--disable-rack-aware'
+        raise "Not all brokers have rack information. Unable to create rack aware shuffle plan."
+      else
+        raise e
+      end
+    end
+  end
+end

data/lib/ktl/topic.rb

@@ -0,0 +1,123 @@
+# encoding: utf-8
+
+module Ktl
+  class Topic < Command
+    desc 'list', 'List current topics'
+    option :zookeeper, aliases: %w[-z], required: true, desc: 'ZooKeeper URI'
+    def list
+      with_zk_client do |zk_client|
+        topic_options = Kafka::Admin.to_topic_options(options.merge(list: nil))
+        Kafka::Admin::TopicCommand.list_topics(zk_client.raw_client, topic_options)
+      end
+    end
+
+    desc 'describe', 'Describe (optionally filtered) topics'
+    option :unavailable, aliases: '-u', desc: 'Describe unavailable partitions for topic(s)'
+    option :with_overrides, aliases: '-w', desc: 'Describe topics with config. overrides'
+    option :under_replicated, aliases: '-r', desc: 'Describe under-replicated partitions for topic(s)'
+    option :zookeeper, aliases: %w[-z], required: true, desc: 'ZooKeeper URI'
+    def describe(regexp=nil)
+      with_zk_client do |zk_client|
+        opts = {describe: nil}
+        opts[:topic] = regexp if regexp
+        opts[:topics_with_overrides] = nil if options.with_overrides?
+        opts[:unavailable_partitions] = nil if options.unavailable?
+        opts[:under_replicated_partitions] = nil if options.under_replicated?
+        topic_options = Kafka::Admin.to_topic_options(opts)
+        Kafka::Admin::TopicCommand.describe_topic(zk_client.raw_client, topic_options)
+      end
+    end
+
+    desc 'create NAMES..', 'Create one or more new topics'
+    option :partitions, aliases: %w[-p], default: '1', desc: 'Partitions for new topic(s)'
+    option :replication_factor, aliases: %w[-r], default: '1', desc: 'Replication factor for new topic(s)'
+    option :replica_assignment, aliases: %w[-a], desc: 'Manual replica assignment'
+    option :disable_rack_aware, desc: 'Disable rack awareness'
+    option :rack_aware_allocation, desc: 'Create partitions with Rack aware + Rendezvous-hashing based shuffle'
+    option :rendezvous_allocation, desc: 'Create partitions with Rendezvous-hashing based shuffle'
+    option :config, aliases: %w[-c], desc: 'Key-value pairs of configuration options', type: :hash, default: {}
+    option :zookeeper, aliases: %w[-z], required: true, desc: 'ZooKeeper URI'
+    def create(*names)
+      with_zk_client do |zk_client|
+        names.each do |name|
+          opts = options.merge(create: nil, topic: name)
+          if options.rack_aware_allocation || options.rendezvous_allocation
+            plan_factory = if options.rack_aware_allocation
+              RackAwareShufflePlan
+            else
+              RendezvousShufflePlan
+            end
+
+            plan = plan_factory.new(zk_client, replication_factor: options.replication_factor.to_i)
+            zk_utils = Kafka::Utils::ZkUtils.new(nil, nil, false)
+            opts.delete(:rack_aware_allocation)
+            opts.delete(:rendezvous_allocation)
+            plan = plan.generate_for_new_topic(name, options.partitions.to_i)
+            opts[:replica_assignment] = plan.map {|broker_list| broker_list.join(':')}.join(',')
+          end
+          topic_options = Kafka::Admin.to_topic_options(opts)
+          silence_scala do
+            Kafka::Admin::TopicCommand.create_topic(zk_client.raw_client, topic_options)
+          end
+          message = %(created topic "#{name}" with #{opts[:partitions]} partition(s))
+          message << %(, and replication factor #{opts[:replication_factor]})
+          message << %(, with replica assignment: #{opts[:replica_assignment]}) if opts[:replica_assignment]
+          message << %(, with config: #{opts[:config]}) unless opts[:config].empty?
+          logger.info(message)
+        end
+      end
+    end
+
+    desc 'add-partitions NAMES..', 'Add partitions to one or more existing topics'
+    option :partitions, aliases: %w[-p], required: true, desc: 'New number of partitions'
+    option :zookeeper, aliases: %w[-z], required: true, desc: 'ZooKeeper URI'
+    def add_partitions(*names)
+      with_zk_client do |zk_client|
+        names.each do |name|
+          opts = options.merge(alter: nil, topic: name)
+          topic_options = Kafka::Admin.to_topic_options(opts)
+          logger.warn %(if "#{name}" uses keyed messages, the partition logic or ordering of the messages will be affected)
+          silence_scala do
+            Kafka::Admin::TopicCommand.alter_topic(zk_client.raw_client, topic_options)
+          end
+          logger.info %(increased partitions to #{options.partitions} for "#{name}")
+        end
+      end
+    end
+
+    desc 'delete REGEXP', 'Delete topics matching given regexp'
+    option :zookeeper, aliases: %w[-z], required: true, desc: 'ZooKeeper URI'
+    def delete(regexp)
+      with_zk_client do |zk_client|
+        topics = zk_client.all_topics
+        topics = topics.filter { |t| !!t.match(regexp) }
+        logger.info %(about to mark #{topics.size} topics for deletion)
+        topics.foreach do |topic|
+          Kafka::Utils.delete_topic(zk_client.raw_client, topic)
+          logger.debug %(successfully marked "#{topic}" for deletion)
+        end
+      end
+    end
+
+    desc 'alter REGEXP', 'Alter topic configuration'
+    option :add, aliases: %w[-a], desc: 'Key-value pairs of config options to add', type: :hash, default: {}
+    option :remove, aliases: %w[-r], desc: 'Key-value pairs of config options to remove', type: :array, default: []
+    option :zookeeper, aliases: %w[-z], required: true, desc: 'ZooKeeper URI'
+    def alter(regexp)
+      with_zk_client do |zk_client|
+        opts = {zookeeper: options.zookeeper, topic: regexp}
+        opts[:config] = options.add.dup unless options.add.empty?
+        opts[:delete_config] = options.remove.dup unless options.remove.empty?
+        if opts[:config] || opts[:delete_config]
+          topic_options = Kafka::Admin.to_topic_options(opts)
+          silence_scala do
+            Kafka::Admin::TopicCommand.alter_topic(zk_client.raw_client, topic_options)
+          end
+          logger.info %(updated configuration for topics matching "#{regexp}")
+        else
+          raise ArgumentError, 'missing --add or --remove option'
+        end
+      end
+    end
+  end
+end

data/lib/ktl/version.rb

@@ -0,0 +1,5 @@
+# encoding: utf-8
+
+module Ktl
+  VERSION = '1.0.0'.freeze
+end

data/lib/ktl/zookeeper_client.rb

@@ -0,0 +1,111 @@
+# encoding: utf-8
+
+module Ktl
+  class ZookeeperClient
+    attr_reader :utils
+
+    def initialize(uri, options={})
+      @uri = uri
+      @threadpool = options[:threadpool] || JavaConcurrent::Executors.new_fixed_thread_pool(CONCURRENCY)
+      @utils = options[:utils] || Kafka::Utils::ZkUtils.apply(@uri, 5000, 5000, false)
+    end
+
+    def setup
+      @submit = @threadpool.java_method(:submit, [java.lang.Class.for_name('java.util.concurrent.Callable')])
+      self
+    end
+
+    def close
+      @threadpool.shutdown_now if @threadpool
+      @utils.close
+    end
+
+    def raw_client
+      @utils
+    end
+
+    def all_partitions
+      @utils.get_all_partitions
+    end
+
+    def all_topics
+      @utils.get_all_topics
+    end
+
+    def brokers
+      @utils.get_all_brokers_in_cluster
+    end
+
+    def broker_ids
+      @utils.get_sorted_broker_list
+    end
+
+    def leader_and_isr_for(partitions)
+      @utils.get_partition_leader_and_isr_for_topics(@utils.class.create_zk_client(@uri, 5_000, 5_000), partitions)
+    end
+
+    def partitions_for_topics(topics)
+      request(:get_partitions_for_topics, topics)
+    end
+
+    def replica_assignment_for_topics(topics)
+      request(:get_replica_assignment_for_topics, topics)
+    end
+
+    def partitions_being_reassigned
+      @utils.get_partitions_being_reassigned
+    end
+
+    def reassign_partitions(json)
+      @utils.create_persistent_path(@utils.class.reassign_partitions_path, json, no_acl)
+    end
+
+    def create_znode(path, data='')
+      @utils.create_persistent_path(path, data, no_acl)
+    end
+
+    def delete_znode(path, options={})
+      if options[:recursive]
+        @utils.delete_path_recursive(path)
+      else
+        @utils.delete_path(path)
+      end
+    end
+
+    def read_data(path)
+      @utils.read_data(path)
+    end
+
+    def get_children(path)
+      @utils.get_children(path)
+    end
+
+    def exists?(path)
+      @utils.path_exists(path)
+    end
+
+    private
+
+    CONCURRENCY = 8
+
+    def no_acl
+      Kafka::Utils::ZkUtils::DefaultAcls(false)
+    end
+
+    def request(method, input)
+      chunk_size = [(input.size.to_f / CONCURRENCY).round, 1].max
+      groups = ScalaEnumerable.new(input.grouped(chunk_size).to_seq)
+      futures = groups.map do |slice|
+        @submit.call { @utils.send(method, slice) }
+      end
+      merge(futures.map(&:get))
+    end
+
+    def merge(results)
+      result = Scala::Collection::Map.empty
+      results.reduce(result) do |acc, v|
+        acc.send('++', v)
+      end
+    end
+  end
+end

metadata

@@ -0,0 +1,101 @@
+--- !ruby/object:Gem::Specification
+name: ktl
+version: !ruby/object:Gem::Version
+  version: 1.0.0
+platform: java
+authors:
+- Burt Platform Team
+autorequire:
+bindir: bin
+cert_chain: []
+date: 2017-02-15 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - '='
+      - !ruby/object:Gem::Version
+        version: 0.10.0.1
+  name: kafka-jars
+  prerelease: false
+  type: :runtime
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - '='
+      - !ruby/object:Gem::Version
+        version: 0.10.0.1
+- !ruby/object:Gem::Dependency
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ~>
+      - !ruby/object:Gem::Version
+        version: '0'
+    - - <
+      - !ruby/object:Gem::Version
+        version: '1.0'
+  name: thor
+  prerelease: false
+  type: :runtime
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ~>
+      - !ruby/object:Gem::Version
+        version: '0'
+    - - <
+      - !ruby/object:Gem::Version
+        version: '1.0'
+description: |-
+  ktl is a tool that attempts to make it easier
+                        to manage Kafka clusers that host a lot of topics
+email:
+- david@burtcorp.com
+executables:
+- ktl
+extensions: []
+extra_rdoc_files: []
+files:
+- LICENSE.txt
+- README.md
+- bin/ktl
+- lib/ext/kafka.rb
+- lib/ext/thor.rb
+- lib/ktl.rb
+- lib/ktl/cli.rb
+- lib/ktl/cluster.rb
+- lib/ktl/cluster_stats_task.rb
+- lib/ktl/command.rb
+- lib/ktl/decommission_plan.rb
+- lib/ktl/migration_plan.rb
+- lib/ktl/reassigner.rb
+- lib/ktl/reassignment_progress.rb
+- lib/ktl/reassignment_task.rb
+- lib/ktl/shell_formatter.rb
+- lib/ktl/shuffle_plan.rb
+- lib/ktl/topic.rb
+- lib/ktl/version.rb
+- lib/ktl/zookeeper_client.rb
+homepage: http://github.com/burtcorp/ktl
+licenses:
+- BSD-3-Clause
+metadata: {}
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - '>='
+    - !ruby/object:Gem::Version
+      version: '0'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - '>='
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubyforge_project:
+rubygems_version: 2.4.5
+signing_key:
+specification_version: 4
+summary: Management tool for Kafka clusters
+test_files: []