Java 中间件Kafka 分区策略(自定义分区器实现负载均衡)

public interface Partitioner extends Configurable, Closeable {
    int partition(String topic, Object key, byte[] keyBytes,
                  Object value, byte[] valueBytes, Cluster cluster);
    void close();
    void configure(Map<String, ?> configs);
}

<dependency>
    <groupId>org.apache.kafka</groupId>
    <artifactId>kafka-clients</artifactId>
    <version>3.6.0</version>
</dependency>

import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
public class UserIdPartitioner implements Partitioner {
    // VIP 用户列表（可从配置或数据库加载）
    private static final Set<String> VIP_USERS = Set.of("1001", "2005", "9999");
    // 缓存 VIP 用户的专属分区，避免频繁计算
    private final Map<String, Integer> vipPartitionCache = new ConcurrentHashMap<>();
    @Override
    public void configure(Map<String, ?> configs) {
        // 可在此处读取自定义配置，如 VIP 列表、分区偏移量等
        System.out.println("UserIdPartitioner configured.");
    }
    @Override
    public int partition(String topic, Object key, byte[] keyBytes,
                         Object value, byte[] valueBytes, Cluster cluster) {
        // 获取该 topic 的所有分区信息
        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
        int numPartitions = partitions.size();
        if (key == null) {
            // 无 key 时，使用轮询（简化版）
            return Math.abs(key.hashCode()) % numPartitions;
        }
        String userId = key.toString();
        if (VIP_USERS.contains(userId)) {
            // VIP 用户：固定分配到前 N 个分区（例如前 2 个）
            // 为每个 VIP 用户分配唯一分区，避免冲突
            return vipPartitionCache.computeIfAbsent(userId, k -> {
                int index = VIP_USERS.stream().toList().indexOf(k);
                return index % Math.min(2, numPartitions); // 最多使用 2 个分区
            });
        } else {
            // 普通用户：使用标准哈希
            return Math.abs(userId.hashCode()) % numPartitions;
        }
    }
    @Override
    public void close() {
        vipPartitionCache.clear();
        System.out.println("UserIdPartitioner closed.");
    }
}

Properties props = new Properties();
props.put("bootstrap.servers", "localhost:9092");
props.put("key.serializer", "org.apache.kafka.common.serialization.StringSerializer");
props.put("value.serializer", "org.apache.kafka.common.serialization.StringSerializer");
// 关键配置：指定自定义分区器
props.put("partitioner.class", "com.example.UserIdPartitioner");
KafkaProducer<String, String> producer = new KafkaProducer<>(props);
// 发送消息
producer.send(new ProducerRecord<>("user-events", "1001", "VIP login"));
producer.send(new ProducerRecord<>("user-events", "5001", "Normal user action"));

public class LoadAwarePartitioner implements Partitioner {
    private final Map<Integer, Long> partitionLoad = new ConcurrentHashMap<>();
    private final Random random = new Random();
    @Override
    public void configure(Map<String, ?> configs) {}
    @Override
    public int partition(String topic, Object key, byte[] keyBytes,
                         Object value, byte[] valueBytes, Cluster cluster) {
        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
        int numPartitions = partitions.size();
        if (key == null) {
            // 无 key：选择负载最低的分区
            return findLeastLoadedPartition(numPartitions);
        }
        // 有 key：仍需保证相同 key 进同一分区（顺序性）
        // 但我们可以记录该分区的负载
        int targetPartition = Math.abs(key.hashCode()) % numPartitions;
        updateLoad(targetPartition);
        return targetPartition;
    }
    private int findLeastLoadedPartition(int numPartitions) {
        long minLoad = Long.MAX_VALUE;
        int bestPartition = 0;
        for (int i = 0; i < numPartitions; i++) {
            long load = partitionLoad.getOrDefault(i, 0L);
            if (load < minLoad) {
                minLoad = load;
                bestPartition = i;
            }
        }
        // 更新负载（模拟）
        updateLoad(bestPartition);
        return bestPartition;
    }
    private void updateLoad(int partition) {
        partitionLoad.compute(partition, (k, v) -> (v == null) ? 1L : v + 1);
    }
    @Override
    public void close() {
        partitionLoad.clear();
    }
}

