add tests to cover connecting to durable object via websocket using the hibernatable apis (#495)

Author: j-white

Cargo.lock

@@ -37,6 +37,7 @@ rand = "0.8.5"
 uuid = { version = "1.3.3", features = ["v4", "serde"] }
 serde-wasm-bindgen = "0.6.1"
 md5 = "0.7.0"
+tokio-stream = "0.1.15"
 
 [dependencies.axum]
 version = "0.7"

worker-sandbox/Cargo.toml

@@ -1,4 +1,5 @@
 use std::time::Duration;
+use tokio_stream::{StreamExt, StreamMap};
 
 use worker::*;
 
@@ -69,3 +70,60 @@ pub async fn handle_put_raw(req: Request, env: Env, _data: SomeSharedData) -> Re
     let stub = id.get_stub()?;
     stub.fetch_with_request(req).await
 }
+
+#[worker::send]
+pub async fn handle_websocket(_req: Request, env: Env, _data: SomeSharedData) -> Result<Response> {
+    // Accept / handle a websocket connection
+    let pair = WebSocketPair::new()?;
+    let server = pair.server;
+    server.accept()?;
+
+    // Connect to Durable Object via WS
+    let namespace = env
+        .durable_object("COUNTER")
+        .expect("failed to get namespace");
+    let stub = namespace.id_from_name("A")?.get_stub()?;
+    let mut req = Request::new("https://fake-host/ws", Method::Get)?;
+    req.headers_mut()?.set("upgrade", "websocket")?;
+
+    let res = stub.fetch_with_request(req).await?;
+    let do_ws = res.websocket().expect("server did not accept websocket");
+    do_ws.accept()?;
+
+    wasm_bindgen_futures::spawn_local(async move {
+        let event_stream = server.events().expect("could not open stream");
+        let do_event_stream = do_ws.events().expect("could not open stream");
+
+        let mut map = StreamMap::new();
+        map.insert("client", event_stream);
+        map.insert("durable", do_event_stream);
+
+        while let Some((key, event)) = map.next().await {
+            match key {
+                "client" => match event.expect("received error in websocket") {
+                    WebsocketEvent::Message(msg) => {
+                        if let Some(text) = msg.text() {
+                            do_ws.send_with_str(text).expect("could not relay text");
+                        }
+                    }
+                    WebsocketEvent::Close(_) => {
+                        let _res = do_ws.close(Some(1000), Some("client closed".to_string()));
+                    }
+                },
+                "durable" => match event.expect("received error in websocket") {
+                    WebsocketEvent::Message(msg) => {
+                        if let Some(text) = msg.text() {
+                            server.send_with_str(text).expect("could not relay text");
+                        }
+                    }
+                    WebsocketEvent::Close(_) => {
+                        let _res = server.close(Some(1000), Some("durable closed".to_string()));
+                    }
+                },
+                _ => unreachable!(),
+            }
+        }
+    });
+
+    Response::from_websocket(pair.client)
+}

worker-sandbox/src/alarm.rs

@@ -19,12 +19,27 @@ impl DurableObject for Counter {
         }
     }
 
-    async fn fetch(&mut self, _req: Request) -> Result<Response> {
+    async fn fetch(&mut self, req: Request) -> Result<Response> {
         if !self.initialized {
             self.initialized = true;
             self.count = self.state.storage().get("count").await.unwrap_or(0);
         }
 
+        if req.path().eq("/ws") {
+            let pair = WebSocketPair::new()?;
+            let server = pair.server;
+            // accept websocket with hibernation api
+            self.state.accept_web_socket(&server);
+            server
+                .serialize_attachment("hello")
+                .expect("failed to serialize attachment");
+
+            return Ok(Response::empty()
+                .unwrap()
+                .with_status(101)
+                .with_websocket(Some(pair.client)));
+        }
+
         self.count += 10;
         self.state.storage().put("count", self.count).await?;
 
@@ -34,4 +49,36 @@ impl DurableObject for Counter {
             self.env.secret("SOME_SECRET")?.to_string()
         ))
     }
