Understanding reactor + reactor-kafka instrumentations

[galbash]

Hi all, I am trying to instrument a reactor kafka pipeline and have a few questions:

Connecting child spans of operations to the correct kafka consume span

My initial code looks something like this:

    private Mono<MessageBuilder<InputPojo>> mySingleAsyncOperation(MessageBuilder<InputPojo> msgs) {
        HttpClient client = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).build();
        HttpRequest request = HttpRequest.newBuilder(URI.create("http://localhost:8080/employees")).GET().build();
        CompletableFuture<HttpResponse<String>> response = client.sendAsync(request,HttpResponse.BodyHandlers.ofString());
        return Mono.fromFuture(response).map(responses -> msgs);
    }
    
 kafkaReceiver.receive().publishOn(Schedulers.parallel())
                .doOnNext(consumerRecord -> log.info("received key={}, value={} from topic={}, offset={}",
                        consumerRecord.key(),
                        consumerRecord.value(),
                        consumerRecord.topic(),
                        consumerRecord.offset())
                )
                .map(ConsumerRecord::value)
                .doOnNext(fakeConsumerDTO -> log.info("successfully consumed {}={}", InputPojo.class.getSimpleName(), fakeConsumerDTO))
                .map(msg -> {
                    try {
                        return MessageBuilder.withPayload(objectMapper.readValue(msg, InputPojo.class));
                    } catch (JsonProcessingException e) {
                        throw new RuntimeException(e);
                    }
                })
                .concatMap(this::mySingleAsyncOperation)
                .doOnError(throwable -> log.error("something bad happened while consuming : {}", throwable.getMessage()))
                .subscribe();

My ideal result would be: for each messages that passes through the kafka a span reporting reading from kafka, and a child span to the correct kafka span that represent the http client request.
However when I run the application with a lot of messages in the Kafka, the child spans are randomly aggregated to be children of some of the spans that represent the kafka consumer, without a correct correlation between them.
Am I doing something wrong? is there a way to structure my code so my operations will be associated with the correct kafka consumer span?

Working with buffers

As a next step, I am performing some batch aggregations, and the code looks something like this:

@Slf4j
@Service
@RequiredArgsConstructor
public class Processor implements ApplicationRunner {
    private final KafkaReceiver<String, String> kafkaReceiver;
    private final ObjectMapper objectMapper = new ObjectMapper();

    private Mono<List<MessageBuilder<InputPojo>>> myAsyncOperation(List<MessageBuilder<InputPojo>> msgs) {
        HttpClient client = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).build();
        HttpRequest request = HttpRequest.newBuilder(URI.create("http://localhost:8080/employees")).GET().build();
        CompletableFuture<HttpResponse<String>> response = client.sendAsync(request,HttpResponse.BodyHandlers.ofString());
        return Mono.fromFuture(response).map(responses -> msgs);
    }
    private Mono<MessageBuilder<InputPojo>> mySingleAsyncOperation(MessageBuilder<InputPojo> msgs) {
        HttpClient client = HttpClient.newBuilder().version(HttpClient.Version.HTTP_2).build();
        HttpRequest request = HttpRequest.newBuilder(URI.create("http://localhost:8080/employees")).GET().build();
        CompletableFuture<HttpResponse<String>> response = client.sendAsync(request,HttpResponse.BodyHandlers.ofString());
        return Mono.fromFuture(response).map(responses -> msgs);
    }

    @Override
    public void run(ApplicationArguments args) throws Exception {
        kafkaReceiver.receive()//.publishOn(Schedulers.parallel())
                //.retryWhen(Retry.max(3))
                .doOnNext(consumerRecord -> log.info("received key={}, value={} from topic={}, offset={}",
                        consumerRecord.key(),
                        consumerRecord.value(),
                        consumerRecord.topic(),
                        consumerRecord.offset())
                )
                .map(ConsumerRecord::value)
                .doOnNext(fakeConsumerDTO -> log.info("successfully consumed {}={}", InputPojo.class.getSimpleName(), fakeConsumerDTO))
                .map(msg -> {
                    try {
                        return MessageBuilder.withPayload(objectMapper.readValue(msg, InputPojo.class));
                    } catch (JsonProcessingException e) {
                        throw new RuntimeException(e);
                    }
                })
                .concatMap(this::mySingleAsyncOperation)
                .buffer(5)
                .concatMap(this::myAsyncOperation)
                .doOnError(throwable -> log.error("something bad happened while consuming : {}", throwable.getMessage()))
                .subscribe();
    }
}

For me, the ideal result here would be two types of traces:

1. a trace that is similar to the one in the first section (a kafka consume span -> an http client span)
2. a trace representing the batched operation, which has one span of an http client from the myAsyncOperation function, with links to all the traces that was buffered in the buffer.
   I know that for this one I would have to do some manual instrumentation, and I was looking for some guidance. My best attempt so far was to add a map step before the buffer to add the current SpanContext to the messages and then create a span after the buffer named "batch" with links to the different contexts in the buffer, but again the async operation got associated with randomly wrong batch spans, not the correct ones (may be the same behavior as in the first section)

If anyone can shed some light on how the reactor instrumentation is expected to work and clues to help me get to the desired behavior i'd be grateful!
Here is a repository with a full reproduction + automatic setup of Kafka & Jaeger for your convenience:
https://github.com/almogtavor/spring-boot-reactor-kafka-example