Spark Streaming Websocket Receiver

Job.scala

val ssc = new StreamingContext("local", "datastream", Seconds(15))

// create InputDStream

ssc.registerInputStream(stream)

// interact with stream

ssc.start()

PriceEcho.scala

class PriceEcho extends PriceWebSocketClient {
  createSocket(println)
}

PriceReceiver.scala

import spark.streaming.dstream.NetworkReceiver
import spark.storage.StorageLevel

class PriceReceiver extends NetworkReceiver[PriceUpdate] with PriceWebSocketClient {

  lazy val blockGenerator = new BlockGenerator(StorageLevel.MEMORY_ONLY_SER)
  
  protected override def onStart() {
    blockGenerator.start
    createSocket(m => {
      val priceUpdate = PriceUpdate(m)
      blockGenerator += priceUpdate
    })  
  }

  protected override def onStop() {
    blockGenerator.stop
    websocket.shutdown
  }
}

PriceUpdate.scala

import scala.util.parsing.json.JSON
import scala.collection.JavaConversions
import java.util.TreeMap

case class PriceUpdate(id: String, price: Double, lastPrice: Double)

object PriceUpdate {
  // No native Scala TreeMap as yet, so I'll borrow Java's
  val lastPrices = JavaConversions.asMap(new TreeMap[String,Double])

  def apply(text: String): PriceUpdate = {
    val (id, price) = getIdAndPriceFromJSON(text)
    val lastPrice: Double = lastPrices.getOrElse(id, price)
    lastPrices.put(id, price)
    PriceUpdate(id, price, lastPrice)
  }

  def getIdAndPriceFromJSON(text: String) = // snip - simple JSON processing

}

PriceWebSocketClient.scala

import scalawebsocket.WebSocket

trait PriceWebSocketClient {
  import Listings._

  def createSocket(handleMessage: String => Unit) = {
    websocket = WebSocket().open("ws://localhost:8080/1.0/marketDataWs").onTextMessage(m => {
      handleMessage(m)
    })
    subscriptions.foreach(listing => websocket.sendText("{\"subscribe\":{" + listing + "}}"))
  }

  var websocket: WebSocket = _
}

Stream.scala

object Stream extends NetworkInputDStream[PriceUpdate](ssc) {
  override def getReceiver(): NetworkReceiver[PriceUpdate] = {
    new PriceReceiver()
  }
}