From b018d0ff141636d69a8054194675fe0050cf3aa9 Mon Sep 17 00:00:00 2001
From: Aaron Piotrowski <aaron@trowski.com>
Date: Thu, 20 Jul 2023 15:16:11 -0500
Subject: [PATCH] Fix Fiber leaks if Suspension never resumed, v2 (#82)

---
 .php-cs-fixer.dist.php                      |  2 +-
 src/EventLoop/Internal/AbstractDriver.php   | 13 ++++++++++--
 src/EventLoop/Internal/DriverSuspension.php | 23 ++++++---------------
 test/EventLoopTest.php                      |  3 +--
 4 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/.php-cs-fixer.dist.php b/.php-cs-fixer.dist.php
index a0c8dc4..81766c9 100644
--- a/.php-cs-fixer.dist.php
+++ b/.php-cs-fixer.dist.php
@@ -28,6 +28,7 @@ public function getRules(): array
                 "allow_single_line_closure" => true,
             ],
             "array_syntax" => ["syntax" => "short"],
+            "blank_lines_before_namespace" => true,
             "cast_spaces" => true,
             "combine_consecutive_unsets" => true,
             "declare_strict_types" => true,
@@ -73,7 +74,6 @@ public function getRules(): array
             "psr_autoloading" => ['dir' => $this->src],
             "return_type_declaration" => ["space_before" => "none"],
             "short_scalar_cast" => true,
-            "single_blank_line_before_namespace" => true,
             "line_ending" => true,
         ];
     }
diff --git a/src/EventLoop/Internal/AbstractDriver.php b/src/EventLoop/Internal/AbstractDriver.php
index 78578bd..d04483c 100644
--- a/src/EventLoop/Internal/AbstractDriver.php
+++ b/src/EventLoop/Internal/AbstractDriver.php
@@ -285,12 +285,21 @@ public function getSuspension(): Suspension
         \assert($fiber !== $this->fiber);
 
         // Use current object in case of {main}
-        return $this->suspensions[$fiber ?? $this] ??= new DriverSuspension(
+        $suspension = ($this->suspensions[$fiber ?? $this] ?? null)?->get();
+        if ($suspension) {
+            return $suspension;
+        }
+
+        $suspension = new DriverSuspension(
             $this->runCallback,
             $this->queueCallback,
             $this->interruptCallback,
-            $this->suspensions
+            $this->suspensions,
         );
+
+        $this->suspensions[$fiber ?? $this] = \WeakReference::create($suspension);
+
+        return $suspension;
     }
 
     public function setErrorHandler(?\Closure $errorHandler): void
diff --git a/src/EventLoop/Internal/DriverSuspension.php b/src/EventLoop/Internal/DriverSuspension.php
index 5f7c968..a678c08 100644
--- a/src/EventLoop/Internal/DriverSuspension.php
+++ b/src/EventLoop/Internal/DriverSuspension.php
@@ -23,25 +23,15 @@ final class DriverSuspension implements Suspension
 
     private bool $pending = false;
 
-    private readonly \WeakReference $suspensions;
-
-    /**
-     * @param \Closure $run
-     * @param \Closure $queue
-     * @param \Closure $interrupt
-     *
-     * @internal
-     */
     public function __construct(
         private readonly \Closure $run,
         private readonly \Closure $queue,
         private readonly \Closure $interrupt,
-        \WeakMap $suspensions
+        private readonly \WeakMap $suspensions,
     ) {
         $fiber = \Fiber::getCurrent();
 
         $this->fiberRef = $fiber ? \WeakReference::create($fiber) : null;
-        $this->suspensions = \WeakReference::create($suspensions);
     }
 
     public function resume(mixed $value = null): void
@@ -101,13 +91,12 @@ public function suspend(): mixed
             $this->pending = false;
             $result && $result(); // Unwrap any uncaught exceptions from the event loop
 
-            $info = '';
-            $suspensions = $this->suspensions->get();
-            if ($suspensions) {
-                \gc_collect_cycles();
+            \gc_collect_cycles(); // Collect any circular references before dumping pending suspensions.
 
-                /** @var self $suspension */
-                foreach ($suspensions as $suspension) {
+            $info = '';
+            foreach ($this->suspensions as $suspensionRef) {
+                if ($suspension = $suspensionRef->get()) {
+                    \assert($suspension instanceof self);
                     $fiber = $suspension->fiberRef?->get();
                     if ($fiber === null) {
                         continue;
diff --git a/test/EventLoopTest.php b/test/EventLoopTest.php
index 9420b9f..0fb1ab2 100644
--- a/test/EventLoopTest.php
+++ b/test/EventLoopTest.php
@@ -269,8 +269,7 @@ public function testSuspensionWithinCallbackGarbageCollectionSuspended(): void
 
         \gc_collect_cycles();
 
-        // This documents an expected failure, should actually be true, but suspensions have to be resumed currently.
-        self::assertNull($finally);
+        self::assertTrue($finally);
     }
 
     public function testSuspensionWithinQueue(): void