[GitHub] [spark] ScrapCodes commented on a change in pull request #19096: [SPARK-21869][SS] A cached Kafka producer should not be closed if any task is using it - adds inuse tracking.

[GitBox <gi...@apache.org>]

ScrapCodes commented on a change in pull request #19096: [SPARK-21869][SS] A cached Kafka producer should not be closed if any task is using it - adds inuse tracking.
URL: https://github.com/apache/spark/pull/19096#discussion_r263753699
 
 

 ##########
 File path: external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/CachedKafkaProducer.scala
 ##########
 @@ -29,51 +30,76 @@ import scala.util.control.NonFatal
 import org.apache.spark.SparkEnv
 import org.apache.spark.internal.Logging
 
-private[kafka010] object CachedKafkaProducer extends Logging {
+private[kafka010] case class CachedKafkaProducer(id: String, inUseCount: AtomicInteger,
+    kafkaProducer: KafkaProducer[Array[Byte], Array[Byte]]) extends Logging {
+  private var closed: Boolean = false
+  private def close(): Unit = this.synchronized {
+    if (!closed) {
+      closed = true
+      kafkaProducer.close()
+      logInfo(s"Closed kafka producer: $kafkaProducer")
+    }
+  }
+
+  def flush(): Unit = {
+    kafkaProducer.flush()
+  }
+
+  private[kafka010] def isClosed: Boolean = closed
+}
 
-  private type Producer = KafkaProducer[Array[Byte], Array[Byte]]
+private[kafka010] object CachedKafkaProducer extends Logging {
 
   private lazy val cacheExpireTimeout: Long =
     SparkEnv.get.conf.getTimeAsMs("spark.kafka.producer.cache.timeout", "10m")
 
-  private val cacheLoader = new CacheLoader[Seq[(String, Object)], Producer] {
-    override def load(config: Seq[(String, Object)]): Producer = {
+  private val cacheLoader = new CacheLoader[Seq[(String, Object)], CachedKafkaProducer] {
+    override def load(config: Seq[(String, Object)]): CachedKafkaProducer = {
       val configMap = config.map(x => x._1 -> x._2).toMap.asJava
       createKafkaProducer(configMap)
     }
   }
 
-  private val removalListener = new RemovalListener[Seq[(String, Object)], Producer]() {
+  private val closeQueue = new ConcurrentLinkedQueue[CachedKafkaProducer]()
+
+  private val removalListener = new RemovalListener[Seq[(String, Object)], CachedKafkaProducer]() {
     override def onRemoval(
-        notification: RemovalNotification[Seq[(String, Object)], Producer]): Unit = {
-      val paramsSeq: Seq[(String, Object)] = notification.getKey
-      val producer: Producer = notification.getValue
-      logDebug(
-        s"Evicting kafka producer $producer params: $paramsSeq, due to ${notification.getCause}")
-      close(paramsSeq, producer)
+        notification: RemovalNotification[Seq[(String, Object)], CachedKafkaProducer]): Unit = {
+      val producer: CachedKafkaProducer = notification.getValue
+      logDebug(s"Evicting kafka producer $producer, due to ${notification.getCause}")
+      if (producer.inUseCount.intValue() > 0) {
+        // When a inuse producer is evicted we wait for it to be released before finally closing it.
+        closeQueue.add(producer)
+      } else {
+        close(producer)
 
 Review comment:
   UUID is printed on close.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
users@infra.apache.org


With regards,
Apache Git Services

---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscribe@spark.apache.org
For additional commands, e-mail: reviews-help@spark.apache.org