Spark Streaming: Custom Receiver kryo registration with Google Pubsub

Question

I'm using Spark 2.0.2 with Kryo serialization.

I'm attempting to implement a custom receiver for ingesting messages from Google PubSub into Spark Streaming:

class PubSubReceiver(project: String, topic: String, subscription: String)
  extends Receiver[Array[Byte]](StorageLevel.MEMORY_AND_DISK_2) with Logging {
  val projectFullName = ProjectName.create(project)
  val topicName = TopicName.create(project, topic)
  val subscriptionName = SubscriptionName.create(project, subscription)
  val subscriber = Subscriber.defaultBuilder(subscriptionName, new receiver).build

  def onStart() {
    new Thread() {
      override def run() {
        subscriber.startAsync()
        //ensure subscriber is running as well as spark receiver
        while (subscriber.isRunning && !isStopped()) {
          logger.info(s"${subscriber.getSubscriptionName} receiver running")
          //sleep 10s
          Thread.sleep(10000)
        }
        logger.info(s"${subscriber.getSubscriptionName} receiver stopping")
      }
    }.start()
  }

  def onStop(): Unit = {
    // There is nothing much to do as the thread calling receive()
    // is designed to stop by itself if isStopped() returns false
  }

  private class receiver extends MessageReceiver {
    override def receiveMessage(message: PubsubMessage, consumer: AckReplyConsumer): Unit = {
      store(ArrayBuffer(message.getData.toByteArray), message.getAttributesMap)
    }
  }

}

However when running a Spark job that utilizes this receiver, it seems that I have to serialized the job itself, which doesnt seem correct (the spark context would then be serialized).

object PubSubStreamingIngestionJob extends App {
  //... setup

  lazy val ssc = new StreamingContext(spark.sparkContext, batchInterval)

  lazy val pubsubUnionStream =the stream
    ssc.receiverStream(new PubSubReceiver(projectName, topicName, subscriptionName))

  pubsubUnionStream.map( messageBytes => ...business logic... )

  ssc.start()
  ssc.awaitTermination()

}

The following error is thrown:

java.io.IOException: com.esotericsoftware.kryo.KryoException: java.lang.IllegalArgumentException: Class is not registered: com.c2fo.atlas.jobs.streaming.gcp.PubSubStreamingIngestionJob
Note: To register this class use: kryo.register(com.mycompany.package.PubSubStreamingIngestionJob.class);
Serialization trace:
classes (sun.misc.Launcher$AppClassLoader)
contextClassLoader (java.lang.Thread)
threads (java.lang.ThreadGroup)
parent (java.lang.ThreadGroup)
group (java.util.concurrent.Executors$DefaultThreadFactory)
val$backingThreadFactory (com.google.common.util.concurrent.ThreadFactoryBuilder$1)
threadFactory (java.util.concurrent.ScheduledThreadPoolExecutor)
e (java.util.concurrent.Executors$DelegatedScheduledExecutorService)
executor (com.google.cloud.pubsub.spi.v1.Subscriber)
subscriber (com.mycompany.package.PubSubReceiver)
array (scala.collection.mutable.WrappedArray$ofRef)

Is there a better way of implementing this?

Answer 1

The issue was the Subscriber instance needed to be thread local to prevent the entire closure from being serialized.

package org.apache.spark.streaming.gcp

import com.c2fo.atlas.util.LazyLogging
import com.google.cloud.pubsub.spi.v1._
import com.google.iam.v1.ProjectName
import com.google.pubsub.v1._
import org.apache.spark.storage.StorageLevel
import org.apache.spark.streaming.receiver.Receiver

import scala.collection.mutable.ArrayBuffer

class PubSubReceiver(project: String, topic: String, subscription: String)
  extends Receiver[PubsubMessage](StorageLevel.MEMORY_AND_DISK_2) with LazyLogging{
  val projectFullName = ProjectName.create(project)
  val topicName = TopicName.create(project, topic)
  val subscriptionName = SubscriptionName.create(project, subscription)
  def onStart() {
    new Thread() {
      **//crucial change below**    
      val subscriber = Subscriber.defaultBuilder(subscriptionName, new receiver).build
      override def run() {
        subscriber.startAsync()
        //ensure subscriber is running as well as spark receiver
        while (subscriber.isRunning && !isStopped()) {
          logger.info(s"${subscriber.getSubscriptionName} receiver running")
          //sleep 10s
          Thread.sleep(10000)
        }
        logger.info(s"${subscriber.getSubscriptionName} receiver stopping")
      }
    }.start()
  }

  def onStop(): Unit = {
    // There is nothing much to do as the thread calling receive()
    // is designed to stop by itself if isStopped() returns false
  }

  class receiver extends MessageReceiver {
    override def receiveMessage(message: PubsubMessage, consumer: AckReplyConsumer): Unit = {
      store(ArrayBuffer(message), message.getAttributesMap)
    }
  }
}