HeartSaVioR 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_r272460956
##########
File path:
external/kafka-0-10-sql/src/main/scala/org/apache/spark/sql/kafka010/CachedKafkaProducer.scala
##########
@@ -18,20 +18,70 @@
package org.apache.spark.sql.kafka010
import java.{util => ju}
-import java.util.concurrent.{ConcurrentMap, ExecutionException, TimeUnit}
+import java.util.concurrent.{ConcurrentLinkedQueue, ConcurrentMap,
ExecutionException, TimeUnit}
+import java.util.concurrent.atomic.AtomicInteger
+
+import scala.collection.JavaConverters._
+import scala.util.control.NonFatal
import com.google.common.cache._
import com.google.common.util.concurrent.{ExecutionError,
UncheckedExecutionException}
import org.apache.kafka.clients.producer.KafkaProducer
-import scala.collection.JavaConverters._
-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(
+ private val id: String = ju.UUID.randomUUID().toString,
+ private val inUseCount: AtomicInteger = new AtomicInteger(0),
+ private val kafkaParams: Seq[(String, Object)]) extends Logging {
+
+ private val configMap = kafkaParams.map(x => x._1 -> x._2).toMap.asJava
+
+ private def updatedAuthConfigIfNeeded(kafkaParamsMap: ju.Map[String,
Object]) =
+ KafkaConfigUpdater("executor", kafkaParamsMap.asScala.toMap)
+ .setAuthenticationConfigIfNeeded()
+ .build()
+
+ lazy val kafkaProducer: KafkaProducer[Array[Byte], Array[Byte]] = {
+ val producer = new KafkaProducer[Array[Byte],
Array[Byte]](updatedAuthConfigIfNeeded(configMap))
+ logDebug(s"Created a new instance of KafkaProducer for " +
+ s"$kafkaParams with Id: $id")
+ closed = false
+ producer
+ }
+ @volatile
+ private var isCached: Boolean = true
+ private var closed: Boolean = true
+ private def close(): Unit = {
+ try {
+ this.synchronized {
+ if (!closed) {
+ closed = true
+ kafkaProducer.close()
+ logInfo(s"Closed kafka producer: $this")
+ }
+ }
+ } catch {
+ case NonFatal(e) =>
+ logWarning(s"Error while closing kafka producer with params:
$kafkaParams", e)
+ }
+ }
+
+ private def inUse(): Boolean = inUseCount.get() > 0
+
+ private def unCache(): Unit = isCached = false
+
+ private[kafka010] def getInUseCount: Int = inUseCount.get()
- private type Producer = KafkaProducer[Array[Byte], Array[Byte]]
+ private[kafka010] def getKafkaParams: Seq[(String, Object)] = kafkaParams
+
+ private[kafka010] def flush(): Unit = kafkaProducer.flush()
+
+ private[kafka010] def isClosed: Boolean = closed
Review comment:
"in non test code" sounds like you tend to agree it's not properly guarded
if we include test code, and it's being used as an assertion in multiple places
in test code - which might lead intermittent test failure.
----------------------------------------------------------------
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:
[email protected]
With regards,
Apache Git Services
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
