reactor.kafka.sender.internals.DefaultKafkaSender Maven / Gradle / Ivy

/*
 * Copyright (c) 2016-2023 VMware Inc. or its affiliates, All Rights Reserved.
 *
 * 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
 *
 *   https://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 reactor.kafka.sender.internals;

import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.reactivestreams.Publisher;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.CoreSubscriber;
import reactor.core.publisher.Flux;
import reactor.core.publisher.FluxOperator;
import reactor.core.publisher.Mono;
import reactor.core.publisher.SignalType;
import reactor.core.publisher.Sinks;
import reactor.core.publisher.Sinks.EmitFailureHandler;
import reactor.core.scheduler.Scheduler;
import reactor.core.scheduler.Schedulers;
import reactor.kafka.sender.KafkaOutbound;
import reactor.kafka.sender.KafkaSender;
import reactor.kafka.sender.SenderOptions;
import reactor.kafka.sender.SenderRecord;
import reactor.kafka.sender.SenderResult;
import reactor.kafka.sender.TransactionManager;

import java.lang.reflect.InvocationHandler;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Proxy;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.function.Function;

/**
 * Reactive producer that sends messages to Kafka topic partitions. The producer is thread-safe
 * and can be used to send messages to multiple partitions. It is recommended that a single
 * producer is shared for each message type in a client.
 *
 * @param  outgoing message key type
 * @param  outgoing message value type
 */
public class DefaultKafkaSender implements KafkaSender, EmitFailureHandler {

    static final Logger log = LoggerFactory.getLogger(DefaultKafkaSender.class.getName());

    /** Note: Methods added to this set should also be included in javadoc for {@link KafkaSender#doOnProducer(Function)} */
    private static final Set DELEGATE_METHODS = new HashSet<>(Arrays.asList(
            "sendOffsetsToTransaction",
            "partitionsFor",
            "metrics",
            "flush"
        ));

    private final Scheduler scheduler;
    private final Mono> producerMono;
    private final AtomicBoolean hasProducer;
    final SenderOptions senderOptions;
    private final TransactionManager transactionManager;
    private final String producerId;
    private Producer producerProxy;

    /**
     * Constructs a reactive Kafka producer with the specified configuration properties. All Kafka
     * producer properties are supported. The underlying Kafka producer is created lazily when required.
     */
    public DefaultKafkaSender(ProducerFactory producerFactory, SenderOptions options) {
        producerId =
            Optional.ofNullable(options.clientId())
                .filter(clientId -> !clientId.isEmpty())
                .orElse("reactor-kafka-sender-" + System.identityHashCode(this));

        this.scheduler =
            Schedulers.newSingle(r -> {
                Thread thread = new Thread(r);
                thread.setName(producerId);
                return thread;
            });
        this.hasProducer = new AtomicBoolean();
        this.senderOptions = options.scheduler(options.isTransactional()
                                        ? Schedulers.newSingle(options.transactionalId())
                                        : options.scheduler()
                                    );
        boolean transactional = this.senderOptions.isTransactional();
        this.producerMono = Mono
                .fromCallable(() -> {
                    Producer producer = producerFactory.createProducer(senderOptions);
                    SenderOptions.ProducerListener producerListener = senderOptions.producerListener();
                    if (producerListener != null) {
                        producerListener.producerAdded(producerId, producer);
                    }
                    if (transactional) {
                        log.info("Initializing transactions for producer {}",
                                senderOptions.transactionalId());
                        producer.initTransactions();
                    }
                    hasProducer.set(true);
                    return producer;
                })
                .publishOn(senderOptions.isTransactional() ? this.scheduler : senderOptions.scheduler())
                .cache()
            .as(flux ->
                senderOptions.isTransactional()
                    ? flux.publishOn(senderOptions.isTransactional() ? this.scheduler : senderOptions.scheduler())
                    : flux);

        if (transactional) {
            this.producerMono.subscribe().dispose();
        }
        this.transactionManager = transactional
            ? new DefaultTransactionManager<>(producerMono, senderOptions)
            : null;
    }

    @Override
    public  Flux> send(Publisher> records) {
        return doSend(records);
    }

     Flux> doSend(Publisher> records) {
        return producerMono
            .flatMapMany(producer -> {
                return Flux.from(records)
                    // Producer#send is blocking
                    .publishOn(scheduler)
                    .as(flux -> new FluxOperator, SenderResult>(flux) {
                        @Override
                        public void subscribe(CoreSubscriber> s) {
                            source.subscribe(new SendSubscriber<>(senderOptions, producer, producerId, s));
                        }
                    });
            })
            .doOnError(e -> log.trace("Send failed with exception", e))
            .publishOn(senderOptions.scheduler(), senderOptions.maxInFlight());
    }

    @Override
    public KafkaOutbound createOutbound() {
        return new DefaultKafkaOutbound<>(this);
    }

    @Override
    public  Flux>> sendTransactionally(Publisher>> transactionRecords) {
        Sinks.Many