Spring Kafka async-acks improvement : opportunistic contiguous commits

[djouvin]

Expected Behavior

The idea is to follow an opportunistic approch to manual commits, independant of the Kafka batch size, with a specific parameter governing the maximum number of messages in transit (as does the Reactor Kafka library).

Current Behavior

Currently the asyncAcks mode follows a "commit per batch" strategy, and waits for a batch messages to be fully acknowledged before commiting them and going on with the next batch.

Context

The problem with the current asyncAcks approach is that it couples too much the batching strategy (batch size, etc.) with the commit strategy, and also with the maximum number of messages you want to accept "in transit" (uncommited), which also determines the maximum number of duplicates you may have when a crash occurs.

The proposed approach is closer to what is doing Reactor Kafka (which is now deprecated unfortunately), and is better because you can control your commit sliding window indepently of the Kafka batching strategy. And you never wait between Kafka batches if you remain inside your max sliding window of "in transit" messages, so performances are much better.

This also allows to better control the Kafka polling independently of the duration of the message processing (by pausing the Kafka subscription when necessary and doing a poll to keep the consumer alive), and to manage a hysteresis parameter to avoid "flickering state changes".

This is not that difficult to implement, using for exemple a LinkedHashMap index for each topic partition (to keep message aknowledgement state and maintain ordering, and index messages by Kafka offset), and a background thread that will cleanup the "in transit" messages sequentially and maintain an accumulator with the "most recent commitable offset", and then commiting this offset as soon as possible, possibly with other commit conditions (to allow to plug a custom commit strategy to control commit frequency).

[djouvin]

To illustrate, here is a prototype Kafka async consumer implementation, in Groovy.

KafkaAsyncConsumer.groovy.txt

[artembilan]

Can you, please, share the code as a snippet here according to GH Markdown style: https://docs.github.com/en/get-started/writing-on-github/working-with-advanced-formatting/creating-and-highlighting-code-blocks ?
Or point us to working project somewhere on GH, so we can play with it on our side and be sure we are on the same page.

[artembilan]

Sounds like a good idea.
Feel free to contribute such a feature.
I also wonder how do you see it might effect existing behavior.
Do we want to keep it yet, and have your solution as an opt-in feature?
Or do we want to fully replace existing conservative approach?

Thanks

[djouvin]

Here is the quoted code in Groovy.

import org.apache.kafka.clients.consumer.*
import org.apache.kafka.common.TopicPartition
import java.time.Duration
import java.util.concurrent.*
import java.util.concurrent.atomic.AtomicInteger
import java.util.concurrent.ConcurrentHashMap

import groovy.transform.Log4j2
import groovy.transform.AutoFinal

/**
 * High performance non-blocking Kafka consumer allowing asynchronous message processing,
 * with an optimized opportunistic contiguous commit strategy (similar to Reactor Kafka)
 */
@Log4j2
@AutoFinal
class KafkaAsyncConsumer implements AutoCloseable {
    private final KafkaConsumer<String, String> kafkaConsumer
    private final ScheduledExecutorService commitScheduler = Executors.newSingleThreadScheduledExecutor()
    private final AtomicInteger inFlightMessages = new AtomicInteger(0)
    private final Map<TopicPartition, PartitionState> partitionStates = new ConcurrentHashMap<>()

    private final int highWatermark
    private final int lowWatermark
    private final long commitIntervalMs
    private final long pollTimeoutMs

    private volatile boolean running = true
    private boolean isPaused = false
    private Closure<Mono> processing

    static class PartitionState {
        final Map<Long, MessageContext> pendingMap = [:] // LinkedHashMap
        long latestConsumableOffset = -1
        long lastCommittedOffset = -1
        int consecutiveFailures = 0
        long lastFailureTime = 0
    }

    static class MessageContext {
        final long offset
        boolean isDone = false
        MessageContext(long offset) { this.offset = offset }
    }

    KafkaAsyncConsumer(Properties kafkaProps, Map<String, Object> options = [:]) {
        this.kafkaConsumer = new KafkaConsumer<>(kafkaProps)
        this.highWatermark = (options.highWatermark ?: 1000) as int
        this.lowWatermark = (options.lowWatermark ?: 500) as int
        this.commitIntervalMs = (options.commitIntervalMs ?: 500L) as long
        this.pollTimeoutMs = (options.pollTimeoutMs ?: 100L) as long

        commitScheduler.scheduleAtFixedRate(this.&runCommitJob,
                this.commitIntervalMs, this.commitIntervalMs, TimeUnit.MILLISECONDS)
    }

