Skip to content

Commit 845ee7b

Browse files
OxyjunOxyjun
authored
[DO] Updating code for websocket hibernation (#24234)
* Updating code for websocket hibernation * Changing 4-space indentation to 2-space indentation
1 parent 67a10b9 commit 845ee7b

File tree

1 file changed

+121
-75
lines changed

1 file changed

+121
-75
lines changed

‎src/content/docs/durable-objects/examples/websocket-hibernation-server.mdx

Lines changed: 121 additions & 75 deletions
Original file line numberDiff line numberDiff line change
@@ -24,88 +24,134 @@ WebSocket Hibernation is unavailable for outgoing WebSocket use cases. Hibernati
2424

2525
<TypeScriptExample>
2626
```ts
27-
import { DurableObject } from "cloudflare:workers";
28-
29-
export interface Env {
30-
WEBSOCKET_HIBERNATION_SERVER: DurableObjectNamespace<WebSocketHibernationServer>;
31-
}
27+
import { DurableObject } from 'cloudflare:workers';
3228

3329
// Worker
3430
export default {
35-
async fetch(
36-
request: Request,
37-
env: Env,
38-
ctx: ExecutionContext,
39-
): Promise<Response> {
40-
if (request.url.endsWith("/websocket")) {
41-
// Expect to receive a WebSocket Upgrade request.
42-
// If there is one, accept the request and return a WebSocket Response.
43-
const upgradeHeader = request.headers.get("Upgrade");
44-
if (!upgradeHeader || upgradeHeader !== "websocket") {
45-
return new Response("Durable Object expected Upgrade: websocket", {
46-
status: 426,
47-
});
48-
}
49-
50-
// This example will refer to the same Durable Object,
51-
// since the name "foo" is hardcoded.
52-
let id = env.WEBSOCKET_HIBERNATION_SERVER.idFromName("foo");
53-
let stub = env.WEBSOCKET_HIBERNATION_SERVER.get(id);
54-
55-
return stub.fetch(request);
56-
}
57-
58-
return new Response(null, {
59-
status: 400,
60-
statusText: "Bad Request",
61-
headers: {
62-
"Content-Type": "text/plain",
63-
},
64-
});
65-
},
31+
async fetch(request: Request, env: Env, ctx: ExecutionContext): Promise<Response> {
32+
if (request.url.endsWith('/websocket')) {
33+
// Expect to receive a WebSocket Upgrade request.
34+
// If there is one, accept the request and return a WebSocket Response.
35+
const upgradeHeader = request.headers.get('Upgrade');
36+
if (!upgradeHeader || upgradeHeader !== 'websocket') {
37+
return new Response('Worker expected Upgrade: websocket', {
38+
status: 426,
39+
});
40+
}
41+
42+
if (request.method !== 'GET') {
43+
return new Response('Worker expected GET method', {
44+
status: 400,
45+
});
46+
}
47+
48+
// Since we are hard coding the Durable Object ID by providing the constant name 'foo',
49+
// all requests to this Worker will be sent to the same Durable Object instance.
50+
let id = env.WEBSOCKET_HIBERNATION_SERVER.idFromName('foo');
51+
let stub = env.WEBSOCKET_HIBERNATION_SERVER.get(id);
52+
53+
return stub.fetch(request);
54+
}
55+
56+
return new Response(
57+
`Supported endpoints:
58+
/websocket: Expects a WebSocket upgrade request`,
59+
{
60+
status: 200,
61+
headers: {
62+
'Content-Type': 'text/plain',
63+
},
64+
}
65+
);
66+
},
6667
};
6768

6869
// Durable Object
6970
export class WebSocketHibernationServer extends DurableObject {
70-
async fetch(request: Request): Promise<Response> {
71-
// Creates two ends of a WebSocket connection.
72-
const webSocketPair = new WebSocketPair();
73-
const [client, server] = Object.values(webSocketPair);
74-
75-
// Calling `acceptWebSocket()` informs the runtime that this WebSocket is to begin terminating
76-
// request within the Durable Object. It has the effect of "accepting" the connection,
77-
// and allowing the WebSocket to send and receive messages.
78-
// Unlike `ws.accept()`, `state.acceptWebSocket(ws)` informs the Workers Runtime that the WebSocket
79-
// is "hibernatable", so the runtime does not need to pin this Durable Object to memory while
80-
// the connection is open. During periods of inactivity, the Durable Object can be evicted
81-
// from memory, but the WebSocket connection will remain open. If at some later point the
82-
// WebSocket receives a message, the runtime will recreate the Durable Object
83-
// (run the `constructor`) and deliver the message to the appropriate handler.
84-
this.ctx.acceptWebSocket(server);
85-
86-
return new Response(null, {
87-
status: 101,
88-
webSocket: client,
89-
});
90-
}
91-
92-
async webSocketMessage(ws: WebSocket, message: ArrayBuffer | string) {
93-
// Upon receiving a message from the client, the server replies with the same message,
94-
// and the total number of connections with the "[Durable Object]: " prefix
95-
ws.send(
96-
`[Durable Object] message: ${message}, connections: ${this.ctx.getWebSockets().length}`,
97-
);
98-
}
99-
100-
async webSocketClose(
101-
ws: WebSocket,
102-
code: number,
103-
reason: string,
104-
wasClean: boolean,
105-
) {
106-
// If the client closes the connection, the runtime will invoke the webSocketClose() handler.
107-
ws.close(code, "Durable Object is closing WebSocket");
108-
}
71+
// Keeps track of all WebSocket connections
72+
// When the DO hibernates, gets reconstructed in the constructor
73+
sessions: Map<WebSocket, { [key: string]: string }>;
74+
75+
constructor(ctx: DurableObjectState, env: Env) {
76+
super(ctx, env);
77+
this.sessions = new Map();
78+
79+
// As part of constructing the Durable Object,
80+
// we wake up any hibernating WebSockets and
81+
// place them back in the `sessions` map.
82+
83+
// Get all WebSocket connections from the DO
84+
this.ctx.getWebSockets().forEach((ws) => {
85+
let attachment = ws.deserializeAttachment();
86+
if (attachment) {
87+
// If we previously attached state to our WebSocket,
88+
// let's add it to `sessions` map to restore the state of the connection.
89+
this.sessions.set(ws, { ...attachment });
90+
}
91+
});
92+
93+
// Sets an application level auto response that does not wake hibernated WebSockets.
94+
this.ctx.setWebSocketAutoResponse(new WebSocketRequestResponsePair('ping', 'pong'));
95+
}
96+
97+
async fetch(request: Request): Promise<Response> {
98+
// Creates two ends of a WebSocket connection.
99+
const webSocketPair = new WebSocketPair();
100+
const [client, server] = Object.values(webSocketPair);
101+
102+
// Calling `acceptWebSocket()` informs the runtime that this WebSocket is to begin terminating
103+
// request within the Durable Object. It has the effect of "accepting" the connection,
104+
// and allowing the WebSocket to send and receive messages.
105+
// Unlike `ws.accept()`, `this.ctx.acceptWebSocket(ws)` informs the Workers Runtime that the WebSocket
106+
// is "hibernatable", so the runtime does not need to pin this Durable Object to memory while
107+
// the connection is open. During periods of inactivity, the Durable Object can be evicted
108+
// from memory, but the WebSocket connection will remain open. If at some later point the
109+
// WebSocket receives a message, the runtime will recreate the Durable Object
110+
// (run the `constructor`) and deliver the message to the appropriate handler.
111+
this.ctx.acceptWebSocket(server);
112+
113+
// Generate a random UUID for the session.
114+
const id = crypto.randomUUID();
115+
116+
// Attach the session ID to the WebSocket connection and serialize it.
117+
// This is necessary to restore the state of the connection when the Durable Object wakes up.
118+
server.serializeAttachment({ id });
119+
120+
// Add the WebSocket connection to the map of active sessions.
121+
this.sessions.set(server, { id });
122+
123+
return new Response(null, {
124+
status: 101,
125+
webSocket: client,
126+
});
127+
}
128+
129+
async webSocketMessage(ws: WebSocket, message: ArrayBuffer | string) {
130+
// Get the session associated with the WebSocket connection.
131+
const session = this.sessions.get(ws)!;
132+
133+
// Upon receiving a message from the client, the server replies with the same message, the session ID of the connection,
134+
// and the total number of connections with the "[Durable Object]: " prefix
135+
ws.send(`[Durable Object] message: ${message}, from: ${session.id}. Total connections: ${this.sessions.size}`);
136+
137+
// Send a message to all WebSocket connections, loop over all the connected WebSockets.
138+
this.sessions.forEach((attachment, session) => {
139+
session.send(`[Durable Object] message: ${message}, from: ${attachment.id}. Total connections: ${this.sessions.size}`);
140+
});
141+
142+
// Send a message to all WebSocket connections except the connection (ws),
143+
// loop over all the connected WebSockets and filter out the connection (ws).
144+
this.sessions.forEach((attachment, session) => {
145+
if (session !== ws) {
146+
session.send(`[Durable Object] message: ${message}, from: ${attachment.id}. Total connections: ${this.sessions.size}`);
147+
}
148+
});
149+
}
150+
151+
async webSocketClose(ws: WebSocket, code: number, reason: string, wasClean: boolean) {
152+
// If the client closes the connection, the runtime will invoke the webSocketClose() handler.
153+
ws.close(code, 'Durable Object is closing WebSocket');
154+
}
109155
}
110156
```
111157
</TypeScriptExample>

0 commit comments

Comments
 (0)