From b4122a1871d46a1a9adeb68e857754d1998b6222 Mon Sep 17 00:00:00 2001
From: Allison Guilhem <allison.guilhem@gmail.com>
Date: Wed, 17 Apr 2024 17:30:02 +1000
Subject: [PATCH] take into account call to db

---
 DependencyInjection/Configuration.php     |  3 +-
 DependencyInjection/DoctrineExtension.php | 34 +++++++--------
 Middleware/ConnectionKeepAlive.php        | 52 +++++++++++++++++++++++
 Resources/config/dbal.xml                 |  3 +-
 Resources/config/middlewares.xml          |  3 ++
 5 files changed, 73 insertions(+), 22 deletions(-)
 create mode 100644 Middleware/ConnectionKeepAlive.php

diff --git a/DependencyInjection/Configuration.php b/DependencyInjection/Configuration.php
index 68e40fc1..65e4259c 100644
--- a/DependencyInjection/Configuration.php
+++ b/DependencyInjection/Configuration.php
@@ -221,10 +221,9 @@ private function getDbalConnectionsNode(): ArrayNodeDefinition
                 ->end()
                 ->booleanNode('disable_type_comments')->end()
                 ->scalarNode('server_version')->end()
-                ->integerNode('check_connection_timing')->end()
+                ->integerNode('check_connection_timing')->defaultValue(30)->end()
                 ->scalarNode('driver_class')->end()
                 ->scalarNode('wrapper_class')->end()
-                ->scalarNode('check_connection_frequency')->defaultValue(30)->end()
                 ->booleanNode('keep_slave')
                     ->setDeprecated(
                         'doctrine/doctrine-bundle',
diff --git a/DependencyInjection/DoctrineExtension.php b/DependencyInjection/DoctrineExtension.php
index f08570e8..75876797 100644
--- a/DependencyInjection/DoctrineExtension.php
+++ b/DependencyInjection/DoctrineExtension.php
@@ -174,12 +174,6 @@ protected function dbalLoad(array $config, ContainerBuilder $container)
             $config['default_connection'] = reset($keys);
         }
 
-        if (! $container->hasParameter('kernel.runtime_mode') && ! $container->hasParameter('kernel.runtime_mode.worker')) {
-            $container->removeDefinition('doctrine.listeners.doctrine_connection_listener');
-        } else {
-            $container->getDefinition('doctrine.orm.listeners.doctrine_connection_listener')->setArgument(1, $config['check_connection_timing']);
-        }
-
         $this->defaultConnection = $config['default_connection'];
 
         $container->setAlias('database_connection', sprintf('doctrine.dbal.%s_connection', $this->defaultConnection));
@@ -203,12 +197,7 @@ protected function dbalLoad(array $config, ContainerBuilder $container)
         $connWithProfiling = [];
         $connWithBacktrace = [];
         $timingByConnection = [];
-        $skipTiming = false;
-
-        if (! $container->hasParameter('kernel.runtime_mode') && ! $container->hasParameter('kernel.runtime_mode.worker')) {
-            $container->removeDefinition('doctrine.listeners.doctrine_connection_listener');
-            $skipTiming = true;
-        }
+        $connWithTimingCheck = [];
 
         foreach ($config['connections'] as $name => $connection) {
             if ($connection['logging']) {
@@ -223,16 +212,16 @@ protected function dbalLoad(array $config, ContainerBuilder $container)
                 }
             }
 
-            if (! $skipTiming) {
-                $timingByConnection[] = ['name' => $connections[$name], 'timing' => $connection['check_connection_timing']];
+            if ($connection['check_connection_timing']) {
+                $connWithTimingCheck[] = $name;
             }
 
+            $timingByConnection[$connections[$name]] = $connection['check_connection_timing'];
+
             $this->loadDbalConnection($name, $connection, $container);
         }
 
-        if (! $skipTiming) {
-            $container->getDefinition('doctrine.orm.listeners.doctrine_connection_listener')->setArgument(1, $timingByConnection);
-        }
+        $container->getDefinition('doctrine.listeners.doctrine_connection_listener')->setArgument(2, $timingByConnection);
 
         $container->registerForAutoconfiguration(MiddlewareInterface::class)->addTag('doctrine.middleware');
 
@@ -250,7 +239,7 @@ protected function dbalLoad(array $config, ContainerBuilder $container)
             }
         });
 
