Prevent reuse of connection after prepare

Prevents overlap of commands at the cost of a long-held prepare may block commands.

lib/AbstractPool.php

@@ -211,11 +211,12 @@ abstract class AbstractPool implements Pool {
         try {
             /** @var \Amp\Postgres\Statement $statement */
             $statement = yield $connection->prepare($sql);
-        } finally {
+        } catch (\Throwable $exception) {
             $this->push($connection);
+            throw $exception;
         }
 
-        return $statement;
+        return new Internal\PooledStatement($connection, $statement, $this->push);
     }
 
     /**

lib/Internal/PooledStatement.php

@@ -0,0 +1,40 @@
+<?php
+
+namespace Amp\Postgres\Internal;
+
+use Amp\Postgres\Connection;
+use Amp\Postgres\Statement;
+use Amp\Promise;
+
+class PooledStatement implements Statement {
+    /** @var \Amp\Postgres\Connection */
+    private $connection;
+
+    /** @var \Amp\Postgres\Statement */
+    private $statement;
+
+    /** @var callable */
+    private $push;
+
+    /**
+     * @param \Amp\Postgres\Connection $connection
+     * @param \Amp\Postgres\Statement $statement
+     * @param callable $push
+     */
+    public function __construct(Connection $connection, Statement $statement, callable $push) {
+        $this->connection = $connection;
+        $this->statement = $statement;
+        $this->push = $push;
+    }
+
+    public function __destruct() {
+        ($this->push)($this->connection);
+    }
+
+    /**
+     * {@inheritdoc}
+     */
+    public function execute(...$params): Promise {
+        return $this->statement->execute(...$params);
+    }
+}

test/AbstractPoolTest.php

@@ -5,7 +5,7 @@ namespace Amp\Postgres\Test;
 use Amp\Loop;
 use Amp\Postgres\CommandResult;
 use Amp\Postgres\Connection;
-use Amp\Postgres\Statement;
+use Amp\Postgres\Listener;
 use Amp\Postgres\Transaction;
 use Amp\Postgres\TupleResult;
 use Amp\Promise;
@@ -50,8 +50,9 @@ abstract class AbstractPoolTest extends TestCase {
         return [
             [3, 'query', TupleResult::class, "SELECT * FROM test"],
             [2, 'query', CommandResult::class, "INSERT INTO test VALUES (1, 7)"],
-            [1, 'prepare', Statement::class, "SELECT * FROM test WHERE id=\$1"],
+            [5, 'listen', Listener::class, "test"],
             [4, 'execute', TupleResult::class, "SELECT * FROM test WHERE id=\$1 AND time>\$2", 1, time()],
+            [4, 'notify', CommandResult::class, "test", "payload"],
         ];
     }
 
@@ -78,10 +79,9 @@ abstract class AbstractPoolTest extends TestCase {
 
         $pool = $this->createPool($connections);
 
-        Loop::run(function () use ($method, $pool, $params, $result) {
+        Loop::run(function () use ($method, $pool, $params, $result, $resultClass) {
             $return = yield $pool->{$method}(...$params);
-
-            $this->assertSame($result, $return);
+            $this->assertInstanceOf($resultClass, $return);
         });
     }