symfony · fabpot · Oct 22, 2019 · Oct 22, 2019
diff --git a/src/Symfony/Bridge/Doctrine/Messenger/DoctrineClearEntityManagerMiddleware.php b/src/Symfony/Bridge/Doctrine/Messenger/DoctrineClearEntityManagerMiddleware.php
@@ -14,6 +14,7 @@
 use Doctrine\ORM\EntityManagerInterface;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Middleware\StackInterface;
+use Symfony\Component\Messenger\Stamp\ReceivedStamp;
 
 /**
  * Clears entity manager after calling all handlers.
@@ -27,7 +28,9 @@ protected function handleForManager(EntityManagerInterface $entityManager, Envel
         try {
             return $stack->next()->handle($envelope, $stack);
         } finally {
-            $entityManager->clear();
+            if (null !== $envelope->last(ReceivedStamp::class)) {
+                $entityManager->clear();
+            }
         }
     }
 }
diff --git a/src/Symfony/Bridge/Doctrine/Messenger/DoctrineCloseConnectionMiddleware.php b/src/Symfony/Bridge/Doctrine/Messenger/DoctrineCloseConnectionMiddleware.php
@@ -14,6 +14,7 @@
 use Doctrine\ORM\EntityManagerInterface;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Middleware\StackInterface;
+use Symfony\Component\Messenger\Stamp\ReceivedStamp;
 
 /**
  * Closes connection and therefore saves number of connections.
@@ -29,7 +30,9 @@ protected function handleForManager(EntityManagerInterface $entityManager, Envel
 
             return $stack->next()->handle($envelope, $stack);
         } finally {
-            $connection->close();
+            if (null !== $envelope->last(ReceivedStamp::class)) {
+                $connection->close();
+            }
         }
     }
 }
diff --git a/src/Symfony/Bridge/Doctrine/Messenger/DoctrinePingConnectionMiddleware.php b/src/Symfony/Bridge/Doctrine/Messenger/DoctrinePingConnectionMiddleware.php
@@ -14,6 +14,7 @@
 use Doctrine\ORM\EntityManagerInterface;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Middleware\StackInterface;
+use Symfony\Component\Messenger\Stamp\ReceivedStamp;
 
 /**
  * Checks whether the connection is still open or reconnects otherwise.
@@ -23,6 +24,15 @@
 class DoctrinePingConnectionMiddleware extends AbstractDoctrineMiddleware
 {
     protected function handleForManager(EntityManagerInterface $entityManager, Envelope $envelope, StackInterface $stack): Envelope
+    {
+        if (null !== $envelope->last(ReceivedStamp::class)) {
+            $this->pingConnection($entityManager);
+        }
+
+        return $stack->next()->handle($envelope, $stack);
+    }
+
+    private function pingConnection(EntityManagerInterface $entityManager)
     {
         $connection = $entityManager->getConnection();
 
@@ -34,7 +44,5 @@ protected function handleForManager(EntityManagerInterface $entityManager, Envel
         if (!$entityManager->isOpen()) {
             $this->managerRegistry->resetManager($this->entityManagerName);
         }
-
-        return $stack->next()->handle($envelope, $stack);
     }
 }
diff --git a/src/Symfony/Bridge/Doctrine/Tests/Messenger/DoctrineClearEntityManagerMiddlewareTest.php b/src/Symfony/Bridge/Doctrine/Tests/Messenger/DoctrineClearEntityManagerMiddlewareTest.php
@@ -16,6 +16,7 @@
 use Symfony\Bridge\Doctrine\Messenger\DoctrineClearEntityManagerMiddleware;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Exception\UnrecoverableMessageHandlingException;
+use Symfony\Component\Messenger\Stamp\ReceivedStamp;
 use Symfony\Component\Messenger\Test\Middleware\MiddlewareTestCase;
 
 class DoctrineClearEntityManagerMiddlewareTest extends MiddlewareTestCase
@@ -34,7 +35,10 @@ public function testMiddlewareClearEntityManager()
 
         $middleware = new DoctrineClearEntityManagerMiddleware($managerRegistry, 'default');
 
-        $middleware->handle(new Envelope(new \stdClass()), $this->getStackMock());
+        $envelope = new Envelope(new \stdClass(), [
+            new ReceivedStamp('async'),
+        ]);
+        $middleware->handle($envelope, $this->getStackMock());
     }
 
     public function testInvalidEntityManagerThrowsException()
@@ -51,4 +55,22 @@ public function testInvalidEntityManagerThrowsException()
 
         $middleware->handle(new Envelope(new \stdClass()), $this->getStackMock(false));
     }
+
+    public function testMiddlewareDoesNotClearInNonWorkerContext()
+    {
+        $entityManager = $this->createMock(EntityManagerInterface::class);
+        $entityManager->expects($this->never())
+            ->method('clear');
+
+        $managerRegistry = $this->createMock(ManagerRegistry::class);
+        $managerRegistry
+            ->method('getManager')
+            ->with('default')
+            ->willReturn($entityManager);
+
+        $middleware = new DoctrineClearEntityManagerMiddleware($managerRegistry, 'default');
+
+        $envelope = new Envelope(new \stdClass());
+        $middleware->handle($envelope, $this->getStackMock());
+    }
 }
diff --git a/src/Symfony/Bridge/Doctrine/Tests/Messenger/DoctrineCloseConnectionMiddlewareTest.php b/src/Symfony/Bridge/Doctrine/Tests/Messenger/DoctrineCloseConnectionMiddlewareTest.php
@@ -17,6 +17,7 @@
 use Symfony\Bridge\Doctrine\Messenger\DoctrineCloseConnectionMiddleware;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Exception\UnrecoverableMessageHandlingException;
+use Symfony\Component\Messenger\Stamp\ReceivedStamp;
 use Symfony\Component\Messenger\Test\Middleware\MiddlewareTestCase;
 
 class DoctrineCloseConnectionMiddlewareTest extends MiddlewareTestCase
@@ -49,7 +50,10 @@ public function testMiddlewareCloseConnection()
             ->method('close')
         ;
 
-        $this->middleware->handle(new Envelope(new \stdClass()), $this->getStackMock());
+        $envelope = new Envelope(new \stdClass(), [
+            new ReceivedStamp('async'),
+        ]);
+        $this->middleware->handle($envelope, $this->getStackMock());
     }
 
     public function testInvalidEntityManagerThrowsException()
@@ -66,4 +70,14 @@ public function testInvalidEntityManagerThrowsException()
 
         $middleware->handle(new Envelope(new \stdClass()), $this->getStackMock(false));
     }
+
+    public function testMiddlewareNotCloseInNonWorkerContext()
+    {
+        $this->connection->expects($this->never())
+            ->method('close')
+        ;
+
+        $envelope = new Envelope(new \stdClass());
+        $this->middleware->handle($envelope, $this->getStackMock());
+    }
 }
diff --git a/src/Symfony/Bridge/Doctrine/Tests/Messenger/DoctrinePingConnectionMiddlewareTest.php b/src/Symfony/Bridge/Doctrine/Tests/Messenger/DoctrinePingConnectionMiddlewareTest.php
@@ -17,6 +17,7 @@
 use Symfony\Bridge\Doctrine\Messenger\DoctrinePingConnectionMiddleware;
 use Symfony\Component\Messenger\Envelope;
 use Symfony\Component\Messenger\Exception\UnrecoverableMessageHandlingException;
+use Symfony\Component\Messenger\Stamp\ReceivedStamp;
 use Symfony\Component\Messenger\Test\Middleware\MiddlewareTestCase;
 
 class DoctrinePingConnectionMiddlewareTest extends MiddlewareTestCase
@@ -56,7 +57,10 @@ public function testMiddlewarePingOk()
             ->method('connect')
         ;
 
-        $this->middleware->handle(new Envelope(new \stdClass()), $this->getStackMock());
+        $envelope = new Envelope(new \stdClass(), [
+            new ReceivedStamp('async'),
+        ]);
+        $this->middleware->handle($envelope, $this->getStackMock());
     }
 
     public function testMiddlewarePingResetEntityManager()
@@ -70,7 +74,10 @@ public function testMiddlewarePingResetEntityManager()
             ->with($this->entityManagerName)
         ;
 
-        $this->middleware->handle(new Envelope(new \stdClass()), $this->getStackMock());
+        $envelope = new Envelope(new \stdClass(), [
+            new ReceivedStamp('async'),
+        ]);
+        $this->middleware->handle($envelope, $this->getStackMock());
     }
 
     public function testInvalidEntityManagerThrowsException()
@@ -87,4 +94,21 @@ public function testInvalidEntityManagerThrowsException()
 
         $middleware->handle(new Envelope(new \stdClass()), $this->getStackMock(false));
     }
+
+    public function testMiddlewareNoPingInNonWorkerContext()
+    {
+        $this->connection->expects($this->never())
+            ->method('ping')
+            ->willReturn(false);
+
+        $this->connection->expects($this->never())
+            ->method('close')
+        ;
+        $this->connection->expects($this->never())
+            ->method('connect')
+        ;
+
+        $envelope = new Envelope(new \stdClass());
+        $this->middleware->handle($envelope, $this->getStackMock());
+    }
 }