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

Show more of this group Show more artifacts with this name
Show all versions of reactor-kafka Show documentation
Reactor Kafka: A reactive API for Apache Kafka
There is a newer version: 1.3.23
/*
 * Copyright (c) 2020-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 io.micrometer.observation.contextpropagation.ObservationThreadLocalAccessor;
import org.apache.kafka.clients.producer.Callback;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.reactivestreams.Subscription;
import reactor.core.CoreSubscriber;
import reactor.core.publisher.Operators;
import reactor.kafka.sender.SenderOptions;
import reactor.kafka.sender.SenderRecord;
import reactor.kafka.sender.SenderResult;
import reactor.kafka.sender.observation.KafkaRecordSenderContext;
import reactor.kafka.sender.observation.KafkaSenderObservation;
import reactor.util.context.Context;

import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.atomic.AtomicReference;

/**
 * This is basically an optimized flatMapDelayError(Function<ProducerRecord,Mono<SenderResult>>), without prefetching
 * and with an unlimited concurrency (implicitly limited by {@link Producer#send(ProducerRecord)}).
 * 
 * The requests are passed to the upstream "as is".
 *
 */
class SendSubscriber implements CoreSubscriber> {

    enum State {
        INIT,
        ACTIVE,
        INBOUND_DONE,
        COMPLETE
    }

    private final CoreSubscriber> actual;
    private final Producer producer;
    private final AtomicInteger inflight = new AtomicInteger();
    private final AtomicReference firstException = new AtomicReference<>();
    private final AtomicReference state = new AtomicReference<>(State.INIT);
    private final SenderOptions senderOptions;

    private final String producerId;

    SendSubscriber(SenderOptions senderOptions, Producer producer, String producerId,
                   CoreSubscriber> actual) {

        this.senderOptions = senderOptions;
        this.producer = producer;
        this.producerId = producerId;
        this.actual = actual;
    }

    @Override
    public Context currentContext() {
        return actual.currentContext();
    }

    @Override
    public void onSubscribe(Subscription s) {
        state.set(State.ACTIVE);
        actual.onSubscribe(s);
    }

    @Override
    public void onNext(ProducerRecord record) {
        if (state.get() == State.COMPLETE) {
            Operators.onNextDropped(record, currentContext());
            return;
        }
        inflight.incrementAndGet();

        if (senderOptions.isTransactional()) {
            DefaultKafkaSender.log.trace("Transactional send initiated for producer {} in state {} inflight {}: {}", senderOptions.transactionalId(), state, inflight, record);
        }

        @SuppressWarnings("unchecked")
        C correlationMetadata = record instanceof SenderRecord
            ? ((SenderRecord) record).correlationMetadata()
            : null;

        Callback callback = (metadata, exception) -> {
            if (senderOptions.isTransactional()) {
                DefaultKafkaSender.log.trace("Transactional send completed for producer {} in state {} inflight {}: {}", senderOptions.transactionalId(), state, inflight, record);
            }

            if (state.get() == State.COMPLETE) {
                return;
            }

            if (exception != null) {
                DefaultKafkaSender.log.trace("Sender failed: ", exception);
                firstException.compareAndSet(null, exception);
                if (senderOptions.stopOnError() || senderOptions.fatalException(exception)) {
                    onError(exception);
                    return;
                }
            }

            actual.onNext(new Response<>(metadata, exception, correlationMetadata));
            if (inflight.decrementAndGet() == 0) {
                maybeComplete();
            }
        };
        try {
            KafkaSenderObservation.SENDER_OBSERVATION.observation(senderOptions.observationConvention(),
                    KafkaSenderObservation.DefaultKafkaSenderObservationConvention.INSTANCE,
                    () -> new KafkaRecordSenderContext(record,
                        producerId,
                        senderOptions.bootstrapServers()),
                    senderOptions.observationRegistry())
                .parentObservation(currentContext().getOrDefault(ObservationThreadLocalAccessor.KEY, null))
                .observe(() -> producer.send(record, callback));
        } catch (Exception e) {
            callback.onCompletion(null, e);
        }
    }

    @Override
    public void onError(Throwable t) {
        DefaultKafkaSender.log.trace("Sender failed with exception", t);
        if (state.getAndSet(State.COMPLETE) == State.COMPLETE) {
            Operators.onErrorDropped(t, currentContext());
            return;
        }

        actual.onError(t);
    }

    @Override
    public void onComplete() {
        if (state.compareAndSet(State.ACTIVE, State.INBOUND_DONE)) {
            if (inflight.get() == 0) {
                maybeComplete();
            }
        }
    }

    private void maybeComplete() {
        if (state.compareAndSet(State.INBOUND_DONE, State.COMPLETE)) {
            Throwable exception = firstException.get();
            if (exception != null) {
                actual.onError(exception);
            } else {
                actual.onComplete();
            }
        }
    }

}