[SPARK-27042][SS] Invalidate cached Kafka producer in case of task retry

Author: gaborgsomogyi

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/CachedKafkaProducer.scala

@@ -25,7 +25,7 @@ import scala.util.control.NonFatal
 
 import org.apache.kafka.clients.producer.{Callback, KafkaProducer, ProducerRecord}
 
-import org.apache.spark.SparkEnv
+import org.apache.spark.{SparkEnv, TaskContext}
 import org.apache.spark.internal.Logging
 import org.apache.spark.kafka010.{KafkaConfigUpdater, KafkaRedactionUtil}
 import org.apache.spark.sql.kafka010.InternalKafkaProducerPool._
@@ -93,6 +93,10 @@ private[kafka010] object CachedKafkaProducer extends Logging {
         .setAuthenticationConfigIfNeeded()
         .build()
     val key = toCacheKey(updatedKafkaParams)
+    if (TaskContext.get != null && TaskContext.get.attemptNumber >= 1) {
+      logDebug(s"Task re-attempt detected, invalidating producers with key $key")
+      producerPool.invalidateKey(key)
+    }
     producerPool.borrowObject(key, updatedKafkaParams)
   }
 

external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/KafkaDataConsumer.scala

@@ -628,6 +628,8 @@ private[kafka010] object KafkaDataConsumer extends Logging {
     if (TaskContext.get != null && TaskContext.get.attemptNumber >= 1) {
       val cacheKey = new CacheKey(topicPartition, kafkaParams)
 
+      logDebug(s"Task re-attempt detected, invalidating consumers with key $cacheKey")
+
       // If this is reattempt at running the task, then invalidate cached consumer if any.
       consumerPool.invalidateKey(cacheKey)
 

external/kafka-0-10-sql/src/test/scala/org/apache/spark/sql/kafka010/CachedKafkaProducerSuite.scala

@@ -97,6 +97,27 @@ class CachedKafkaProducerSuite extends SharedSparkSession with PrivateMethodTest
     assert(producerPool.size(toCacheKey(kafkaParams2)) === 1)
   }
 
+  test("acquire should return a new instance with Task retry") {
+    try {
+      val kafkaParams = getKafkaParams()
+
+      val context1 = new TaskContextImpl(0, 0, 0, 0, 0, null, null, null)
+      TaskContext.setTaskContext(context1)
+      val producer1 = CachedKafkaProducer.acquire(kafkaParams)
+      CachedKafkaProducer.release(producer1)
+
+      val context2 = new TaskContextImpl(0, 0, 0, 0, 1, null, null, null)
+      TaskContext.setTaskContext(context2)
+      val producer2 = CachedKafkaProducer.acquire(kafkaParams)
+      CachedKafkaProducer.release(producer2)
+
+      assert(producer1 !== producer2)
+      assert(producerPool.size(toCacheKey(kafkaParams)) === 1)
+    } finally {
+      TaskContext.unset()
+    }
+  }
+
   test("Concurrent use of CachedKafkaProducer") {
     val data = (1 to 1000).map(_.toString)
     testUtils.createTopic(topic, 1)