使用Flux groupBy为每个数据批次生成新traceId的实现方案

我想实现基于Reactive Kafka的消费者，参考分区有序并发处理的写法：

Scheduler scheduler = Schedulers.newElastic("sample", 60, true);
KafkaReceiver.create(receiverOptions)
         .receive()
         .groupBy(m -> m.receiverOffset().topicPartition())
         .flatMap(partitionFlux ->
             partitionFlux.publishOn(scheduler)
                          .map(r -> processRecord(partitionFlux.key(), r))
                          .sample(Duration.ofMillis(5000))
                          .concatMap(offset -> offset.commit()));

现在需要对每次poll操作做链路追踪，但用groupBy时，每个分组只会分配一次traceId，不符合需求。

我写了不含Kafka的测试代码验证这个问题：

Flux
    .interval(Duration.ofMillis(100)) // 模拟kafka消息
    .filter {
        it % 2L == 0L // 简化为只保留一个分组
    }
    .groupBy {
        it % 2L
    }
//    .windowTimeout(2, Duration.ofMillis(10))
    .flatMap { group ->
        val traceId1 = UUID.randomUUID().toString() // 用于区分的两个traceId
        Flux.deferContextual {
            Mono.just(it.get<String>("traceId2")) 
        }.flatMap { traceId2 ->
            group.bufferTimeout(2, Duration.ofMillis(10))
                .concatMap {
                    // 处理批次
                    Mono.just(it).delayElement(Duration.ofMillis(50))
                }.flatMap {
                    // 提交操作
                    println("$it - $traceId1 - $traceId2")
                    Mono.just(it)
                }
        }
            .contextWrite {
                it.put("traceId2", UUID.randomUUID().toString())
            }
    }
    .blockLast()

输出结果里traceId1和traceId2始终不变：

[0] - d69da7b2-205b-4742-b538-308192af29d6 - 529c50a5-4966-4b52-bd63-18d4b879d8d4
[2] - d69da7b2-205b-4742-b538-308192af29d6 - 529c50a5-4966-4b52-bd63-18d4b879d8d4
[4] - d69da7b2-205b-4742-b538-308192af29d6 - 529c50a5-4966-4b52-bd63-18d4b879d8d4
[6] - d69da7b2-205b-4742-b538-308192af29d6 - 529c50a5-4966-4b52-bd63-18d4b879d8d4
[8] - d69da7b2-205b-4742-b538-308192af29d6 - 529c50a5-4966-4b52-bd63-18d4b879d8d4
[10] - d69da7b2-205b-4742-b538-308192af29d6 - 529c50a5-4966-4b52-bd63-18d4b879d8d4

但改用windowTimeout分组后，每个批次的traceId都不同，符合预期：

Flux
    .interval(Duration.ofMillis(100)) // 模拟kafka消息
    .filter {
        it % 2L == 0L // 简化为只保留一个分组
    }
/*
    .groupBy {
        it % 2L
    }
*/
    .windowTimeout(2, Duration.ofMillis(10))
    .flatMap { group ->
        val traceId1 = UUID.randomUUID().toString() // 用于区分的两个traceId
        Flux.deferContextual {
            Mono.just(it.get<String>("traceId2")) 
        }.flatMap { traceId2 ->
            group.bufferTimeout(2, Duration.ofMillis(10))
                .concatMap {
                    // 处理批次
                    Mono.just(it).delayElement(Duration.ofMillis(50))
                }.flatMap {
                    // 提交操作
                    println("$it - $traceId1 - $traceId2")
                    Mono.just(it)
                }
        }
            .contextWrite {
                it.put("traceId2", UUID.randomUUID().toString())
            }
    }
    .blockLast()

输出：

[0] - d571b08e-d14c-425e-9062-49de555b6b6e - edef203f-4727-4b9b-bc27-82743e8c16f3
[2] - 8b1b2998-ff83-4b09-990c-93b3ad8e386b - e4b00092-e4b7-4f54-aa8a-90ac4882dea3
[4] - 323450bf-6d2c-4e39-8534-7d7ea63fbc31 - 1be0022c-88f4-4339-a00f-7e9f4f5ed45d
[6] - 8db7d902-5047-4772-9ebb-9b8bef983126 - ce92b804-ac72-4af1-b353-5c7576d71a01
[8] - 838e230c-fdb2-4e33-901e-96b7cbf2543e - f355382e-54ef-4a02-b2ed-a32bf855563c
[10] - d8bfc951-98a1-428a-b697-ebc447a78a40 - 916d2824-4e76-47fb-8573-653e954841eb

