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());
+    }
 }