com.opentable.kafka.builders.KafkaProducerBaseBuilder Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of otj-kafka Show documentation
Kafka integrations component
The newest version!
/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.opentable.kafka.builders;

import java.time.Duration;
import java.util.Map;
import java.util.Optional;
import java.util.OptionalInt;
import java.util.OptionalLong;

import com.codahale.metrics.MetricRegistry;

import org.apache.kafka.clients.CommonClientConfigs;
import org.apache.kafka.clients.producer.Partitioner;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerInterceptor;
import org.apache.kafka.clients.producer.internals.DefaultPartitioner;
import org.apache.kafka.common.serialization.Serializer;

import com.opentable.kafka.logging.LoggingProducerInterceptor;

@SuppressWarnings("OptionalUsedAsFieldOrParameterType")
public abstract class KafkaProducerBaseBuilder, K, V> extends KafkaBaseBuilder {

    private Optional ackType = Optional.empty();
    private OptionalInt retries = OptionalInt.empty();
    private OptionalInt batchSize = OptionalInt.empty();
    private OptionalInt maxInfFlight = OptionalInt.empty();
    private OptionalLong lingerMS = OptionalLong.empty();
    private OptionalLong bufferMemory = OptionalLong.empty();
    private Optional compressionType = Optional.empty();
    private Class partitioner = DefaultPartitioner.class;
    private Class> keySe;
    private Class> valueSe;
    protected Serializer keySerializer;
    protected Serializer valueSerializer;

    protected KafkaProducerBaseBuilder(Map props, EnvironmentProvider environmentProvider) {
        super(props, environmentProvider);
    }


    public SELF withAcks(KafkaProducerBuilder.AckType val) {
        this.ackType = Optional.ofNullable(val);
        return self();
    }

    public SELF withCompressionType(Compression val) {
        this.compressionType = Optional.ofNullable(val);
        return self();
    }

    public SELF withRetries(int val) {
        this.retries = OptionalInt.of(val);
        return self();
    }

    public SELF withMaxInFlightRequests(int val) {
        this.maxInfFlight = OptionalInt.of(val);
        return self();
    }

    public SELF withBatchSize(int val) {
        this.batchSize = OptionalInt.of(val);
        return self();
    }

    /**
     * Provide an class. If you have a no-args constructor use this
     * @param keySer key serializer
     * @param valSer value serializer
     * @param  The type of the key expected by serializer
     * @param  The type of the value expected by serializer
     * @return this
     */
    @SuppressWarnings("unchecked")
    protected  KafkaProducerBaseBuilder withSerializers(Class> keySer, Class> valSer) {
        KafkaProducerBaseBuilder res = (KafkaProducerBaseBuilder) this;
        res.keySe =  keySer;
        res.valueSe =  valSer;
        res.keySerializer = null;
        res.valueSerializer = null;
        return res;
    }

    /**
     * Provide an instance. If you don't have a no-args constructor use this
     * @param keySer key serializer
     * @param valSer value serializer
     * @param  The type of the key expected by serializer
     * @param  The type of the value expected by serializer
     * @return this
     */
    @SuppressWarnings("unchecked")
    protected  KafkaProducerBaseBuilder withSerializers(Serializer keySer, Serializer valSer) {
        KafkaProducerBaseBuilder res = (KafkaProducerBaseBuilder) this;
        res.keySerializer = keySer;
        res.valueSerializer = valSer;
        res.keySe = null;
        res.valueSe = null;
        return res;
    }

    public SELF withPartitioner(Class partitioner) {
        if (partitioner != null) {
            this.partitioner = partitioner;
        }
        return self();
    }


    public SELF withLingerMs(Duration duration) {
        if (duration != null) {
            lingerMS = OptionalLong.of(duration.toMillis());
        } else {
            lingerMS = OptionalLong.empty();
        }
        return self();
    }

    public SELF withBufferMemory(long size) {
        this.bufferMemory = OptionalLong.of(size);
        return self();
    }


    public SELF withMetricRegistry(MetricRegistry metricRegistry, String metricsPrefix) {
        this.withMetricRegistry(metricRegistry);
        this.withPrefix(metricsPrefix);
        return self();
    }

    public SELF withInterceptor(Class> clazz) {
        this.addInterceptor(clazz.getName());
        return self();
    }

    protected void internalBuild() {
        this.addProperty(ProducerConfig.PARTITIONER_CLASS_CONFIG, partitioner.getName());
        this.setupInterceptors(ProducerConfig.INTERCEPTOR_CLASSES_CONFIG, LoggingProducerInterceptor.class.getName());
        maxInfFlight.ifPresent(m -> this.addProperty(ProducerConfig.MAX_IN_FLIGHT_REQUESTS_PER_CONNECTION, m));
        bufferMemory.ifPresent(b -> this.addProperty(ProducerConfig.BUFFER_MEMORY_CONFIG, b));
        batchSize.ifPresent(b -> this.addProperty(ProducerConfig.BATCH_SIZE_CONFIG, b));
        lingerMS.ifPresent(l -> this.addProperty(ProducerConfig.LINGER_MS_CONFIG, l));
        ackType.ifPresent(ack -> this.addProperty(ProducerConfig.ACKS_CONFIG, ack.value));
        retries.ifPresent(retries -> this.addProperty(CommonClientConfigs.RETRIES_CONFIG, retries));
        compressionType.ifPresent(ct -> this.addProperty(ProducerConfig.COMPRESSION_TYPE_CONFIG, ct.getKafkaCompressionType()));

        if (keySe != null) {
            this.addProperty(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, keySe.getName());
        }
        if (valueSe != null) {
            this.addProperty(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, valueSe.getName());
        }

        // merge in common and seed properties
        this.finishBuild();
    }

    public Map buildProperties() {
        internalBuild();
        return this.getFinalProperties();
    }

    public enum AckType {
        all("all"),
        none("0"),
        atleastOne("1");

        final String value;

        AckType(String value) {
            this.value = value;
        }
    }

    public enum Compression {
        NONE(org.apache.kafka.common.record.CompressionType.NONE),
        SNAPPY(org.apache.kafka.common.record.CompressionType.SNAPPY),
        LZ4(org.apache.kafka.common.record.CompressionType.LZ4),
        ZSTD(org.apache.kafka.common.record.CompressionType.ZSTD);

        private final org.apache.kafka.common.record.CompressionType compressionType;

        Compression(org.apache.kafka.common.record.CompressionType compressionType) {
            this.compressionType = compressionType;
        }

        public String getKafkaCompressionType() {
            return compressionType.name;
        }
    }
}