请问怎么用groupBy实现每个批次生成新traceId的需求？或者有没有groupBy的替代方案？

核心问题分析 #

groupBy会为每个分组生成一个持久的Flux，你在flatMap外层生成traceId1、设置上下文的操作只会执行一次（每个分组初始化时），所以后续所有批次都会复用同一个traceId。而windowTimeout是每次窗口触发时生成新的Flux，所以每次都会重新生成traceId。

基于groupBy的改造方案 #

要在groupBy的流里为每个批次生成新traceId，需要把traceId的生成逻辑放到批次处理的流内部，而不是flatMap的外层。具体来说：

1. 把traceId1的生成移到bufferTimeout之后的操作里，确保每个批次都生成新的ID；
2. 上下文的设置也要绑定到每个批次的流，而不是整个分组的流。

改造后的测试代码：

Flux
    .interval(Duration.ofMillis(100)) // 模拟kafka消息
    .filter { it % 2L == 0L }
    .groupBy { it % 2L }
    .flatMap { group ->
        group.bufferTimeout(2, Duration.ofMillis(10))
            // 为每个批次生成新的traceId1
            .flatMap { batch ->
                val traceId1 = UUID.randomUUID().toString()
                // 为当前批次设置新的上下文traceId2
                Mono.just(batch)
                    .contextWrite { it.put("traceId2", UUID.randomUUID().toString()) }
                    .delayElement(Duration.ofMillis(50)) // 模拟处理
                    .flatMap {
                        Flux.deferContextual { ctx ->
                            val traceId2 = ctx.get<String>("traceId2")
                            println("$it - $traceId1 - $traceId2")
                            Mono.just(it)
                        }
                    }
            }
    }
    .blockLast()

适配Reactive Kafka的写法 #

回到你的Kafka消费者场景，需要把traceId的生成逻辑放到sample（或你用来触发批次的操作）之后，确保每次poll/批次处理都生成新的traceId：

Scheduler scheduler = Schedulers.newElastic("sample", 60, true);
KafkaReceiver.create(receiverOptions)
         .receive()
         .groupBy(m -> m.receiverOffset().topicPartition())
         .flatMap(partitionFlux ->
             partitionFlux.publishOn(scheduler)
                          .map(r -> processRecord(partitionFlux.key(), r))
                          .sample(Duration.ofMillis(5000))
                          // 为每个批次生成新traceId
                          .flatMap(offsetBatch -> {
                              String traceId = UUID.randomUUID().toString();
                              // 绑定traceId到上下文或直接传入处理逻辑
                              return processWithTrace(offsetBatch, traceId)
                                      .concatMap(offset -> offset.commit());
                          }));

替代方案：groupBy+windowTimeout组合 #

如果你需要同时保留分区顺序和批次级的traceId，可以在groupBy之后对每个分组的流再应用windowTimeout，这样既保证了分区内的顺序，又能为每个窗口（批次）生成新的traceId：

Flux
    .interval(Duration.ofMillis(100))
    .filter { it % 2L == 0L }
    .groupBy { it % 2L }
    .flatMap { group ->
        group.windowTimeout(2, Duration.ofMillis(10))
            .flatMap { window ->
                val traceId1 = UUID.randomUUID().toString()
                Flux.deferContextual { ctx ->
                    val traceId2 = ctx.get<String>("traceId2")
                    window.buffer()
                        .delayElement(Duration.ofMillis(50))
                        .flatMap {
                            println("$it - $traceId1 - $traceId2")
                            Mono.just(it)
                        }
                }
                .contextWrite { it.put("traceId2", UUID.randomUUID().toString()) }
            }
    }
    .blockLast()

这种方式既保留了groupBy的分区有序性，又实现了每个批次的traceId独立。

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