feat: Simplify acquiring and resuing connections

Author: pratikshazalte69

src/Bolt/ConnectionPool.php

@@ -24,6 +24,7 @@
 use Laudis\Neo4j\Databags\DriverConfiguration;
 use Laudis\Neo4j\Databags\SessionConfiguration;
 
+use Laudis\Neo4j\Exception\ConnectionPoolException;
 use function method_exists;
 use function microtime;
 
@@ -44,6 +45,7 @@ public function __construct(
         private readonly BoltFactory $factory,
         private readonly ConnectionRequestData $data,
         private readonly ?Neo4jLogger $logger,
+        private readonly float $acquireConnectionTimeout
     ) {
     }
 
@@ -63,33 +65,39 @@ public static function create(
                 $conf->getUserAgent(),
                 $conf->getSslConfiguration()
             ),
-            $conf->getLogger()
+            $conf->getLogger(),
+            $conf->getAcquireConnectionTimeout()
         );
     }
 
     public function acquire(SessionConfiguration $config): Generator
     {
-        $generator = $this->semaphore->wait();
-        $start = microtime(true);
+        return (function () use ($config) {
+            $connection = $this->reuseConnectionIfPossible($config);
+            if ($connection !== null) {
+                return $connection;
+            }
 
-        return (function () use ($generator, $start, $config) {
+            $generator = $this->semaphore->wait();
             // If the generator is valid, it means we are waiting to acquire a new connection.
             // This means we can use this time to check if we can reuse a connection or should throw a timeout exception.
             while ($generator->valid()) {
-                /** @var bool $continue */
-                $continue = yield microtime(true) - $start;
-                $generator->send($continue);
-                if ($continue === false) {
-                    return null;
-                }
+                $waitTime = $generator->current();
+                if ($waitTime <= $this->acquireConnectionTimeout) {
+                    yield $waitTime;
+
+                    $connection = $this->reuseConnectionIfPossible($config);
+                    if ($connection !== null) {
+                        return $connection;
+                    }
 
-                $connection = $this->returnAnyAvailableConnection($config);
-                if ($connection !== null) {
-                    return $connection;
+                    $generator->next();
+                } else {
+                    throw new ConnectionPoolException('Connection acquire timeout reached: ' . $waitTime);
                 }
             }
 
-            $connection = $this->returnAnyAvailableConnection($config);
+            $connection = $this->reuseConnectionIfPossible($config);
             if ($connection !== null) {
                 return $connection;
             }
@@ -119,53 +127,15 @@ public function getLogger(): ?Neo4jLogger
         return $this->logger;
     }
 
-    /**
-     * @return BoltConnection|null
-     */
-    private function returnAnyAvailableConnection(SessionConfiguration $config): ?ConnectionInterface
+    private function reuseConnectionIfPossible(SessionConfiguration $config): ?ConnectionInterface
     {
-        $streamingConnection = null;
-        $requiresReconnectConnection = null;
         // Ensure random connection reuse before picking one.
         shuffle($this->activeConnections);
-
         foreach ($this->activeConnections as $activeConnection) {
             // We prefer a connection that is just ready
-            if ($activeConnection->getServerState() === 'READY') {
-                if ($this->factory->canReuseConnection($activeConnection, $this->data, $config)) {
-                    return $this->factory->reuseConnection($activeConnection, $config);
-                } else {
-                    $requiresReconnectConnection = $activeConnection;
-                }
+            if ($activeConnection->getServerState() === 'READY' && $this->factory->canReuseConnection($activeConnection, $config)) {
+                return $this->factory->reuseConnection($activeConnection, $config);
             }
-
-            // We will store any streaming connections, so we can use that one
-            // as we can force the subscribed result sets to consume the results
-            // and become ready again.
-            // This code will make sure we never get stuck if the user has many
-            // results open that aren't consumed yet.
-            // https://github.com/neo4j-php/neo4j-php-client/issues/146
-            // NOTE: we cannot work with TX_STREAMING as we cannot force the transaction to implicitly close.
-            if ($streamingConnection === null && $activeConnection->getServerState() === 'STREAMING') {
-                if ($this->factory->canReuseConnection($activeConnection, $this->data, $config)) {
-                    $streamingConnection = $activeConnection;
-                    if (method_exists($streamingConnection, 'consumeResults')) {
-                        $streamingConnection->consumeResults(); // State should now be ready
-                    }
-                } else {
-                    $requiresReconnectConnection = $activeConnection;
-                }
-            }
-        }
-
-        if ($streamingConnection) {
-            return $this->factory->reuseConnection($streamingConnection, $config);
-        }
-
-        if ($requiresReconnectConnection) {
-            $this->release($requiresReconnectConnection);
-
-            return $this->factory->createConnection($this->data, $config);
         }
 
         return null;

src/Bolt/Messages/BoltCommitMessage.php

@@ -40,6 +40,7 @@ public function send(): BoltCommitMessage
     {
         $this->logger?->log(LogLevel::DEBUG, 'COMMIT');
         $response = $this->protocol->commit()->getResponse();
+        $this->protocol->serverState = ServerState::READY;
 
         /** @var array{bookmark?: string} $content */
         $content = $response->content;

src/BoltFactory.php

@@ -79,7 +79,7 @@ public function createConnection(ConnectionRequestData $data, SessionConfigurati
         return new BoltConnection($protocol, $connection, $data->getAuth(), $data->getUserAgent(), $config, $this->logger);
     }
 
-    public function canReuseConnection(ConnectionInterface $connection, ConnectionRequestData $data, SessionConfiguration $config): bool
+    public function canReuseConnection(ConnectionInterface $connection, SessionConfiguration $config): bool
     {
         if (!$connection->isOpen()) {
             return false;
@@ -88,12 +88,7 @@ public function canReuseConnection(ConnectionInterface $connection, ConnectionRe
         $databaseInfo = $connection->getDatabaseInfo();
         $database = $databaseInfo?->getName();
 
-        return $connection->getServerAddress()->getHost() === $data->getUri()->getHost()
-               && $connection->getServerAddress()->getPort() === $data->getUri()->getPort()
-               && $connection->getAuthentication()->toString($data->getUri()) === $data->getAuth()->toString($data->getUri())
-               && $connection->getEncryptionLevel() === $this->sslConfigurationFactory->create($data->getUri(), $data->getSslConfig())[0]
-               && $connection->getUserAgent() === $data->getUserAgent()
-            && $connection->getAccessMode() === $config->getAccessMode()
+        return $connection->getAccessMode() === $config->getAccessMode()
                && $database === $config->getDatabase();
     }
 

src/Exception/ConnectionPoolException.php

@@ -0,0 +1,8 @@
+<?php
+
+namespace Laudis\Neo4j\Exception;
+
+final class ConnectionPoolException extends \RuntimeException
+{
+
+}

src/Neo4j/Neo4jConnectionPool.php

@@ -74,6 +74,7 @@ public function __construct(
         private readonly CacheInterface $cache,
         private readonly AddressResolverInterface $resolver,
         private readonly ?Neo4jLogger $logger,
+        private readonly float $acquireConnectionTimeout,
     ) {
     }
 
@@ -96,7 +97,8 @@ public static function create(
             ),
             Cache::getInstance(),
             $resolver,
-            $conf->getLogger()
+            $conf->getLogger(),
+            $conf->getAcquireConnectionTimeout()
         );
     }
 
@@ -112,7 +114,7 @@ public function createOrGetPool(string $hostname, UriInterface $uri): Connection
 
         $key = $this->createKey($data);
         if (!array_key_exists($key, self::$pools)) {
-            self::$pools[$key] = new ConnectionPool($this->semaphore, $this->factory, $data, $this->logger);
+            self::$pools[$key] = new ConnectionPool($this->semaphore, $this->factory, $data, $this->logger, $this->acquireConnectionTimeout);
         }
 
         return self::$pools[$key];

tests/Integration/EdgeCasesTest.php

@@ -102,18 +102,17 @@ public function testRunALotOfStatements(): void
             }
         }
 
-        $statements = [];
+        $count = 0;
         foreach ($personIds as $personId) {
             foreach ($movieIds as $movieId) {
-                $statements[] = Statement::create(
+                $count++;
+                $this->getSession()->runStatement(Statement::create(
                     'MATCH (a), (b) WHERE id(a) = $ida AND id(b) = $idb MERGE (a) <-[r:ACTED_IN]- (b) RETURN id(r)',
                     ['ida' => $personId, 'idb' => $movieId]
-                );
+                ));
             }
         }
-
-        $this->getSession()->runStatements($statements);
-        self::assertCount(4978, $statements);
+        self::assertEquals(4978, $count);
     }
 
     public function testGettingKeysFromArraylist(): void

tests/Unit/BoltConnectionPoolTest.php

@@ -171,6 +171,7 @@ private function setupPool(Generator $semaphoreGenerator): void
                 SslConfiguration::default()
             ),
             null,
+            10.0
         );
     }
 }