Kafka自定义分区策略实战避坑指南

概要

kafka生产者发送消息默认根据总分区数和设置的key计算哈希取余数，key不变就默认存放在一个分区，没有key则随机数分区，明显默认的是最不好用的，那kafka也提供了一个轮询分区策略，我自己使用的是一言难尽，具体我也没有深究下去，那么针对业务硬性要求消息按照升序或降序轮询分区，就需要我们自己定义分区策略了。

有多少小伙伴第一次配置自定义分区策略时，发现分区总是按照倍数分区，并没有按照指定的规则去分区呢？嘿嘿，相信没阅读过源码的都应该踩过这一个坑，原因在于生产者发送消息时，kafka会先去分区策略那里逛一圈，拿到本次分区值，再去执行下一步流程，而在真正执行发送消息之前，kafka会再次进入分区策略内拿取本次的分区值，那么轮询策略一般按照依次递增或递减，致使发送消息时都会拿到自增两次后的分区值。

好，知道了问题所在，那就简单了，修改逻辑就行了呗，这一块考虑到使用分区策略一般是应对多个消息的产生同时发送，所以就涉及到并发了，那么并发就要考虑线程安全，这里推荐使用原子自增类和原子Boolean(非必要)，能不使用锁就不使用锁，具体根据各位的业务而定吧，那话不多说，上代码。

代码示例

package org.example.springkafkademo.config;import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.utils.Utils;import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;public class CustomerPartitioner implements Partitioner {//针对并发设计，使分区数量原子自增private static AtomicInteger nextPartition  = new AtomicInteger(0);//二次进入判断机制private static AtomicBoolean flag = new AtomicBoolean(false);@Overridepublic int partition(String topic, Object key, byte[] bytes, Object o1, byte[] keyBytes, Cluster cluster) {List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);//最大自增值int numPartitions = partitions.size();if (key == null) {//二次判断机制为true则说明自增过一次，需要返回自增之前的值if (flag.get()){flag.set(false);return nextPartition.get()-1;}//原子类将旧值返回再自增int next = nextPartition.getAndIncrement();//如果自增后与大于最大值或相等则直接cas赋值0，使下一次的轮询从0开始if (next >= numPartitions) {nextPartition.compareAndSet(numPartitions, 0);}//标记已经进入过一次flag.set(true);System.out.println("分区值：" + next);return next;} else {// 如果key不为null，则使用默认的分区策略return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;}}@Overridepublic void close() {}@Overridepublic void configure(Map<String, ?> map) {}
}

小结

本文分享kafka实现自定义轮询策略，在应对需要将大量的消息轮询发送给分区的场景时，可以采纳本文的代码逻辑，但是并不是适配所有分区轮询，毕竟业务逻辑不是定死的，各位小伙伴一定要结合实际业务逻辑，针对性的对代码进行修改扩展。
有哪里不懂得小伙伴可留言或私信，如与本文章有不同观点欢迎讨论留言，大家一起进步。