// ❌ 错误！相同 key 可能进入不同分区
public int partition(...) {
    if (isVip(key)) {
        return random.nextInt(numPartitions); // 随机分配 VIP
    }
    return hash(key) % numPartitions;
}

// ✅ 正确：相同 key 始终进入同一分区
public int partition(...) {
    if (isVip(key)) {
        // 所有 VIP 用户进入分区 0
        return 0;
    }
    return hash(key) % numPartitions;
}

// ✅ 使用 ConcurrentHashMap
private final Map<String, Integer> cache = new ConcurrentHashMap<>();
// ❌ 非线程安全
private Map<String, Integer> cache = new HashMap<>();

@Test
public void testVipUserRouting() {
    UserIdPartitioner partitioner = new UserIdPartitioner();
    partitioner.configure(Collections.emptyMap());
    Cluster cluster = mockClusterWithPartitions(4); // 模拟 4 分区集群
    int p1 = partitioner.partition("test", "1001", null, null, null, cluster);
    int p2 = partitioner.partition("test", "1001", null, null, null, cluster);
    assertEquals(p1, p2); // 同一 VIP 应进入同一分区
    assertTrue(p1 < 2);   // 且应在前 2 个分区
}

{
  "order_id": "O12345",
  "user_id": "U789",
  "region": "CN-EAST",
  "amount": 299.99
}

public class RegionAwareOrderPartitioner implements Partitioner {
    private static final Map<String, Integer> REGION_PARTITION_OFFSET = Map.of(
        "CN-EAST", 0,
        "CN-NORTH", 2,
        "CN-SOUTH", 4
    );
    private static final int PARTITIONS_PER_REGION = 2;
    @Override
    public int partition(String topic, Object key, byte[] keyBytes,
                         Object value, byte[] valueBytes, Cluster cluster) {
        // 假设 value 是 JSON 字符串
        String jsonValue = (String) value;
        String region = extractRegionFromJson(jsonValue); // 解析 region
        int basePartition = REGION_PARTITION_OFFSET.getOrDefault(region, 0);
        int totalPartitions = cluster.partitionsForTopic(topic).size();
        // 计算该 region 的可用分区范围
        int start = basePartition;
        int end = Math.min(basePartition + PARTITIONS_PER_REGION, totalPartitions);
        if (start >= totalPartitions) {
            // fallback to default
            return Math.abs(key.hashCode()) % totalPartitions;
        }
        // 在 region 内部按 user_id 哈希
        int userHash = Math.abs(key.hashCode());
        int regionPartition = start + (userHash % (end - start));
        return regionPartition;
    }
    private String extractRegionFromJson(String json) {
        // 简化：实际应使用 JSON 解析库
        int start = json.indexOf("\"region\":\"") + 11;
        int end = json.indexOf("\"", start);
        return json.substring(start, end);
    }
    @Override
    public void configure(Map<String, ?> configs) {}
    @Override
    public void close() {}
}

渲染错误: Mermaid 渲染失败: Parsing failed: unexpected character: ->“<- at offset: 29, skipped 6 characters. unexpected character: ->：<- at offset: 36, skipped 1 characters. unexpected character: ->“<- at offset: 44, skipped 6 characters. unexpected character: ->：<- at offset: 51, skipped 1 characters. unexpected character: ->“<- at offset: 59, skipped 6 characters. unexpected character: ->：<- at offset: 66, skipped 1 characters. Expecting token of type 'EOF' but found `2`. Expecting token of type 'EOF' but found `2`. Expecting token of type 'EOF' but found `2`.

kafka-run-class.sh kafka.tools.GetOffsetShell \
  --broker-list localhost:9092 \
  --topic user-events

user-events:0:15000
user-events:1:5000
user-events:2:5000