From 28e0fb049b6c681dc71b13871aa4ed01bcb876b3 Mon Sep 17 00:00:00 2001
From: Alejandro Celaya <alejandro@alejandrocelaya.com>
Date: Thu, 9 Apr 2020 11:54:54 +0200
Subject: [PATCH] Added check to ensure DB connection is gracefully recovered
 on swoole task workers

---
 .../EventDispatcher/LocateShortUrlVisit.php   | 38 +++++++++++++------
 .../LocateShortUrlVisitTest.php               |  5 +++
 2 files changed, 31 insertions(+), 12 deletions(-)

diff --git a/module/Core/src/EventDispatcher/LocateShortUrlVisit.php b/module/Core/src/EventDispatcher/LocateShortUrlVisit.php
index 6abbe02b..ff382770 100644
--- a/module/Core/src/EventDispatcher/LocateShortUrlVisit.php
+++ b/module/Core/src/EventDispatcher/LocateShortUrlVisit.php
@@ -9,6 +9,7 @@ use Psr\EventDispatcher\EventDispatcherInterface;
 use Psr\Log\LoggerInterface;
 use Shlinkio\Shlink\CLI\Exception\GeolocationDbUpdateFailedException;
 use Shlinkio\Shlink\CLI\Util\GeolocationDbUpdaterInterface;
+use Shlinkio\Shlink\Common\Doctrine\ReopeningEntityManager;
 use Shlinkio\Shlink\Core\Entity\Visit;
 use Shlinkio\Shlink\Core\Entity\VisitLocation;
 use Shlinkio\Shlink\IpGeolocation\Exception\WrongIpException;
@@ -41,22 +42,35 @@ class LocateShortUrlVisit
 
     public function __invoke(ShortUrlVisited $shortUrlVisited): void
     {
+        // FIXME Temporarily handling DB connection reset here to fix https://github.com/shlinkio/shlink/issues/717
+        //       Remove when https://github.com/shlinkio/shlink-event-dispatcher/issues/23 is implemented
+        if ($this->em instanceof ReopeningEntityManager) {
+            $this->em->open();
+        }
+
         $visitId = $shortUrlVisited->visitId();
 
-        /** @var Visit|null $visit */
-        $visit = $this->em->find(Visit::class, $visitId);
-        if ($visit === null) {
-            $this->logger->warning('Tried to locate visit with id "{visitId}", but it does not exist.', [
-                'visitId' => $visitId,
-            ]);
-            return;
-        }
+        try {
+            /** @var Visit|null $visit */
+            $visit = $this->em->find(Visit::class, $visitId);
+            if ($visit === null) {
+                $this->logger->warning('Tried to locate visit with id "{visitId}", but it does not exist.', [
+                    'visitId' => $visitId,
+                ]);
+                return;
+            }
 
-        if ($this->downloadOrUpdateGeoLiteDb($visitId)) {
-            $this->locateVisit($visitId, $shortUrlVisited->originalIpAddress(), $visit);
-        }
+            if ($this->downloadOrUpdateGeoLiteDb($visitId)) {
+                $this->locateVisit($visitId, $shortUrlVisited->originalIpAddress(), $visit);
+            }
 
-        $this->eventDispatcher->dispatch(new VisitLocated($visitId));
+            $this->eventDispatcher->dispatch(new VisitLocated($visitId));
+        } finally {
+            // FIXME Temporarily handling DB connection reset here to fix https://github.com/shlinkio/shlink/issues/717
+            //       Remove when https://github.com/shlinkio/shlink-event-dispatcher/issues/23 is implemented
+            $this->em->getConnection()->close();
+            $this->em->clear();
+        }
     }
 
     private function downloadOrUpdateGeoLiteDb(string $visitId): bool
diff --git a/module/Core/test/EventDispatcher/LocateShortUrlVisitTest.php b/module/Core/test/EventDispatcher/LocateShortUrlVisitTest.php
index 087c0e0b..e35e7921 100644
--- a/module/Core/test/EventDispatcher/LocateShortUrlVisitTest.php
+++ b/module/Core/test/EventDispatcher/LocateShortUrlVisitTest.php
@@ -4,6 +4,7 @@ declare(strict_types=1);
 
 namespace ShlinkioTest\Shlink\Core\EventDispatcher;
 
+use Doctrine\DBAL\Connection;
 use Doctrine\ORM\EntityManagerInterface;
 use PHPUnit\Framework\TestCase;
 use Prophecy\Argument;
@@ -37,6 +38,10 @@ class LocateShortUrlVisitTest extends TestCase
     {
         $this->ipLocationResolver = $this->prophesize(IpLocationResolverInterface::class);
         $this->em = $this->prophesize(EntityManagerInterface::class);
+        $conn = $this->prophesize(Connection::class);
+        $this->em->getConnection()->willReturn($conn->reveal());
+        $this->em->clear()->will(function (): void {
+        });
         $this->logger = $this->prophesize(LoggerInterface::class);
         $this->dbUpdater = $this->prophesize(GeolocationDbUpdaterInterface::class);
         $this->eventDispatcher = $this->prophesize(EventDispatcherInterface::class);