Fix onClose handler on socket disconnect

amphp · Sep 14, 2023 · f40de82 · f40de82 · bennnjamin · Sep 14, 2023
1 parent 227e285
commit f40de82
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 5 deletions.
diff --git a/src/Internal/Rfc6455FrameHandler.php b/src/Internal/Rfc6455FrameHandler.php
@@ -7,6 +7,7 @@
 use Amp\DeferredFuture;
 use Amp\ForbidCloning;
 use Amp\ForbidSerialization;
+use Amp\Future;
 use Amp\Pipeline\ConcurrentIterator;
 use Amp\Pipeline\DisposedException;
 use Amp\Pipeline\Queue;
@@ -299,6 +300,20 @@ public function close(int $code, string $reason = '', bool $byPeer = false): voi
         $this->socket->close();
     }
 
+    /**
+     * @param \Closure(int, WebsocketCloseInfo):void $onClose
+     */
+    public function onClose(\Closure $onClose): void
+    {
+        $future = $this->closeDeferred?->getFuture() ?? Future::complete();
+
+        $metadata = $this->metadata;
+        $future->finally(static function () use ($onClose, $metadata): void {
+            \assert($metadata->closeInfo !== null, 'Client was not closed when onClose invoked');
+            $onClose($metadata->id, $metadata->closeInfo);
+        });
+    }
+
     private static function createMessage(WebsocketFrameType $frameType, ReadableStream|string $stream): WebsocketMessage
     {
         if ($frameType === WebsocketFrameType::Binary) {

diff --git a/src/Rfc6455Client.php b/src/Rfc6455Client.php
@@ -307,10 +307,6 @@ public function close(int $code = WebsocketCloseCode::NORMAL_CLOSE, string $reas
 
     public function onClose(\Closure $onClose): void
     {
-        $metadata = $this->metadata;
-        $this->socket->onClose(static function () use ($onClose, $metadata): void {
-            \assert($metadata->closeInfo !== null, 'Client was not closed when onClose invoked');
-            $onClose($metadata->id, $metadata->closeInfo);
-        });
+        $this->frameHandler->onClose($onClose);
     }
 }