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);
    @Override
    public 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;
        }
    }

    @Override
    public void close() {

    }

    @Override
    public void configure(Map<String, ?> map) {

    }
}

小结

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