大数据-67 Kafka 高级特性 分区 分配策略 Ranger、RoundRobin、Sticky、自定义分区器

import org.apache.kafka.clients.consumer.ConsumerPartitionAssignor;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.TopicPartition;

import java.nio.ByteBuffer;
import java.util.*;

public class WeightedPartitionAssignor implements ConsumerPartitionAssignor {

    @Override
    public Subscription subscription(Set<String> topics) {
        // 在这里添加权重信息到 userData
        ByteBuffer buffer = ByteBuffer.allocate(4);
        buffer.putInt(getWeight());
        buffer.flip();
        return new Subscription(new ArrayList<>(topics), buffer);
    }

    @Override
    public Map<String, Assignment> assign(Cluster metadata, Map<String, Subscription> subscriptions) {
        Map<String, Assignment> assignments = new HashMap<>();
        Map<TopicPartition, List<String>> partitionConsumers = new HashMap<>();

        // 遍历所有订阅的topics
        for (String topic : metadata.topics()) {
            List<TopicPartition> partitions = metadata.partitionsForTopic(topic);
            for (TopicPartition partition : partitions) {
                partitionConsumers.putIfAbsent(partition, new ArrayList<>());
            }
        }

        // 根据权重分配分区
        for (Map.Entry<String, Subscription> subscriptionEntry : subscriptions.entrySet()) {
            String consumerId = subscriptionEntry.getKey();
            Subscription subscription = subscriptionEntry.getValue();
            int weight = subscription.userData().getInt();

            for (String topic : subscription.topics()) {
                List<TopicPartition> partitions = metadata.partitionsForTopic(topic);
                for (TopicPartition partition : partitions) {
                    List<String> consumers = partitionConsumers.get(partition);
                    for (int i = 0; i < weight; i++) {
                        consumers.add(consumerId);  // 权重高的消费者多次添加，增加选中的机会
                    }
                }
            }
        }

        // 随机分配分区给消费者
        Random random = new Random();
        for (Map.Entry<TopicPartition, List<String>> entry : partitionConsumers.entrySet()) {
            List<String> consumers = entry.getValue();
            String assignedConsumer = consumers.get(random.nextInt(consumers.size()));
            assignments.computeIfAbsent(assignedConsumer, k -> new Assignment(new ArrayList<>()))
                       .partitions().add(entry.getKey());
        }

        return assignments;
    }

    @Override
    public void onAssignment(Assignment assignment, Cluster metadata) {
        // 可以在这里处理分配后的逻辑，比如保存当前分配的快照
    }

    @Override
    public String name() {
        return "weighted";
    }

    private int getWeight() {
        // 获取权重，可以从配置文件或环境变量中获取
        return 10; // 默认权重为10
    }
}

Properties props = new Properties();
props.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:9092");
props.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringDeserializer");
props.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, "org.apache.kafka.common.serialization.StringDeserializer");
props.put(ConsumerConfig.GROUP_ID_CONFIG, "test-group");
props.put(ConsumerConfig.PARTITION_ASSIGNMENT_STRATEGY_CONFIG, WeightedPartitionAssignor.class.getName());
// 配置其他消费者属性

KafkaConsumer<String, String> consumer = new KafkaConsumer<>(props);
consumer.subscribe(Arrays.asList("topic1", "topic2"));