Fix onClose handler on socket disconnect

Author: trowski

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) {

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);
     }
 }