如何为Kafka Streams实现适配StreamPartitioner的RoundRobinPartitioner

我有一个基于Kafka Streams的简单应用，当key为null时，默认分区器底层会使用粘性分区器。我希望将DefaultStreamPartitioner替换为RoundRobinPartitioner，但遇到了接口不兼容的问题：默认的RoundRobinPartitioner实现了Partitioner接口（方法包含Cluster参数），而DefaultStreamPartitioner实现的是StreamPartitioner接口（无Cluster实例）。

常规分区器与Stream分区器的方法签名如下：

// 常规分区器
public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster)

// Stream分区器
Integer partition(String topic, K key, V value, int numPartitions);

我的问题是：如何为StreamPartitioner实现RoundRobinPartitioner？

我查看了DefaultStreamPartitioner的实现，发现它持有Cluster实例，但我不知道该从何处获取这个实例。

DefaultStreamPartitioner的实现代码：

public class DefaultStreamPartitioner<K, V> implements StreamPartitioner<K, V> {

    private final Cluster cluster;
    private final Serializer<K> keySerializer;
    private final DefaultPartitioner defaultPartitioner;

    public DefaultStreamPartitioner(final Serializer<K> keySerializer, final Cluster cluster) {
        this.cluster = cluster;
        this.keySerializer = keySerializer;
        this.defaultPartitioner = new DefaultPartitioner();
    }

    @Override
    public Integer partition(final String topic, final K key, final V value, final int numPartitions) {
        final byte[] keyBytes = keySerializer.serialize(topic, key);
        return defaultPartitioner.partition(topic, key, keyBytes, value, null, cluster);
    }
}

方案一：直接实现独立的轮询Stream分区器（推荐） #

不需要依赖Cluster实例，利用StreamPartitioner提供的numPartitions参数，通过原子计数器实现线程安全的轮询逻辑，同时保留非null key的哈希分区一致性。

import org.apache.kafka.streams.processor.StreamPartitioner;
import java.util.concurrent.atomic.AtomicInteger;

public class RoundRobinStreamPartitioner<K, V> implements StreamPartitioner<K, V> {

    // 原子计数器保证多线程环境下的线程安全
    private final AtomicInteger counter = new AtomicInteger(0);

    @Override
    public Integer partition(String topic, K key, V value, int numPartitions) {
        // key不为null时，保持原有哈希分区逻辑，保证数据一致性
        if (key != null) {
            return Math.abs(key.hashCode()) % numPartitions;
        }
        // key为null时，轮询分配分区
        return counter.getAndIncrement() % numPartitions;
    }
}

使用方式：在Topology的输出节点指定该分区器即可

Topology topology = new Topology();
topology.addSource("source", "input-topic")
        // 业务处理逻辑
        .to("output-topic", 
            Produced.with(Serdes.String(), Serdes.String())
                   .withStreamPartitioner(new RoundRobinStreamPartitioner<>()));

方案二：复用原生RoundRobinPartitioner（适配接口） #

如果需要复用Kafka原生的RoundRobinPartitioner，可以模仿DefaultStreamPartitioner的实现方式，通过构造器传入Cluster和序列化器，适配StreamPartitioner接口。

1. 实现适配类

import org.apache.kafka.clients.producer.RoundRobinPartitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.streams.processor.StreamPartitioner;

public class WrappedRoundRobinStreamPartitioner<K, V> implements StreamPartitioner<K, V> {

    private final Cluster cluster;
    private final Serializer<K> keySerializer;
    private final RoundRobinPartitioner roundRobinPartitioner;

    public WrappedRoundRobinStreamPartitioner(Serializer<K> keySerializer, Cluster cluster) {
        this.keySerializer = keySerializer;
        this.cluster = cluster;
        this.roundRobinPartitioner = new RoundRobinPartitioner();
    }

    @Override
    public Integer partition(String topic, K key, V value, int numPartitions) {
        byte[] keyBytes = keySerializer.serialize(topic, key);
        return roundRobinPartitioner.partition(topic, key, keyBytes, value, null, cluster);
    }
}

2. 获取Cluster实例

通过AdminClient获取集群元数据：

import org.apache.kafka.clients.admin.AdminClient;
import org.apache.kafka.clients.admin.AdminClientConfig;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.Node;
import java.util.Properties;
import java.util.List;

public Cluster getCluster(Properties streamsConfig) {
    Properties adminProps = new Properties();
    adminProps.put(AdminClientConfig.BOOTSTRAP_SERVERS_CONFIG, streamsConfig.get("bootstrap.servers"));
    try (AdminClient adminClient = AdminClient.create(adminProps)) {
        List<Node> nodes = adminClient.describeCluster().nodes().get();
        return new Cluster(streamsConfig.getProperty("bootstrap.servers"), nodes, List.of(), List.of(), List.of());
    } catch (Exception e) {
        throw new RuntimeException("获取集群元数据失败", e);
    }
}

3. 使用适配后的分区器

Properties props = new Properties();
props.put(StreamsConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
// 其他Streams配置...

Cluster cluster = getCluster(props);
Serializer<String> keySerializer = new StringSerializer(); // 替换为你的key序列化器
StreamPartitioner<String, String> partitioner = new WrappedRoundRobinStreamPartitioner<>(keySerializer, cluster);

Topology topology = new Topology();
topology.addSource("source", "input-topic")
        // 业务处理逻辑
        .to("output-topic", 
            Produced.with(Serdes.String(), Serdes.String())
                   .withStreamPartitioner(partitioner));

内容的提问来源于stack exchange，提问作者Mykeo