Fix simultaneous accept cancel and server close

amphp · Mar 20, 2024 · 051211a · 051211a
1 parent efd71b3
commit 051211a
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 14 deletions.
diff --git a/src/Ipc/SocketIpcHub.php b/src/Ipc/SocketIpcHub.php
@@ -4,7 +4,7 @@
 
 use Amp\Cache\LocalCache;
 use Amp\Cancellation;
-use Amp\CancelledException;
+use Amp\DeferredFuture;
 use Amp\ForbidCloning;
 use Amp\ForbidSerialization;
 use Amp\NullCancellation;
@@ -25,14 +25,15 @@ final class SocketIpcHub implements IpcHub
     /** @var non-empty-string */
     private readonly string $uri;
 
-    /** @var array<string, EventLoop\Suspension> */
+    /** @var array<string, DeferredFuture> */
     private array $waitingByKey = [];
 
     /** @var \Closure(): void */
     private readonly \Closure $accept;
 
     private bool $queued = false;
 
+    /** @var LocalCache<ResourceSocket> */
     private LocalCache $clientsByKey;
 
     /**
@@ -68,8 +69,8 @@ public function __construct(
                 if (!$client) {
                     $queued = false;
                     $exception = new Socket\SocketException('IPC socket closed before the client connected');
-                    foreach ($waitingByKey as $suspension) {
-                        $suspension->throw($exception);
+                    foreach ($waitingByKey as $deferred) {
+                        $deferred->error($exception);
                     }
                     return;
                 }
@@ -82,7 +83,7 @@ public function __construct(
                 }
 
                 if (isset($waitingByKey[$received])) {
-                    $waitingByKey[$received]->resume($client);
+                    $waitingByKey[$received]->complete($client);
                     unset($waitingByKey[$received]);
                 } else {
                     $clientsByKey->set($received, $client);
@@ -112,8 +113,8 @@ public function close(): void
         }
 
         $exception = new Socket\SocketException('IPC socket closed before the client connected');
-        foreach ($this->waitingByKey as $suspension) {
-            $suspension->throw($exception);
+        foreach ($this->waitingByKey as $deferred) {
+            $deferred->error($exception);
         }
     }
 
@@ -165,19 +166,19 @@ public function accept(string $key, ?Cancellation $cancellation = null): Resourc
             $this->queued = true;
         }
 
-        $cancellation = $cancellation ?? new NullCancellation();
-        $cancellation->throwIfRequested();
+        $cancellation ??= new NullCancellation();
 
-        $this->waitingByKey[$key] = $suspension = EventLoop::getSuspension();
-        $cancellationId = $cancellation->subscribe(function (CancelledException $exception) use ($suspension) {
-            $suspension->throw($exception);
+        $this->waitingByKey[$key] = $deferred = new DeferredFuture();
+
+        $waitingByKey = &$this->waitingByKey;
+        $cancellationId = $cancellation->subscribe(static function () use (&$waitingByKey, $key): void {
+            unset($waitingByKey[$key]);
         });
 
         try {
-            $client = $suspension->suspend();
+            $client = $deferred->getFuture()->await($cancellation);
         } finally {
             $cancellation->unsubscribe($cancellationId);
-            unset($this->waitingByKey[$key]);
         }
 
         return $client;

diff --git a/test/Ipc/SocketIpcHubTest.php b/test/Ipc/SocketIpcHubTest.php
@@ -50,4 +50,23 @@ public function testAcceptAfterCancel(): void
 
         self::assertSame($this->server->getAddress()->toString(), $client->getLocalAddress()->toString());
     }
+
+    public function testCancelledAcceptAndCloseServer(): void
+    {
+        $deferredCancellation = new DeferredCancellation();
+        $future = async(fn () => $this->ipcHub->accept(
+            $this->ipcHub->generateKey(),
+            $deferredCancellation->getCancellation(),
+        ));
+
+        async(function () use ($deferredCancellation): void {
+            $this->server->close();
+            $deferredCancellation->cancel();
+        });
+
+        $this->expectException(Socket\SocketException::class);
+        $this->expectExceptionMessage('closed before');
+
+        $future->await();
+    }
 }