feat(clients): add new interface ConsumedErrorHandler (#1)

This commit adds the new interface ConsumedErrorHandler
to handle error thrown while processing records.

In addition, it adds two built-in implementations:
 - CloseTaskOnConsumedError
 - LogAndCommitOnConsumedError

Additional changes:
 - ConsumerTask interface now exposes the two methods: commitAsync, commitSync
 - ConsumerWorker.Builder has a new onConsumedError method
 - ConsumerBatchRecordsListener is only invoked on non-empty list of records
 - add some unit-tests for KafkaConsumerTask

Resolves: #1

Author: fhussonnois

clients/src/main/kotlin/io/streamthoughts/kafka/clients/consumer/ConsumerTask.kt

@@ -19,6 +19,9 @@
 package io.streamthoughts.kafka.clients.consumer
 
 import org.apache.kafka.clients.consumer.Consumer
+import org.apache.kafka.clients.consumer.OffsetAndMetadata
+import org.apache.kafka.common.TopicPartition
+import java.time.Duration
 
 interface ConsumerTask{
 
@@ -72,11 +75,17 @@ interface ConsumerTask{
     fun resume()
 
     /**
-     * Shutdowns the [ConsumerTask]
+     * Shutdowns the [ConsumerTask] and wait for completion.
      * @see org.apache.kafka.clients.consumer.Consumer.close
      */
     fun shutdown()
 
+    /**
+     * Shutdowns the [ConsumerTask] and wait for completion until the given [timeout].
+     * @see org.apache.kafka.clients.consumer.Consumer.close
+     */
+    fun shutdown(timeout: Duration)
+
     /**
      * @return the [State] of this [ConsumerTask].
      */
@@ -86,4 +95,20 @@ interface ConsumerTask{
      * Executes the given [action] with the underlying [Consumer].
      */
     fun <T> execute(action: (consumer: Consumer<ByteArray, ByteArray>) -> T): T
+
+    /**
+     * Commits asynchronously the positions of the internal [Consumer] for the given [offsets].
+     * If passed [offsets] is {@code null} then commit the [Consumer] positions for its current partition assignments.
+     *
+     * @see [Consumer.commitAsync]
+     */
+    fun commitAsync(offsets: Map<TopicPartition, OffsetAndMetadata>? = null)
+
+    /**
+     * Commits synchronously the positions of the internal [Consumer] for the given offsets.
+     * If passed [offsets] is {@code null} then commit the [Consumer] positions for its current partition assignments.
+     *
+     * @see [Consumer.commitAsync]
+     */
+    fun commitSync(offsets: Map<TopicPartition, OffsetAndMetadata>? = null)
 }

clients/src/main/kotlin/io/streamthoughts/kafka/clients/consumer/ConsumerWorker.kt

@@ -18,6 +18,7 @@
  */
 package io.streamthoughts.kafka.clients.consumer
 
+import io.streamthoughts.kafka.clients.consumer.error.ConsumedErrorHandler
 import io.streamthoughts.kafka.clients.consumer.error.serialization.DeserializationErrorHandler
 import io.streamthoughts.kafka.clients.consumer.listener.ConsumerBatchRecordsListener
 import java.util.regex.Pattern
@@ -41,30 +42,39 @@ interface ConsumerWorker<K, V> {
         fun factory(consumerFactory: ConsumerFactory): Builder<K, V>
 
         /**
-         * Sets the [RebalanceListener] to invoke when a rebalance is in progress and partitions are assigned.
+         * Sets the [listener] to invoke when a rebalance is in progress and partitions are assigned.
          */
         fun onPartitionsAssigned(listener: RebalanceListener): Builder<K, V>
 
         /**
-         * Sets the [RebalanceListener] to invoke when a rebalance is in progress and partitions are revoked.
+         * Sets the [listener] to invoke when a rebalance is in progress and partitions are revoked.
          */
         fun onPartitionsRevokedBeforeCommit(listener: RebalanceListener): Builder<K, V>
 
         /**
-         * Sets the [RebalanceListener] to invoke when a rebalance is in progress and partitions are revoked.
+         * Sets the [listener] to invoke when a rebalance is in progress and partitions are revoked.
          */
         fun onPartitionsRevokedAfterCommit(listener: RebalanceListener): Builder<K, V>
 
         /**
-         * Sets the [RebalanceListener] to invoke when a rebalance is in progress and partitions are lost.
+         * Sets the [listener] to invoke when a rebalance is in progress and partitions are lost.
          */
         fun onPartitionsLost(listener: RebalanceListener): Builder<K, V>
 
         /**
-         * Sets the [DeserializationErrorHandler] to invoke when a exception happen while de-serializing a record.
+         * Sets the [handler] to invoke when a exception happen while de-serializing a record.
          */
         fun onDeserializationError(handler: DeserializationErrorHandler<K, V>): Builder<K, V>
 
+        /**
+         * Sets the [handler] to invoked when a error is thrown while processing last records returned from the
+         * the [org.apache.kafka.clients.consumer.Consumer.poll] method, i.e. an exception thrown by the provided
+         * [ConsumerBatchRecordsListener].
+         *
+         * @see [onConsumed]
+         */
+        fun onConsumedError(handler: ConsumedErrorHandler): Builder<K, V>
+
         /**
          * Sets the [ConsumerBatchRecordsListener] to invoke when a non-empty batch of records is returned from
          * the [org.apache.kafka.clients.consumer.Consumer.poll] method.

clients/src/main/kotlin/io/streamthoughts/kafka/clients/consumer/KafkaConsumerTask.kt

@@ -20,6 +20,7 @@ package io.streamthoughts.kafka.clients.consumer
 
 import ch.qos.logback.classic.Level
 import io.streamthoughts.kafka.clients.consumer.ConsumerTask.State
+import io.streamthoughts.kafka.clients.consumer.error.ConsumedErrorHandler
 import io.streamthoughts.kafka.clients.consumer.error.serialization.DeserializationErrorHandler
 import io.streamthoughts.kafka.clients.consumer.listener.ConsumerBatchRecordsListener
 import io.streamthoughts.kafka.clients.loggerFor
@@ -38,8 +39,9 @@ import org.apache.kafka.common.errors.WakeupException
 import org.apache.kafka.common.serialization.ByteArrayDeserializer
 import org.apache.kafka.common.serialization.Deserializer
 import java.time.Duration
-import java.util.*
+import java.util.LinkedList
 import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicBoolean
 import kotlin.collections.HashMap
 import kotlin.math.max
@@ -53,11 +55,16 @@ class KafkaConsumerTask<K, V>(
     private val listener: ConsumerBatchRecordsListener<K, V>,
     private var clientId: String = "",
     private val deserializationErrorHandler: DeserializationErrorHandler<K, V>,
+    private val consumedErrorHandler: ConsumedErrorHandler? = null,
     private val consumerAwareRebalanceListener : ConsumerAwareRebalanceListener? = null
 ) : ConsumerTask {
 
     companion object {
         private val Log = loggerFor(KafkaConsumerTask::class.java)
+
+        private fun <K, V> flatten(records: Map<TopicPartition, List<ConsumerRecord<K?, V?>>>): List<ConsumerRecord<K?, V?>> {
+            return records.flatMap { (_, v) -> v }.toList()
+        }
     }
 
     @Volatile
@@ -134,10 +141,10 @@ class KafkaConsumerTask<K, V>(
         } catch (e: WakeupException) {
             if (!isShutdown.get()) throw e
             else {
-                logWithConsumerInfo(Level.INFO, "Stop polling due to the io.streamthoughts.kafka.clients.consumer-task is being closed")
+                logWithConsumerInfo(Level.INFO, "Stop polling due to the consumer-task is being closed")
             }
         } catch (e: CancellationException) {
-            logWithConsumerInfo(Level.INFO, "Stop polling due to the io.streamthoughts.kafka.clients.consumer-task has been canceled")
+            logWithConsumerInfo(Level.INFO, "Stop polling due to the consumer-task has been canceled")
             throw e
         } finally {
             state = State.PENDING_SHUTDOWN
@@ -160,21 +167,36 @@ class KafkaConsumerTask<K, V>(
     }
 
     private fun pollOnce() {
-        val records: ConsumerRecords<ByteArray, ByteArray> = consumer.poll(pollTime)
+        val rawRecords: ConsumerRecords<ByteArray, ByteArray> = consumer.poll(pollTime)
 
         if (state == State.PARTITIONS_ASSIGNED) {
             state = State.RUNNING
         }
 
-        // deserialize all records using user-provided Deserializer
-        val deserialized : Map<TopicPartition, List<ConsumerRecord<K?, V?>>> =
-            records.partitions()
-                .map { Pair(it, deserialize(records.records(it))) }
-                .toMap()
+        if (!rawRecords.isEmpty) {
+            // deserialize all records using user-provided Deserializer
+            val recordsPerPartitions: Map<TopicPartition, List<ConsumerRecord<K?, V?>>> =
+                rawRecords.partitions()
+                    .map { Pair(it, deserialize(rawRecords.records(it))) }
+                    .toMap()
+            try {
+                processBatchRecords(ConsumerRecords(recordsPerPartitions))
+                updateConsumedOffsets(rawRecords) // only update once all records from batch have been processed.
+                mayCommitAfterBatch()
+            } catch (e: Exception) {
+                mayHandleConsumedError(recordsPerPartitions, e)
+            }
+        }
+    }
 
-        processBatchRecords(ConsumerRecords(deserialized))
-        updateConsumedOffsets(records) // only update once all records from batch have been processed.
-        mayCommitAfterBatch()
+    private fun mayHandleConsumedError(recordsPerPartitions: Map<TopicPartition, List<ConsumerRecord<K?, V?>>>,
+                                       thrownException: Exception
+    ) {
+        consumedErrorHandler?.handle(
+            this,
+            flatten(recordsPerPartitions),
+            thrownException
+        )
     }
 
     private fun processBatchRecords(records: ConsumerRecords<K?, V?>) {
@@ -200,6 +222,19 @@ class KafkaConsumerTask<K, V>(
         shutdownLatch.await()
     }
 
+    override fun shutdown(timeout: Duration) {
+        logWithConsumerInfo(Level.INFO, "Closing")
+        isShutdown.set(true)
+        consumer.wakeup()
+        if (timeout != Duration.ZERO) {
+            try {
+                shutdownLatch.await(timeout.toMillis(), TimeUnit.MILLISECONDS)
+            } catch (e: InterruptedException) {
+                logWithConsumerInfo(Level.WARN, "Failed to close consumer before timeout")
+            }
+        }
+    }
+
     private fun deserialize(records: List<ConsumerRecord<ByteArray, ByteArray>>): List<ConsumerRecord<K?, V?>> {
         val deserialized = LinkedList<ConsumerRecord<K?, V?>>()
         for (record : ConsumerRecord<ByteArray, ByteArray> in records) {
@@ -253,7 +288,7 @@ class KafkaConsumerTask<K, V>(
                 state = State.PARTITIONS_REVOKED
                 consumerAwareRebalanceListener?.onPartitionsRevokedBeforeCommit(consumer, partitions)
 
-                doCommitSync(offsetAndMetadataToCommit())
+                commitSync(offsetAndMetadataToCommit())
 
                 consumerAwareRebalanceListener?.onPartitionsRevokedAfterCommit(consumer, partitions)
                 assignedPartitions.clear()
@@ -267,11 +302,11 @@ class KafkaConsumerTask<K, V>(
         }
     }
 
-    private fun offsetAndMetadataToCommit() = consumedOffsets.map { Pair(it.key, OffsetAndMetadata(it.value)) }.toMap()
+    private fun offsetAndMetadataToCommit() = consumedOffsets.map { Pair(it.key, OffsetAndMetadata(it.value + 1)) }.toMap()
 
     private fun mayCommitAfterBatch() {
         if (!isAutoCommitEnabled && consumedOffsets.isNotEmpty()) {
-            doCommitAsync(offsetAndMetadataToCommit())
+            commitAsync(offsetAndMetadataToCommit())
             consumedOffsets.clear()
         }
     }
@@ -282,11 +317,11 @@ class KafkaConsumerTask<K, V>(
                 val offset = consumer.position(topicPartition)
                 Pair(topicPartition, OffsetAndMetadata(offset))
             }.toMap()
-            doCommitSync(positionsToCommit)
+            commitSync(positionsToCommit)
         }
     }
 
-    private fun doCommitAsync(offsets: Map<TopicPartition, OffsetAndMetadata>? = null) {
+    override fun commitAsync(offsets: Map<TopicPartition, OffsetAndMetadata>?) {
         logWithConsumerInfo(Level.INFO, "Committing offsets async-synchronously for positions: $offsets")
         consumer.commitAsync(offsets) {
             _, exception -> if (exception != null)    {
@@ -295,8 +330,8 @@ class KafkaConsumerTask<K, V>(
         }
     }
 
-    private fun doCommitSync(offsets: Map<TopicPartition, OffsetAndMetadata>? = null) {
-        if (consumer.assignment().isEmpty()) return // no need to commit if no partition is assign to this io.streamthoughts.kafka.clients.consumer
+    override fun commitSync(offsets: Map<TopicPartition, OffsetAndMetadata>?) {
+        if (consumer.assignment().isEmpty()) return // no need to commit if no partition is assign to this consumer
         try {
             if (offsets == null) {
                 logWithConsumerInfo(Level.WARN, "Committing offsets synchronously for consumed records")
@@ -307,7 +342,7 @@ class KafkaConsumerTask<K, V>(
             }
             logWithConsumerInfo(Level.WARN, "Offsets committed for partitions: $assignedPartitions")
         } catch (e: RetriableCommitFailedException) {
-            doCommitSync(offsets)
+            commitSync(offsets)
         } catch (e : RebalanceInProgressException) {
             logWithConsumerInfo(Level.WARN, "Error while committing offsets due to a rebalance in progress. Ignored")
         }

clients/src/main/kotlin/io/streamthoughts/kafka/clients/consumer/KafkaConsumerWorker.kt

@@ -19,6 +19,8 @@
 package io.streamthoughts.kafka.clients.consumer
 
 import io.streamthoughts.kafka.clients.consumer.KafkaConsumerWorker.KafkaConsumerWorker
+import io.streamthoughts.kafka.clients.consumer.error.ConsumedErrorHandler
+import io.streamthoughts.kafka.clients.consumer.error.ConsumedErrorHandlers.closeTaskOnConsumedError
 import io.streamthoughts.kafka.clients.consumer.error.serialization.DeserializationErrorHandler
 import io.streamthoughts.kafka.clients.consumer.error.serialization.DeserializationErrorHandlers
 import io.streamthoughts.kafka.clients.consumer.listener.ConsumerBatchRecordsListener
@@ -32,7 +34,6 @@ import kotlinx.coroutines.joinAll
 import kotlinx.coroutines.launch
 import org.apache.kafka.clients.consumer.Consumer
 import org.apache.kafka.clients.consumer.ConsumerConfig
-import org.apache.kafka.clients.consumer.ConsumerRecords
 import org.apache.kafka.common.TopicPartition
 import org.apache.kafka.common.serialization.Deserializer
 import java.util.concurrent.ExecutorService
@@ -50,6 +51,7 @@ class KafkaConsumerWorker<K, V> (
     private val valueDeserializer: Deserializer<V>,
     private val consumerRebalanceListener: ConsumerAwareRebalanceListener,
     private val batchRecordListener: ConsumerBatchRecordsListener<K, V>,
+    private val onConsumedError: ConsumedErrorHandler,
     private val onDeserializationError: DeserializationErrorHandler<K, V>,
     private val consumerFactory: ConsumerFactory = ConsumerFactory.DefaultConsumerFactory
     ): ConsumerWorker<K, V> {
@@ -104,7 +106,8 @@ class KafkaConsumerWorker<K, V> (
                 batchRecordListener,
                 clientId = computeClientId(taskId),
                 consumerAwareRebalanceListener = consumerRebalanceListener,
-                deserializationErrorHandler = onDeserializationError
+                deserializationErrorHandler = onDeserializationError,
+                consumedErrorHandler = onConsumedError
             )
         }
         doStart()
@@ -165,7 +168,8 @@ class KafkaConsumerWorker<K, V> (
         var onPartitionsLost: RebalanceListener? = null,
         var batchRecordListener: ConsumerBatchRecordsListener<K, V>? = null,
         var onDeserializationError: DeserializationErrorHandler<K, V>? = null,
-        var consumerFactory: ConsumerFactory? = null
+        var consumerFactory: ConsumerFactory? = null,
+        var onConsumedError: ConsumedErrorHandler? = null
     ) : ConsumerWorker.Builder<K, V> {
 
         override fun configure(init: KafkaConsumerConfigs.() -> Unit) {
@@ -190,6 +194,9 @@ class KafkaConsumerWorker<K, V> (
         override fun onDeserializationError(handler : DeserializationErrorHandler<K, V>)  =
             apply { onDeserializationError = handler }
 
+        override fun onConsumedError(handler : ConsumedErrorHandler)  =
+            apply { onConsumedError = handler }
+
         override fun onConsumed(listener: ConsumerBatchRecordsListener<K, V>) =
             apply { this.batchRecordListener = listener }
 
@@ -200,6 +207,7 @@ class KafkaConsumerWorker<K, V> (
                 valueDeserializer,
                 SimpleConsumerAwareRebalanceListener(),
                 batchRecordListener ?: noop(),
+                onConsumedError ?: closeTaskOnConsumedError(),
                 onDeserializationError ?: DeserializationErrorHandlers.logAndFail(),
                 consumerFactory ?: ConsumerFactory.DefaultConsumerFactory
             )

clients/src/main/kotlin/io/streamthoughts/kafka/clients/consumer/error/ConsumedErrorHandler.kt

@@ -0,0 +1,39 @@
+/*
+ * Copyright 2020 StreamThoughts.
+ *
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 io.streamthoughts.kafka.clients.consumer.error
+
+import io.streamthoughts.kafka.clients.consumer.ConsumerTask
+import org.apache.kafka.clients.consumer.ConsumerRecord
+
+/**
+ * Handles errors thrown during the processing of a non-empty batch of [ConsumerRecord]
+ * using a given [io.streamthoughts.kafka.clients.consumer.listener.ConsumerBatchRecordsListener]
+ */
+interface ConsumedErrorHandler {
+
+    /**
+     * This method is invoked when an [thrownException] is thrown while a [consumerTask] is processing
+     * a non-empty batch of [records].
+     *
+     * @param consumerTask    the [ConsumerTask] polling records.
+     * @param records         the remaining [records] to be processed (including the one that failed).
+     * @param thrownException the [Exception] that was thrown while processing [records].
+     */
+    fun handle(consumerTask: ConsumerTask, records: List<ConsumerRecord<*, *>>, thrownException: Exception)
+}