def main(): # Start by loading all the modules action_map = get_module_map(kafka.tools.assigner.actions, kafka.tools.assigner.actions.ActionModule) sizer_map = get_module_map(kafka.tools.assigner.sizers, kafka.tools.assigner.sizers.SizerModule) plugins = get_all_plugins() # Set up and parse all CLI arguments args = set_up_arguments(action_map, sizer_map, plugins) run_plugins_at_step(plugins, 'set_arguments', args) tools_path = get_tools_path(args.tools_path) check_java_home() cluster = Cluster.create_from_zookeeper(args.zookeeper, getattr(args, 'default_retention', 1)) run_plugins_at_step(plugins, 'set_cluster', cluster) # If the module needs the partition sizes, call a size module to get the information check_and_get_sizes(action_map[args.action], args, cluster, sizer_map) run_plugins_at_step(plugins, 'after_sizes') print_leadership("before", cluster, args.leadership) # Clone the cluster, and run the action to generate a new cluster state newcluster = cluster.clone() action_to_run = action_map[args.action](args, newcluster) action_to_run.process_cluster() run_plugins_at_step(plugins, 'set_new_cluster', action_to_run.cluster) print_leadership("after", newcluster, args.leadership) move_partitions = cluster.changed_partitions(action_to_run.cluster) batches = split_partitions_into_batches(move_partitions, batch_size=args.moves, use_class=Reassignment) run_plugins_at_step(plugins, 'set_batches', batches) log.info("Partition moves required: {0}".format(len(move_partitions))) log.info("Number of batches: {0}".format(len(batches))) dry_run = is_dry_run(args) for i, batch in enumerate(batches): log.info("Executing partition reassignment {0}/{1}: {2}".format(i + 1, len(batches), repr(batch))) batch.execute(i + 1, len(batches), args.zookeeper, tools_path, plugins, dry_run) run_plugins_at_step(plugins, 'before_ple') if not args.skip_ple: all_cluster_partitions = [p for p in action_to_run.cluster.partitions(args.exclude_topics)] batches = split_partitions_into_batches(all_cluster_partitions, batch_size=args.ple_size, use_class=ReplicaElection) log.info("Number of replica elections: {0}".format(len(batches))) run_preferred_replica_elections(batches, args, tools_path, plugins, dry_run) run_plugins_at_step(plugins, 'finished') if args.output_json: data = { 'before': cluster.to_dict(), 'after': action_to_run.cluster.to_dict() } sys.stdout.write(json.dumps(data, indent=4, sort_keys=True)) return os.EX_OK
def test_split_batches_proper_class(self): batches = split_partitions_into_batches(self.topic.partitions, batch_size=100, use_class=Reassignment) assert isinstance(batches[0], Reassignment) batches = split_partitions_into_batches(self.topic.partitions, batch_size=100, use_class=ReplicaElection) assert isinstance(batches[0], ReplicaElection)
def main(): # Start by loading all the modules action_map = get_module_map(kafka.tools.assigner.actions, kafka.tools.assigner.actions.ActionModule) sizer_map = get_module_map(kafka.tools.assigner.sizers, kafka.tools.assigner.sizers.SizerModule) plugins = get_all_plugins() # Set up and parse all CLI arguments args = set_up_arguments(action_map, sizer_map, plugins) run_plugins_at_step(plugins, 'set_arguments', args) tools_path = get_tools_path(args.tools_path) check_java_home() cluster = Cluster.create_from_zookeeper(args.zookeeper) run_plugins_at_step(plugins, 'set_cluster', cluster) # If the module needs the partition sizes, call a size module to get the information check_and_get_sizes(action_map[args.action], args, cluster, sizer_map) run_plugins_at_step(plugins, 'after_sizes') print_leadership("before", cluster, args.leadership) # Clone the cluster, and run the action to generate a new cluster state newcluster = cluster.clone() action_to_run = action_map[args.action](args, newcluster) action_to_run.process_cluster() run_plugins_at_step(plugins, 'set_new_cluster', action_to_run.cluster) print_leadership("after", newcluster, args.leadership) move_partitions = cluster.changed_partitions(action_to_run.cluster) batches = split_partitions_into_batches(move_partitions, batch_size=args.moves, use_class=Reassignment) run_plugins_at_step(plugins, 'set_batches', batches) log.info("Partition moves required: {0}".format(len(move_partitions))) log.info("Number of batches: {0}".format(len(batches))) dry_run = is_dry_run(args) for i, batch in enumerate(batches): log.info("Executing partition reassignment {0}/{1}: {2}".format(i + 1, len(batches), repr(batch))) batch.execute(i + 1, len(batches), args.zookeeper, tools_path, plugins, dry_run) run_plugins_at_step(plugins, 'before_ple') if not args.skip_ple: all_cluster_partitions = [p for p in action_to_run.cluster.partitions()] batches = split_partitions_into_batches(all_cluster_partitions, batch_size=args.ple_size, use_class=ReplicaElection) log.info("Number of replica elections: {0}".format(len(batches))) run_preferred_replica_elections(batches, args, tools_path, plugins, dry_run) run_plugins_at_step(plugins, 'finished') return os.EX_OK
def test_split_batches_notenough(self): batches = split_partitions_into_batches(self.topic.partitions, batch_size=20, use_class=Reassignment) partition_count = sum([len(batch.partitions) for batch in batches]) assert len(batches) == 1 assert partition_count == 10
def test_split_batches_empty(self): partitions = [] batches = split_partitions_into_batches(partitions, batch_size=1, use_class=Reassignment) assert len(batches) == 0
def main(): # Start by loading all the modules action_map = get_action_map() sizer_map = get_sizer_map() plugins_list = get_plugins_list() # Instantiate all plugins plugins = [plugin() for plugin in plugins_list] # Set up and parse all CLI arguments args = set_up_arguments(action_map, sizer_map, plugins) for plugin in plugins: plugin.set_arguments(args) tools_path = get_tools_path(args.tools_path) check_java_home() cluster = Cluster.create_from_zookeeper(args.zookeeper) for plugin in plugins: plugin.set_cluster(cluster) # If the module needs the partition sizes, call a size module to get the information check_and_get_sizes(action_map[args.action], args, cluster, sizer_map) for plugin in plugins: plugin.after_sizes() if args.leadership: log.info("Cluster Leadership Balance (before):") cluster.log_broker_summary() # Clone the cluster, and run the action to generate a new cluster state newcluster = cluster.clone() action_to_run = action_map[args.action](args, newcluster) action_to_run.process_cluster() for plugin in plugins: plugin.set_new_cluster(action_to_run.cluster) if args.leadership: log.info("Cluster Leadership Balance (after):") newcluster.log_broker_summary() move_partitions = cluster.changed_partitions(action_to_run.cluster) batches = split_partitions_into_batches(move_partitions, batch_size=args.moves, use_class=Reassignment) for plugin in plugins: plugin.set_batches(batches) log.info("Partition moves required: {0}".format(len(move_partitions))) log.info("Number of batches: {0}".format(len(batches))) dry_run = args.generate or not args.execute if dry_run: log.info("--execute flag NOT specified. DRY RUN ONLY") for i, batch in enumerate(batches): log.info("Executing partition reassignment {0}/{1}: {2}".format(i + 1, len(batches), repr(batch))) batch.execute(i + 1, len(batches), args.zookeeper, tools_path, plugins, dry_run) for plugin in plugins: plugin.before_ple() if not args.skip_ple: batches = split_partitions_into_batches(move_partitions, batch_size=args.moves, use_class=ReplicaElection) log.info("Number of replica elections: {0}".format(len(batches))) run_preferred_replica_elections(batches, args, tools_path, plugins, dry_run) for plugin in plugins: plugin.finished() return 0