kafka组件traceId增强

问题背景

kafka 组件在排查生产者和消费者问题的时候，经常出现日志不匹配的情况，为了解决这个问题，本文实现 traceId 在kafka 组件的生产者和消费者之间传递，达到日志匹配，快速排查问题目的。

实现方案

1、新增生产者拦截器，读取 MDC 中的 traceId 并放入 kafka 消息的 headers 中；

2、新增消费者拦截器，读取 kafka 消息 headers 中的 traceId，覆写 MDC 的 traceId。

新增两个拦截器的方案对代码几乎没有侵入性，很优雅的解决了我们的问题。

生产者拦截器

1、kafka 消息发送前，进入到生产者拦截器，如果 MDC.get("trace_id") 存在值，则读取当前 MDC 中的 traceId 并放入消息的 headers 中；

2、如果 MDC.get("trace_id") 不存在，则生成一个 traceId，并写入消息的 headers 中。

java 复制代码

import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.producer.ProducerInterceptor;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.apache.kafka.common.header.Header;
import org.slf4j.MDC;
import org.springframework.stereotype.Component;

import java.nio.charset.StandardCharsets;
import java.util.Map;
import java.util.Optional;
import java.util.UUID;

@Slf4j
@Component
public class KafkaTraceIdProducerInterceptor implements ProducerInterceptor<Integer, String> {

    public static final String TRACE_ID = "trace_id";
    @Override
    public ProducerRecord<Integer, String> onSend(ProducerRecord<Integer, String> producerRecord) {
        producerRecord.headers()
                .add(new Header() {
                    @Override
                    public String key() {
                        return TRACE_ID;
                    }

                    @Override
                    public byte[] value() {
                        return getOrGenerateTraceId()
                                .getBytes(StandardCharsets.UTF_8);
                    }
                });
        return producerRecord;
    }

    @Override
    public void onAcknowledgement(RecordMetadata metadata, Exception exception) {
    }

    @Override
    public void close() {
    }

    @Override
    public void configure(Map<String, ?> configs) {
    }

    /**
     * 获取当前请求TraceId或生成新的traceId
     *
     * @return traceId
     */
    public static String getOrGenerateTraceId() {
        return Optional.ofNullable(MDC.get(TRACE_ID))
                .orElseGet(UUID.randomUUID()::toString);
    }
}

消费者拦截器

读取消息 headers 中的 traceId，覆写 MDC 的 traceId，消息消费完成后清理相关的 MDC。

java 复制代码

import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.common.header.Header;
import org.slf4j.MDC;
import org.springframework.kafka.listener.RecordInterceptor;
import org.springframework.stereotype.Component;

import java.nio.charset.StandardCharsets;
import java.util.UUID;

@Slf4j
@Component
public class KafkaConsumerInterceptor<K, V> implements RecordInterceptor<K, V> {

    private final static String TRACE_ID = "trace_id";


    @Override
    public ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> consumerRecord) {
        return consumerRecord;
    }

    @Override
    public ConsumerRecord<K, V> intercept(ConsumerRecord<K, V> record, Consumer<K, V> consumer) {
        try {
            Header traceHeader = record.headers().lastHeader(TRACE_ID);
            String traceId = (traceHeader != null && traceHeader.value() != null)
                    ? new String(traceHeader.value(), StandardCharsets.UTF_8)
                    : UUID.randomUUID().toString();
            MDC.put(TRACE_ID, traceId);
        } catch (Exception e) {
            log.error("处理Kafka消息头异常", e);
        }
        return record;
    }

    @Override
    public void success(ConsumerRecord<K, V> record, Consumer<K, V> consumer) {
        MDC.clear();
    }

    @Override
    public void failure(ConsumerRecord<K, V> record, Exception exception, Consumer<K, V> consumer) {
        MDC.clear();
    }

}

问题延伸

如果遇到需要打印生产者和消费者接口日志或者遇到安全问题，生产者消息加密，消费者消息解密，都可以通过上述增加生产者拦截器和消费者拦截器来解决问题。