From 53c779a8228d54272f4d30f537c4f21e4c3983d6 Mon Sep 17 00:00:00 2001 From: Harshil Agrawal Date: Mon, 21 Oct 2024 16:00:21 +0200 Subject: [PATCH] add serializeAttachment API to Hibernation example --- .../examples/websocket-hibernation-server.mdx | 328 +++++++++++------- 1 file changed, 202 insertions(+), 126 deletions(-) diff --git a/src/content/docs/durable-objects/examples/websocket-hibernation-server.mdx b/src/content/docs/durable-objects/examples/websocket-hibernation-server.mdx index 5ac63fd8184f47..e9670f39fa427d 100644 --- a/src/content/docs/durable-objects/examples/websocket-hibernation-server.mdx +++ b/src/content/docs/durable-objects/examples/websocket-hibernation-server.mdx @@ -10,19 +10,16 @@ sidebar: order: 3 description: Build a WebSocket server using WebSocket Hibernation on Durable Objects and Workers. - --- -import { TabItem, Tabs } from "~/components" +import { TabItem, Tabs } from "~/components"; This example is similar to the [Build a WebSocket server](/durable-objects/examples/websocket-server/) example, but uses the WebSocket Hibernation API. The WebSocket Hibernation API should be preferred for WebSocket server applications built on Durable Objects, since it significantly decreases duration charge, and provides additional features that pair well with WebSocket applications. For more information, refer to [Use Durable Objects with WebSockets](/durable-objects/reference/websockets/). :::note - WebSocket Hibernation is unavailable for outgoing WebSocket use cases. Hibernation is only supported when the Durable Object acts as a server. For use cases where outgoing WebSockets are required, refer to [Write a WebSocket client](/workers/examples/websockets/#write-a-websocket-client). - ::: @@ -32,71 +29,106 @@ import { DurableObject } from "cloudflare:workers"; // Worker export default { - async fetch(request, env, ctx) { - if (request.url.endsWith("/websocket")) { - // Expect to receive a WebSocket Upgrade request. - // If there is one, accept the request and return a WebSocket Response. - const upgradeHeader = request.headers.get('Upgrade'); - if (!upgradeHeader || upgradeHeader !== 'websocket') { - return new Response('Durable Object expected Upgrade: websocket', { status: 426 }); - } - - // This example will refer to the same Durable Object, - // since the name "foo" is hardcoded. - let id = env.WEBSOCKET_HIBERNATION_SERVER.idFromName("foo"); - let stub = env.WEBSOCKET_HIBERNATION_SERVER.get(id); - - return stub.fetch(request); - } - - return new Response(null, { - status: 400, - statusText: 'Bad Request', - headers: { - 'Content-Type': 'text/plain', - }, - }); - } + async fetch(request, env, ctx) { + if (request.url.endsWith("/websocket")) { + // Expect to receive a WebSocket Upgrade request. + // If there is one, accept the request and return a WebSocket Response. + const upgradeHeader = request.headers.get("Upgrade"); + if (!upgradeHeader || upgradeHeader !== "websocket") { + return new Response("Durable Object expected Upgrade: websocket", { + status: 426, + }); + } + + // This example will refer to the same Durable Object, + // since the name "foo" is hardcoded. + let id = env.WEBSOCKET_HIBERNATION_SERVER.idFromName("foo"); + let stub = env.WEBSOCKET_HIBERNATION_SERVER.get(id); + + return stub.fetch(request); + } + + return new Response(null, { + status: 400, + statusText: "Bad Request", + headers: { + "Content-Type": "text/plain", + }, + }); + }, }; // Durable Object export class WebSocketHibernationServer extends DurableObject { - - async fetch(request) { - // Creates two ends of a WebSocket connection. - const webSocketPair = new WebSocketPair(); - const [client, server] = Object.values(webSocketPair); - - // Calling `acceptWebSocket()` informs the runtime that this WebSocket is to begin terminating - // request within the Durable Object. It has the effect of "accepting" the connection, - // and allowing the WebSocket to send and receive messages. - // Unlike `ws.accept()`, `state.acceptWebSocket(ws)` informs the Workers Runtime that the WebSocket - // is "hibernatable", so the runtime does not need to pin this Durable Object to memory while - // the connection is open. During periods of inactivity, the Durable Object can be evicted - // from memory, but the WebSocket connection will remain open. If at some later point the - // WebSocket receives a message, the runtime will recreate the Durable Object - // (run the `constructor`) and deliver the message to the appropriate handler. - this.ctx.acceptWebSocket(server); - - return new Response(null, { - status: 101, - webSocket: client, - }); - } - - async webSocketMessage(ws, message) { - // Upon receiving a message from the client, reply with the same message, - // but will prefix the message with "[Durable Object]: " and return the - // total number of connections. - ws.send(`[Durable Object] message: ${message}, connections: ${this.ctx.getWebSockets().length}`); - } - - async webSocketClose(ws, code, reason, wasClean) { - // If the client closes the connection, the runtime will invoke the webSocketClose() handler. - ws.close(code, "Durable Object is closing WebSocket"); - } + // Keep track of all WebSocket connections + sessions = new Map(); + + async fetch(request) { + // Creates two ends of a WebSocket connection. + const webSocketPair = new WebSocketPair(); + const [client, server] = Object.values(webSocketPair); + + // Calling `acceptWebSocket()` informs the runtime that this WebSocket is to begin terminating + // request within the Durable Object. It has the effect of "accepting" the connection, + // and allowing the WebSocket to send and receive messages. + // Unlike `ws.accept()`, `state.acceptWebSocket(ws)` informs the Workers Runtime that the WebSocket + // is "hibernatable", so the runtime does not need to pin this Durable Object to memory while + // the connection is open. During periods of inactivity, the Durable Object can be evicted + // from memory, but the WebSocket connection will remain open. If at some later point the + // WebSocket receives a message, the runtime will recreate the Durable Object + // (run the `constructor`) and deliver the message to the appropriate handler. + this.ctx.acceptWebSocket(server); + + // Keep a copy of value in memory to survive hibernation. + this.sessions.set(server, {}); + + return new Response(null, { + status: 101, + webSocket: client, + }); + } + + async webSocketMessage(sender, message) { + // Upon receiving a message, get the session associated with the WebSocket connection. + const session = this.sessions.get(sender); + + // If it is a new connection, generate a new ID for the session. + if (!session.id) { + session.id = crypto.randomUUID(); + sender.serializeAttachment({ + ...sender.deserializeAttachment(), + id: session.id, + }); + } + + // Upon receiving a message from the client, the server replies with the same message, + // and the total number of connections with the "[Durable Object]: " prefix + sender.send( + `[Durable Object] message: ${message}, from: ${session.id}. Total connections: ${this.ctx.getWebSockets().length}`, + ); + + // Send a message to all WebSocket connections, loop over all the connected WebSockets. + this.ctx.getWebSockets().forEach((ws) => { + ws.send( + `[Durable Object] message: ${message}, from: ${session.id}. Total connections: ${this.ctx.getWebSockets().length}`, + ); + }); + + // Send a message to all WebSocket connections except the sender, loop over all the connected WebSockets and filter out the sender. + this.ctx.getWebSockets().forEach((ws) => { + if (ws !== sender) { + ws.send( + `[Durable Object] message: ${message}, from: ${session.id}. Total connections: ${this.ctx.getWebSockets().length}`, + ); + } + }); + } + + async webSocketClose(ws, code, reason, wasClean) { + // If the client closes the connection, the runtime will invoke the webSocketClose() handler. + ws.close(code, "Durable Object is closing WebSocket"); + } } - ``` @@ -104,74 +136,118 @@ export class WebSocketHibernationServer extends DurableObject { ```ts import { DurableObject } from "cloudflare:workers"; -export interface Env { - WEBSOCKET_HIBERNATION_SERVER: DurableObjectNamespace; -} +// Use npm run cf-typegen to generate the type definitions for the Durable Object // Worker export default { - async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise { - if (request.url.endsWith("/websocket")) { - // Expect to receive a WebSocket Upgrade request. - // If there is one, accept the request and return a WebSocket Response. - const upgradeHeader = request.headers.get('Upgrade'); - if (!upgradeHeader || upgradeHeader !== 'websocket') { - return new Response('Durable Object expected Upgrade: websocket', { status: 426 }); - } - - // This example will refer to the same Durable Object, - // since the name "foo" is hardcoded. - let id = env.WEBSOCKET_HIBERNATION_SERVER.idFromName("foo"); - let stub = env.WEBSOCKET_HIBERNATION_SERVER.get(id); - - return stub.fetch(request); - } - - return new Response(null, { - status: 400, - statusText: 'Bad Request', - headers: { - 'Content-Type': 'text/plain', - }, - }); - } + async fetch( + request: Request, + env: Env, + ctx: ExecutionContext, + ): Promise { + if (request.url.endsWith("/websocket")) { + // Expect to receive a WebSocket Upgrade request. + // If there is one, accept the request and return a WebSocket Response. + const upgradeHeader = request.headers.get("Upgrade"); + if (!upgradeHeader || upgradeHeader !== "websocket") { + return new Response("Durable Object expected Upgrade: websocket", { + status: 426, + }); + } + + // This example will refer to the same Durable Object, + // since the name "foo" is hardcoded. + let id = env.WEBSOCKET_HIBERNATION_SERVER.idFromName("foo"); + let stub = env.WEBSOCKET_HIBERNATION_SERVER.get(id); + + return stub.fetch(request); + } + + return new Response(null, { + status: 400, + statusText: "Bad Request", + headers: { + "Content-Type": "text/plain", + }, + }); + }, }; // Durable Object export class WebSocketHibernationServer extends DurableObject { - - async fetch(request: Request): Promise { - // Creates two ends of a WebSocket connection. - const webSocketPair = new WebSocketPair(); - const [client, server] = Object.values(webSocketPair); - - // Calling `acceptWebSocket()` informs the runtime that this WebSocket is to begin terminating - // request within the Durable Object. It has the effect of "accepting" the connection, - // and allowing the WebSocket to send and receive messages. - // Unlike `ws.accept()`, `state.acceptWebSocket(ws)` informs the Workers Runtime that the WebSocket - // is "hibernatable", so the runtime does not need to pin this Durable Object to memory while - // the connection is open. During periods of inactivity, the Durable Object can be evicted - // from memory, but the WebSocket connection will remain open. If at some later point the - // WebSocket receives a message, the runtime will recreate the Durable Object - // (run the `constructor`) and deliver the message to the appropriate handler. - this.ctx.acceptWebSocket(server); - - return new Response(null, { - status: 101, - webSocket: client, - }); - } - - async webSocketMessage(ws: WebSocket, message: ArrayBuffer | string) { - // Upon receiving a message from the client, the server replies with the same message, - // and the total number of connections with the "[Durable Object]: " prefix - ws.send(`[Durable Object] message: ${message}, connections: ${this.ctx.getWebSockets().length}`); - } - - async webSocketClose(ws: WebSocket, code: number, reason: string, wasClean: boolean) { - // If the client closes the connection, the runtime will invoke the webSocketClose() handler. - ws.close(code, "Durable Object is closing WebSocket"); - } + // Keep track of all WebSocket connections + sessions = new Map(); + + async fetch(request: Request): Promise { + // Creates two ends of a WebSocket connection. + const webSocketPair = new WebSocketPair(); + const [client, server] = Object.values(webSocketPair); + + // Calling `acceptWebSocket()` informs the runtime that this WebSocket is to begin terminating + // request within the Durable Object. It has the effect of "accepting" the connection, + // and allowing the WebSocket to send and receive messages. + // Unlike `ws.accept()`, `state.acceptWebSocket(ws)` informs the Workers Runtime that the WebSocket + // is "hibernatable", so the runtime does not need to pin this Durable Object to memory while + // the connection is open. During periods of inactivity, the Durable Object can be evicted + // from memory, but the WebSocket connection will remain open. If at some later point the + // WebSocket receives a message, the runtime will recreate the Durable Object + // (run the `constructor`) and deliver the message to the appropriate handler. + this.ctx.acceptWebSocket(server); + + // Keep a copy of value in memory to survive hibernation. + this.sessions.set(server, {}); + + return new Response(null, { + status: 101, + webSocket: client, + }); + } + + async webSocketMessage(sender: WebSocket, message: ArrayBuffer | string) { + // Upon receiving a message, get the session associated with the WebSocket connection. + const session = this.sessions.get(sender); + + // If it is a new connection, generate a new ID for the session. + if (!session.id) { + session.id = crypto.randomUUID(); + sender.serializeAttachment({ + ...sender.deserializeAttachment(), + id: session.id, + }); + } + + // Upon receiving a message from the client, the server replies with the same message, + // and the total number of connections with the "[Durable Object]: " prefix + sender.send( + `[Durable Object] message: ${message}, from: ${session.id}. Total connections: ${this.ctx.getWebSockets().length}`, + ); + + // Send a message to all WebSocket connections, loop over all the connected WebSockets. + this.ctx.getWebSockets().forEach((ws) => { + ws.send( + `[Durable Object] message: ${message}, from: ${session.id}. Total connections: ${this.ctx.getWebSockets().length}`, + ); + }); + + // Send a message to all WebSocket connections except the sender, loop over all the connected WebSockets and filter out the sender. + this.ctx.getWebSockets().forEach((ws) => { + if (ws !== sender) { + ws.send( + `[Durable Object] message: ${message}, from: ${session.id}. Total connections: ${this.ctx.getWebSockets().length}`, + ); + } + }); + } + + async webSocketClose( + ws: WebSocket, + code: number, + reason: string, + wasClean: boolean, + ) { + // If the client closes the connection, the runtime will invoke the webSocketClose() handler. + ws.close(code, "Durable Object is closing WebSocket"); + } } ``` @@ -193,4 +269,4 @@ new_classes = ["WebSocketHibernationServer"] ### Related resources -* [Durable Objects: Edge Chat Demo with Hibernation](https://github.com/cloudflare/workers-chat-demo/). +- [Durable Objects: Edge Chat Demo with Hibernation](https://github.com/cloudflare/workers-chat-demo/).