refactor(op_crates/websocket): refactor event loop (#9079)

2 files changed, 70 insertions, 49 deletions

diff --git a/op_crates/websocket/01_websocket.js b/op_crates/websocket/01_websocket.js
index 6cc23eb95..741e643e5 100644
--- a/op_crates/websocket/01_websocket.js
+++ b/op_crates/websocket/01_websocket.js
@@ -294,59 +294,78 @@
     }
 
     async #eventLoop() {
-      if (this.#readyState === OPEN) {
+      while (this.#readyState === OPEN) {
         const message = await core.jsonOpAsync(
           "op_ws_next_event",
           { rid: this.#rid },
         );
-        if (message.type === "string" || message.type === "binary") {
-          let data;
 
-          if (message.type === "string") {
-            data = message.data;
-          } else {
+        switch (message.kind) {
+          case "string": {
+            const event = new MessageEvent("message", {
+              data: message.data,
+              origin: this.#url,
+            });
+            event.target = this;
+            this.dispatchEvent(event);
+
+            break;
+          }
+
+          case "binary": {
+            let data;
+
             if (this.binaryType === "blob") {
               data = new Blob([new Uint8Array(message.data)]);
             } else {
               data = new Uint8Array(message.data).buffer;
             }
+
+            const event = new MessageEvent("message", {
+              data,
+              origin: this.#url,
+            });
+            event.target = this;
+            this.dispatchEvent(event);
+
+            break;
           }
 
-          const event = new MessageEvent("message", {
-            data,
-            origin: this.#url,
-          });
-          event.target = this;
-          this.dispatchEvent(event);
+          case "ping":
+            core.jsonOpAsync("op_ws_send", {
+              rid: this.#rid,
+              kind: "pong",
+            });
 
-          this.#eventLoop();
-        } else if (message.type === "ping") {
-          core.jsonOpAsync("op_ws_send", {
-            rid: this.#rid,
-            kind: "pong",
-          });
+            break;
 
-          this.#eventLoop();
-        } else if (message.type === "close") {
-          this.#readyState = CLOSED;
-          const event = new CloseEvent("close", {
-            wasClean: true,
-            code: message.code,
-            reason: message.reason,
-          });
-          event.target = this;
-          this.dispatchEvent(event);
-        } else if (message.type === "error") {
-          this.#readyState = CLOSED;
+          case "close": {
+            this.#readyState = CLOSED;
+
+            const event = new CloseEvent("close", {
+              wasClean: true,
+              code: message.data.code,
+              reason: message.data.reason,
+            });
+            event.target = this;
+            this.dispatchEvent(event);
 
-          const errorEv = new ErrorEvent("error");
-          errorEv.target = this;
-          this.dispatchEvent(errorEv);
+            break;
+          }
 
-          this.#readyState = CLOSED;
-          const closeEv = new CloseEvent("close");
-          closeEv.target = this;
-          this.dispatchEvent(closeEv);
+          case "error": {
+            this.#readyState = CLOSED;
+
+            const errorEv = new ErrorEvent("error");
+            errorEv.target = this;
+            this.dispatchEvent(errorEv);
+
+            const closeEv = new CloseEvent("close");
+            closeEv.target = this;
+            this.dispatchEvent(closeEv);
+
+            break;
+          }
         }
       }
     }
diff --git a/op_crates/websocket/lib.rs b/op_crates/websocket/lib.rs
index b688fe9fd..e15bab827 100644
--- a/op_crates/websocket/lib.rs
+++ b/op_crates/websocket/lib.rs
@@ -88,7 +88,7 @@ struct CheckPermissionArgs {
 }
 
 // This op is needed because creating a WS instance in JavaScript is a sync
-// operation and should throw error when permissions are not fullfiled,
+// operation and should throw error when permissions are not fulfilled,
 // but actual op that connects WS is async.
 pub fn op_ws_check_permission<WP>(
   state: &mut OpState,
@@ -155,7 +155,7 @@ where
   let try_socket = TcpStream::connect(addr).await;
   let tcp_socket = match try_socket.map_err(TungsteniteError::Io) {
     Ok(socket) => socket,
-    Err(_) => return Ok(json!({"success": false})),
+    Err(_) => return Ok(json!({ "success": false })),
   };
 
   let socket: MaybeTlsStream = match uri.scheme_str() {
@@ -305,28 +305,30 @@ pub async fn op_ws_next_event(
   let val = rx.next().or_cancel(cancel).await?;
   let res = match val {
     Some(Ok(Message::Text(text))) => json!({
-      "type": "string",
+      "kind": "string",
       "data": text
     }),
     Some(Ok(Message::Binary(data))) => {
       // TODO(ry): don't use json to send binary data.
       json!({
-        "type": "binary",
+        "kind": "binary",
         "data": data
       })
     }
     Some(Ok(Message::Close(Some(frame)))) => json!({
-      "type": "close",
-      "code": u16::from(frame.code),
-      "reason": frame.reason.as_ref()
+      "kind": "close",
+      "data": {
+        "code": u16::from(frame.code),
+        "reason": frame.reason.as_ref()
+      }
     }),
-    Some(Ok(Message::Close(None))) => json!({ "type": "close" }),
-    Some(Ok(Message::Ping(_))) => json!({"type": "ping"}),
-    Some(Ok(Message::Pong(_))) => json!({"type": "pong"}),
-    Some(Err(_)) => json!({"type": "error"}),
+    Some(Ok(Message::Close(None))) => json!({ "kind": "close" }),
+    Some(Ok(Message::Ping(_))) => json!({ "kind": "ping" }),
+    Some(Ok(Message::Pong(_))) => json!({ "kind": "pong" }),
+    Some(Err(_)) => json!({ "kind": "error" }),
     None => {
       state.borrow_mut().resource_table.close(args.rid).unwrap();
-      json!({"type": "closed"})
+      json!({ "kind": "closed" })
     }
   };
   Ok(res)