-        $this->registerDbalMiddlewares($container, $connWithLogging, $connWithProfiling, $connWithBacktrace);
+        $this->registerDbalMiddlewares($container, $connWithLogging, $connWithProfiling, $connWithBacktrace, $connWithTimingCheck);
     }
 
     /**
@@ -1201,7 +1190,8 @@ private function registerDbalMiddlewares(
         ContainerBuilder $container,
         array $connWithLogging,
         array $connWithProfiling,
-        array $connWithBacktrace
+        array $connWithBacktrace,
+        array $connWithTimingCheck
     ): void {
         $loader = new XmlFileLoader($container, new FileLocator(__DIR__ . '/../Resources/config'));
         $loader->load('middlewares.xml');
@@ -1217,5 +1207,11 @@ private function registerDbalMiddlewares(
             $debugMiddlewareAbstractDef
                 ->addTag('doctrine.middleware', ['connection' => $connName]);
         }
+
+        $checkTimingMiddlewareAbstractDef = $container->getDefinition('doctrine.connection.keep.alive_middleware');
+        foreach ($connWithTimingCheck as $connName) {
+            $checkTimingMiddlewareAbstractDef
+                ->addTag('doctrine.middleware', ['connection' => $connName, 'priority' => 10]);
+        }
     }
 }
diff --git a/Middleware/ConnectionKeepAlive.php b/Middleware/ConnectionKeepAlive.php
new file mode 100644
index 00000000..f36f0776
--- /dev/null
+++ b/Middleware/ConnectionKeepAlive.php
@@ -0,0 +1,52 @@
+<?php
+
+namespace Doctrine\Bundle\DoctrineBundle\Middleware;
+
+use Doctrine\DBAL\Driver;
+use Doctrine\DBAL\Driver\Connection as DriverConnection;
+use Doctrine\DBAL\Driver\Middleware;
+use Doctrine\DBAL\Driver\Middleware\AbstractDriverMiddleware;
+use Symfony\Bridge\Doctrine\ConnectionTimingWeakMap;
+
+class ConnectionKeepAlive implements Middleware, ConnectionNameAwareInterface
+{
+    private ConnectionTimingWeakMap $connectionTimingWeakMap;
+    private string $connectionName;
+
+    public function __construct(ConnectionTimingWeakMap $connectionTimingWeakMap, string $connectionName = 'default')
+    {
+        $this->connectionName = $connectionName;
+        $this->connectionTimingWeakMap = $connectionTimingWeakMap;
+    }
+
+    public function setConnectionName(string $name): void
+    {
+        $this->connectionName = $name;
+    }
+
+    public function wrap(Driver $driver): Driver
+    {
+        return new class($driver, $this->connectionTimingWeakMap, $this->connectionName) extends AbstractDriverMiddleware {
+            private Driver $driver;
+            private ConnectionTimingWeakMap $connectionTimingWeakMap;
+            private string $connectionName;
+
+            public function __construct(Driver $driver, ConnectionTimingWeakMap $connectionTimingWeakMap, string $connectionName = 'default')
+            {
+                $this->connectionName = $connectionName;
+                $this->connectionTimingWeakMap = $connectionTimingWeakMap;
+                $this->driver = $driver;
+
+                parent::__construct($driver);
+            }
+
+            public function connect(array $params): DriverConnection
+            {
+                $connection = parent::connect($params);
+                $this->connectionTimingWeakMap->setTimingPerConnection($connection, $this->connectionName);
+
+                return $connection;
+            }
+        };
+    }
+}
diff --git a/Resources/config/dbal.xml b/Resources/config/dbal.xml
index 3057d3b3..2c0f6d4e 100644
--- a/Resources/config/dbal.xml
+++ b/Resources/config/dbal.xml
@@ -101,7 +101,8 @@
             <tag name="controller.service_arguments" />
         </service>
 
-        <service id="doctrine.listeners.doctrine_connection_listener" class="Symfony\Bridge\Doctrine\Listener\ConnectionListener">
+        <service id="doctrine.listeners.doctrine_connection_listener" class="Symfony\Bridge\Doctrine\ConnectionListener">
+            <argument type="service" id="connection.weak_map" />
             <argument type="service" id="service_container" />
             <argument /> <!-- check timing -->
             <tag name="kernel.event_subscriber" />
diff --git a/Resources/config/middlewares.xml b/Resources/config/middlewares.xml
index d6bf92bc..4a92e52d 100644
--- a/Resources/config/middlewares.xml
+++ b/Resources/config/middlewares.xml
@@ -17,5 +17,8 @@
             <argument type="service" id="doctrine.debug_data_holder" />
             <argument type="service" id="debug.stopwatch" on-invalid="null" />
         </service>
+        <service id="doctrine.connection.keep.alive_middleware" class="Doctrine\Bundle\DoctrineBundle\Middleware\ConnectionKeepAlive" abstract="true">
+            <argument type="service" id="connection.weak_map" />
+        </service>
     </services>
 </container>