[ BUG ] Fix a race condition and a mutable state closure on a future

jeanadrien · Dec 12, 2017 · 9849cc2 · 9849cc2
1 parent 6d4124a
commit 9849cc2
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 7 deletions.
diff --git a/scalastyle-config.xml b/scalastyle-config.xml
@@ -85,7 +85,7 @@
  </check>
  <check level="warning" class="org.scalastyle.scalariform.CyclomaticComplexityChecker" enabled="true">
   <parameters>
-   <parameter name="maximum"><![CDATA[10]]></parameter>
+   <parameter name="maximum"><![CDATA[15]]></parameter>
   </parameters>
  </check>
  <check level="warning" class="org.scalastyle.scalariform.UppercaseLChecker" enabled="true"></check>

diff --git a/src/main/scala/com/github/jeanadrien/gatling/mqtt/client/MqttClient.scala b/src/main/scala/com/github/jeanadrien/gatling/mqtt/client/MqttClient.scala
@@ -19,11 +19,12 @@ abstract class MqttClient(val gatlingMqttId : String) extends Actor with LazyLog
 
     import context._
 
-    // FIXME: Timeout and feebackListener polution
     private var feedbackListener : Map[String, List[(FeedbackFunction, ActorRef)]] = Map()
 
     private var waitForMessagesReceivedListeners : List[ActorRef] = Nil
 
+    private var delayIncomingMessages : Boolean = false
+
     protected def connect(replyTo : ActorRef) : Unit
 
     protected def subscribe(topics : List[(String, MqttQoS)], replyTo : ActorRef) : Unit
@@ -51,6 +52,9 @@ abstract class MqttClient(val gatlingMqttId : String) extends Actor with LazyLog
                     fn(payload)
                 }
                 // fire matching listeners
+                if (matching.nonEmpty) {
+                    logger.debug(s"Client ${gatlingMqttId} : message on topic $topic matches ${matching.length} awaiting listener(s)")
+                }
                 matching.foreach { case (_, replyTo) =>
                     replyTo ! FeedbackReceived
                 }
@@ -74,15 +78,20 @@ abstract class MqttClient(val gatlingMqttId : String) extends Actor with LazyLog
         replyTo         : ActorRef
     ) = {
         implicit val timeout = Timeout(1 minute)
+        delayIncomingMessages = true
         self ? MqttCommands.Publish(
             topic = topic,
             payload = payload,
             mqttQoS = qos,
             retain = retain
         ) andThen {
             case Success(_) =>
-                addFeedbackListener(topic, (payloadFeedback, replyTo))
+                // Register listener
+                self ! PublishAckRegisterFeedback(
+                    topic, payloadFeedback, replyTo
+                )
             case Failure(th) =>
+                delayIncomingMessages = false
                 replyTo ! akka.actor.Status.Failure(th)
         }
     }
@@ -100,6 +109,11 @@ abstract class MqttClient(val gatlingMqttId : String) extends Actor with LazyLog
         waitForMessagesReceivedListeners = Nil
     }
 
+    private def publishAckRegisterFeedback(topic : String, payloadFeedback : FeedbackFunction, listener : ActorRef): Unit = {
+        delayIncomingMessages = false
+        addFeedbackListener(topic, (payloadFeedback, listener))
+    }
+
     override def postStop() = {
         super.postStop()
         close()
@@ -114,10 +128,16 @@ abstract class MqttClient(val gatlingMqttId : String) extends Actor with LazyLog
             publish(topic, payload, mqttQoS, retain, sender())
         case PublishAndWait(topic, payload, payloadFeedback, mqttQoS, retain) =>
             publishAndWait(topic, payload, payloadFeedback, mqttQoS, retain, sender())
-        case OnPublish(topic, payload) =>
-            onPublish(topic, payload)
-            if (feedbackListener.isEmpty) {
-                fireAllWaitForMessageListeners
+        case PublishAckRegisterFeedback(topic, payloadFeedback, listener) =>
+            publishAckRegisterFeedback(topic, payloadFeedback, listener)
+        case msg @ OnPublish(topic, payload) =>
+            if (delayIncomingMessages) {
+                system.scheduler.scheduleOnce(1 milliseconds, self, msg)
+            } else {
+                onPublish(topic, payload)
+                if (feedbackListener.isEmpty) {
+                    fireAllWaitForMessageListeners
+                }
             }
         case WaitForMessages =>
             waitForMessages(sender())

diff --git a/src/main/scala/com/github/jeanadrien/gatling/mqtt/client/MqttCommands.scala b/src/main/scala/com/github/jeanadrien/gatling/mqtt/client/MqttCommands.scala
@@ -1,5 +1,6 @@
 package com.github.jeanadrien.gatling.mqtt.client
 
+import akka.actor.ActorRef
 import com.github.jeanadrien.gatling.mqtt.client.MqttClient.FeedbackFunction
 import com.github.jeanadrien.gatling.mqtt.client.MqttQoS.MqttQoS
 
@@ -28,6 +29,10 @@ object MqttCommands {
         topic : String, payload : Array[Byte], payloadFeedback : FeedbackFunction, qos : MqttQoS, retain : Boolean
     ) extends MqttCommands
 
+    case class PublishAckRegisterFeedback(
+        topic : String, payloadFeedback : FeedbackFunction, listener : ActorRef
+    ) extends MqttCommands
+
     case object FeedbackReceived extends MqttCommands
 
     case object WaitForMessages extends MqttCommands