springboot kafka 发送和接受消息 租户id 集中统一处理

发送消息端

写一个 ProducerInterceptor

public class TenantProducerInterceptor implements ProducerInterceptor {

    private Logger logger = LoggerFactory.getLogger(TenantProducerInterceptor.class);



    @Override
    public ProducerRecord onSend(ProducerRecord producerRecord) {
        TenantAcquireService tenantAcquireService = ApplicationContentHolder.applicationContext.getBean(TenantAcquireService.class);
        String moduleName = ApplicationContentHolder.applicationContext.getEnvironment().getProperty("spring.application.name","none");
        //获取租户id
        String tenantId = tenantAcquireService.acquire();
        if (tenantId == null) {
            throw new IllegalArgumentException("tenantId is null");
        }
        logger.info("producerInterceptor_tenantId:{}",tenantId);
       // 设置租户id
        producerRecord.headers().add(new RecordHeader("tenantId", tenantId.getBytes(StandardCharsets.UTF_8)));
        producerRecord.headers().add(new RecordHeader("moduleName", moduleName.getBytes(StandardCharsets.UTF_8)));
        return producerRecord;
    }

    @Override
    public void onAcknowledgement(RecordMetadata recordMetadata, Exception e) {

    }

    @Override
    public void close() {

    }

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

    }
}

配置ProducerInterceptor

application.properties 添加配置 spring.kafka.producer.interceptor.classes =com.unitechs.tenantacquire.kafka.TenantProducerInterceptor

spring容器注入ProducerFactory

private final KafkaProperties properties;

@Value(("${spring.kafka.producer.interceptor.classes:none}"))
private String interceptor;

public KafkaTenantConfig(KafkaProperties properties) {
    this.properties = properties;
}

@Bean
public ProducerFactory<?, ?> kafkaProducerFactory() {
    Map<String, Object> map = this.properties.buildProducerProperties();
    if(!interceptor.equals("none")){
        map.put(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, Arrays.asList(interceptor.split(",")));
    }
    DefaultKafkaProducerFactory<?, ?> factory = new DefaultKafkaProducerFactory<>(map
            );
    String transactionIdPrefix = this.properties.getProducer().getTransactionIdPrefix();
    if (transactionIdPrefix != null) {
        factory.setTransactionIdPrefix(transactionIdPrefix);
    }

    return factory;
}

发送端就完成了

消费端

创建代理类 实现对messagelistener消息处理增强

import org.springframework.kafka.config.ConcurrentKafkaListenerContainerFactory;
import org.springframework.kafka.config.KafkaListenerEndpoint;
import org.springframework.kafka.listener.ConcurrentMessageListenerContainer;
import org.springframework.kafka.listener.ContainerProperties;
import org.springframework.kafka.support.TopicPartitionOffset;

import java.util.Collection;

public class TenantConcurrentKafkaListenerContainerFactory<K, V> extends ConcurrentKafkaListenerContainerFactory<K, V> {


    @Override
    protected ConcurrentMessageListenerContainer<K, V> createContainerInstance(KafkaListenerEndpoint endpoint) {
        TopicPartitionOffset[] topicPartitions = endpoint.getTopicPartitionsToAssign();
        if (topicPartitions != null && topicPartitions.length > 0) {
            ContainerProperties properties = new ContainerProperties(topicPartitions);
            return new TenantConcurrentMessageListenerContainer<>(getConsumerFactory(), properties);
        } else {
            Collection<String> topics = endpoint.getTopics();
            if (!topics.isEmpty()) {
                ContainerProperties properties = new ContainerProperties(topics.toArray(new String[0]));
                return new TenantConcurrentMessageListenerContainer<>(getConsumerFactory(), properties);
            } else {
                ContainerProperties properties = new ContainerProperties(endpoint.getTopicPattern());
                return new TenantConcurrentMessageListenerContainer<>(getConsumerFactory(), properties);
            }
        }
    }

}

 

 

import org.springframework.kafka.core.ConsumerFactory;
import org.springframework.kafka.listener.ConcurrentMessageListenerContainer;
import org.springframework.kafka.listener.ContainerProperties;

import java.lang.reflect.Proxy;

