Mqtt Source

mqtt/build.sbt

@@ -1,5 +1,8 @@
 name := "akka-stream-contrib-mqtt"
 
 libraryDependencies ++= Seq(
-  "org.eclipse.paho" % "org.eclipse.paho.client.mqttv3" % "1.1.0"
+  "org.eclipse.paho" % "org.eclipse.paho.client.mqttv3" % "1.1.0",       // Eclipse Public License Version 1.0
+  "io.moquette"      % "moquette-broker"                % "0.8.1" % Test // Apache License Version 2.0
 )
+
+resolvers += "moquette" at "http://dl.bintray.com/andsel/maven/"

mqtt/src/main/scala/Mqtt.scala

@@ -0,0 +1,67 @@
+/*
+ * Copyright (C) 2016 Lightbend Inc. <http://www.lightbend.com>
+ */
+package akka.stream.contrib.mqtt
+
+import akka.util.ByteString
+import akka.stream.stage._
+import org.eclipse.paho.client.mqttv3.{ MqttMessage => PahoMqttMessage, _ }
+import scala.util._
+import scala.language.implicitConversions
+
+case class MqttSourceSettings(
+  connectionSettings: MqttConnectionSettings,
+  topics:             Map[String, Int]
+)
+
+case class MqttConnectionSettings(
+  broker:      String,
+  clientId:    String,
+  persistence: MqttClientPersistence
+)
+
+case class MqttMessage(topic: String, payload: ByteString)
+
+trait MqttConnectorLogic { this: GraphStageLogic =>
+
+  import MqttConnectorLogic._
+
+  def connectionSettings: MqttConnectionSettings
+  def onConnect: AsyncCallback[IMqttAsyncClient]
+  def onMessage: AsyncCallback[MqttMessage]
+  def onConnectionLost: AsyncCallback[Throwable]
+
+  final override def preStart(): Unit = {
+    val client = new MqttAsyncClient(
+      connectionSettings.broker,
+      connectionSettings.clientId,
+      connectionSettings.persistence
+    )
+
+    client.connect((), connectHandler(client))
+    client.setCallback(new MqttCallback {
+      def messageArrived(topic: String, message: PahoMqttMessage) =
+        onMessage.invoke(MqttMessage(topic, ByteString(message.getPayload)))
+
+      def deliveryComplete(token: IMqttDeliveryToken) =
+        println(s"Delivery complete $token")
+
+      def connectionLost(cause: Throwable) =
+        onConnectionLost.invoke(cause)
+    })
+  }
+
+  private val connectHandler: IMqttAsyncClient => Try[IMqttToken] => Unit = client => {
+    case Success(_)  => onConnect.invoke(client)
+    case Failure(ex) => onConnectionLost.invoke(ex)
+  }
+}
+
+object MqttConnectorLogic {
+
+  implicit def funcToMqttActionListener(func: Try[IMqttToken] => Unit): IMqttActionListener = new IMqttActionListener {
+    def onSuccess(token: IMqttToken) = func(Success(token))
+    def onFailure(token: IMqttToken, ex: Throwable) = func(Failure(ex))
+  }
+
+}

mqtt/src/main/scala/MqttSource.scala

@@ -0,0 +1,87 @@
+/*
+ * Copyright (C) 2016 Lightbend Inc. <http://www.lightbend.com>
+ */
+package akka.stream.contrib.mqtt
+
+import akka.NotUsed
+import akka.stream.scaladsl._
+import akka.stream.stage._
+import akka.stream._
+import scala.collection.mutable
+import scala.concurrent._
+import org.eclipse.paho.client.mqttv3.IMqttAsyncClient
+import org.eclipse.paho.client.mqttv3.IMqttToken
+import scala.util.Try
+
+object MqttSource {
+
+  /**
+   * Scala API:
+   */
+  def apply(settings: MqttSourceSettings, bufferSize: Int): Source[MqttMessage, Future[NotUsed]] =
+    Source.fromGraph(new MqttSource(settings, bufferSize))
+
+  /**
+   * Java API:
+   */
+  def create(settings: MqttSourceSettings, bufferSize: Int): akka.stream.javadsl.Source[MqttMessage, Future[NotUsed]] =
+    akka.stream.javadsl.Source.fromGraph(new MqttSource(settings, bufferSize))
+}
+
+final class MqttSource(settings: MqttSourceSettings, bufferSize: Int) extends GraphStageWithMaterializedValue[SourceShape[MqttMessage], Future[NotUsed]] {
+
+  import MqttConnectorLogic._
+
+  val out = Outlet[MqttMessage]("MqttSource.out")
+  override val shape: SourceShape[MqttMessage] = SourceShape.of(out)
+  override protected def initialAttributes: Attributes = Attributes.name("MqttSource")
+
+  private val subscriptionPromise = Promise[NotUsed]
+
+  override def createLogicAndMaterializedValue(inheritedAttributes: Attributes) = (new GraphStageLogic(shape) with MqttConnectorLogic {
+
+    override val connectionSettings = settings.connectionSettings
+
+    override val onConnect = getAsyncCallback[IMqttAsyncClient](handleConnection)
+    override val onMessage = getAsyncCallback[MqttMessage](handleMessage)
+    override val onConnectionLost = getAsyncCallback[Throwable](handleConnectionLost)
+
+    private val queue = mutable.Queue[MqttMessage]()
+    private val mqttSubscriptionCallback: Try[IMqttToken] => Unit = conn =>
+      subscriptionPromise.complete(conn.map { _ => NotUsed })
+
+    setHandler(out, new OutHandler {
+      override def onPull(): Unit = {
+        if (queue.nonEmpty) {
+          pushAndAckMessage(queue.dequeue())
+        }
+      }
+    })
+
+    def handleConnection(client: IMqttAsyncClient) = {
+      val (topics, qos) = settings.topics.unzip
+      client.subscribe(topics.toArray, qos.toArray, (), mqttSubscriptionCallback)
+    }
+
+    def handleMessage(message: MqttMessage): Unit = {
+      if (isAvailable(out)) {
+        pushAndAckMessage(message)
+      } else {
+        if (queue.size + 1 > bufferSize) {
+          failStage(new RuntimeException(s"Reached maximum buffer size $bufferSize"))
+        } else {
+          queue.enqueue(message)
+        }
+      }
+    }
+
+    def pushAndAckMessage(message: MqttMessage): Unit = {
+      push(out, message)
+    }
+
+    def handleConnectionLost(ex: Throwable) =
+      failStage(ex)
+
+  }, subscriptionPromise.future)
+
+}

