Add isAlive() method to connection handles

Used in pool to drop dead connections from the pool automatically.

lib/AbstractConnection.php

@@ -37,6 +37,13 @@ abstract class AbstractConnection implements Connection {
         $this->release = $this->callableFromInstanceMethod("release");
     }
 
+    /**
+     * {@inheritdoc}
+     */
+    public function isAlive(): bool {
+        return $this->handle->isAlive();
+    }
+
     /**
      * @param string $methodName Method to execute.
      * @param mixed ...$args Arguments to pass to function.

lib/AbstractPool.php

@@ -83,6 +83,10 @@ abstract class AbstractPool implements Pool {
 
         $this->connections->attach($connection);
         $this->idle->push($connection);
+
+        if ($this->deferred instanceof Deferred) {
+            $this->deferred->resolve($connection);
+        }
     }
 
     /**
@@ -101,12 +105,12 @@ abstract class AbstractPool implements Pool {
             }
         }
 
-        if ($this->idle->isEmpty()) {
+        while ($this->idle->isEmpty()) { // While loop to ensure an idle connection is available after promises below are resolved.
             try {
                 if ($this->connections->count() >= $this->getMaxConnections()) {
                     // All possible connections busy, so wait until one becomes available.
                     $this->deferred = new Deferred;
-                    yield $this->promise = $this->deferred->promise();
+                    yield $this->promise = $this->deferred->promise(); // May be resolved with defunct connection.
                 } else {
                     // Max connection count has not been reached, so open another connection.
                     $this->promise = $this->createConnection();
@@ -128,11 +132,13 @@ abstract class AbstractPool implements Pool {
      * @throws \Error If the connection is not part of this pool.
      */
     private function push(Connection $connection) {
-        if (!isset($this->connections[$connection])) {
-            throw new \Error('Connection is not part of this pool');
-        }
+        \assert(isset($this->connections[$connection]), 'Connection is not part of this pool');
 
-        $this->idle->push($connection);
+        if ($connection->isAlive()) {
+            $this->idle->push($connection);
+        } else {
+            $this->connections->detach($connection);
+        }
 
         if ($this->deferred instanceof Deferred) {
             $this->deferred->resolve($connection);

lib/ConnectionException.php

@@ -0,0 +1,6 @@
+<?php
+
+namespace Amp\Postgres;
+
+class ConnectionException extends FailureException {
+}

lib/Handle.php

@@ -3,6 +3,13 @@
 namespace Amp\Postgres;
 
 interface Handle extends Executor {
+    /**
+     * Indicates if the connection to the database is still alive.
+     *
+     * @return bool
+     */
+    public function isAlive(): bool;
+
     /**
      * Quotes (escapes) the given string for use as a string literal or identifier in a query. This method wraps the
      * string in single quotes, so additional quotes should not be added in the query.

lib/Internal/PooledConnection.php

@@ -28,6 +28,13 @@ class PooledConnection implements Connection {
         ($this->push)($this->connection);
     }
 
+    /**
+     * {@inheritdoc}
+     */
+    public function isAlive(): bool {
+        return $this->connection->isAlive();
+    }
+
     /**
      * {@inheritdoc}
      */

lib/PgSqlHandle.php

@@ -49,14 +49,25 @@ class PgSqlHandle implements Handle {
     public function __construct($handle, $socket) {
         $this->handle = $handle;
 
+        $handle = &$this->handle;
         $deferred = &$this->deferred;
         $listeners = &$this->listeners;
 
-        $this->poll = Loop::onReadable($socket, static function ($watcher) use (&$deferred, &$listeners, $handle) {
+        $this->poll = Loop::onReadable($socket, static function ($watcher) use (&$deferred, &$listeners, &$handle) {
             if (!\pg_consume_input($handle)) {
-                if ($deferred !== null) {
-                    $deferred->fail(new FailureException(\pg_last_error($handle)));
+                $handle = null; // Marks connection as dead.
+                Loop::disable($watcher);
+
+                $exception = new ConnectionException(\pg_last_error($handle));
+
+                foreach ($listeners as $listener) {
+                    $listener->fail($exception);
                 }
+
+                if ($deferred !== null) {
+                    $deferred->fail($exception);
+                }
+
                 return;
             }
 
@@ -123,8 +134,13 @@ class PgSqlHandle implements Handle {
     }
 
     /**
-     * @coroutine
-     *
+     * {@inheritdoc}
+     */
+    public function isAlive(): bool {
+        return $this->handle !== null;
+    }
+
+    /**
      * @param callable $function Function name to execute.
      * @param mixed ...$args Arguments to pass to function.
      *
@@ -144,6 +160,10 @@ class PgSqlHandle implements Handle {
             }
         }
 
+        if (!$this->handle) {
+            throw new ConnectionException("The connection to the database has been lost");
+        }
+
         $result = $function($this->handle, ...$args);
 
         if ($result === false) {

lib/PqHandle.php

@@ -60,13 +60,25 @@ class PqHandle implements Handle {
     public function __construct(pq\Connection $handle) {
         $this->handle = $handle;
 
+        $handle = &$this->handle;
         $deferred = &$this->deferred;
         $listeners = &$this->listeners;
 
-        $this->poll = Loop::onReadable($this->handle->socket, static function ($watcher) use (&$deferred, &$listeners, $handle) {
+        $this->poll = Loop::onReadable($this->handle->socket, static function ($watcher) use (&$deferred, &$listeners, &$handle) {
             if ($handle->poll() === pq\Connection::POLLING_FAILED) {
-                $deferred->fail(new FailureException($handle->errorMessage));
+                $handle = null; // Marks connection as dead.
                 Loop::disable($watcher);
+
+                $exception = new ConnectionException($handle->errorMessage);
+
+                foreach ($listeners as $listener) {
+                    $listener->fail($exception);
+                }
+
+                if ($deferred !== null) {
+                    $deferred->fail($exception);
+                }
+
                 return;
             }
 
@@ -111,6 +123,13 @@ class PqHandle implements Handle {
         Loop::cancel($this->await);
     }
 
+    /**
+     * {@inheritdoc}
+     */
+    public function isAlive(): bool {
+        return $this->handle !== null;
+    }
+
     /**
      * @param callable $method Method to execute.
      * @param mixed ...$args Arguments to pass to function.
@@ -130,6 +149,10 @@ class PqHandle implements Handle {
             }
         }
 
+        if (!$this->handle) {
+            throw new ConnectionException("The connection to the database has been lost");
+        }
+
         try {
             $handle = $method(...$args);
 
@@ -229,6 +252,10 @@ class PqHandle implements Handle {
     }
 
     private function deallocate(string $name) {
+        if (!$this->handle) {
+            return; // Connection dead.
+        }
+
         \assert(isset($this->statements[$name]), "Named statement not found when deallocating");
 
         $storage = $this->statements[$name];

lib/Transaction.php

@@ -55,6 +55,13 @@ class Transaction implements Handle, Operation {
         $this->queue->onComplete($onComplete);
     }
 
+    /**
+     * {@inheritdoc}
+     */
+    public function isAlive(): bool {
+        return $this->handle !== null && $this->handle->isAlive();
+    }
+
     /**
      * @return bool True if the transaction is active, false if it has been committed or rolled back.
      */

test/AbstractConnectionTest.php

@@ -42,6 +42,10 @@ abstract class AbstractConnectionTest extends TestCase {
         $this->connection = $this->createConnection('host=localhost user=postgres');
     }
 
+    public function testIsAlive() {
+        $this->assertTrue($this->connection->isAlive());
+    }
+
     public function testQueryWithTupleResult() {
         Loop::run(function () {
             /** @var \Amp\Postgres\TupleResult $result */
@@ -376,6 +380,7 @@ abstract class AbstractConnectionTest extends TestCase {
 
             $data = $this->getData()[0];
 
+            $this->assertTrue($transaction->isAlive());
             $this->assertTrue($transaction->isActive());
             $this->assertSame($isolation, $transaction->getIsolationLevel());
 
@@ -387,6 +392,7 @@ abstract class AbstractConnectionTest extends TestCase {
 
             yield $transaction->commit();
 
+            $this->assertFalse($transaction->isAlive());
             $this->assertFalse($transaction->isActive());
 
             try {

test/AbstractPoolTest.php

@@ -24,8 +24,12 @@ abstract class AbstractPoolTest extends TestCase {
     /**
      * @return \PHPUnit_Framework_MockObject_MockObject|\Amp\Postgres\Connection
      */
-    private function createConnection() {
-        return $this->createMock(Connection::class);
+    protected function createConnection() {
+        $mock = $this->createMock(Connection::class);
+        $mock->method('isAlive')
+            ->willReturn(true);
+
+        return $mock;
     }
 
     /**

test/AggregatePoolTest.php

@@ -27,4 +27,23 @@ class AggregatePoolTest extends AbstractPoolTest {
 
         return $mock;
     }
+
+    public function testGetMaxConnections() {
+        $pool = $this->createPool([$this->createConnection()]);
+        $this->assertSame(1, $pool->getMaxConnections());
+        $pool->addConnection($this->createConnection());
+        $this->assertSame(2, $pool->getMaxConnections());
+    }
+
+    public function testGetConnectionCount() {
+        $pool = $this->createPool([$this->createConnection(), $this->createConnection()]);
+        $this->assertSame(2, $pool->getConnectionCount());
+    }
+
+    public function testGetIdleConnectionCount() {
+        $pool = $this->createPool([$this->createConnection(), $this->createConnection()]);
+        $this->assertSame(2, $pool->getIdleConnectionCount());
+        $promise = $pool->query("SELECT 1");
+        $this->assertSame(1, $pool->getIdleConnectionCount());
+    }
 }