Merge pull request #20225 from akka/wip-17037-stash-patriknw

str: Support stash in ActorPublisher, #17037

akka-stream-tests/src/test/scala/akka/stream/actor/ActorPublisherSpec.scala

@@ -10,10 +10,10 @@ import akka.stream.testkit._
 import akka.stream.testkit.Utils._
 import akka.testkit.TestEvent.Mute
 import akka.testkit.{ AkkaSpec, EventFilter, ImplicitSender, TestProbe }
-
 import scala.annotation.tailrec
 import scala.concurrent.duration._
 import scala.util.control.NoStackTrace
+import akka.actor.Stash
 
 object ActorPublisherSpec {
 
@@ -29,6 +29,12 @@ object ActorPublisherSpec {
     else p
   }
 
+  def testPublisherWithStashProps(probe: ActorRef, useTestDispatcher: Boolean = true): Props = {
+    val p = Props(new TestPublisherWithStash(probe))
+    if (useTestDispatcher) p.withDispatcher("akka.test.stream-dispatcher")
+    else p
+  }
+
   case class TotalDemand(elements: Long)
   case class Produce(elem: String)
   case class Err(reason: String)
@@ -53,6 +59,19 @@ object ActorPublisherSpec {
     }
   }
 
+  class TestPublisherWithStash(probe: ActorRef) extends TestPublisher(probe) with Stash {
+
+    override def receive = stashing
+
+    def stashing: Receive = {
+      case "unstash" ⇒
+        unstashAll()
+        context.become(super.receive)
+      case _ ⇒ stash()
+    }
+
+  }
+
   def senderProps: Props = Props[Sender].withDispatcher("akka.test.stream-dispatcher")
 
   class Sender extends ActorPublisher[Int] {
@@ -445,6 +464,22 @@ class ActorPublisherSpec extends AkkaSpec(ActorPublisherSpec.config) with Implic
       expectMsgType[String] should include("my-dispatcher1")
     }
 
+    "handle stash" in {
+      val probe = TestProbe()
+      val ref = system.actorOf(testPublisherWithStashProps(probe.ref))
+      val p = ActorPublisher[String](ref)
+      val s = TestSubscriber.probe[String]()
+      p.subscribe(s)
+      s.request(2)
+      s.request(3)
+      ref ! "unstash"
+      probe.expectMsg(TotalDemand(5))
+      probe.expectMsg(TotalDemand(5))
+      s.request(4)
+      probe.expectMsg(TotalDemand(9))
+      s.cancel()
+    }
+
   }
 
 }

akka-stream/src/main/scala/akka/stream/actor/ActorPublisher.scala

@@ -45,7 +45,18 @@ object ActorPublisherMessage {
    * more elements.
    * @param n number of requested elements
    */
-  final case class Request(n: Long) extends ActorPublisherMessage with NoSerializationVerificationNeeded
+  final case class Request(n: Long) extends ActorPublisherMessage with NoSerializationVerificationNeeded {
+    private var processed = false
+    /**
+     * INTERNAL API: needed for stash support
+     */
+    private[akka] def markProcessed(): Unit = processed = true
+
+    /**
+     * INTERNAL API: needed for stash support
+     */
+    private[akka] def isProcessed(): Boolean = processed
+  }
 
   /**
    * This message is delivered to the [[ActorPublisher]] actor when the stream subscriber cancels the
@@ -256,7 +267,11 @@ trait ActorPublisher[T] extends Actor {
    * INTERNAL API
    */
   protected[akka] override def aroundReceive(receive: Receive, msg: Any): Unit = msg match {
-    case Request(n) ⇒
+    case req @ Request(n) ⇒
+      if (req.isProcessed()) {
+        // it's an unstashed Request, demand is already handled
+        super.aroundReceive(receive, req)
+      } else {
         if (n < 1) {
           if (lifecycleState == Active)
             onError(numberOfElementsInRequestMustBePositiveException)
@@ -264,7 +279,9 @@ trait ActorPublisher[T] extends Actor {
           demand += n
           if (demand < 0)
             demand = Long.MaxValue // Long overflow, Reactive Streams Spec 3:17: effectively unbounded
-        super.aroundReceive(receive, msg)
+          req.markProcessed()
+          super.aroundReceive(receive, req)
+        }
       }
 
     case Subscribe(sub: Subscriber[_]) ⇒
@@ -293,8 +310,11 @@ trait ActorPublisher[T] extends Actor {
       }
 
     case Cancel ⇒
+      if (lifecycleState != Canceled) {
+        // possible to receive again in case of stash
         cancelSelf()
         super.aroundReceive(receive, msg)
+      }
 
     case SubscriptionTimeoutExceeded ⇒
       if (!scheduledSubscriptionTimeout.isCancelled) {