reactor.kafka.receiver.internals.CommittableBatch 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.receiver.internals;

import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.OffsetAndMetadata;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.MonoSink;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.SortedSet;
import java.util.TreeSet;

class CommittableBatch {

    private static final Logger log = LoggerFactory.getLogger(CommittableBatch.class);

    final Map consumedOffsets = new HashMap<>();
    final Map> uncommitted = new HashMap<>();
    final Map> deferred = new HashMap<>();
    private final Map latestOffsets = new HashMap<>();
    boolean outOfOrderCommits;
    private int batchSize;
    private List> callbackEmitters = new ArrayList<>();
    private int inPipeline;

    public synchronized int updateOffset(TopicPartition topicPartition, long offset) {
        log.trace("Update offset {}@{}", topicPartition, offset);
        if (this.outOfOrderCommits) {
            SortedSet uncommittedThisTP = this.uncommitted.get(topicPartition);
            if (uncommittedThisTP != null && uncommittedThisTP.contains(offset)) {
                SortedSet offsets = this.deferred.computeIfAbsent(topicPartition, tp -> new TreeSet<>());
                if (offsets.add(offset)) {
                    batchSize++;
                    this.inPipeline--;
                }
            } else {
                log.debug("No uncomitted offset for {}@{}, partition revoked?", topicPartition, offset);
            }
        } else if (!((Long) offset).equals(consumedOffsets.put(topicPartition, offset))) {
            batchSize++;
            this.inPipeline--;
        }
        return batchSize;
    }

    public synchronized void addCallbackEmitter(MonoSink emitter) {
        callbackEmitters.add(emitter);
    }

    public synchronized boolean isEmpty() {
        return batchSize == 0;
    }

    public synchronized int batchSize() {
        return batchSize;
    }

    public synchronized void addUncommitted(ConsumerRecords records) {
        if (this.outOfOrderCommits) {
            records.partitions().forEach(tp -> {
                SortedSet offsets = this.uncommitted.computeIfAbsent(tp, part -> new TreeSet<>());
                records.records(tp).forEach(rec -> offsets.add(rec.offset()));
            });
        }
        this.inPipeline += records.count();
    }

    public synchronized void partitionsRevoked(Collection revoked) {
        revoked.forEach(part -> {
            this.uncommitted.remove(part);
            this.deferred.remove(part);
        });
    }

    public synchronized int deferredCount() {
        int count = 0;
        for (SortedSet offsets : this.deferred.values()) {
            count += offsets.size();
        }
        return count;
    }

    public synchronized int getInPipeline() {
        return this.inPipeline;
    }

    public synchronized CommitArgs getAndClearOffsets() {
        Map offsetMap = new HashMap<>();
        if (this.outOfOrderCommits) {
            this.deferred.forEach((tp, offsets) -> {
                if (offsets.size() > 0) {
                    long lastThisPart = -1;
                    Iterator uncommittedIterator = this.uncommitted.get(tp).iterator();
                    Iterator deferredIterator = offsets.iterator();

                    while (deferredIterator.hasNext()) {
                        Long earliestDeferredOffset = deferredIterator.next();
                        Long earliestUncommittedOffset = uncommittedIterator.next();
                        if (!earliestDeferredOffset.equals(earliestUncommittedOffset)) {
                            break;
                        }

                        lastThisPart = earliestDeferredOffset;
                        uncommittedIterator.remove();
                        deferredIterator.remove();
                    }

                    if (lastThisPart >= 0) {
                        offsetMap.put(tp, new OffsetAndMetadata(lastThisPart + 1));
                    }
                }
            });
            batchSize = deferredCount();
        } else {
            latestOffsets.putAll(consumedOffsets);
            Iterator> iterator = consumedOffsets.entrySet().iterator();
            while (iterator.hasNext()) {
                Map.Entry entry = iterator.next();
                offsetMap.put(entry.getKey(), new OffsetAndMetadata(entry.getValue() + 1));
                iterator.remove();
            }
            batchSize = 0;
        }

        List> currentCallbackEmitters;
        if (!callbackEmitters.isEmpty()) {
            currentCallbackEmitters = callbackEmitters;
            callbackEmitters = new ArrayList<>();
        } else
            currentCallbackEmitters = null;

        return new CommitArgs(offsetMap, currentCallbackEmitters);
    }

    public synchronized void restoreOffsets(CommitArgs commitArgs, boolean restoreCallbackEmitters) {
        // Restore offsets that haven't been updated.
        if (outOfOrderCommits) {
            commitArgs.offsets.forEach((tp, offset) -> {
                this.deferred.get(tp).add(offset.offset() - 1);
                this.uncommitted.get(tp).add(offset.offset() - 1);
            });
        } else {
            for (Map.Entry entry : commitArgs.offsets.entrySet()) {
                TopicPartition topicPart = entry.getKey();
                long offset = entry.getValue().offset();
                Long latestOffset = latestOffsets.get(topicPart);
                if (latestOffset == null || latestOffset <= offset - 1)
                    consumedOffsets.putIfAbsent(topicPart, offset - 1);
            }
        }
        // If Mono is being failed after maxAttempts or due to fatal error, callback emitters
        // are not restored. Mono#retry will generate new callback emitters. If Mono status
        // is not being updated because commits are attempted again by KafkaReceiver, restore
        // the emitters for the next attempt.
        if (restoreCallbackEmitters && commitArgs.callbackEmitters != null)
            this.callbackEmitters = commitArgs.callbackEmitters;
    }

    @Override
    public synchronized String toString() {
        return String.valueOf(consumedOffsets);
    }

    public static class CommitArgs {
        private final Map offsets;
        private final List> callbackEmitters;
        CommitArgs(Map offsets, List> callbackEmitters) {
            this.offsets = offsets;
            this.callbackEmitters = callbackEmitters;
        }

        public Map offsets() {
            return offsets;
        }
        List> callbackEmitters() {
            return callbackEmitters;
        }
    }
}