Kafka façade for use in a Spark job (adapted liberally from elsewhere)

KafkaSink.scala

import java.util.concurrent.Future
import java.util.concurrent.atomic.AtomicReference

import org.apache.kafka.clients.producer.{Callback, KafkaProducer, ProducerRecord, RecordMetadata}

import scala.collection.JavaConversions._

/** This construct allows for a KafkaPublisher to be reused in an executor.
  *
  * Establishing communication to the Kafka cluster is expensive, so this should
  * improve performance and reduce the risk of connection errors.
  */
object KafkaSink {

  def apply(config: Map[String, Object], topic: String): KafkaSink = {
    val create = () => {
      val producer = new KafkaProducer[String, Array[Byte]](config)

      sys.addShutdownHook {
        producer.close()
      }

      producer
    }
    new KafkaSink(create, topic)
  }
}

class KafkaSink(createProducer: () => KafkaProducer[String, Array[Byte]], topic: String) extends Serializable {

  // Lazy producer allows it to be created by executor and correctly serialize.
  lazy val producer = createProducer()

  def send(key: String, bytes: Array[Byte], callback: Callback): Future[RecordMetadata] =
    producer.send(new ProducerRecord(topic, key, bytes), callback)
}

class KafkaSinkExceptionHandler extends Callback {

  private val lastException = new AtomicReference[Option[Exception]](None)

  override def onCompletion(metadata: RecordMetadata, exception: Exception): Unit =
    lastException.set(Option(exception))

  def throwExceptionIfAny(): Unit =
    lastException.getAndSet(None).foreach(ex => throw ex)
}

And an example usage:

    records.foreachRDD { rdd =>
      rdd.foreachPartition { records =>
        val logger = Logger.getLogger(getClass.getName)
        val callback = new KafkaSinkExceptionHandler

        val metadata = records.foreach {rec => {
          callback.throwExceptionIfAny()
          kafkaSink.value.send(filter.filterId.toString, rec, callback)
        })

        // Wait on ack from Kafka publisher before continuing, or records may be lost.
        if (metadata.nonEmpty) {
          logger.info("Waiting on metadata to ack publish to Kafka.")
          metadata.foreach { metadata => metadata.get() }
          callback.throwExceptionIfAny()
        }
      }
    }