From 1e8ae4337216c75f26de60bdfd55795dfa8f3b81 Mon Sep 17 00:00:00 2001
From: Maxime Steinhausser <maxime.steinhausser@elao.com>
Date: Fri, 28 Aug 2020 12:11:33 +0200
Subject: [PATCH] [Messenger] Don't prevent dispatch out of another bus

---
 src/Symfony/Component/Messenger/CHANGELOG.md  |  1 +
 .../DispatchAfterCurrentBusMiddleware.php     |  9 ++++----
 .../DispatchAfterCurrentBusMiddlewareTest.php | 22 +++++++++++++++++++
 3 files changed, 28 insertions(+), 4 deletions(-)

diff --git a/src/Symfony/Component/Messenger/CHANGELOG.md b/src/Symfony/Component/Messenger/CHANGELOG.md
index 39947599a3546..8a0e0b1dec057 100644
--- a/src/Symfony/Component/Messenger/CHANGELOG.md
+++ b/src/Symfony/Component/Messenger/CHANGELOG.md
@@ -6,6 +6,7 @@ CHANGELOG
 
 * Added `FlattenExceptionNormalizer` to give more information about the exception on Messenger background processes. The `FlattenExceptionNormalizer` has a higher priority than `ProblemNormalizer` and it is only used when the Messenger serialization context is set.
 * Added factory methods to `DelayStamp`.
+* Removed the exception when dispatching a message with a `DispatchAfterCurrentBusStamp` and not in a context of another dispatch call
 
 5.1.0
 -----
diff --git a/src/Symfony/Component/Messenger/Middleware/DispatchAfterCurrentBusMiddleware.php b/src/Symfony/Component/Messenger/Middleware/DispatchAfterCurrentBusMiddleware.php
index 05ee86ffeb77f..a088140b7c784 100644
--- a/src/Symfony/Component/Messenger/Middleware/DispatchAfterCurrentBusMiddleware.php
+++ b/src/Symfony/Component/Messenger/Middleware/DispatchAfterCurrentBusMiddleware.php
@@ -44,12 +44,13 @@ class DispatchAfterCurrentBusMiddleware implements MiddlewareInterface
     public function handle(Envelope $envelope, StackInterface $stack): Envelope
     {
         if (null !== $envelope->last(DispatchAfterCurrentBusStamp::class)) {
-            if (!$this->isRootDispatchCallRunning) {
-                throw new \LogicException(sprintf('You can only use a "%s" stamp in the context of a message handler.', DispatchAfterCurrentBusStamp::class));
+            if ($this->isRootDispatchCallRunning) {
+                $this->queue[] = new QueuedEnvelope($envelope, $stack);
+
+                return $envelope;
             }
-            $this->queue[] = new QueuedEnvelope($envelope, $stack);
 
-            return $envelope;
+            $envelope = $envelope->withoutAll(DispatchAfterCurrentBusStamp::class);
         }
 
         if ($this->isRootDispatchCallRunning) {
diff --git a/src/Symfony/Component/Messenger/Tests/Middleware/DispatchAfterCurrentBusMiddlewareTest.php b/src/Symfony/Component/Messenger/Tests/Middleware/DispatchAfterCurrentBusMiddlewareTest.php
index b2407f1d1089c..c6a1a34da75d4 100644
--- a/src/Symfony/Component/Messenger/Tests/Middleware/DispatchAfterCurrentBusMiddlewareTest.php
+++ b/src/Symfony/Component/Messenger/Tests/Middleware/DispatchAfterCurrentBusMiddlewareTest.php
@@ -256,6 +256,28 @@ public function testHandleDelayedEventFromQueue()
         $messageBus->dispatch($message);
     }
 
+    public function testDispatchOutOfAnotherHandlerDispatchesAndRemoveStamp()
+    {
+        $event = new DummyEvent('First event');
+
+        $middleware = new DispatchAfterCurrentBusMiddleware();
+        $handlingMiddleware = $this->createMock(MiddlewareInterface::class);
+
+        $handlingMiddleware
+            ->method('handle')
+            ->with($this->expectHandledMessage($event))
+            ->will($this->willHandleMessage());
+
+        $eventBus = new MessageBus([
+            $middleware,
+            $handlingMiddleware,
+        ]);
+
+        $enveloppe = $eventBus->dispatch($event, [new DispatchAfterCurrentBusStamp()]);
+
+        self::assertNull($enveloppe->last(DispatchAfterCurrentBusStamp::class));
+    }
+
     private function expectHandledMessage($message): Callback
     {
         return $this->callback(function (Envelope $envelope) use ($message) {