Fix #105: Quickly unload the client when closed by remote

amphp · Dec 22, 2016 · de8bc2f · de8bc2f
1 parent 5066f71
commit de8bc2f
Showing 1 changed file with 10 additions and 5 deletions.
diff --git a/lib/Websocket/Rfc6455Endpoint.php b/lib/Websocket/Rfc6455Endpoint.php
@@ -276,7 +276,7 @@ private function doClose(Rfc6455Client $client, int $code, string $reason): \Gen
         $promise = $this->sendCloseFrame($client, $code, $reason);
         yield from $this->tryAppOnClose($client->id, $code, $reason);
         return $promise;
-        // Don't unload the client here, it will be unloaded upon timeout
+        // Don't unload the client here, it will be unloaded upon timeout or last data written if closed by client or OP_CLOSE received by client
     }
 
     private function sendCloseFrame(Rfc6455Client $client, $code, $msg): Promise {
@@ -507,11 +507,16 @@ public function onWritable($watcherId, $socket, Rfc6455Client $client) {
                 }
 
             } elseif ($client->closedAt) {
-                @stream_socket_shutdown($socket, STREAM_SHUT_WR);
-                \Amp\cancel($watcherId);
-                $client->writeWatcher = null;
-                $client->writeDeferred = null;
                 $client->writeBuffer = "";
+                $client->writeDeferred = null;
+                if ($client->readWatcher) { // readWatcher exists == we are still waiting for an OP_CLOSE frame
+                    @stream_socket_shutdown($socket, STREAM_SHUT_WR);
+                    \Amp\cancel($watcherId);
+                    $client->writeWatcher = null;
+                } else {
+                    unset($this->closeTimeouts[$client->id]);
+                    $this->unloadClient($client);
+                }
             } elseif ($client->writeDataQueue) {
                 $key = key($client->writeDataQueue);
                 $client->writeBuffer = $client->writeDataQueue[$key];