public class TenantConcurrentMessageListenerContainer<K, V> extends ConcurrentMessageListenerContainer {
    /**
     * Construct an instance with the supplied configuration properties.
     * The topic partitions are distributed evenly across the delegate
     *
     * @param consumerFactory     the consumer factory.
     * @param containerProperties the container properties.
     */
    public TenantConcurrentMessageListenerContainer(ConsumerFactory consumerFactory, ContainerProperties containerProperties) {
        super(consumerFactory, containerProperties);
    }

    public void setupMessageListener(Object messageListener) {
//代理messageListener
        Object listener = Proxy.newProxyInstance(messageListener.getClass().getClassLoader(), messageListener.getClass().getInterfaces(), new TenantListenerProxy(messageListener));
        super.setupMessageListener(listener);
    }
}

 

 

public class TenantListenerProxy implements InvocationHandler {
    private Object messageListener;

    private Logger logger = LoggerFactory.getLogger(TenantListenerProxy.class);

    public TenantListenerProxy(Object messageListener) {

        this.messageListener = messageListener;
    }

    @Override
    public Object invoke(Object o, Method method, Object[] objects) throws Throwable {
        boolean onMessageMethod = isOnMessageMethod(o, method, objects);
        if (onMessageMethod) {
//设置租户id
            setTenantId(objects);
        }
        try {
            return method.invoke(messageListener, objects);
        } finally {
            if (onMessageMethod) {
                logger.info("清理当前线程的租户id");
//清除租户id
                TenantThreadHolder.clear();
            }
        }
    }

    public void setTenantId(Object[] objects) {
        MessageDto messageDto = getMessageDto(objects);
        logger.info("设置当前线程的租户信息 messageInfo:{}", messageDto);
        TenantThreadHolder.set(messageDto.tenantId);
    }

    public MessageDto getMessageDto(Object[] objects) {
        String tenantId = null;
        String topic = null;
        ConsumerRecord consumerRecord = null;
        for (Object object : objects) {
            if (object instanceof ConsumerRecord) {
                consumerRecord = (ConsumerRecord) object;
                break;
            }
        }
        if (consumerRecord == null) {
            return new MessageDto(null, null, null);
        }
        return new MessageDto(consumerRecord.topic(), getHeadValue(consumerRecord, "tenantId"), getHeadValue(consumerRecord, "moduleName"));
    }

    public String getHeadValue(ConsumerRecord consumerRecord, String headName) {
        Headers headers = consumerRecord.headers();
        Iterable<Header> headerIterable = headers.headers(headName);
        for (Header header : headerIterable) {
            byte[] value = header.value();
            return new String(value);
        }
        return null;
    }

    public boolean isOnMessageMethod(Object o, Method method, Object[] objects) {
        if (method.getName().equals("onMessage")) {
            for (Object object : objects) {
                if (object instanceof ConsumerRecord) {
                    return true;
                }
            }
            return false;
        }
        return false;
    }


    public static class MessageDto {
        private String topic;
        private String tenantId;
        private String moduleName;

        public MessageDto(String topic, String tenantId, String moduleName) {
            this.topic = topic;
            this.tenantId = tenantId;
            this.moduleName = moduleName;
        }

        public String getModuleName() {
            return moduleName;
        }

        public String getTopic() {
            return topic;
        }

        public String getTenantId() {
            return tenantId;
        }

        @Override
        public String toString() {
            return "{" +
                    "'topic':'" + topic + '\'' +
                    ", 'tenantId':'" + tenantId + '\'' +
                    ", 'moduleName':'" + moduleName + '\'' +
                    '}';
        }
    }

}

使我们的代理类生效 添加ConcurrentKafkaListenerContainerFactory到spring容器

@Bean
ConcurrentKafkaListenerContainerFactory<?, ?> kafkaListenerContainerFactory(ConcurrentKafkaListenerContainerFactoryConfigurer configurer, ObjectProvider<ConsumerFactory<Object, Object>> kafkaConsumerFactory) {
    ConcurrentKafkaListenerContainerFactory<Object, Object> factory = new TenantConcurrentKafkaListenerContainerFactory();
    configurer.configure(factory, (ConsumerFactory)kafkaConsumerFactory.getIfAvailable(() -> {
        return new DefaultKafkaConsumerFactory(this.properties.buildConsumerProperties());
    }));
    return factory;
}

 

这里的完毕了 可以启动试一下 可不可以使用