MQTT streaming: Termination signal handling (#1592)

Fixed a couple of behaviours that were not handling the termination signal where it was possible to receive it.
akka · Mar 19, 2019 · bd737f7 · bd737f7
1 parent b860c64
commit bd737f7
Show file tree

Hide file tree

Showing 2 changed files with 62 additions and 55 deletions.
diff --git a/mqtt-streaming/src/main/scala/akka/stream/alpakka/mqtt/streaming/impl/ClientState.scala b/mqtt-streaming/src/main/scala/akka/stream/alpakka/mqtt/streaming/impl/ClientState.scala
@@ -193,61 +193,67 @@ import scala.util.{Failure, Success}
   private val ConsumerNamePrefix = "consumer-"
   private val ProducerNamePrefix = "producer-"
 
-  def disconnected(data: Disconnected)(implicit mat: Materializer): Behavior[Event] = Behaviors.receivePartial {
-    case (context, ConnectReceivedLocally(connect, connectData, remote)) =>
-      val (queue, source) = Source
-        .queue[ForwardConnectCommand](1, OverflowStrategy.dropHead)
-        .toMat(BroadcastHub.sink)(Keep.both)
-        .run()
-      remote.success(source)
-
-      queue.offer(ForwardConnect)
-      data.stash.foreach(context.self.tell)
-
-      if (connect.connectFlags.contains(ConnectFlags.CleanSession)) {
-        context.children.foreach(context.stop)
-        serverConnect(
-          ConnectReceived(
-            connect,
-            connectData,
-            queue,
-            Vector.empty,
-            Map.empty,
-            Map.empty,
-            Vector.empty,
-            Vector.empty,
-            data.consumerPacketRouter,
-            data.producerPacketRouter,
-            data.subscriberPacketRouter,
-            data.unsubscriberPacketRouter,
-            data.settings
-          )
-        )
-      } else {
-        serverConnect(
-          ConnectReceived(
-            connect,
-            connectData,
-            queue,
-            Vector.empty,
-            data.activeConsumers,
-            data.activeProducers,
-            data.pendingLocalPublications,
-            data.pendingRemotePublications,
-            data.consumerPacketRouter,
-            data.producerPacketRouter,
-            data.subscriberPacketRouter,
-            data.unsubscriberPacketRouter,
-            data.settings
-          )
-        )
+  def disconnected(data: Disconnected)(implicit mat: Materializer): Behavior[Event] =
+    Behaviors
+      .receivePartial[Event] {
+        case (context, ConnectReceivedLocally(connect, connectData, remote)) =>
+          val (queue, source) = Source
+            .queue[ForwardConnectCommand](1, OverflowStrategy.dropHead)
+            .toMat(BroadcastHub.sink)(Keep.both)
+            .run()
+          remote.success(source)
+
+          queue.offer(ForwardConnect)
+          data.stash.foreach(context.self.tell)
+
+          if (connect.connectFlags.contains(ConnectFlags.CleanSession)) {
+            context.children.foreach(context.stop)
+            serverConnect(
+              ConnectReceived(
+                connect,
+                connectData,
+                queue,
+                Vector.empty,
+                Map.empty,
+                Map.empty,
+                Vector.empty,
+                Vector.empty,
+                data.consumerPacketRouter,
+                data.producerPacketRouter,
+                data.subscriberPacketRouter,
+                data.unsubscriberPacketRouter,
+                data.settings
+              )
+            )
+          } else {
+            serverConnect(
+              ConnectReceived(
+                connect,
+                connectData,
+                queue,
+                Vector.empty,
+                data.activeConsumers,
+                data.activeProducers,
+                data.pendingLocalPublications,
+                data.pendingRemotePublications,
+                data.consumerPacketRouter,
+                data.producerPacketRouter,
+                data.subscriberPacketRouter,
+                data.unsubscriberPacketRouter,
+                data.settings
+              )
+            )
 
+          }
+        case (_, ConnectionLost) =>
+          Behavior.same
+        case (_, e) =>
+          disconnected(data.copy(stash = data.stash :+ e))
+      }
+      .receiveSignal {
+        case (_, _: Terminated) =>
+          Behaviors.same
       }
-    case (_, ConnectionLost) =>
-      Behavior.same
-    case (_, e) =>
-      disconnected(data.copy(stash = data.stash :+ e))
-  }
 
   def disconnect(context: ActorContext[Event], remote: SourceQueueWithComplete[ForwardConnectCommand], data: Data)(
       implicit mat: Materializer
@@ -319,6 +325,8 @@ import scala.util.{Failure, Success}
             serverConnect(data.copy(stash = data.stash :+ e))
         }
         .receiveSignal {
+          case (_, _: Terminated) =>
+            Behaviors.same
           case (_, PostStop) =>
             data.remote.complete()
             Behaviors.same

diff --git a/mqtt-streaming/src/test/java/docs/javadsl/MqttFlowTest.java b/mqtt-streaming/src/test/java/docs/javadsl/MqttFlowTest.java
@@ -207,8 +207,7 @@ public void establishServerBidirectionalConnectionAndSubscribeToATopic()
                                 JavaConverters.asJavaCollectionConverter(subscribe.topicFilters())
                                     .asJavaCollection();
                             List<Integer> flags =
-                                topicFilters
-                                    .stream()
+                                topicFilters.stream()
                                     .map(x -> x._2().underlying())
                                     .collect(Collectors.toList());
                             queue.offer(