+
+    async fn websocket_message(
+        &mut self,
+        ws: WebSocket,
+        _message: WebSocketIncomingMessage,
+    ) -> Result<()> {
+        let _attach: String = ws
+            .deserialize_attachment()?
+            .expect("websockets should have an attachment");
+        // get and increment storage by 10
+        let mut count: usize = self.state.storage().get("count").await.unwrap_or(0);
+        count += 10;
+        self.state.storage().put("count", count).await?;
+        // send value to client
+        ws.send_with_str(format!("{}", count))
+            .expect("failed to send value to client");
+        Ok(())
+    }
+
+    async fn websocket_close(
+        &mut self,
+        _ws: WebSocket,
+        _code: usize,
+        _reason: String,
+        _was_clean: bool,
+    ) -> Result<()> {
+        Ok(())
+    }
+
+    async fn websocket_error(&mut self, _ws: WebSocket, _error: Error) -> Result<()> {
+        Ok(())
+    }
 }

worker-sandbox/src/counter.rs

@@ -107,6 +107,7 @@ pub fn make_router(data: SomeSharedData, env: Env) -> axum::Router {
         .route("/durable/alarm", get(handler!(alarm::handle_alarm)))
         .route("/durable/:id", get(handler!(alarm::handle_id)))
         .route("/durable/put-raw", get(handler!(alarm::handle_put_raw)))
+        .route("/durable/websocket", get(handler!(alarm::handle_websocket)))
         .route("/var", get(handler!(request::handle_var)))
         .route("/secret", get(handler!(request::handle_secret)))
         .route("/kv/:key/:value", post(handler!(kv::handle_post_key_value)))
@@ -241,6 +242,7 @@ pub fn make_router<'a>(data: SomeSharedData) -> Router<'a, SomeSharedData> {
         .get_async("/durable/alarm", handler!(alarm::handle_alarm))
         .get_async("/durable/:id", handler!(alarm::handle_id))
         .get_async("/durable/put-raw", handler!(alarm::handle_put_raw))
+        .get_async("/durable/websocket", handler!(alarm::handle_websocket))
         .get_async("/secret", handler!(request::handle_secret))
         .get_async("/var", handler!(request::handle_var))
         .post_async("/kv/:key/:value", handler!(kv::handle_post_key_value))

worker-sandbox/src/router.rs

@@ -1,9 +1,46 @@
-import { describe, test, expect } from "vitest";
+import {describe, test, expect, vi} from "vitest";
 import { mf } from "./mf";
+import {MessageEvent} from "miniflare";
 
 describe("durable", () => {
   test("put-raw", async () => {
     const resp = await mf.dispatchFetch("https://fake.host/durable/put-raw");
     expect(await resp.text()).toBe("ok");
   });
+
+  test("websocket-to-durable", async () => {
+    const resp = await mf.dispatchFetch("http://fake.host/durable/websocket", {
+      headers: {
+        upgrade: "websocket",
+      },
+    });
+    expect(resp.webSocket).not.toBeNull();
+
+    const socket = resp.webSocket!;
+    socket.accept();
+
+    const handlers = {
+      messageHandler: (event: MessageEvent) => {
+        expect(event.data).toMatch(/^10|20|30$/);
+      },
+      close(event: CloseEvent) {},
+    };
+
+    const messageHandlerWrapper = vi.spyOn(handlers, "messageHandler");
+    const closeHandlerWrapper = vi.spyOn(handlers, "messageHandler");
+    socket.addEventListener("message", handlers.messageHandler);
+    socket.addEventListener("close", handlers.close);
+
+    socket.send("hi, can you ++?");
+    await new Promise((resolve) => setTimeout(resolve, 500));
+    expect(messageHandlerWrapper).toHaveBeenCalledTimes(1);
+
+    socket.send("hi again, more ++?");
+    await new Promise((resolve) => setTimeout(resolve, 500));
+    expect(messageHandlerWrapper).toHaveBeenCalledTimes(2);
+
+    socket.close();
+    expect(closeHandlerWrapper).toBeCalled();
+  });
 });
+