1 files changed, 41 insertions, 34 deletions
diff --git a/std/ws/example_client.ts b/std/ws/example_client.ts
index d680c6fef..4213025f4 100644
--- a/std/ws/example_client.ts
+++ b/std/ws/example_client.ts
@@ -11,42 +11,49 @@ import { blue, green, red, yellow } from "../fmt/colors.ts";
 
 const endpoint = Deno.args[0] || "ws://127.0.0.1:8080";
 /** simple websocket cli */
-const sock = await connectWebSocket(endpoint);
-console.log(green("ws connected! (type 'close' to quit)"));
-(async function (): Promise<void> {
-  for await (const msg of sock.receive()) {
-    if (typeof msg === "string") {
-      console.log(yellow("< " + msg));
-    } else if (isWebSocketPingEvent(msg)) {
-      console.log(blue("< ping"));
-    } else if (isWebSocketPongEvent(msg)) {
-      console.log(blue("< pong"));
-    } else if (isWebSocketCloseEvent(msg)) {
-      console.log(red(`closed: code=${msg.code}, reason=${msg.reason}`));
+try {
+  const sock = await connectWebSocket(endpoint);
+  console.log(green("ws connected! (type 'close' to quit)"));
+
+  const messages = async (): Promise<void> => {
+    for await (const msg of sock) {
+      if (typeof msg === "string") {
+        console.log(yellow(`< ${msg}`));
+      } else if (isWebSocketPingEvent(msg)) {
+        console.log(blue("< ping"));
+      } else if (isWebSocketPongEvent(msg)) {
+        console.log(blue("< pong"));
+      } else if (isWebSocketCloseEvent(msg)) {
+        console.log(red(`closed: code=${msg.code}, reason=${msg.reason}`));
+      }
     }
-  }
-})();
+  };
 
-const tpr = new TextProtoReader(new BufReader(Deno.stdin));
-while (true) {
-  await Deno.stdout.write(encode("> "));
-  const line = await tpr.readLine();
-  if (line === null) {
-    break;
-  }
-  if (line === "close") {
-    break;
-  } else if (line === "ping") {
-    await sock.ping();
-  } else {
-    await sock.send(line);
+  const cli = async (): Promise<void> => {
+    const tpr = new TextProtoReader(new BufReader(Deno.stdin));
+    while (true) {
+      await Deno.stdout.write(encode("> "));
+      const line = await tpr.readLine();
+      if (line === null) {
+        break;
+      }
+      if (line === "close") {
+        break;
+      } else if (line === "ping") {
+        await sock.ping();
+      } else {
+        await sock.send(line);
+      }
+    }
+  };
+
+  await Promise.race([messages(), cli()]).catch(console.error);
+
+  if (!sock.isClosed) {
+    await sock.close(1000).catch(console.error);
   }
-  // FIXME: Without this,
-  // sock.receive() won't resolved though it is readable...
-  await new Promise((resolve): void => {
-    setTimeout(resolve, 0);
-  });
+} catch (err) {
+  console.error(red(`Could not connect to WebSocket: '${err}'`));
 }
-await sock.close(1000);
-// FIXME: conn.close() won't shutdown process...
+
 Deno.exit(0);