Kafka为我们提供了默认的分区策略和自定义消息分区策略两种方式管理producer端消息分区。
一、默认分区
Kafka默认分区使用的是DefaultPartitioner，它实现了Partitioner接口。 首先我们看一下Partitioner接口。
Partitioner接口
在2.4.0之前，Partitioner接口中有一下方法：
/**
* Compute the partition for the given record.
*
* @param topic The topic name
* @param key The key to partition on (or null if no key)
* @param keyBytes The serialized key to partition on( or null if no key)
* @param value The value to partition on or null
* @param valueBytes The serialized value to partition on or null
* @param cluster The current cluster metadata
*/
public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster);
/**
* This is called when partitioner is closed.
*/
public void close();
很显然，partition方法就是用来计算消息分区的。
在2.4.0版本上，Kafka增加了onNewBatch方法，这个方法主要用来更新一个新的发送batch分区数的，下面会详细讲解。
/**
* Notifies the partitioner a new batch is about to be created. When using the sticky partitioner,
* this method can change the chosen sticky partition for the new batch.
* @param topic The topic name
* @param cluster The current cluster metadata
* @param prevPartition The partition previously selected for the record that triggered a new batch
*/
default public void onNewBatch(String topic, Cluster cluster, int prevPartition) {
}
DefaultPartitioner
看一下DefaultPartitioner实现，主要看一下partition方法，它实现如何计算partition值。 Kafka 2.4.0之前代码：
public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
//根据topic获取partition数量
List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
int numPartitions = partitions.size();
if (keyBytes == null) {
int nextValue = nextValue(topic);
//获取可用的分区数量
List<PartitionInfo> availablePartitions = cluster.availablePartitionsForTopic(topic);
if (availablePartitions.size() > 0) {
//partition = nextValue值对可用分区数取模，nextValue是一个自增的AtomicInteger，这就相当于没有指定key时，轮训所有可用分区
int part = Utils.toPositive(nextValue) % availablePartitions.size();
return availablePartitions.get(part).partition();
} else {
//没有可用分区，随机分配到一个partition
return Utils.toPositive(nextValue) % numPartitions;
}
} else {
//消息指定了key，partition = key的hash值对分区数取模
return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
}
}
private int nextValue(String topic) {
//根据topic获取counter值，没有就创建一个新的，有就自增返回
AtomicInteger counter = topicCounterMap.get(topic);
if (null == counter) {
counter = new AtomicInteger(ThreadLocalRandom.current().nextInt());
AtomicInteger currentCounter = topicCounterMap.putIfAbsent(topic, counter);
if (currentCounter != null) {
counter = currentCounter;
}
}
//自增返回
return counter.getAndIncrement();
}
从partition方法中可以看到，如果消息指定了key，就采用key的hash值对分区数取模的方式得到分区；如果没有指定key，则采用轮训的方式得到分区数。
在Kafka 2.4.0之后，Kafka对未指定key的分区逻辑做了修改,引入了新的粘分区缓存类StickyPartitionCache，没有key时，程序会去缓存中获取分区数。
public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster,
int numPartitions) {
//没有key，去粘分区中获取partition
if (keyBytes == null) {
return stickyPartitionCache.partition(topic, cluster);
}
// hash the keyBytes to choose a partition
//有key，去序列化后的key的hash，模分区数
return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
}
public void onNewBatch(String topic, Cluster cluster, int prevPartition) {
//更新粘分区缓存中的分区数，详情看nextPartition方法。
stickyPartitionCache.nextPartition(topic, cluster, prevPartition);
}
看一下StickyPartitionCache代码：
public int partition(String topic, Cluster cluster) {
//缓存中有，直接返回，如果缓存中没有，就用nextPartition方法更新缓存
Integer part = indexCache.get(topic);
if (part == null) {
return nextPartition(topic, cluster, -1);
}
return part;
}
public int nextPartition(String topic, Cluster cluster, int prevPartition) {
//更新缓存，这个逻辑本质上也是轮训，与老版本逻辑类似
List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
Integer oldPart = indexCache.get(topic);
Integer newPart = oldPart;
if (oldPart == null || oldPart == prevPartition) {
List<PartitionInfo> availablePartitions = cluster.availablePartitionsForTopic(topic);
//可用分区数小于1，则使用总分区数随机分配分区
if (availablePartitions.size() < 1) {
Integer random = Utils.toPositive(ThreadLocalRandom.current().nextInt());
newPart = random % partitions.size();
} else if (availablePartitions.size() == 1) {
//可用分区为1
newPart = availablePartitions.get(0).partition();
} else {
//多个可用分区，随机分配
while (newPart == null || newPart.equals(oldPart)) {
Integer random = Utils.toPositive(ThreadLocalRandom.current().nextInt());
newPart = availablePartitions.get(random % availablePartitions.size()).partition();
}
}
//cache中没有分区数，或者是新的batch时，更新分区数
if (oldPart == null) {
indexCache.putIfAbsent(topic, newPart);
} else {
indexCache.replace(topic, prevPartition, newPart);
}
return indexCache.get(topic);
}
return indexCache.get(topic);
}
上面的nextPartition方法时用来更新缓存的，两种情况会用到，一是缓存丢失或者第一次还没有缓存时，二是当accumulator中产生新的batch时，我们看一下，它是怎么被调用的：
private Future<RecordMetadata> doSend(ProducerRecord<K, V> record, Callback callback) {
... ...
if (result.abortForNewBatch) { //产生新的batch
int prevPartition = partition;
//当record进入一个新的batch的时候，修改粘分区的partition数
partitioner.onNewBatch(record.topic(), cluster, prevPartition);
partition = partition(record, serializedKey, serializedValue, cluster);
... ...
}
... ...
}
StickyPartitionCache作用
StickyPartitionCache是为了保证不产生过多的小batch导致producer发送过多的请求。
老版本分区方式，各个消息拥有不同的分区号，会被轮训到每一个分区中，因为发送时一个batch中分区号必须相同，这样就会产生很更多的网络请求。而StickyPartitionCache保证了没有可以的消息使用相同的分区号，发送时以一个整体发送网络请求到相同的分区中，可以提高Kafka的吞吐量。
二、自定义Partitioner
自定义Partitioner只需要实现Partitioner接口，重写partition方法就行，具体的分区逻辑在partition方法中实现，example：
public class CustomPartitioner implements Partitioner {
@Override public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes1, Cluster cluster) {
if (Integer.parseInt((String)key)%3==1) {
return 0;
}else if (Integer.parseInt((String)key)%3==2) {
return 1;
}else{
return 2;
}
}
@Override public void close() {
}
@Override public void configure(Map<String, ?> map) {
}
}
