com.hframework.common.springext.kafkatemplate.DynamicKafkaTemplate Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of hframe-common Show documentation
There is a newer version: 1.0.2
package com.hframework.common.springext.kafkatemplate;

import org.apache.kafka.clients.producer.Callback;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.clients.producer.RecordMetadata;
import org.springframework.kafka.core.KafkaOperations;
import org.springframework.kafka.core.KafkaProducerException;
import org.springframework.kafka.support.LoggingProducerListener;
import org.springframework.kafka.support.ProducerListener;
import org.springframework.kafka.support.SendResult;
import org.springframework.util.concurrent.ListenableFuture;
import org.springframework.util.concurrent.SettableListenableFuture;

import java.util.HashMap;
import java.util.Map;

/**
 * 重写一个动态的kafkaTemplate
 * 注意：DynamicKafkaTemplate与DynamicKafkaProducerFactory需要同时使用
 * Created by zhangquanhong on 2016/10/11.
 */
public class DynamicKafkaTemplate extends org.springframework.kafka.core.KafkaTemplate implements KafkaOperations {

    private DynamicKafkaProducerFactory defaultProducerFactory;
    private  boolean autoFlush;
    private volatile ProducerListener producerListener = new LoggingProducerListener();

    private volatile Map> producers = new HashMap>();

    public DynamicKafkaTemplate(DynamicKafkaProducerFactory producerFactory) {
        this(producerFactory, false);
    }

    public DynamicKafkaTemplate(DynamicKafkaProducerFactory producerFactory, boolean autoFlush) {
        super(producerFactory, autoFlush);
        defaultProducerFactory = producerFactory;
        this.autoFlush = autoFlush;
    }

    public ListenableFuture> send(String servers, String topic, V data) {
        ProducerRecord producerRecord = new ProducerRecord(topic, data);
        return doSend(servers, producerRecord);
    }

    public ListenableFuture> send(String servers, String topic, K key, V data) {
        ProducerRecord producerRecord = new ProducerRecord(topic, key, data);
        return doSend(servers, producerRecord);
    }

    public ListenableFuture> send(String servers, String topic, int partition, V data) {
        ProducerRecord producerRecord = new ProducerRecord(topic, partition, null, data);
        return doSend(servers, producerRecord);
    }

    public ListenableFuture> send(String servers, String topic, int partition, K key, V data) {
        ProducerRecord producerRecord = new ProducerRecord(topic, partition, key, data);
        return doSend(servers, producerRecord);
    }

    /**
     * Send the producer record.
     * @param producerRecord the producer record.
     * @return a Future for the {@link RecordMetadata}.
     */
    protected ListenableFuture> doSend(String servers, final ProducerRecord producerRecord) {
        getTheProducer(servers);
        if (this.logger.isTraceEnabled()) {
            this.logger.trace("Sending: " + producerRecord);
        }
        final SettableListenableFuture> future = new SettableListenableFuture>();
        getTheProducer(servers).send(producerRecord, new Callback() {

            public void onCompletion(RecordMetadata metadata, Exception exception) {
                if (exception == null) {
                    future.set(new SendResult(producerRecord, metadata));
                    if (DynamicKafkaTemplate.this.producerListener != null
                            && DynamicKafkaTemplate.this.producerListener.isInterestedInSuccess()) {
                        DynamicKafkaTemplate.this.producerListener.onSuccess(producerRecord.topic(),
                                producerRecord.partition(), producerRecord.key(), producerRecord.value(), metadata);
                    }
                }
                else {
                    future.setException(new KafkaProducerException(producerRecord, "Failed to send", exception));
                    if (DynamicKafkaTemplate.this.producerListener != null) {
                        DynamicKafkaTemplate.this.producerListener.onError(producerRecord.topic(),
                                producerRecord.partition(), producerRecord.key(), producerRecord.value(), exception);
                    }
                }
            }

        });
        if (this.autoFlush) {
            flush();
        }
        if (this.logger.isTraceEnabled()) {
            this.logger.trace("Sent: " + producerRecord);
        }
        return future;
    }

    private Producer getTheProducer(String servers) {
        if(!producers.containsKey(servers)) {
            synchronized (this) {
                if(!producers.containsKey(servers)) {
                    producers.put(servers, createProducer(servers));
                }
            }
        }
        return producers.get(servers);
    }

    private Producer createProducer(String servers) {
        return defaultProducerFactory.createProducer(servers);
    }

    public void setProducerListener(ProducerListener producerListener) {
        this.producerListener = producerListener;
    }
}