Fix unsetting {main} suspension

revoltphp · Nov 30, 2023 · 25de49a · 25de49a
1 parent fce6063
commit 25de49a
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 5 deletions.
diff --git a/src/EventLoop/Internal/AbstractDriver.php b/src/EventLoop/Internal/AbstractDriver.php
@@ -61,6 +61,7 @@ abstract class AbstractDriver implements Driver
     private bool $idle = false;
     private bool $stopped = false;
 
+    /** @var \WeakMap<object, \WeakReference<DriverSuspension>> */
     private \WeakMap $suspensions;
 
     public function __construct()
@@ -291,8 +292,10 @@ public function getSuspension(): Suspension
         // User callbacks are always executed outside the event loop fiber, so this should always be false.
         \assert($fiber !== $this->fiber);
 
-        // Use current object in case of {main}
-        $suspension = ($this->suspensions[$fiber ?? $this] ?? null)?->get();
+        // Use queue closure in case of {main}, which can be unset by DriverSuspension after an uncaught exception.
+        $key = $fiber ?? $this->queueCallback;
+
+        $suspension = ($this->suspensions[$key] ?? null)?->get();
         if ($suspension) {
             return $suspension;
         }
@@ -304,7 +307,7 @@ public function getSuspension(): Suspension
             $this->suspensions,
         );
 
-        $this->suspensions[$fiber ?? $this] = \WeakReference::create($suspension);
+        $this->suspensions[$key] = \WeakReference::create($suspension);
 
         return $suspension;
     }
@@ -395,7 +398,6 @@ final protected function error(\Closure $closure, \Throwable $exception): void
             $this->interrupt = static fn () => $exception instanceof UncaughtThrowable
                 ? throw $exception
                 : throw UncaughtThrowable::throwingCallback($closure, $exception);
-            unset($this->suspensions[$this]); // Remove suspension for {main}
             return;
         }
 
@@ -629,7 +631,6 @@ private function createErrorCallback(): void
                 $this->interrupt = static fn () => $exception instanceof UncaughtThrowable
                     ? throw $exception
                     : throw UncaughtThrowable::throwingErrorHandler($errorHandler, $exception);
-                unset($this->suspensions[$this]); // Remove suspension for {main}
             }
         };
     }

diff --git a/src/EventLoop/Internal/DriverSuspension.php b/src/EventLoop/Internal/DriverSuspension.php
@@ -27,6 +27,9 @@ final class DriverSuspension implements Suspension
 
     private bool $deadMain = false;
 
+    /**
+     * @param \WeakMap<object, \WeakReference<DriverSuspension>> $suspensions
+     */
     public function __construct(
         private readonly \Closure $run,
         private readonly \Closure $queue,
@@ -118,6 +121,9 @@ public function suspend(): mixed
             // This is now a dead {main} suspension.
             $this->deadMain = true;
 
+            // Unset suspension for {main} using queue closure.
+            unset($this->suspensions[$this->queue]);
+
             $result && $result(); // Unwrap any uncaught exceptions from the event loop
 
             \gc_collect_cycles(); // Collect any circular references before dumping pending suspensions.