    void start(List<String> topics, Closure<Mono> processing) {
        this.processing = processing
        try {
            final listener = new ConsumerRebalanceListener() {
                void onPartitionsRevoked(Collection<TopicPartition> partitions) {
                    log.info "Revoked partitions: ${partitions}"
                    partitions.each { partitionStates.remove(it) } // memory cleanup
                }
                void onPartitionsAssigned(Collection<TopicPartition> partitions) {
                    log.info "Assigned partitions: ${partitions}"
                }
            }
            kafkaConsumer.subscribe(topics, listener)
            log.info "Main loop consuming $topics Kafka topics"
            while(running) {
                manageBackPressure()
                final records = kafkaConsumer.poll(Duration.ofMillis(this.pollTimeoutMs))
                records.each { record -> processAsync(record) }
            }
        }
        finally { close() }
    }

    private void manageBackPressure() {
        final load = inFlightMessages.get()
        if(!isPaused && load > this.highWatermark) {
            log.warn "Backpressure ON: load $load > $highWatermark -> pausing subscription"
            kafkaConsumer.pause(kafkaConsumer.assignment())
            isPaused = true
        }
        else if(isPaused && load < this.lowWatermark) {
            log.info "Backpressure OFF: load $load < $lowWatermark -> resuming subscription"
            kafkaConsumer.resume(kafkaConsumer.assignment())
            isPaused = false
        }
    }

    private void processAsync(ConsumerRecord<String, String> record) {
        final tp = new TopicPartition(record.topic(), record.partition())
        final state = partitionStates.computeIfAbsent(tp) { new PartitionState() }
        synchronized(state) {
            if(record.offset() <= state.latestConsumableOffset) return // Ignore duplicates
            state.pendingMap[record.offset()] = new MessageContext(record.offset())
        }
        inFlightMessages.incrementAndGet()
        
        // Closure returns a Mono (Reactor), can easily be adapted to other kind of async processing
        processing(record)
            .doFinally { signal ->
                markAsDone(tp, record.offset())
                inFlightMessages.decrementAndGet()
            }
            .subscribe()
    }

    private void markAsDone(TopicPartition tp, long offset) {
        final state = partitionStates[tp]
        if(!state) return
        synchronized(state) { state.pendingMap[offset]?.isDone = true }
    }

    /** Optimized contiguous commit scheduling */
    private void runCommitJob() {
        if(partitionStates.isEmpty()) return
        final Map<TopicPartition, OffsetAndMetadata> batchToCommit = [:]
        partitionStates.each { tp, state ->
            synchronized(state) {
                // drain : cleanup of the LinkedHashMap
                final iterator = state.pendingMap.iterator()
                while(iterator.hasNext()) {
                    final entry = iterator.next()
                    if(entry.value.isDone) {
                        state.latestConsumableOffset = entry.key
                        iterator.remove() // O(1)
                    }
                    else break // hole detected -> stop
                }
                // commit preparation
                if(state.latestConsumableOffset > state.lastCommittedOffset) {
                    batchToCommit[tp] = new OffsetAndMetadata(state.latestConsumableOffset + 1)
                }
            }
        }

        if(batchToCommit) {
            kafkaConsumer.commitAsync(batchToCommit) { offsets, exception ->
                if(!exception) handleCommitSuccess(offsets)
                else handleCommitFailure(offsets, exception)
            }
        }
    }

    private void handleCommitSuccess(Map<TopicPartition, OffsetAndMetadata> offsets) {
        offsets.each { tp, meta ->
            final state = partitionStates[tp]
            if(state) {
                synchronized(state) {
                    state.lastCommittedOffset = meta.offset() - 1
                    state.consecutiveFailures = 0
                }
            }
        }
    }

    private void handleCommitFailure(Map<TopicPartition, OffsetAndMetadata> offsets, Exception exception) {
        log.warn "Async commit failure (implicit retry at next cycle): ${exception.message}"
        offsets.keySet().each { tp ->
            def state = partitionStates[tp]
            if(state) {
                synchronized(state) {
                    state.consecutiveFailures ++
                    state.lastFailureTime = System.currentTimeMillis()
                    if (state.consecutiveFailures > 10) {
                        log.error "Partition $tp blocked at offset ${state.latestConsumableOffset
                                } since ${state.consecutiveFailures} cycles"
                    }
                }
            }
        }
    }

    @Override
    void close() {
        running = false
        commitScheduler.shutdown()
        kafkaConsumer.close()
    }
}

I will try to propose a correct Spring Kafka integration in Java as a MR next week.

I think the best would be a replacement of the current asyncAcks mode by this new solution : in fact, the current solution has the drawback of being closer to the auto-commit mode, in the sense that commits are done at each polling / batch transition. Howvere asyncAcks is supposed to be an option of manual commits, where commits are fully decorrelated/decoupled from the Kafka polling and batching, to allow a custom commit policy/strategy.