Check for next result before releasing connection

Author: trowski

Diff for: src/PooledResultSet.php

@@ -10,19 +10,75 @@ final class PooledResultSet extends SqlPooledResultSet implements ResultSet
     /** @var ConnectionResultSet */
     private $result;
 
+    /** @var callable */
+    private $release;
+
+    /** @var Promise|null a*/
+    private $nextResultPromise;
+
     /**
      * @param ResultSet $result
-     * @param callable            $release
+     * @param callable  $release
      */
     public function __construct(ResultSet $result, callable $release)
     {
-        parent::__construct($result, $release);
         $this->result = $result;
+        $this->release = &$release;
+
+        parent::__construct($this->result, static function () use (&$release) {
+            if ($release !== null) {
+                $release();
+            }
+        });
+    }
+
+    public function advance(): Promise
+    {
+        $promise = $this->result->advance();
+
+        $promise->onResolve(function (\Throwable $exception = null, bool $moreResults = null) {
+            if ($moreResults || $this->release === null) {
+                return;
+            }
+
+            $this->nextResultPromise = $this->result->nextResultSet();
+            $this->nextResultPromise->onResolve(function (\Throwable $exception = null, bool $moreResults = null) {
+                $this->nextResultPromise = null;
+
+                if ($moreResults || $this->release === null) {
+                    return;
+                }
+
+                $release = $this->release;
+                $this->release = null;
+                $release();
+            });
+        });
+
+        return $promise;
     }
 
     public function nextResultSet(): Promise
     {
-        return $this->result->nextResultSet();
+        if ($this->nextResultPromise !== null) {
+            $nextResultPromise = $this->nextResultPromise;
+            $this->nextResultPromise = null;
+            return $nextResultPromise;
+        }
+
+        $promise = $this->result->nextResultSet();
+
+        $promise->onResolve(function (\Throwable $exception = null, bool $moreResults = null) {
+            if ($moreResults || $this->release === null) {
+                return;
+            }
+
+            $release = $this->release;
+            $this->release = null;
+            $release();
+        });
+
+        return $promise;
     }
 
     public function getFields(): Promise

Diff for: test/LinkTest.php

@@ -93,13 +93,14 @@ public function testMultiStmt()
             $this->assertSame([[6]], $got);
             $this->assertTrue(yield $resultset->nextResultSet());
 
+            $fields = yield $resultset->getFields();
+
             $got = [];
             while (yield $resultset->advance()) {
                 $got[] = $resultset->getCurrent();
             }
             $this->assertSame([["d" => 5, "c" => 5], ["d" => 6, "c" => 6]], $got);
 
-            $fields = yield $resultset->getFields();
             $this->assertCount(2, $fields);
             $this->assertSame($fields[0]["original_name"], "b");
             $this->assertSame($fields[0]["name"], "d");