Refactored the NetworkReceiver API for future stability.

Author: tdas

external/flume/src/main/scala/org/apache/spark/streaming/flume/FlumeInputDStream.scala

@@ -34,6 +34,7 @@ import org.apache.spark.util.Utils
 import org.apache.spark.storage.StorageLevel
 import org.apache.spark.streaming.StreamingContext
 import org.apache.spark.streaming.dstream._
+import org.apache.spark.Logging
 
 private[streaming]
 class FlumeInputDStream[T: ClassTag](
@@ -115,13 +116,13 @@ private[streaming] object SparkFlumeEvent {
 private[streaming]
 class FlumeEventServer(receiver : FlumeReceiver) extends AvroSourceProtocol {
   override def append(event : AvroFlumeEvent) : Status = {
-    receiver.blockGenerator += SparkFlumeEvent.fromAvroFlumeEvent(event)
+    receiver.store(SparkFlumeEvent.fromAvroFlumeEvent(event))
     Status.OK
   }
 
   override def appendBatch(events : java.util.List[AvroFlumeEvent]) : Status = {
     events.foreach (event =>
-      receiver.blockGenerator += SparkFlumeEvent.fromAvroFlumeEvent(event))
+      receiver.store(SparkFlumeEvent.fromAvroFlumeEvent(event)))
     Status.OK
   }
 }
@@ -133,23 +134,21 @@ class FlumeReceiver(
     host: String,
     port: Int,
     storageLevel: StorageLevel
-  ) extends NetworkReceiver[SparkFlumeEvent] {
+  ) extends NetworkReceiver[SparkFlumeEvent](storageLevel) with Logging {
 
-  lazy val blockGenerator = new BlockGenerator(storageLevel)
+  lazy val responder = new SpecificResponder(
+    classOf[AvroSourceProtocol], new FlumeEventServer(this))
+  lazy val server = new NettyServer(responder, new InetSocketAddress(host, port))
 
-  protected override def onStart() {
-    val responder = new SpecificResponder(
-      classOf[AvroSourceProtocol], new FlumeEventServer(this))
-    val server = new NettyServer(responder, new InetSocketAddress(host, port))
-    blockGenerator.start()
+  def onStart() {
     server.start()
     logInfo("Flume receiver started")
   }
 
-  protected override def onStop() {
-    blockGenerator.stop()
+  def onStop() {
+    server.close()
     logInfo("Flume receiver stopped")
   }
 
-  override def getLocationPreference = Some(host)
+  override def preferredLocation = Some(host)
 }

external/kafka/src/main/scala/org/apache/spark/streaming/kafka/KafkaInputDStream.scala

@@ -70,21 +70,15 @@ class KafkaReceiver[
     kafkaParams: Map[String, String],
     topics: Map[String, Int],
     storageLevel: StorageLevel
-  ) extends NetworkReceiver[Any] {
+  ) extends NetworkReceiver[Any](storageLevel) with Logging {
 
-  // Handles pushing data into the BlockManager
-  lazy protected val blockGenerator = new BlockGenerator(storageLevel)
   // Connection to Kafka
   var consumerConnector : ConsumerConnector = null
 
-  def onStop() {
-    blockGenerator.stop()
-  }
+  def onStop() { }
 
   def onStart() {
 
-    blockGenerator.start()
-
     // In case we are using multiple Threads to handle Kafka Messages
     val executorPool = Executors.newFixedThreadPool(topics.values.reduce(_ + _))
 
@@ -130,7 +124,7 @@ class KafkaReceiver[
     def run() {
       logInfo("Starting MessageHandler.")
       for (msgAndMetadata <- stream) {
-        blockGenerator += (msgAndMetadata.key, msgAndMetadata.message)
+        store((msgAndMetadata.key, msgAndMetadata.message))
       }
     }
   }

external/mqtt/src/main/scala/org/apache/spark/streaming/mqtt/MQTTInputDStream.scala

@@ -49,38 +49,34 @@ import org.apache.spark.streaming.dstream._
  */
 
 private[streaming] 
-class MQTTInputDStream[T: ClassTag](
+class MQTTInputDStream(
     @transient ssc_ : StreamingContext,
     brokerUrl: String,
     topic: String,
     storageLevel: StorageLevel
-  ) extends NetworkInputDStream[T](ssc_) with Logging {
+  ) extends NetworkInputDStream[String](ssc_) with Logging {
 
-  def getReceiver(): NetworkReceiver[T] = {
-    new MQTTReceiver(brokerUrl, topic, storageLevel).asInstanceOf[NetworkReceiver[T]]
+  def getReceiver(): NetworkReceiver[String] = {
+    new MQTTReceiver(brokerUrl, topic, storageLevel)
   }
 }
 
 private[streaming] 
-class MQTTReceiver(brokerUrl: String,
-  topic: String,
-  storageLevel: StorageLevel
-  ) extends NetworkReceiver[Any] {
-  lazy protected val blockGenerator = new BlockGenerator(storageLevel)
-  
-  def onStop() {
-    blockGenerator.stop()
-  }
+class MQTTReceiver(
+    brokerUrl: String,
+    topic: String,
+    storageLevel: StorageLevel
+  ) extends NetworkReceiver[String](storageLevel) {
+
+  def onStop() { }
 
   def onStart() {
 
-    blockGenerator.start()
-
     // Set up persistence for messages 
-    var peristance: MqttClientPersistence = new MemoryPersistence()
+    val peristance: MqttClientPersistence = new MemoryPersistence()
 
     // Initializing Mqtt Client specifying brokerUrl, clientID and MqttClientPersistance
-    var client: MqttClient = new MqttClient(brokerUrl, MqttClient.generateClientId(), peristance)
+    val client: MqttClient = new MqttClient(brokerUrl, MqttClient.generateClientId(), peristance)
 
     // Connect to MqttBroker    
     client.connect()
@@ -89,18 +85,19 @@ class MQTTReceiver(brokerUrl: String,
     client.subscribe(topic)
 
     // Callback automatically triggers as and when new message arrives on specified topic
-    var callback: MqttCallback = new MqttCallback() {
+    val callback: MqttCallback = new MqttCallback() {
 
       // Handles Mqtt message 
       override def messageArrived(arg0: String, arg1: MqttMessage) {
-        blockGenerator += new String(arg1.getPayload())
+        store(new String(arg1.getPayload()))
       }
 
       override def deliveryComplete(arg0: IMqttDeliveryToken) {
       }
 
       override def connectionLost(arg0: Throwable) {
-        logInfo("Connection lost " + arg0)
+        store("Connection lost " + arg0)
+        stopOnError(new Exception(arg0))
       }
     }
 

external/mqtt/src/main/scala/org/apache/spark/streaming/mqtt/MQTTUtils.scala

@@ -37,7 +37,7 @@ object MQTTUtils {
       topic: String,
       storageLevel: StorageLevel = StorageLevel.MEMORY_AND_DISK_SER_2
     ): DStream[String] = {
-    new MQTTInputDStream[String](ssc, brokerUrl, topic, storageLevel)
+    new MQTTInputDStream(ssc, brokerUrl, topic, storageLevel)
   }
 
   /**

external/twitter/src/main/scala/org/apache/spark/streaming/twitter/TwitterInputDStream.scala

@@ -25,6 +25,7 @@ import twitter4j.auth.OAuthAuthorization
 import org.apache.spark.streaming._
 import org.apache.spark.streaming.dstream._
 import org.apache.spark.storage.StorageLevel
+import org.apache.spark.Logging
 
 /* A stream of Twitter statuses, potentially filtered by one or more keywords.
 *
@@ -59,17 +60,15 @@ class TwitterReceiver(
     twitterAuth: Authorization,
     filters: Seq[String],
     storageLevel: StorageLevel
-  ) extends NetworkReceiver[Status] {
+  ) extends NetworkReceiver[Status](storageLevel) with Logging {
 
   var twitterStream: TwitterStream = _
-  lazy val blockGenerator = new BlockGenerator(storageLevel)
 
-  protected override def onStart() {
-    blockGenerator.start()
+  def onStart() {
     twitterStream = new TwitterStreamFactory().getInstance(twitterAuth)
     twitterStream.addListener(new StatusListener {
       def onStatus(status: Status) = {
-        blockGenerator += status
+        store(status)
       }
       // Unimplemented
       def onDeletionNotice(statusDeletionNotice: StatusDeletionNotice) {}
@@ -89,8 +88,7 @@ class TwitterReceiver(
     logInfo("Twitter receiver started")
   }
 
-  protected override def onStop() {
-    blockGenerator.stop()
+  def onStop() {
     twitterStream.shutdown()
     logInfo("Twitter receiver stopped")
   }