Kafka 自定义指定消息partition策略规则及DefaultPartitioner源码分析

一.概述

kafka默认使用DefaultPartitioner类作为默认的partition策略规则,具体默认设置是在ProducerConfig类中(如下图)

二.DefaultPartitioner.class 源码分析

1.类关系图

2.源码分析

1public class DefaultPartitioner implements Partitioner {
2	//缓存map key->topic  value->RandomNumber 随机数 
3    private final ConcurrentMap<String, AtomicInteger> topicCounterMap = new ConcurrentHashMap<>();
4	
5	//实现Configurable接口里configure方法,
6    public void configure(Map<String, ?> configs) {}
7
8	//策略核心方法
9    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
10    	//根据topic获取对应的partition
11        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
12        int numPartitions = partitions.size();
13        //如果key是null的话
14        if (keyBytes == null) {
15            int nextValue = nextValue(topic);
16            //获取可用的分区数量
17            List<PartitionInfo> availablePartitions = cluster.availablePartitionsForTopic(topic);
18            //如果存在可用的分区
19            if (availablePartitions.size() > 0) {
20            	//消息随机分布到topic的可用partition中
21                int part = Utils.toPositive(nextValue) % availablePartitions.size();
22                return availablePartitions.get(part).partition();
23            } else {
24				//不存在可用分区 随机分配一个不可用的partition中            		
25                return Utils.toPositive(nextValue) % numPartitions;
26            }
27        } else {
28        	//使用自己的 hash 算法对 key 取 hash 值，使用 hash 值与 partition 数量取模，从而确定发送到哪个分区。
29            return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
30        }
31    }
32
33    private int nextValue(String topic) {
34        //获取topic 获取对应的随机数
35        AtomicInteger counter = topicCounterMap.get(topic);
36        if (null == counter) {
37            //获取一个随机值
38            counter = new AtomicInteger(ThreadLocalRandom.current().nextInt());
39            //缓存到topicCounterMap中
40            AtomicInteger currentCounter = topicCounterMap.putIfAbsent(topic, counter);
41            if (currentCounter != null) {
42                counter = currentCounter;
43            }
44        }
45        //获取 并且实现自增,最终效果是实现轮训插入partition
46        return counter.getAndIncrement();
47    }
48
49    public void close() {}
50
51}

三.自定义Partition

自定义selfPartitioner类,并且实现Partitioner接口,重写partition和close方法

1public class SelfPartitioner implements Partitioner {
2
3    @Override
4    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
5        //TODO 自定义数据分区策略
6        return 0;
7    }
8
9    @Override
10    public void close() {
11
12    }
13
14    @Override
15    public void configure(Map<String, ?> configs) {
16
17    }
18}

四.生产者中使用自定义的Partition

在初始化producer时候,配置项中指定对应的partitioner

props.put("partitioner.class", "org.jake.partitioner.selfPartitioner");