From fc382dea1e4b7edf8c5c8d71231e55caf81454a1 Mon Sep 17 00:00:00 2001
From: Allison Guilhem <allison.guilhem@gmail.com>
Date: Thu, 4 Apr 2024 23:45:51 +1100
Subject: [PATCH] In long-running processes, allows checking the database
 connection by pinging it, thus avoiding exceptions. Additionally, it provides
 the flexibility to configure the frequency of the check to be performed in
 the case of long-running processes.

---
 src/Configuration.php                         | 12 ++++++
 src/Connection.php                            | 38 ++++++++++++++++-
 .../Connection/ConnectionReactivatedTest.php  | 41 +++++++++++++++++++
 tests/FunctionalTestCase.php                  |  9 +++-
 tests/TestUtil.php                            |  6 ++-
 5 files changed, 103 insertions(+), 3 deletions(-)
 create mode 100644 tests/Functional/Connection/ConnectionReactivatedTest.php

diff --git a/src/Configuration.php b/src/Configuration.php
index 9aa001db483..6f7a9214836 100644
--- a/src/Configuration.php
+++ b/src/Configuration.php
@@ -36,6 +36,8 @@ class Configuration
 
     private ?SchemaManagerFactory $schemaManagerFactory = null;
 
+    private ?float $checkConnectionTiming = null;
+
     public function __construct()
     {
         $this->schemaAssetsFilter = static function (): bool {
@@ -153,4 +155,14 @@ public function setDisableTypeComments(bool $disableTypeComments): self
 
         return $this;
     }
+
+    public function setCheckConnectionTiming(float $timing): void
+    {
+        $this->checkConnectionTiming = $timing;
+    }
+
+    public function getCheckConnectionTiming(): ?float
+    {
+        return $this->checkConnectionTiming;
+    }
 }
diff --git a/src/Connection.php b/src/Connection.php
index 3bb951445ff..8c8025c3291 100644
--- a/src/Connection.php
+++ b/src/Connection.php
@@ -42,6 +42,7 @@
 use function is_string;
 use function key;
 use function sprintf;
+use function time;
 
 /**
  * A database abstraction-level connection that implements features like transaction isolation levels,
@@ -99,6 +100,12 @@ class Connection implements ServerVersionProvider
 
     private SchemaManagerFactory $schemaManagerFactory;
 
+    private bool $isChecking = false;
+
+    private ?int $lastCheckedAt = null;
+
+    private ?int $heartbeat;
+
     /**
      * Initializes a new instance of the Connection class.
      *
@@ -119,6 +126,8 @@ public function __construct(
         $this->params     = $params;
         $this->autoCommit = $this->_config->getAutoCommit();
 
+        $this->heartbeat = $this->_config->getCheckConnectionTiming();
+
         $this->schemaManagerFactory = $this->_config->getSchemaManagerFactory()
             ?? new DefaultSchemaManagerFactory();
     }
@@ -210,7 +219,23 @@ public function createExpressionBuilder(): ExpressionBuilder
     protected function connect(): DriverConnection
     {
         if ($this->_conn !== null) {
-            return $this->_conn;
+            $isTimeToCheck = $this->lastCheckedAt === null || time() - $this->lastCheckedAt >= $this->heartbeat;
+            $noCheckNeeded = $this->heartbeat === null || $this->isChecking;
+
+            if ($noCheckNeeded || ! $isTimeToCheck) {
+                return $this->_conn;
+            }
+
+            $this->isChecking = true;
+
+            $isAvailable = $this->reconnectOnFailure();
+
+            $this->lastCheckedAt = time();
+            $this->isChecking    = false;
+
+            if ($isAvailable) {
+                return $this->_conn;
+            }
         }
 
         try {
@@ -1371,4 +1396,15 @@ private function handleDriverException(
 
         return $exception;
     }
+
+    private function reconnectOnFailure(): bool
+    {
+        try {
+            $this->executeQuery($this->getDatabasePlatform($this)->getDummySelectSQL());
+
+            return true;
+        } catch (ConnectionLost) {
+            return false;
+        }
+    }
 }
diff --git a/tests/Functional/Connection/ConnectionReactivatedTest.php b/tests/Functional/Connection/ConnectionReactivatedTest.php
new file mode 100644
index 00000000000..88cdca94c46
--- /dev/null
+++ b/tests/Functional/Connection/ConnectionReactivatedTest.php
@@ -0,0 +1,41 @@
+<?php
+
+declare(strict_types=1);
+
+namespace Doctrine\DBAL\Tests\Functional\Connection;
+
+use Doctrine\DBAL\Platforms\AbstractMySQLPlatform;
+use Doctrine\DBAL\Tests\FunctionalTestCase;
+
+use function sleep;
+
+class ConnectionReactivatedTest extends FunctionalTestCase
+{
+    public static function setUpBeforeClass(): void
+    {
+        self::markConnectionWithHeartBeat();
+    }
+
+    protected function setUp(): void
+    {
+        if ($this->connection->getDatabasePlatform() instanceof AbstractMySQLPlatform) {
+            return;
+        }
+
+        self::markTestSkipped('Currently only supported with MySQL');
+    }
+
+    public function testConnectionReactivated(): void
+    {
+        $this->connection->executeStatement('SET SESSION wait_timeout=1');
+
+        sleep(2);
+
+        $query = $this->connection->getDatabasePlatform()
+            ->getDummySelectSQL();
+
+        $this->connection->executeQuery($query);
+
+        self::assertEquals(1, $this->connection->fetchOne($query));
+    }
+}
diff --git a/tests/FunctionalTestCase.php b/tests/FunctionalTestCase.php
index 29428454ccf..2d96a6712b1 100644
--- a/tests/FunctionalTestCase.php
+++ b/tests/FunctionalTestCase.php
@@ -26,6 +26,8 @@ abstract class FunctionalTestCase extends TestCase
      */
     private bool $isConnectionReusable = true;
 
+    protected static bool $hasHeartBeat = false;
+
     /**
      * Mark shared connection not reusable for subsequent tests.
      *
@@ -37,11 +39,16 @@ protected function markConnectionNotReusable(): void
         $this->isConnectionReusable = false;
     }
 
+    protected static function markConnectionWithHeartBeat(): void
+    {
+        self::$hasHeartBeat = true;
+    }
+
     #[Before]
     final protected function connect(): void
     {
         if (self::$sharedConnection === null) {
-            self::$sharedConnection = TestUtil::getConnection();
+            self::$sharedConnection = TestUtil::getConnection(self::$hasHeartBeat);
         }
 
         $this->connection = self::$sharedConnection;
diff --git a/tests/TestUtil.php b/tests/TestUtil.php
index 354ce38ef87..6c50f333eb0 100644
--- a/tests/TestUtil.php
+++ b/tests/TestUtil.php
@@ -62,10 +62,14 @@ class TestUtil
      *
      * @return Connection The database connection instance.
      */
-    public static function getConnection(): Connection
+    public static function getConnection($hasHeartBeat = false): Connection
     {
         $params = self::getConnectionParams();
 
+        if ($hasHeartBeat) {
+            $params['check_connection_frequency'] = 1;
+        }
+
         if (empty($params['memory']) && ! self::$initialized) {
             self::initializeDatabase();
             self::$initialized = true;