mqtt/src/test/scala/MqttSourceSpec.scala

@@ -0,0 +1,82 @@
+/*
+ * Copyright (C) 2016 Lightbend Inc. <http://www.lightbend.com>
+ */
+package akka.stream.contrib.mqtt
+
+import akka.actor.ActorSystem
+import akka.stream._
+import akka.stream.scaladsl.Keep
+import akka.stream.testkit.scaladsl.TestSink
+import akka.util.ByteString
+import io.moquette.server.Server
+import io.moquette.proto.messages.PublishMessage
+import io.moquette.proto.messages.AbstractMessage.QOSType
+import org.eclipse.paho.client.mqttv3.persist.MemoryPersistence
+import org.scalatest._
+import org.scalatest.concurrent.ScalaFutures
+import scala.concurrent.duration._
+
+class MqttSourceSpec extends WordSpec with Matchers with ScalaFutures {
+
+  "mqtt source" should {
+    "receive a message from a topic" in withBroker(Map("topic1" -> 0)) { settings => implicit server => implicit sys => implicit mat =>
+      val (subscriptionFuture, probe) = MqttSource(settings, 8).toMat(TestSink.probe)(Keep.both).run()
+      whenReady(subscriptionFuture) { _ =>
+        publish("topic1", "ohi")
+        probe.requestNext shouldBe MqttMessage("topic1", ByteString("ohi"))
+      }
+    }
+
+    "receive messages from multiple topics" in withBroker(Map("topic1" -> 0, "topic2" -> 0)) { settings => implicit server => implicit sys => implicit mat =>
+      val (subscriptionFuture, probe) = MqttSource(settings, 8).toMat(TestSink.probe)(Keep.both).run()
+      whenReady(subscriptionFuture) { _ =>
+        publish("topic1", "ohi")
+        publish("topic2", "hello again")
+        probe.requestNext shouldBe MqttMessage("topic1", ByteString("ohi"))
+        probe.requestNext shouldBe MqttMessage("topic2", ByteString("hello again"))
+      }
+    }
+
+    "fail stream when disconnected" in withBroker(Map("topic1" -> 0)) { settings => implicit server => implicit sys => implicit mat =>
+      val (subscriptionFuture, probe) = MqttSource(settings, 8).toMat(TestSink.probe)(Keep.both).run()
+      whenReady(subscriptionFuture) { _ =>
+        publish("topic1", "ohi")
+        probe.requestNext shouldBe MqttMessage("topic1", ByteString("ohi"))
+
+        server.stopServer()
+        probe.expectError.getMessage should be("Connection lost")
+      }
+    }
+  }
+
+  def publish(topic: String, payload: String)(implicit server: Server) = {
+    val msg = new PublishMessage()
+    msg.setPayload(ByteString(payload).toByteBuffer)
+    msg.setTopicName(topic)
+    msg.setQos(QOSType.MOST_ONE)
+    server.internalPublish(msg)
+  }
+
+  def withBroker(subscriptions: Map[String, Int])(test: MqttSourceSettings => Server => ActorSystem => Materializer => Any) = {
+    implicit val sys = ActorSystem("MqttSourceSpec")
+    val mat = ActorMaterializer()
+
+    val settings = MqttSourceSettings(
+      MqttConnectionSettings(
+        "tcp://localhost:1883",
+        "test-client",
+        new MemoryPersistence
+      ),
+      subscriptions
+    )
+
+    val server = new Server()
+    server.startServer()
+    try {
+      test(settings)(server)(sys)(mat)
+    } finally {
+      server.stopServer()
+    }
+  }
+
+}